[Mlir-commits] [mlir] e86fe36 - [MLIR] Allow Affine scalar replacement to handle inner operations

[William S. Moses]

Author: William S. Moses
Date: 2021-07-01T15:12:59-04:00
New Revision: e86fe368db11e43b4b7e84065ba5e8439fc24840

URL: https://github.com/llvm/llvm-project/commit/e86fe368db11e43b4b7e84065ba5e8439fc24840
DIFF: https://github.com/llvm/llvm-project/commit/e86fe368db11e43b4b7e84065ba5e8439fc24840.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Analysis/Utils.cpp
    mlir/test/Dialect/Affine/scalrep.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index e87ecdac2d6c..fc14e198c5b2 100644
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -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());

diff  --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 308186fe676d..31e52323ed78 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -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:      }