Merge branch 'develop' into ck_tile/fa_asm_bwd

include/ck_tile/ops/fmha/pipeline/block_fmha_bwd_pipeline_default_policy.hpp

@@ -25,14 +25,15 @@ struct BlockFmhaBwdPipelineDefaultPolicy
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetQKBlockGemm()
     {
-        using BlockGemmProblem =
-            BlockGemmPipelineProblem<typename Problem::QDataType,
+        using BlockGemmProblem = BlockGemmPipelineProblem<
+            typename Problem::QDataType,
             typename Problem::KDataType,
             typename Problem::AccDataType,
-                                     Problem::kBlockSize,
-                                     TileGemmShape<Problem::BlockFmhaShape::kM0,
+            TileGemmShape<sequence<Problem::BlockFmhaShape::kM0,
                                    Problem::BlockFmhaShape::kN0,
-                                                   Problem::BlockFmhaShape::kK0>>;
+                                   Problem::BlockFmhaShape::kK0>,
+                          typename Problem::BlockFmhaShape::Gemm0BlockWarps,
+                          typename Problem::BlockFmhaShape::Gemm0WarpTile>>;
 
         using WarpGemm = WarpGemmMfmaDispatcher<
             typename Problem::QDataType,
@@ -57,14 +58,15 @@ struct BlockFmhaBwdPipelineDefaultPolicy
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetPTOGradTBlockGemm()
     {
-        using BlockGemmProblem =
-            BlockGemmPipelineProblem<typename Problem::GemmDataType,
+        using BlockGemmProblem = BlockGemmPipelineProblem<
+            typename Problem::GemmDataType,
             typename Problem::OGradDataType,
             typename Problem::AccDataType,
-                                     Problem::kBlockSize,
-                                     TileGemmShape<Problem::BlockFmhaShape::kN0,
+            TileGemmShape<sequence<Problem::BlockFmhaShape::kN0,
                                    Problem::BlockFmhaShape::kVHeaddim,
-                                                   Problem::BlockFmhaShape::kK1>>;
+                                   Problem::BlockFmhaShape::kK1>,
+                          typename Problem::BlockFmhaShape::Gemm1BlockWarps,
+                          typename Problem::BlockFmhaShape::Gemm1WarpTile>>;
 
         using WarpGemm =
             WarpGemmMfmaDispatcher<typename Problem::GemmDataType,
@@ -88,14 +90,15 @@ struct BlockFmhaBwdPipelineDefaultPolicy
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetOGradVBlockGemm()
     {
-        using BlockGemmProblem =
-            BlockGemmPipelineProblem<typename Problem::OGradDataType,
+        using BlockGemmProblem = BlockGemmPipelineProblem<
+            typename Problem::OGradDataType,
             typename Problem::VDataType,
             typename Problem::AccDataType,
-                                     Problem::kBlockSize,
-                                     TileGemmShape<Problem::BlockFmhaShape::kM0,
+            TileGemmShape<sequence<Problem::BlockFmhaShape::kM0,
                                    Problem::BlockFmhaShape::kN0,
-                                                   Problem::BlockFmhaShape::kK2>>;
+                                   Problem::BlockFmhaShape::kK2>,
+                          typename Problem::BlockFmhaShape::Gemm2BlockWarps,
+                          typename Problem::BlockFmhaShape::Gemm2WarpTile>>;
 
         using WarpGemm = WarpGemmMfmaDispatcher<
             typename Problem::OGradDataType,
@@ -120,14 +123,15 @@ struct BlockFmhaBwdPipelineDefaultPolicy
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSGradTQTBlockGemm()
     {
-        using BlockGemmProblem =
-            BlockGemmPipelineProblem<typename Problem::GemmDataType,
+        using BlockGemmProblem = BlockGemmPipelineProblem<
+            typename Problem::GemmDataType,
             typename Problem::QDataType,
             typename Problem::AccDataType,
-                                     Problem::kBlockSize,
-                                     TileGemmShape<Problem::BlockFmhaShape::kN0,
+            TileGemmShape<sequence<Problem::BlockFmhaShape::kN0,
                                    Problem::BlockFmhaShape::kQKHeaddim,
-                                                   Problem::BlockFmhaShape::kK3>>;
+                                   Problem::BlockFmhaShape::kK3>,
+                          typename Problem::BlockFmhaShape::Gemm3BlockWarps,
+                          typename Problem::BlockFmhaShape::Gemm3WarpTile>>;
 
         using WarpGemm =
             WarpGemmMfmaDispatcher<typename Problem::GemmDataType,
@@ -151,14 +155,15 @@ struct BlockFmhaBwdPipelineDefaultPolicy
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSGradKTBlockGemm()
     {
-        using BlockGemmProblem =
-            BlockGemmPipelineProblem<typename Problem::GemmDataType,
+        using BlockGemmProblem = BlockGemmPipelineProblem<
+            typename Problem::GemmDataType,
             typename Problem::KDataType,
             typename Problem::AccDataType,
-                                     Problem::kBlockSize,
-                                     TileGemmShape<Problem::BlockFmhaShape::kM0,
+            TileGemmShape<sequence<Problem::BlockFmhaShape::kM0,
                                    Problem::BlockFmhaShape::kQKHeaddim,
-                                                   Problem::BlockFmhaShape::kK4>>;
+                                   Problem::BlockFmhaShape::kK4>,
+                          typename Problem::BlockFmhaShape::Gemm4BlockWarps,
+                          typename Problem::BlockFmhaShape::Gemm4WarpTile>>;
 
         using WarpGemm =
             WarpGemmMfmaDispatcher<typename Problem::GemmDataType,

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp

@@ -215,8 +215,8 @@ struct BlockFmhaPipelineQRKSVS
 
         const auto num_total_loop = integer_divide_ceil(seqlen_k_end - seqlen_k_start, kN0);
 
-        // check early exit if masked and no work to do.
-        if constexpr(FmhaMask::IsMasking)
+        // check early exit if no work to do
+        if constexpr(FmhaMask::IsMasking || kPadSeqLenK)
         {
             if(num_total_loop <= 0)
             {

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp

@@ -268,7 +268,7 @@ struct BlockFmhaPipelineQRKSVSAsync
 
         const auto num_total_loop = integer_divide_ceil(seqlen_k_end - seqlen_k_start, kN0);
 
-        // check early exit
+        // check early exit if no work to do
         if constexpr(FmhaMask::IsMasking || kPadSeqLenK)
         {
             if(num_total_loop <= 0)

include/ck_tile/ops/gemm.hpp

@@ -21,6 +21,8 @@
 #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp"
 #include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_default_policy.hpp"
 #include "ck_tile/ops/gemm/block/block_gemm_problem.hpp"
+#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp"
+#include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp"
 #include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1.hpp"
 #include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp"
 #include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp"

include/ck_tile/ops/gemm/block/block_gemm_areg_bgmem_creg_v1.hpp

@@ -4,7 +4,8 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/gemm/block/block_gemm_areg_bsmem_creg_v1_default_policy.hpp"
+#include "ck_tile/ops/gemm/block/block_gemm_areg_bgmem_creg_v1_default_policy.hpp"
+#include "ck_tile/ops/gemm/block/block_gemm_problem.hpp"
 
 namespace ck_tile {
 
@@ -27,9 +28,9 @@ struct BlockGemmARegBGmemCRegV1
     static constexpr index_t kBlockSize = Problem::kBlockSize;
 
     // use BlockGemmARegBSmemCRegV1 as the underlying block-GEMM implementation
-    using BlockGemmARegBSmemCRegImpl = BlockGemmARegBSmemCRegV1<
+    using BlockGemmARegBGmemCRegImpl = BlockGemmARegBGmemCRegV1<
         BlockGemmProblem<ADataType, BDataType, CDataType, kBlockSize, BlockGemmShape>,
-        BlockGemmARegBSmemCRegV1DefaultPolicy>;
+        BlockGemmARegBGmemCRegV1DefaultPolicy>;
 
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetStaticLdsSize()
     {
@@ -82,7 +83,7 @@ struct BlockGemmARegBGmemCRegV1
         block_sync_lds();
 
         // block GEMM
-        BlockGemmARegBSmemCRegImpl{}(c_block_tensor, a_block_tensor, b_block_smem_window);
+        BlockGemmARegBGmemCRegImpl{}(c_block_tensor, a_block_tensor, b_block_smem_window);
     }
 
     // C = A * B
@@ -128,7 +129,7 @@ struct BlockGemmARegBGmemCRegV1
         block_sync_lds();
 
         // block GEMM
-        return BlockGemmARegBSmemCRegImpl{}(a_block_tensor, b_block_smem_window);
+        return BlockGemmARegBGmemCRegImpl{}(a_block_tensor, b_block_smem_window);
     }
 };
 

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp

@@ -4,6 +4,7 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
+#include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
 
 namespace ck_tile {