feat(cpu): support concat negative axis

spike-zhu · spike-zhu · commit a828bfdaab3f · 2024-12-23T18:41:41.000+08:00
diff --git a/include/ops/concat/concat.h b/include/ops/concat/concat.h
@@ -4,29 +4,24 @@
 #include "../../export.h"
 #include "../../operators.h"
 
-// Concat描述符结构
 typedef struct ConcatDescriptor {
-    Device device;  // 设备类型（例如 DevCpu、DevNvGpu）
-    uint64_t axis;  // 拼接轴（从0开始）
+    Device device;  
 } ConcatDescriptor;
 
 typedef ConcatDescriptor *infiniopConcatDescriptor_t;
 
-// 创建Concat描述符
 __C __export infiniopStatus_t infiniopCreateConcatDescriptor(infiniopHandle_t handle,
                                                              infiniopConcatDescriptor_t *desc_ptr,
                                                              infiniopTensorDescriptor_t y,
                                                              infiniopTensorDescriptor_t *x,
                                                              uint64_t num_inputs,
-                                                             uint64_t axis);
+                                                             int64_t axis);
 
-// 执行Concat操作
 __C __export infiniopStatus_t infiniopConcat(infiniopConcatDescriptor_t desc,
                                              void *y,
                                              void const **x,
                                              void *stream);
-
-// 销毁Concat描述符
+                                             
 __C __export infiniopStatus_t infiniopDestroyConcatDescriptor(infiniopConcatDescriptor_t desc);
 
 #endif
diff --git a/operatorspy/liboperators.py b/operatorspy/liboperators.py
@@ -8,7 +8,7 @@
 Device = c_int
 Optype = c_int
 
-LIB_OPERATORS_DIR = os.path.join(os.environ.get("INFINI_ROOT"))
+LIB_OPERATORS_DIR = os.path.join(os.environ.get("INFINI_ROOT"), "lib")
 
 class TensorDescriptor(Structure):
     _fields_ = [
diff --git a/operatorspy/tests/concat.py b/operatorspy/tests/concat.py
@@ -64,9 +64,6 @@ def test(
         c = torch.zeros(c_shape, dtype=tensor_dtype).to(torch_device)
     
     ans = concat_py(*inputs, dim=axis)
-
-    print("ans:",ans)
-    print("-" * 50)  
     
     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)
@@ -97,11 +94,7 @@ def test(
             None  
         )
     )
-    
-    print("c2:",c)
-    print("-" * 50)  
 
-    # 验证结果
     assert torch.allclose(c, ans, atol=0, rtol=0), "Concat result does not match PyTorch's result."
     
     check_error(lib.infiniopDestroyConcatDescriptor(descriptor))
