save code

Author: xiaolil1

csrc/xpu_cutlass_fusion.cpp

@@ -262,22 +262,21 @@ class gemm_4bit_cutlass_kernel {
     Tensor frag_copy_B_b = thr_copy_B.retile_D(dequant_frag_b);
     Tensor frag_copy_Scale_b = thr_copy_scale.retile_D(fragment_scale_b);
 
-//    cute::tuple<decltype(mma_A_a), decltype(mma_A_b)> mma_A(mma_A_a, mma_A_b);
-//    cute::tuple<decltype(mma_B_a), decltype(mma_B_b)> mma_B(mma_B_a, mma_B_b);
-//    cute::tuple<decltype(dequant_frag_a), decltype(dequant_frag_b)> dequant_frag(dequant_frag_a, dequant_frag_b);
-//    cute::tuple<decltype(fragment_scale_a), decltype(fragment_scale_b)> fragment_scale(fragment_scale_a, fragment_scale_b);
-//    cute::tuple<decltype(frag_copy_A_a), decltype(frag_copy_A_b)> frag_copy_A(frag_copy_A_a, frag_copy_A_b);
-//    cute::tuple<decltype(frag_copy_B_a), decltype(frag_copy_B_b)> frag_copy_B(frag_copy_B_a, frag_copy_B_b);
-//    cute::tuple<decltype(frag_copy_Scale_a), decltype(frag_copy_Scale_b)> frag_copy_Scale(frag_copy_Scale_a, frag_copy_Scale_b);
-////auto& mma_A_0 = cute::get<0>(mma_A_tuple);  // 引用 mma_A_a
-////auto& mma_A_1 = cute::get<1>(mma_A_tuple);  // 引用 mma_A_b
-    decltype(mma_A_a)* mma_A[] = {&mma_A_a, &mma_A_b};
-    decltype(mma_B_a)* mma_B[] = {&mma_B_a, &mma_B_b};
-    decltype(dequant_frag_a)* dequant_frag[] = {&dequant_frag_a, &dequant_frag_b};
-    decltype(fragment_scale_a)* fragment_scale[] = {&fragment_scale_a, &fragment_scale_b};
-    decltype(frag_copy_A_a)* frag_copy_A[] = {&frag_copy_A_a, &frag_copy_A_b};
-    decltype(frag_copy_B_a)* frag_copy_B[] = {&frag_copy_B_a, &frag_copy_B_b};
-    decltype(frag_copy_Scale_a)* frag_copy_Scale[] = {&frag_copy_Scale_a, &frag_copy_Scale_b};
+    //decltype(mma_A_a)* mma_A[] = {&mma_A_a, &mma_A_b};
+    //decltype(mma_B_a)* mma_B[] = {&mma_B_a, &mma_B_b};
+    //decltype(dequant_frag_a)* dequant_frag[] = {&dequant_frag_a, &dequant_frag_b};
+    //decltype(fragment_scale_a)* fragment_scale[] = {&fragment_scale_a, &fragment_scale_b};
+    //decltype(frag_copy_A_a)* frag_copy_A[] = {&frag_copy_A_a, &frag_copy_A_b};
+    //decltype(frag_copy_B_a)* frag_copy_B[] = {&frag_copy_B_a, &frag_copy_B_b};
+    //decltype(frag_copy_Scale_a)* frag_copy_Scale[] = {&frag_copy_Scale_a, &frag_copy_Scale_b};
+
+    cute::array<decltype(mma_A_a), 2> mma_A = {mma_A_a, mma_A_b};
+    cute::array<decltype(mma_B_a), 2> mma_B = {mma_B_a, mma_B_b};
+    cute::array<decltype(dequant_frag_a), 2> dequant_frag = {dequant_frag_a, dequant_frag_b};
+    cute::array<decltype(fragment_scale_a), 2> fragment_scale = {fragment_scale_a, fragment_scale_b};
+    cute::array<decltype(frag_copy_A_a), 2> frag_copy_A = {frag_copy_A_a, frag_copy_A_b};
+    cute::array<decltype(frag_copy_B_a), 2> frag_copy_B = {frag_copy_B_a, frag_copy_B_b};
+    cute::array<decltype(frag_copy_Scale_a), 2> frag_copy_Scale = {frag_copy_Scale_a, frag_copy_Scale_b};
 
 #endif    
     Tensor tAgA = thr_copy_A.retile_S(tCgA);
@@ -315,8 +314,8 @@ class gemm_4bit_cutlass_kernel {
 
 #if 1
     auto dequant = [&](int start_lut_id, const int buffer_idx) {
-      constexpr int N = decltype(cute::size<1>(*mma_B[buffer_idx]))::value;
-      constexpr int K = decltype(cute::size(*mma_B[buffer_idx]))::value / N;
+      constexpr int N = decltype(cute::size<1>(mma_B[buffer_idx]))::value;
+      constexpr int K = decltype(cute::size(mma_B[buffer_idx]))::value / N;
 
       using src_compress_type = uint32_t;
       using dst_compress_type = uint32_t;
@@ -340,13 +339,14 @@ class gemm_4bit_cutlass_kernel {
 
           #pragma unroll
           for (int v = 0; v < src_vec_size; v++) {
-            src_compress_type src_value = reinterpret_cast<sycl::vec<src_compress_type, src_vec_size>*>(cute::raw_pointer_cast(dequant_frag[buffer_idx]->data()))[n*src_loop_num + l][v];
+            //src_compress_type src_value = reinterpret_cast<sycl::vec<src_compress_type, src_vec_size>*>(cute::raw_pointer_cast(dequant_frag[buffer_idx]->data()))[n*src_loop_num + l][v];
+            src_compress_type src_value = reinterpret_cast<sycl::vec<src_compress_type, src_vec_size>*>(cute::raw_pointer_cast(dequant_frag[buffer_idx].data()))[n*src_loop_num + l][v];
             int dst_base_idx = l * src_vec_size * src_compress_size + v * src_compress_size;
 
             #pragma unroll
             for (int c = 0; c < src_compress_size; c++) {
               uint8_t bit_value = (src_value >> (4 * (((c + 1) & 1) + (c >> 1) * 2))) & 0xF;
-              float scale_value = (*fragment_scale[buffer_idx])((n * BLK_K  + dst_base_idx + c) >> (31 - std::countl_zero<unsigned int>(GROUP_SIZE))); 
+              float scale_value = (fragment_scale[buffer_idx])((n * BLK_K  + dst_base_idx + c) >> (31 - std::countl_zero<unsigned int>(GROUP_SIZE))); 
 
               dst[dst_base_idx + c] = static_cast<ElementMMA>(quant_map_[lut_id][bit_value] * scale_value);
               lut_id = (lut_id + 1) % LUT_NUM;
@@ -356,14 +356,14 @@ class gemm_4bit_cutlass_kernel {
 
         #pragma unroll
         for (int l = 0; l < dst_loop_num; l++) {
-          reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(cute::raw_pointer_cast(mma_B[buffer_idx]->data()))[n * dst_loop_num + l] = reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(dst)[l];
+          reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(cute::raw_pointer_cast(mma_B[buffer_idx].data()))[n * dst_loop_num + l] = reinterpret_cast<sycl::vec<dst_compress_type, dst_vec_size>*>(dst)[l];
         }
       }
     };
 
-    copy(params.tiled_copy_b, tBgB(_,_,_,k_start_idx), *frag_copy_B[0]); 
-    copy(params.tiled_copy_scale, tSgS(_,_,_,k_start_idx * BLK_K/params.group_size), *frag_copy_Scale[0]);
-    copy(params.tiled_copy_a, tAgA(_,_,_,k_start_idx), *frag_copy_A[0]);
+    copy(params.tiled_copy_b, tBgB(_,_,_,k_start_idx), frag_copy_B[0]); 
+    copy(params.tiled_copy_scale, tSgS(_,_,_,k_start_idx * BLK_K/params.group_size), frag_copy_Scale[0]);
+    copy(params.tiled_copy_a, tAgA(_,_,_,k_start_idx), frag_copy_A[0]);
 
     if (prefetch_k < k_tile_count) {
       prefetch(tiled_prefetch_a, pAgA(_,_,_,prefetch_k));
@@ -376,19 +376,19 @@ class gemm_4bit_cutlass_kernel {
 
       dequant(start_lut_id, 1 - buf_idx);
 
-      copy(params.tiled_copy_b, tBgB(_,_,_,k_tile), *frag_copy_B[buf_idx]);
-      copy(params.tiled_copy_scale, tSgS(_,_,_,(k_start_idx+k_s)*BLK_K/params.group_size), *frag_copy_Scale[buf_idx]);
-      copy(params.tiled_copy_a, tAgA(_,_,_,k_tile), *frag_copy_A[buf_idx]);
+      copy(params.tiled_copy_b, tBgB(_,_,_,k_tile), frag_copy_B[buf_idx]);
+      copy(params.tiled_copy_scale, tSgS(_,_,_,(k_start_idx+k_s)*BLK_K/params.group_size), frag_copy_Scale[buf_idx]);
+      copy(params.tiled_copy_a, tAgA(_,_,_,k_tile), frag_copy_A[buf_idx]);
 
       if (prefetch_k < k_tile_count) {
         prefetch(tiled_prefetch_a, pAgA(_,_,_,prefetch_k));
         prefetch(tiled_prefetch_b, pBgB(_,_,_,prefetch_k));
       }
 
-      cute::gemm(tiled_mma, *mma_A[1 - buf_idx], *mma_B[1 - buf_idx], accumulators);
+      cute::gemm(tiled_mma, mma_A[1 - buf_idx], mma_B[1 - buf_idx], accumulators);
       barrier_wait(3);
     }
-    cute::gemm(tiled_mma, *mma_A[1], *mma_B[1], accumulators);
+    cute::gemm(tiled_mma, mma_A[1], mma_B[1], accumulators);
 
 #else