[Codegen] Forward distribution inner-tile alignment hints in forall tiling

Threads the distribution-level inner-tile alignment hints (set on scalable
`linalg.pack`/`linalg.unpack` ops during lowering-strategy selection) into
`TileAndDistributeToWorkgroupsUsingForall`, via
`makeInnerTileAlignmentFn(DistributionTiles)` on both the tiling options and the
consumer-fusion step. With a `Multiple` hint the workgroup tiling folds a
scalable/dynamic inner tile to a clean `tile ceildiv innerTile` division instead
of leaving the pack/unpack untiled.

Adds lit coverage: two scalable pack/unpack cases (hint forwarded and preserved
through tiling) and a dynamic-inner-tile case where the hint is load-bearing —
the pack is only tiled into the forall because the hint asserts the multiple
relationship the driver cannot infer from the IR.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Ege Beysel <beyselege@gmail.com>

Author: egebeysel

compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp

@@ -267,10 +267,18 @@ void TileAndDistributeToWorkgroupsUsingForallOpPass::runOnOperation() {
           context, funcOp.getLoc(), deviceMappingAttribute))) {
     return signalPassFailure();
   }
+  // Forward the distribution-level inner-tile alignment hints recorded on
+  // tiled/fused `linalg.pack`/`linalg.unpack` ops, so the workgroup tiling can
+  // fold their scalable inner tiles to a static shape where a hint guarantees
+  // the distribution tile is a whole multiple of the inner tile.
+  scf::InnerTileAlignmentFnTy innerTileAlignmentFn =
+      makeInnerTileAlignmentFn(IREE::CPU::TilingLevel::DistributionTiles);
+
   scf::SCFTilingOptions tilingOptions;
   tilingOptions.setTileSizes(tilingInfo->tileSizes);
   tilingOptions.setInterchange(tilingInfo->interchange);
   tilingOptions.setMapping(deviceMappingAttribute);
+  tilingOptions.setInnerTileAlignmentFn(innerTileAlignmentFn);
 
   IREE::Codegen::WorkgroupReorderingAttrInterface workgroupReorderingStrategy =
       getLoweringConfig(tilingInfo->tilableOp).getWorkgroupReorderingStrategy();
