Merge pull request #139 from InfiniTensor/issue/9

feat: 添加昇腾swiglu算子

src/infiniop/devices/ascend/CMakeLists.txt

@@ -3,7 +3,7 @@ cmake_minimum_required(VERSION 3.16.0)
 # project information
 project(Ascend_C)
 set(SOC_VERSION "Ascend910B3" CACHE STRING "system on chip type")
-set(ASCEND_CANN_PACKAGE_PATH $ENV{ASCEND_HOME} CACHE PATH "ASCEND CANN package installation directory")
+set(ASCEND_CANN_PACKAGE_PATH $ENV{ASCEND_TOOLKIT_HOME} CACHE PATH "ASCEND CANN package installation directory")
 set(RUN_MODE "npu" CACHE STRING "run mode: npu")
 set(CMAKE_BUILD_TYPE "Release" CACHE STRING "Build type Release/Debug (default Debug)" FORCE)
 set(CMAKE_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/out" CACHE STRING "path for install()" FORCE)
@@ -19,10 +19,14 @@ else()
 endif()
 
 include(${ASCENDC_CMAKE_DIR}/ascendc.cmake)
+include_directories(
+    ${CMAKE_SOURCE_DIR}/../../../../include/infiniop/ 
+)
+
 
 ascendc_library(ascend_kernels STATIC
-    ../../ops/swiglu/ascend/swiglu_kernel.cpp
-    ../../ops/rotary_embedding/ascend/rotary_embedding_kernel.cpp
-    ../../ops/random_sample/ascend/random_sample_kernel.cpp
+    ../../ops/swiglu/ascend/swiglu_ascend_kernel.cpp
+    # ../../ops/rotary_embedding/ascend/rotary_embedding_kernel.cpp
+    # ../../ops/random_sample/ascend/random_sample_kernel.cpp
 )
 

src/infiniop/devices/ascend/ascend_kernel_common.h

+#ifndef __INFINIOP_ASCEND_KERNEL_COMMON_H__
+#define __INFINIOP_ASCEND_KERNEL_COMMON_H__
+
+#include "../../../../include/infinicore.h"
+#include "kernel_operator.h"
+
+constexpr int32_t BLOCK_NUM = 8;
+constexpr int32_t BUFFER_NUM = 2;
+constexpr int32_t BYTE_ALIGN = 32;
+
+#endif

src/infiniop/ops/causal_softmax/ascend/causal_softmax_aclnn.cc → src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.cc

-#include "causal_softmax_aclnn.h"
+#include "causal_softmax_ascend.h"
 #include "../../../devices/ascend/common_ascend.h"
 #include <aclnnop/aclnn_masked_fill_tensor.h>
 #include <aclnnop/aclnn_softmax.h>

src/infiniop/ops/causal_softmax/operator.cc

@@ -9,7 +9,7 @@
 #include "cuda/causal_softmax_cuda.cuh"
 #endif
 #ifdef ENABLE_ASCEND_API
-#include "ascend/causal_softmax_aclnn.h"
+#include "ascend/causal_softmax_ascend.h"
 #endif
 
 __C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor(

src/infiniop/ops/swiglu/ascend/swiglu_ascend.cc

+#include "swiglu_ascend.h"
+#include "../../../devices/ascend/common_ascend.h"
+
+namespace op::swiglu::ascend {
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(infiniopHandle_t handle, Descriptor **desc_ptr,
+                                  infiniopTensorDescriptor_t c_desc,
+                                  std::vector<infiniopTensorDescriptor_t> input_descs) {
+    auto handle_ascend = reinterpret_cast<device::ascend::Handle *>(handle);
+
+    auto dtype = c_desc->dtype();
+    CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32);
+
+    const auto &a_desc = input_descs[0];
+    const auto &b_desc = input_descs[1];
+
+    auto result = SwigluInfo::create(c_desc, a_desc, b_desc);
+    CHECK_RESULT(result);
+    SwigluInfo info = result.take();
+
+    // https://www.hiascend.com/document/detail/zh/canncommercial/800/apiref/ascendcopapi/atlasascendc_api_07_0777.html
+    size_t workspace_size = 0;
+
+    *desc_ptr = new Descriptor(std::move(info), workspace_size, handle_ascend->device, handle_ascend->device_id);
+    return INFINI_STATUS_SUCCESS;
+}
+
+extern "C" infiniStatus_t swiglu_kernel_launch(
+    void *c, void *a, void *b,
+    infiniDtype_t dtype, size_t batch, size_t seq, size_t hd,
+    ptrdiff_t stride_batch_c, ptrdiff_t stride_batch_a, ptrdiff_t stride_batch_b,
+    ptrdiff_t stride_seq_c, ptrdiff_t stride_seq_a, ptrdiff_t stride_seq_b, void *stream);
+
+infiniStatus_t Descriptor::calculate(void *workspace,
+                                     size_t workspace_size,
+                                     void *c,
+                                     std::vector<const void *> inputs,
+                                     void *stream) const {
+    auto batch = _info.ndim == 2 ? 1 : _info.shape[0];
+    auto seq_len = _info.ndim == 2 ? _info.shape[0] : _info.shape[1];
+    auto hidden_size = _info.shape[_info.ndim - 1];
+    auto stride_batch_c = _info.ndim == 2 ? 1 : _info.c_strides[0];
+    auto stride_batch_a = _info.ndim == 2 ? 1 : _info.a_strides[0];
+    auto stride_batch_b = _info.ndim == 2 ? 1 : _info.b_strides[0];
+    auto stride_seq_c = _info.ndim == 2 ? _info.c_strides[0] : _info.c_strides[1];
+    auto stride_seq_a = _info.ndim == 2 ? _info.a_strides[0] : _info.a_strides[1];
+    auto stride_seq_b = _info.ndim == 2 ? _info.b_strides[0] : _info.b_strides[1];
+    auto status = swiglu_kernel_launch(c, (void *)inputs[0], (void *)inputs[1], _info.dtype, batch, seq_len, hidden_size, stride_batch_c, stride_batch_a, stride_batch_b, stride_seq_c, stride_seq_a, stride_seq_b, stream);
+    return status;
+}
+
+} // namespace op::swiglu::ascend

src/infiniop/ops/swiglu/ascend/swiglu_ascend.h

+#ifndef __ACLNN_SWIGLU_H__
+#define __ACLNN_SWIGLU_H__
+
+#include "../../../../utils.h"
+#include "../../../../utils/check.h"
+#include "../../../operator.h"
+#include "../../../tensor.h"
+
+namespace op::swiglu::ascend {
+class SwigluInfo {
+
+    SwigluInfo() = default;
+
+public:
+    infiniDtype_t dtype;
+    std::vector<size_t> shape;
+    int32_t ndim;
+    std::vector<ptrdiff_t> c_strides;
+    std::vector<ptrdiff_t> a_strides;
+    std::vector<ptrdiff_t> b_strides;
+
+    static utils::Result<SwigluInfo> create(infiniopTensorDescriptor_t c_desc, infiniopTensorDescriptor_t a_desc, infiniopTensorDescriptor_t b_desc) {
+        CHECK_OR_RETURN(c_desc && a_desc && b_desc, INFINI_STATUS_BAD_PARAM);
+        CHECK_OR_RETURN(!c_desc->hasBroadcastDim(), INFINI_STATUS_BAD_TENSOR_STRIDES);
+        CHECK_OR_RETURN(c_desc->ndim() == a_desc->ndim()
+                            && c_desc->ndim() == b_desc->ndim()
+                            && (c_desc->ndim() == 2 || c_desc->ndim() == 3),
+                        INFINI_STATUS_BAD_TENSOR_SHAPE);
+        CHECK_SAME_SHAPE(c_desc->shape(), a_desc->shape(), b_desc->shape());
+        int32_t ndim = c_desc->ndim();
+        CHECK_OR_RETURN(c_desc->stride(ndim - 1) == 1
+                            && a_desc->stride(ndim - 1) == 1
+                            && b_desc->stride(ndim - 1) == 1,
+                        INFINI_STATUS_BAD_TENSOR_STRIDES);
+        CHECK_OR_RETURN(c_desc->dtype() == a_desc->dtype()
+                            && c_desc->dtype() == b_desc->dtype(),
+                        INFINI_STATUS_BAD_TENSOR_DTYPE);
+
+        return utils::Result<SwigluInfo>(SwigluInfo{
+            c_desc->dtype(),
+            c_desc->shape(),
+            ndim,
+            c_desc->strides(),
+            a_desc->strides(),
+            b_desc->strides(),
+        });
+    }
+};
+
+class Descriptor final : public InfiniopDescriptor {
+    SwigluInfo _info;
+    size_t _workspace_size;
+
+    Descriptor(SwigluInfo info, size_t workspace_size, infiniDevice_t device_type, int device_id) : InfiniopDescriptor{device_type, device_id},
+                                                                                                    _info(info), _workspace_size(workspace_size) {}
+
+public:
+    ~Descriptor();
+    static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr,
+                                 infiniopTensorDescriptor_t c_desc,
+                                 std::vector<infiniopTensorDescriptor_t> input_descs);
+    size_t workspaceSize() const { return _workspace_size; }
+
+    infiniStatus_t calculate(
+        void *workspace,
+        size_t workspace_size,
+        void *c,
+        std::vector<const void *> inputs,
+        void *stream) const;
+};
+
+} // namespace op::swiglu::ascend
+#endif // __ACLNN_SWIGLU_H__

