issue/11: add random sample ascend

src/infiniop/devices/ascend/CMakeLists.txt

@@ -3,9 +3,21 @@ 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")
+if(DEFINED ENV{USER} AND "$ENV{ASCEND_HOME}" STREQUAL "root")
+    set(DEFAULT_ASCEND_CANN_PACKAGE_PATH "/usr/local/Ascend/ascend-toolkit/latest" CACHE PATH "ASCEND CANN package default installation directory for root user")
+else()
+    set(DEFAULT_ASCEND_CANN_PACKAGE_PATH "$ENV{HOME}/Ascend/ascend-toolkit/latest" CACHE PATH "ASCEND CANN package default installation directory for other user")
+endif()
+
+if(DEFINED ASCEND_CANN_PACKAGE_PATH)
+elseif(DEFINED ENV{ASCEND_HOME_PATH})
+    set(ASCEND_CANN_PACKAGE_PATH "$ENV{ASCEND_HOME_PATH}" CACHE PATH "ASCEND CANN package installation directory" FORCE)
+else()
+    set(ASCEND_CANN_PACKAGE_PATH "${DEFAULT_ASCEND_CANN_PACKAGE_PATH}" CACHE PATH "ASCEND CANN package installation directory")
+endif()
+
 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_BUILD_TYPE "Debug" CACHE STRING "Build type Release/Debug (default Debug)" FORCE)
 set(CMAKE_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/out" CACHE STRING "path for install()" FORCE)
 
 if(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/tools/tikcpp/ascendc_kernel_cmake)
@@ -20,9 +32,11 @@ endif()
 
 include(${ASCENDC_CMAKE_DIR}/ascendc.cmake)
 
+
 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_kernel.cpp
+# ../../ops/rotary_embedding/ascend/rotary_embedding_kernel.cpp
+../../ops/random_sample/ascend/random_sample_kernel.cpp
 )
 
+target_include_directories(ascend_kernels PRIVATE ../../../../include)

src/infiniop/ops/random_sample/ascend/random_sample_ascend.cc