diff --git a/src/ops/concat/cpu/concat_cpu.cc b/src/ops/concat/cpu/concat_cpu.cc
@@ -8,22 +8,25 @@ infiniopStatus_t cpuCreateConcatDescriptor(
     infiniopTensorDescriptor_t y,
     infiniopTensorDescriptor_t *x,
     uint64_t num_inputs,
-    uint64_t axis) {
+    int64_t axis) {
     if (y == nullptr || x == nullptr || desc_ptr == nullptr || num_inputs == 0) {
         return STATUS_BAD_PARAM;
     }
 
-    uint64_t ndim = y->ndim;  // 输出张量维度
-    if (axis >= ndim) {
-        return STATUS_BAD_TENSOR_SHAPE;
+    int64_t ndim = y->ndim;  
+    if (axis >= ndim || axis < -ndim) {
+        return STATUS_BAD_PARAM;
+    }
+
+    if(axis < 0){
+        axis = axis + ndim;
     }
 
-    uint64_t total_size = 0;  // 拼接轴的总大小
-    std::vector<std::vector<uint64_t>> input_shapes(num_inputs);  // 输入张量形状
+    uint64_t total_size = 0;  
+    std::vector<std::vector<uint64_t>> input_shapes(num_inputs);  
 
     std::vector<uint64_t> output_shape(y->shape, y->shape + ndim);
 
-    // 验证输入张量的形状和步长
     for (size_t i = 0; i < num_inputs; ++i) {
 
         if (x[i]->dt != y->dt) {
@@ -41,12 +44,9 @@ infiniopStatus_t cpuCreateConcatDescriptor(
         }
 
         input_shapes[i] = std::vector<uint64_t>(x[i]->shape, x[i]->shape + ndim);
-
-        // 累加拼接轴的总大小
         total_size += x[i]->shape[axis];
     }
 
-    // 验证输出张量形状是否匹配
     if (total_size != y->shape[axis]) {
         return STATUS_BAD_TENSOR_SHAPE;
     }
@@ -72,8 +72,7 @@ template <typename T>
 infiniopStatus_t concatCompute(const ConcatCpuDescriptor_t& desc,
                                T* y,
                                void const** x) {
-    // 获取描述符中的信息
-    uint64_t axis = desc->axis;
+    int64_t axis = desc->axis;
     uint64_t num_inputs = desc->num_inputs;
     const std::vector<std::vector<uint64_t>>& input_shapes = desc->input_shapes;
     const std::vector<uint64_t>& output_shape = desc->output_shape;
@@ -84,7 +83,6 @@ infiniopStatus_t concatCompute(const ConcatCpuDescriptor_t& desc,
     }
     size_t blockOffset = output_shape[axis] * blockOffsetInner;
 
-    // concat
     for (size_t i = 0; i < num_inputs; ++i) {
         const std::vector<uint64_t>& input_shape = input_shapes[i];
 
@@ -104,7 +102,6 @@ infiniopStatus_t concatCompute(const ConcatCpuDescriptor_t& desc,
             inSize *= dim;
         }
 
-        // 获取输入和输出的数据指针
         T* input_data = static_cast<T*>(const_cast<void*>(x[i]));
 
         #pragma omp parallel for
@@ -120,16 +117,15 @@ infiniopStatus_t concatCompute(const ConcatCpuDescriptor_t& desc,
     return STATUS_SUCCESS; 
 }
 
-// 主拼接函数
 infiniopStatus_t cpuConcat(ConcatCpuDescriptor_t desc,
                            void *y,
                            void const **x,
                            void *stream) {
-    // 根据数据类型调用相应的模板实例
+
     switch (desc->dtype.size) {
         case sizeof(float): // FLOAT32
             return concatCompute<float>(desc, reinterpret_cast<float*>(y), x);
-        // 可以根据需要添加更多数据类型
+        // add other data.type
         default:
             return STATUS_SUCCESS;
     }
diff --git a/src/ops/concat/cpu/concat_cpu.h b/src/ops/concat/cpu/concat_cpu.h
@@ -4,26 +4,23 @@
 #include <vector>
 #include <cstring>
 
-// 支持高维拼接的CPU-specific Concat描述符
 struct ConcatCpuDescriptor {
     Device device;                                
     DT dtype;                                    
-    uint64_t axis;                               
+    int64_t axis;                               
     uint64_t num_inputs;                        
-    std::vector<std::vector<uint64_t>> input_shapes;  // 输入张量的形状
-    std::vector<uint64_t> output_shape;              // 输出张量的形状
+    std::vector<std::vector<uint64_t>> input_shapes;  
+    std::vector<uint64_t> output_shape;              
 };
 
-
-
 typedef struct ConcatCpuDescriptor *ConcatCpuDescriptor_t;
 
 infiniopStatus_t cpuCreateConcatDescriptor(infiniopHandle_t handle,
                                            ConcatCpuDescriptor_t *desc_ptr,
                                            infiniopTensorDescriptor_t y,
                                            infiniopTensorDescriptor_t *x,
                                            uint64_t num_inputs,
-                                           uint64_t axis);
+                                           int64_t axis);
 
 infiniopStatus_t cpuConcat(ConcatCpuDescriptor_t desc,
                            void *y,
diff --git a/src/ops/concat/operator.cc b/src/ops/concat/operator.cc
@@ -16,7 +16,7 @@ __C infiniopStatus_t infiniopCreateConcatDescriptor(
     infiniopTensorDescriptor_t y,
     infiniopTensorDescriptor_t *x,
     uint64_t num_inputs,
-    uint64_t axis) {
+    int64_t axis) {
     switch (handle->device) {
 #ifdef ENABLE_CPU
         case DevCpu:
