maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[MLIR] Allow Affine scalar replacement to handle inner operations 

Affine scalar replacement (and other affine passes, though not fixed here) don't properly handle operations with nested regions. This patch fixes the pass and two affine utilities to function properly given a non-affine internal region

This patch prevents the pass from throwing an internal compiler error when running on the added test case.

Differential Revision: https://reviews.llvm.org/D105058
This commit is contained in:
William S. Moses 2021-07-01 15:05:27 -04:00
parent 78e70cee0d
commit e86fe368db
2 changed files with 30 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
mlir/lib/Analysis/Utils.cpp
View File

@ -35,9 +35,8 @@ void mlir::getLoopIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops) {
AffineForOp currAffineForOp;
// Traverse up the hierarchy collecting all 'affine.for' operation while
// skipping over 'affine.if' operations.
while (currOp && ((currAffineForOp = dyn_cast<AffineForOp>(currOp)) ||
isa<AffineIfOp>(currOp))) {
if (currAffineForOp)
while (currOp) {
if (AffineForOp currAffineForOp = dyn_cast<AffineForOp>(currOp))
loops->push_back(currAffineForOp);
currOp = currOp->getParentOp();
}
@ -54,8 +53,9 @@ void mlir::getEnclosingAffineForAndIfOps(Operation &op,
// Traverse up the hierarchy collecting all `affine.for` and `affine.if`
// operations.
while (currOp && (isa<AffineIfOp, AffineForOp>(currOp))) {
ops->push_back(currOp);
while (currOp) {
if (isa<AffineIfOp, AffineForOp>(currOp))
ops->push_back(currOp);
currOp = currOp->getParentOp();
}
std::reverse(ops->begin(), ops->end());

26
mlir/test/Dialect/Affine/scalrep.mlir
View File

@ -670,10 +670,34 @@ func @redundant_store_elim_fail(%out : memref<512xf32>) {
}
return
}
// CHECK: affine.for
// CHECK-NEXT: affine.store
// CHECK-NEXT: "test.use"
// CHECK-NEXT: affine.store
// CHECK-NEXT: }
// CHECK-LABEL: @with_inner_ops
func @with_inner_ops(%arg0: memref<?xf64>, %arg1: memref<?xf64>, %arg2: i1) {
%cst = constant 0.000000e+00 : f64
%cst_0 = constant 3.140000e+00 : f64
%cst_1 = constant 1.000000e+00 : f64
affine.for %arg3 = 0 to 28 {
affine.store %cst, %arg1[%arg3] : memref<?xf64>
affine.store %cst_0, %arg1[%arg3] : memref<?xf64>
%0 = scf.if %arg2 -> (f64) {
scf.yield %cst_1 : f64
} else {
%1 = affine.load %arg1[%arg3] : memref<?xf64>
scf.yield %1 : f64
}
affine.store %0, %arg0[%arg3] : memref<?xf64>
}
return
}
// CHECK: %[[pi:.+]] = constant 3.140000e+00 : f64
// CHECK: %{{.*}} = scf.if %arg2 -> (f64) {
// CHECK: scf.yield %{{.*}} : f64
// CHECK: } else {
// CHECK: scf.yield %[[pi]] : f64
// CHECK: }