[CK Tile] Spatially local GEMM tile partitioner. (#1843)

* Add spatially local tile partitioner

* Use 1D Grid size & create partitioner object.

* Docs & use 1D partitioner in example.

* Clang format.

* Change kernel grid size

Now: X is the # of output C-tiles,
     Y is the batch count
     Z is the splitK

* Formatting & more doc.

* Clang format.

* Fix batched gemm test. Use 1d partitioner.

* Move condition.

* FIx ctor.

* clang-format.
This commit is contained in:
Adam Osewski
2025-01-31 00:10:16 +01:00
committed by GitHub
parent e6d4180498
commit ce448002ee
10 changed files with 285 additions and 88 deletions

View File

@@ -40,7 +40,7 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config&
ck_tile::sequence<M_Warp, N_Warp, K_Warp>,
ck_tile::sequence<M_Warp_Tile, N_Warp_Tile, K_Warp_Tile>>;
using TilePartitioner = ck_tile::GemmTile2DPartitioner<CodegenGemmShape>;
using TilePartitioner = ck_tile::GemmTile1DPartitioner<CodegenGemmShape>;
using CodegenGemmTraits =
ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, CLayout>;