add support for negative axis in concat opreation

Author: spike-zhu

operatorspy/tests/concat.py

@@ -1,4 +1,4 @@
-from ctypes import POINTER, Structure, c_int32, c_void_p, c_uint64
+from ctypes import POINTER, Structure, c_int32, c_void_p, c_uint64, c_int64
 import ctypes
 import sys
 import os
@@ -22,7 +22,6 @@
 
 class Inplace(Enum):
     OUT_OF_PLACE = auto()
-    # 对于 concat 算子，通常不支持 in-place 操作，保留 OUT_OF_PLACE
 
 class ConcatDescriptor(Structure):
     _fields_ = [("device", c_int32),]
@@ -32,7 +31,6 @@ class ConcatDescriptor(Structure):
 
 
 def concat_py(*tensors, dim=0):
-    """使用 PyTorch 进行拼接的辅助函数"""
     return torch.cat(tensors, dim=dim)
 
 
@@ -63,7 +61,6 @@ def test(
     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)
 
     ans = concat_py(*inputs, dim=axis)
@@ -74,10 +71,8 @@ def test(
     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)
 
-    # 创建 Concat 描述符
     descriptor = infiniopConcatDescriptor_t()
-    
-    # 准备输入描述符数组
+
     num_inputs = len(input_tensors)
     input_desc_array_type = infiniopTensorDescriptor_t * num_inputs
     input_desc_array = input_desc_array_type(*[t.descriptor for t in input_tensors])
@@ -89,7 +84,7 @@ def test(
             c_tensor.descriptor,            
             input_desc_array,              
             c_uint64(num_inputs),
-            c_uint64(axis),
+            c_int64(axis),
         )
     )
 
@@ -139,34 +134,57 @@ def test_bang(lib, test_cases):
 
 
 if __name__ == "__main__":
-    # 定义测试用例
+
     test_cases = [
-        # (output_shape, axis, input_shapes, inplace)
-        
-        # 一维张量拼接
-        ((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轴拼接
+
+        ((6,), 0, [(2,), (4,)], Inplace.OUT_OF_PLACE),  
+
+        ((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),  
+
+        ((4, 3, 6), 0, [(3, 3, 6), (1, 3, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 6, 3), 1, [(2, 3, 3), (2, 3, 3)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 6), 2, [(2, 3, 3), (2, 3, 3)], Inplace.OUT_OF_PLACE),  
+
+        ((4, 3, 5, 6), 0, [(1, 3, 5, 6), (3, 3, 5, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 5, 5, 6), 1, [(2, 3, 5, 6), (2, 2, 5, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 5, 6), 2, [(2, 3, 2, 6), (2, 3, 3, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 5, 6), 3, [(2, 3, 5, 3), (2, 3, 5, 3)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 5, 15), 3, [(2, 3, 5, 3), (2, 3, 5, 3), (2, 3, 5, 9)], Inplace.OUT_OF_PLACE),  
+
+        ((4, 2, 3, 4, 5), 0, [(1, 2, 3, 4, 5), (3, 2, 3, 4, 5)], Inplace.OUT_OF_PLACE),  
+        ((2, 4, 3, 2, 5), 1, [(2, 2, 3, 2, 5), (2, 2, 3, 2, 5)], Inplace.OUT_OF_PLACE),  
+        ((1, 2, 4, 4, 5), 2, [(1, 2, 2, 4, 5), (1, 2, 2, 4, 5)], Inplace.OUT_OF_PLACE),  
+        ((1, 2, 3, 8, 5), 3, [(1, 2, 3, 4, 5), (1, 2, 3, 4, 5)], Inplace.OUT_OF_PLACE), 
+        ((1, 2, 3, 4, 5), 4, [(1, 2, 3, 4, 3), (1, 2, 3, 4, 2)], Inplace.OUT_OF_PLACE),  
+        ((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),  
+
+
+        ((6,), -1, [(2,), (4,)], Inplace.OUT_OF_PLACE),  
+
+        ((6, 3), -2, [(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), -1, [(3, 3, 4), (3, 3, 6)], Inplace.OUT_OF_PLACE),  
+
+        ((4, 3, 6), -3, [(3, 3, 6), (1, 3, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 6, 3), -2, [(2, 3, 3), (2, 3, 3)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 6), -1, [(2, 3, 3), (2, 3, 3)], Inplace.OUT_OF_PLACE),  
+
+        ((4, 3, 5, 6), -4, [(1, 3, 5, 6), (3, 3, 5, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 5, 5, 6), -3, [(2, 3, 5, 6), (2, 2, 5, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 5, 6), -2, [(2, 3, 2, 6), (2, 3, 3, 6)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 5, 6), -1, [(2, 3, 5, 3), (2, 3, 5, 3)], Inplace.OUT_OF_PLACE),  
+        ((2, 3, 5, 15), -1, [(2, 3, 5, 3), (2, 3, 5, 3), (2, 3, 5, 9)], Inplace.OUT_OF_PLACE),  
+
+        ((4, 2, 3, 4, 5), -5, [(1, 2, 3, 4, 5), (3, 2, 3, 4, 5)], Inplace.OUT_OF_PLACE),  
+        ((2, 4, 3, 2, 5), -4, [(2, 2, 3, 2, 5), (2, 2, 3, 2, 5)], Inplace.OUT_OF_PLACE),  
+        ((1, 2, 4, 4, 5), -3, [(1, 2, 2, 4, 5), (1, 2, 2, 4, 5)], Inplace.OUT_OF_PLACE),  
+        ((1, 2, 3, 8, 5), -2, [(1, 2, 3, 4, 5), (1, 2, 3, 4, 5)], Inplace.OUT_OF_PLACE), 
+        ((1, 2, 3, 4, 5), -1, [(1, 2, 3, 4, 3), (1, 2, 3, 4, 2)], Inplace.OUT_OF_PLACE),  
+        ((4, 14, 3, 4, 5), -4, [(4, 3, 3, 4, 5), (4, 5, 3, 4, 5), (4, 6, 3, 4, 5)], Inplace.OUT_OF_PLACE),     
     ]
 
     args = get_args()
@@ -179,7 +197,7 @@ def test_bang(lib, test_cases):
         infiniopTensorDescriptor_t,  
         POINTER(infiniopTensorDescriptor_t),  
         c_uint64,  # nums_input
-        c_uint64,  # axis
+        c_int64,  # axis
     ]
 
     lib.infiniopConcat.restype = c_int32