diff --git a/projects/composablekernel/dispatcher/codegen/codegen_common.py b/projects/composablekernel/dispatcher/codegen/codegen_common.py index 0c4e38c7f970..4f8cee7f0ee6 100644 --- a/projects/composablekernel/dispatcher/codegen/codegen_common.py +++ b/projects/composablekernel/dispatcher/codegen/codegen_common.py @@ -783,7 +783,9 @@ def normalize_gfx_arch(arch: str) -> str: # Pipelines / GEMM variants whose problem and epilogue honour fixed vector sizes. # Not preshuffle: its pre-shuffled B needs K and N to be warp-tile multiples, # which already implies native alignment, so narrower widths never help it. -VECTOR_SIZE_PIPELINES: FrozenSet[str] = frozenset({"mem", "compv3", "compv4", "compv5"}) +VECTOR_SIZE_PIPELINES: FrozenSet[str] = frozenset( + {"mem", "compv3", "compv4", "compv5", "comp_async"} +) VECTOR_SIZE_VARIANTS: FrozenSet[str] = frozenset( {"standard", "batched", "grouped", "multi_d", "multi_abd", "stream_k"} ) @@ -1098,8 +1100,23 @@ def resolve_gemm_vector_sizes( return eff, f"variant {variant} does not support fixed vector sizes" if pipeline not in VECTOR_SIZE_PIPELINES: return eff, f"pipeline {pipeline} does not support fixed vector sizes" - if epilogue != "cshuffle": + # The fixed A/B widths live in the pipeline problem, so any epilogue keeps + # them. Only CShuffle can also store C narrower than its native width; the + # default epilogue's C width is fixed by the warp-tile layout. + if epilogue == "default" and eff[2] != native[2]: + return eff, f"epilogue {epilogue} cannot use a fixed vector_size_c" + if epilogue not in ("cshuffle", "default"): return eff, f"epilogue {epilogue} cannot use fixed vector sizes" + # comp_async copies A/B global->LDS asynchronously, and that copy only exists + # in 4/8/12/16-byte chunks plus a 1-byte one (amd_async_global_load_to_lds). + if pipeline == "comp_async": + for name, width, elem in (("a", eff[0], ea), ("b", eff[1], eb)): + if width * elem < 4 and not (elem == 1 and width == 1): + return ( + eff, + f"comp_async async copy cannot move {width * elem} bytes " + f"(vector_size_{name}={width})", + ) # Stream-K reduces partial C tiles with buffer atomics, which need >= 4 bytes. if variant == "stream_k" and eff[2] * ec < 4: return ( @@ -1129,7 +1146,7 @@ def resolve_gemm_vector_sizes( f"{name} LDS layout needs more warps for vector_size_{name.lower()}={width}", ) per_thread = warp_tile_m * warp_tile_n // warp_size - if per_thread > eff[2]: + if epilogue == "cshuffle" and per_thread > eff[2]: shuffles = per_thread // eff[2] if per_thread % eff[2] or (tile_m if c_row else tile_n) % shuffles: return ( diff --git a/projects/composablekernel/dispatcher/tests/test_gemm_vector_sizes.py b/projects/composablekernel/dispatcher/tests/test_gemm_vector_sizes.py index 91603e4de9ca..dff42ae359bd 100644 --- a/projects/composablekernel/dispatcher/tests/test_gemm_vector_sizes.py +++ b/projects/composablekernel/dispatcher/tests/test_gemm_vector_sizes.py @@ -185,13 +185,40 @@ def test_rejects_non_power_of_two(self): def test_rejects_unsupported_scope(self): for kw, word in [ - (dict(pipeline="comp_async"), "pipeline"), - (dict(epilogue="default"), "epilogue"), + (dict(pipeline="preshufflev2"), "pipeline"), (dict(variant="grouped_quant"), "variant"), (dict(variant="preshuffle", pipeline="preshufflev2"), "variant"), ]: self.assertIn(word, _resolve("rcr", (1, 1, 8), **kw)[1], kw) + def test_default_epilogue_narrows_ab_only(self): + # Fixed A/B widths live in the pipeline problem; the default epilogue keeps + # its own C width (1 for row-major C), so only a narrower C is refused. + self.assertEqual( + _resolve("rcr", (1, 1, 8), epilogue="default"), ((1, 1, 1), None) + ) + self.assertEqual( + _resolve("rrr", (8, 2, 8), epilogue="default"), ((8, 2, 1), None) + ) + self.assertIn( + "vector_size_c", _resolve("rcc", (8, 2, 2), epilogue="default")[1] + ) + + def test_comp_async_widths_follow_the_async_copy(self): + # amd_async_global_load_to_lds moves 4/8/12/16 bytes, or 1 byte for 8-bit types. + self.assertEqual( + _resolve("rcr", (2, 4, 8), pipeline="comp_async"), ((2, 4, 8), None) + ) + self.assertIn( + "async copy", _resolve("rcr", (1, 8, 8), pipeline="comp_async")[1] + ) + self.assertIn( + "async copy", _resolve("rcr", (8, 1, 8), pipeline="comp_async")[1] + ) + fp8 = dict(dtype_a="fp8", dtype_b="fp8", dtype_c="fp16", pipeline="comp_async") + self.assertIsNone(_resolve("rcr", (1, 4, 8), **fp8)[1]) + self.assertIn("async copy", _resolve("rcr", (2, 16, 8), **fp8)[1]) + def test_native_width_skips_sub_element_steps(self): # Like CK: fp16 never uses the 4-byte (2-wide) step. self.assertEqual(_native_ab_vector_size(2, 64, 2, 64, 256), 1) @@ -580,10 +607,12 @@ def test_fallback_configs_and_rejects(self): native = [c for c in cfgs if not any(c.vector_sizes)] fixed = [c for c in cfgs if any(c.vector_sizes)] self.assertTrue(native and fixed) - self.assertTrue(all(c.epilogue == "cshuffle" for c in fixed)) - self.assertTrue(all(c.name.endswith("_vec1_1_8") for c in fixed)) - # The CI config also sweeps the default epilogue, which cannot take fixed widths. - self.assertTrue(any("epilogue default" in r for r in rejects)) + # CShuffle narrows only what the problem needs; the default epilogue (also in + # the CI config) takes the A/B widths and keeps its own C width, 1 for row C. + suffix = {"cshuffle": "_vec1_1_8", "default": "_vec1_1_1"} + self.assertEqual({c.epilogue for c in fixed}, set(suffix)) + self.assertTrue(all(c.name.endswith(suffix[c.epilogue]) for c in fixed)) + self.assertFalse(any("epilogue default" in r for r in rejects)) self.assertEqual(len({c.name for c in cfgs}), len(cfgs)) def test_max_kernels_keeps_fixed_width_variants(self):