feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all layouts)

projects/composablekernel/dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp

@@ -20,6 +20,7 @@
 #include <memory>
 #include <sstream>
 #include <string>
+#include <type_traits>
 
 #include "ck_tile/dispatcher/dispatcher.hpp"
 #include "ck_tile/dispatcher/registry.hpp"
@@ -65,15 +66,75 @@ int dispatcher_initialize()
         return 0; // Already initialized
     }
 
-    // Create kernel key from the force-included kernel header
+    // Create kernel key from the force-included kernel header.
+    //
+    // The GEMM_KEY_* macros are emitted by the codegen into the force-included
+    // header (see unified_gemm_codegen.py, CK_TILE_SINGLE_KERNEL_INCLUDE block).
+    // Building the key from them makes the registry entry truthful: it reflects
+    // THIS kernel's real dtypes/layouts/tile/traits instead of a hard-coded
+    // fp16/rcr/128x128x32 default. Enum fields use the string_to_* helpers from
+    // kernel_key.hpp, whose accepted strings match the codegen's emitted values
+    // byte-for-byte.
     KernelKey key;
-    key.signature.dtype_a             = DataType::FP16;
-    key.signature.dtype_b             = DataType::FP16;
-    key.signature.dtype_c             = DataType::FP16;
-    key.signature.dtype_acc           = DataType::FP32;
-    key.signature.layout_a            = LayoutTag::RowMajor;
-    key.signature.layout_b            = LayoutTag::ColMajor;
-    key.signature.layout_c            = LayoutTag::RowMajor;
+#ifdef GEMM_KEY_DTYPE_A
+    key.signature.dtype_a             = string_to_dtype(GEMM_KEY_DTYPE_A);
+    key.signature.dtype_b             = string_to_dtype(GEMM_KEY_DTYPE_B);
+    key.signature.dtype_c             = string_to_dtype(GEMM_KEY_DTYPE_C);
+    key.signature.dtype_acc           = string_to_dtype(GEMM_KEY_DTYPE_ACC);
+    key.signature.layout_a            = string_to_layout(GEMM_KEY_LAYOUT_A);
+    key.signature.layout_b            = string_to_layout(GEMM_KEY_LAYOUT_B);
+    key.signature.layout_c            = string_to_layout(GEMM_KEY_LAYOUT_C);
+    key.signature.transpose_a         = false;
+    key.signature.transpose_b         = false;
+    key.signature.grouped             = (GEMM_KEY_GROUPED != 0);
+    key.signature.split_k             = GEMM_KEY_SPLIT_K;
+    key.signature.elementwise_op      = "PassThrough";
+    key.signature.num_d_tensors       = 0;
+    key.signature.structured_sparsity = false;
+
+    key.algorithm.tile_shape      = {GEMM_KEY_TILE_M, GEMM_KEY_TILE_N, GEMM_KEY_TILE_K};
+    key.algorithm.wave_shape      = {GEMM_KEY_WAVE_M, GEMM_KEY_WAVE_N, GEMM_KEY_WAVE_K};
+    key.algorithm.warp_tile_shape = {
+        GEMM_KEY_WARP_TILE_M, GEMM_KEY_WARP_TILE_N, GEMM_KEY_WARP_TILE_K};
+    key.algorithm.pipeline        = string_to_pipeline(GEMM_KEY_PIPELINE);
+    key.algorithm.scheduler       = string_to_scheduler(GEMM_KEY_SCHEDULER);
+    key.algorithm.epilogue        = string_to_epilogue(GEMM_KEY_EPILOGUE);
+    key.algorithm.block_size      = GEMM_KEY_BLOCK_SIZE;
+    key.algorithm.double_buffer   = (GEMM_KEY_DOUBLE_BUFFER != 0);
+    key.algorithm.persistent      = (GEMM_KEY_PERSISTENT != 0);
+    key.algorithm.preshuffle      = (GEMM_KEY_PRESHUFFLE != 0);
+    key.algorithm.transpose_c     = (GEMM_KEY_TRANSPOSE_C != 0);
+    key.algorithm.num_wave_groups = GEMM_KEY_NUM_WAVE_GROUPS;
+    // pad_m/n/k participate in both the key's hash/equality and the kernel
+    // name, so they must be derived from the codegen macros too -- otherwise a
+    // kernel built with padding disabled would register under a key claiming
+    // pad=true and disagree with its own name.
+    key.algorithm.pad_m = (GEMM_KEY_PAD_M != 0);
+    key.algorithm.pad_n = (GEMM_KEY_PAD_N != 0);
+    key.algorithm.pad_k = (GEMM_KEY_PAD_K != 0);
+    key.gfx_arch        = GFX_ARCH;
+#else
+    // Fallback default for headers generated before GEMM_KEY_* macros existed
+    // (fp16 / rcr / compv4-cshuffle-intrawave, 128x128x32). The macro path
+    // above is the source of truth for any freshly generated kernel.
+    key.signature.dtype_a   = DataType::FP16;
+    key.signature.dtype_b   = DataType::FP16;
+    key.signature.dtype_c   = DataType::FP16;
+    key.signature.dtype_acc = DataType::FP32;
+    // Derive A/B/C layouts from the force-included kernel's own layout types
+    // instead of hardcoding rcr. The dispatcher's supports() gate is layout-aware
+    // (it only constrains a dimension that an operand's inner axis maps to), so a
+    // wrong key layout makes it reject valid problems -- e.g. a crr kernel does not
+    // gate K, but with a hardcoded rcr key supports() would apply rcr's K-gate and
+    // reject TileK=192 problems that Old-TE runs. ALayout/BLayout/CLayout are the
+    // global aliases exported by the kernel header under CK_TILE_SINGLE_KERNEL_INCLUDE.
+    using RowMajorLayout = ck_tile::tensor_layout::gemm::RowMajor;
+    key.signature.layout_a =
+        std::is_same_v<ALayout, RowMajorLayout> ? LayoutTag::RowMajor : LayoutTag::ColMajor;
+    key.signature.layout_b =
+        std::is_same_v<BLayout, RowMajorLayout> ? LayoutTag::RowMajor : LayoutTag::ColMajor;
+    key.signature.layout_c =
+        std::is_same_v<CLayout, RowMajorLayout> ? LayoutTag::RowMajor : LayoutTag::ColMajor;
     key.signature.transpose_a         = false;
     key.signature.transpose_b         = false;
     key.signature.grouped             = false;
