diff --git a/tile_engine/CMakeLists.txt b/tile_engine/CMakeLists.txt index cd1a192a74..6e92b2125a 100644 --- a/tile_engine/CMakeLists.txt +++ b/tile_engine/CMakeLists.txt @@ -3,3 +3,5 @@ include_directories(BEFORE ) add_subdirectory(ops) + +add_subdirectory(01_gemm) diff --git a/tile_engine/ops/gemm/codegen_utils.py b/tile_engine/ops/gemm/codegen_utils.py index 98595933b8..0487d3d798 100644 --- a/tile_engine/ops/gemm/codegen_utils.py +++ b/tile_engine/ops/gemm/codegen_utils.py @@ -71,7 +71,9 @@ CSHUFFLE_EPILOGUE = """ """ PIPELINE_MAP = { - "mem": ["ck_tile::BaseGemmPipelineAgBgCrMem", "ck_tile::GemmPipelineAgBgCrMem"], + "mem": ["ck_tile::BaseGemmPipelineAgBgCrMem", + "ck_tile::GemmPipelineAgBgCrMem" + ], "compv3": [ "ck_tile::BaseGemmPipelineAgBgCrCompV3", "ck_tile::GemmPipelineAgBgCrCompV3", diff --git a/tile_engine/ops/gemm/single_instance/generate_kernel_instance.py b/tile_engine/ops/gemm/single_instance/generate_kernel_instance.py index 073acc3f7b..12d1319e54 100644 --- a/tile_engine/ops/gemm/single_instance/generate_kernel_instance.py +++ b/tile_engine/ops/gemm/single_instance/generate_kernel_instance.py @@ -1,8 +1,28 @@ - class GemmKernelBuilder: def _generate_kernel_instance(self, is_header=True): + # Map pipeline names to base pipeline for hot loop detection + base_pipeline_map = { + "mem": "ck_tile::BaseGemmPipelineAgBgCrMem", + "compv3": "ck_tile::BaseGemmPipelineAgBgCrCompV3", + "compv4": "ck_tile::BaseGemmPipelineAgBgCrCompV4", + } + + # Map pipeline names to the correct pipeline implementation + pipeline_impl_map = { + "mem": "ck_tile::GemmPipelineAgBgCrMem", + "compv3": "ck_tile::GemmPipelineAgBgCrCompV3", + "compv4": "ck_tile::GemmPipelineAgBgCrCompV4", + } + + # Map scheduler names to the correct enum values + scheduler_type_map = { + "intrawave": "ck_tile::GemmPipelineScheduler::Intrawave", + "interwave": "ck_tile::GemmPipelineScheduler::Interwave", + "default": "ck_tile::GemmPipelineScheduler::Default", + } + a_type = "ck_tile::fp16_t" b_type = "ck_tile::fp16_t" c_type = "ck_tile::fp16_t" @@ -14,19 +34,36 @@ class GemmKernelBuilder: b_layout = "ck_tile::tensor_layout::gemm::ColumnMajor" c_layout = "ck_tile::tensor_layout::gemm::RowMajor" - # BlkTile: 128x256x64 + # BlkTile: 128x256x64, WaveMap: 4x4, WaveTile: 16x16, tile_config = { "tile_m": 128, "tile_n": 256, "tile_k": 64, - "warp_m": warp_m, - "warp_n": warp_n, - "warp_k": warp_k, - "warp_tile_m": warp_tile_m, - "warp_tile_n": warp_tile_n, - "warp_tile_k": warp_tile_k, + "warp_m": 4, + "warp_n": 4, + "warp_k": 1, + "warp_tile_m": 16, + "warp_tile_n": 16, + "warp_tile_k": 16, } + pad_m = "false" + pad_n = "false" + pad_k = "false" + + persistent = "false" + + # BlkGemmPipelineVersion: v1, -> ????? + pipeline="mem" + + # BlkGemmPipelineScheduler: Intrawave + scheduler = "intrawave" + + # DeviceGemm_Wmma_CShuffleV3 + epilogue = "cshuffle" + + # VmemReadVec: 8x8 -> ????? + kernel_name = f"gemm" # Generate kernel instance code using the correct API diff --git a/tile_engine/ops/gemm/validation_utils.py b/tile_engine/ops/gemm/validation_utils.py index c0e109bf11..65f69ab7de 100644 --- a/tile_engine/ops/gemm/validation_utils.py +++ b/tile_engine/ops/gemm/validation_utils.py @@ -128,10 +128,11 @@ WARP_SUPPORTED_COMBINATIONS = { [4, 1, 1], ], "gfx1201": [ - [2, 4, 1], - [1, 8, 1], - [8, 1, 1], - [4, 2, 1], + [4, 4, 1], + # [2, 4, 1], + # [1, 8, 1], + # [8, 1, 1], + # [4, 2, 1], ], }