@@ -0,0 +1,148 @@
+#include "random_sample_ascend.h"
+
+InfiniopRandomSampleAscendDescriptor::InfiniopRandomSampleAscendDescriptor(infiniDevice_t device_) {
+    device = device_;
+    device_id = 0;
+    pDesc = new aclnnTensorDescriptor();
+    topkIdxDesc = new aclnnTensorDescriptor();
+    topkValDesc = new aclnnTensorDescriptor();
+    resDesc = new aclnnTensorDescriptor();
+}
+
+infiniopStatus_t ascendCreateRandomSampleDescriptor(infiniopAscendHandle_t handle,
+                                                    infiniopRandomSampleAscendDescriptor_t *desc_ptr,
+                                                    infiniopTensorDescriptor_t result,
+                                                    infiniopTensorDescriptor_t probs) {
+    if (probs->ndim != 1) {
+        return INFINIOP_STATUS_BAD_TENSOR_SHAPE;
+    }
+    if (result->ndim != 1 && result->shape[0] != 1) {
+        return INFINIOP_STATUS_BAD_TENSOR_SHAPE;
+    }
+
+    (*desc_ptr) = new InfiniopRandomSampleAscendDescriptor(handle->device);
+    (*desc_ptr)->device_id = handle->device_id;
+
+    CHECK_STATUS((*desc_ptr)->pDesc->fromInfiniOpTensorDescriptor(probs), INFINIOP_STATUS_SUCCESS);
+    CHECK_STATUS((*desc_ptr)->resDesc->fromInfiniOpTensorDescriptor(result), INFINIOP_STATUS_SUCCESS);
+    // Ascend aclnnTopk doesn't support U64 type
+    (*desc_ptr)->resDesc->dataType = aclDataType::ACL_INT64;
+
+    return INFINIOP_STATUS_SUCCESS;
+}
+
+infiniopStatus_t ascendGetRandomSampleWorkspaceSize(infiniopRandomSampleAscendDescriptor_t desc,
+                                                    uint64_t *size) {
+    auto &pDesc = desc->pDesc;
+    *size = numElements(pDesc->shape.data(), pDesc->ndim) * aclDataTypeSize(pDesc->dataType)
+          + numElements(pDesc->shape.data(), pDesc->ndim) * infiniSizeof(infiniDtype_t::INFINI_DTYPE_I64);
+
+    return INFINIOP_STATUS_SUCCESS;
+}
+
+infiniopStatus_t ascendRandomSample(infiniopRandomSampleAscendDescriptor_t desc,
+                                    void *workspace,
+                                    uint64_t workspace_size,
+                                    void *result,
+                                    void const *probs,
+                                    float random_val,
+                                    float topp,
+                                    int topk,
+                                    float temperature,
+                                    void *stream) {
+    if (topk <= 0 || topp < 0 || topp > 1.0) {
+        return INFINIOP_STATUS_BAD_PARAM;
+    }
+
+    if (random_val < 0 || random_val > 1.0) {
+        return INFINIOP_STATUS_BAD_PARAM;
+    }
+
+    auto &pDesc = desc->pDesc;
+    auto &topkIdxDesc = desc->topkIdxDesc;
+    auto &topkValDesc = desc->topkValDesc;
+    auto ndim = static_cast<int64_t>(pDesc->ndim);
+    auto voc = pDesc->shape[0];
+    auto topk_ = topk <= voc ? topk : voc;
+    bool doSample = topk_ > 1 && temperature != 0 && topp != 0;
+
+    auto topkShape = std::vector<int64_t>(pDesc->shape);
+    topkShape[ndim - 1] = doSample ? topk_ : 1;
+
+    auto topkStrides = std::vector<int64_t>(pDesc->strides);
+    // Infer contiguous strides
+    topkStrides[ndim - 1] = 1;
+    for (int64_t i = ndim - 2; i >= 0; --i) {
+        topkStrides[i] = topkStrides[i + 1] * topkShape[i + 1];
+    }
+
+    CHECK_STATUS(topkValDesc->setDescriptor(pDesc->dataType, topkShape, topkStrides), INFINIOP_STATUS_SUCCESS);
+    CHECK_STATUS(topkIdxDesc->setDescriptor(aclDataType::ACL_INT64, topkShape, topkStrides), INFINIOP_STATUS_SUCCESS);
+
+    // Infer data ptr
+    auto workspaceTmp = workspace;
+    auto topkValAddr = workspaceTmp;
+    workspaceTmp = (void *)((uint8_t *)workspace + numElements(topkValDesc->shape.data(), topkValDesc->ndim) * aclDataTypeSize(topkValDesc->dataType));
+    auto topkIdxAddr = workspaceTmp;
+    auto pAddr = (void *)probs;
+
+    // Create aclTensor
+    CHECK_STATUS(pDesc->createTensor(pAddr), INFINIOP_STATUS_SUCCESS);
+    CHECK_STATUS(topkValDesc->createTensor(topkValAddr), INFINIOP_STATUS_SUCCESS);
+    CHECK_STATUS(topkIdxDesc->createTensor(topkIdxAddr), INFINIOP_STATUS_SUCCESS);
+    if (!doSample) {
+        CHECK_STATUS(desc->resDesc->createTensor(result), INFINIOP_STATUS_SUCCESS);
+    }
+
+    // Do Topk calculate
+    uint64_t topkWorkspaceSize = 0;
+    aclOpExecutor *topkExecutor = nullptr;
+    auto ret = aclnnTopkGetWorkspaceSize(pDesc->t,
+                                         topkShape[ndim - 1],
+                                         ndim - 1,
+                                         true,
+                                         true,
+                                         topkValDesc->t,
+                                         doSample ? topkIdxDesc->t
+                                                  : desc->resDesc->t,
+                                         &topkWorkspaceSize,
+                                         &topkExecutor);
+    CHECK_RET(ret == ACL_SUCCESS,
+              LOG_PRINT("aclnnTopkGetWorkspaceSize failed ERROR: %d\n", ret);
+              return INFINIOP_STATUS_INTERNAL_ERROR);
+    void *topkWorkspace;
+    CHECK_STATUS(mallocWorkspace(&topkWorkspace, topkWorkspaceSize), INFINIOP_STATUS_SUCCESS);
+    ret = aclnnTopk(topkWorkspace,
+                    topkWorkspaceSize,
+                    topkExecutor,
+                    stream);
+    CHECK_RET(ret == ACL_SUCCESS,
+              LOG_PRINT("aclnnTopk failed ERROR: %d\n", ret);
+              return INFINIOP_STATUS_INTERNAL_ERROR);
+    CHECK_STATUS(freeWorkspace(topkWorkspace), INFINIOP_STATUS_SUCCESS);
+
+    if (doSample) {
+        // Do softmax and topp random sample
+        random_sample_do(
+            pAddr,
+            result,
+            topkValAddr,
+            topkIdxAddr,
+            topk,
+            static_cast<int>(pDesc->shape[0]),
+            topp,
+            temperature,
+            random_val,
+            pDesc->dataType,
+            stream);
+    }
+    return INFINIOP_STATUS_SUCCESS;
+}
+
+infiniopStatus_t ascendDestroyRandomSampleDescriptor(infiniopRandomSampleAscendDescriptor_t desc) {
+    delete desc->pDesc;
+    delete desc->topkIdxDesc;
+    delete desc->topkValDesc;
+    delete desc;
+    return INFINIOP_STATUS_SUCCESS;
+}