@@ -95,6 +156,7 @@ int dispatcher_initialize()
     key.algorithm.transpose_c     = false;
     key.algorithm.num_wave_groups = 1;
     key.gfx_arch                  = GFX_ARCH;
+#endif // GEMM_KEY_DTYPE_A
 
     // Register kernel using types from force-included header
     auto kernel =
@@ -310,10 +372,40 @@ int dispatcher_run_gemm(
 }
 
 /**
- * Get kernel information
+ * Get kernel information (legacy single-kernel ABI).
+ *
+ * Returns the compile-time KERNEL_NAME of the force-included kernel header.
+ * Kept for backward compatibility with one-kernel-per-.so callers.
  */
 const char* dispatcher_get_kernel_name() { return KERNEL_NAME; }
 
+/**
+ * Get the name of the kernel at a given registry index (multi-kernel ABI).
+ *
+ * Mirrors the conv/fmha ctypes libs: copies the index-th registered kernel's
+ * name into the caller-provided buffer so one .so can report a whole batch and
+ * be selected by name at runtime. Returns 0 on success, -1 on bad args or
+ * out-of-range index.
+ */
+int dispatcher_get_kernel_name_at(int index, char* buffer, int buffer_size)
+{
+    if(!buffer || buffer_size <= 0)
+    {
+        return -1;
+    }
+
+    auto kernels = Registry::instance().get_all();
+    if(index < 0 || index >= static_cast<int>(kernels.size()))
+    {
+        return -1;
+    }
+
+    std::string name = kernels[index]->get_name();
+    std::strncpy(buffer, name.c_str(), static_cast<size_t>(buffer_size) - 1);
+    buffer[buffer_size - 1] = '\0';
+    return 0;
+}
+
 /**
  * Initialize dispatcher (alias)
  */
@@ -398,4 +490,4 @@ void dispatcher_cleanup()
     g_initialized = false;
 }
 
-} // extern "C"
+} // extern "C"

projects/composablekernel/dispatcher/codegen/unified_gemm_codegen.py

@@ -187,6 +187,10 @@ def is_preshuffle_config_valid(
 log = logging.getLogger(__name__)
 
 
+def _is_power_of_two(x: int) -> bool:
+    return x > 0 and (x & (x - 1)) == 0
+
+
 # ============================================================================
 # Configuration and Data Structures
 # ============================================================================
@@ -520,6 +524,43 @@ def _selected_kernel_struct(self, config: KernelConfig, kernel_name: str) -> str
 using BDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.datatype]};
 using CDataType = {self.tm.DTYPE_TO_CK_QUALIFIED[self.tm.get_output_dtype(self.datatype)]};
 using AccDataType = float;
