[Mlir-commits] [mlir] d423fc3 - Add RegionBranchOpInterface on affine.for op
Uday Bondhugula
llvmlistbot at llvm.org
Wed Apr 20 05:19:45 PDT 2022
Author: Uday Bondhugula
Date: 2022-04-20T17:46:07+05:30
New Revision: d423fc372466c304dfec530d2d9cc559ad01a762
URL: https://github.com/llvm/llvm-project/commit/d423fc372466c304dfec530d2d9cc559ad01a762
DIFF: https://github.com/llvm/llvm-project/commit/d423fc372466c304dfec530d2d9cc559ad01a762.diff
LOG: Add RegionBranchOpInterface on affine.for op
Add RegionBranchOpInterface on affine.for op so that transforms relying
on RegionBranchOpInterface can support affine.for. E.g.:
buffer-deallocation pass.
Reviewed By: herhut
Differential Revision: https://reviews.llvm.org/D123568
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
mlir/test/Transforms/sccp-structured.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 854c582a42b58..a44f13da4f0fb 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -110,7 +110,9 @@ def AffineForOp : Affine_Op<"for",
[AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
- "getSingleUpperBound"]>]> {
+ "getSingleUpperBound"]>,
+ DeclareOpInterfaceMethods<RegionBranchOpInterface,
+ ["getSuccessorEntryOperands"]>]> {
let summary = "for operation";
let description = [{
Syntax:
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 60ada4eaa9b52..d6534900a2546 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1723,6 +1723,57 @@ void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
results.add<AffineForEmptyLoopFolder>(context);
}
+/// Return operands used when entering the region at 'index'. These operands
+/// correspond to the loop iterator operands, i.e., those excluding the
+/// induction variable. AffineForOp only has one region, so zero is the only
+/// valid value for `index`.
+OperandRange AffineForOp::getSuccessorEntryOperands(unsigned index) {
+ assert(index == 0 && "invalid region index");
+
+ // The initial operands map to the loop arguments after the induction
+ // variable.
+ return getIterOperands();
+}
+
+/// Given the region at `index`, or the parent operation if `index` is None,
+/// return the successor regions. These are the regions that may be selected
+/// during the flow of control. `operands` is a set of optional attributes that
+/// correspond to a constant value for each operand, or null if that operand is
+/// not a constant.
+void AffineForOp::getSuccessorRegions(
+ Optional<unsigned> index, ArrayRef<Attribute> operands,
+ SmallVectorImpl<RegionSuccessor> ®ions) {
+ assert((!index.hasValue() || index.getValue() == 0) &&
+ "expected loop region");
+ // The loop may typically branch back to its body or to the parent operation.
+ // If the predecessor is the parent op and the trip count is known to be at
+ // least one, branch into the body using the iterator arguments. And in cases
+ // we know the trip count is zero, it can only branch back to its parent.
+ Optional<uint64_t> tripCount = getTrivialConstantTripCount(*this);
+ if (!index.hasValue() && tripCount.hasValue()) {
+ if (tripCount.getValue() > 0) {
+ regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+ return;
+ }
+ if (tripCount.getValue() == 0) {
+ regions.push_back(RegionSuccessor(getResults()));
+ return;
+ }
+ }
+
+ // From the loop body, if the trip count is one, we can only branch back to
+ // the parent.
+ if (index.hasValue() && tripCount.hasValue() && tripCount.getValue() == 1) {
+ regions.push_back(RegionSuccessor(getResults()));
+ return;
+ }
+
+ // In all other cases, the loop may branch back to itself or the parent
+ // operation.
+ regions.push_back(RegionSuccessor(&getLoopBody(), getRegionIterArgs()));
+ regions.push_back(RegionSuccessor(getResults()));
+}
+
/// Returns true if the affine.for has zero iterations in trivial cases.
static bool hasTrivialZeroTripCount(AffineForOp op) {
Optional<uint64_t> tripCount = getTrivialConstantTripCount(op);
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
index e7219b5c8cb7a..c50b053b25bab 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir
@@ -1092,6 +1092,24 @@ func @loop_nested_alloc(
// -----
+// CHECK-LABEL: func @affine_loop
+func @affine_loop() {
+ %buffer = memref.alloc() : memref<1024xf32>
+ %sum_init_0 = arith.constant 0.0 : f32
+ %res = affine.for %i = 0 to 10 step 2 iter_args(%sum_iter = %sum_init_0) -> f32 {
+ %t = affine.load %buffer[%i] : memref<1024xf32>
+ %sum_next = arith.addf %sum_iter, %t : f32
+ affine.yield %sum_next : f32
+ }
+ // CHECK: %[[M:.*]] = memref.alloc
+ // CHECK: affine.for
+ // CHECK: }
+ // CHECK-NEXT: memref.dealloc %[[M]]
+ return
+}
+
+// -----
+
// Test Case: explicit control-flow loop with a dynamically allocated buffer.
// The BufferDeallocation transformation should fail on this explicit
// control-flow loop since they are not supported.
diff --git a/mlir/test/Transforms/sccp-structured.mlir b/mlir/test/Transforms/sccp-structured.mlir
index 32184bfdfda11..eea683a3b93e0 100644
--- a/mlir/test/Transforms/sccp-structured.mlir
+++ b/mlir/test/Transforms/sccp-structured.mlir
@@ -149,3 +149,33 @@ func @loop_region_branch_terminator_op(%arg1 : i32) {
}
return
}
+
+/// Check that propgation happens for affine.for -- tests its region branch op
+/// interface as well.
+
+// CHECK-LABEL: func @affine_loop_one_iter(
+func @affine_loop_one_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+ // CHECK: %[[C1:.*]] = arith.constant 1 : i32
+ %s0 = arith.constant 0 : i32
+ %s1 = arith.constant 1 : i32
+ %result = affine.for %i = 0 to 1 iter_args(%si = %s0) -> (i32) {
+ %sn = arith.addi %si, %s1 : i32
+ affine.yield %sn : i32
+ }
+ // CHECK: return %[[C1]] : i32
+ return %result : i32
+}
+
+// CHECK-LABEL: func @affine_loop_zero_iter(
+func @affine_loop_zero_iter(%arg0 : index, %arg1 : index, %arg2 : index) -> i32 {
+ // This exposes a crash in sccp/forward data flow analysis: https://github.com/llvm/llvm-project/issues/54928
+ // CHECK: %[[C0:.*]] = arith.constant 0 : i32
+ %s0 = arith.constant 0 : i32
+ // %result = affine.for %i = 0 to 0 iter_args(%si = %s0) -> (i32) {
+ // %sn = arith.addi %si, %si : i32
+ // affine.yield %sn : i32
+ // }
+ // return %result : i32
+ // CHECK: return %[[C0]] : i32
+ return %s0 : i32
+}
More information about the Mlir-commits
mailing list