@@ -370,9 +378,11 @@ void TileAndDistributeToWorkgroupsUsingForallOpPass::runOnOperation() {
     FailureOr<std::queue<Operation *>> newFusionOpportunities =
         fuseConsumersIntoForall(
             rewriter, tileAndFuseResult->tiledAndFusedOps.getArrayRef(),
-            tilingLoops, [&tiledAndFusedOps](Operation *op) {
+            tilingLoops,
+            [&tiledAndFusedOps](Operation *op) {
               return tiledAndFusedOps.contains(op);
-            });
+            },
+            innerTileAlignmentFn);
     if (failed(newFusionOpportunities)) {
       // Continue the work if the failure is allowed.
       if (!verifyComputeOpsAfterDistribution(funcOp)) {

compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_workgroups_using_forall.mlir

@@ -1498,3 +1498,112 @@ func.func @arg_compare_fold_broadcast(
 //   CHECK-NOT:     linalg.broadcast
 //       CHECK:     scf.forall {{.*}} shared_outs(%{{.*}} = %[[INIT_F16]], %{{.*}} = %[[INIT_I32]])
 //       CHECK:       iree_linalg_ext.arg_compare
+
+// -----
+
+// Distribution of a scalable linalg.pack (inner tiles [8*vscale, 1]) with an
+// elementwise-add producer fused into the workgroup forall. The pack carries a
+// `distribution = [Multiple, Unknown]` inner-tile alignment hint recorded during
+// lowering-strategy selection, asserting that the distribution tile (64) on the
+// packed dim is a whole multiple of the scalable inner tile. The forall tiling
+// forwards that hint, so the fused pack tiles cleanly: its outer packed dim is
+// `64 ceildiv (8 * vscale)`.
+#config = #iree_codegen.lowering_config<tile_sizes = [[64, 64]]>
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @scalable_pack_distribute_with_producer(%arg0: tensor<384x512xf32>, %arg1: tensor<384x512xf32>) -> tensor<?x512x?x1xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c8 = arith.constant 8 : index
+  %vscale = vector.vscale
+  %c8_vscale = arith.muli %vscale, %c8 : index
+  %0 = tensor.empty() : tensor<384x512xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<384x512xf32>, tensor<384x512xf32>) outs(%0 : tensor<384x512xf32>) attrs = {lowering_config = #config} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %3 = arith.addf %in, %in_0 : f32
+    linalg.yield %3 : f32
+  } -> tensor<384x512xf32>
+  %mouter = affine.apply affine_map<()[s0] -> (384 ceildiv s0)>()[%c8_vscale]
+  %2 = tensor.empty(%mouter, %c8_vscale) : tensor<?x512x?x1xf32>
+  %pack = linalg.pack %1 padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, 1] into %2 {inner_tile_alignments = #iree_cpu.inner_tile_alignments<distribution = [Multiple, Unknown]>} : tensor<384x512xf32> -> tensor<?x512x?x1xf32>
+  return %pack : tensor<?x512x?x1xf32>
+}
+// CHECK-LABEL: func.func @scalable_pack_distribute_with_producer
+// CHECK:         %[[VS:.+]] = vector.vscale
+// CHECK:         %[[C8VS:.+]] = arith.muli %[[VS]], %{{.+}}
+// CHECK:         scf.forall {{.*}} = (0, 0) to (384, 512) step (64, 64)
+// CHECK:           linalg.generic
+// CHECK:             arith.addf
+// CHECK:           %[[OUTER:.+]] = affine.apply affine_map<()[s0] -> (64 ceildiv s0)>()[%[[C8VS]]]
+// CHECK:           tensor.extract_slice %{{.+}}[%{{.+}}, %{{.+}}, 0, 0] [%[[OUTER]], 64, %[[C8VS]], 1]
+// CHECK:           linalg.pack
+// CHECK-SAME:        inner_tiles = [%[[C8VS]], 1]
+// CHECK-SAME:        inner_tile_alignments = #iree_cpu.inner_tile_alignments<distribution = [Multiple, Unknown]>
+// CHECK:         mapping = [#iree_codegen.workgroup_mapping<y>, #iree_codegen.workgroup_mapping<x>]
+
+// -----
+
+// Distribution of a scalable linalg.unpack (inner tiles [7, 8*vscale]) fused as
+// the producer of an elementwise-add consumer. The unpack's
+// `distribution = [Unknown, Multiple]` hint asserts that the distribution tile
+// (64) on the unpacked N dim is a whole multiple of the scalable inner tile, so
+// the fused unpack slices its packed source with `64 ceildiv (8 * vscale)`.
+#config = #iree_codegen.lowering_config<tile_sizes = [[84, 64]]>
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @scalable_unpack_distribute_with_consumer(%arg0: tensor<12x?x7x?xf32>, %arg1: tensor<80x320xf32>) -> tensor<80x320xf32> {
+  %c8 = arith.constant 8 : index
+  %vscale = vector.vscale
+  %c8_vscale = arith.muli %vscale, %c8 : index
+  %0 = tensor.empty() : tensor<80x320xf32>
+  %unpack = linalg.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [7, %c8_vscale] into %0 {inner_tile_alignments = #iree_cpu.inner_tile_alignments<distribution = [Unknown, Multiple]>} : tensor<12x?x7x?xf32> -> tensor<80x320xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg1, %unpack : tensor<80x320xf32>, tensor<80x320xf32>) outs(%0 : tensor<80x320xf32>) attrs = {lowering_config = #config} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %2 = arith.addf %in, %in_0 : f32
+    linalg.yield %2 : f32
+  } -> tensor<80x320xf32>
+  return %1 : tensor<80x320xf32>
+}
+// CHECK-LABEL: func.func @scalable_unpack_distribute_with_consumer
+// CHECK:         %[[VS:.+]] = vector.vscale
+// CHECK:         %[[C8VS:.+]] = arith.muli %[[VS]], %{{.+}}
+// CHECK:         scf.forall {{.*}} = (0, 0) to (80, 320) step (84, 64)
+// CHECK:           %[[NOUTER:.+]] = affine.apply affine_map<()[s0] -> (64 ceildiv s0)>()[%[[C8VS]]]
+// CHECK:           tensor.extract_slice %{{.+}}[0, %{{.+}}, 0, 0] [12, %[[NOUTER]], 7, %[[C8VS]]]
+// CHECK:           linalg.unpack
+// CHECK-SAME:        inner_tiles = [7, %[[C8VS]]]
+// CHECK-SAME:        inner_tile_alignments = #iree_cpu.inner_tile_alignments<distribution = [Unknown, Multiple]>
+// CHECK:           linalg.generic
+// CHECK:             arith.addf
+// CHECK:         mapping = [#iree_codegen.workgroup_mapping<y>, #iree_codegen.workgroup_mapping<x>]
+
+// -----
+
+// A dynamic (non-vscale) inner tile: the tiling driver cannot infer the multiple
+// relationship from the IR, so the `distribution = [Multiple, Unknown]` hint is
+// what lets the pack be tiled and fused into the workgroup forall (its outer
+// packed dim becomes `64 ceildiv %d`). Without the hint the pack stays untiled.
+#config = #iree_codegen.lowering_config<tile_sizes = [[64, 64]]>
+#map = affine_map<(d0, d1) -> (d0, d1)>
+func.func @dyn_pack_distribute_multiple_hint(%arg0: tensor<384x512xf32>, %arg1: tensor<384x512xf32>, %d: index) -> tensor<?x512x?x1xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<384x512xf32>
+  %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<384x512xf32>, tensor<384x512xf32>) outs(%0 : tensor<384x512xf32>) attrs = {lowering_config = #config} {
+  ^bb0(%in: f32, %in_0: f32, %out: f32):
+    %3 = arith.addf %in, %in_0 : f32
+    linalg.yield %3 : f32
+  } -> tensor<384x512xf32>
+  %mouter = affine.apply affine_map<()[s0] -> (384 ceildiv s0)>()[%d]
+  %2 = tensor.empty(%mouter, %d) : tensor<?x512x?x1xf32>
+  %pack = linalg.pack %1 padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [%d, 1] into %2 {inner_tile_alignments = #iree_cpu.inner_tile_alignments<distribution = [Multiple, Unknown]>} : tensor<384x512xf32> -> tensor<?x512x?x1xf32>
+  return %pack : tensor<?x512x?x1xf32>
+}
+// CHECK-LABEL: func.func @dyn_pack_distribute_multiple_hint
+// CHECK-SAME:      %[[D:[A-Za-z0-9_]+]]: index
+// CHECK:         scf.forall {{.*}} = (0, 0) to (384, 512) step (64, 64)
+// CHECK:           linalg.generic
+// CHECK:             arith.addf
+// CHECK:           %[[OUTER:.+]] = affine.apply affine_map<()[s0] -> (64 ceildiv s0)>()[%[[D]]]
+// CHECK:           tensor.extract_slice %{{.+}}[%{{.+}}, %{{.+}}, 0, 0] [%[[OUTER]], 64, %[[D]], 1]
+// CHECK:           linalg.pack
+// CHECK-SAME:        inner_tiles = [%[[D]], 1]
+// CHECK-SAME:        inner_tile_alignments = #iree_cpu.inner_tile_alignments<distribution = [Multiple, Unknown]>
+// CHECK:           scf.forall.in_parallel
+// CHECK:         mapping = [#iree_codegen.workgroup_mapping<y>, #iree_codegen.workgroup_mapping<x>]