src/infiniop/ops/random_sample/ascend/random_sample_ascend.h

@@ -0,0 +1,25 @@
+#ifndef __RANDOM_SAMPLE_ASCEND_H__
+#define __RANDOM_SAMPLE_ASCEND_H__
+
+#include "../../../devices/ascend/tensor_aclnn.h"
+#include "../../utils.h"
+#include "random_sample_ascend_api.h"
+#include <acl/acl.h>
+#include <acl/acl_base.h>
+#include <acl/acl_rt.h>
+#include <aclnnop/aclnn_topk.h>
+
+struct InfiniopRandomSampleAscendDescriptor {
+    infiniDevice_t device;
+    int device_id;
+    aclnnTensorDescriptor_t pDesc, topkValDesc, topkIdxDesc, resDesc;
+
+    InfiniopRandomSampleAscendDescriptor(infiniDevice_t device_);
+};
+
+extern "C" void
+random_sample_do(void *p, void *res, void *topkAddr, void *topkIdxAddr,
+                 int32_t topk, int32_t voc, float topp, float temper,
+                 float random, int dtype, void *stream);
+
+#endif

src/infiniop/ops/random_sample/ascend/random_sample_ascend_api.h

@@ -0,0 +1,31 @@
+#ifndef __RANDOM_SAMPLE_ASCEND_API_H__
+#define __RANDOM_SAMPLE_ASCEND_API_H__
+
+#include "../../../devices/ascend/ascend_handle.h"
+#include "infiniop/operator.h"
+
+struct InfiniopRandomSampleAscendDescriptor;
+typedef struct InfiniopRandomSampleAscendDescriptor *infiniopRandomSampleAscendDescriptor_t;
+
+infiniopStatus_t ascendCreateRandomSampleDescriptor(infiniopAscendHandle_t handle,
+                                                    infiniopRandomSampleAscendDescriptor_t *desc_ptr,
+                                                    infiniopTensorDescriptor_t results,
+                                                    infiniopTensorDescriptor_t probs);
+
+infiniopStatus_t ascendGetRandomSampleWorkspaceSize(infiniopRandomSampleAscendDescriptor_t desc,
+                                                    uint64_t *size);
+
+infiniopStatus_t ascendRandomSample(infiniopRandomSampleAscendDescriptor_t desc,
+                                    void *workspace,
+                                    uint64_t workspace_size,
+                                    void *result,
+                                    void const *probs,
+                                    float random_val,
+                                    float topp,
+                                    int topk,
+                                    float temperature,
+                                    void *stream);
+
+infiniopStatus_t ascendDestroyRandomSampleDescriptor(infiniopRandomSampleAscendDescriptor_t desc);
+
+#endif