issue/497 - simplified infinicore tensor creation from torch

Author: wooway777

test/infinicore/framework/__init__.py

@@ -1,11 +1,10 @@
 from .base import TensorSpec, TestConfig, TestRunner, TestCase, BaseOperatorTest
 from .utils import (
-    create_infinicore_tensor,
-    create_strided_infinicore_tensor,
     compare_results,
     create_test_comparator,
     debug,
     get_tolerance,
+    infinicore_tensor_from_torch,
     profile_operation,
     rearrange_tensor,
     convert_infinicore_to_torch,
@@ -24,17 +23,16 @@
     "BaseOperatorTest",
     "ParameterMapping",
     "create_test_cases",
-    "create_infinicore_tensor",
-    "create_strided_infinicore_tensor",
     "compare_results",
     "create_test_comparator",
     "convert_infinicore_to_torch",
     "debug",
+    "get_args",
+    "get_test_devices",
     "get_tolerance",
+    "infinicore_tensor_from_torch",
     "profile_operation",
     "rearrange_tensor",
-    "get_test_devices",
-    "get_args",
     "InfiniDeviceEnum",
     "InfiniDeviceNames",
     "torch_device_map",

test/infinicore/framework/base.py

@@ -7,9 +7,8 @@
 from .datatypes import to_torch_dtype, to_infinicore_dtype
 from .devices import InfiniDeviceNames, torch_device_map
 from .utils import (
-    create_infinicore_tensor,
-    create_strided_infinicore_tensor,
     create_test_comparator,
+    infinicore_tensor_from_torch,
     profile_operation,
     rearrange_tensor,
     synchronize_device,
@@ -356,16 +355,7 @@ def _run_single_test(self, device, test_case, dtype_config, config, mode_name):
         infini_inputs = []
         for inp in inputs:
             if isinstance(inp, torch.Tensor):
-                if not inp.is_contiguous():
-                    infini_tensor = infinicore.strided_from_blob(
-                        inp.data_ptr(),
-                        list(inp.shape),
-                        list(inp.stride()),
-                        dtype=to_infinicore_dtype(inp.dtype),
-                        device=infinicore.device(device_str, 0),
-                    )
-                else:
-                    infini_tensor = create_infinicore_tensor(inp, device_str)
+                infini_tensor = infinicore_tensor_from_torch(inp)
                 infini_inputs.append(infini_tensor)
             else:
                 infini_inputs.append(inp)
@@ -439,13 +429,12 @@ def torch_op_inplace():
             torch_dummy = torch.zeros(
                 output_shape, dtype=output_dtype, device=device_str
             )
-            if test_case.output.is_contiguous or test_case.output.strides is None:
-                infini_output = create_infinicore_tensor(torch_dummy, device_str)
-            else:
+            if (
+                not test_case.output.is_contiguous
+                and not test_case.output.strides is None
+            ):
                 rearrange_tensor(torch_dummy, list(torch_output.stride()))
-                infini_output = create_strided_infinicore_tensor(
-                    torch_dummy, device_str
-                )
+            infini_output = infinicore_tensor_from_torch(torch_dummy)
 
             def infini_op_inplace():
                 self.infinicore_operator(*infini_inputs, out=infini_output, **kwargs)

test/infinicore/framework/utils.py

@@ -136,29 +136,23 @@ def get_tolerance(tolerance_map, tensor_dtype, default_atol=0, default_rtol=1e-3
     return tolerance["atol"], tolerance["rtol"]
 
 
-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 create_strided_infinicore_tensor(torch_tensor, device_str):
-    """Create strided infinicore tensor from PyTorch tensor"""
-    infini_device = infinicore.device(device_str, 0)
-
-    return infinicore.strided_from_blob(
-        torch_tensor.data_ptr(),
-        list(torch_tensor.shape),
-        list(torch_tensor.stride()),
-        dtype=to_infinicore_dtype(torch_tensor.dtype),
-        device=infini_device,
-    )
+def infinicore_tensor_from_torch(torch_tensor):
+    infini_device = infinicore.device(torch_tensor.device.type, 0)
+    if torch_tensor.is_contiguous():
+        return infinicore.from_blob(
+            torch_tensor.data_ptr(),
+            list(torch_tensor.shape),
+            dtype=to_infinicore_dtype(torch_tensor.dtype),
+            device=infini_device,
+        )
+    else:
+        return infinicore.strided_from_blob(
+            torch_tensor.data_ptr(),
+            list(torch_tensor.shape),
+            list(torch_tensor.stride()),
+            dtype=to_infinicore_dtype(torch_tensor.dtype),
+            device=infini_device,
+        )
 
 
 def convert_infinicore_to_torch(infini_result, torch_reference):
@@ -179,16 +173,7 @@ def convert_infinicore_to_torch(infini_result, torch_reference):
         dtype=to_torch_dtype(infini_result.dtype),
         device=infini_result.device.type,
     )
-    if infini_result.is_contiguous():
-        temp_tensor = create_infinicore_tensor(
-            torch_result_from_infini, infini_result.device.type
-        )
-    else:
-        rearrange_tensor(torch_result_from_infini, list(torch_reference.stride()))
-        temp_tensor = create_strided_infinicore_tensor(
-            torch_result_from_infini, infini_result.device.type
-        )
-
+    temp_tensor = infinicore_tensor_from_torch(torch_result_from_infini)
     temp_tensor.copy_(infini_result)
     return torch_result_from_infini