issue/346: 增加CublasLt支持

Author: qinyiqun

src/infiniop/devices/nvidia/nvidia_common.cu

@@ -49,6 +49,18 @@ infiniStatus_t Handle::Internal::useCudnn(cudaStream_t stream, const Fn<cudnnHan
 }
 #endif
 
+#ifdef ENABLE_CUBLASLT_API
+infiniStatus_t Handle::Internal::useCublasLt(cudaStream_t stream, const Fn<cublasLtHandle_t> &f) const {
+    auto handle = blaslt_handles.pop();
+    if (!handle) {
+        CHECK_CUBLASLT(cublasLtCreate(&(*handle)));
+    }
+    CHECK_STATUS(f(*handle));
+    blaslt_handles.push(std::move(*handle));
+    return INFINI_STATUS_SUCCESS;
+}
+#endif
+
 int Handle::Internal::warpSize() const { return _warp_size; }
 int Handle::Internal::maxThreadsPerBlock() const { return _max_threads_per_block; }
 int Handle::Internal::blockSizeX() const { return _block_size[0]; }

src/infiniop/devices/nvidia/nvidia_handle.cuh

@@ -11,8 +11,13 @@
 #include <cudnn.h>
 #endif
 
+#ifdef ENABLE_CUBLASLT_API
+#include <cublasLt.h>
+#endif
+
 #define CHECK_CUBLAS(API) CHECK_INTERNAL(API, CUBLAS_STATUS_SUCCESS)
 #define CHECK_CUDNN(API) CHECK_INTERNAL(API, CUDNN_STATUS_SUCCESS)
+#define CHECK_CUBLASLT(API) CHECK_INTERNAL(API, CUBLAS_STATUS_SUCCESS)
 
 namespace device::nvidia {
 
@@ -21,6 +26,9 @@ class Handle::Internal {
 #ifdef ENABLE_CUDNN_API
     Pool<cudnnHandle_t> dnn_handles;
 #endif
+#ifdef ENABLE_CUBLASLT_API
+    Pool<cublasLtHandle_t> blaslt_handles;
+#endif
 
     int _warp_size,
         _max_threads_per_block,
@@ -37,6 +45,9 @@ public:
 #ifdef ENABLE_CUDNN_API
     infiniStatus_t useCudnn(cudaStream_t stream, const Fn<cudnnHandle_t> &f) const;
 #endif
+#ifdef ENABLE_CUBLASLT_API
+    infiniStatus_t useCublasLt(cudaStream_t stream, const Fn<cublasLtHandle_t> &f) const;
+#endif
 
     int warpSize() const;
     int maxThreadsPerBlock() const;

xmake.lua

@@ -66,6 +66,16 @@ if has_config("cudnn") then
     add_defines("ENABLE_CUDNN_API")
 end
 
+option("cublaslt")
+    set_default(true)
+    set_showmenu(true)
+    set_description("Whether to compile cublaslt for Nvidia GPU")
+option_end()
+
+if has_config("cublaslt") then
+    add_defines("ENABLE_CUBLASLT_API")
+end
+
 -- 寒武纪
 option("cambricon-mlu")
     set_default(false)
@@ -244,6 +254,20 @@ target("infiniop")
     if has_config("iluvatar-gpu") then
         add_deps("infiniop-iluvatar")
     end
+    if has_config("sugon-dcu") then
+        local builddir = string.format(
+            "build/%s/%s/%s",
+            get_config("plat"),
+            get_config("arch"),
+            get_config("mode")
+        )
+        add_shflags("-s", "-shared", "-fPIC")
+        add_links("cublas", "cublaslt", "cudnn", "cudadevrt", "cudart_static", "rt", "pthread", "dl")
+        -- Using -linfiniop-nvidia will fail, manually link the target using full path
+        add_deps("nv-gpu", {inherit = false})
+        add_links(builddir.."/libinfiniop-nvidia.a")
+        set_toolchains("sugon-dcu-linker")
+    end
 
     if has_config("cambricon-mlu") then
         add_deps("infiniop-cambricon")