[MLIR] Allow Affine scalar replacement to handle inner operations

[wsmoses]

Summary

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.

[ftynse]

The relevant test looks broken

[ayzhuang]

mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp

155 ↗

(On Diff #355006)

should be >. A loop nest of depth 2 can have dependence at 3 different levels. Here all the stores have dependence into the following load, at different depths. First store: depth 1, second store: depth 2, third store: depth 3:

affine.for %i0 = 0 to 9 {
    affine.for %i1 = 0 to 9 {
      %a0 = affine.apply affine_map<(d0, d1) -> (d0 + 1)> (%i0, %i1)
      affine.store %c7, %m[%a0, %i1] : memref<20x20xf32>
      %v0 = affine.load %m[%i0, %i1] : memref<20x20xf32>
    }
  }
  
  affine.for %i0 = 0 to 9 {
    affine.for %i1 = 0 to 9 {
      %a1 = affine.apply affine_map<(d0, d1) -> (d1 + 1)> (%i0, %i1)
      affine.store %c7, %m[%i0, %a1] : memref<20x20xf32>
      %v0 = affine.load %m[%i0, %i1] : memref<20x20xf32>
    }
  }
  
  affine.for %i0 = 0 to 9 {
    affine.for %i1 = 0 to 9 {
      affine.store %c7, %m[%i0, %i1] : memref<20x20xf32>
      %v0 = affine.load %m[%i0, %i1] : memref<20x20xf32>
    }
  }

You can take a look at addOrderingConstraints function in lib/Analysis/AffineAnalysis.cpp for better understanding.

[wsmoses]

Fix depth inequality

[ayzhuang]

Looks good, thanks.

mlir/lib/Analysis/Utils.cpp

 
 /// Populates 'loops' with IVs of the loops surrounding 'op' ordered from
 /// the outermost 'affine.for' operation to the innermost one.
 void mlir::getLoopIVs(Operation &op, SmallVectorImpl<AffineForOp> *loops) {
   auto *currOp = op.getParentOp();
   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();
   }
   std::reverse(loops->begin(), loops->end());
 }
 
 /// Populates 'ops' with IVs of the loops surrounding `op`, along with
 /// `affine.if` operations interleaved between these loops, ordered from the
 /// outermost `affine.for` operation to the innermost one.
 void mlir::getEnclosingAffineForAndIfOps(Operation &op,
                                          SmallVectorImpl<Operation *> *ops) {
   ops->clear();
   Operation *currOp = op.getParentOp();
 
   // Traverse up the hierarchy collecting all `affine.for` and `affine.if`
   // operations.
-  while (currOp && (isa<AffineIfOp, AffineForOp>(currOp))) {
+  while (currOp) {
+    if (isa<AffineIfOp, AffineForOp>(currOp))
       ops->push_back(currOp);
     currOp = currOp->getParentOp();
   }
   std::reverse(ops->begin(), ops->end());
 }
 
 // Populates 'cst' with FlatAffineConstraints which represent original domain of
 // the loop bounds that define 'ivs'.
 LogicalResult

mlir/test/Dialect/Affine/scalrep.mlir

func @redundant_store_elim_fail(%out : memref<512xf32>) {
   %cf2 = constant 2.0 : f32
   affine.for %i = 0 to 16 {
     affine.store %cf1, %out[32*%i] : memref<512xf32>
     "test.use"(%out) : (memref<512xf32>) -> ()
     affine.store %cf2, %out[32*%i] : 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:      }