InfiniTensor · PanZezhong1725 · Oct 16, 2024 · Oct 16, 2024 · Oct 16, 2024 · Oct 16, 2024
diff --git a/src/ops/matmul/bang/matmul_cnnl.cc b/src/ops/matmul/bang/matmul_cnnl.cc
@@ -3,17 +3,18 @@
 #include "../../../devices/bang/handle_pool.h"
 #include "../../utils.h"
 #include "cnrt.h"
-
-MatmulBangDescriptor::MatmulBangDescriptor(Device device) {
-    this->device = device;
-    get_cnnl_pool();
-}
-
-void matmul_cnnl_f16(Tensor c, float beta, Tensor a, Tensor b, float alpha, void *stream) {
-    auto info = MatmulInfo(c, a, b, false);
-
-    int32_t use_stride = true;
-
+infiniopStatus_t bangCreateMatmulDescriptor(BangHandle_t handle,
+                                            MatmulBangDescriptor_t *desc_ptr,
+                                            infiniopTensorDescriptor_t c_desc,
+                                            float alpha,
+                                            infiniopTensorDescriptor_t a_desc,
+                                            infiniopTensorDescriptor_t b_desc,
+                                            float beta) {
+    infiniopStatus_t *status = new infiniopStatus_t{STATUS_EXECUTION_FAILED};
+    auto info = MatmulInfo(c_desc, a_desc, b_desc, status);
+    if (*status != STATUS_SUCCESS) {
+        return *status;
+    }
     cnnlTensorDescriptor_t aDesc, bDesc, cDesc;
     cnnlCreateTensorDescriptor(&aDesc);
     cnnlCreateTensorDescriptor(&bDesc);
@@ -29,35 +30,73 @@ void matmul_cnnl_f16(Tensor c, float beta, Tensor a, Tensor b, float alpha, void
     cnnlMatMulDescCreate(&opDesc);
     cnnlMatMulAlgoCreate(&algo);
     cnnlCreateMatMulHeuristicResult(&algoResult);
-
+    int32_t use_stride = true;
     cnnlSetMatMulDescAttr(opDesc, CNNL_MATMUL_USE_STRIDE, &use_stride,
                           sizeof(int32_t));
+    *desc_ptr = new MatmulBangDescriptor{
+        handle->device,
+        handle->device_id,
+        info,
+        alpha,
+        beta,
+        c_desc->dt,
+        handle->cnnl_handles,
+        aDesc,
+        bDesc,
+        cDesc,
+        opDesc,
+        algo,
+        algoResult};
+    return STATUS_SUCCESS;
+}
+infiniopStatus_t bangGetMatmulWorkspaceSize(MatmulBangDescriptor_t desc, uint64_t *size) {
+    *size = 0;
+    return STATUS_SUCCESS;
+}
 
+infiniopStatus_t bangDestroyMatmulDescriptor(MatmulBangDescriptor_t desc) {
+    desc->cnnl_handles = nullptr;
+    cnnlDestroyTensorDescriptor(desc->aDesc);
+    cnnlDestroyTensorDescriptor(desc->bDesc);
+    cnnlDestroyTensorDescriptor(desc->cDesc);
+    cnnlMatMulDescDestroy(desc->opDesc);
+    cnnlMatMulAlgoDestroy(desc->algo);
+    cnnlDestroyMatMulHeuristicResult(desc->algoResult);
+    delete desc;
+    return STATUS_SUCCESS;
+}
 
-    void *workspace;
+void matmul_cnnl_f16(MatmulBangDescriptor_t desc, void *workspace, void *c, float beta, void const *a, void const *b, float alpha, void *stream) {
+    auto info = desc->info;
+    if (info.is_transed) {
+        std::swap(a, b);
+    }
 
-    use_cnnl((cnrtQueue_t) stream,
+    use_cnnl(desc->cnnl_handles, desc->device_id, (cnrtQueue_t) stream,
              [&](cnnlHandle_t handle) {
                  int count = 0;
-                 cnnlGetBatchMatMulAlgoHeuristic(handle, opDesc, aDesc,
-                                                 bDesc, cDesc,
-                                                 NULL, 1, &algoResult, &count);
+                 cnnlGetBatchMatMulAlgoHeuristic(handle, desc->opDesc, desc->aDesc,
+                                                 desc->bDesc, desc->cDesc,
+                                                 NULL, 1, &desc->algoResult, &count);
                  size_t wsSize;
-                 cnnlGetBatchMatMulHeuristicResult(algoResult, algo, &wsSize);
+                 cnnlGetBatchMatMulHeuristicResult(desc->algoResult, desc->algo, &wsSize);
                  cnrtMalloc(&workspace, wsSize);
-                 cnnlBatchMatMulBCast_v2(handle, opDesc, algo,
-                                         &alpha, aDesc, info.a_ptr,
-                                         bDesc, info.b_ptr,
-                                         &beta, cDesc, info.c_ptr,
+                 cnnlBatchMatMulBCast_v2(handle, desc->opDesc, desc->algo,
+                                         &alpha, desc->aDesc, a,
+                                         desc->bDesc, b,
+                                         &beta, desc->cDesc, c,
                                          workspace, wsSize);
              });
-
-    cnrtFree(workspace);
-
-    cnnlDestroyTensorDescriptor(aDesc);
-    cnnlDestroyTensorDescriptor(bDesc);
-    cnnlDestroyTensorDescriptor(cDesc);
-    cnnlMatMulDescDestroy(opDesc);
-    cnnlMatMulAlgoDestroy(algo);
-    cnnlDestroyMatMulHeuristicResult(algoResult);
+}
+infiniopStatus_t bangMatmul(MatmulBangDescriptor_t desc, void *workspace, uint64_t workspace_size, void *c, void const *a, void const *b, void *stream) {
+    if (cnrtSetDevice(desc->device_id) != cnrtSuccess) {
+        return STATUS_BAD_DEVICE;
+    }
+    float alpha = desc->alpha;
+    float beta = desc->beta;
+    if (dtype_eq(desc->dtype, F16)) {
+        matmul_cnnl_f16(desc, workspace, c, beta, a, b, alpha, stream);
+        return STATUS_SUCCESS;
+    }
+    return STATUS_BAD_TENSOR_DTYPE;
 }
diff --git a/src/ops/matmul/bang/matmul_cnnl.h b/src/ops/matmul/bang/matmul_cnnl.h
@@ -1,15 +1,41 @@
 #ifndef __CNNL_MATMUL_H__
 #define __CNNL_MATMUL_H__
-
+#include "../../../devices/bang/bang_handle.h"
 #include "../blas.h"
 #include "cnnl.h"
 #include "cnnl_extra.h"
 #include "operators.h"
 
 struct MatmulBangDescriptor {
     Device device;
-    MatmulBangDescriptor(Device device);
+    int device_id;
+    MatmulInfo info;
+    float alpha;
+    float beta;
+    DT dtype;
+    std::shared_ptr<Pool<cnnlHandle_t>> cnnl_handles;
+    cnnlTensorDescriptor_t aDesc;
+    cnnlTensorDescriptor_t bDesc;
+    cnnlTensorDescriptor_t cDesc;
+    cnnlMatMulDescriptor_t opDesc;
+    cnnlMatMulAlgo_t algo;
+    cnnlMatMulHeuristicResult_t algoResult;
 };
+typedef struct MatmulBangDescriptor *MatmulBangDescriptor_t;
+
+infiniopStatus_t bangCreateMatmulDescriptor(BangHandle_t handle,
+                                            MatmulBangDescriptor_t *desc_ptr,
+                                            infiniopTensorDescriptor_t c_desc,
+                                            float alpha,
+                                            infiniopTensorDescriptor_t a_desc,
+                                            infiniopTensorDescriptor_t b_desc,
+                                            float beta);
+
+infiniopStatus_t bangGetMatmulWorkspaceSize(MatmulBangDescriptor_t desc, uint64_t *size);
+
+infiniopStatus_t bangMatmul(MatmulBangDescriptor_t desc, void *workspace, uint64_t workspace_size, void *c, void const *a, void const *b, void *stream);
+
+infiniopStatus_t bangDestroyMatmulDescriptor(MatmulBangDescriptor_t desc);
 
 inline void setMatrixTensorEx(cnnlTensorDescriptor_t desc, const BlasMatrix &matrix, bool trans = false) {
     int ndim = matrix.ndim;
@@ -33,6 +59,5 @@ inline void setMatrixTensorEx(cnnlTensorDescriptor_t desc, const BlasMatrix &mat
     }
 }
 
-void matmul_cnnl_f16(Tensor c, float beta, Tensor a, Tensor b, float alpha, void *stream);
 
 #endif// __CNNL_MATMUL_H__
diff --git a/src/ops/matmul/operator.cc b/src/ops/matmul/operator.cc
@@ -30,7 +30,9 @@ __C infiniopStatus_t infiniopCreateMatmulDescriptor(infiniopHandle_t handle,
         }
 #endif
 #ifdef ENABLE_CAMBRICON_MLU
-        // TODO
+        case DevCambriconMlu: {
+            return bangCreateMatmulDescriptor((BangHandle_t) handle, (MatmulBangDescriptor_t *) desc_ptr, c_desc, alpha, a_desc, b_desc, beta);
+        }
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -49,7 +51,9 @@ __C infiniopStatus_t infiniopGetMatmulWorkspaceSize(infiniopMatmulDescriptor_t d
 
 #endif
 #ifdef ENABLE_CAMBRICON_MLU
-        // TODO
+        case DevCambriconMlu: {
+            return bangGetMatmulWorkspaceSize((MatmulBangDescriptor_t) desc, size);
+        }
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -66,7 +70,9 @@ __C infiniopStatus_t infiniopMatmul(infiniopMatmulDescriptor_t desc, void *works
             return cudaMatmul((MatmulCudaDescriptor_t) desc, workspace, workspace_size, c, a, b, stream);
 #endif
 #ifdef ENABLE_CAMBRICON_MLU
-            // TODO
+        case DevCambriconMlu: {
+            return bangMatmul((MatmulBangDescriptor_t) desc, workspace, workspace_size, c, a, b, stream);
+        }
 #endif
     }
     return STATUS_BAD_DEVICE;
@@ -85,7 +91,9 @@ infiniopStatus_t infiniopDestroyMatmulDescriptor(infiniopMatmulDescriptor_t desc
 
 #endif
 #ifdef ENABLE_CAMBRICON_MLU
-        // TODO
+        case DevCambriconMlu: {
+            return bangDestroyMatmulDescriptor((MatmulBangDescriptor_t) desc);
+        }
 #endif
     }
     return STATUS_BAD_DEVICE;

diff --git a/src/ops/utils.h b/src/ops/utils.h
@@ -86,7 +86,7 @@ inline bool getBroadcastShape(const uint64_t *shape1, uint64_t ndim1,
     std::copy(shape2, shape2 + ndim2, padded_shape2 + max_rank - ndim2);
 
     // compute broadcasted shape
-    for (int i = 0; i < max_rank; ++i) {
+    for (size_t i = 0; i < max_rank; ++i) {
         if (padded_shape1[i] == padded_shape2[i] || padded_shape1[i] == 1 || padded_shape2[i] == 1) {
             broadcast_shape[i] = std::max(padded_shape1[i], padded_shape2[i]);
         } else {