From b27f4df3461b1fbed736132b049876a4ca4b7075 Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Fri, 14 Aug 2026 16:15:28 +0200 Subject: [PATCH 1/2] [schedule] Space Filling Curve (SFC) iteration space remapping Adds SFC-based transform that remaps contractions' parallel iteration space to improve data locality improving overall GEMM performance. The remapping is currently limited to 2D iteration spaces. Based on: arXiv:2601.16294 Assisted-by: Copilot --- .../transform/transform_ext/__init__.py | 2 + .../transform_ext/ops/sfc_remap_forall.py | 258 ++++++++ .../transform/transform_ext/utils/sfc.py | 62 ++ .../descriptors/x86_64/matmul-like.yaml | 3 +- lighthouse/schedule/sfc.py | 30 + test/transform/test_sfc_remap_forall.py | 574 ++++++++++++++++++ 6 files changed, 928 insertions(+), 1 deletion(-) create mode 100644 lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py create mode 100644 lighthouse/dialects/transform/transform_ext/utils/sfc.py create mode 100644 lighthouse/schedule/sfc.py create mode 100644 test/transform/test_sfc_remap_forall.py diff --git a/lighthouse/dialects/transform/transform_ext/__init__.py b/lighthouse/dialects/transform/transform_ext/__init__.py index 363bd6b4..3d230f48 100644 --- a/lighthouse/dialects/transform/transform_ext/__init__.py +++ b/lighthouse/dialects/transform/transform_ext/__init__.py @@ -25,6 +25,7 @@ from .ops.clear_tile_and_fuse_annotations import clear_tile_and_fuse_annotations from .ops.get_fusion_roots import get_fusion_roots from .ops.propagate_tile_sizes import propagate_tile_sizes +from .ops.sfc_remap_forall import sfc_remap_forall __all__ = [ "TransformExtensionDialect", @@ -51,6 +52,7 @@ "replace", "replace_with_fused_attention", "reverse_handles", + "sfc_remap_forall", "trace_producers", "update_address_space", "wrap_in_benching_func", diff --git a/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py b/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py new file mode 100644 index 00000000..7bf3abc0 --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py @@ -0,0 +1,258 @@ +from array import array + +from mlir import ir +from mlir.dialects import arith, ext, linalg, scf, tensor, transform +from mlir.dialects.transform import DiagnosedSilenceableFailure + +from lighthouse.dialects.transform.transform_ext import TransformExtensionDialect +from lighthouse.dialects.transform.transform_ext.utils.sfc import gilbert2d + + +def _static_bounds(loop: scf.ForallOp) -> tuple[list[int], list[int], list[int]] | None: + """Return static forall bounds, or None when any bound is dynamic.""" + lower = list(loop.staticLowerBound) + upper = list(loop.staticUpperBound) + steps = list(loop.staticStep) + if any( + value == ir.ShapedType.get_dynamic_size() for value in lower + upper + steps + ): + return None + return lower, upper, steps + + +def _uses_value(value: ir.Value, target: ir.BlockArgument, visited: set[int]) -> bool: + """Return whether value depends on the given block argument.""" + if value == target: + return True + if isinstance(value, ir.BlockArgument): + return False + owner = value.owner + if id(owner) in visited: + return False + visited.add(id(owner)) + return any(_uses_value(operand, target, visited) for operand in owner.operands) + + +def _used_loop_ivs( + loop: scf.ForallOp, offsets: list[ir.Value] +) -> list[ir.BlockArgument]: + """Return loop IVs used by at least one of the given offset values.""" + return [ + iv + for iv in loop.induction_variables + if any(_uses_value(offset, iv, set()) for offset in offsets) + ] + + +def _relevant_forall( + target: ir.Operation, +) -> tuple[scf.ForallOp, list[ir.BlockArgument]] | None: + """Find the enclosing two-dimensional forall defining both tile IVs.""" + offsets = [] + for operand in target.operands: + for extract in _find_slices(operand, set()): + offsets.extend(extract.offsets) + + owner = target.parent + while owner is not None: + if isinstance(owner.opview, scf.ForallOp): + loop = owner.opview + if len(loop.induction_variables) == 2: + used_ivs = _used_loop_ivs(loop, offsets) + if len(used_ivs) == 2: + return loop, used_ivs + owner = owner.parent + return None + + +def _find_slices(value: ir.Value, visited: set[int]) -> list[tensor.ExtractSliceOp]: + """Collect extract-slice operations reachable from an SSA value.""" + if isinstance(value, ir.BlockArgument): + return [] + owner = value.owner + if id(owner) in visited: + return [] + visited.add(id(owner)) + slices = [] + if isinstance(owner.opview, tensor.ExtractSliceOp): + slices.append(owner.opview) + for operand in owner.operands: + slices.extend(_find_slices(operand, visited)) + return slices + + +def _replace_operands( + operation: ir.Operation, + value_map: dict[ir.Value, ir.Value], + operation_map: dict[ir.Operation, ir.Operation], +) -> ir.WalkResult: + """Remap operands in a cloned operation and its nested operations.""" + for index, operand in enumerate(operation.operands): + replacement = value_map.get(operand) + if replacement is None and not isinstance(operand, ir.BlockArgument): + producer = operation_map.get(operand.owner) + if producer is not None: + replacement = producer.results[operand.result_number] + if replacement is not None: + operation.operands[index] = replacement + return ir.WalkResult.ADVANCE + + +def _clone_body(old_loop: scf.ForallOp, new_loop: scf.ForallOp, ivs: list[ir.Value]): + """Clone a forall body while remapping IVs, outputs, and local results.""" + old_block = old_loop.region.blocks[0] + new_block = new_loop.region.blocks[0] + + old_ivs = list(old_loop.induction_variables) + new_ivs = list(new_loop.induction_variables) + value_map: dict[ir.Value, ir.Value] = { + old_iv: new_iv for old_iv, new_iv in zip(old_ivs, ivs) + } + value_map.update( + { + old_out: new_out + for old_out, new_out in zip( + old_block.arguments[len(old_ivs) :], + new_block.arguments[len(new_ivs) :], + ) + } + ) + + operation_map = {} + with ir.InsertionPoint(new_block): + for old_operation in old_block.operations: + new_operation = old_operation.clone() + operation_map[old_operation] = new_operation + new_operation.walk( + lambda nested: _replace_operands(nested, value_map, operation_map) + ) + value_map.update(dict(zip(old_operation.results, new_operation.results))) + + +def _constant_table(values: list[int], location: ir.Location) -> ir.Value: + """Create an i64 tensor constant containing SFC coordinates.""" + element_type = ir.IntegerType.get_signless(64) + tensor_type = ir.RankedTensorType.get([len(values)], element_type) + dense = ir.DenseElementsAttr.get(array("q", values), type=element_type) + return arith.ConstantOp(tensor_type, dense, loc=location).result + + +def _rewrite( + old_loop: scf.ForallOp, + used_ivs: list[ir.BlockArgument], + rewriter: transform.TransformRewriter, +) -> scf.ForallOp | None: + """Replace a two-dimensional forall with its one-dimensional SFC traversal.""" + bounds = _static_bounds(old_loop) + if bounds is None: + return None + lower, upper, steps = bounds + if len(lower) != 2 or lower != [0, 0] or steps != [1, 1]: + return None + + old_ivs = list(old_loop.induction_variables) + if len(old_ivs) != 2 or used_ivs != old_ivs: + return None + + axis0, axis1 = used_ivs + height = upper[old_ivs.index(axis0)] + width = upper[old_ivs.index(axis1)] + points = list(gilbert2d(width, height)) + if len(points) != height * width: + return None + # The SFC is generated in (x, y) coordinates. Keep the row and column + # tables explicit while rewriting the flattened loop to avoid binding the + # indices to the wrong axes when the original 2D forall is collapsed into 1D. + m_indices = [row for _, row in points] + n_indices = [col for col, _ in points] + assert all(0 <= row < height for row in m_indices) + assert all(0 <= col < width for col in n_indices) + + with ir.InsertionPoint(old_loop.operation), old_loop.location: + m_table = _constant_table(m_indices, old_loop.location) + n_table = _constant_table(n_indices, old_loop.location) + new_loop = scf.ForallOp( + [0], + [height * width], + [1], + shared_outs=list(old_loop.outputs), + loc=old_loop.location, + ) + loop_iv = list(new_loop.induction_variables)[0] + with ir.InsertionPoint(new_loop.region.blocks[0]): + m_value = tensor.ExtractOp(m_table, [loop_iv], loc=old_loop.location).result + n_value = tensor.ExtractOp(n_table, [loop_iv], loc=old_loop.location).result + m_index = arith.IndexCastOp( + ir.IndexType.get(), m_value, loc=old_loop.location + ).result + n_index = arith.IndexCastOp( + ir.IndexType.get(), n_value, loc=old_loop.location + ).result + _clone_body(old_loop, new_loop, [m_index, n_index]) + + rewriter.replace_op(old_loop.operation, new_loop.operation) + return new_loop + + +class SfcRemapForallOp(TransformExtensionDialect.Operation, name="sfc_remap_forall"): + """Remap a tiled contraction's two-dimensional forall with a 2D SFC.""" + + target: ext.Operand[transform.AnyOpType] + remapped_loop: ext.Result[transform.AnyOpType[()]] = ext.infer_result() + + @classmethod + def attach_interface_impls(cls, context=None): + cls.TransformOpInterfaceModel.attach(cls.OPERATION_NAME, context=context) + cls.MemoryEffectsOpInterfaceModel.attach(cls.OPERATION_NAME, context=context) + + class TransformOpInterfaceModel(transform.TransformOpInterface): + @staticmethod + def apply( + op: "SfcRemapForallOp", + _rewriter: transform.TransformRewriter, + results: transform.TransformResults, + state: transform.TransformState, + ) -> DiagnosedSilenceableFailure: + loops = [] + targets = list(state.get_payload_ops(op.target)) + + # Gather all unique loop ops as rewriting might invalidate + # target handles and to avoid duplicate work. + loops_by_relevant_ivs = [] + seen_loops: set[ir.Operation] = set() + for target in targets: + if not linalg.isa_contraction_op(target): + continue + relevant = _relevant_forall(target.operation) + if relevant is not None: + loop, used_ivs = relevant + if loop.operation in seen_loops: + continue + seen_loops.add(loop.operation) + loops_by_relevant_ivs.append((loop, used_ivs)) + + # Try to remap the gathered loops. + for loop, used_ivs in loops_by_relevant_ivs: + loop = _rewrite(loop, used_ivs, _rewriter) + if loop is not None: + loops.append(loop) + results.set_ops(op.remapped_loop, loops) + return DiagnosedSilenceableFailure.Success + + @staticmethod + def allow_repeated_handle_operands(_op: "SfcRemapForallOp") -> bool: + return False + + class MemoryEffectsOpInterfaceModel(ir.MemoryEffectsOpInterface): + @staticmethod + def get_effects(op: "SfcRemapForallOp", effects): + transform.consumes_handle(op.op_operands, effects) + transform.produces_handle(op.results, effects) + transform.modifies_payload(effects) + + +def sfc_remap_forall( + target: ir.Value[transform.AnyOpType], +) -> ir.Value[transform.AnyOpType]: + """Create an SFC forall-remapping transform operation.""" + return SfcRemapForallOp(target=target).remapped_loop diff --git a/lighthouse/dialects/transform/transform_ext/utils/sfc.py b/lighthouse/dialects/transform/transform_ext/utils/sfc.py new file mode 100644 index 00000000..b512b20f --- /dev/null +++ b/lighthouse/dialects/transform/transform_ext/utils/sfc.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: BSD-2-Clause +# Copyright (c) 2018 Jakub Cervený + +"""Compile-time space-filling curve generators.""" + + +def gilbert2d(width: int, height: int): + """Generate a generalized Hilbert traversal of a 2D rectangle.""" + if width <= 0 or height <= 0: + raise ValueError("SFC dimensions must be positive") + if width >= height: + yield from _generate2d(0, 0, width, 0, 0, height) + else: + yield from _generate2d(0, 0, 0, height, width, 0) + + +def _sign(value: int) -> int: + return -1 if value < 0 else (1 if value > 0 else 0) + + +def _generate2d(x: int, y: int, ax: int, ay: int, bx: int, by: int): + width = abs(ax + ay) + height = abs(bx + by) + dax, day = _sign(ax), _sign(ay) + dbx, dby = _sign(bx), _sign(by) + + if height == 1: + for _ in range(width): + yield x, y + x, y = x + dax, y + day + return + + if width == 1: + for _ in range(height): + yield x, y + x, y = x + dbx, y + dby + return + + ax2, ay2 = ax // 2, ay // 2 + bx2, by2 = bx // 2, by // 2 + width2 = abs(ax2 + ay2) + height2 = abs(bx2 + by2) + + if 2 * width > 3 * height: + if width2 % 2 and width > 2: + ax2, ay2 = ax2 + dax, ay2 + day + yield from _generate2d(x, y, ax2, ay2, bx, by) + yield from _generate2d(x + ax2, y + ay2, ax - ax2, ay - ay2, bx, by) + return + + if height2 % 2 and height > 2: + bx2, by2 = bx2 + dbx, by2 + dby + yield from _generate2d(x, y, bx2, by2, ax2, ay2) + yield from _generate2d(x + bx2, y + by2, ax, ay, bx - bx2, by - by2) + yield from _generate2d( + x + (ax - dax) + (bx2 - dbx), + y + (ay - day) + (by2 - dby), + -bx2, + -by2, + -(ax - ax2), + -(ay - ay2), + ) diff --git a/lighthouse/pipeline/descriptors/x86_64/matmul-like.yaml b/lighthouse/pipeline/descriptors/x86_64/matmul-like.yaml index 2ef6a16c..860b4502 100644 --- a/lighthouse/pipeline/descriptors/x86_64/matmul-like.yaml +++ b/lighthouse/pipeline/descriptors/x86_64/matmul-like.yaml @@ -4,8 +4,9 @@ Pipeline: block_factors=$block_factors tile_size=$tile_size } + - schedule: "sfc.py[gen=remap]" - ## CPU specific register tiling (depends on uArch & data type) + # CPU specific register tiling (depends on uArch & data type) - schedule: "x86/register_tiling.py[gen=matmul_register_tiling] { target=linalg.contract reg_tile_batch=$reg_tile_batch diff --git a/lighthouse/schedule/sfc.py b/lighthouse/schedule/sfc.py new file mode 100644 index 00000000..8907f8d9 --- /dev/null +++ b/lighthouse/schedule/sfc.py @@ -0,0 +1,30 @@ +from mlir import ir +from mlir.dialects import transform +from mlir.dialects.transform import structured + +from lighthouse.dialects.transform import transform_ext +from lighthouse.schedule.builders import schedule_boilerplate +import lighthouse.transform as lh_transform + + +def remap( + target_op: str | list[str] | None = None, +) -> ir.Module: + """ + Remaps `target_op` parent loop iteration space using + the space-filling curve strategy. + + Args: + target_op: Op(s) to consider. Defaults to all linalg ops. + + Returns: + Schedule + """ + if target_op is None: + target_op = structured.MatchInterfaceEnum.LinalgOp + + with schedule_boilerplate() as (schedule, named_seq): + ops = lh_transform.match_op(named_seq.bodyTarget, target_op) + transform_ext.sfc_remap_forall(ops) + transform.yield_() + return schedule diff --git a/test/transform/test_sfc_remap_forall.py b/test/transform/test_sfc_remap_forall.py new file mode 100644 index 00000000..da9bee29 --- /dev/null +++ b/test/transform/test_sfc_remap_forall.py @@ -0,0 +1,574 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir import ir +from mlir.dialects import transform + +import lighthouse.dialects as lh_dialects +import lighthouse.transform as lh_transform +from lighthouse.dialects.transform import transform_ext +from lighthouse.dialects.transform.transform_ext.utils.sfc import gilbert2d +from lighthouse.schedule.builders import schedule_boilerplate +from lighthouse.schedule.sfc import remap as remap_sfc + + +def run(name: str, payload_str: str, *schedules): + """Parse a payload, apply the given schedules in order and print it.""" + 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) + payload.operation.verify() + print(payload) + + +def remap_matmul(): + with schedule_boilerplate() as (schedule, named): + target = lh_transform.match_op(named.bodyTarget, "linalg.matmul") + transform_ext.sfc_remap_forall(target) + transform.yield_() + return schedule + + +def remap_fill(): + """Apply SFC remap on a non-contraction op handle.""" + with schedule_boilerplate() as (schedule, named): + target = lh_transform.match_op(named.bodyTarget, "linalg.fill") + transform_ext.sfc_remap_forall(target) + transform.yield_() + return schedule + + +for width, height in ((1, 1), (2, 3), (15, 12)): + points = list(gilbert2d(width, height)) + assert len(points) == width * height + assert len(set(points)) == len(points) + assert set(points) == {(x, y) for x in range(width) for y in range(height)} + + +# An 8x2 by 2x4 GEMM tiled 1x1, so the C-tile grid is 8x4 (Mb x Nb). The A/B +# slice offsets on the two forall IVs must resolve to SFC lookup tables and a +# single flattened 1D loop over all 32 tiles. +MATMUL = """ +module { + func.func @main(%a: tensor<8x2xf32>, %b: tensor<2x4xf32>) -> tensor<8x4xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<8x4xf32> + %result = scf.forall (%i, %j) in (8, 4) shared_outs(%out = %empty) + -> (tensor<8x4xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<8x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x4xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<8x4xf32> + } + } + return %result : tensor<8x4xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_zero_fill +# CHECK: arith.constant dense<[0, 1, 1, 0, 0, 0, 1, 1, 2, 2, 3, 3, 3, 2, 2, 3, 4, 5, 5, 4, 4, 4, 5, 5, 6, 6, 7, 7, 7, 6, 6, 7]> : tensor<32xi64> +# CHECK: arith.constant dense<[0, 0, 1, 1, 2, 3, 3, 2, 2, 3, 3, 2, 1, 1, 0, 0, 0, 0, 1, 1, 2, 3, 3, 2, 2, 3, 3, 2, 1, 1, 0, 0]> : tensor<32xi64> +# CHECK: scf.forall ({{.*}}) in (32) +# CHECK: tensor.extract +# CHECK: linalg.fill +# CHECK: linalg.matmul +run("matmul_zero_fill", MATMUL, remap_matmul) + + +# Nested 1D IV is unrelated to the tile offsets; the outer forall is relevant. +NESTED_MATMUL = """ +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>) -> tensor<2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2xf32> + %result = scf.forall (%i, %j) in (2, 2) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %inner_empty = tensor.empty() : tensor<1x1xf32> + %inner = scf.forall (%k) in (1) shared_outs(%inner_out = %inner_empty) + -> (tensor<1x1xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %inner_out[0, 0] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<1x1xf32> + } + } + scf.forall.in_parallel { + tensor.parallel_insert_slice %inner into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %result : tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: nested_matmul +# CHECK: scf.forall ({{.*}}) in (4) +# CHECK: scf.forall ({{.*}}) in (1) +run("nested_matmul", NESTED_MATMUL, remap_matmul) + + +# IV declaration order sets orientation; swapping it preserves coverage. +SWAPPED_MATMUL = """ +module { + func.func @main(%a: tensor<8x2xf32>, %b: tensor<2x4xf32>) -> tensor<8x4xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<8x4xf32> + %result = scf.forall (%j, %i) in (4, 8) shared_outs(%out = %empty) + -> (tensor<8x4xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<8x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x4xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<8x4xf32> + } + } + return %result : tensor<8x4xf32> + } +} +""" + +# CHECK-LABEL: Test: swapped_matmul +# CHECK: scf.forall ({{.*}}) in (32) +run("swapped_matmul", SWAPPED_MATMUL, remap_matmul) + + +# Both matmuls share one parent forall. The transform must gather both relevant +# operations before replacing that parent and must clone the loop body once. +TWO_MATMULS_ONE_FORALL = """ +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>, + %c: tensor<2x2xf32>, %d: tensor<2x2xf32>) + -> (tensor<2x2xf32>, tensor<2x2xf32>) { + %cst = arith.constant 0.000000e+00 : f32 + %empty0 = tensor.empty() : tensor<2x2xf32> + %empty1 = tensor.empty() : tensor<2x2xf32> + %result0, %result1 = scf.forall (%i, %j) in (2, 2) + shared_outs(%out0 = %empty0, %out1 = %empty1) + -> (tensor<2x2xf32>, tensor<2x2xf32>) { + %as0 = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs0 = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty0 = tensor.empty() : tensor<1x1xf32> + %init0 = linalg.fill ins(%cst : f32) outs(%init_empty0 : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product0 = linalg.matmul ins(%as0, %bs0 : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init0 : tensor<1x1xf32>) -> tensor<1x1xf32> + %as1 = tensor.extract_slice %c[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs1 = tensor.extract_slice %d[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty1 = tensor.empty() : tensor<1x1xf32> + %init1 = linalg.fill ins(%cst : f32) outs(%init_empty1 : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product1 = linalg.matmul ins(%as1, %bs1 : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init1 : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product0 into %out0[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + tensor.parallel_insert_slice %product1 into %out1[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %result0, %result1 : tensor<2x2xf32>, tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: two_matmuls_one_forall +# CHECK-COUNT-1: scf.forall ({{.*}}) in (4) +# CHECK-COUNT-2: linalg.matmul +run("two_matmuls_one_forall", TWO_MATMULS_ONE_FORALL, remap_matmul) + +# Exercise the production schedule, which must match only contractions before +# replacing their parent foralls. Matching every linalg op would leave sibling +# fill/elementwise handles pointing into an erased forall. +# CHECK-LABEL: Test: matmul_production_schedule +# CHECK: scf.forall ({{.*}}) in (32) +# CHECK: linalg.matmul +run("matmul_production_schedule", MATMUL, remap_sfc) + + +# A GEMM with an elementwise prologue (scaling the A slice) and an elementwise +# epilogue (doubling the accumulator). Offsets reach the forall IVs through the +# wrapping linalg.generic ops, so SFC-remapping must still trace through them +# and flatten the loop, keeping the prologue/epilogue ops in the cloned body. +PROLOGUE_EPILOGUE = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>) -> tensor<2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %scale = arith.constant 2.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2xf32> + %result = scf.forall (%i, %j) in (2, 2) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %prologue = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%as : tensor<1x2xf32>) outs(%as : tensor<1x2xf32>) { + ^bb0(%in: f32, %o: f32): + %scaled = arith.mulf %in, %scale : f32 + linalg.yield %scaled : f32 + } -> tensor<1x2xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%prologue, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + %epilogue = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%product : tensor<1x1xf32>) outs(%product : tensor<1x1xf32>) { + ^bb0(%in: f32, %o: f32): + %doubled = arith.addf %in, %in : f32 + linalg.yield %doubled : f32 + } -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %epilogue into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %result : tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_prologue_epilogue +# CHECK: arith.constant dense<[0, 1, 1, 0]> : tensor<4xi64> +# CHECK: arith.constant dense<[0, 0, 1, 1]> : tensor<4xi64> +# CHECK: scf.forall ({{.*}}) in (4) +# CHECK: arith.mulf +# CHECK: linalg.matmul +# CHECK: arith.addf +run("matmul_prologue_epilogue", PROLOGUE_EPILOGUE, remap_matmul) + + +# Two dependent contractions must be collected before any parent loop is +# replaced. The transform must not revisit a target handle after mutation. +MULTIPLE_DEPENDENT_MATMULS = """ +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>) -> tensor<2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2xf32> + %first = scf.forall (%i, %j) in (2, 2) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + %second = scf.forall (%i, %j) in (2, 2) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %as = tensor.extract_slice %first[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %second : tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: multiple_dependent_matmuls +# CHECK-COUNT-2: scf.forall ({{.*}}) in (4) +run("multiple_dependent_matmuls", MULTIPLE_DEPENDENT_MATMULS, remap_matmul) + + +# The matmul reads the same %a[0, 0]/%b[0, 0] tile every iteration, so its +# operands cannot be traced to the forall's induction variables. SFC-remapping +# must skip it, leaving the 2D forall untouched. +INDEPENDENT_MATMUL = """ +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>) -> tensor<2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2xf32> + %result = scf.forall (%i, %j) in (2, 2) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %as = tensor.extract_slice %a[0, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, 0] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %result : tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_independent_of_forall +# CHECK: scf.forall ({{.*}}, {{.*}}) in (2, 2) +# CHECK-NOT: arith.constant dense +# CHECK-NOT: tensor.extract % +run("matmul_independent_of_forall", INDEPENDENT_MATMUL, remap_matmul) + + +# The enclosing forall is 3D. The transform requires an enclosing 2D forall, +# so it must skip rewriting. +THREED_FORALL_MATMUL = """ +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>) -> tensor<2x2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2x2xf32> + %result = scf.forall (%k, %i, %j) in (2, 2, 2) shared_outs(%out = %empty) + -> (tensor<2x2x2xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%k, %i, %j] [1, 1, 1] [1, 1, 1] + : tensor<1x1xf32> into tensor<2x2x2xf32> + } + } + return %result : tensor<2x2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_in_3d_forall +# CHECK: scf.forall ({{.*}}, {{.*}}, {{.*}}) in (2, 2, 2) +# CHECK-NOT: arith.constant dense +run("matmul_in_3d_forall", THREED_FORALL_MATMUL, remap_matmul) + + +# The matmul depends on only one forall IV (%i), with %j constant in B-slices. +# The transform requires both IVs to be relevant, so it must skip rewriting. +ONE_IV_MATMUL = """ +module { + func.func @main(%a: tensor<2x2xf32>, %b: tensor<2x2xf32>) -> tensor<2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2xf32> + %result = scf.forall (%i, %j) in (2, 2) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, 0] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %result : tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_one_relevant_iv +# CHECK: scf.forall ({{.*}}, {{.*}}) in (2, 2) +# CHECK-NOT: arith.constant dense +run("matmul_one_relevant_iv", ONE_IV_MATMUL, remap_matmul) + + +# The transform may be invoked on non-contraction handles, but must not rewrite +# the parent forall in that case. +# CHECK-LABEL: Test: non_contraction_target +# CHECK: scf.forall ({{.*}}, {{.*}}) in (8, 4) +# CHECK-NOT: arith.constant dense +run("non_contraction_target", MATMUL, remap_fill) + + +# Dynamic loop bounds are unsupported by this transform, so the loop should +# remain unchanged. +DYNAMIC_BOUNDS_MATMUL = """ +module { + func.func @main(%m: index, %n: index, %a: tensor<2x2xf32>, %b: tensor<2x2xf32>) + -> tensor<2x2xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<2x2xf32> + %result = scf.forall (%i, %j) in (%m, %n) shared_outs(%out = %empty) + -> (tensor<2x2xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<2x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x2xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<2x2xf32> + } + } + return %result : tensor<2x2xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_dynamic_bounds +# CHECK: scf.forall ({{.*}}, {{.*}}) in ({{.*}}, {{.*}}) +# CHECK-NOT: arith.constant dense +run("matmul_dynamic_bounds", DYNAMIC_BOUNDS_MATMUL, remap_matmul) + + +# Non-zero lower bounds are outside the supported shape, so the loop should +# remain unchanged. +NONZERO_LOWER_BOUNDS_MATMUL = """ +module { + func.func @main(%a: tensor<4x2xf32>, %b: tensor<2x4xf32>) -> tensor<4x4xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<4x4xf32> + %result = scf.forall (%i, %j) = (1, 0) to (3, 2) step (1, 1) + shared_outs(%out = %empty) -> (tensor<4x4xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<4x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x4xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<4x4xf32> + } + } + return %result : tensor<4x4xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_nonzero_lower_bounds +# CHECK: scf.forall ({{.*}}, {{.*}}) = (1, 0) to (3, 2) step (1, 1) +# CHECK-NOT: arith.constant dense +run("matmul_nonzero_lower_bounds", NONZERO_LOWER_BOUNDS_MATMUL, remap_matmul) + + +# Non-unit step sizes are outside the supported shape, so the loop should +# remain unchanged. +NONUNIT_STEP_MATMUL = """ +module { + func.func @main(%a: tensor<4x2xf32>, %b: tensor<2x4xf32>) -> tensor<4x4xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<4x4xf32> + %result = scf.forall (%i, %j) = (0, 0) to (4, 4) step (2, 1) + shared_outs(%out = %empty) -> (tensor<4x4xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<4x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x4xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<4x4xf32> + } + } + return %result : tensor<4x4xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_nonunit_step +# CHECK: scf.forall ({{.*}}, {{.*}}) = (0, 0) to (4, 4) step (2, 1) +# CHECK-NOT: arith.constant dense +run("matmul_nonunit_step", NONUNIT_STEP_MATMUL, remap_matmul) + + +# Dynamic step operands are unsupported by this transform, so the loop should +# remain unchanged. +DYNAMIC_STEP_MATMUL = """ +module { + func.func @main(%si: index, %a: tensor<4x2xf32>, %b: tensor<2x4xf32>) + -> tensor<4x4xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %empty = tensor.empty() : tensor<4x4xf32> + %result = scf.forall (%i, %j) = (0, 0) to (4, 4) step (%si, 1) + shared_outs(%out = %empty) -> (tensor<4x4xf32>) { + %as = tensor.extract_slice %a[%i, 0] [1, 2] [1, 1] + : tensor<4x2xf32> to tensor<1x2xf32> + %bs = tensor.extract_slice %b[0, %j] [2, 1] [1, 1] + : tensor<2x4xf32> to tensor<2x1xf32> + %init_empty = tensor.empty() : tensor<1x1xf32> + %init = linalg.fill ins(%cst : f32) outs(%init_empty : tensor<1x1xf32>) + -> tensor<1x1xf32> + %product = linalg.matmul ins(%as, %bs : tensor<1x2xf32>, tensor<2x1xf32>) + outs(%init : tensor<1x1xf32>) -> tensor<1x1xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %product into %out[%i, %j] [1, 1] [1, 1] + : tensor<1x1xf32> into tensor<4x4xf32> + } + } + return %result : tensor<4x4xf32> + } +} +""" + +# CHECK-LABEL: Test: matmul_dynamic_step +# CHECK: scf.forall ({{.*}}, {{.*}}) = (0, 0) to (4, 4) step ({{.*}}, 1) +# CHECK-NOT: arith.constant dense +run("matmul_dynamic_step", DYNAMIC_STEP_MATMUL, remap_matmul) From 7e368184b41fb60e320ab10e8a34d9d89107fb44 Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Fri, 21 Aug 2026 11:36:39 +0200 Subject: [PATCH 2/2] Fix after rebase --- .../transform/transform_ext/ops/sfc_remap_forall.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py b/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py index 7bf3abc0..b9f7825f 100644 --- a/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py +++ b/lighthouse/dialects/transform/transform_ext/ops/sfc_remap_forall.py @@ -245,10 +245,12 @@ def allow_repeated_handle_operands(_op: "SfcRemapForallOp") -> bool: class MemoryEffectsOpInterfaceModel(ir.MemoryEffectsOpInterface): @staticmethod - def get_effects(op: "SfcRemapForallOp", effects): - transform.consumes_handle(op.op_operands, effects) - transform.produces_handle(op.results, effects) - transform.modifies_payload(effects) + def get_effects(op: "SfcRemapForallOp"): + return ( + transform.consumes_handle(op.op_operands) + + transform.produces_handle(op.results) + + transform.modifies_payload() + ) def sfc_remap_forall(