feat(ck-tile): comp_tdm_pc producer/consumer GEMM TE to dispatcher bridge (gfx1250)

projects/composablekernel/dispatcher/codegen/arch_filter.py

@@ -50,6 +50,8 @@
 from codegen_common import (
     CommonTypeMappings,
     gfx1250_fp32_tile_reject_reason,
+    gfx1250_lds_operand_bytes,
+    gfx1250_lds_operand_transposed,
     gfx1250_pipeline_reject_reason,
 )
 
@@ -278,6 +280,8 @@ class OperatorType(Enum):
         "preshuffle_tdm": 32768,
         # Mandatory double buffering (num_lds_buffers = 2), so half the budget.
         "comp_async": 32768,
+        # Ring of four LDS slots (producer/consumer TDM), so a quarter.
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "default": 65536,
     }
@@ -314,6 +318,9 @@ def get_lds_limit(
         ("comp_tdm_v2", "tdm", "interwave"),
         ("comp_tdm_v2", "cshuffle", "interwave"),
         ("comp_tdm_v2", "default", "interwave"),
+        ("comp_tdm_pc", "tdm", "interwave"),
+        ("comp_tdm_pc", "cshuffle", "interwave"),
+        ("comp_tdm_pc", "default", "interwave"),
     }
 
     DTYPE_COMBINATIONS = {
@@ -331,7 +338,7 @@ def get_lds_limit(
 
 # Tensor Data Mover pipelines. They exist only on gfx1250; elsewhere the TDM
 # instructions are no-ops and the kernel silently produces zeros.
-TDM_PIPELINES = ("comp_tdm", "comp_tdm_v2")
+TDM_PIPELINES = ("comp_tdm", "comp_tdm_v2", "comp_tdm_pc")
 TDM_ARCH = "gfx1250"
 
 
@@ -856,16 +863,18 @@ def _validate_trait_combo(self, config: KernelConfig, result: ValidationResult):
     def _validate_gfx1250_pipeline(
         self, config: KernelConfig, result: ValidationResult
     ):
-        """Validate the gfx1250 pipelines (comp_async, comp_tdm, comp_tdm_v2)
+        """Validate the gfx1250 pipelines (comp_async, comp_tdm*)
         and the TDM epilogue.
 
         The rules come from codegen_common.gfx1250_pipeline_reject_reason, the
         single source of truth shared with unified_gemm_codegen and
         python/gemm_utils: TDM pipelines are gfx1250-only (off gfx1250 the TDM
         instructions compile to no-ops and the kernel silently writes zeros),
         need the TDM epilogue (and vice versa), intrawave, unpadded tiles and,
-        for comp_tdm_v2, exactly four waves; only the plain and preshuffle GEMM
-        operators can use them. Non-MX comp_async GEMM on gfx1250 needs the cshuffle
+        for comp_tdm_v2, exactly four waves, for comp_tdm_pc, at most 256 fp32
+        accumulators per lane; only the plain and preshuffle GEMM operators can
+        use them (comp_tdm_pc is plain only). Non-MX comp_async GEMM on gfx1250
+        needs the cshuffle
         epilogue, an rc A/B layout, pad_m=pad_n=pad_k=True and, for fp8/bf8,
         warp_tile_k >= 128 and a 16x16 warp tile (32x32 does not compile).
 
@@ -902,6 +911,8 @@ def _validate_gfx1250_pipeline(
             variant_supported=is_gemm,
             variant_name=f"operator={config.operator.value}",
             pads=None if None in pads else tuple(bool(p) for p in pads),
+            tile_mn=(config.tile_m, config.tile_n),
+            warp_k=config.warp_k,
         )
         if reason:
             result.add_error(reason)
@@ -928,14 +939,21 @@ def _validate_lds_capacity(self, config: KernelConfig, result: ValidationResult)
             and config.operator in (OperatorType.GEMM, OperatorType.GEMM_PRESHUFFLE)
         )
         if uses_gfx1250_base_lds:
-            # TDM and non-MX comp_async use the gfx1250 non-transposed base LDS
-            # descriptors: each group of rows spanning at least 256 bytes is
-            # followed by 16 padding bytes (same accounting as the Tile Engine
-            # validate_lds_capacity).
-            a_lds_layer = max(1, 256 // (config.tile_k * elem_size_a))
-            b_lds_layer = max(1, 256 // (config.tile_k * elem_size_b))
-            matrix_a_size += max(0, config.tile_m // a_lds_layer - 1) * 16
-            matrix_b_size += max(0, config.tile_n // b_lds_layer - 1) * 16
+            # TDM and non-MX comp_async use the gfx1250 base LDS descriptors,
+            # whose padding depends on whether the operand is transpose-loaded
+            # (same accounting as the Tile Engine validate_lds_capacity).
+            matrix_a_size = gfx1250_lds_operand_bytes(
+                config.tile_m,
+                config.tile_k,
+                elem_size_a,
+                gfx1250_lds_operand_transposed("a", config.datatype_a, config.layout),
+            )
+            matrix_b_size = gfx1250_lds_operand_bytes(
+                config.tile_n,
+                config.tile_k,
+                elem_size_b_staged,
+                gfx1250_lds_operand_transposed("b", config.datatype_b, config.layout),
+            )
         total_lds = matrix_a_size + matrix_b_size
 
         # The budget depends on the target, not just the pipeline: a tile that

projects/composablekernel/dispatcher/codegen/arch_specs.json

@@ -348,6 +348,17 @@
         "pipeline_lds_budget_aliases."
       ]
     },
+    "comp_tdm_pc": {
+      "basis": "fraction",
+      "value": 0.25,
+      "_reason": [
+        "Ring of NumSlots LDS buffers, no configuration choice:",
+        "GemmPipelineAgBgCrCompTDMProducerConsumer::GetSmemSize() returns",
+        "NumSlots * Policy::GetSmemSize() with the policy default NumSlots_ = 4.",
+        "Emitted to C++ as Pipeline::CompTDMPC. Matches the Tile Engine",
+        "PIPELINE_SMEM_SLOTS entry; raise both if NumSlots_ changes."
+      ]
+    },
     "wavelet": {
       "basis": "fraction",
       "value": 1.0,
@@ -410,7 +421,10 @@
       ["comp_tdm", "tdm", "interwave"],
       ["comp_tdm_v2", "cshuffle", "interwave"],
       ["comp_tdm_v2", "default", "interwave"],
-      ["comp_tdm_v2", "tdm", "interwave"]
+      ["comp_tdm_v2", "tdm", "interwave"],
+      ["comp_tdm_pc", "cshuffle", "interwave"],
+      ["comp_tdm_pc", "default", "interwave"],
+      ["comp_tdm_pc", "tdm", "interwave"]
     ]
   },
 

projects/composablekernel/dispatcher/codegen/arch_specs_generated.py

@@ -4,7 +4,7 @@
 AUTO-GENERATED FILE - DO NOT EDIT DIRECTLY!
 
 Generated from: arch_specs.json
-Generated at: 2026-10-25T20:57:27.671022
+Generated at: 2026-10-30T14:59:00.519447
 
 To update this file:
 1. Edit arch_specs.json
@@ -162,6 +162,7 @@
         "preshufflev2": 32768,
         "preshuffle_tdm": 32768,
         "comp_async": 32768,
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "compv6": 32768,
         "preshufflev1": 32768,
@@ -178,6 +179,7 @@
         "preshufflev2": 32768,
         "preshuffle_tdm": 32768,
         "comp_async": 32768,
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "compv6": 32768,
         "preshufflev1": 32768,
@@ -194,6 +196,7 @@
         "preshufflev2": 32768,
         "preshuffle_tdm": 32768,
         "comp_async": 32768,
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "compv6": 32768,
         "preshufflev1": 32768,
@@ -210,6 +213,7 @@
         "preshufflev2": 81920,
         "preshuffle_tdm": 81920,
         "comp_async": 81920,
+        "comp_tdm_pc": 40960,
         "wavelet": 163840,
         "compv6": 81920,
         "preshufflev1": 81920,
@@ -226,6 +230,7 @@
         "preshufflev2": 32768,
         "preshuffle_tdm": 32768,
         "comp_async": 32768,
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "compv6": 32768,
         "preshufflev1": 32768,
@@ -242,6 +247,7 @@
         "preshufflev2": 32768,
         "preshuffle_tdm": 32768,
         "comp_async": 32768,
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "compv6": 32768,
         "preshufflev1": 32768,
@@ -258,6 +264,7 @@
         "preshufflev2": 32768,
         "preshuffle_tdm": 32768,
         "comp_async": 32768,
+        "comp_tdm_pc": 16384,
         "wavelet": 65536,
         "compv6": 32768,
         "preshufflev1": 32768,
@@ -274,6 +281,7 @@
         "preshufflev2": 163840,
         "preshuffle_tdm": 163840,
         "comp_async": 163840,
+        "comp_tdm_pc": 81920,
         "wavelet": 327680,
         "compv6": 163840,
         "preshufflev1": 163840,
@@ -319,6 +327,9 @@
     ("comp_tdm_v2", "cshuffle", "interwave"),
     ("comp_tdm_v2", "default", "interwave"),
     ("comp_tdm_v2", "tdm", "interwave"),
+    ("comp_tdm_pc", "cshuffle", "interwave"),
+    ("comp_tdm_pc", "default", "interwave"),
+    ("comp_tdm_pc", "tdm", "interwave"),
 }
 
 # Valid dtype combinations: (A_dtype, B_dtype) -> acc_dtype and notes