From ab7d154a821e60ce7f39c17d5616388df32dbf33 Mon Sep 17 00:00:00 2001 From: Adam Siemieniuk Date: Thu, 13 Aug 2026 14:32:58 +0200 Subject: [PATCH] [transform] Propagate tile sizes through loops Extends tile size propagation to optionally reach producers and consumers through loops. That propagation mode is mostly useful when selecting tiles in later stages of multi-level tiling e.g., toward register unrolling. --- .../transform_ext/ops/propagate_tile_sizes.py | 192 +++++- .../transform_ext/utils/tile_propagation.py | 128 +++- lighthouse/schedule/tile_and_fuse.py | 14 +- test/transform/test_propagate_tile_sizes.py | 627 ++++++++++++++++++ test/transform/test_tile_and_fuse.py | 64 ++ 5 files changed, 993 insertions(+), 32 deletions(-) create mode 100644 test/transform/test_propagate_tile_sizes.py diff --git a/lighthouse/dialects/transform/transform_ext/ops/propagate_tile_sizes.py b/lighthouse/dialects/transform/transform_ext/ops/propagate_tile_sizes.py index c2d6a116..6d77c585 100644 --- a/lighthouse/dialects/transform/transform_ext/ops/propagate_tile_sizes.py +++ b/lighthouse/dialects/transform/transform_ext/ops/propagate_tile_sizes.py @@ -1,6 +1,8 @@ from mlir import ir from mlir.dialects import ext, transform from mlir.dialects.transform import DiagnosedSilenceableFailure +from collections import deque +from collections.abc import Iterator, Sequence from lighthouse.dialects.transform.transform_ext import TransformExtensionDialect from lighthouse.dialects.transform.transform_ext.utils import tile_size_analysis as tsa @@ -30,13 +32,21 @@ class PropagateTileSizesOp( Barriers are never re-tiled. Reduction dimensions are never tiled, and already-annotated ops keep their sizes. + By default, `scf.for` loops act as barriers: propagation does not cross into or + out of a loop body. Set `propagate_through_loops` to bridge a value through a + loop's iter args/results and continue propagating on the other side. + Args: root: Handle to annotated anchor op(s). + propagate_through_loops: Optional bool (default: false). When true, tile + sizes are propagated through `scf.for` loop-carried values instead of + stopping at the loop. Return: Handle to all annotated ops after propagation (roots plus newly annotated). """ root: ext.Operand[transform.AnyOpType] + propagate_through_loops: ext.Operand[transform.AnyParamType] | None = None annotated: ext.Result[transform.AnyOpType[()]] = ext.infer_result() @classmethod @@ -52,23 +62,56 @@ def apply( results: transform.TransformResults, state: transform.TransformState, ) -> DiagnosedSilenceableFailure: + """Run the two-phase (forward then backward) tile-size propagation.""" root_ops = list(state.get_payload_ops(op.root)) + through_loops = False + if op.propagate_through_loops is not None: + param_attr = state.get_params(op.propagate_through_loops) + if len(param_attr) == 1 and isinstance(param_attr[0], ir.BoolAttr): + through_loops = bool(param_attr[0]) + # An op is "visited" once it carries an annotation, which prevents # re-processing and acts as a barrier. Track ordered, de-duplicated # annotated ops for the result handle. annotated: list[ir.Operation] = [] - seen: set = set() + seen: set[ir.Operation] = set() def remember(target_op: ir.Operation) -> None: - key = target_op.operation.__hash__() - if key not in seen: - seen.add(key) + """Record `target_op` in `annotated`, once, for the result handle.""" + if target_op not in seen: + seen.add(target_op) annotated.append(target_op) def claim( - src: ir.Operation, src_sizes, shared: ir.Value, dst: ir.Operation + src: ir.Operation, + src_sizes: Sequence[int], + src_shared: ir.Value, + dst_shared: ir.Value, + dst: ir.Operation, ) -> ir.Operation | None: + """Try to tile `dst` from `src`'s sizes across their shared tensor. + + `src`'s sizes are projected onto the shared tensor and mapped into + `dst`'s iteration space, so both sides tile the tensor identically. + + `dst` is left untouched when it is a barrier, when it already carries + an annotation, or when no non-zero sizes can be derived for it. An + already-annotated `dst` that disagrees with `src` on the shared tensor + belongs to a different fusion group: the consumer side is marked as a + fusion boundary so grouping can split the two cheaply later. + + Args: + src: Annotated op the tile sizes are propagated from. + src_sizes: `src`'s tile sizes, in its loop order. + src_shared: Value of the shared tensor as seen by `src`. + dst_shared: Value of the shared tensor as seen by `dst`. May + differ from `src_shared` when the tensor crosses a wrapper + op such as `scf.yield`. + dst: Neighboring op to annotate. + Returns: + `dst` if it was newly annotated, else None. + """ dst_sizes = tsa.get_tile_sizes_attr(dst) if dst_sizes is not None: # `dst` already carries a tiling. If it disagrees with `src` @@ -76,21 +119,123 @@ def claim( # fusion groups; mark the consumer side (the op reading the # shared tensor) as a boundary so grouping can split them # cheaply without recomputing compatibility. - if not tp.compatible_on_value( - src, src_sizes, dst, dst_sizes, shared + if not tp.compatible_on_values( + src, + src_sizes, + src_shared, + dst, + dst_sizes, + dst_shared, ): - consumer = src if any(r == shared for r in dst.results) else dst + consumer = ( + src if any(r == dst_shared for r in dst.results) else dst + ) fa.mark_fusion_boundary(consumer) return None if not tp.is_propagatable(dst): return None - dst_sizes = tp.propagate_through_value(src, src_sizes, shared, dst) + dst_sizes = tp.propagate_through_values( + src, + src_sizes, + src_shared, + dst_shared, + dst, + ) if dst_sizes is None or not any(dst_sizes): return None tsa.set_tile_sizes_attr(dst, dst_sizes) remember(dst) return dst + def forward_loop_passthrough( + value: ir.Value, user: ir.Operation + ) -> ir.Value | None: + """Follow `value` forward through a loop's body to the loop's result.""" + # Bridge `%v` -> `scf.yield %v` -> `%for_result` for scf.for. + if user.operation.name != "scf.yield": + return None + parent = user.operation.parent + if parent is None or parent.name != "scf.for": + return None + operands = list(user.opview.operands) + try: + idx = operands.index(value) + except ValueError: + return None + results = list(parent.opview.results) + return results[idx] if idx < len(results) else None + + def backward_loop_passthrough(value: ir.Value) -> ir.Value | None: + """Follow `value` backward from a loop-carried argument to its initial value.""" + # Bridge iter arg -> its initial value in the enclosing scf.for. + block = value.owner + if not isinstance(block, ir.Block): + return None + parent_op = block.owner.operation + if parent_op.name != "scf.for": + return None + # arg 0 is the induction variable; iter args start at 1. + args = list(block.arguments) + try: + arg_num = args.index(value) + except ValueError: + return None + if arg_num == 0: + return None + # scf.for operands: lb(0), ub(1), step(2), iter_init_0(3), ... + init_index = arg_num + 2 + operands = list(parent_op.opview.operands) + if init_index >= len(operands): + return None + return operands[init_index] + + def forward_neighbors( + src_shared: ir.Value, + ) -> Iterator[tuple[ir.Value, ir.Value, ir.Operation]]: + """Yield `(src_shared, dst_shared, user)` for each propagatable + consumer of `src_shared`, transparently following it through loops. + """ + queue = deque([src_shared]) + seen_values = {src_shared} + while queue: + value = queue.popleft() + for user in op_users(value): + if tp.is_propagatable(user): + yield src_shared, value, user + continue + if not through_loops: + continue + next_from_loop = forward_loop_passthrough(value, user) + if next_from_loop is not None: + if next_from_loop in seen_values: + continue + seen_values.add(next_from_loop) + queue.append(next_from_loop) + + def backward_neighbors( + src_shared: ir.Value, + ) -> Iterator[tuple[ir.Value, ir.Value, ir.Operation]]: + """Yield `(src_shared, dst_shared, producer)` for each propagatable + producer of `src_shared`, transparently following it through loops. + """ + queue = deque([src_shared]) + seen_values = {src_shared} + while queue: + value = queue.popleft() + producer = defining_op(value) + if producer is None: + if not through_loops: + continue + prev_from_loop = backward_loop_passthrough(value) + if prev_from_loop is not None: + if prev_from_loop in seen_values: + continue + seen_values.add(prev_from_loop) + queue.append(prev_from_loop) + continue + if tp.is_propagatable(producer): + yield src_shared, value, producer + seeds = [r for r in root_ops if tsa.get_tile_sizes_attr(r) is not None] for seed in seeds: remember(seed) @@ -103,8 +248,8 @@ def claim( idx += 1 src_sizes = tsa.get_tile_sizes_attr(src) for result in src.opview.results: - for user in op_users(result): - dst = claim(src, src_sizes, result, user) + for src_shared, dst_shared, user in forward_neighbors(result): + dst = claim(src, src_sizes, src_shared, dst_shared, user) if dst is not None: forward.append(dst) @@ -119,18 +264,17 @@ def claim( idx += 1 src_sizes = tsa.get_tile_sizes_attr(src) for operand in src.opview.operands: - producer = defining_op(operand) - if producer is None: - continue - dst = claim(src, src_sizes, operand, producer) - if dst is not None: - backward.append(dst) + for src_shared, dst_shared, producer in backward_neighbors(operand): + dst = claim(src, src_sizes, src_shared, dst_shared, producer) + if dst is not None: + backward.append(dst) results.set_ops(op.annotated, annotated) return DiagnosedSilenceableFailure.Success @staticmethod def allow_repeated_handle_operands(_op: "PropagateTileSizesOp") -> bool: + """Disallow the same payload op from being passed via multiple handles.""" return False class MemoryEffectsOpInterfaceModel(ir.MemoryEffectsOpInterface): @@ -145,13 +289,25 @@ def get_effects(op: ir.Operation): def propagate_tile_sizes( root: ir.Value[transform.AnyOpType], + propagate_through_loops: bool | ir.Value = False, ) -> ir.Value: """ snake_case wrapper to create a PropagateTileSizesOp. Args: root: Handle to annotated anchor op(s). + propagate_through_loops: Whether to propagate through `scf.for` loop-carried + values instead of treating loops as barriers (default: False). Returns: Handle to all ops carrying a tile-size annotation after propagation. """ - return PropagateTileSizesOp(root=root).annotated + if isinstance(propagate_through_loops, bool): + param_attr = ir.BoolAttr.get(propagate_through_loops) + propagate_through_loops = transform.ParamConstantOp( + transform.AnyParamType.get(), param_attr + ) + + return PropagateTileSizesOp( + root=root, + propagate_through_loops=propagate_through_loops, + ).annotated diff --git a/lighthouse/dialects/transform/transform_ext/utils/tile_propagation.py b/lighthouse/dialects/transform/transform_ext/utils/tile_propagation.py index 210b0d37..fdec1886 100644 --- a/lighthouse/dialects/transform/transform_ext/utils/tile_propagation.py +++ b/lighthouse/dialects/transform/transform_ext/utils/tile_propagation.py @@ -17,6 +17,11 @@ def is_propagatable(op: ir.Operation | ir.OpView) -> bool: True for any structured linalg op that is not a fusion barrier; non-linalg ops have no indexing maps to translate tiles through and are excluded. + + Args: + op: Candidate op to annotate. + Returns: + True if `op` can receive propagated tile sizes. """ return indexing_maps(op) is not None and not fa.is_fusion_barrier(op) @@ -24,7 +29,18 @@ def is_propagatable(op: ir.Operation | ir.OpView) -> bool: def _map_for_value( op: ir.OpView, value: ir.Value, maps: Sequence[ir.AffineMap] ) -> ir.AffineMap | None: - """Return the indexing map associated with `value` on `op`.""" + """Return the indexing map associated with `value` on `op`. + + `value` is matched against `op`'s inputs, then its outputs, then its results + (a DPS result shares the map of its positionally matching output operand). + + Args: + op: Structured linalg op owning the maps. + value: Operand or result of `op` to look up. + maps: `op`'s indexing maps, in operand order. + Returns: + The map indexing `value`, or None if `op` does not touch `value`. + """ inputs = linalg_inputs(op) outputs = linalg_outputs(op) if inputs is None or outputs is None: @@ -49,13 +65,18 @@ def tiles_on_value( ) -> list[int] | None: """Per-dimension tiles that `op`, tiled by `sizes`, induces on `value`. - Returns one entry per dimension of `value` (tensor-dim order); 0 means the - dimension is left untiled / not constrained by `op`. Returns None when `op` - is not a structured linalg op or does not touch `value`. - Projecting both a producer's and a consumer's sizes onto the shared tensor this way makes tile comparisons robust to transposition: differently ordered iteration spaces still agree when they tile the shared tensor identically. + + Args: + op: Structured linalg op the tiles originate from. + sizes: `op`'s tile sizes, in its loop order. + value: Operand or result of `op` to project the tiles onto. + Returns: + One entry per dimension of `value` (tensor-dim order), where 0 means the + dimension is left untiled / not constrained by `op`; None when `op` is not + a structured linalg op or does not touch `value`. """ ov = opview(op) maps = indexing_maps(ov) @@ -86,13 +107,58 @@ def compatible_on_value( sizes; a zero (untiled / broadcast / unconstrained) side is a wildcard and never conflicts. - Returns True when the tiles cannot be determined, so grouping errs toward - fusion rather than over-splitting. + Args: + src_op: First op sharing the tensor. + src_sizes: `src_op`'s tile sizes, in its loop order. + dst_op: Second op sharing the tensor. + dst_sizes: `dst_op`'s tile sizes, in its loop order. + shared: Tensor value both ops use. + Returns: + True if the two tilings agree on `shared`. Also True when the tiles + cannot be determined, so grouping errs toward fusion rather than + over-splitting. """ - a = tiles_on_value(src_op, src_sizes, shared) - b = tiles_on_value(dst_op, dst_sizes, shared) + return compatible_on_values( + src_op, + src_sizes, + shared, + dst_op, + dst_sizes, + shared, + ) + + +def compatible_on_values( + src_op: ir.Operation | ir.OpView, + src_sizes: Sequence[int], + src_shared: ir.Value, + dst_op: ir.Operation | ir.OpView, + dst_sizes: Sequence[int], + dst_shared: ir.Value, +) -> bool: + """Check compatibility when source and destination see aliased values. + + This is used when the shared tensor crosses through lightweight wrapper ops + (for example scf.yield), so the producer and consumer do not use the exact + same SSA value. + + Args: + src_op: First op sharing the tensor. + src_sizes: `src_op`'s tile sizes, in its loop order. + src_shared: Shared tensor as seen by `src_op`. + dst_op: Second op sharing the tensor. + dst_sizes: `dst_op`'s tile sizes, in its loop order. + dst_shared: Shared tensor as seen by `dst_op`. + Returns: + True if the two tilings agree on the shared tensor, or if they cannot + be determined. + """ + a = tiles_on_value(src_op, src_sizes, src_shared) + b = tiles_on_value(dst_op, dst_sizes, dst_shared) if a is None or b is None: return True + if len(a) != len(b): + return True return all(x == y for x, y in zip(a, b) if x != 0 and y != 0) @@ -107,7 +173,43 @@ def propagate_through_value( The shared tensor's per-dimension tiles are derived from `src_op`'s sizes and mapped onto `dst_op`'s iteration space; reduction dims of `dst_op` stay untiled. - Returns `dst_op`'s tile sizes (loop order), or None if not possible. + Args: + src_op: Annotated op the tile sizes come from. + src_sizes: `src_op`'s tile sizes, in its loop order. + shared: Tensor value both ops use. + dst_op: Op to derive tile sizes for. + Returns: + `dst_op`'s tile sizes in its loop order, or None if not possible. + """ + return propagate_through_values( + src_op, + src_sizes, + shared, + shared, + dst_op, + ) + + +def propagate_through_values( + src_op: ir.Operation | ir.OpView, + src_sizes: Sequence[int], + src_shared: ir.Value, + dst_shared: ir.Value, + dst_op: ir.Operation | ir.OpView, +) -> list[int] | None: + """Propagate tile sizes across possibly-aliased shared values. + + `src_shared` and `dst_shared` may be different SSA values that represent the + same logical tensor across wrapper ops. + + Args: + src_op: Annotated op the tile sizes come from. + src_sizes: `src_op`'s tile sizes, in its loop order. + src_shared: Shared tensor as seen by `src_op`. + dst_shared: Shared tensor as seen by `dst_op`. + dst_op: Op to derive tile sizes for. + Returns: + `dst_op`'s tile sizes in its loop order, or None if not possible. """ src = opview(src_op) dst = opview(dst_op) @@ -117,10 +219,12 @@ def propagate_through_value( return None # Tile size per dimension of the shared tensor, as induced by the source. - tensor_tiles = tiles_on_value(src, src_sizes, shared) - dst_map = _map_for_value(dst, shared, dst_maps) + tensor_tiles = tiles_on_value(src, src_sizes, src_shared) + dst_map = _map_for_value(dst, dst_shared, dst_maps) if tensor_tiles is None or dst_map is None: return None + if len(tensor_tiles) != len(dst_map.results): + return None if len(list(dst.results)) != 1: return None diff --git a/lighthouse/schedule/tile_and_fuse.py b/lighthouse/schedule/tile_and_fuse.py index 147000e2..763f8b19 100644 --- a/lighthouse/schedule/tile_and_fuse.py +++ b/lighthouse/schedule/tile_and_fuse.py @@ -28,6 +28,7 @@ def assign_and_propagate_tile_sizes( tile_size: int = 32, strategy: str = "cache", propagate: bool = True, + propagate_through_loops: bool = False, ) -> ir.Module: """ Assign tile sizes to anchor ops and propagate them to their neighbors. @@ -42,6 +43,8 @@ def assign_and_propagate_tile_sizes( tile_size: Tiling size hint. strategy: Tiling strategy. propagate: Whether to propagate the tile sizes to neighboring ops. + propagate_through_loops: Whether to propagate through loop-carried + values instead of treating loops as barriers (default: False). Returns: Schedule """ @@ -56,7 +59,9 @@ def assign_and_propagate_tile_sizes( strategy=strategy, ) if propagate: - transform_ext.propagate_tile_sizes(annotated) + transform_ext.propagate_tile_sizes( + annotated, propagate_through_loops=propagate_through_loops + ) transform.yield_() return sched @@ -65,6 +70,7 @@ def assign_elementwise_tile_sizes( tile_size: int = 32, strategy: str = "cache", propagate: bool = True, + propagate_through_loops: bool = False, ) -> ir.Module: """ Anchor tiling on elementwise ops only. @@ -76,6 +82,8 @@ def assign_elementwise_tile_sizes( tile_size: Tiling size hint. strategy: Tiling strategy. propagate: Whether to propagate the tile sizes to neighboring ops. + propagate_through_loops: Whether to propagate through loop-carried + values instead of treating loops as barriers (default: False). Returns: Schedule """ @@ -90,7 +98,9 @@ def assign_elementwise_tile_sizes( strategy=strategy, ) if propagate: - transform_ext.propagate_tile_sizes(annotated) + transform_ext.propagate_tile_sizes( + annotated, propagate_through_loops=propagate_through_loops + ) transform.yield_() return sched diff --git a/test/transform/test_propagate_tile_sizes.py b/test/transform/test_propagate_tile_sizes.py new file mode 100644 index 00000000..37829bb6 --- /dev/null +++ b/test/transform/test_propagate_tile_sizes.py @@ -0,0 +1,627 @@ +# RUN: %PYTHON %s | FileCheck %s + +from mlir import ir +from mlir.dialects import transform + +import lighthouse.dialects as lh_dialects +from lighthouse import transform as lh_transform +from lighthouse.dialects.transform.transform_ext import propagate_tile_sizes +from lighthouse.schedule.builders import schedule_boilerplate + + +def run(name: str, payload_str: str, build_schedule): + print(f"Test: {name}", flush=True) + with ir.Context(), ir.Location.unknown(): + lh_dialects.register_and_load() + payload = ir.Module.parse(payload_str) + sched = build_schedule() + sched.body.operations[0].apply(payload.operation) + print(payload) + + +def build_propagate_schedule(anchor_op_name: str, propagate_through_loops: bool): + with schedule_boilerplate() as (sched, named_seq): + anchors = lh_transform.match_op(named_seq.bodyTarget, anchor_op_name) + propagate_tile_sizes(anchors, propagate_through_loops=propagate_through_loops) + transform.yield_() + return sched + + +# Forward propagation: matmul -> elementwise consumer +FORWARD_SIMPLE = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + %out = tensor.empty() : tensor<128x128xf32> + %relu = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%mm : tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) { + ^bb0(%i: f32, %o: f32): + %r = arith.maxnumf %i, %cst : f32 + linalg.yield %r : f32 + } -> tensor<128x128xf32> + return %relu : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: forward_simple +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# Elementwise consumer picks up the M and N tile sizes from the matmul result. +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run( + "forward_simple", + FORWARD_SIMPLE, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Backward propagation: fill is tiled to match the matmul it initialises +BACKWARD_FILL = """ +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + return %mm : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: backward_fill +# The fill (prologue) must be annotated with [32, 32] to match the matmul output tile. +# CHECK: linalg.fill {transform_ext.tile_sizes = array} +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +run( + "backward_fill", + BACKWARD_FILL, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Forward + backward: both epilogue and prologue get tiled +FORWARD_AND_BACKWARD = """ +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>, + %bias: tensor<128x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + %out = tensor.empty() : tensor<128x128xf32> + %add = linalg.add + ins(%mm, %bias : tensor<128x128xf32>, tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) -> tensor<128x128xf32> + return %add : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: forward_and_backward +# CHECK: linalg.fill {transform_ext.tile_sizes = array} +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.add {transform_ext.tile_sizes = array} +# The add (epilogue) receives tile sizes propagated forward from the matmul. +run( + "forward_and_backward", + FORWARD_AND_BACKWARD, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Forward propagation stops at tensor.cast (negative test) +CAST_STOPS_FORWARD = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>) -> tensor { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + %casted = tensor.cast %mm : tensor<128x128xf32> to tensor + %out = tensor.empty() : tensor<128x128xf32> + %out_cast = tensor.cast %out : tensor<128x128xf32> to tensor + %relu = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%casted : tensor) + outs(%out_cast : tensor) { + ^bb0(%i: f32, %o: f32): + %r = arith.maxnumf %i, %cst : f32 + linalg.yield %r : f32 + } -> tensor + %result = tensor.cast %relu : tensor to tensor + return %result : tensor +} +""" + +# CHECK-LABEL: Test: cast_stops_forward +# tensor.cast is opaque; propagation stops there, the consumer stays unannotated. +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic {indexing_maps +# CHECK-NOT: transform_ext.tile_sizes +# CHECK: return +run( + "cast_stops_forward", + CAST_STOPS_FORWARD, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Forward propagation through an scf.for loop: annotated matmul inside the loop +# yields its result; the consumer outside is annotated via the for-result bridge. +FORWARD_THROUGH_FOR = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, + %arg2: tensor<128x128xf32>) -> tensor<128x128xf32> { + %c0 = arith.constant 0 : index + %c128 = arith.constant 128 : index + %c1 = arith.constant 1 : index + %0 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %arg2) -> (tensor<128x128xf32>) { + %col = tensor.extract_slice %arg0[0, %arg3] [128, 1] [1, 1] : tensor<128x128xf32> to tensor<128x1xf32> + %row = tensor.extract_slice %arg1[%arg3, 0] [1, 128] [1, 1] : tensor<128x128xf32> to tensor<1x128xf32> + %update = linalg.matmul {transform_ext.tile_sizes = array} + ins(%col, %row : tensor<128x1xf32>, tensor<1x128xf32>) + outs(%arg4 : tensor<128x128xf32>) -> tensor<128x128xf32> + scf.yield %update : tensor<128x128xf32> + } + %out = tensor.empty() : tensor<128x128xf32> + %relu = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%0 : tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) { + ^bb0(%in: f32, %o: f32): + %cst = arith.constant 0.0 : f32 + %r = arith.maxnumf %in, %cst : f32 + linalg.yield %r : f32 + } -> tensor<128x128xf32> + return %relu : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: forward_through_for +# The forward bridge traverses scf.for: matmul yields into the loop, and the +# loop result's consumer (relu) gets annotated with the translated tile sizes. +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run( + "forward_through_for", + FORWARD_THROUGH_FOR, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=True), +) + +# CHECK-LABEL: Test: loop_is_barrier_by_default +# Without opting in, scf.for stops propagation: the consumer stays unannotated. +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-NOT: transform_ext.tile_sizes +run( + "loop_is_barrier_by_default", + FORWARD_THROUGH_FOR, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Chained forward: two consumers in sequence both get annotated +CHAINED_FORWARD = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + %mid = tensor.empty() : tensor<128x128xf32> + %act1 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%mm : tensor<128x128xf32>) + outs(%mid : tensor<128x128xf32>) { + ^bb0(%i: f32, %o: f32): + %r = arith.maxnumf %i, %cst : f32 + linalg.yield %r : f32 + } -> tensor<128x128xf32> + %out = tensor.empty() : tensor<128x128xf32> + %act2 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%act1 : tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) { + ^bb0(%i: f32, %o: f32): + linalg.yield %i : f32 + } -> tensor<128x128xf32> + return %act2 : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: chained_forward +# Propagation cascades: mm -> act1 -> act2 all end up annotated. +# CHECK: linalg.fill {transform_ext.tile_sizes = array} +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +run( + "chained_forward", + CHAINED_FORWARD, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Barrier stops propagation: a contraction op between two elementwise ops +# acts as a fusion barrier; propagation does not cross it. +BARRIER_STOPS_PROPAGATION = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>, + %x: tensor<128x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + %out = tensor.empty() : tensor<128x128xf32> + %epilogue = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"], + transform_ext.tile_sizes = array} + ins(%mm : tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) { + ^bb0(%i: f32, %o: f32): + %r = arith.maxnumf %i, %cst : f32 + linalg.yield %r : f32 + } -> tensor<128x128xf32> + return %epilogue : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: barrier_stops_propagation +# Backward propagation from the epilogue reaches the matmul (a contraction / +# barrier), but must not annotate it, and must not continue to the fill either. +# CHECK: linalg.fill ins +# CHECK-NOT: transform_ext.tile_sizes +# CHECK: linalg.matmul ins +# CHECK-NOT: transform_ext.tile_sizes +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +run( + "barrier_stops_propagation", + BARRIER_STOPS_PROPAGATION, + lambda: build_propagate_schedule("linalg.generic", propagate_through_loops=False), +) + + +# Pre-annotated op acts as barrier: already-annotated downstream op with +# different tile sizes prevents re-annotation and marks a fusion boundary. +PREANNOTED_BARRIER = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%a: tensor<128x64xf32>, %b: tensor<64x128xf32>) -> tensor<128x128xf32> { + %cst = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %f = linalg.fill ins(%cst : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%a, %b : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%f : tensor<128x128xf32>) -> tensor<128x128xf32> + %out = tensor.empty() : tensor<128x128xf32> + %relu = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"], + transform_ext.tile_sizes = array} + ins(%mm : tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) { + ^bb0(%i: f32, %o: f32): + %r = arith.maxnumf %i, %cst : f32 + linalg.yield %r : f32 + } -> tensor<128x128xf32> + return %relu : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: preannotated_barrier +# The downstream generic already carries tile sizes that disagree with matmul's [32, 32]. +# Propagation must not overwrite it; the conflicting consumer is marked as a boundary. +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run( + "preannotated_barrier", + PREANNOTED_BARRIER, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Reduction dims are never tiled during propagation +REDUCTION_UNTILED = """ +#map3 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> +#reduce = affine_map<(d0, d1, d2) -> (d0, d1)> +func.func @main(%a: tensor<128x64x32xf32>) -> tensor<128x64xf32> { + %cst = arith.constant 0.0 : f32 + %src_empty = tensor.empty() : tensor<128x64x32xf32> + %src = linalg.generic {indexing_maps = [#map3, #map3], + iterator_types = ["parallel", "parallel", "parallel"], + transform_ext.tile_sizes = array} + ins(%a : tensor<128x64x32xf32>) + outs(%src_empty : tensor<128x64x32xf32>) { + ^bb0(%i: f32, %o: f32): + linalg.yield %i : f32 + } -> tensor<128x64x32xf32> + %red_empty = tensor.empty() : tensor<128x64xf32> + %red_init = linalg.fill ins(%cst : f32) outs(%red_empty : tensor<128x64xf32>) -> tensor<128x64xf32> + %reduce = linalg.generic { + indexing_maps = [#map3, #reduce], + iterator_types = ["parallel", "parallel", "reduction"]} + ins(%src : tensor<128x64x32xf32>) + outs(%red_init : tensor<128x64xf32>) { + ^bb0(%i: f32, %o: f32): + %sum = arith.addf %i, %o : f32 + linalg.yield %sum : f32 + } -> tensor<128x64xf32> + return %reduce : tensor<128x64xf32> +} +""" + +# CHECK-LABEL: Test: reduction_untiled +# Propagation from the elementwise source reaches the reduction. +# The reduction dim (d2) must stay untiled (0) when annotating the source. +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +# CHECK: linalg.fill {{{.*}}transform_ext.tile_sizes = array +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +run( + "reduction_untiled", + REDUCTION_UNTILED, + lambda: build_propagate_schedule("linalg.generic", propagate_through_loops=False), +) + + +# Backward propagation through scf.for: the fill initialising the iter arg +# is reached by bridging the iter arg back to its initial value on the for. +BACKWARD_THROUGH_FOR = """ +func.func @main(%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>) -> tensor<128x128xf32> { + %c0 = arith.constant 0 : index + %c128 = arith.constant 128 : index + %c1 = arith.constant 1 : index + %zero = arith.constant 0.0 : f32 + %e = tensor.empty() : tensor<128x128xf32> + %fill = linalg.fill ins(%zero : f32) outs(%e : tensor<128x128xf32>) -> tensor<128x128xf32> + %0 = scf.for %arg3 = %c0 to %c128 step %c1 iter_args(%arg4 = %fill) -> (tensor<128x128xf32>) { + %col = tensor.extract_slice %arg0[0, %arg3] [128, 1] [1, 1] : tensor<128x128xf32> to tensor<128x1xf32> + %row = tensor.extract_slice %arg1[%arg3, 0] [1, 128] [1, 1] : tensor<128x128xf32> to tensor<1x128xf32> + %update = linalg.matmul {transform_ext.tile_sizes = array} + ins(%col, %row : tensor<128x1xf32>, tensor<1x128xf32>) + outs(%arg4 : tensor<128x128xf32>) -> tensor<128x128xf32> + scf.yield %update : tensor<128x128xf32> + } + return %0 : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: backward_through_for +# CHECK: linalg.fill {{{.*}}transform_ext.tile_sizes = array +# CHECK: linalg.matmul {{{.*}}transform_ext.tile_sizes = array +run( + "backward_through_for", + BACKWARD_THROUGH_FOR, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=True), +) + +# CHECK-LABEL: Test: backward_loop_is_barrier_by_default +# Without opting in, the iter-arg bridge is off too: the fill stays unannotated. +# CHECK: linalg.fill ins +# CHECK-NOT: transform_ext.tile_sizes +# CHECK: linalg.matmul {{{.*}}transform_ext.tile_sizes = array +run( + "backward_loop_is_barrier_by_default", + BACKWARD_THROUGH_FOR, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=False), +) + + +# Nested scf.for loops: forward propagation cascades through two yields +# inner matmul -> inner scf.yield -> outer scf.yield -> consumer +FORWARD_NESTED_FORS = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, + %arg2: tensor<128x128xf32>) -> tensor<128x128xf32> { + %c0 = arith.constant 0 : index + %c128 = arith.constant 128 : index + %c1 = arith.constant 1 : index + %outer = scf.for %i = %c0 to %c128 step %c1 iter_args(%outer_carry = %arg2) -> (tensor<128x128xf32>) { + %col = tensor.extract_slice %arg0[0, %i] [128, 1] [1, 1] : tensor<128x128xf32> to tensor<128x1xf32> + %inner = scf.for %j = %c0 to %c128 step %c1 iter_args(%inner_carry = %outer_carry) -> (tensor<128x128xf32>) { + %row = tensor.extract_slice %arg1[%i, 0] [1, 128] [1, 1] : tensor<128x128xf32> to tensor<1x128xf32> + %update = linalg.matmul {transform_ext.tile_sizes = array} + ins(%col, %row : tensor<128x1xf32>, tensor<1x128xf32>) + outs(%inner_carry : tensor<128x128xf32>) -> tensor<128x128xf32> + scf.yield %update : tensor<128x128xf32> + } + scf.yield %inner : tensor<128x128xf32> + } + %out = tensor.empty() : tensor<128x128xf32> + %consumer = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%outer : tensor<128x128xf32>) + outs(%out : tensor<128x128xf32>) { + ^bb0(%in: f32, %o: f32): + linalg.yield %in : f32 + } -> tensor<128x128xf32> + return %consumer : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: forward_nested_fors +# Forward bridge cascades: mm -> inner scf.yield -> inner result -> outer scf.yield +# -> outer result -> consumer, annotating the consumer with translated tile sizes. +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run( + "forward_nested_fors", + FORWARD_NESTED_FORS, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=True), +) + + +# Loop carrying two values: only the chain through the second iter arg is annotated. +# Exercises the index mapping of both bridges (yield operand -> loop result, and +# iter arg -> its init operand), which a single-iter-arg loop cannot distinguish. +MULTI_ITER_ARGS = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>) + -> (tensor<128x128xf32>, tensor<128x128xf32>) { + %c0 = arith.constant 0 : index + %c128 = arith.constant 128 : index + %c1 = arith.constant 1 : index + %zero = arith.constant 0.0 : f32 + %e0 = tensor.empty() : tensor<128x128xf32> + %fill0 = linalg.fill ins(%zero : f32) outs(%e0 : tensor<128x128xf32>) -> tensor<128x128xf32> + %e1 = tensor.empty() : tensor<128x128xf32> + %fill1 = linalg.fill ins(%zero : f32) outs(%e1 : tensor<128x128xf32>) -> tensor<128x128xf32> + %r:2 = scf.for %i = %c0 to %c128 step %c1 iter_args(%a0 = %fill0, %a1 = %fill1) + -> (tensor<128x128xf32>, tensor<128x128xf32>) { + %col = tensor.extract_slice %arg0[0, %i] [128, 1] [1, 1] : tensor<128x128xf32> to tensor<128x1xf32> + %row = tensor.extract_slice %arg1[%i, 0] [1, 128] [1, 1] : tensor<128x128xf32> to tensor<1x128xf32> + %update = linalg.matmul {transform_ext.tile_sizes = array} + ins(%col, %row : tensor<128x1xf32>, tensor<1x128xf32>) + outs(%a1 : tensor<128x128xf32>) -> tensor<128x128xf32> + scf.yield %a0, %update : tensor<128x128xf32>, tensor<128x128xf32> + } + %o0 = tensor.empty() : tensor<128x128xf32> + %use0 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%r#0 : tensor<128x128xf32>) + outs(%o0 : tensor<128x128xf32>) { + ^bb0(%in: f32, %o: f32): + linalg.yield %in : f32 + } -> tensor<128x128xf32> + %o1 = tensor.empty() : tensor<128x128xf32> + %use1 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%r#1 : tensor<128x128xf32>) + outs(%o1 : tensor<128x128xf32>) { + ^bb0(%in: f32, %o: f32): + linalg.yield %in : f32 + } -> tensor<128x128xf32> + return %use0, %use1 : tensor<128x128xf32>, tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: multi_iter_args_index_mapping +# The first iter arg is only forwarded, so neither its init fill nor its consumer +# may be annotated; only the second chain (fill1 -> matmul -> %r#1) is. +# CHECK: linalg.fill ins +# CHECK-NOT: transform_ext.tile_sizes +# CHECK: linalg.fill {transform_ext.tile_sizes = array} +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-NOT: transform_ext.tile_sizes +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +run( + "multi_iter_args_index_mapping", + MULTI_ITER_ARGS, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=True), +) + + +# Only scf.for is bridged; scf.forall stays a barrier even when opted in. +FORALL_STAYS_BARRIER = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @main(%arg0: tensor<128x64xf32>, %arg1: tensor<64x128xf32>, + %arg2: tensor<128x128xf32>) -> tensor<128x128xf32> { + %0 = scf.forall (%i) in (4) shared_outs(%o = %arg2) -> (tensor<128x128xf32>) { + %sl = tensor.extract_slice %o[0, 0] [128, 128] [1, 1] : tensor<128x128xf32> to tensor<128x128xf32> + %mm = linalg.matmul {transform_ext.tile_sizes = array} + ins(%arg0, %arg1 : tensor<128x64xf32>, tensor<64x128xf32>) + outs(%sl : tensor<128x128xf32>) -> tensor<128x128xf32> + scf.forall.in_parallel { + tensor.parallel_insert_slice %mm into %o[0, 0] [128, 128] [1, 1] + : tensor<128x128xf32> into tensor<128x128xf32> + } + } + %e = tensor.empty() : tensor<128x128xf32> + %relu = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%0 : tensor<128x128xf32>) + outs(%e : tensor<128x128xf32>) { + ^bb0(%in: f32, %o: f32): + linalg.yield %in : f32 + } -> tensor<128x128xf32> + return %relu : tensor<128x128xf32> +} +""" + +# CHECK-LABEL: Test: forall_stays_barrier +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-NOT: transform_ext.tile_sizes +run( + "forall_stays_barrier", + FORALL_STAYS_BARRIER, + lambda: build_propagate_schedule("linalg.matmul", propagate_through_loops=True), +) + + +# An anchor with two propagatable producers and two consumers: every neighbor is +# claimed, not just the first one found in each direction. The second consumer +# reads the anchor transposed, so its tiles must be remapped, not copied. +FAN_IN_OUT = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +#tin = affine_map<(d0, d1) -> (d1, d0)> +func.func @main(%a: tensor<128x256xf32>, %b: tensor<128x256xf32>) + -> (tensor<128x256xf32>, tensor<256x128xf32>) { + %cst = arith.constant 0.0 : f32 + %ep0 = tensor.empty() : tensor<128x256xf32> + %p0 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%a : tensor<128x256xf32>) outs(%ep0 : tensor<128x256xf32>) { + ^bb0(%i: f32, %o: f32): + %e = math.exp %i : f32 + linalg.yield %e : f32 + } -> tensor<128x256xf32> + %ep1 = tensor.empty() : tensor<128x256xf32> + %p1 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%b : tensor<128x256xf32>) outs(%ep1 : tensor<128x256xf32>) { + ^bb0(%i: f32, %o: f32): + %e = math.exp %i : f32 + linalg.yield %e : f32 + } -> tensor<128x256xf32> + %eo = tensor.empty() : tensor<128x256xf32> + %anchor = linalg.add {transform_ext.tile_sizes = array} + ins(%p0, %p1 : tensor<128x256xf32>, tensor<128x256xf32>) + outs(%eo : tensor<128x256xf32>) -> tensor<128x256xf32> + %ec0 = tensor.empty() : tensor<128x256xf32> + %use0 = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%anchor : tensor<128x256xf32>) outs(%ec0 : tensor<128x256xf32>) { + ^bb0(%i: f32, %o: f32): + %r = arith.maxnumf %i, %cst : f32 + linalg.yield %r : f32 + } -> tensor<128x256xf32> + %ec1 = tensor.empty() : tensor<256x128xf32> + %use1 = linalg.generic {indexing_maps = [#tin, #map], + iterator_types = ["parallel", "parallel"]} + ins(%anchor : tensor<128x256xf32>) outs(%ec1 : tensor<256x128xf32>) { + ^bb0(%i: f32, %o: f32): + %n = arith.negf %i : f32 + linalg.yield %n : f32 + } -> tensor<256x128xf32> + return %use0, %use1 : tensor<128x256xf32>, tensor<256x128xf32> +} +""" + +# CHECK-LABEL: Test: fan_in_and_fan_out +# Both producers (backward) and both consumers (forward) get the anchor's tiling. +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +# CHECK: linalg.add {transform_ext.tile_sizes = array} +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +# The transposing consumer tiles the same tensor dims, so its loop-order sizes swap. +# CHECK: linalg.generic {{{.*}}transform_ext.tile_sizes = array +run( + "fan_in_and_fan_out", + FAN_IN_OUT, + lambda: build_propagate_schedule("linalg.add", propagate_through_loops=False), +) diff --git a/test/transform/test_tile_and_fuse.py b/test/transform/test_tile_and_fuse.py index c55bcd0e..a783177f 100644 --- a/test/transform/test_tile_and_fuse.py +++ b/test/transform/test_tile_and_fuse.py @@ -33,6 +33,12 @@ def assign_gemm(): return tf.assign_and_propagate_tile_sizes(tile_size=32, strategy="cache") +def assign_gemm_through_loops(): + return tf.assign_and_propagate_tile_sizes( + tile_size=32, propagate_through_loops=True + ) + + def assign_elementwise(): return tf.assign_elementwise_tile_sizes(tile_size=32, strategy="cache") @@ -298,6 +304,53 @@ def tile_and_fuse_for(): } """ +# A K-split matmul in an scf.for accumulation loop, with elementwise consumers +# outside the loop. Tile propagation must cross the loop boundary from the +# yielded matmul result to the loop result and then to the consumers. +K_LOOP_GEMM_OUTER_CHAIN = """ +#map = affine_map<(d0, d1) -> (d0, d1)> +module { + func.func @main(%a: tensor<32x64xf32>, %b: tensor<64x128xf32>) -> tensor<32x128xf32> { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %c64 = arith.constant 64 : index + %cst = arith.constant 0.0 : f32 + %init = tensor.empty() : tensor<32x128xf32> + %acc0 = linalg.fill ins(%cst : f32) outs(%init : tensor<32x128xf32>) -> tensor<32x128xf32> + %acc = scf.for %k = %c0 to %c64 step %c32 iter_args(%iter = %acc0) -> (tensor<32x128xf32>) { + %a_slice = tensor.extract_slice %a[0, %k] [32, 32] [1, 1] + : tensor<32x64xf32> to tensor<32x32xf32> + %b_slice = tensor.extract_slice %b[%k, 0] [32, 128] [1, 1] + : tensor<64x128xf32> to tensor<32x128xf32> + %mm = linalg.matmul ins(%a_slice, %b_slice : tensor<32x32xf32>, tensor<32x128xf32>) + outs(%iter : tensor<32x128xf32>) -> tensor<32x128xf32> + scf.yield %mm : tensor<32x128xf32> + } + %out0 = tensor.empty() : tensor<32x128xf32> + %relu = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%acc : tensor<32x128xf32>) + outs(%out0 : tensor<32x128xf32>) { + ^bb0(%in: f32, %out: f32): + %c = arith.cmpf ugt, %in, %cst : f32 + %s = arith.select %c, %in, %cst : f32 + linalg.yield %s : f32 + } -> tensor<32x128xf32> + %out1 = tensor.empty() : tensor<32x128xf32> + %exp = linalg.generic {indexing_maps = [#map, #map], + iterator_types = ["parallel", "parallel"]} + ins(%relu : tensor<32x128xf32>) + outs(%out1 : tensor<32x128xf32>) { + ^bb0(%in: f32, %out: f32): + %e = math.exp %in : f32 + linalg.yield %e : f32 + } -> tensor<32x128xf32> + return %exp : tensor<32x128xf32> + } +} +""" + # A high-dimensional contraction. HIGH_DIM_CONTRACT = """ #mapA = affine_map<(b, m0, m1, n, k) -> (b, m0, m1, k)> @@ -666,6 +719,17 @@ def tile_and_fuse_for(): run("dynamic_tile_and_fuse", DYN, assign_gemm, tile_and_fuse) +# Propagation crosses the scf.for loop boundary: a matmul inside a K loop drives +# tile annotations on elementwise consumers outside the loop. +# CHECK-LABEL: Test: k_loop_matmul_propagation +# CHECK: linalg.matmul {transform_ext.tile_sizes = array} +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +# CHECK: linalg.generic +# CHECK-SAME: transform_ext.tile_sizes = array +run("k_loop_matmul_propagation", K_LOOP_GEMM_OUTER_CHAIN, assign_gemm_through_loops) + + # High-dimensional contraction: only the innermost two parallel (M, N) output # dims are tiled with the tile size; the outer parallel dims (batch and the # outer M dim) get a unit tile and the reduction (K) dim is left untiled, just