+
+// KernelKey field descriptors for the force-included kernel.
+// The ctypes library builds the registry KernelKey from these so the
+// registered entry reflects this kernel's real traits (not a hard-coded
+// fp16/rcr default). Enum-valued fields are emitted as the exact strings
+// consumed by string_to_dtype/layout/pipeline/scheduler/epilogue in
+// kernel_key.hpp; shape/flag fields are emitted as numeric/0-1 literals.
+#define GEMM_KEY_DTYPE_A "{self.datatype}"
+#define GEMM_KEY_DTYPE_B "{self.datatype}"
+#define GEMM_KEY_DTYPE_C "{output_dtype}"
+#define GEMM_KEY_DTYPE_ACC "fp32"
+#define GEMM_KEY_LAYOUT_A "{self.layout[0]}"
+#define GEMM_KEY_LAYOUT_B "{self.layout[1]}"
+#define GEMM_KEY_LAYOUT_C "{self.layout[2]}"
+#define GEMM_KEY_PIPELINE "{tr.pipeline}"
+#define GEMM_KEY_SCHEDULER "{tr.scheduler}"
+#define GEMM_KEY_EPILOGUE "{tr.epilogue}"
+#define GEMM_KEY_TILE_M {t.tile_m}
+#define GEMM_KEY_TILE_N {t.tile_n}
+#define GEMM_KEY_TILE_K {t.tile_k}
+#define GEMM_KEY_WAVE_M {t.warp_m}
+#define GEMM_KEY_WAVE_N {t.warp_n}
+#define GEMM_KEY_WAVE_K {t.warp_k}
+#define GEMM_KEY_WARP_TILE_M {t.warp_tile_m}
+#define GEMM_KEY_WARP_TILE_N {t.warp_tile_n}
+#define GEMM_KEY_WARP_TILE_K {t.warp_tile_k}
+#define GEMM_KEY_BLOCK_SIZE {config.block_size}
+#define GEMM_KEY_NUM_WAVE_GROUPS {config.num_wave_groups}
+#define GEMM_KEY_PAD_M {int(tr.pad_m)}
+#define GEMM_KEY_PAD_N {int(tr.pad_n)}
+#define GEMM_KEY_PAD_K {int(tr.pad_k)}
+#define GEMM_KEY_PERSISTENT {int(tr.persistent)}
+#define GEMM_KEY_DOUBLE_BUFFER {int(tr.pipeline == "compv4" or tr.pipeline == "preshufflev2")}
+#define GEMM_KEY_PRESHUFFLE {int(config.preshuffle)}
+#define GEMM_KEY_TRANSPOSE_C 0
+#define GEMM_KEY_GROUPED 0
+#define GEMM_KEY_SPLIT_K 1
 #endif // CK_TILE_SINGLE_KERNEL_INCLUDE
 """
 
@@ -1014,6 +1055,31 @@ def _get_preselected_configs(self) -> List[KernelConfig]:
             log.error(f"Invalid preselected set: {e}")
             return []
 
+    @staticmethod
+    def _cshuffle_repeat_ok(tile: TileConfig) -> bool:
+        """CShuffle-store correctness gate.
+
+        The CShuffle epilogue stores the accumulator back through LDS in
+        power-of-two MRepeat/NRepeat chunks, so a tile whose per-wave repeat
+        count -- tile / (warp * warp_tile) -- is not a power of two is
+        mis-stored and yields numerically WRONG results at runtime. The kernel
+        still compiles (the epilogue's static_asserts only check divisibility,
+        which such tiles satisfy), so it must be filtered in codegen. Observed
+        on MI350 for tile_m=192 (MRepeat = 192 / (2*32) = 3): verified incorrect
+        on BOTH the bridge and Tile Engine at every shape, including shapes
+        divisible by 192. Power-of-two tiles (64/128/256) are unaffected.
+
+        This is CShuffle-specific: the "default" (DefaultGemm2DEpilogue) path
+        stores directly (not through the LDS repack) and is numerically correct
+        for non-pow2 repeats -- verified on gfx942 at tile_m=192/MRepeat=3
+        (max_rel ~5e-4 across shapes divisible by 192, while the same tile under
+        CShuffle returns garbage, max_rel ~1.3). Only call this for kernels
+        whose resolved epilogue is "cshuffle".
+        """
+        m_repeat = tile.tile_m // (tile.warp_m * tile.warp_tile_m)
+        n_repeat = tile.tile_n // (tile.warp_n * tile.warp_tile_n)
+        return _is_power_of_two(m_repeat) and _is_power_of_two(n_repeat)
+
     def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]:
         """Get all configurations for a variant
 
@@ -1030,12 +1096,24 @@ def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]:
         trait_configs = self._get_trait_configs()
 
         for tile, trait in itertools.product(tile_configs, trait_configs):
-            # Perform variant-specific architecture validation
+            # Perform variant-specific architecture validation against the
+            # trait's ACTUAL pipeline/scheduler (not a hard-coded compv4).
             if self.arch_filter and HAS_ARCH_FILTER:
-                if not self._is_tile_arch_valid(tile, variant):
+                if not self._is_tile_arch_valid(
+                    tile,
+                    variant,
+                    pipeline=trait.pipeline,
+                    scheduler=trait.scheduler,
+                ):
                     continue
 
             if variant == GemmVariant.STANDARD:
