fix the issue of batch gemm

codeplaysoftware · Jul 18, 2024 · 7739df6 · 7739df6
1 parent 962766b
commit 7739df6
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 9 deletions.
diff --git a/examples/sycl/pvc/pvc_gemm.cpp b/examples/sycl/pvc/pvc_gemm.cpp
@@ -491,7 +491,6 @@ void collective_gemm(int M, int K, int N, int L = 1) {
 }
 
 int main() {
-  auto gmem_size = syclcompat::get_current_device().get_global_mem_size();
   collective_gemm<256, 256, 32, 64, 32>(4096, 4096, 4096);
   collective_gemm<256, 256, 32, 64, 32>(8192, 8192, 8192);
   collective_gemm<256, 256, 32, 64, 32>(1, 5120, 13824);

diff --git a/include/cute/atom/copy_traits_xe.hpp b/include/cute/atom/copy_traits_xe.hpp
@@ -51,7 +51,7 @@ struct XE_2D_LD_Unpack {
         int W = size<1>(traits.tensor)
                 * sizeof(typename Copy_Traits::CopyInternalType);
         auto [y, x, z] = src.data().coord_;
-        CopyOp::copy(traits.tensor.data() + z * W * H / sizeof(typename Copy_Traits::CopyInternalType), W, H, W, intel::coord_t {x, y},
+        CopyOp::copy(traits.tensor.data() + z, W, H, W, intel::coord_t {x, y},
                 &*dst.data());
     }
 
@@ -99,7 +99,7 @@ struct XE_2D_PF_Unpack {
         int H = size<0>(traits.tensor);
         int W = size<1>(traits.tensor) * sizeof(T);
         auto [y, x, z] = src.data().coord_;
-        CopyOp::template copy<T>(traits.tensor.data() + z * W * H / sizeof(T), W, H, W,
+        CopyOp::template copy<T>(traits.tensor.data() + z, W, H, W,
                 intel::coord_t {static_cast<int>(x), static_cast<int>(y)});
     }
 };
@@ -416,7 +416,7 @@ struct XE_2D_ST_Unpack {
                 * sizeof(typename Copy_Traits::CopyInternalType);
         auto [y, x, z] = dst.data().coord_;
 
-        CopyOp::copy(traits.tensor.data() + z * W * H / sizeof(typename Copy_Traits::CopyInternalType), W, H, W, intel::coord_t {x, y},
+        CopyOp::copy(traits.tensor.data() + z, W, H, W, intel::coord_t {x, y},
                 &*src.data());
     }
 

diff --git a/include/cutlass/gemm/kernel/intel_pvc_gemm.hpp b/include/cutlass/gemm/kernel/intel_pvc_gemm.hpp
@@ -239,7 +239,7 @@ class GemmUniversal<
     const auto tile_coord = make_coord(m_coord, n_coord, _, l_coord);
 
     Tensor tAi = params.mainloop.gmem_tiled_copy_a.get_pvc_tensor(
-        make_coord(m_coord, 0, l_coord), make_shape(_1{}, K, _1{}),
+        make_coord(m_coord, 0, 0), make_shape(_1{}, K, L),
         make_stride(Int<FragsM * DpasM>{}, _1{}));
     constexpr int version =
         is_same_v<typename CollectiveMainloop::GmemTiledCopyB,
@@ -248,8 +248,8 @@ class GemmUniversal<
             : 2;
 
     Tensor tBi = params.mainloop.gmem_tiled_copy_b.get_pvc_tensor(
-        make_coord(0, n_coord, l_coord),
-        make_shape(K, Int<FragsN / version>{}, _1{}),
+        make_coord(0, n_coord, 0),
+        make_shape(K, Int<FragsN / version>{}, L),
         make_stride(_1{}, Int<version * DpasN>{}));
 
     // Compute tile residues for predication
@@ -271,8 +271,8 @@ class GemmUniversal<
     CollectiveMainloop collective_mma;
     collective_mma(
       accumulators,
-      tAi(_,_,_,0),
-      tBi(_,_,_,0),
+      tAi(_,_,_,l_coord),
+      tBi(_,_,_,l_coord),
       accumulators,
       k_tile_iter, k_tile_count,
       residue_mnk,