Calculate tile start/end in device kernel.

21fbf2ce · Adam Osewski · 1de7e9bc · 21fbf2ce · 21fbf2ce · 21fbf2ce
Commit 21fbf2ce authored Jul 17, 2023 by Adam Osewski
3 changed files
--- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle_tile_loop.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle_tile_loop.hpp
--- a/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp
+++ b/include/ck/tensor_operation/gpu/grid/block_to_ctile_map.hpp
@@ -269,7 +269,8 @@ struct BlockToCTileMap_KSplit_M00_N0_M01Adapt
    {
    }

-    __host__ constexpr index_t CalculateGridSize(const CGridDesc_M_N& c_grid_desc_m_n) const
+    __host__ __device__ constexpr index_t
+    CalculateGridSize(const CGridDesc_M_N& c_grid_desc_m_n) const
    {
        const auto M0 = math::integer_divide_ceil(c_grid_desc_m_n.GetLength(I0), MPerBlock);
        const auto N0 = math::integer_divide_ceil(c_grid_desc_m_n.GetLength(I1), NPerBlock);

--- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r4r2.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_v2r4r2.hpp
@@ -610,9 +610,7 @@ struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_v2r4r2
            p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());

        // divide block work by [KBatch, M, N]
-        const auto block_work_idx =
-            // block_2_ctile_map.CalculateBottomIndex(make_multi_index(get_block_1d_id()));
-            block_2_ctile_map.CalculateBottomIndex();
+        const auto block_work_idx = block_2_ctile_map.CalculateBottomIndex();

        if(!block_2_ctile_map.ValidCTileIndex(
               block_work_idx,