update

Author: spike-zhu

operatorspy/liboperators.py

@@ -8,8 +8,7 @@
 Device = c_int
 Optype = c_int
 
-LIB_OPERATORS_DIR = os.path.join(os.environ.get("INFINI_ROOT"), "lib")
-
+LIB_OPERATORS_DIR = os.path.join(os.environ.get("INFINI_ROOT"))
 
 class TensorDescriptor(Structure):
     _fields_ = [
@@ -19,10 +18,8 @@ class TensorDescriptor(Structure):
         ("pattern", POINTER(c_int64)),
     ]
 
-
 infiniopTensorDescriptor_t = ctypes.POINTER(TensorDescriptor)
 
-
 class CTensor:
     def __init__(self, desc, data):
         self.descriptor = desc

operatorspy/tests/concat.py

@@ -3,7 +3,6 @@
 import sys
 import os
 
-# 调整路径以导入 operatorspy 模块
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
 from operatorspy import (
     open_lib,
@@ -24,10 +23,6 @@
 class Inplace(Enum):
     OUT_OF_PLACE = auto()
     # 对于 concat 算子，通常不支持 in-place 操作，因此这里只保留 OUT_OF_PLACE
-    # 你可以根据实际需求扩展其他选项
-    # INPLACE_A = auto()
-    # INPLACE_B = auto()
-
 
 class ConcatDescriptor(Structure):
     _fields_ = [("device", c_int32),]
@@ -58,28 +53,24 @@ def test(
         f"Testing Concat on {torch_device} with output_shape:{c_shape}, input_shapes:{input_shapes}, axis:{axis}, dtype:{tensor_dtype}, inplace: {inplace.name}"
     )
 
-    # 创建输入张量
     inputs = [torch.rand(shape, dtype=tensor_dtype).to(torch_device) for shape in input_shapes]
 
     for idx, tensor in enumerate(inputs):
         print(f"Input {idx}:")
         print(tensor)
         print("-" * 50)  
 
-    # 创建输出张量
     if inplace == Inplace.OUT_OF_PLACE:
         c = torch.zeros(c_shape, dtype=tensor_dtype).to(torch_device)
     else:
         # 对于 concat，通常不支持 in-place 操作，因此这里简化为 OUT_OF_PLACE
         c = torch.zeros(c_shape, dtype=tensor_dtype).to(torch_device)
 
-    # 使用 PyTorch 进行拼接，作为参考答案
     ans = concat_py(*inputs, dim=axis)
 
     print("ans:",ans)
     print("-" * 50)  
 
-    # 将张量转换为 infiniop 所需的格式
     input_tensors = [to_tensor(t, lib) for t in inputs]
     c_tensor = to_tensor(c, lib) if inplace == Inplace.OUT_OF_PLACE else to_tensor(c, lib)
 
@@ -91,22 +82,17 @@ def test(
     input_desc_array_type = infiniopTensorDescriptor_t * num_inputs
     input_desc_array = input_desc_array_type(*[t.descriptor for t in input_tensors])
 
-    # 创建描述符
     check_error(
         lib.infiniopCreateConcatDescriptor(
             handle,
             ctypes.byref(descriptor),
-            c_tensor.descriptor,            # 使用 c_tensor 的描述符
-            input_desc_array,               # 输入张量描述符数组
+            c_tensor.descriptor,            
+            input_desc_array,              
             c_uint64(num_inputs),
             c_uint64(axis),
         )
     )
 
-    print("c1:",c)
-    print("-" * 50)  
-
-    # 执行拼接操作
     input_data_ptrs = (c_void_p * num_inputs)(*[t.data for t in input_tensors])
     check_error(
         lib.infiniopConcat(
@@ -121,9 +107,8 @@ def test(
     print("-" * 50)  
 
     # 验证结果
-    assert torch.allclose(c, ans, atol=0, rtol=1e-5), "Concat result does not match PyTorch's result."
+    assert torch.allclose(c, ans, atol=0, rtol=0), "Concat result does not match PyTorch's result."
 
-    # 销毁描述符
     check_error(lib.infiniopDestroyConcatDescriptor(descriptor))
 
 
@@ -157,50 +142,59 @@ def test_bang(lib, test_cases):
     # 定义测试用例
     test_cases = [
         # (output_shape, axis, input_shapes, inplace)
-
-        ((6, 3), 0, [(2, 3), (4, 3)], Inplace.OUT_OF_PLACE),
-        # ((3, 6), 1, [(3, 2), (3, 4)], Inplace.OUT_OF_PLACE),
-        # ((3, 7), 1, [(3, 2), (3, 4), (3,1)], Inplace.OUT_OF_PLACE),
-        # ((3, 3, 10), 2, [(3, 3, 4), (3, 3, 6)], Inplace.OUT_OF_PLACE),
-        # ((1, 1), 0, [(1, 1)], Inplace.OUT_OF_PLACE),
-        # ((4, 5, 6), 0, [(1, 5, 6), (3, 5, 6)], Inplace.OUT_OF_PLACE),
-        # ((2, 3, 6), 2, [(2, 3, 2), (2, 3, 4)], Inplace.OUT_OF_PLACE),
-
-        # 添加更多测试用例以覆盖不同的维度和拼接轴
-        # ((2, 10, 3), 1, [(2, 5, 3), (2, 2, 3),(2,3,3)], Inplace.OUT_OF_PLACE),  # 拼接沿第二维
+        
+        # 一维张量拼接
+        ((6,), 0, [(2,), (4,)], Inplace.OUT_OF_PLACE),  # 沿第0轴拼接
+        # 二维张量拼接
+        ((6, 3), 0, [(2, 3), (4, 3)], Inplace.OUT_OF_PLACE),  # 沿第0轴拼接
+        ((3, 6), 1, [(3, 2), (3, 4)], Inplace.OUT_OF_PLACE),  # 沿第1轴拼接
+        ((3, 7), 1, [(3, 2), (3, 4), (3, 1)], Inplace.OUT_OF_PLACE),  # 沿第1轴拼接
+        ((3, 3, 10), 2, [(3, 3, 4), (3, 3, 6)], Inplace.OUT_OF_PLACE),  # 沿第2轴拼接
+        # 三维张量拼接
+        ((4, 3, 6), 0, [(3, 3, 6), (1, 3, 6)], Inplace.OUT_OF_PLACE),  # 沿第0轴拼接
+        ((2, 6, 3), 1, [(2, 3, 3), (2, 3, 3)], Inplace.OUT_OF_PLACE),  # 沿第1轴拼接
+        ((2, 3, 6), 2, [(2, 3, 3), (2, 3, 3)], Inplace.OUT_OF_PLACE),  # 沿第2轴拼接
+        # 四维张量拼接
+        ((4, 3, 5, 6), 0, [(1, 3, 5, 6), (3, 3, 5, 6)], Inplace.OUT_OF_PLACE),  # 沿第0轴拼接
+        ((2, 5, 5, 6), 1, [(2, 3, 5, 6), (2, 2, 5, 6)], Inplace.OUT_OF_PLACE),  # 沿第1轴拼接
+        ((2, 3, 5, 6), 2, [(2, 3, 2, 6), (2, 3, 3, 6)], Inplace.OUT_OF_PLACE),  # 沿第2轴拼接
+        ((2, 3, 5, 6), 3, [(2, 3, 5, 3), (2, 3, 5, 3)], Inplace.OUT_OF_PLACE),  # 沿第3轴拼接
+        ((2, 3, 5, 15), 3, [(2, 3, 5, 3), (2, 3, 5, 3), (2, 3, 5, 9)], Inplace.OUT_OF_PLACE),  # 沿第3轴拼接
+        # 五维张量拼接
+        ((4, 2, 3, 4, 5), 0, [(1, 2, 3, 4, 5), (3, 2, 3, 4, 5)], Inplace.OUT_OF_PLACE),  # 沿第0轴拼接
+        ((2, 4, 3, 2, 5), 1, [(2, 2, 3, 2, 5), (2, 2, 3, 2, 5)], Inplace.OUT_OF_PLACE),  # 沿第1轴拼接
+        ((1, 2, 4, 4, 5), 2, [(1, 2, 2, 4, 5), (1, 2, 2, 4, 5)], Inplace.OUT_OF_PLACE),  # 沿第2轴拼接
+        ((1, 2, 3, 8, 5), 3, [(1, 2, 3, 4, 5), (1, 2, 3, 4, 5)], Inplace.OUT_OF_PLACE),  # 沿第3轴拼接
+        ((1, 2, 3, 4, 5), 4, [(1, 2, 3, 4, 3), (1, 2, 3, 4, 2)], Inplace.OUT_OF_PLACE),  # 沿第4轴拼接
+        ((4, 14, 3, 4, 5), 1, [(4, 3, 3, 4, 5), (4, 5, 3, 4, 5), (4, 6, 3, 4, 5)], Inplace.OUT_OF_PLACE),  # 沿第1轴拼接
     ]
 
     args = get_args()
     lib = open_lib()
 
-    # 绑定 C++ 函数
-    # 创建 Concat 描述符
     lib.infiniopCreateConcatDescriptor.restype = c_int32
     lib.infiniopCreateConcatDescriptor.argtypes = [
         infiniopHandle_t,
         POINTER(infiniopConcatDescriptor_t),
-        infiniopTensorDescriptor_t,  # 输出张量描述符
-        POINTER(infiniopTensorDescriptor_t),  # 输入张量描述符数组
-        c_uint64,  # 输入张量数量
-        c_uint64,  # 拼接轴
+        infiniopTensorDescriptor_t,  
+        POINTER(infiniopTensorDescriptor_t),  
+        c_uint64,  # nums_input
+        c_uint64,  # axis
     ]
 
-    # 执行 Concat
     lib.infiniopConcat.restype = c_int32
     lib.infiniopConcat.argtypes = [
         infiniopConcatDescriptor_t,
-        c_void_p,  # 输出数据指针
-        POINTER(c_void_p),  # 输入数据指针数组
-        c_void_p,  # 流（假设为 NULL）
+        c_void_p,  
+        POINTER(c_void_p),  
+        c_void_p, 
     ]
 
-    # 销毁 Concat 描述符
     lib.infiniopDestroyConcatDescriptor.restype = c_int32
     lib.infiniopDestroyConcatDescriptor.argtypes = [
         infiniopConcatDescriptor_t,
     ]
 
-    # 根据命令行参数执行测试
     if args.cpu:
         test_cpu(lib, test_cases)
     if args.cuda: