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