ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-10 00:57:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

fixes after merge

Browse Source
This commit is contained in:
Jakub Piasecki
2025-05-05 17:59:06 +00:00
parent 608fd3a638
commit 4baaa1bbe1
10 changed files with 51 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
example/ck_tile/03_gemm/run_gemm_example.inc
View File

@@ -52,6 +52,8 @@ void permute_tensor_b(Tensor& tensor)
GemmConfig::kPadN,
GemmConfig::kPadK,
GemmConfig::DoubleSmemBuffer,
GemmConfig::SkipALds,
GemmConfig::SkipBLds,
ALayout,
BLayout,
CLayout,

2
example/ck_tile/03_gemm/universal_gemm.cpp
View File

@@ -43,6 +43,8 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config&
GemmConfig::kPadN,
GemmConfig::kPadK,
GemmConfig::DoubleSmemBuffer,
GemmConfig::SkipALds,
GemmConfig::SkipBLds,
ALayout,
BLayout,
CLayout,

25
include/ck_tile/ops/gemm/block/block_universal_gemm_ar_bs_cr.hpp
View File

@@ -129,6 +129,31 @@ struct BlockUniversalGemmArBsCr
using I0 = number<0>;
using I1 = number<1>;
CK_TILE_DEVICE static constexpr auto MakeABlockDistributionEncode()
{
constexpr index_t KPerThread = Traits::KPerThread;
constexpr index_t NumMacClusters = Traits::InterWaveSchedulingMacClusters;
constexpr index_t KPerInnerLoop =
ck_tile::max(KPerThread / NumMacClusters, WarpGemm::kKPerThread);
constexpr index_t KIterInterwave = KPerInnerLoop / WarpGemm::kKPerThread;
using KIterSeq = std::conditional_t<Scheduler == GemmPipelineScheduler::Interwave,
sequence<KIterInterwave>,
sequence<KIterPerWarp>>;
constexpr auto a_block_outer_dstr_encoding =
tile_distribution_encoding<sequence<NWarp>,
tuple<sequence<MIterPerWarp, MWarp>, KIterSeq>,
tuple<sequence<1, 0>>,
tuple<sequence<1, 0>>,
sequence<1, 2>,
sequence<0, 0>>{};
constexpr auto a_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
a_block_outer_dstr_encoding, typename WarpGemm::AWarpDstrEncoding{});
return a_block_dstr_encode;
}
CK_TILE_DEVICE static constexpr auto MakeBBlockDistributionEncode()
{
constexpr index_t KPerThread = Traits::KPerThread;

4
include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp
View File

@@ -693,6 +693,10 @@ struct GemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem<Problem>
Base::LocalPrefill(b_copy_lds_window, b_block_tiles.get(I0{}), b_element_func);
}
}
// TODO add encoding and support for BRowMajor for SkipBLds, current
// MakeShuffledBRegTileDistribution takes into account shuffling encoding which is used
// for [Global -> Vgpr -> Lds -> Vgpr] reading, but for skipping lds we need to have
// different shuffled layout for [Global -> Vgpr] reads, similar for AColMajor
// Global prefetch [1, PrefetchStages]
static_for<1, PrefetchStages, 1>{}([&](auto prefetch_idx) {

5
include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp
View File

@@ -51,14 +51,13 @@ struct TileGemmUniversalTraits
static constexpr bool DoubleSmemBuffer = DoubleSmemBuffer_;
static constexpr bool SkipALds = SkipALds_;
static constexpr bool SkipBLds = SkipBLds_;
using ALayout = ALayout_;
using BLayout = BLayout_;
using CLayout = CLayout_;
static constexpr bool TransposeC = TransposeC_;
static constexpr bool SkipALds = SkipALds_;
static constexpr bool SkipBLds = SkipBLds_;
static constexpr bool UseStructuredSparsity = UseStructuredSparsity_;
};

2
test/ck_tile/gemm/CMakeLists.txt
View File

@@ -16,8 +16,6 @@ if(CK_USE_OCP_FP8)
list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8)
if(GPU_TARGETS MATCHES "gfx94" OR GPU_TARGETS MATCHES "gfx95")
add_gtest_executable(test_ck_tile_gemm_pipeline_mem test_gemm_pipeline_mem.cpp)
add_gtest_executable(test_ck_tile_gemm_pipeline_mem_skip_a_lds test_gemm_pipeline_mem_skip_a_lds.cpp)
add_gtest_executable(test_ck_tile_gemm_pipeline_mem_skip_b_lds test_gemm_pipeline_mem_skip_b_lds.cpp)
add_gtest_executable(test_ck_tile_gemm_pipeline_compv3 test_gemm_pipeline_compv3.cpp)
add_gtest_executable(test_ck_tile_gemm_pipeline_compv4 test_gemm_pipeline_compv4.cpp)

37
test/ck_tile/gemm/test_gemm_pipeline_kernel_types.hpp
View File

