feat(dispatcher/codegen): fixed A/B vector widths for the default epilogue and comp_async

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 (

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):