[MLIR] Fix hasNoInterveningEffect in the presence of ops from different affine scopes

[bondhugula]

Summary

Fix hasNoInterveningEffect in the presence of ops from different affine
scopes. Also, correctly check for dependence failures as well instead of
just for the existence of a dependence.

mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h

DependenceResult checkMemrefAccessDependence(
     bool allowRAR = false);
 
 /// Utility function that returns true if the provided DependenceResult
 /// corresponds to a dependence result.
 inline bool hasDependence(DependenceResult result) {
   return result.value == DependenceResult::HasDependence;
 }
 
+/// Returns true if the provided DependenceResult corresponds to the absence of
+/// a dependence.
+inline bool noDependence(DependenceResult result) {
+  return result.value == DependenceResult::NoDependence;
+}
+
 /// Returns in 'depCompsVec', dependence components for dependences between all
 /// load and store ops in loop nest rooted at 'forOp', at loop depths in range
 /// [1, maxLoopDepth].
 void getDependenceComponents(
     AffineForOp forOp, unsigned maxLoopDepth,
     std::vector<SmallVector<DependenceComponent, 2>> *depCompsVec);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H

mlir/lib/Dialect/Affine/Utils/Utils.cpp

if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(op)) {
       if (!opMayHaveEffect)
         return;
 
       // If the side effect comes from an affine read or write, try to
       // prove the side effecting `op` cannot reach `memOp`.
       if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op)) {
         MemRefAccess srcAccess(op);
         MemRefAccess destAccess(memOp);
-        // Dependence analysis is only correct if both ops operate on the same
-        // memref.
-        if (srcAccess.memref == destAccess.memref) {
-          FlatAffineValueConstraints dependenceConstraints;
-
+        // Affine dependence analysis here is applicable only if both ops
+        // operate on the same memref and if `op`, `memOp`, and `start` are in
+        // the same AffineScope.
+        if (srcAccess.memref == destAccess.memref &&
+            getAffineScope(op) == getAffineScope(memOp) &&
+            getAffineScope(op) == getAffineScope(start)) {
           // Number of loops containing the start op and the ending operation.
           unsigned minSurroundingLoops =
               getNumCommonSurroundingLoops(*start, *memOp);
 
           // 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 >
           // minSurrounding loops since `start` would overwrite any store with a
           // smaller number of surrounding loops before.
           unsigned d;
+          FlatAffineValueConstraints dependenceConstraints;
           for (d = nsLoops + 1; d > minSurroundingLoops; d--) {
             DependenceResult result = checkMemrefAccessDependence(
                 srcAccess, destAccess, d, &dependenceConstraints,
                 /*dependenceComponents=*/nullptr);
-            if (hasDependence(result)) {
+            // A dependence failure or the presence of a dependence implies a
+            // side effect.
+            if (!noDependence(result)) {
               hasSideEffect = true;
               return;
             }
           }
 
           // No side effect was seen, simply return.
           return;
         }
+        // TODO: Check here if the memrefs alias: there is no side effect if
+        // `srcAccess.memref` and `destAccess.memref` don't alias.
       }
+      // We have an op with a memory effect and we cannot prove if it
+      // intervenes.
       hasSideEffect = true;
       return;
     }
 
     if (op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
       // Recurse into the regions for this op and check whether the internal
       // operations may have the side effect `EffectType` on memOp.
       for (Region &region : op->getRegions())

mlir/test/Dialect/Affine/scalrep.mlir

func.func @affine_load_store_in_different_scopes() -> memref<1xf32> {
   }
   %v = affine.load %A[0] : memref<1xf32>
   // CHECK:      affine.store
   // CHECK-NEXT: test.affine_scope
   // CHECK:        affine.store
   // CHECK:      affine.load
   return %A : memref<1xf32>
 }
+
+// No forwarding should again happen here.
+
+// CHECK-LABEL: func.func @no_forwarding_across_scopes
+func.func @no_forwarding_across_scopes() -> memref<1xf32> {
+  %A = memref.alloc() : memref<1xf32>
+  %cf0 = arith.constant 0.0 : f32
+  %cf5 = arith.constant 5.0 : f32
+  %c0 = arith.constant 0 : index
+  %c100 = arith.constant 100 : index
+  %c1 = arith.constant 1 : index
+
+  // Store shouldn't be forwarded to the load.
+  affine.store %cf0, %A[0] : memref<1xf32>
+  // CHECK:      test.affine_scope
+  // CHECK-NEXT:   affine.load
+  test.affine_scope {
+    %l = affine.load %A[0] : memref<1xf32>
+    %s = arith.addf %l, %cf5 : f32
+    affine.store %s, %A[0] : memref<1xf32>
+    "terminator"() : () -> ()
+  }
+  return %A : memref<1xf32>
+}