@@ -8,20 +8,18 @@
#include "ck_tile/host.hpp"
#include "test_gemm_pipeline_util.hpp"
using F16 = ck_tile::half_t;
using F32 = float;
using F8 = ck_tile::fp8_t;
using Row = ck_tile::tensor_layout::gemm::RowMajor;
using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
using Intrawave = ck_tile::integral_constant<ck_tile::GemmPipelineScheduler,
using F16 = ck_tile::half_t;
using F32 = float;
using F8 = ck_tile::fp8_t;
using Row = ck_tile::tensor_layout::gemm::RowMajor;
using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
using Intrawave = ck_tile::integral_constant<ck_tile::GemmPipelineScheduler,
ck_tile::GemmPipelineScheduler::Intrawave>;
using Interwave = ck_tile::integral_constant<ck_tile::GemmPipelineScheduler,
using Interwave = ck_tile::integral_constant<ck_tile::GemmPipelineScheduler,
ck_tile::GemmPipelineScheduler::Interwave>;
using Mem = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::Mem>;
using MemSkipALds = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::MemSkipALds>;
using MemSkipBLds = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::MemSkipBLds>;
using CompV3 = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::CompV3>;
using CompV4 = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::CompV4>;
using Mem = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::Mem>;
using CompV3 = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::CompV3>;
using CompV4 = ck_tile::integral_constant<GemmPipelineType, GemmPipelineType::CompV4>;
// clang-format off
using KernelTypesMem = ::testing::Types<
@@ -43,21 +41,6 @@ using KernelTypesMem = ::testing::Types<
std::tuple< Col, Col, Row, F8, F8, F32, F16, Interwave, Mem>
>;
using KernelTypesMemSkipALds = ::testing::Types<
std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave, MemSkipALds>,
std::tuple< Row, Row, Row, F16, F16, F32, F16, Interwave, MemSkipALds>,
std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave, MemSkipALds>,
std::tuple< Row, Col, Row, F16, F16, F32, F16, Interwave, MemSkipALds>
>;
using KernelTypesMemSkipBLds = ::testing::Types<
std::tuple< Row, Col, Row, F16, F16, F32, F16, Intrawave, MemSkipBLds>,
std::tuple< Row, Col, Row, F16, F16, F32, F16, Interwave, MemSkipBLds>,
std::tuple< Col, Col, Row, F16, F16, F32, F16, Intrawave, MemSkipBLds>,
std::tuple< Col, Col, Row, F16, F16, F32, F16, Interwave, MemSkipBLds>
>;
using KernelTypesCompV3 = ::testing::Types<
std::tuple< Row, Row, Row, F16, F16, F32, F16, Intrawave, CompV3>,
std::tuple< Row, Row, Row, F8, F8, F32, F16, Intrawave, CompV3>,

16
test/ck_tile/gemm/test_gemm_pipeline_mem_skip_a_lds.cpp
View File

@@ -1,16 +0,0 @@
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_util.hpp"
#include "gtest/gtest.h"
template <typename T>
class TestCkTileGemmPipelineMemSkipALds : public TestCkTileGemmPipeline<T>
{
};
#define TEST_SUITE_NAME TestCkTileGemmPipelineMemSkipALds
TYPED_TEST_SUITE(TestCkTileGemmPipelineMemSkipALds, KernelTypesMemSkipALds);
#include "test_gemm_pipeline_ut_cases.inc"
#undef TEST_SUITE_NAME

16
test/ck_tile/gemm/test_gemm_pipeline_mem_skip_b_lds.cpp
View File

@@ -1,16 +0,0 @@
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_util.hpp"
#include "gtest/gtest.h"
template <typename T>
class TestCkTileGemmPipelineMemSkipBLds : public TestCkTileGemmPipeline<T>
{
};
#define TEST_SUITE_NAME TestCkTileGemmPipelineMemSkipBLds
TYPED_TEST_SUITE(TestCkTileGemmPipelineMemSkipBLds, KernelTypesMemSkipBLds);
#include "test_gemm_pipeline_ut_cases.inc"
#undef TEST_SUITE_NAME

23
test/ck_tile/gemm/test_gemm_pipeline_util.hpp
View File

@@ -51,20 +51,6 @@ struct GemmPipelineTypeSelector<GemmPipelineType::Mem, Problem>
using pipeline = ck_tile::GemmPipelineAgBgCrMem<Problem>;
};
template <typename Problem>
struct GemmPipelineTypeSelector<GemmPipelineType::MemSkipALds, Problem>
{
using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrMem<Problem>;
using pipeline = ck_tile::GemmPipelineAgBgCrMemSkipALds<Problem>;
};
template <typename Problem>
struct GemmPipelineTypeSelector<GemmPipelineType::MemSkipBLds, Problem>
{
using base_pipeline = ck_tile::BaseGemmPipelineAgBgCrMem<Problem>;
using pipeline = ck_tile::GemmPipelineAgBgCrMemSkipBLds<Problem>;
};
template <typename Problem>
struct GemmPipelineTypeSelector<GemmPipelineType::CompV3, Problem>
{
@@ -116,6 +102,9 @@ class TestCkTileGemmPipeline : public ::testing::Test
constexpr bool DoubleSmemBuffer = (PipelineType == GemmPipelineType::CompV4) ? true : false;
constexpr bool SkipALds = false;
constexpr bool SkipBLds = false;
// TODO: For now - but this should also be a test parameter
constexpr bool TransposeC = false;
@@ -137,6 +126,8 @@ class TestCkTileGemmPipeline : public ::testing::Test
kPadN,
kPadK,
DoubleSmemBuffer,
SkipALds,
SkipBLds,
ALayout,
BLayout,
CLayout,
@@ -230,9 +221,7 @@ class TestCkTileGemmPipeline : public ::testing::Test
}
}
if constexpr(PipelineType == GemmPipelineType::Mem ||
PipelineType == GemmPipelineType::MemSkipBLds ||
PipelineType == GemmPipelineType::MemSkipALds)
if constexpr(PipelineType == GemmPipelineType::Mem)
{
// Tail pipeline One to Seven
if(tail_num == ck_tile::TailNumber::One)