[LLVMCPU] thread inner tile alignment hint for vector-level tiling (#24698) Introduces a helper function that will generate the inner tile alignment hint control function introduced in https://github.com/llvm/llvm-project/pull/204007. The control function is passed to the SCF tiling driver and is invoked per-op before tiling/fusion. It merely forwards the inner tile alignment hints set during tile size selection, added in the stack of https://github.com/iree-org/iree/pull/24806. Also utilizes this in the `TileAndFuseProducerConsumer` pass. Assisted-by: Claude Code --------- Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp index aef1767..d62fabf 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp +++ b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp
@@ -10,6 +10,7 @@ #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" #include "llvm/Support/Debug.h" #include "mlir/Analysis/TopologicalSortUtils.h" @@ -20,7 +21,9 @@ #include "mlir/Dialect/SCF/Utils/Utils.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/Dominance.h" +#include "mlir/Interfaces/TilingInterface.h" #include <cassert> @@ -28,6 +31,26 @@ namespace mlir::iree_compiler { +scf::InnerTileAlignmentFnTy +makeInnerTileAlignmentFn(IREE::CPU::TilingLevel tilingLevel) { + return [tilingLevel]( + TilingInterface op, ArrayRef<OpFoldResult>, + ArrayRef<Operation *>) -> SmallVector<mlir::InnerTileAlignment> { + // Tile-size selection already recorded this op's per-dimension inner-tile + // alignments per tiling level. We simply forward those along. + auto hint = + IREE::CPU::InnerTileAlignmentsAttr::getFromOp(op.getOperation()); + if (!hint) { + return {}; + } + auto alignments = hint.getAlignments().getAs<DenseI64ArrayAttr>( + IREE::CPU::getTilingLevelName(tilingLevel)); + return alignments + ? mlir::convertInnerTileAlignments(alignments.asArrayRef()) + : SmallVector<mlir::InnerTileAlignment>{}; + }; +} + void fuseProducersOfSlices(RewriterBase &rewriter, std::queue<Operation *> &worklist, scf::SCFTileAndFuseOptions &options, @@ -54,7 +77,9 @@ // values produced by operations that implement the `TilingInterface`. // Add these operations to the worklist. std::optional<scf::SCFFuseProducerOfSliceResult> fusedResult = - scf::tileAndFuseProducerOfSlice(rewriter, candidateSlice, loops); + scf::tileAndFuseProducerOfSlice( + rewriter, candidateSlice, loops, + options.tilingOptions.innerTileAlignmentFn); if (!fusedResult) { continue; } @@ -109,10 +134,11 @@ }; } // namespace -FailureOr<std::queue<Operation *>> -fuseConsumersIntoForall(RewriterBase &rewriter, ArrayRef<Operation *> tiledOps, - MutableArrayRef<LoopLikeOpInterface> loops, - std::function<bool(Operation *)> filterFn) { +FailureOr<std::queue<Operation *>> fuseConsumersIntoForall( + RewriterBase &rewriter, ArrayRef<Operation *> tiledOps, + MutableArrayRef<LoopLikeOpInterface> loops, + std::function<bool(Operation *)> filterFn, + const scf::InnerTileAlignmentFnTy &innerTileAlignmentFn) { // Collect the candidate slices which can be potential consumers that can be // fused. Keep them in a vector reverse-sorted by dominance: the candidate // dominating others comes last (so it can be cheaply popped from the vector). @@ -213,7 +239,8 @@ ConsumerFusionQueueEntry entry = candidates.pop_back_val(); FailureOr<scf::SCFFuseConsumerOfSliceResult> fusedResult = - mlir::scf::tileAndFuseConsumer(rewriter, entry.fusableUser, loops); + mlir::scf::tileAndFuseConsumer(rewriter, entry.fusableUser, loops, + innerTileAlignmentFn); if (failed(fusedResult)) { return failure(); }
diff --git a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h index bae549c..ebd18e2 100644 --- a/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h +++ b/compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h
@@ -7,6 +7,7 @@ #ifndef IREE_COMPILER_CODEGEN_COMMON_TILEANDFUSEUTILS_H_ #define IREE_COMPILER_CODEGEN_COMMON_TILEANDFUSEUTILS_H_ +#include "iree/compiler/Codegen/Dialect/CPU/IR/IREECPUTypes.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h" #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUEnums.h" #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h" @@ -31,15 +32,21 @@ void collectTiledAndFusedOps(Operation *rootOp, llvm::SmallDenseSet<Operation *> &result); +/// Returns an inner tile alignment control function for scalable tiling at +/// `tilingLevel`. It forwards the per-dimension `InnerTileAlignment` hints that +/// tile-size selection set on `linalg.pack`/`linalg.unpack` ops. +scf::InnerTileAlignmentFnTy +makeInnerTileAlignmentFn(IREE::CPU::TilingLevel tilingLevel); + /// Fuse all consumers of the given `tiledOps` into the surrounding `scf.forall` /// unless specified otherwise by `filterFn`. Returns a list of new /// `tensor.extract_slice` ops with new fusion opportunities. FailureOr<std::queue<Operation *>> fuseConsumersIntoForall( RewriterBase &rewriter, ArrayRef<Operation *> tiledOps, MutableArrayRef<LoopLikeOpInterface> loops, - std::function<bool(Operation *)> filterFn = [](Operation *) { - return true; - }); + std::function<bool(Operation *)> filterFn = + [](Operation *) { return true; }, + const scf::InnerTileAlignmentFnTy &innerTileAlignmentFn = nullptr); /// Apply a tile and fuse transformation to all payload ops and store both the /// tiled operation as well as the created tile loops.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp index 1295c8b..1ce9113 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/LLVMCPUTileAndFuseProducerConsumer.cpp
@@ -187,9 +187,16 @@ tileSizes.resize(numLoops, 0); tileScalableFlags.resize(numLoops, false); + // Lets a scalable pack/unpack's tiled outer dims fold to a static single + // inner tile, using the alignment hints tile-size selection recorded on the + // op. + scf::InnerTileAlignmentFnTy innerTileAlignmentFn = + makeInnerTileAlignmentFn(tilingLevel); + scf::SCFTilingOptions tilingOptions; setSCFTileSizes(tilingOptions, rootOp, std::move(tileSizes), std::move(tileScalableFlags)); + tilingOptions.setInnerTileAlignmentFn(innerTileAlignmentFn); // onlyFuseProducerInputOperands implies reduction tiling. if (!onlyFuseProducerInputOperands) { @@ -262,7 +269,8 @@ [&tiledAndFusedOps, &unfusableOps](Operation *op) { return tiledAndFusedOps.contains(op) && !unfusableOps.contains(op); - }); + }, + innerTileAlignmentFn); if (failed(newFusionOpportunities)) { LDBG() << "failed to fuse consumers, skip";
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir index 220f9fa..bce3301 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/tile_and_fuse_producer_consumer_anchoring_root_op.mlir
@@ -416,3 +416,176 @@ // CHECK: linalg.generic // CHECK: scf.forall.in_parallel // CHECK: linalg.pack + +// ----- + +// Scalable linalg.pack with an elementwise-add producer. +// The `Equal` hint folds the scalable inner tile's outer dim to a static 1. + +#config = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [[8], 1]> +#config1 = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]> +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @scalable_pack_with_producer(%arg0: tensor<384x512xf32>, %arg1: tensor<384x512xf32>, %arg2 : 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 + %generic = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} + ins(%arg0, %arg1 : tensor<384x512xf32>, tensor<384x512xf32>) + outs(%arg2 : 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 %generic 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<vector_common_parallel = [Equal, Unknown]>, + lowering_config = #config1} : tensor<384x512xf32> -> tensor<?x512x?x1xf32> + return %pack : tensor<?x512x?x1xf32> +} +// CHECK-LABEL: func.func @scalable_pack_with_producer +// CHECK: %[[VSCALE:.+]] = vector.vscale +// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index +// CHECK: scf.forall {{.*}} = (0, 0) to (384, 512) step (%[[C8VS]], 1) +// CHECK: %[[GEN:.+]] = linalg.generic +// CHECK: linalg.pack %[[GEN]] +// CHECK-SAME: inner_tiles = [%[[C8VS]], 1] +// The outermost dimension has to equal 1 to confirm the Equal hint is taken into account. +// CHECK-SAME: -> tensor<1x1x?x1xf32> +// CHECK: scf.forall.in_parallel + +// ----- + +// Scalable linalg.unpack fused as the producer of an elementwise-add consumer. +// The `Equal` hint folds the scalable inner tile's outer dim to a static 1. + +#config = #iree_cpu.lowering_config<vector_common_parallel = [7, [8]]> +#config1 = #iree_cpu.lowering_config<distribution = [84, 64], vector_common_parallel = [7, [8]]> +#map = affine_map<(d0, d1) -> (d0, d1)> +func.func @scalable_unpack_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<vector_common_parallel = [Unknown, Equal]>, + lowering_config = #config} : tensor<12x?x7x?xf32> -> tensor<80x320xf32> + %generic = 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 = #config1} { + ^bb0(%in: f32, %in_0: f32, %out: f32): + %2 = arith.addf %in, %in_0 : f32 + linalg.yield %2 : f32 + } -> tensor<80x320xf32> + return %generic : tensor<80x320xf32> +} +// CHECK-LABEL: func.func @scalable_unpack_with_consumer +// CHECK: %[[VSCALE:.+]] = vector.vscale +// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index +// CHECK: scf.forall {{.*}} = (0, 0) to (80, 320) step (7, %[[C8VS]]) +// CHECK: %[[UNPACK:.+]] = linalg.unpack +// CHECK-SAME: inner_tiles = [7, %[[C8VS]]] +// The outermost dimensions have to equal 1 to confirm the Equal hint is taken into account. +// CHECK-SAME: tensor<1x1x7x?xf32> -> tensor<?x?xf32> +// CHECK: linalg.generic +// CHECK-SAME: ins(%{{.*}}, %[[UNPACK]] +// CHECK: scf.forall.in_parallel + +// ----- + +// Scalable linalg.mmt4d (N0 = 8*vscale) with a linalg.unpack consumer fused into +// the mmt4d's tiling. The `Equal` hint folds the scalable inner tile's outer dim to a static 1. + +#config = #iree_cpu.lowering_config<distribution = [5, 8, 0, 0, 0, 0], vector_common_parallel = [1, 1, 0, 7, [8], 0]> +#config1 = #iree_cpu.lowering_config<vector_common_parallel = [7, [8]]> +func.func @scalable_mmt4d_with_unpack_consumer(%lhs: tensor<55x512x7x1xf32>, %rhs: tensor<?x512x?x1xf32>) -> tensor<385x?xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %c0 = arith.constant 0 : index + %c8 = arith.constant 8 : index + %vscale = vector.vscale + %c8_vscale = arith.muli %vscale, %c8 : index + %nouter = tensor.dim %rhs, %c0 : tensor<?x512x?x1xf32> + %acc = tensor.empty(%nouter, %c8_vscale) : tensor<55x?x7x?xf32> + %fill = linalg.fill ins(%cst : f32) outs(%acc : tensor<55x?x7x?xf32>) -> tensor<55x?x7x?xf32> + %mmt4d = linalg.mmt4d {lowering_config = #config} ins(%lhs, %rhs : tensor<55x512x7x1xf32>, tensor<?x512x?x1xf32>) outs(%fill : tensor<55x?x7x?xf32>) -> tensor<55x?x7x?xf32> + %n = arith.muli %nouter, %c8_vscale : index + %out = tensor.empty(%n) : tensor<385x?xf32> + %unpack = linalg.unpack %mmt4d outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [7, %c8_vscale] into %out {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Unknown, Equal]>, lowering_config = #config1} : tensor<55x?x7x?xf32> -> tensor<385x?xf32> + return %unpack : tensor<385x?xf32> +} +// CHECK-LABEL: func.func @scalable_mmt4d_with_unpack_consumer +// CHECK: %[[VSCALE:.+]] = vector.vscale +// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index +// CHECK: scf.forall {{.*}} step (1, 1, %[[C8VS]]) +// CHECK: %[[FILL:.+]] = linalg.fill +// CHECK: %[[MMT4D:.+]] = linalg.mmt4d +// CHECK-SAME: outs(%[[FILL]] +// CHECK: linalg.unpack %[[MMT4D]] +// CHECK-SAME: inner_tiles = [7, %[[C8VS]]] +// The outermost dimensions have to equal 1 to confirm the Equal hint is taken into account. +// CHECK-SAME: tensor<1x1x7x?xf32> -> tensor<7x?xf32> +// CHECK: scf.forall.in_parallel + +// ----- + +// Scalable linalg.mmt4d (M0 = N0 = 8*vscale) with fused generic, unpack and pack +// consumers. Each scalable inner tile dim is expected to fold to a static 1. + +#config = #iree_cpu.lowering_config<distribution = [5, 7, 0, 0, 0, 0], vector_common_parallel = [1, 1, 0, [8], [8], 0]> +#config1 = #iree_cpu.lowering_config<vector_common_parallel = [1, 1, [8], [8]]> +#config2 = #iree_cpu.lowering_config<vector_common_parallel = [[8], [8]]> +#config3 = #iree_cpu.lowering_config<vector_common_parallel = [[8], 1]> +#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> +func.func @scalable_mmt4d_generic_unpack_pack(%lhs: tensor<5x512x?x1xf32>, %rhs: tensor<7x512x?x1xf32>) -> tensor<?x?x?x1xf32> { + %cst = arith.constant 0.000000e+00 : f32 + %c5 = arith.constant 5 : index + %c7 = arith.constant 7 : index + %c8 = arith.constant 8 : index + %vscale = vector.vscale + %c8_vscale = arith.muli %vscale, %c8 : index + %acc = tensor.empty(%c8_vscale, %c8_vscale) : tensor<5x7x?x?xf32> + %fill = linalg.fill ins(%cst : f32) outs(%acc : tensor<5x7x?x?xf32>) -> tensor<5x7x?x?xf32> + %mmt4d = linalg.mmt4d {lowering_config = #config} + ins(%lhs, %rhs : tensor<5x512x?x1xf32>, tensor<7x512x?x1xf32>) + outs(%fill : tensor<5x7x?x?xf32>) -> tensor<5x7x?x?xf32> + %genout = tensor.empty(%c8_vscale, %c8_vscale) : tensor<5x7x?x?xf32> + %gen = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} + ins(%mmt4d : tensor<5x7x?x?xf32>) outs(%genout : tensor<5x7x?x?xf32>) + attrs = {lowering_config = #config1} { + ^bb0(%in: f32, %out: f32): + %v = arith.maximumf %in, %cst : f32 + linalg.yield %v : f32 + } -> tensor<5x7x?x?xf32> + %m = arith.muli %c5, %c8_vscale : index + %n = arith.muli %c7, %c8_vscale : index + %uout = tensor.empty(%m, %n) : tensor<?x?xf32> + %unpack = linalg.unpack %gen outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] + inner_tiles = [%c8_vscale, %c8_vscale] into %uout + {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Equal, Equal]>, + lowering_config = #config2} : tensor<5x7x?x?xf32> -> tensor<?x?xf32> + %m_outer = affine.apply affine_map<()[s0, s1] -> (s0 ceildiv s1)>()[%m, %c8_vscale] + %pout = tensor.empty(%m_outer, %n, %c8_vscale) : tensor<?x?x?x1xf32> + %pack = linalg.pack %unpack padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] + inner_tiles = [%c8_vscale, 1] into %pout + {inner_tile_alignments = #iree_cpu.inner_tile_alignments<vector_common_parallel = [Equal, Multiple]>, + lowering_config = #config3} : tensor<?x?xf32> -> tensor<?x?x?x1xf32> + return %pack : tensor<?x?x?x1xf32> +} +// CHECK-LABEL: func.func @scalable_mmt4d_generic_unpack_pack +// CHECK: %[[VSCALE:.+]] = vector.vscale +// CHECK: %[[C8VS:.+]] = arith.muli %[[VSCALE]], %{{.+}} : index +// CHECK: scf.forall {{.*}} step (1, 1, %[[C8VS]], %[[C8VS]]) +// CHECK: %[[MMT4D:.+]] = linalg.mmt4d +// CHECK: %[[GEN:.+]] = linalg.generic +// CHECK-SAME: ins(%[[MMT4D]] +// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[GEN]] +// CHECK-SAME: inner_tiles = [%[[C8VS]], %[[C8VS]]] +// CHECK-SAME: tensor<1x1x?x?xf32> -> tensor<?x?xf32> +// CHECK: linalg.pack %[[UNPACK]] +// CHECK-SAME: inner_tiles = [%[[C8VS]], 1] +// The outermost dimension has to equal 1 to confirm the Equal hint is taken into account. +// CHECK-SAME: tensor<?x?xf32> -> tensor<1x?x?x1xf32> +// CHECK: scf.forall.in_parallel