issue/461 - slightly simplified infinicore test framework

Author: wooway777

test/infinicore/framework/base.py

@@ -1,6 +1,6 @@
 import torch
 import infinicore
-from .devices import InfiniDeviceNames, torch_device_map
+from .devices import InfiniDeviceNames
 from .utils import synchronize_device
 
 

test/infinicore/framework/utils.py

@@ -2,12 +2,10 @@
 import time
 import infinicore
 from .datatypes import to_infinicore_dtype, to_torch_dtype
-from .devices import torch_device_map
 
 
-def create_infinicore_tensor(torch_tensor, device_enum):
+def create_infinicore_tensor(torch_tensor, device_str):
     """Create infinicore tensor from PyTorch tensor"""
-    device_str = torch_device_map[device_enum]
     infini_device = infinicore.device(device_str, 0)
 
     return infinicore.from_blob(
@@ -152,7 +150,7 @@ def get_tolerance(tolerance_map, tensor_dtype, default_atol=0, default_rtol=1e-3
 
 
 def compare_results(
-    infini_result, torch_result, dtype, config, device_str, device, tolerance_map=None
+    infini_result, torch_result, dtype, config, device_str, tolerance_map=None
 ):
     """
     Compare infinicore result with PyTorch reference result
@@ -173,7 +171,7 @@ def compare_results(
     torch_result_from_infini = torch.zeros(
         torch_result.shape, dtype=to_torch_dtype(dtype), device=device_str
     )
-    temp_tensor = create_infinicore_tensor(torch_result_from_infini, device)
+    temp_tensor = create_infinicore_tensor(torch_result_from_infini, device_str)
     temp_tensor.copy_(infini_result)
 
     # Retrieve tolerance - use provided map or config's map

test/infinicore/op/matmul.py

@@ -83,8 +83,8 @@ def torch_matmul():
     torch_result = torch_matmul()
 
     # Create infinicore tensors
-    infini_a = create_infinicore_tensor(torch_a, device)
-    infini_b = create_infinicore_tensor(torch_b, device)
+    infini_a = create_infinicore_tensor(torch_a, device_str)
+    infini_b = create_infinicore_tensor(torch_b, device_str)
 
     # Out-of-place matmul
     def infini_matmul():
@@ -93,9 +93,7 @@ def infini_matmul():
     infini_result = infini_matmul()
 
     # Validate results using common method
-    is_valid = compare_results(
-        infini_result, torch_result, dtype, config, device_str, device
-    )
+    is_valid = compare_results(infini_result, torch_result, dtype, config, device_str)
     assert is_valid, "Matmul test failed"
 
     # Performance test
@@ -152,8 +150,8 @@ def torch_matmul_inplace():
     torch_matmul_inplace()
 
     # Create infinicore tensors
-    infini_a = create_infinicore_tensor(torch_a, device)
-    infini_b = create_infinicore_tensor(torch_b, device)
+    infini_a = create_infinicore_tensor(torch_a, device_str)
+    infini_b = create_infinicore_tensor(torch_b, device_str)
     infini_c = infinicore.empty(
         result_shape, dtype=dtype, device=infinicore.device(device_str, 0)
     )
@@ -166,9 +164,7 @@ def infini_matmul_inplace():
     infini_matmul_inplace()
 
     # Validate results using common method
-    is_valid = compare_results(
-        infini_c, torch_preallocated, dtype, config, device_str, device
-    )
+    is_valid = compare_results(infini_c, torch_preallocated, dtype, config, device_str)
     assert is_valid, "In-place matmul test failed"
 
     # Performance test