issue/497 - simplified infinicore test functions

Author: wooway777

test/infinicore/framework/__init__.py

@@ -2,10 +2,12 @@
 from .utils import (
     create_infinicore_tensor,
     compare_results,
+    create_test_comparator,
     debug,
     get_tolerance,
     profile_operation,
     rearrange_tensor,
+    convert_infinicore_to_torch,
 )
 from .config import get_test_devices, get_args
 from .devices import InfiniDeviceEnum, InfiniDeviceNames, torch_device_map
@@ -17,6 +19,8 @@
     "TestCase",
     "create_infinicore_tensor",
     "compare_results",
+    "create_test_comparator",
+    "convert_infinicore_to_torch",
     "debug",
     "get_tolerance",
     "profile_operation",

test/infinicore/framework/utils.py

@@ -4,18 +4,6 @@
 from .datatypes import to_infinicore_dtype, to_torch_dtype
 
 
-def create_infinicore_tensor(torch_tensor, device_str):
-    """Create infinicore tensor from PyTorch tensor"""
-    infini_device = infinicore.device(device_str, 0)
-
-    return infinicore.from_blob(
-        torch_tensor.data_ptr(),
-        list(torch_tensor.shape),
-        dtype=to_infinicore_dtype(torch_tensor.dtype),
-        device=infini_device,
-    )
-
-
 def synchronize_device(torch_device):
     """Device synchronization"""
     if torch_device == "cuda":
@@ -149,44 +137,95 @@ def get_tolerance(tolerance_map, tensor_dtype, default_atol=0, default_rtol=1e-3
     return tolerance["atol"], tolerance["rtol"]
 
 
-def compare_results(
-    infini_result, torch_result, dtype, config, device_str, tolerance_map=None
-):
+def create_infinicore_tensor(torch_tensor, device_str):
+    """Create infinicore tensor from PyTorch tensor"""
+    infini_device = infinicore.device(device_str, 0)
+
+    return infinicore.from_blob(
+        torch_tensor.data_ptr(),
+        list(torch_tensor.shape),
+        dtype=to_infinicore_dtype(torch_tensor.dtype),
+        device=infini_device,
+    )
+
+
+def convert_infinicore_to_torch(infini_result, torch_reference):
     """
-    Compare infinicore result with PyTorch reference result
+    Convert infinicore tensor to PyTorch tensor for comparison
 
     Args:
         infini_result: infinicore tensor result
-        torch_result: PyTorch tensor reference result
+        torch_reference: PyTorch tensor reference (for shape and device)
         dtype: infinicore data type
-        config: test config
         device_str: torch device string
-        device: device enum
-        tolerance_map: optional tolerance map (defaults to config's tolerance_map)
 
     Returns:
-        bool: True if results match within tolerance
+        torch.Tensor: PyTorch tensor with infinicore data
     """
-    # Convert infinicore result to PyTorch tensor for comparison
     torch_result_from_infini = torch.zeros(
-        torch_result.shape, dtype=to_torch_dtype(dtype), device=device_str
+        torch_reference.shape,
+        dtype=to_torch_dtype(infini_result.dtype),
+        device=infini_result.device.type,
+    )
+    temp_tensor = create_infinicore_tensor(
+        torch_result_from_infini, infini_result.device.type
     )
-    temp_tensor = create_infinicore_tensor(torch_result_from_infini, device_str)
     temp_tensor.copy_(infini_result)
+    return torch_result_from_infini
 
-    # Retrieve tolerance - use provided map or config's map
-    if tolerance_map is None:
-        tolerance_map = config.tolerance_map
-    atol, rtol = get_tolerance(tolerance_map, dtype)
+
+def compare_results(
+    infini_result, torch_result, atol=1e-5, rtol=1e-5, debug_mode=False
+):
+    """
+    Generic function to compare infinicore result with PyTorch reference result
+
+    Args:
+        infini_result: infinicore tensor result
+        torch_result: PyTorch tensor reference result
+        atol: absolute tolerance
+        rtol: relative tolerance
+        debug_mode: whether to enable debug output
+
+    Returns:
+        bool: True if results match within tolerance
+    """
+    # Convert infinicore result to PyTorch tensor for comparison
+    torch_result_from_infini = convert_infinicore_to_torch(infini_result, torch_result)
 
     # Debug mode: detailed comparison
-    if config.debug:
+    if debug_mode:
         debug(torch_result_from_infini, torch_result, atol=atol, rtol=rtol)
 
     # Check if results match within tolerance
     return torch.allclose(torch_result_from_infini, torch_result, atol=atol, rtol=rtol)
 
 
+def create_test_comparator(config, dtype, tolerance_map=None):
+    """
+    Create a test-specific comparison function that handles test configuration
+
+    Args:
+        config: test configuration
+        dtype: infinicore data type
+        tolerance_map: optional tolerance map (defaults to config's tolerance_map)
+
+    Returns:
+        callable: function that takes (infini_result, torch_result) and returns bool
+    """
+    if tolerance_map is None:
+        tolerance_map = config.tolerance_map
+
+    atol, rtol = get_tolerance(tolerance_map, dtype)
+
+    def compare_test_results(infini_result, torch_result):
+        return compare_results(
+            infini_result, torch_result, atol=atol, rtol=rtol, debug_mode=config.debug
+        )
+
+    return compare_test_results
+
+
 def rearrange_tensor(tensor, new_strides):
     """
     Given a PyTorch tensor and a list of new strides, return a new PyTorch tensor with the given strides.

test/infinicore/op/matmul.py

@@ -11,7 +11,7 @@
     TestRunner,
     TestCase,
     create_infinicore_tensor,
-    compare_results,
+    create_test_comparator,
     get_args,
     get_test_devices,
     profile_operation,
@@ -92,8 +92,11 @@ def infini_matmul():
 
     infini_result = infini_matmul()
 
-    # Validate results using common method
-    is_valid = compare_results(infini_result, torch_result, dtype, config, device_str)
+    # Create test-specific comparator
+    compare_fn = create_test_comparator(config, dtype)
+
+    # Validate results using the test-specific comparator
+    is_valid = compare_fn(infini_result, torch_result)
     assert is_valid, "Matmul test failed"
 
     # Performance test
@@ -163,8 +166,11 @@ def infini_matmul_inplace():
     # Execute in-place operation
     infini_matmul_inplace()
 
-    # Validate results using common method
-    is_valid = compare_results(infini_c, torch_preallocated, dtype, config, device_str)
+    # Create test-specific comparator
+    compare_fn = create_test_comparator(config, dtype)
+
+    # Validate results using the test-specific comparator
+    is_valid = compare_fn(infini_c, torch_preallocated)
     assert is_valid, "In-place matmul test failed"
 
     # Performance test