Fully retire CanonicalizedSequenceOp (#12467)
This revision fully retires the CanonicalizedSequenceOp and makes the
need for canonicalization / tiling_canonicalization / licm / singleIterationLoopPromotion
explicit in the transform IR.
This removes automatic blanket pattern applications which is expected to be beneficial for:
1. reducing compile time
2. reducing surprise (explicit being better than implicit)
3. help surface issues similar to #12444
In the process, plumb the isolatedFromAbove through the API to make the
requirement explicit until we are able to remove it in some future (i.e. with better targeted
canonicalization and greedy pattern rewriter that does not require an isolated from above).
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir
index e827c0b..4913b20 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/set_transform_strategy.mlir
@@ -29,15 +29,18 @@
}
// CHECK-LABEL: func.func @group_reduction
-// CHECK: transform.structured.canonicalized_sequence failures(propagate)
+// CHECK: transform.sequence failures(propagate)
// CHECK: transform.iree.match_callback failures(propagate) "reduction"(%{{.+}})
// CHECK: transform.iree.take_first
// CHECK: transform.iree.tile_to_forall_and_workgroup_count_region {{.*}} tile_sizes [1](mapping = [#gpu.block<x>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK-COUNT-2: transform.structured.fuse_into_containing_op
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.iree.take_first
// CHECK: tile_reduction_using_forall {{.*}} by num_threads = [0, 64], tile_sizes = [0, 1], mapping = [#gpu.thread<x>]
// CHECK: transform.structured.fuse_into_containing_op
// CHECK: transform.structured.tile_to_forall_op %{{.*}} tile_sizes [1](mapping = [#gpu.thread<y>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: cast %{{.*}} : !pdl.operation to !transform.op<"scf.forall">
// CHECK: transform.iree.share_forall_operands %{{.*}} share_operands = [0] : (!transform.op<"scf.forall">) -> !transform.op<"scf.forall">
// CHECK: transform.structured.match ops{["func.func"]} in %arg0
@@ -91,7 +94,7 @@
// Checking only the tile sizes.
// CHECK-LABEL: func.func @group_reduction_128
-// CHECK: transform.structured.canonicalized_sequence failures(propagate)
+// CHECK: transform.sequence failures(propagate)
// CHECK: transform.structured.tile_reduction_using_forall %{{.*}} by num_threads = [0, 32], tile_sizes = [0, 4], mapping = [#gpu.thread<x>]
// CHECK: transform.iree.map_nested_forall_to_gpu_threads %{{.*}} {workgroup_size = [32, 1, 1]}
// CHECK: transform.iree.vector.to_warp_execute_on_lane_0 %{{.*}} {warp_size = 32 : i64}
@@ -133,7 +136,7 @@
// Checking only the tile sizes.
// CHECK-LABEL: func.func @group_reduction_D
-// CHECK: transform.structured.canonicalized_sequence failures(propagate)
+// CHECK: transform.sequence failures(propagate)
// CHECK: transform.structured.tile_reduction_using_forall %{{.*}} by num_threads = [0, 256], tile_sizes = [0, 1], mapping = [#gpu.thread<x>]
// CHECK: transform.iree.map_nested_forall_to_gpu_threads %{{.*}} {workgroup_size = [256, 1, 1]}
// CHECK: transform.iree.vector.to_warp_execute_on_lane_0 %{{.*}} {warp_size = 256 : i64}
@@ -173,12 +176,15 @@
// Checking only the tile sizes.
// CHECK-LABEL: func.func @group_reduction_34
-// CHECK: transform.structured.canonicalized_sequence failures(propagate)
+// CHECK: transform.sequence failures(propagate)
// CHECK: transform.iree.tile_to_forall_and_workgroup_count_region %{{.*}} num_threads [] tile_sizes [64](mapping = [#gpu.block<x>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.structured.tile_to_forall_op %{{.*}} num_threads [64] tile_sizes [](mapping = [#gpu.thread<x>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK-COUNT-4: transform.structured.scalarize %{{.*}}
// CHECK: transform.structured.split %{{.*}} after 32 {dimension = 1 : i64}
// CHECK: transform.structured.tile %{{.*}}[0, 4]
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.iree.map_nested_forall_to_gpu_threads %{{.*}} {workgroup_size = [64, 1, 1]}
// CHECK-NOT: transform.iree.vector.to_warp_execute_on_lane_0
@@ -226,12 +232,17 @@
}
// CHECK-LABEL: func.func @group_reduction_12345
-// CHECK: transform.structured.canonicalized_sequence failures(propagate)
+// CHECK: transform.sequence failures(propagate)
// CHECK: transform.iree.tile_to_forall_and_workgroup_count_region %{{.*}} num_threads [] tile_sizes [1](mapping = [#gpu.block<x>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.structured.tile_to_forall_op %{{.*}} num_threads [] tile_sizes [1](mapping = [#gpu.thread<y>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.structured.split %{{.*}} after 8192 {dimension = 1 : i64}
// CHECK: transform.structured.tile %{{.*}}[0, 8192]
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.structured.tile_to_forall_op %{{.*}} num_threads [0, 1024] tile_sizes [](mapping = [#gpu.thread<x>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.structured.tile_to_forall_op %{{.*}} num_threads [0, 1024] tile_sizes [](mapping = [#gpu.thread<x>])
+// CHECK: transform.iree.apply_patterns %{{.*}} {canonicalization, cse, licm, tiling_canonicalization}
// CHECK: transform.iree.map_nested_forall_to_gpu_threads %{{.*}} {workgroup_size = [1024, 1, 1]}
// CHECK: transform.iree.vector.to_warp_execute_on_lane_0{{.*}}{warp_size = 1024 : i64}
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/Common.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/Common.cpp
index 07bc9ba..e5bf994 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/Common.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/Common.cpp
@@ -31,6 +31,7 @@
using iree_compiler::cpu::CPUModel;
using iree_compiler::cpu::ReductionConfig;
using iree_compiler::cpu::ReductionStrategy;
+using iree_compiler::IREE::transform_dialect::ApplyPatternsOpPatterns;
using iree_compiler::IREE::transform_dialect::ForallToWorkgroupOp;
using transform::LowerVectorsOp;
using transform::MatchOp;
@@ -53,9 +54,24 @@
std::pair<Value, Value> mlir::iree_compiler::cpu::buildCommonTrailingStrategy(
ImplicitLocOpBuilder &b, Value variantH,
const vector::LowerVectorsOptions &lowerVectorsOpts) {
- // Step N-2. Bufferize and drop HAL descriptor from memref ops.
Value funcH = b.create<MatchOp>(variantH, func::FuncOp::getOperationName());
+
+ // Step N-5. Fold tensor.empty to avoid large allocations.
+ ApplyPatternsOpPatterns configuration;
+ configuration.foldTensorEmptyExtract = true;
+
+ // Step N-4. Perform a pass of canonicalization + enabling after tiling.
+ funcH = mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, funcH);
funcH = iree_compiler::buildVectorize(b, funcH);
+
+ // Step N-3. Perform a pass of canonicalization + enabling after vectorization
+ // as well as hoisting subset operations such as vector.transfer_read/write.
+ funcH = mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, funcH);
+ funcH = iree_compiler::buildHoisting(b, funcH);
+
+ // Step N-2. Bufferize and drop HAL descriptor from memref ops.
variantH = iree_compiler::buildBufferize(b, variantH);
// Step N-1. Post-bufferization mapping to blocks only.
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/ReductionStrategy.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/ReductionStrategy.cpp
index 3bd6764..6d11188 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/ReductionStrategy.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/CPU/ReductionStrategy.cpp
@@ -74,7 +74,7 @@
if (rank == 0) continue;
SmallVector<int64_t> tileSizes(rank - 1, 0);
tileSizes.push_back(strategy.getVectorSize());
- buildTileFuseToScfFor(b, val, {},
+ buildTileFuseToScfFor(b, variantH, val, {},
getAsOpFoldResult(b.getI64ArrayAttr(tileSizes)));
}
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
index 1bc37b5..3417d9c 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
@@ -34,6 +34,7 @@
using iree_compiler::IREE::transform_dialect::
TileToForallAndWorkgroupCountRegionOp;
using transform::FuseIntoContainingOp;
+using transform::HoistRedundantTensorSubsetsOp;
using transform::MatchOp;
using transform::MergeHandlesOp;
using transform::PrintOp;
@@ -94,9 +95,10 @@
auto topLevelTransformModule = b.create<ModuleOp>(loc);
Region &topLevelTransformRegion = topLevelTransformModule.getBodyRegion();
b.setInsertionPointToStart(&topLevelTransformRegion.front());
- auto sequence = b.create<::transform_ext::CanonicalizedSequenceOp>(
- loc, transform::FailurePropagationMode::Propagate,
- [&](OpBuilder &b, Location loc, Value variantH) {
+ auto pdlOperationType = pdl::OperationType::get(b.getContext());
+ auto sequence = b.create<transform::SequenceOp>(
+ loc, TypeRange{}, transform::FailurePropagationMode::Propagate,
+ pdlOperationType, [&](OpBuilder &b, Location loc, Value variantH) {
ImplicitLocOpBuilder ib(loc, b);
buildStrategy(ib, variantH);
b.create<transform::YieldOp>(loc);
@@ -119,6 +121,21 @@
for (auto h : handles) b.create<PrintOp>(h);
}
+/// Create an ApplyPatternsOp that performs a set of key canonicalizations and
+/// so-called enabling transformations to normalize the IR.
+/// Take an existing configuration by copy (cheap object) that will be augmented
+/// locally to additionally perform:
+/// canonicalization, tiling_canonicalization, licm and cse (in this order).
+Value mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ ImplicitLocOpBuilder &b, ApplyPatternsOpPatterns configuration,
+ Value variantH) {
+ configuration.canonicalization = true;
+ configuration.cse = true;
+ configuration.licm = true;
+ configuration.tilingCanonicalization = true;
+ return b.create<ApplyPatternsOp>(variantH, configuration);
+}
+
/// Dynamically selects the first non-empty handle; i.e. if (h1, h2) is:
/// - (non-empty, non-empty), returns (h1, h2)
/// - (empty, non-empty), returns (h2, empty)
@@ -136,7 +153,8 @@
}
mlir::iree_compiler::TileToScfForAndFuseResult
-mlir::iree_compiler::buildTileFuseToScfFor(ImplicitLocOpBuilder &b, Value rootH,
+mlir::iree_compiler::buildTileFuseToScfFor(ImplicitLocOpBuilder &b,
+ Value isolatedParentOpH, Value rootH,
ValueRange opsHToFuse,
ArrayRef<OpFoldResult> tileSizes) {
assert(opsHToFuse.empty() && "No fusion supported yet");
@@ -144,6 +162,12 @@
auto tiletoScfForOp = b.create<transform::TileOp>(rootH, tileSizes);
result.forLoops = tiletoScfForOp.getLoops();
result.tiledOpH = tiletoScfForOp.getTiledLinalgOp();
+
+ // Perform a pass of canonicalization + enabling after tiling.
+ ApplyPatternsOpPatterns configuration;
+ isolatedParentOpH =
+ mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, isolatedParentOpH);
return result;
}
@@ -169,7 +193,8 @@
// TODO: apply forwarding pattern.
template <typename TilingTransformOp, typename TileOrNumThreadSpec>
static iree_compiler::TileToForallAndFuseAndDistributeResult
-buildTileAndFuseAndDistributeImpl(ImplicitLocOpBuilder &b, Value rootH,
+buildTileAndFuseAndDistributeImpl(ImplicitLocOpBuilder &b,
+ Value isolatedParentOpH, Value rootH,
ValueRange opsHToFuse,
ArrayRef<OpFoldResult> tileSizesOrNumThreads,
ArrayAttr threadDimMapping) {
@@ -179,6 +204,12 @@
result.forallH = tileToForeachOp.getForallOp();
result.tiledOpH = tileToForeachOp.getTiledOp();
+ // Perform a pass of canonicalization + enabling after tiling.
+ ApplyPatternsOpPatterns configuration;
+ isolatedParentOpH =
+ mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, isolatedParentOpH);
+
// Batch fusion if requested.
if (opsHToFuse.size() > 1) {
Value mergedOpsH =
@@ -196,27 +227,29 @@
// sigh.
template <typename TilingTransformOp>
static iree_compiler::TileToForallAndFuseAndDistributeResult
-buildTileFuseDistWithTileSizes(ImplicitLocOpBuilder &b, Value rootH,
- ValueRange opsHToFuse,
+buildTileFuseDistWithTileSizes(ImplicitLocOpBuilder &b, Value isolatedParentOpH,
+ Value rootH, ValueRange opsHToFuse,
ArrayRef<OpFoldResult> tileSizes,
ArrayAttr threadDimMapping) {
return buildTileAndFuseAndDistributeImpl<TilingTransformOp,
transform::TileSizesSpec>(
- b, rootH, opsHToFuse, tileSizes, threadDimMapping);
+ b, isolatedParentOpH, rootH, opsHToFuse, tileSizes, threadDimMapping);
}
iree_compiler::TileToForallAndFuseAndDistributeResult
mlir::iree_compiler::buildTileFuseDistToForallWithTileSizes(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes, ArrayAttr threadDimMapping) {
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes,
+ ArrayAttr threadDimMapping) {
return buildTileFuseDistWithTileSizes<TileToForallOp>(
- b, rootH, opsHToFuse, tileSizes, threadDimMapping);
+ b, isolatedParentOpH, rootH, opsHToFuse, tileSizes, threadDimMapping);
}
iree_compiler::TileToForallAndFuseAndDistributeResult
mlir::iree_compiler::buildTileFuseDistToForallAndWorkgroupCountWithTileSizes(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes, ArrayAttr threadDimMapping) {
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes,
+ ArrayAttr threadDimMapping) {
return buildTileFuseDistWithTileSizes<TileToForallAndWorkgroupCountRegionOp>(
- b, rootH, opsHToFuse, tileSizes, threadDimMapping);
+ b, isolatedParentOpH, rootH, opsHToFuse, tileSizes, threadDimMapping);
}
/// Call buildTileAndFuseAndDistributeImpl with ArrayRef<int64_t> numThreads.
@@ -224,27 +257,30 @@
// sigh.
template <typename TilingTransformOp>
static iree_compiler::TileToForallAndFuseAndDistributeResult
-buildTileFuseDistWithNumThreads(ImplicitLocOpBuilder &b, Value rootH,
+buildTileFuseDistWithNumThreads(ImplicitLocOpBuilder &b,
+ Value isolatedParentOpH, Value rootH,
ValueRange opsHToFuse,
ArrayRef<OpFoldResult> numThreads,
ArrayAttr threadDimMapping) {
return buildTileAndFuseAndDistributeImpl<TilingTransformOp,
transform::NumThreadsSpec>(
- b, rootH, opsHToFuse, numThreads, threadDimMapping);
+ b, isolatedParentOpH, rootH, opsHToFuse, numThreads, threadDimMapping);
}
iree_compiler::TileToForallAndFuseAndDistributeResult
mlir::iree_compiler::buildTileFuseDistToForallWithNumThreads(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes, ArrayAttr threadDimMapping) {
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes,
+ ArrayAttr threadDimMapping) {
return buildTileFuseDistWithNumThreads<TileToForallOp>(
- b, rootH, opsHToFuse, tileSizes, threadDimMapping);
+ b, isolatedParentOpH, rootH, opsHToFuse, tileSizes, threadDimMapping);
}
iree_compiler::TileToForallAndFuseAndDistributeResult
mlir::iree_compiler::buildTileFuseDistToForallAndWorgroupCountWithNumThreads(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes, ArrayAttr threadDimMapping) {
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes,
+ ArrayAttr threadDimMapping) {
return buildTileFuseDistWithNumThreads<TileToForallAndWorkgroupCountRegionOp>(
- b, rootH, opsHToFuse, tileSizes, threadDimMapping);
+ b, isolatedParentOpH, rootH, opsHToFuse, tileSizes, threadDimMapping);
}
/// Apply patterns and vectorize.
@@ -256,13 +292,21 @@
return b.create<VectorizeOp>(funcH);
}
+/// Hoist redundant subet ops.
+Value mlir::iree_compiler::buildHoisting(ImplicitLocOpBuilder &b, Value funcH) {
+ auto pdlOperationType = pdl::OperationType::get(b.getContext());
+ return b.create<HoistRedundantTensorSubsetsOp>(pdlOperationType, funcH);
+}
+
/// Bufferize and drop HAL descriptor from memref ops.
Value mlir::iree_compiler::buildBufferize(ImplicitLocOpBuilder &b,
Value variantH, bool targetGpu) {
- Value funcH = b.create<MatchOp>(variantH, func::FuncOp::getOperationName());
- ApplyPatternsOpPatterns patterns;
- patterns.foldReassociativeReshapes = true;
- funcH = b.create<ApplyPatternsOp>(funcH, patterns);
+ // Perform a pass of canonicalization + enabling before bufferization to avoid
+ // spurious allocations.
+ ApplyPatternsOpPatterns configuration;
+ configuration.foldReassociativeReshapes = true;
+ variantH =
+ buildCanonicalizationAndEnablingTransforms(b, configuration, variantH);
variantH = b.create<IREEEliminateEmptyTensorsOp>(variantH);
variantH = b.create<IREEBufferizeOp>(variantH, targetGpu);
Value memrefFunc =
@@ -310,9 +354,9 @@
}
auto funcH = b.create<MatchOp>(variantH, func::FuncOp::getOperationName());
- ApplyPatternsOpPatterns patterns;
- patterns.bubbleCollapseExpand = true;
- b.create<ApplyPatternsOp>(funcH, patterns);
+ ApplyPatternsOpPatterns configuration;
+ configuration.bubbleCollapseExpand = true;
+ b.create<ApplyPatternsOp>(funcH, configuration);
std::tie(result.originalFillH, result.splitFillH) =
matchAndUnpack<2>(b, variantH, linalg::FillOp::getOperationName());
if (hasTrailingEltwise) {
@@ -334,8 +378,9 @@
/// Return a triple of handles to (forall, fill, combiner)
std::tuple<Value, Value, Value>
mlir::iree_compiler::buildTileReductionUsingScfForeach(
- ImplicitLocOpBuilder &b, Value reductionH, int64_t reductionRank,
- int64_t tileSize, int64_t reductionVectorSize, Attribute mappingAttr) {
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value reductionH,
+ int64_t reductionRank, int64_t tileSize, int64_t reductionVectorSize,
+ Attribute mappingAttr) {
SmallVector<int64_t> leadingParallelDims(reductionRank - 1, 0);
SmallVector<int64_t> numThreads = leadingParallelDims;
numThreads.push_back(tileSize);
@@ -374,6 +419,7 @@
TileToForallAndFuseAndDistributeResult tileResult =
buildTileFuseDistToForallAndWorkgroupCountWithTileSizes(
/*builder=*/b,
+ /*isolatedParentOpH=*/variantH,
/*rootH=*/fusionTargetH,
/*opsToFuseH=*/fusionGroupH,
/*tileSizes=*/
@@ -384,6 +430,12 @@
fillH = b.create<FuseIntoContainingOp>(fillH, tileResult.forallH);
maybeLeadingH =
b.create<FuseIntoContainingOp>(maybeLeadingH, tileResult.forallH);
+
+ // Perform a pass of canonicalization + enabling after fusion.
+ ApplyPatternsOpPatterns configuration;
+ variantH = mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, variantH);
+
// Step 3. Normalize to reorder results irrespective of emptiness.
auto [blockReductionH, maybeBlockTrailingH] = buildSelectFirstNonEmpty(
b, tileResult.resultingFusedOpsHandles.front(), tileResult.tiledOpH);
@@ -393,9 +445,11 @@
Value mlir::iree_compiler::buildMemoryOptimizations(ImplicitLocOpBuilder &b,
Value funcH) {
- ApplyPatternsOpPatterns patterns;
- patterns.lowerTransferOpPermutations = true;
- patterns.rankReducingVector = true;
- funcH = b.create<ApplyPatternsOp>(funcH, patterns);
+ ApplyPatternsOpPatterns configuration;
+ configuration.lowerTransferOpPermutations = true;
+ configuration.rankReducingVector = true;
+ // Apply canonicalizations and enablings twice as they enable each other.
+ buildCanonicalizationAndEnablingTransforms(b, configuration, funcH);
+ buildCanonicalizationAndEnablingTransforms(b, configuration, funcH);
return b.create<ApplyBufferOptimizationsOp>(funcH);
}
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.h b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.h
index 7d8620f..49bb7ad 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.h
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.h
@@ -14,6 +14,12 @@
namespace mlir {
namespace iree_compiler {
+namespace IREE {
+namespace transform_dialect {
+struct ApplyPatternsOpPatterns;
+} // namespace transform_dialect
+} // namespace IREE
+
struct AbstractReductionStrategy;
//===----------------------------------------------------------------------===//
@@ -56,6 +62,16 @@
/// `handles` is empty.
void buildPrint(ImplicitLocOpBuilder &b, ValueRange handles = {});
+/// Create an ApplyPatternsOp that performs a set of key canonicalizations and
+/// so-called enabling transformations to normalize the IR.
+/// Take an existing configuration by copy (cheap object) that will be augmented
+/// locally to additionally perform:
+/// canonicalization, tiling_canonicalization, licm and cse (in this order).
+Value buildCanonicalizationAndEnablingTransforms(
+ ImplicitLocOpBuilder &b,
+ IREE::transform_dialect::ApplyPatternsOpPatterns configuration,
+ Value variantH);
+
/// Build transform IR to dynamically selects the first non-empty handle; i.e.
/// if (h1, h2) is:
/// - (non-empty, non-empty), returns (h1, h2)
@@ -86,8 +102,8 @@
/// Build transform IR to perform multi-level tile and fuse into an scf.for op.
/// Note: fusion is currently unsupported.
TileToScfForAndFuseResult buildTileFuseToScfFor(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes);
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes);
/// Result of the combined transform performing tiling, fusion and
/// distribution to parallel constructs.
@@ -127,36 +143,43 @@
///
// TODO: if someone knows how to properly export templates go for it .. sigh.
TileToForallAndFuseAndDistributeResult buildTileFuseDistToForallWithTileSizes(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes, ArrayAttr threadDimMapping);
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes,
+ ArrayAttr threadDimMapping);
/// Version of `buildTileFuseDistToForallWithTileSizes` that is aware of
/// IREE's `workgroup_count` region and should be used for the block-level
/// tiling in a dispatch region.
TileToForallAndFuseAndDistributeResult
buildTileFuseDistToForallAndWorkgroupCountWithTileSizes(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> tileSizes, ArrayAttr threadDimMapping);
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> tileSizes,
+ ArrayAttr threadDimMapping);
/// Similar to `buildTileFuseDistWithTileSizes` but using `numThreads` instead
/// of `tileSizes`.
TileToForallAndFuseAndDistributeResult buildTileFuseDistToForallWithNumThreads(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> numThreads, ArrayAttr threadDimMapping);
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> numThreads,
+ ArrayAttr threadDimMapping);
/// Version of `buildTileFuseDistToForallWithNumThreads` that is aware of
/// IREE's `workgroup_count` region and should be used for the block-level
/// tiling in a dispatch region.
TileToForallAndFuseAndDistributeResult
buildTileFuseDistToForallAndWorgroupCountWithNumThreads(
- ImplicitLocOpBuilder &b, Value rootH, ValueRange opsHToFuse,
- ArrayRef<OpFoldResult> numThreads, ArrayAttr threadDimMapping);
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value rootH,
+ ValueRange opsHToFuse, ArrayRef<OpFoldResult> numThreads,
+ ArrayAttr threadDimMapping);
-/// Build transform IR that applies rank-reduction patterns and vectorizes.
+/// Build transform IR that applies rank-reduction patterns and vectorizes.
/// Takes a handle to a func.func and returns an updated handle to a
/// func.func.
Value buildVectorize(ImplicitLocOpBuilder &b, Value funcH);
+/// Build transform IR to hoist redundant subset operations.
+Value buildHoisting(ImplicitLocOpBuilder &b, Value funcH);
+
/// Build transform IR to bufferize and drop HAL descriptor from memref ops.
/// Takes a handle variantOp and returns a handle to the same variant op.
Value buildBufferize(ImplicitLocOpBuilder &b, Value variantH,
@@ -169,8 +192,9 @@
/// resulting containing forall op.
/// Return a triple of handles to (forall, fill, combiner)
std::tuple<Value, Value, Value> buildTileReductionUsingScfForeach(
- ImplicitLocOpBuilder &b, Value reductionH, int64_t reductionRank,
- int64_t tileSize, int64_t reductionVectorSize, Attribute mappingAttr);
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value reductionH,
+ int64_t reductionRank, int64_t tileSize, int64_t reductionVectorSize,
+ Attribute mappingAttr);
//===----------------------------------------------------------------------===//
// Higher-level problem-specific strategy creation APIs, these should favor
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp
index b029779..6802b3d 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.cpp
@@ -175,9 +175,9 @@
//===----------------------------------------------------------------------===//
void mlir::iree_compiler::gpu::
build1DSplittingStrategyWithOptionalThreadMapping(
- ImplicitLocOpBuilder &b, Value opH, int64_t rank, int64_t mostMinorDim,
- SmallVector<int64_t> opSizes, int64_t numThreads, Attribute mappingAttr,
- int64_t maxVectorSize) {
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value opH,
+ int64_t rank, int64_t mostMinorDim, SmallVector<int64_t> opSizes,
+ int64_t numThreads, Attribute mappingAttr, int64_t maxVectorSize) {
// Poor man's handling of optionality in C++. Will need to be converted to
// proper transform dialect filters or handling of emptiness.
if (rank == 0) return;
@@ -210,6 +210,7 @@
if (vectorSize > 1) {
auto res = iree_compiler::buildTileFuseToScfFor(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*rootH=*/opH,
/*opsHToFuse=*/{},
/*tileSizes=*/
@@ -223,6 +224,7 @@
assert(mappingAttr && "must specify a mapping attribute");
iree_compiler::buildTileFuseDistToForallWithNumThreads(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*rootH=*/opH,
/*opsHToFuse=*/{},
/*numThreads=*/getAsOpFoldResult(b.getI64ArrayAttr(foreachTileSizes)),
@@ -235,6 +237,7 @@
if (vectorSize > 1) {
auto res = iree_compiler::buildTileFuseToScfFor(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*rootH=*/opH,
/*opsHToFuse=*/{},
/*tileSizes=*/getAsOpFoldResult(b.getI64ArrayAttr({scfForTileSizes})));
@@ -244,6 +247,7 @@
assert(mappingAttr && "must specify a mapping attribute");
iree_compiler::buildTileFuseDistToForallWithNumThreads(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*rootH=*/opH,
/*opsHToFuse=*/{},
/*numThreads=*/getAsOpFoldResult(b.getI64ArrayAttr(foreachTileSizes)),
@@ -258,22 +262,36 @@
std::pair<Value, Value> mlir::iree_compiler::gpu::buildCommonTrailingStrategy(
ImplicitLocOpBuilder &b, Value variantH,
const AbstractReductionStrategy &strategy) {
- // Step N-1. Bufferize and drop HAL descriptor from memref ops.
Value funcH = b.create<MatchOp>(variantH, func::FuncOp::getOperationName());
- // Fold tensor.empty to avoid large allocations.
- ApplyPatternsOpPatterns patterns;
- patterns.foldTensorEmptyExtract = true;
- funcH = b.create<ApplyPatternsOp>(funcH, patterns);
+ // Step N-5. Fold tensor.empty to avoid large allocations.
+ ApplyPatternsOpPatterns configuration;
+ configuration.foldTensorEmptyExtract = true;
+ // Step N-4. Perform a pass of canonicalization + enabling after tiling.
+ funcH = mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, funcH);
funcH = iree_compiler::buildVectorize(b, funcH);
+
+ // Step N-3. Perform a pass of canonicalization + enabling after vectorization
+ // as well as hoisting subset operations such as vector.transfer_read/write.
+ funcH = mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, funcH);
+ funcH = iree_compiler::buildHoisting(b, funcH);
+
+ // Step N-2. Bufferize and drop HAL descriptor from memref ops.
variantH = iree_compiler::buildBufferize(b, variantH, /*targetGpu=*/true);
- // Step N. Post-bufferization mapping to blocks and threads.
+ // Step N-1. Post-bufferization mapping to blocks and threads.
// Need to match again since bufferize invalidated all handles.
// TODO: assumes a single func::FuncOp to transform, may need hardening.
funcH = b.create<MatchOp>(variantH, func::FuncOp::getOperationName());
funcH = buildMapToBlockAndThreads(b, funcH, strategy.getNumThreadsInBlock());
+
+ // Step N. Perform a final pass of canonicalization + enabling before
+ // returning.
+ variantH = mlir::iree_compiler::buildCanonicalizationAndEnablingTransforms(
+ b, configuration, variantH);
return std::make_pair(variantH, funcH);
}
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.h b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.h
index 7f36138..babb1fc 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.h
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/Common.h
@@ -78,8 +78,8 @@
// in schedule complexity and can be handled with simple padding of the
// underlying allocation.
void build1DSplittingStrategyWithOptionalThreadMapping(
- ImplicitLocOpBuilder& b, Value opH, int64_t rank, int64_t mostMinorDim,
- SmallVector<int64_t> opSizes, int64_t numThreads,
+ ImplicitLocOpBuilder& b, Value isolatedParentOpH, Value opH, int64_t rank,
+ int64_t mostMinorDim, SmallVector<int64_t> opSizes, int64_t numThreads,
Attribute mappingAttr = Attribute(), int64_t maxVectorSize = 4);
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/SmallReductionStrategy.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/SmallReductionStrategy.cpp
index 6cf9c72..17f95e9 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/SmallReductionStrategy.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/SmallReductionStrategy.cpp
@@ -96,14 +96,16 @@
}
static void buildSmallReductionStrategyThreadDistribution(
- ImplicitLocOpBuilder &b, Value maybeLeadingH, Value fillH, Value reductionH,
- Value maybeTrailingH, const AbstractReductionStrategy &strategy) {
+ ImplicitLocOpBuilder &b, Value variantH, Value maybeLeadingH, Value fillH,
+ Value reductionH, Value maybeTrailingH,
+ const AbstractReductionStrategy &strategy) {
auto [fusionTargetH, fusionGroupH] =
iree_compiler::buildSelectFirstNonEmpty(b, maybeTrailingH, reductionH);
ArrayRef<Attribute> allThreadsRef(strategy.allThreadAttrs);
iree_compiler::TileToForallAndFuseAndDistributeResult tileResult =
iree_compiler::buildTileFuseDistToForallWithNumThreads(
/*builder=*/b,
+ /*isolatedParentOpH=*/variantH,
/*rootH=*/fusionTargetH,
/*opsToFuseH=*/fusionGroupH,
/*numThreads=*/
@@ -130,6 +132,7 @@
// part.
build1DSplittingStrategyWithOptionalThreadMapping(
/*b=*/b,
+ /*isolatedParentOpH=*/variantH,
/*opH=*/blockReductionH,
/*rank=*/strategy.captures.reductionRank,
// TODO: capture and generalize mostMinorDim.
@@ -140,6 +143,7 @@
// 3. apply the 1d splitting strategy to the trailing elementwise.
build1DSplittingStrategyWithOptionalThreadMapping(
/*b=*/b,
+ /*isolatedParentOpH=*/variantH,
/*opH=*/maybeBlockTrailingH,
/*rank=*/strategy.captures.maybeTrailingRank,
// TODO: capture and generalize mostMinorDim.
@@ -159,7 +163,7 @@
// Step 2. Apply thread-level part of the strategy, keeps everything fused.
buildSmallReductionStrategyThreadDistribution(
- b, maybeLeadingHBlock, gridFillH, gridReductionH,
+ b, variantH, maybeLeadingHBlock, gridFillH, gridReductionH,
maybeTiledTrailingHBlock, strategy);
// Step 3-4. Common trailing steps.
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/StagedReductionStrategy.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/StagedReductionStrategy.cpp
index 5f1b480..131a9ab 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/StagedReductionStrategy.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/GPU/StagedReductionStrategy.cpp
@@ -124,8 +124,8 @@
}
static void buildStagedReductionStrategyThreadLevel(
- ImplicitLocOpBuilder &b, Value gridReductionH, Value gridFillH,
- Value maybeTiledLeadingH, Value maybeTiledTrailingH,
+ ImplicitLocOpBuilder &b, Value isolatedParentOpH, Value gridReductionH,
+ Value gridFillH, Value maybeTiledLeadingH, Value maybeTiledTrailingH,
const StagedReductionStrategy &strategy) {
// Map the potential maybeTiledLeadingH.
// TODO: Consider fusing leading elementwise into threads.
@@ -136,6 +136,7 @@
assert((vectorSize & (vectorSize - 1)) == 0 && "size must be power of 2");
build1DSplittingStrategyWithOptionalThreadMapping(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*opH=*/maybeTiledLeadingH,
/*rank=*/strategy.captures.maybeLeadingRank,
// TODO: capture and generalize mostMinorDim.
@@ -150,6 +151,7 @@
auto [blockParallelForallOp, blockParallelFillH, blockCombinerOpH] =
buildTileReductionUsingScfForeach(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*reductionH=*/gridReductionH,
/*reductionRank=*/strategy.captures.reductionRank,
/*tileSize=*/strategy.getNumThreadsXInBlock(),
@@ -183,6 +185,7 @@
}
iree_compiler::buildTileFuseDistToForallWithTileSizes(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*rootH=*/root,
/*opsToFuse=*/opsToFuse,
/*tileSizes=*/getAsOpFoldResult(b.getI64ArrayAttr({1})),
@@ -196,6 +199,7 @@
strategy.captures.maybeTrailingOutputElementalTypeBitWidth;
build1DSplittingStrategyWithOptionalThreadMapping(
/*b=*/b,
+ /*isolatedParentOpH=*/isolatedParentOpH,
/*opH=*/maybeTiledTrailingH,
/*rank=*/strategy.captures.maybeTrailingRank,
// TODO: capture and generalize mostMinorDim.
@@ -221,9 +225,10 @@
// Step 2. Split the reduction and tile the pieces to ensure vector
// load/stores and mapping to a single warp with shuffles.
- buildStagedReductionStrategyThreadLevel(b, gridReductionH, gridFillH,
- maybeLeadingHBlock,
- maybeTiledTrailingHBlock, strategy);
+ buildStagedReductionStrategyThreadLevel(
+ b,
+ /*isolatedParentOpH=*/variantH, gridReductionH, gridFillH,
+ maybeLeadingHBlock, maybeTiledTrailingHBlock, strategy);
// Step 3. Make sure we don't create allocation by sharing forall
// output. This amounts to injecting user-defined static information that each
diff --git a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.td b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.td
index 3d16f0e..c39d12f 100644
--- a/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.td
+++ b/llvm-external-projects/iree-dialects/include/iree-dialects/Dialect/LinalgTransform/StructuredTransformOpsExt.td
@@ -16,49 +16,6 @@
include "mlir/IR/OpAsmInterface.td"
include "mlir/IR/OpBase.td"
-def CanonicalizedSequenceOp
- : TransformDialectOp<"structured.canonicalized_sequence",
- [OpAsmOpInterface,
- PossibleTopLevelTransformOpTrait,
- DeclareOpInterfaceMethods<RegionBranchOpInterface,
- ["getSuccessorEntryOperands", "getSuccessorRegions",
- "getRegionInvocationBounds"]>,
- DeclareOpInterfaceMethods<TransformOpInterface>,
- DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
- SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
-
- let summary = "A transformation sequence interspersed with canonicalizations";
- let description = [{
- This op is a copy of `transform.sequence`, but applies canonicalizations
- after each step in the sequence.
- }];
-
- let arguments = (ins FailurePropagationMode:$failure_propagation_mode,
- Optional<PDL_Operation>:$root);
- let results = (outs Variadic<AnyType>:$results);
- let regions = (region SizedRegion<1>:$body);
-
- let assemblyFormat =
- "($root^)? `failures` `(` $failure_propagation_mode `)` attr-dict-with-keyword regions (`:` type($results)^)?";
-
- let builders = [
- OpBuilder<(ins "::mlir::transform::FailurePropagationMode":$failure_propagation_mode,
- CArg<"::llvm::function_ref<void(::mlir::OpBuilder &, ::mlir::Location, ::mlir::Value)>",
- "nullptr">)>
- ];
-
- let extraClassDeclaration = [{
- using BodyBuilderFn =
- ::llvm::function_ref<void(::mlir::OpBuilder &, ::mlir::Location, ::mlir::Value)>;
-
- /// Allow the dialect prefix to be omitted.
- static ::llvm::StringRef getDefaultDialect() { return "transform"; }
- }];
-
- let cppNamespace = "mlir::transform_ext";
- let hasVerifier = 1;
-}
-
def LowerToLLVMOp : Op<Transform_Dialect, "lower_to_llvm",
[FunctionalStyleTransformOpTrait,
MemoryEffectsOpInterface,
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/StructuredTransformOpsExt.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/StructuredTransformOpsExt.cpp
index 264f47e..20865dc 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/StructuredTransformOpsExt.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/IR/StructuredTransformOpsExt.cpp
@@ -508,446 +508,6 @@
}
//===----------------------------------------------------------------------===//
-// CanonicalizedSequenceOp
-//===----------------------------------------------------------------------===//
-
-void ::transform_ext::CanonicalizedSequenceOp::build(
- OpBuilder &builder, OperationState &state,
- transform::FailurePropagationMode failurePropagationMode,
- ::transform_ext::CanonicalizedSequenceOp::BodyBuilderFn bodyBuilder) {
- assert(state.name.isRegistered() && "not registered!");
- assert(bodyBuilder && "requires a body builder");
- MLIRContext *ctx = builder.getContext();
- state.addAttribute(
- CanonicalizedSequenceOp::getFailurePropagationModeAttrName(state.name),
- transform::FailurePropagationModeAttr::get(ctx, failurePropagationMode));
- Region *bodyRegion = state.addRegion();
- bodyRegion->push_back(new Block);
- Block &bodyBlock = bodyRegion->front();
- bodyBlock.addArgument(pdl::OperationType::get(ctx), state.location);
- OpBuilder::InsertionGuard guard(builder);
- builder.setInsertionPointToStart(&bodyBlock);
- bodyBuilder(builder, state.location, bodyBlock.getArgument(0));
-}
-
-/// Run enabling transformations (LICM and its variants, single-iteration loop
-/// removal, CSE) on the given function.
-static LogicalResult performEnablerTransformations(
- func::FuncOp func, RewriterBase::Listener &listener,
- LinalgEnablingOptions options = LinalgEnablingOptions()) {
- MLIRContext *ctx = func->getContext();
- RewritePatternSet patterns(ctx);
- linalg::populateLinalgTilingCanonicalizationPatterns(patterns);
- scf::populateSCFForLoopCanonicalizationPatterns(patterns);
- GreedyRewriteConfig config;
- config.listener = &listener;
- if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns), config)))
- return failure();
-
- // This assumes LICM never removes operations so we don't need tracking.
- if (options.licm) {
- func->walk(
- [](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
- }
-
- func.walk([](Operation *op) {
- (void)llvm::TypeSwitch<Operation *, LogicalResult>(op)
- .Case<AffineForOp, scf::ForOp>(
- [](auto loop) { return promoteIfSingleIteration(loop); })
- .Default([](Operation *) { return success(); });
- });
-
- if (options.hoistRedundantVectorTransfers)
- linalg::hoistRedundantVectorTransfers(func);
- if (options.hoistRedundantVectorTransfersOnTensor)
- linalg::hoistRedundantVectorTransfersOnTensor(func);
-
- return eliminateCommonSubexpressions(func, /*domInfo=*/nullptr, &listener);
-}
-
-/// Run enabling transformations on the given `containerOp` while preserving the
-/// operation tracking information.
-static LogicalResult performEnablerTransformations(
- Operation *containerOp, RewriterBase::Listener &listener,
- LinalgEnablingOptions options = LinalgEnablingOptions()) {
- auto res = containerOp->walk([&](func::FuncOp func) {
- if (failed(performEnablerTransformations(func, listener, options)))
- return WalkResult::interrupt();
- return WalkResult::advance();
- });
- return failure(res.wasInterrupted());
-}
-
-/// Drop the association between payload operations and transform dialect
-/// handles when it is no longer necessary in a canonicalized sequence.
-/// Specifically, drop the association between payload operations and the
-/// operand handles if all handles to them will not be used after the current
-/// `transform`. Also drop the association between payload operations and result
-/// handles if results are never read. Note that the operand part is specific to
-/// sequence-like execution that is not guaranteed in the transform dialect in
-/// general.
-static void
-forgetUnnecessaryHandles(transform::TransformState &state,
- transform_ext::CanonicalizedSequenceOp sequence,
- transform::TransformOpInterface transform) {
- auto *listener = state.getExtension<TrackingListener>();
- assert(transform->getParentOp() == sequence &&
- "only works for transform ops immediately nested in a canonicalized "
- "sequence");
- assert(listener && "expected tracking listener to be present");
-
- // Checks if the operation or its ancestor is before `transform` in its block
- // or is `transform` itself.
- auto userIsBefore = [&](Operation *user) {
- while (user && user->getParentOp() != sequence)
- user = user->getParentOp();
- if (!user)
- return false;
- return user->isBeforeInBlock(transform) || user == transform;
- };
-
- // Drop associations for operands that will not be read again. Ignore consumed
- // operands that have been deassociated already. Consider all handles to each
- // payload operation and only drop the association if all handles pointing to
- // the same operation will are not used after the current transform op. The
- // handle will be erased automatically after the last payload operation is
- // deassociated from it.
- llvm::SmallDenseSet<Operation *> seen;
- llvm::SmallDenseMap<Value, bool> handlesUsedAfterTransform;
- for (Value operand : transform->getOperands()) {
- if (transform::isHandleConsumed(operand, transform))
- continue;
- if (!operand.getType().isa<transform::TransformHandleTypeInterface>())
- continue;
- for (Operation *payload : state.getPayloadOps(operand)) {
- if (!payload || seen.contains(payload))
- continue;
- SmallVector<Value> allHandles;
- (void)state.getHandlesForPayloadOp(payload, allHandles);
- bool allHandlesUnused = llvm::all_of(allHandles, [&](Value handle) {
- if (!handlesUsedAfterTransform.count(handle)) {
- handlesUsedAfterTransform[handle] =
- !llvm::all_of(handle.getUsers(), userIsBefore);
- }
- return !handlesUsedAfterTransform[handle];
- });
- if (allHandlesUnused) {
- // Store for mapping removal, but not remove immediately because it
- // could invalidate other handles.
- seen.insert(payload);
- }
- }
- }
-
- // Drop associations for results that will never be read.
- for (Value result : transform->getResults()) {
- if (!result.getUses().empty())
- continue;
- if (!result.getType().isa<transform::TransformHandleTypeInterface>())
- continue;
- for (Operation *payload : state.getPayloadOps(result)) {
- if (!payload || seen.contains(payload))
- continue;
- // Store for mapping removal, but not remove immediately because it could
- // invalidate other handles.
- seen.insert(payload);
- }
- }
-
- // Remove everything in one bunch.
- for (Operation *payload : seen)
- listener->removeMappings(payload);
-}
-
-DiagnosedSilenceableFailure transform_ext::CanonicalizedSequenceOp::apply(
- transform::TransformResults &results, transform::TransformState &state) {
-
- MLIRContext *ctx = getContext();
- RewritePatternSet patternList(ctx);
- for (Dialect *dialect : ctx->getLoadedDialects())
- dialect->getCanonicalizationPatterns(patternList);
- for (RegisteredOperationName op : ctx->getRegisteredOperations())
- op.getCanonicalizationPatterns(patternList, ctx);
- FrozenRewritePatternSet patterns(std::move(patternList));
-
- transform::TransformState::RegionScope regionScope =
- state.make_region_scope(getBodyRegion());
- auto &listener = state.addExtension<::mlir::TrackingListener>();
- auto detachListener = llvm::make_scope_exit(
- [&] { state.removeExtension<::mlir::TrackingListener>(); });
- if (failed(mapBlockArguments(state)))
- return DiagnosedSilenceableFailure::definiteFailure();
-
- auto checkedListenerTransform =
- [&](function_ref<LogicalResult(Operation *, RewriterBase::Listener &)>
- transform) {
- SmallVector<Operation *> roots;
- if (Value root = getRoot())
- llvm::append_range(roots, state.getPayloadOps(root));
- else
- roots.push_back(state.getTopLevel());
-
- for (Operation *target : roots) {
- // Make sure we always check the error state, no boolean
- // short-circuting.
- if (failed(transform(target, listener))) {
- target->emitOpError("Transform application failed.");
- return failure();
- }
- if (failed(listener.checkErrorState())) {
- target->emitOpError("Listener failed.");
- return failure();
- }
- }
- return success();
- };
-
- auto performCSE = [](Operation *root, RewriterBase::Listener &listener) {
- LogicalResult result =
- eliminateCommonSubexpressions(root, /*domInfo=*/nullptr, &listener);
- LLVM_DEBUG(
- DBGS() << (succeeded(result) ? "successfully performed" : "failed")
- << " CSE\n");
- return result;
- };
- auto performEnabler = [](Operation *root, RewriterBase::Listener &listener) {
- LogicalResult result = performEnablerTransformations(root, listener);
- LLVM_DEBUG(
- DBGS() << (succeeded(result) ? "successfully performed" : "failed")
- << " enabling transformations\n");
- return result;
- };
- auto performCanonicalization = [&patterns](Operation *root,
- RewriterBase::Listener &listener) {
- GreedyRewriteConfig config;
- config.listener = &listener;
- // Manually gather list of ops because the other GreedyPatternRewriteDriver
- // overloads only accepts ops that are isolated from above.
- SmallVector<Operation *> ops;
- root->walk([&](Operation *op) {
- if (op != root)
- ops.push_back(op);
- });
- LogicalResult result = applyOpPatternsAndFold(ops, patterns, config);
- LLVM_DEBUG(
- DBGS() << (succeeded(result) ? "successfully performed" : "failed")
- << " canonicalization\n");
- return result;
- };
-
- LLVM_DEBUG(DBGS() << "begin canonicalizing sequence\n");
- if (failed(checkedListenerTransform(performCSE))) {
- return mlir::emitDefiniteFailure(
- *this, "Failed to performCSE beform transform sequence");
- }
- if (failed(checkedListenerTransform(performCanonicalization))) {
- return mlir::emitDefiniteFailure(
- *this, "Failed to performCanonicalization beform transform sequence");
- }
-
- // Apply the sequenced ops one by one.
- for (Operation &transform : getBodyBlock()->without_terminator()) {
- auto transformOp = cast<transform::TransformOpInterface>(transform);
- DiagnosedSilenceableFailure result = state.applyTransform(transformOp);
- if (result.isDefiniteFailure()) {
- LLVM_DEBUG(DBGS() << "failed: " << transform << "\n");
- return result;
- }
- if (result.isSilenceableFailure()) {
- LLVM_DEBUG(DBGS() << "failed silently: " << transform << "\n");
- if (getFailurePropagationMode() ==
- transform::FailurePropagationMode::Propagate)
- return result;
- (void)result.silence();
- }
- LLVM_DEBUG(DBGS() << "successfully performed: " << transform << "\n");
-
- // Canonicalization may replace payload operations associated with the
- // transform dialect handles. Post-canonicalize reassociation is fragile and
- // may fail. To make this less likely, drop any association that are no
- // longer necessary, i.e., if the operand is no longer used in the sequence
- // or elsewhere or if the result is never read.
- forgetUnnecessaryHandles(state, *this, transformOp);
-
- if (failed(checkedListenerTransform(performCSE))) {
- return mlir::emitDefiniteFailure(&transform,
- "Failed to performCSE after transform");
- }
- if (failed(checkedListenerTransform(performEnabler))) {
- return mlir::emitDefiniteFailure(
- &transform, "Failed to performEnabler after transform");
- }
- if (failed(checkedListenerTransform(performCanonicalization))) {
- return mlir::emitDefiniteFailure(
- &transform, "Failed to performCanonicalization after transform");
- }
- }
-
- // Forward the operation mapping for values yielded from the sequence to the
- // values produced by the sequence op.
- for (const auto &pair :
- llvm::zip(getBodyBlock()->getTerminator()->getOperands(),
- getOperation()->getOpResults())) {
- Value terminatorOperand = std::get<0>(pair);
- OpResult result = std::get<1>(pair);
- results.set(result, state.getPayloadOps(terminatorOperand));
- }
-
- LLVM_DEBUG(DBGS() << "end canonicalizing sequence\n");
- return DiagnosedSilenceableFailure::success();
-}
-
-/// Returns `true` if the given op operand may be consuming the handle value in
-/// the Transform IR. That is, if it may have a Free effect on it.
-static bool isValueUsePotentialConsumer(OpOperand &use) {
- // Conservatively assume the effect being present in absence of the interface.
- auto memEffectInterface = dyn_cast<MemoryEffectOpInterface>(use.getOwner());
- if (!memEffectInterface)
- return true;
-
- SmallVector<MemoryEffects::EffectInstance, 2> effects;
- memEffectInterface.getEffectsOnValue(use.get(), effects);
- return llvm::any_of(effects, [](const MemoryEffects::EffectInstance &effect) {
- return isa<transform::TransformMappingResource>(effect.getResource()) &&
- isa<MemoryEffects::Free>(effect.getEffect());
- });
-}
-
-// TODO: Add declaration to TransformOps.h, then we do not have to duplicate
-// this function.
-static LogicalResult
-checkDoubleConsume(Value value,
- function_ref<InFlightDiagnostic()> reportError) {
- OpOperand *potentialConsumer = nullptr;
- for (OpOperand &use : value.getUses()) {
- if (!isValueUsePotentialConsumer(use))
- continue;
-
- if (!potentialConsumer) {
- potentialConsumer = &use;
- continue;
- }
-
- InFlightDiagnostic diag = reportError()
- << " has more than one potential consumer";
- diag.attachNote(potentialConsumer->getOwner()->getLoc())
- << "used here as operand #" << potentialConsumer->getOperandNumber();
- diag.attachNote(use.getOwner()->getLoc())
- << "used here as operand #" << use.getOperandNumber();
- return diag;
- }
-
- return success();
-}
-
-LogicalResult transform_ext::CanonicalizedSequenceOp::verify() {
- // Check if the block argument has more than one consuming use.
- for (BlockArgument argument : getBodyBlock()->getArguments()) {
- auto report = [&]() {
- return (emitOpError() << "block argument #" << argument.getArgNumber());
- };
- if (failed(checkDoubleConsume(argument, report)))
- return failure();
- }
-
- // Check properties of the nested operations they cannot check themselves.
- for (Operation &child : *getBodyBlock()) {
- if (!isa<transform::TransformOpInterface>(child) &&
- &child != &getBodyBlock()->back()) {
- InFlightDiagnostic diag =
- emitOpError()
- << "expected children ops to implement TransformOpInterface";
- diag.attachNote(child.getLoc()) << "op without interface";
- return diag;
- }
-
- for (OpResult result : child.getResults()) {
- auto report = [&]() {
- return (child.emitError() << "result #" << result.getResultNumber());
- };
- if (failed(checkDoubleConsume(result, report)))
- return failure();
- }
- }
-
- if (getBodyBlock()->getTerminator()->getOperandTypes() !=
- getOperation()->getResultTypes()) {
- InFlightDiagnostic diag = emitOpError()
- << "expects the types of the terminator operands "
- "to match the types of the result";
- diag.attachNote(getBodyBlock()->getTerminator()->getLoc()) << "terminator";
- return diag;
- }
- return success();
-}
-
-void transform_ext::CanonicalizedSequenceOp::getEffects(
- SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
- auto *mappingResource = transform::TransformMappingResource::get();
- // Effects on root if present.
- if (getRoot())
- effects.emplace_back(MemoryEffects::Read::get(), getRoot(),
- mappingResource);
- // Effects on results.
- for (Value result : getResults()) {
- effects.emplace_back(MemoryEffects::Allocate::get(), result,
- mappingResource);
- effects.emplace_back(MemoryEffects::Write::get(), result, mappingResource);
- }
-
- for (Operation &op : *getBodyBlock()) {
- auto iface = dyn_cast<MemoryEffectOpInterface>(&op);
- if (!iface) {
- // TODO: fill all possible effects; or require ops to actually implement
- // the memory effect interface always
- assert(false);
- }
- if (getRoot()) {
- // Carry over all effects on the argument of the entry block as those on
- // the operand, this is the same value just remapped.
- SmallVector<MemoryEffects::EffectInstance, 2> nestedEffects;
- iface.getEffectsOnValue(getBodyBlock()->getArgument(0), nestedEffects);
- for (const auto &effect : nestedEffects)
- effects.emplace_back(effect.getEffect(), getRoot(),
- effect.getResource());
- } else {
- // Otherwise, get all the effects.
- iface.getEffects(effects);
- }
- }
-}
-
-OperandRange transform_ext::CanonicalizedSequenceOp::getSuccessorEntryOperands(
- Optional<unsigned> index) {
- assert(index && index.value() == 0 && "unexpected region index");
- if (getOperation()->getNumOperands() == 1)
- return getOperation()->getOperands();
- return OperandRange(getOperation()->operand_end(),
- getOperation()->operand_end());
-}
-
-void transform_ext::CanonicalizedSequenceOp::getSuccessorRegions(
- Optional<unsigned> index, ArrayRef<Attribute> operands,
- SmallVectorImpl<RegionSuccessor> ®ions) {
- if (!index.has_value()) {
- Region *bodyRegion = &getBody();
- regions.emplace_back(bodyRegion, !operands.empty()
- ? bodyRegion->getArguments()
- : Block::BlockArgListType());
- return;
- }
-
- assert(*index == 0 && "unexpected region index");
- regions.emplace_back(getOperation()->getResults());
-}
-
-void transform_ext::CanonicalizedSequenceOp::getRegionInvocationBounds(
- ArrayRef<Attribute> operands, SmallVectorImpl<InvocationBounds> &bounds) {
- (void)operands;
- bounds.emplace_back(1, 1);
-}
-
-//===----------------------------------------------------------------------===//
// TODO: WILL MIGRATE
//===----------------------------------------------------------------------===//
diff --git a/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/_iree_structured_transform_ops_ext.py b/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/_iree_structured_transform_ops_ext.py
index 4a1b08b..754c9fa 100644
--- a/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/_iree_structured_transform_ops_ext.py
+++ b/llvm-external-projects/iree-dialects/python/iree/compiler/dialects/_iree_structured_transform_ops_ext.py
@@ -49,69 +49,6 @@
return value
-class CanonicalizedSequenceOp:
-
- @overload
- def __init__(self, resultsOrRoot: Sequence[Type],
- optionalRoot: Optional[Union[Operation, Value]]):
- ...
-
- @overload
- def __init__(self, resultsOrRoot: Optional[Union[Operation, Value]],
- optionalRoot: NoneType):
- ...
-
- def __init__(self, resultsOrRoot=None, optionalRoot=None):
- results = resultsOrRoot if isinstance(resultsOrRoot, Sequence) else []
- root = (resultsOrRoot
- if not isinstance(resultsOrRoot, Sequence) else optionalRoot)
- root = _get_op_result_or_value(root) if root else None
- super().__init__(results_=results, root=root)
- self.regions[0].blocks.append(pdl.OperationType.get())
-
- @property
- def body(self) -> Block:
- return self.regions[0].blocks[0]
-
- @property
- def bodyTarget(self) -> Value:
- return self.body.arguments[0]
-
-
-class LowerVectorsOp:
- """Specialization for the LowerVectorsOp class."""
-
- def __init__(self,
- *,
- stages: IntListArg = None,
- contraction_lowering: StringArg = None,
- multireduction_lowering: StringArg = None,
- split_transfers: StringArg = None,
- unroll_vector_transfers: BoolArg = None,
- transpose_lowering: StringArg = None,
- transpose_avx2_lowering: BoolArg = None,
- loc=None,
- ip=None):
- stages = _ensure_int_array_attr(stages, [0, 1, 2, 3, 4, 5, 6])
- contraction_lowering = _ensure_string_attr(contraction_lowering,
- "outerproduct")
- multireduction_lowering = _ensure_string_attr(multireduction_lowering,
- "innerparallel")
- split_transfers = _ensure_string_attr(split_transfers, "linalg-copy")
- unroll_vector_transfers = _ensure_bool_attr(unroll_vector_transfers, True)
- transpose_lowering = _ensure_string_attr(transpose_lowering, "eltwise")
- transpose_avx2_lowering = _ensure_bool_attr(transpose_avx2_lowering, False)
- super().__init__(stages=stages,
- contraction_lowering=contraction_lowering,
- multireduction_lowering=multireduction_lowering,
- split_transfers=split_transfers,
- unroll_vector_transfers=unroll_vector_transfers,
- transpose_lowering=transpose_lowering,
- transpose_avx2_lowering=transpose_avx2_lowering,
- loc=loc,
- ip=ip)
-
-
class LowerToLLVMOp:
"""Specialization for the LowerToLLVMOp class."""
diff --git a/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir b/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir
index 53afe4d..a240e7b 100644
--- a/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir
+++ b/tests/transform_dialect/cuda/reduction_v3_codegen_spec.mlir
@@ -1,9 +1,6 @@
// RUN: iree-opt %s
-// TODO: port this test to transform.sequence, atm some canonicalization patterns
-// ping-pong into oblivion.
-// transform.sequence failures(propagate) {
-transform.structured.canonicalized_sequence failures(propagate) {
+transform.sequence failures(propagate) {
^bb1(%variant_op: !pdl.operation):
%fill = transform.structured.match ops{["linalg.fill"]} in %variant_op
: (!pdl.operation) -> !pdl.operation
@@ -17,6 +14,10 @@
( mapping = [#gpu.block<x>] )
transform.structured.fuse_into_containing_op %fill into %forall_grid
+ // Canonicalizations.
+ transform.iree.apply_patterns %variant_op
+ { canonicalization, tiling_canonicalization, licm, cse }
+
// Step 2. Split the reduction to get meatier parallelism.
// This also parallelizes to threads.
// ===========================================================================
@@ -34,6 +35,10 @@
transform.structured.tile_to_forall_op %block_combiner_op_2 num_threads [1]
( mapping = [#gpu.thread<y>] )
+ // Canonicalizations.
+ transform.iree.apply_patterns %variant_op
+ { canonicalization, tiling_canonicalization, licm, cse }
+
// Step 3. Rank-reduce and vectorize.
// ===========================================================================
%func = transform.structured.match ops{["func.func"]} in %variant_op
@@ -44,13 +49,24 @@
{ rank_reducing_linalg, rank_reducing_vector }
%func_3 = transform.structured.vectorize %func_2
+ // Canonicalizations is necessary to get rid of some tensor.cast that block
+ // hoisting.
+ transform.iree.apply_patterns %variant_op
+ { canonicalization, tiling_canonicalization, licm, cse }
+ %func_4 = transform.structured.hoist_redundant_tensor_subsets %func_3
+ : (!pdl.operation) -> !pdl.operation
+
+
// Step 4. Bufferize and drop HAL descriptor from memref ops.
// ===========================================================================
- %func_4 = transform.iree.apply_patterns %func_3 { fold_reassociative_reshapes }
+ // Canonicalizations required before bufferization to avoid unnecessary allocs.
+ transform.iree.apply_patterns %variant_op
+ { canonicalization, tiling_canonicalization, licm, cse }
+ %func_5 = transform.iree.apply_patterns %func_4 { fold_reassociative_reshapes }
%variant_op_2 = transform.iree.eliminate_empty_tensors %variant_op
- %func_5 = transform.structured.match ops{["func.func"]} in %variant_op_2
+ %func_6 = transform.structured.match ops{["func.func"]} in %variant_op_2
: (!pdl.operation) -> !pdl.operation
- %func_6 = transform.iree.apply_patterns %func_5 { erase_unnecessary_tensor_operands }
+ %func_7 = transform.iree.apply_patterns %func_6 { erase_unnecessary_tensor_operands }
%variant_op_3 = transform.iree.bufferize { target_gpu } %variant_op_2
%memref_func = transform.structured.match ops{["func.func"]} in %variant_op_3
: (!pdl.operation) -> !pdl.operation
@@ -58,18 +74,23 @@
// Step 5. Post-bufferization mapping to blocks and threads.
// ===========================================================================
- %func_7 = transform.structured.match ops{["func.func"]} in %variant_op_3
+ %func_m = transform.structured.match ops{["func.func"]} in %variant_op_3
: (!pdl.operation) -> !pdl.operation
- %func_8 = transform.iree.forall_to_workgroup %func_7
- %func_9 = transform.iree.map_nested_forall_to_gpu_threads %func_8
+ %func_m_2 = transform.iree.forall_to_workgroup %func_m
+ %func_m_3 = transform.iree.map_nested_forall_to_gpu_threads %func_m_2
{ workgroup_size = [1024, 1, 1] }
// Step 6. Post-bufferization vector distribution with rank-reduction.
// ===========================================================================
- %func_10 = transform.iree.apply_patterns %func_9 { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases }
+ %func_m_4 = transform.iree.apply_patterns %func_m_3 { rank_reducing_linalg, rank_reducing_vector, fold_memref_aliases }
%if_op = transform.structured.match ops{["scf.if"]} in %variant_op_3
: (!pdl.operation) -> !pdl.operation
%warp = transform.iree.vector.to_warp_execute_on_lane_0 %if_op { warp_size = 32 }
- transform.iree.vector.warp_distribute %func_10
+ transform.iree.vector.warp_distribute %func_m_4
: (!pdl.operation) -> !pdl.operation
+
+
+ // Late canonicalizations.
+ transform.iree.apply_patterns %variant_op_3
+ { canonicalization, tiling_canonicalization, licm, cse }
}
