[transform] Target-aware register parallel dim tiling strategy

lighthouse/dialects/transform/transform_ext/utils/tiling/common.py

@@ -15,6 +15,50 @@ def parallel_and_reduction_dims(out_map: ir.AffineMap) -> tuple[list[int], list[
     return parallel_dims, reduction_dims
 
 
+def assign_from_tail(dims: list[int], values: list[int], sizes: list[int]) -> None:
+    """Write `values` onto `sizes` at the trailing `dims`, aligning both by their tails."""
+    if not dims or not values:
+        return
+    tail = dims[-len(values) :]
+    vals = values[-len(tail) :]
+    for dim, tile in zip(tail, vals):
+        sizes[dim] = int(tile)
+
+
+def assign_parallel_tiles(
+    parallel_dims: list[int], inner_tiles: list[int], sizes: list[int]
+) -> None:
+    """Outer parallel dims -> 1; innermost parallel tail -> inner_tiles."""
+    inner_count = min(len(inner_tiles), len(parallel_dims))
+    for d in parallel_dims[:-inner_count] if inner_count else []:
+        sizes[d] = 1
+    assign_from_tail(parallel_dims, inner_tiles, sizes)
+
+
+def assign_reduction_tiles(
+    reduction_dims: list[int], red_tiles: list[int], sizes: list[int]
+) -> None:
+    """Innermost reduction tail -> red_tiles; remaining reduction dims -> 1."""
+    assign_from_tail(reduction_dims, red_tiles, sizes)
+    if reduction_dims:
+        protected = set(reduction_dims[-len(red_tiles) :]) if red_tiles else set()
+        unitize_unassigned_dims(reduction_dims, sizes, protected=protected)
+
+
+def unitize_unassigned_dims(
+    dims: list[int],
+    sizes: list[int],
+    protected: set[int] | None = None,
+) -> None:
+    """Set each untiled (zero) dim to 1, leaving `protected` dims untouched."""
+    protected = protected or set()
+    for dim in dims:
+        if dim in protected:
+            continue
+        if sizes[dim] == 0:
+            sizes[dim] = 1
+
+
 def output_tensor_dim_of_iter_dim(out_map: ir.AffineMap) -> dict[int, int]:
     """Map each iteration dim to the output tensor dim it indexes (if any)."""
     mapping: dict[int, int] = {}

lighthouse/dialects/transform/transform_ext/utils/tiling/registry.py

@@ -1,10 +1,12 @@
 from .strategy_base import TilingStrategy
 from .strategy_cache import CacheTilingStrategy
+from .strategy_register_parallel import RegisterParallelTilingStrategy
 
 
 # Maps each canonical strategy name to its implementation class.
 _STRATEGY_REGISTRY: dict[str, type[TilingStrategy]] = {
     "cache": CacheTilingStrategy,
+    "register_parallel": RegisterParallelTilingStrategy,
 }
 
 

lighthouse/dialects/transform/transform_ext/utils/tiling/strategy_register_parallel.py

@@ -0,0 +1,43 @@
+from mlir import ir
+
+from lighthouse.utils.mlir import opview
+
+from .strategy_base import StrategyContext, TilingStrategy
+from .common import (
+    assign_parallel_tiles,
+    disable_small_tiles,
+    parallel_and_reduction_dims,
+)
+from .target_caps import (
+    generic_parallel_tiles,
+    is_amx_bf16_contraction,
+    is_f32_contraction,
+)
+
+
+class RegisterParallelTilingStrategy(TilingStrategy):
+    """Register-level tiling of parallel dimensions; target-derived defaults."""
+
+    def compute(
+        self, op: ir.Operation | ir.OpView, ctx: StrategyContext
+    ) -> list[int] | None:
+        out_map = self.output_map(op)
+        if out_map is None:
+            return None
+
+        sizes = [0] * out_map.n_dims
+        parallel_dims, _ = parallel_and_reduction_dims(out_map)
+        if not parallel_dims:
+            return None
+
+        ov = opview(op)
+        if is_amx_bf16_contraction(ov, ctx.target):
+            inner_tiles = [32, 32]
+        elif is_f32_contraction(ov):
+            inner_tiles = [8, 32]
+        else:
+            inner_tiles = generic_parallel_tiles(ov, out_map, ctx.target)
+
+        assign_parallel_tiles(parallel_dims, inner_tiles, sizes)
+        disable_small_tiles(ov, out_map, sizes, max(inner_tiles, default=ctx.tile_size))
+        return sizes

lighthouse/dialects/transform/transform_ext/utils/tiling/target_caps.py

@@ -0,0 +1,74 @@
+from mlir import ir
+from mlir.dialects import linalg
+
+from lighthouse.execution.target import TargetInfo
+from lighthouse.utils.mlir import linalg_inputs, linalg_outputs, opview
+
+from .common import parallel_and_reduction_dims
+
+
+def _contraction_operand_types(
+    op: ir.Operation | ir.OpView,
+) -> tuple[ir.Type, ir.Type, ir.Type] | None:
+    """(lhs, rhs, acc) element types of a single-output contraction, else None."""
+    ov = opview(op)
+    if not linalg.isa_contraction_op(ov):
+        return None
+    inputs = linalg_inputs(ov)
+    outputs = linalg_outputs(ov)
+    if inputs is None or outputs is None or len(inputs) < 2 or len(outputs) != 1:
+        return None
+    return (
+        ir.ShapedType(inputs[0].type).element_type,
+        ir.ShapedType(inputs[1].type).element_type,
+        ir.ShapedType(outputs[0].type).element_type,
+    )
+
+
+def is_amx_bf16_contraction(
+    op: ir.Operation | ir.OpView, target: TargetInfo | None
+) -> bool:
+    """True for a bf16 -> f32 contraction on an AMX-capable target."""
+    if target is None or not target.is_supported("amx"):
+        return False
+    types = _contraction_operand_types(op)
+    if types is None:
+        return False
+    lhs, rhs, acc = types
+    return (
+        isinstance(lhs, ir.BF16Type)
+        and isinstance(rhs, ir.BF16Type)
+        and isinstance(acc, ir.F32Type)
+    )
+
+
+def is_f32_contraction(op: ir.Operation | ir.OpView) -> bool:
+    """True for a contraction with all-f32 operands (lhs, rhs and acc)."""
+    types = _contraction_operand_types(op)
+    return types is not None and all(isinstance(t, ir.F32Type) for t in types)
+
+
+def _vector_lane_count(target: TargetInfo | None, elem_type: ir.Type) -> int:
+    """SIMD lane count for `elem_type` on `target` (defaults assume 512-bit)."""
+    vector_bits = (
+        target.vector_register_width_bits
+        if target is not None and target.vector_register_width_bits is not None
+        else 512
+    )
+    if isinstance(elem_type, (ir.FloatType, ir.IntegerType)):
+        return max(1, vector_bits // max(1, elem_type.width))
+    return 16
+
+
+def generic_parallel_tiles(
+    op: ir.Operation | ir.OpView,
+    out_map: ir.AffineMap,
+    target: TargetInfo | None,
+) -> list[int]:
+    """SIMD-lane default parallel tiles for ops without a microkernel profile."""
+    parallel_dims, _ = parallel_and_reduction_dims(out_map)
+    if not parallel_dims:
+        return []
+    out_elem = ir.ShapedType(linalg_outputs(op)[0].type).element_type
+    inner = _vector_lane_count(target, out_elem)
+    return [inner] if len(parallel_dims) == 1 else [1, inner]

test/transform/test_assign_tile_sizes.py

@@ -66,6 +66,11 @@ def run(
         print(payload)
 
 
+# CHECK-LABEL: Test: strategy_attr_register_parallel
+# CHECK: linalg.matmul
+# CHECK-SAME: transform_ext.tile_sizes = array<i64: 8, 32, 0>
+run("strategy_attr_register_parallel", PAYLOAD, "linalg.matmul", "register_parallel")
+
 # CHECK-LABEL: Test: strategy_attr_cache
 # CHECK: linalg.matmul
 # CHECK-SAME: transform_ext.tile_sizes = array<i64: 32, 32, 0>

test/transform/test_tile_and_fuse_register_tiling.py

@@ -0,0 +1,147 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir import ir
+
+import lighthouse.dialects as lh_dialects
+from lighthouse.execution.target import TargetInfo
+from lighthouse.schedule import tile_and_fuse as tf
+
+
+def run(name: str, payload_str: str, *schedules):
+    print(f"Test: {name}", flush=True)
+    with ir.Context(), ir.Location.unknown():
+        lh_dialects.register_and_load()
+        payload = ir.Module.parse(payload_str)
+        modules = []
+        for make_schedule in schedules:
+            sched = make_schedule()
+            modules.append(sched)
+            sched.body.operations[0].apply(payload.operation)
+        print(payload)
+
+
+def run_with_target_override(
+    name: str,
+    payload_str: str,
+    *,
+    features: list[str] | None = None,
+    arch: str | None = None,
+    schedules,
+):
+    with TargetInfo.override(features=features, arch=arch):
+        run(name, payload_str, *schedules)
+
+
+def assign_register_parallel():
+    return tf.assign_and_propagate_tile_sizes(
+        tile_size=32,
+        strategy="register_parallel",
+    )
+
+
+def tile_and_fuse_keep():
+    return tf.tile_and_fuse_annotated(clear_annotations=False)
+
+
+MLP = """
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+module {
+  func.func @main(%arg0: tensor<32x64xf32>, %w: tensor<64x128xf32>, %b: tensor<128xf32>)
+        -> tensor<32x128xf32> {
+    %cst = arith.constant 0.000000e+00 : f32
+    %1 = tensor.empty() : tensor<32x128xf32>
+    %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<32x128xf32>) -> tensor<32x128xf32>
+    %3 = linalg.matmul ins(%arg0, %w : tensor<32x64xf32>, tensor<64x128xf32>)
+        outs(%2 : tensor<32x128xf32>) -> tensor<32x128xf32>
+    %4 = linalg.generic {indexing_maps = [#map, #map1, #map],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%3, %b : tensor<32x128xf32>, tensor<128xf32>)
+        outs(%1 : tensor<32x128xf32>) {
+    ^bb0(%in: f32, %in_2: f32, %out: f32):
+      %6 = arith.addf %in, %in_2 : f32
+      linalg.yield %6 : f32
+    } -> tensor<32x128xf32>
+    %5 = linalg.generic {indexing_maps = [#map, #map],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%4 : tensor<32x128xf32>)
+        outs(%1 : tensor<32x128xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %6 = arith.cmpf ugt, %in, %cst : f32
+      %7 = arith.select %6, %in, %cst : f32
+      linalg.yield %7 : f32
+    } -> tensor<32x128xf32>
+    return %5 : tensor<32x128xf32>
+  }
+}
+"""
+
+MLP_BF16 = """
+#map = affine_map<(d0, d1) -> (d0, d1)>
+#map1 = affine_map<(d0, d1) -> (d1)>
+module {
+  func.func @main(%arg0: tensor<32x64xbf16>, %w: tensor<64x128xbf16>, %b: tensor<128xf32>)
+        -> tensor<32x128xf32> {
+    %cst = arith.constant 0.000000e+00 : f32
+    %1 = tensor.empty() : tensor<32x128xf32>
+    %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<32x128xf32>) -> tensor<32x128xf32>
+    %3 = linalg.matmul ins(%arg0, %w : tensor<32x64xbf16>, tensor<64x128xbf16>)
+        outs(%2 : tensor<32x128xf32>) -> tensor<32x128xf32>
+    %4 = linalg.generic {indexing_maps = [#map, #map1, #map],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%3, %b : tensor<32x128xf32>, tensor<128xf32>)
+        outs(%1 : tensor<32x128xf32>) {
+    ^bb0(%in: f32, %in_2: f32, %out: f32):
+      %6 = arith.addf %in, %in_2 : f32
+      linalg.yield %6 : f32
+    } -> tensor<32x128xf32>
+    %5 = linalg.generic {indexing_maps = [#map, #map],
+        iterator_types = ["parallel", "parallel"]}
+        ins(%4 : tensor<32x128xf32>)
+        outs(%1 : tensor<32x128xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %6 = arith.cmpf ugt, %in, %cst : f32
+      %7 = arith.select %6, %in, %cst : f32
+      linalg.yield %7 : f32
+    } -> tensor<32x128xf32>
+    return %5 : tensor<32x128xf32>
+  }
+}
+"""
+
+
+# CHECK-LABEL: Test: register_parallel_tile_and_fuse
+# CHECK: linalg.matmul {transform_ext.tile_sizes = array<i64: 8, 32, 0>}
+# CHECK: linalg.generic
+# CHECK-SAME: transform_ext.tile_sizes = array<i64: 8, 32>
+# CHECK: scf.forall
+run(
+    "register_parallel_tile_and_fuse", MLP, assign_register_parallel, tile_and_fuse_keep
+)
+
+
+# Register-parallel strategy can infer target defaults from features.
+# CHECK-LABEL: Test: register_parallel_target_defaults_amx
+# CHECK: linalg.matmul {transform_ext.tile_sizes = array<i64: 32, 32, 0>}
+# CHECK: linalg.generic
+# CHECK-SAME: transform_ext.tile_sizes = array<i64: 32, 32>
+run_with_target_override(
+    "register_parallel_target_defaults_amx",
+    MLP_BF16,
+    features=["amx_tile"],
+    schedules=(assign_register_parallel, tile_and_fuse_keep),
+)
+
+
+# An AMX-capable target must not change f32 GEMM tiling: the AMX branch only
+# applies to bf16 x bf16 -> f32 contractions, so f32 keeps its own defaults.
+# CHECK-LABEL: Test: register_parallel_f32_under_amx_target
+# CHECK: linalg.matmul {transform_ext.tile_sizes = array<i64: 8, 32, 0>}
+# CHECK: linalg.generic
+# CHECK-SAME: transform_ext.tile_sizes = array<i64: 8, 32>
+run_with_target_override(
+    "register_parallel_f32_under_amx_target",
+    MLP,
+    features=["amx_tile"],
+    schedules=(assign_register_parallel, tile_and_fuse_keep),
+)

test/transform/test_tile_size_ops.py

@@ -193,3 +193,18 @@ def build_get_leading_unit_tile_sizes_schedule(op_name: str):
     BATCH_MATMUL_PAYLOAD,
     lambda: build_assign_schedule("linalg.batch_matmul"),
 )
+
+
+# Register-parallel strategy: batch outer dim tiled to 1, inner M/N tiled by
+# f32 FMA defaults, K untiled.
+# CHECK-LABEL: Test: register_parallel_strategy
+# CHECK: linalg.batch_matmul
+# CHECK-SAME: transform_ext.tile_sizes = array<i64: 1, 8, 32, 0>
+run(
+    "register_parallel_strategy",
+    BATCH_MATMUL_PAYLOAD,
+    lambda: build_assign_schedule(
+        "linalg.batch_matmul",
+        strategy="register_parallel",
+    ),
+)