Merge pull request #98 from YdrMaster/main

issue/97/feat: 实现 rearrange cpu

include/infiniop/ops/rearrange.h

@@ -5,12 +5,19 @@
 
 typedef InfiniopDescriptor *infiniopRearrangeDescriptor_t;
 
-__C __export infiniStatus_t infiniopCreateRearrangeDescriptor(infiniopHandle_t handle,
+__C __export infiniStatus_t infiniopCreateRearrangeDescriptor(
+    infiniopHandle_t handle,
     infiniopRearrangeDescriptor_t *desc_ptr,
     infiniopTensorDescriptor_t dst,
     infiniopTensorDescriptor_t src);
 
-__C __export infiniStatus_t infiniopRearrange(infiniopRearrangeDescriptor_t desc, void *dst, void const *src, void *stream);
+__C __export infiniStatus_t infiniopRearrange(
+    infiniopRearrangeDescriptor_t desc,
+    void *dst,
+    const void *src,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyRearrangeDescriptor(
+    infiniopRearrangeDescriptor_t desc);
 
-__C __export infiniStatus_t infiniopDestroyRearrangeDescriptor(infiniopRearrangeDescriptor_t desc);
 #endif

src/infiniop/ops/rearrange/cpu/rearrange_cpu.cc

+#include "rearrange_cpu.h"
+#include "../../../devices/cpu/common_cpu.h"
+#include "../../../devices/cpu/cpu_handle.h"
+#include "../../../tensor.h"
+
+namespace op::rearrange::cpu {
+
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t y_desc,
+    infiniopTensorDescriptor_t x_desc) {
+    auto handle = reinterpret_cast<device::cpu::Handle *>(handle_);
+    auto dtype = y_desc->dtype();
+    auto ndim = y_desc->ndim();
+    auto shape = y_desc->shape().data();
+
+    CHECK_API_OR(x_desc->dtype(), dtype, return INFINI_STATUS_BAD_TENSOR_DTYPE);
+    CHECK_API_OR(x_desc->ndim(), ndim, return INFINI_STATUS_BAD_TENSOR_SHAPE);
+
+    for (size_t i = 0; i < ndim; ++i) {
+        CHECK_API_OR(x_desc->shape()[i], shape[i], return INFINI_STATUS_BAD_TENSOR_SHAPE);
+    }
+
+    auto dst_strides = y_desc->strides().data();
+    auto src_strides = x_desc->strides().data();
+    auto element_size = infiniSizeOf(dtype);
+    auto meta = utils::RearrangeMeta::create(shape, dst_strides, src_strides, ndim, element_size);
+
+    if (!meta) {
+        return INFINI_STATUS_BAD_TENSOR_STRIDES;
+    }
+
+    *desc_ptr = new Descriptor(
+        std::move(*meta),
+        nullptr,
+        handle->device,
+        handle->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *y,
+    const void *x,
+    void *stream) const {
+    _meta.launch(y, x);
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace op::rearrange::cpu

src/infiniop/ops/rearrange/cpu/rearrange_cpu.h

+#ifndef __REARRANGE_CPU_H__
+#define __REARRANGE_CPU_H__
+
+#include "../rearrange.h"
+
+DESCRIPTOR(cpu)
+
+#endif // __REARRANGE_CPU_H__

src/infiniop/ops/rearrange/operator.cc

@@ -2,120 +2,78 @@
 #include "../../handle.h"
 #include "infiniop/ops/rearrange.h"
 
+#ifdef ENABLE_CPU_API
+#include "cpu/rearrange_cpu.h"
+#endif
+
 __C infiniStatus_t infiniopCreateRearrangeDescriptor(
     infiniopHandle_t handle,
     infiniopRearrangeDescriptor_t *desc_ptr,
     infiniopTensorDescriptor_t dst,
     infiniopTensorDescriptor_t src) {
+
+#define CREATE(CASE, NAMESPACE)                                                  \
+    case CASE:                                                                   \
+        return op::rearrange::NAMESPACE::Descriptor::create(                     \
+            handle,                                                              \
+            reinterpret_cast<op::rearrange::NAMESPACE::Descriptor **>(desc_ptr), \
+            dst,                                                                 \
+            src)
+
     switch (handle->device) {
-#ifdef ENABLE_CPU
-    case DevCpu:
-        return cpuCreateRearrangeDescriptor(handle, (RearrangeCpuDescriptor_t *)desc_ptr, dst, src);
-#endif
-#ifdef ENABLE_NV_GPU
-    case DevNvGpu: {
-        return cudaCreateRearrangeDescriptor((CudaHandle_t)handle, (RearrangeCudaDescriptor_t *)desc_ptr, dst, src);
-    }
 
+#ifdef ENABLE_CPU_API
+        CREATE(INFINI_DEVICE_CPU, cpu);
 #endif
-#ifdef ENABLE_CAMBRICON_MLU
-    case DevCambriconMlu: {
-        return bangCreateRearrangeDescriptor((BangHandle_t)handle, (RearrangeBangDescriptor_t *)desc_ptr, dst, src);
-    }
-#endif
-#ifdef ENABLE_ASCEND_NPU
-    case DevAscendNpu: {
-        return aclnnCreateRearrangeDescriptor((AscendHandle_t)handle,
-                                              (RearrangeAclnnDescriptor_t *)desc_ptr,
-                                              dst,
-                                              src);
-    }
-#endif
-#ifdef ENABLE_METAX_GPU
-    case DevMetaxGpu: {
-        return macaCreateRearrangeDescriptor((MacaHandle_t)handle, (RearrangeMacaDescriptor_t *)desc_ptr, dst, src);
-    }
-#endif
-#ifdef ENABLE_MTHREADS_GPU
-    case DevMthreadsGpu: {
-        return musaCreateRearrangeDescriptor((MusaHandle_t)handle, (RearrangeMusaDescriptor_t *)desc_ptr, dst, src);
-    }
-#endif
-    }
+
+    default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef CREATE
 }
 
-__C infiniStatus_t infiniopRearrange(infiniopRearrangeDescriptor_t desc, void *dst, const void *src, void *stream) {
+__C infiniStatus_t infiniopRearrange(
+    infiniopRearrangeDescriptor_t desc,
+    void *dst,
+    const void *src,
+    void *stream) {
+
+#define CALCULATE(CASE, NAMESPACE)                                                  \
+    case CASE:                                                                      \
+        return reinterpret_cast<const op::rearrange::NAMESPACE::Descriptor *>(desc) \
+            ->calculate(dst, src, stream)
+
     switch (desc->device_type) {
-#ifdef ENABLE_CPU
-    case DevCpu:
-        return cpuRearrange((RearrangeCpuDescriptor_t)desc, dst, src, stream);
-#endif
-#ifdef ENABLE_NV_GPU
-    case DevNvGpu: {
-        return cudaRearrange((RearrangeCudaDescriptor_t)desc, dst, src, stream);
-    }
 
+#ifdef ENABLE_CPU_API
+        CALCULATE(INFINI_DEVICE_CPU, cpu);
 #endif
-#ifdef ENABLE_CAMBRICON_MLU
-    case DevCambriconMlu: {
-        return bangRearrange((RearrangeBangDescriptor_t)desc, dst, src, stream);
-    }
-#endif
-#ifdef ENABLE_ASCEND_NPU
-    case DevAscendNpu: {
-        return aclnnRearrange((RearrangeAclnnDescriptor_t)desc,
-                              dst,
-                              src,
-                              stream);
-    }
-#endif
-#ifdef ENABLE_METAX_GPU
-    case DevMetaxGpu: {
-        return macaRearrange((RearrangeMacaDescriptor_t)desc, dst, src, stream);
-    }
-#endif
-#ifdef ENABLE_MTHREADS_GPU
-    case DevMthreadsGpu: {
-        return musaRearrange((RearrangeMusaDescriptor_t)desc, dst, src, stream);
-    }
-#endif
-    }
+
+    default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef CALCULATE
 }
 
-__C infiniStatus_t infiniopDestroyRearrangeDescriptor(infiniopRearrangeDescriptor_t desc) {
+__C infiniStatus_t infiniopDestroyRearrangeDescriptor(
+    infiniopRearrangeDescriptor_t desc) {
+
+#define DELETE(CASE, NAMESPACE)                                                      \
+    case CASE:                                                                       \
+        delete reinterpret_cast<const op::rearrange::NAMESPACE::Descriptor *>(desc); \
+        return INFINI_STATUS_SUCCESS;
+
     switch (desc->device_type) {
-#ifdef ENABLE_CPU
-    case DevCpu:
-        return cpuDestroyRearrangeDescriptor((RearrangeCpuDescriptor_t)desc);
-#endif
-#ifdef ENABLE_NV_GPU
-    case DevNvGpu: {
-        return cudaDestroyRearrangeDescriptor((RearrangeCudaDescriptor_t)desc);
-    }
 
+#ifdef ENABLE_CPU_API
+        DELETE(INFINI_DEVICE_CPU, cpu);
 #endif
-#ifdef ENABLE_CAMBRICON_MLU
-    case DevCambriconMlu: {
-        return bangDestroyRearrangeDescriptor((RearrangeBangDescriptor_t)desc);
-    }
-#endif
-#ifdef ENABLE_ASCEND_NPU
-    case DevAscendNpu: {
-        return aclnnDestroyRearrangeDescriptor((RearrangeAclnnDescriptor_t)desc);
-    }
-#endif
-#ifdef ENABLE_METAX_GPU
-    case DevMetaxGpu: {
-        return macaDestroyRearrangeDescriptor((RearrangeMacaDescriptor_t)desc);
-    }
-#endif
-#ifdef ENABLE_MTHREADS_GPU
-    case DevMthreadsGpu: {
-        return musaDestroyRearrangeDescriptor((RearrangeMusaDescriptor_t)desc);
-    }
-#endif
-    }
+
+    default:
         return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
+    }
+
+#undef DELETE
 }

src/infiniop/ops/rearrange/rearrange.h

+#ifndef __REARRANGE_H__
+#define __REARRANGE_H__
+
+#include "../../../utils.h"
+#include "../../operator.h"
+
+#define DESCRIPTOR(NAMESPACE)                             \
+                                                          \
+    namespace op::rearrange::NAMESPACE {                  \
+    class Descriptor final : public InfiniopDescriptor {  \
+        struct Opaque;                                    \
+        Opaque *_opaque;                                  \
+        utils::RearrangeMeta _meta;                       \
+                                                          \
+        Descriptor(                                       \
+            utils::RearrangeMeta meta,                    \
+            Opaque *opaque,                               \
+            infiniDevice_t device_type,                   \
+            int device_id)                                \
+            : InfiniopDescriptor{device_type, device_id}, \
+              _opaque(opaque),                            \
+              _meta(meta) {}                              \
+                                                          \
+    public:                                               \
+        ~Descriptor();                                    \
+                                                          \
+        static infiniStatus_t create(                     \
+            infiniopHandle_t handle,                      \
+            Descriptor **desc_ptr,                        \
+            infiniopTensorDescriptor_t y_desc,            \
+            infiniopTensorDescriptor_t x_desc);           \
+                                                          \
+        infiniStatus_t calculate(                         \
+            void *y,                                      \
+            const void *x,                                \
+            void *stream) const;                          \
+    };                                                    \
+    }
+
+#endif // __REARRANGE_H__