+                # CShuffle-store correctness gate: skip non-pow2 repeat tiles
+                # only for the cshuffle epilogue (see _cshuffle_repeat_ok). The
+                # "default" epilogue is correct with non-pow2 repeats, so it is
+                # NOT gated here.
+                if trait.epilogue == "cshuffle" and not self._cshuffle_repeat_ok(tile):
+                    continue
                 configs.append(KernelConfig(tile=tile, trait=trait, variant=variant))
 
             elif variant == GemmVariant.PRESHUFFLE:
@@ -1052,7 +1130,13 @@ def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]:
                 )
                 # Only generate one preshuffle config per tile (not per trait)
                 # since preshuffle has fixed pipeline/scheduler
-                if trait.pipeline == "compv3" and trait.scheduler == "intrawave":
+                # Preshuffle always uses the cshuffle epilogue, so the
+                # CShuffle-store pow2 repeat gate always applies here.
+                if (
+                    trait.pipeline == "compv3"
+                    and trait.scheduler == "intrawave"
+                    and self._cshuffle_repeat_ok(tile)
+                ):
                     configs.append(
                         KernelConfig(
                             tile=tile,
@@ -1063,6 +1147,10 @@ def _get_configs_for_variant(self, variant: GemmVariant) -> List[KernelConfig]:
                     )
 
             elif variant == GemmVariant.MULTI_D:
+                # CShuffle-store correctness gate: applies only when the
+                # (swept) epilogue is cshuffle; the default epilogue is exempt.
+                if trait.epilogue == "cshuffle" and not self._cshuffle_repeat_ok(tile):
+                    continue
                 multi_d = self.config.get("multi_d_config", {})
                 for ew_op, num_d in itertools.product(
                     multi_d.get("elementwise_ops", ["MultiDAdd"]),
@@ -1105,9 +1193,28 @@ def _get_tile_configs(self) -> List[TileConfig]:
                 rejected_count += 1
                 continue
 
-            # Architecture-specific validation
+            # NOTE: the CShuffle-store pow2 MRepeat/NRepeat correctness gate is
+            # NOT applied here. It is epilogue-specific (only the CShuffle
+            # epilogue mis-stores non-pow2 repeats; the "default" epilogue is
+            # correct), so it is applied per (tile, trait) in
+            # _get_configs_for_variant once the epilogue is known. See
+            # _cshuffle_repeat_ok.
+
+            # Architecture-specific validation. This is a pre-filter run before
+            # tiles are paired with traits, so keep a tile if it is legal under
+            # ANY configured pipeline/scheduler; the precise per-trait check
+            # happens later in _get_configs_for_variant. Filtering here with a
+            # single hard-coded pipeline (compv4) wrongly dropped tiles that are
+            # legal under mem/compv3.
             if self.arch_filter and HAS_ARCH_FILTER:
-                if not self._is_tile_arch_valid(tile):
+                trait_cfg = self.config.get("trait_config", {})
+                pipelines = trait_cfg.get("pipeline") or ["compv4"]
+                schedulers = trait_cfg.get("scheduler") or ["intrawave"]
+                if not any(
+                    self._is_tile_arch_valid(tile, pipeline=pl, scheduler=sc)
+                    for pl in pipelines
+                    for sc in schedulers
+                ):
                     rejected_count += 1
                     continue
 
@@ -1119,13 +1226,23 @@ def _get_tile_configs(self) -> List[TileConfig]:
         return configs
 
     def _is_tile_arch_valid(
-        self, tile: TileConfig, variant: GemmVariant = None
+        self,
+        tile: TileConfig,
+        variant: GemmVariant = None,
+        pipeline: str = None,
+        scheduler: str = None,
     ) -> bool:
         """Check if tile configuration is valid for target architecture
 
         Args:
             tile: Tile configuration to validate
             variant: GEMM variant (affects operator-specific constraints)
+            pipeline: Trait pipeline to validate against. Pass the config's
+                actual pipeline -- omitting it falls back to ``compv4``, whose
+                MFMA constraints are stricter than ``mem``/``compv3`` and would
+                wrongly reject tiles that are legal under those pipelines.
+            scheduler: Trait scheduler to validate against (defaults to
+                ``intrawave`` for the same reason).
         """
         if not self.arch_filter or not HAS_ARCH_FILTER:
             return True
@@ -1146,8 +1263,10 @@ def _is_tile_arch_valid(
 
         # Map GEMM variant to operator type for validation
         operator = None
-        pipeline = "compv4"  # Default
-        scheduler = "intrawave"  # Default
+        if pipeline is None:
+            pipeline = "compv4"  # Default (representative compute pipeline)
+        if scheduler is None:
+            scheduler = "intrawave"  # Default
 
         if OperatorType is not None and variant is not None:
             variant_to_operator = {