CK-Tile Grouped GEMM refactor and post PR fixes #1756

mozga-amd · 2024-12-17T12:25:21Z

This pull-request contains changes as following:

Standardizing the function naming,
Moving the get_stride function, etc., to a different location,
Generalizing the 1D partitioner function for general use and single responsibility,
Adding an offset function to calculate the shift.

aosewski · 2024-12-18T09:48:46Z

include/ck_tile/core/arch/arch.hpp

+#define CK_CONSTANT_ADDRESS_SPACE __attribute__((address_space(4)))
+
+template <typename T>
+__device__ T* cast_pointer_to_generic_address_space(T CK_CONSTANT_ADDRESS_SPACE* p)
+{
+    // cast a pointer in "Constant" address space (4) to "Generic" address space (0)
+    // only c-style pointer cast seems be able to be compiled
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wold-style-cast"
+    return (T*)p; // NOLINT(old-style-cast)
+#pragma clang diagnostic pop
+}
+
+template <typename T>
+__host__ __device__ T CK_CONSTANT_ADDRESS_SPACE* cast_pointer_to_constant_address_space(T* p)
+{
+    // cast a pointer in "Generic" address space (0) to "Constant" address space (4)
+    // only c-style pointer cast seems be able to be compiled
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wold-style-cast"
+    return (T CK_CONSTANT_ADDRESS_SPACE*)p; // NOLINT(old-style-cast)
+#pragma clang diagnostic pop
+}


It seems like here the address space id is not consistent with the one here:
https://github.com/ROCm/composable_kernel/pull/1756/files#diff-5fbd9be40988c2586d9a6d3568593c4a9c5a5fb3a718eded871578acd7f34b8cR15-R22

Please make it consistent and add constant to the enum.

include/ck_tile/core/arch/arch.hpp

aosewski · 2024-12-18T09:51:44Z

include/ck_tile/core/arch/arch.hpp

@@ -109,4 +109,28 @@ CK_TILE_DEVICE void s_nop(index_t cnt = 0)
 #endif
 }

+#define CK_CONSTANT_ADDRESS_SPACE __attribute__((address_space(4)))


Is it possible to use enumerator address_space_enum here?

aosewski · 2024-12-18T09:52:56Z

include/ck_tile/host/host_tensor.hpp

+{
+    using namespace ck_tile::literals;
+
+    if constexpr(std::is_same_v<decltype(layout), ck_tile::tensor_layout::gemm::RowMajor>)


We're already in ck_tile namespace.

IMO, it's good practice to maintain namespaces even inside them. This indicates its exact place of origin.

aosewski · 2024-12-18T09:59:30Z

include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp

+        index_t iM = __builtin_amdgcn_readfirstlane(blockIdx.x * MPerBlock);
+        index_t iN = __builtin_amdgcn_readfirstlane(blockIdx.x * NPerBlock);


This is incorrect.

aosewski · 2024-12-18T10:01:53Z

include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp

    using TilePartitioner                    = remove_cvref_t<TilePartitioner_>;
    using GemmPipeline                       = remove_cvref_t<GemmPipeline_>;
    using EpiloguePipeline                   = remove_cvref_t<EpiloguePipeline_>;
    using ALayout                            = remove_cvref_t<typename GemmPipeline::ALayout>;
    using BLayout                            = remove_cvref_t<typename GemmPipeline::BLayout>;
    using CLayout                            = remove_cvref_t<typename GemmPipeline::CLayout>;
+    using Block2ETileMap                     = OffsettedBlockToCTileMap<TilePartitioner>;


Please use CK-Tile naming convention and not the old CK's one.

aosewski · 2024-12-18T10:02:15Z

include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp

+    // Block2CTileMap configuration parameter.
+    static constexpr index_t B2E_M01 = 8;


Currently this shouldn't be needed.

aosewski · 2024-12-18T10:20:24Z

include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp


            const index_t block_start = grid_size;
            const index_t block_end   = grid_size + grid_size_grp;

            grid_size += grid_size_grp;

+            auto grouped_block_2_ctile_map = Block2ETileMap(B2E_M01, M, N);


This should be the offseted TilePartitioner. Take a look here: https://github.com/ROCm/composable_kernel/blob/develop/include/ck/tensor_operation/gpu/device/impl/device_grouped_gemm_xdl_splitk_cshuffle.hpp

aosewski · 2024-12-18T10:40:39Z

include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp

@@ -62,12 +62,10 @@ struct GemmTile1DPartitioner
        return integer_divide_ceil(K, KPerBlock);
    }

-    CK_TILE_DEVICE auto operator()(index_t blockOffset, index_t NBlockSize)


Right now the function in https://github.com/ROCm/composable_kernel/pull/1756/files#diff-1452d09e42d9aed38e087f552e940394aae58d72f2991fa842fa15caced37854R55 is not used anywhere.

@bartekxk

Will be off from tomorrow till 2nd Jan. Please pass on review to ie. @bartekxk

carlushuang · 2024-12-22T09:25:02Z

include/ck_tile/ops/gemm.hpp

@@ -26,6 +26,7 @@
 #include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.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/kernel/gemm_offset_block.hpp"


run python3 remod.py under include/ck_tile, you may find we will sorting the headers before include into this gemm.hpp. Actually headers under include/ck_tile/ops/ better use this script to be autogenerated :) in case ppl argue about alphbetical order

Grouped gemm simple code refactor

44fae39

mozga-amd requested review from junliume, illsilin, carlushuang, qianfengz, aosewski, poyenc, geyyer, bartekxk and andriy-ca as code owners December 17, 2024 12:25

aosewski changed the title ~~Apply refactor to universal gemm to bwd_weight_cshuffle~~ CK-Tile Grouped GEMM refactor and post PR fixes Dec 18, 2024

aosewski previously requested changes Dec 18, 2024

View reviewed changes

carlushuang reviewed Dec 22, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

CK-Tile Grouped GEMM refactor and post PR fixes #1756

CK-Tile Grouped GEMM refactor and post PR fixes #1756

mozga-amd commented Dec 17, 2024

aosewski Dec 18, 2024

aosewski Dec 18, 2024

aosewski Dec 18, 2024

mozga-amd Dec 18, 2024

aosewski Dec 18, 2024

aosewski Dec 18, 2024

aosewski Dec 18, 2024

aosewski Dec 18, 2024

aosewski Dec 18, 2024

carlushuang Dec 22, 2024 •

edited

Loading

		index_t iM = __builtin_amdgcn_readfirstlane(blockIdx.x * MPerBlock);
		index_t iN = __builtin_amdgcn_readfirstlane(blockIdx.x * NPerBlock);

		// Block2CTileMap configuration parameter.
		static constexpr index_t B2E_M01 = 8;

CK-Tile Grouped GEMM refactor and post PR fixes #1756

Are you sure you want to change the base?

CK-Tile Grouped GEMM refactor and post PR fixes #1756

Conversation

mozga-amd commented Dec 17, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

carlushuang Dec 22, 2024 • edited Loading

Choose a reason for hiding this comment

carlushuang Dec 22, 2024 •

edited

Loading