src/infiniop/ops/swiglu/ascend/swiglu_ascend_kernel.cpp

+#include "../../../devices/ascend/ascend_kernel_common.h"
+
+using namespace AscendC;
+
+template <typename T>
+class SwigluKernel {
+public:
+    __aicore__ inline SwigluKernel() {}
+    __aicore__ inline void init(GM_ADDR c, GM_ADDR a, GM_ADDR b, int64_t batch_, int64_t seq, int64_t hd,
+                                int64_t stride_batch_c, int64_t stride_batch_a, int64_t stride_batch_b,
+                                int64_t stride_seq_c, int64_t stride_seq_a, int64_t stride_seq_b);
+    __aicore__ inline void process();
+
+private:
+    __aicore__ inline void copyIn(int64_t i);
+    __aicore__ inline void compute(int64_t i);
+    __aicore__ inline void copyOut(int64_t i);
+
+private:
+    GlobalTensor<T> _c_gm, _a_gm, _b_gm;
+    TQue<QuePosition::VECIN, BUFFER_NUM> _in_queue_a, _in_queue_b;
+    TQue<QuePosition::VECOUT, BUFFER_NUM> _out_queue_c;
+
+    TPipe _pipe;
+    float _beta_value = 1.0f;
+    int64_t _block_idx, _tile_len, _copy_len,
+        _batch, _seq_len, _hidden_size,
+        _stride_seq_a, _stride_seq_b, _stride_seq_c;
+    int64_t _stride_batch_a = 1, _stride_batch_b = 1, _stride_batch_c = 1;
+};
+
+template <typename T>
+__aicore__ inline void SwigluKernel<T>::init(GM_ADDR c, GM_ADDR a, GM_ADDR b, int64_t batch_, int64_t seq, int64_t hd,
+                                             int64_t stride_batch_c, int64_t stride_batch_a, int64_t stride_batch_b,
+                                             int64_t stride_seq_c, int64_t stride_seq_a, int64_t stride_seq_b) {
+    // Init Shape & StrideVariables
+    _batch = batch_;
+    _seq_len = seq;
+    _hidden_size = hd;
+    _stride_batch_a = stride_batch_a;
+    _stride_batch_b = stride_batch_b;
+    _stride_batch_c = stride_batch_c;
+    _stride_seq_a = stride_seq_a;
+    _stride_seq_b = stride_seq_b;
+    _stride_seq_c = stride_seq_c;
+
+    _block_idx = GetBlockIdx();
+    _tile_len = _block_idx < (_hidden_size % BLOCK_NUM) ? (_hidden_size / BLOCK_NUM) + 1 : (_hidden_size / BLOCK_NUM);
+    _copy_len = (_tile_len * sizeof(T)) % BYTE_ALIGN == 0 ? _tile_len : (_tile_len * sizeof(T) + (BYTE_ALIGN - _tile_len * sizeof(T) % BYTE_ALIGN)) / sizeof(T);
+
+    // Set global tensor
+    _a_gm.SetGlobalBuffer((__gm__ T *)a);
+    _b_gm.SetGlobalBuffer((__gm__ T *)b);
+    _c_gm.SetGlobalBuffer((__gm__ T *)c);
+
+    // _pipe alloc memory to queue, the unit is bytes
+    _pipe.InitBuffer(_in_queue_a, BUFFER_NUM, _copy_len * sizeof(T));
+    _pipe.InitBuffer(_in_queue_b, BUFFER_NUM, _copy_len * sizeof(T));
+    _pipe.InitBuffer(_out_queue_c, BUFFER_NUM, _copy_len * sizeof(T));
+}
+
+template <typename T>
+__aicore__ inline void SwigluKernel<T>::copyIn(int64_t i) {
+    // Alloc tensor from queue memory
+    LocalTensor<T> aLocal = _in_queue_a.AllocTensor<T>();
+    LocalTensor<T> bLocal = _in_queue_b.AllocTensor<T>();
+    // Get idx of current tile
+    auto batch_idx = _batch == 1 ? 0 : i / _seq_len;
+    auto seq_idx = _batch == 1 ? i : i % _seq_len;
+
+    int64_t idxa = batch_idx * _stride_batch_a + seq_idx * _stride_seq_a + _block_idx * _tile_len;
+    int64_t idxb = batch_idx * _stride_batch_b + seq_idx * _stride_seq_b + _block_idx * _tile_len;
+    // Copy process_th tile from global tensor to local tensor
+    DataCopy(aLocal, _a_gm[idxa], _copy_len);
+    DataCopy(bLocal, _b_gm[idxb], _copy_len);
+
+    // Enque input tensor to VECIN queue
+    _in_queue_a.EnQue(aLocal);
+    _in_queue_b.EnQue(bLocal);
+}
+
+template <typename T>
+__aicore__ inline void SwigluKernel<T>::compute(int64_t i) {
+    // Deque input tensors from VECIN queue
+    LocalTensor<T> aLocal = _in_queue_a.DeQue<T>();
+    LocalTensor<T> bLocal = _in_queue_b.DeQue<T>();
+    LocalTensor<T> cLocal = _out_queue_c.AllocTensor<T>();
+    // Call SwiGLU ascend api
+    SwiGLU<T, false>(cLocal, aLocal, bLocal, _beta_value, _copy_len);
+    // Enque result and free input
+    _out_queue_c.EnQue<T>(cLocal);
+    _in_queue_a.FreeTensor(aLocal);
+    _in_queue_b.FreeTensor(bLocal);
+}
+
+template <typename T>
+__aicore__ inline void SwigluKernel<T>::copyOut(int64_t i) {
+    // Deque output tensor from VECOUT queue
+    LocalTensor<T> cLocal = _out_queue_c.DeQue<T>();
+    auto batch_idx = _batch == 1 ? 0 : i / _seq_len;
+    auto seq_idx = _batch == 1 ? i : i % _seq_len;
+    int64_t idxc = batch_idx * _stride_batch_c + seq_idx * _stride_seq_c + _block_idx * _tile_len;
+    // Copy progress_th tile from local tensor to global tensor
+    if (_tile_len * sizeof(T) % BYTE_ALIGN != 0) {
+        DataCopyExtParams dcep = {1, static_cast<uint32_t>(_tile_len * sizeof(T)), 0, 0, 0};
+        DataCopyPad(_c_gm[idxc], cLocal, dcep);
+    } else {
+        DataCopy(_c_gm[idxc], cLocal, _tile_len);
+    }
+    // Free output Local tensor
+    _out_queue_c.FreeTensor(cLocal);
+}
+
+template <typename T>
+__aicore__ inline void SwigluKernel<T>::process() {
+    for (int64_t i = 0; i < _batch * _seq_len; ++i) {
+        copyIn(i);
+        compute(i);
+        copyOut(i);
+    }
+}
+
+#define DEFINE_SWIGLU_KERNEL(KERNEL_NAME, TYPE)                                    \
+    __global__ __aicore__ void KERNEL_NAME(GM_ADDR c, GM_ADDR a, GM_ADDR b,        \
+                                           int64_t batch, int64_t seq, int64_t hd, \
+                                           int64_t stride_batch_c,                 \
+                                           int64_t stride_batch_a,                 \
+                                           int64_t stride_batch_b,                 \
+                                           int64_t stride_seq_c,                   \
+                                           int64_t stride_seq_a,                   \
+                                           int64_t stride_seq_b) {                 \
+        SwigluKernel<TYPE> op;                                                     \
+        op.init(c, a, b,                                                           \
+                batch, seq, hd,                                                    \
+                stride_batch_c, stride_batch_a, stride_batch_b,                    \
+                stride_seq_c, stride_seq_a, stride_seq_b);                         \
+        op.process();                                                              \
+    }
+
+DEFINE_SWIGLU_KERNEL(swiglu_kernel_half, half)
+DEFINE_SWIGLU_KERNEL(swiglu_kernel_float, float)
+
+#undef DEFINE_SWIGLU_KERNEL
+
+extern "C" infiniStatus_t swiglu_kernel_launch(
+    void *c, void *a, void *b,
+    infiniDtype_t dtype, size_t batch, size_t seq, size_t hd,
+    ptrdiff_t stride_batch_c, ptrdiff_t stride_batch_a, ptrdiff_t stride_batch_b,
+    ptrdiff_t stride_seq_c, ptrdiff_t stride_seq_a, ptrdiff_t stride_seq_b, void *stream) {
+
+#define LAUNCH_SWIGLU_KERNEL(DTYPE_ENUM, KERNEL_NAME)       \
+    case DTYPE_ENUM:                                        \
+        KERNEL_NAME<<<BLOCK_NUM, nullptr, stream>>>(        \
+            c, a, b,                                        \
+            static_cast<int64_t>(batch),                    \
+            static_cast<int64_t>(seq),                      \
+            static_cast<int64_t>(hd),                       \
+            stride_batch_c, stride_batch_a, stride_batch_b, \
+            stride_seq_c, stride_seq_a, stride_seq_b);      \
+        return INFINI_STATUS_SUCCESS;
+
+    switch (dtype) {
+        LAUNCH_SWIGLU_KERNEL(INFINI_DTYPE_F16, swiglu_kernel_half)
+        LAUNCH_SWIGLU_KERNEL(INFINI_DTYPE_F32, swiglu_kernel_float)
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+#undef LAUNCH_SWIGLU_KERNEL
+}

src/infiniop/ops/swiglu/operator.cc

@@ -11,6 +11,9 @@
 #ifdef ENABLE_KUNLUN_API
 #include "kunlun/swiglu_kunlun.h"
 #endif
+#ifdef ENABLE_ASCEND_API
+#include "ascend/swiglu_ascend.h"
+#endif
 
 __C infiniStatus_t infiniopCreateSwiGLUDescriptor(
     infiniopHandle_t handle,
@@ -46,11 +49,8 @@ __C infiniStatus_t infiniopCreateSwiGLUDescriptor(
                                           c_desc, a_desc, b_desc);
     }
 #endif
-#ifdef ENABLE_ASCEND_NPU
-    case DevAscendNpu:
-        return ascendCreateSwiGLUDescriptor(
-            (AscendHandle_t)handle, (SwiGLUAscendDescriptor_t *)desc_ptr,
-            c_desc, a_desc, b_desc);
+#ifdef ENABLE_ASCEND_API
+        CREATE(INFINI_DEVICE_ASCEND, ascend);
 #endif
 #ifdef ENABLE_METAX_GPU
     case DevMetaxGpu: {
@@ -95,7 +95,7 @@ __C infiniStatus_t infiniopGetSwiGLUWorkspaceSize(infiniopSwiGLUDescriptor_t des
     }
 #endif
 #ifdef ENABLE_ASCEND_API
-    // GET(INFINI_DEVICE_ASCEND, ascend)
+        GET(INFINI_DEVICE_ASCEND, ascend)
 #endif
 #ifdef ENABLE_METAX_GPU
     case DevMetaxGpu: {
@@ -144,9 +144,8 @@ __C infiniStatus_t infiniopSwiGLU(
         return bangSwiGLU((SwiGLUBangDescriptor_t)desc, c, a, b, stream);
     }
 #endif
-#ifdef ENABLE_ASCEND_NPU
-    case DevAscendNpu:
-        return ascendSwiGLU((SwiGLUAscendDescriptor_t)desc, c, a, b, stream);
+#ifdef ENABLE_ASCEND_API
+        CALCULATE(INFINI_DEVICE_ASCEND, ascend);
 #endif
 #ifdef ENABLE_METAX_GPU
     case DevMetaxGpu:
@@ -188,9 +187,8 @@ infiniopDestroySwiGLUDescriptor(infiniopSwiGLUDescriptor_t desc) {
         return bangDestroySwiGLUDescriptor((SwiGLUBangDescriptor_t)desc);
     }
 #endif
-#ifdef ENABLE_ASCEND_NPU
-    case DevAscendNpu:
-        return ascendDestroySwiGLUDescriptor((SwiGLUAscendDescriptor_t)desc);
+#ifdef ENABLE_ASCEND_API
+        DELETE(INFINI_DEVICE_ASCEND, ascend)
 #endif
 #ifdef ENABLE_METAX_GPU
     case DevMetaxGpu:

xmake/ascend.lua

@@ -50,9 +50,8 @@ target("infiniop-ascend")
     add_files("$(projectdir)/src/infiniop/devices/ascend/*.cc", "$(projectdir)/src/infiniop/ops/*/ascend/*.cc")
 
     -- Add operator
-    -- TODO: add it back after ascend-kernels is fixed
-    -- add_rules("ascend-kernels")
-    -- add_links(builddir.."/libascend_kernels.a")
+    add_rules("ascend-kernels")
+    add_links(builddir.."/libascend_kernels.a")
 target_end()
 
 target("infinirt-ascend")