issue/66: 重构7个算子的测试脚本

test/infiniop/causal_softmax.py

@@ -1,25 +1,47 @@
-from ctypes import POINTER, Structure, c_int32, c_uint64, c_void_p
+import torch
 import ctypes
-import sys
-import os
-
-
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
-from operatorspy import (
-    open_lib,
-    to_tensor,
-    DeviceEnum,
+from ctypes import POINTER, Structure, c_int32, c_size_t, c_uint64, c_void_p, c_float
+from libinfiniop import (
     infiniopHandle_t,
     infiniopTensorDescriptor_t,
-    create_handle,
-    destroy_handle,
+    open_lib,
+    to_tensor,
+    get_test_devices,
     check_error,
-    rearrange_tensor,
+    rearrange_if_needed,
     create_workspace,
+    test_operator,
+    get_args,
+    debug,
+    get_tolerance,
+    profile_operation,
 )
 
-from operatorspy.tests.test_utils import get_args
-import torch
+# ==============================================================================
+#  Configuration (Internal Use Only)
+# ==============================================================================
+# These are not meant to be imported from other modules
+_TEST_CASES = [
+    # x_shape, x_stride
+    ((32, 512), None),
+    ((32, 512), (1024, 1)),
+    ((32, 5, 5), None),
+    ((32, 20, 512), None),
+    ((32, 20, 512), (20480, 512, 1)),  # Ascend 暂不支持非连续
+]
+
+# Data types used for testing
+_TENSOR_DTYPES = [torch.float16]
+
+# Tolerance map for different data types
+_TOLERANCE_MAP = {
+    torch.float16: {"atol": 0, "rtol": 1e-2},
+}
+
+DEBUG = False
+PROFILE = False
+NUM_PRERUN = 10
+NUM_ITERATIONS = 1000
 
 
 class CausalSoftmaxDescriptor(Structure):
@@ -37,101 +59,82 @@ def causal_softmax(x):
     return torch.nn.functional.softmax(masked, dim=-1).to(type)
 
 
-def test(lib, handle, torch_device, x_shape, x_stride=None, x_dtype=torch.float16):
+def test(lib, handle, torch_device, x_shape, x_stride=None, dtype=torch.float16):
     print(
-        f"Testing CausalSoftmax on {torch_device} with x_shape:{x_shape} x_stride:{x_stride} dtype:{x_dtype}"
+        f"Testing CausalSoftmax on {torch_device} with x_shape:{x_shape} x_stride:{x_stride} dtype:{dtype}"
     )
-    x = torch.rand(x_shape, dtype=x_dtype).to(torch_device)
-    if x_stride is not None:
-        x = rearrange_tensor(x, x_stride)
+
+    x = torch.rand(x_shape, dtype=dtype).to(torch_device)
+
     ans = causal_softmax(x)
+
+    x = rearrange_if_needed(x, x_stride)
+
     x_tensor = to_tensor(x, lib)
+
     descriptor = infiniopCausalSoftmaxDescriptor_t()
     check_error(
         lib.infiniopCreateCausalSoftmaxDescriptor(
             handle, ctypes.byref(descriptor), x_tensor.descriptor
         )
     )
-    workspace_size = c_uint64(0)
-    check_error(
-        lib.infiniopGetCausalSoftmaxWorkspaceSize(
-            descriptor, ctypes.byref(workspace_size)
-        )
-    )
 
     # Invalidate the shape and strides in the descriptor to prevent them from being directly used by the kernel
     x_tensor.descriptor.contents.invalidate()
 
-    workspace = create_workspace(workspace_size.value, x.device)
+    workspace_size = c_uint64(0)
     check_error(
-        lib.infiniopCausalSoftmax(
-            descriptor,
-            workspace.data_ptr() if workspace is not None else None,
-            workspace_size.value,
-            x_tensor.data,
-            None,
+        lib.infiniopGetCausalSoftmaxWorkspaceSize(
+            descriptor, ctypes.byref(workspace_size)
         )
     )
-    assert torch.allclose(x, ans, atol=0, rtol=1e-2)
-    check_error(lib.infiniopDestroyCausalSoftmaxDescriptor(descriptor))
-
-
-def test_cpu(lib, test_cases):
-    device = DeviceEnum.DEVICE_CPU
-    handle = create_handle(lib, device)
-    for x_shape, x_stride in test_cases:
-        test(lib, handle, "cpu", x_shape, x_stride)
-    destroy_handle(lib, handle)
-
-
-def test_cuda(lib, test_cases):
-    device = DeviceEnum.DEVICE_CUDA
-    handle = create_handle(lib, device)
-    for x_shape, x_stride in test_cases:
-        test(lib, handle, "cuda", x_shape, x_stride)
-    destroy_handle(lib, handle)
-
-
-def test_bang(lib, test_cases):
-    import torch_mlu
+    workspace = create_workspace(workspace_size.value, x.device)
 
-    device = DeviceEnum.DEVICE_BANG
-    handle = create_handle(lib, device)
-    for x_shape, x_stride in test_cases:
-        test(lib, handle, "mlu", x_shape, x_stride)
-    destroy_handle(lib, handle)
+    def lib_causal_softmax():
+        check_error(
+            lib.infiniopCausalSoftmax(
+                descriptor,
+                workspace.data_ptr() if workspace is not None else None,
+                workspace_size.value,
+                x_tensor.data,
+                None,
+            )
+        )
 
+    lib_causal_softmax()
 
-def test_ascend(lib, test_cases):
-    import torch_npu
+    atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype)
+    if DEBUG:
+        debug(x, ans, atol=atol, rtol=rtol)
+    assert torch.allclose(x, ans, atol=atol, rtol=rtol)
 
-    device = DeviceEnum.DEVICE_ASCEND
-    handle = create_handle(lib, device)
-    for x_shape, x_stride in test_cases:
-        test(lib, handle, "npu", x_shape, x_stride)
+    # Profiling workflow
+    if PROFILE:
+        # fmt: off
+        profile_operation("PyTorch", lambda: causal_softmax(x), torch_device, NUM_PRERUN, NUM_ITERATIONS)
+        profile_operation("    lib", lambda: lib_causal_softmax(), torch_device, NUM_PRERUN, NUM_ITERATIONS)
+        # fmt: on
 
-    destroy_handle(lib, handle)
+    check_error(lib.infiniopDestroyCausalSoftmaxDescriptor(descriptor))
 
 
 if __name__ == "__main__":
-    test_cases = [
-        # x_shape, x_stride
-        ((32, 20, 512), None),
-        ((32, 20, 512), (20480, 512, 1)),  # Ascend 暂不支持非连续
-    ]
     args = get_args()
     lib = open_lib()
+
     lib.infiniopCreateCausalSoftmaxDescriptor.restype = c_int32
     lib.infiniopCreateCausalSoftmaxDescriptor.argtypes = [
         infiniopHandle_t,
         POINTER(infiniopCausalSoftmaxDescriptor_t),
         infiniopTensorDescriptor_t,
     ]
+
     lib.infiniopGetCausalSoftmaxWorkspaceSize.restype = c_int32
     lib.infiniopGetCausalSoftmaxWorkspaceSize.argtypes = [
         infiniopCausalSoftmaxDescriptor_t,
         POINTER(c_uint64),
     ]
+
     lib.infiniopCausalSoftmax.restype = c_int32
     lib.infiniopCausalSoftmax.argtypes = [
         infiniopCausalSoftmaxDescriptor_t,
@@ -140,19 +143,19 @@ def test_ascend(lib, test_cases):
         c_void_p,
         c_void_p,
     ]
+
     lib.infiniopDestroyCausalSoftmaxDescriptor.restype = c_int32
     lib.infiniopDestroyCausalSoftmaxDescriptor.argtypes = [
         infiniopCausalSoftmaxDescriptor_t,
     ]
 
-    if args.cpu:
-        test_cpu(lib, test_cases)
-    if args.cuda:
-        test_cuda(lib, test_cases)
-    if args.bang:
-        test_bang(lib, test_cases)
-    if args.ascend:
-        test_ascend(lib, test_cases)
-    if not (args.cpu or args.cuda or args.bang or args.ascend):
-        test_cpu(lib, test_cases)
+    # Configure testing options
+    DEBUG = args.debug
+    PROFILE = args.profile
+    NUM_PRERUN = args.num_prerun
+    NUM_ITERATIONS = args.num_iterations
+
+    for device in get_test_devices(args):
+        test_operator(lib, device, test, _TEST_CASES, _TENSOR_DTYPES)
+
     print("\033[92mTest passed!\033[0m")