[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

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/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp

if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(op)) {
           // Number of loops containing the operation `op` which has the
           // potential memory side effect and can occur on a path between
           // `start` and `memOp`.
           unsigned nsLoops = getNumCommonSurroundingLoops(*op, *memOp);
 
           // For ease, let's consider the case that `op` is a store and we're
           // looking for other potential stores (e.g `op`) that overwrite memory
           // after `start`, and before being read in `memOp`. In this case, we
-          // only need to consider other potential stores with depth >
+          // only need to consider other potential stores with depth >=
           // minSurrounding loops since `start` would overwrite any store with a
           // smaller number of surrounding loops before.
           unsigned d;
-          for (d = nsLoops + 1; d > minSurroundingLoops; d--) {
+          for (d = nsLoops + 1; d >= minSurroundingLoops; d--) {
             DependenceResult result = checkMemrefAccessDependence(
                 srcAccess, destAccess, d, &dependenceConstraints,
                 /*dependenceComponents=*/nullptr);
             if (hasDependence(result)) {
               hasSideEffect = true;
               return;
             }
           }

mlir/test/Dialect/Affine/scalrep.mlir

 // CHECK-NEXT:      %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  %{{.*}} = affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:  return %{{.*}} : f32
 }
 
+// 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:      }