[MLIR] Extend isLoopMemoryParallel to account for locally allocated memrefs

[bondhugula]

Summary

Extend isLoopMemoryParallel check to include locally allocated memrefs.
This strengthens and also speeds up the dependence check used by the
utility by excluding locally allocated memrefs where appropriate.

Additional memref dialect ops can be supported exhaustively via proper
interfaces.

[bondhugula]

Fix test case spacing.

[dcaballe]

Thanks, Uday. A couple of comments that would need attention. LGTM. You can move forward after addressing them.

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

122

Use Operation::IsAncestor instead?

146–149

I think it's important that we move this information to the utility documentation above. We may also want to rename it so that the name state that it only checks for memory dependences. I would expect isLoopMemoryParallel to also check for memrefs that escape or used as iter args.

[ayzhuang]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

137

Do we need isNestedWithin? How about !isDefinedOutsideOfLoop since enclosingOp is essentially AffineForOp?

[bondhugula]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

122

Thanks! isProperAncestor is the one that I missed using here.

137

The isNestedWithin is the same as "not being defined outside of loop". It's defined inside the loop if and only if enclosingOp is a proper ancestor of defOp. Any situation you had in mind that's not covered?

146–149

I think it's important that we move this information to the utility documentation

I kept the detailed comment here and added an additional line on the util doc comment clarifying that even memrefs on iter_args aren't checked.

We may also want to rename it so that the name state that

This is a good point. I was thinking about this. The current doc comment already states:

This function doesn't check iter_args and should be used
only as a building block for full parallel-checking functions.

How about actually checking for memref types being passed as iter_args in this utility? That way, all memory dependences would be accounted for and the name isLoopMemoryParallel would be accurate and it would still be conservatively detecting parallel loops irrespective of any memref typed iter_args. (You have the other method isLoopParallel that gives the semblances of a full check given the latter's general name.)

[bondhugula]

Address some of the review comments.

[bondhugula]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

146–149

Added an extra to doc comment that memrefs allocated inside forOp don't impact dependences and parallelism. While we can add a check for memrefs among iter arguments, it doesn't defend against custom/unknown types (say out-of-tree types) that somehow alias with memrefs -- but the latter part should be clear since "memory" for this method is associated with built-in memrefs. isLoopParallel is expected to take care of all iter_args.

[Hardcode84]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

134

Use ViewLikeOpInterface?

[bondhugula]

Use ViewLikeOpInterface.

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

134

Thanks - this was what I was missing!

[ayzhuang]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

137

No, I meant that we didn't need to define a new function.

[dcaballe]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

146–149

I guess it's ok as long as it's properly documented.

How about actually checking for memref types being passed as iter_args in this utility?

That sounds good to me and I would also consider adding checks for memrefs that escape. That would keep all the memory checks in one place and would make isLoopMemoryParallel less surprising to the user, I think. However, I'm ok with the current implementation as long as it is properly documented.

[bondhugula]

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

146–149

Done. (Note that since the result types of the affine.for match 1:1 with iter_args/iter operands, just checking for memref types in result types is sufficient to check for both cross-iter memref flow and escapes out of loop.)

[bondhugula]

Check iter_args for memref types.

[bondhugula]

Adjust comments and include.

[bondhugula]

Rebase and resolve conflict.

[bondhugula]

Trim extra include.

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

 /// Returns true if `forOp' is a parallel loop. If `parallelReductions` is
 /// provided, populates it with descriptors of the parallelizable reductions and
 /// treats them as not preventing parallelization.
 bool isLoopParallel(
     AffineForOp forOp,
     SmallVectorImpl<LoopReduction> *parallelReductions = nullptr);
 
 /// Returns true if `forOp' doesn't have memory dependences preventing
-/// parallelization. This function doesn't check iter_args and should be used
-/// only as a building block for full parallel-checking functions.
+/// parallelization. Memrefs that are allocated inside `forOp` do not impact its
+/// dependences and parallelism. This function does not check iter_args (for
+/// values other than memref types) and should be used only as a building block
+/// for complete parallelism-checking functions.
 bool isLoopMemoryParallel(AffineForOp forOp);
 
 /// Returns in `affineApplyOps`, the sequence of those AffineApplyOp
 /// Operations that are reachable via a search starting from `operands` and
 /// ending at those operands that are not the result of an AffineApplyOp.
 void getReachableAffineApplyOps(ArrayRef<Value> operands,
                                 SmallVectorImpl<Operation *> &affineApplyOps);
 

mlir/lib/Dialect/Affine/Analysis/AffineAnalysis.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
-#include "mlir/Support/MathExtras.h"
-#include "llvm/ADT/DenseMap.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "affine-analysis"
 
 using namespace mlir;
 
if (parallelReductions) {
     if (parallelReductions->size() != numIterArgs)
       return false;
   }
 
   // Check memory dependences.
   return isLoopMemoryParallel(forOp);
 }
 
-/// Returns true if `forOp' doesn't have memory dependences preventing
-/// parallelization. This function doesn't check iter_args and should be used
-/// only as a building block for full parallel-checking functions.
+/// Returns true if `op` is an alloc-like op, i.e., one allocating memrefs.
+static bool isAllocLikeOp(Operation *op) {
+  auto memEffects = dyn_cast<MemoryEffectOpInterface>(op);
+  return memEffects && memEffects.hasEffect<MemoryEffects::Allocate>();
+}
+
+/// Returns true if `v` is allocated locally to `enclosingOp` -- i.e., it is
+/// allocated by an operation nested within `enclosingOp`.
+static bool isLocallyDefined(Value v, Operation *enclosingOp) {
+  Operation *defOp = v.getDefiningOp();
+  if (!defOp)
+    return false;
+
+  if (isAllocLikeOp(defOp) && enclosingOp->isProperAncestor(defOp))
+    return true;
+
+  // Aliasing ops.
+  auto viewOp = dyn_cast<ViewLikeOpInterface>(defOp);
+  return viewOp && isLocallyDefined(viewOp.getViewSource(), enclosingOp);
+}
+
 bool mlir::isLoopMemoryParallel(AffineForOp forOp) {
+  // Any memref-typed iteration arguments are treated as serializing.
+  if (llvm::any_of(forOp.getResultTypes(),
+                   [](Type type) { return type.isa<BaseMemRefType>(); }))
+    return false;
+
   // Collect all load and store ops in loop nest rooted at 'forOp'.
   SmallVector<Operation *, 8> loadAndStoreOps;
   auto walkResult = forOp.walk([&](Operation *op) -> WalkResult {
-    if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
+    if (auto readOp = dyn_cast<AffineReadOpInterface>(op)) {
+      // Memrefs that are allocated inside `forOp` need not be considered.
+      if (!isLocallyDefined(readOp.getMemRef(), forOp))
+        loadAndStoreOps.push_back(op);
+    } else if (auto writeOp = dyn_cast<AffineWriteOpInterface>(op)) {
+      // Filter out stores the same way as above.
+      if (!isLocallyDefined(writeOp.getMemRef(), forOp))
         loadAndStoreOps.push_back(op);
-    else if (!isa<AffineForOp, AffineYieldOp, AffineIfOp>(op) &&
-             !MemoryEffectOpInterface::hasNoEffect(op))
+    } else if (!isa<AffineForOp, AffineYieldOp, AffineIfOp>(op) &&
+               !isAllocLikeOp(op) &&
+               !MemoryEffectOpInterface::hasNoEffect(op)) {
+      // Alloc-like ops inside `forOp` are fine (they don't impact parallelism)
+      // as long as they don't escape the loop (which has been checked above).
       return WalkResult::interrupt();
+    }
 
     return WalkResult::advance();
   });
 
   // Stop early if the loop has unknown ops with side effects.
   if (walkResult.wasInterrupted())
     return false;
 

mlir/test/Dialect/Affine/parallelize.mlir

affine.for %j = 0 to 10 {
     // CHECK:                          to (min(%[[UB0]], %[[UB1]]))
     affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %j)
             to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
       affine.load %m[%i] : memref<?xf32>
     }
   }
   return
 }
+
+// Test in the presence of locally allocated memrefs.
+
+// CHECK: func @local_alloc
+func @local_alloc() {
+  %cst = arith.constant 0.0 : f32
+  affine.for %i = 0 to 100 {
+    %m = memref.alloc() : memref<1xf32>
+    %ma = memref.alloca() : memref<1xf32>
+    affine.store %cst, %m[0] : memref<1xf32>
+  }
+  // CHECK: affine.parallel
+  return
+}
+
+// CHECK: func @local_alloc_cast
+func @local_alloc_cast() {
+  %cst = arith.constant 0.0 : f32
+  affine.for %i = 0 to 100 {
+    %m = memref.alloc() : memref<128xf32>
+    affine.for %j = 0 to 128 {
+      affine.store %cst, %m[%j] : memref<128xf32>
+    }
+    affine.for %j = 0 to 128 {
+      affine.store %cst, %m[0] : memref<128xf32>
+    }
+    %r = memref.reinterpret_cast %m to offset: [0], sizes: [8, 16],
+           strides: [16, 1] : memref<128xf32> to memref<8x16xf32>
+    affine.for %j = 0 to 8 {
+      affine.store %cst, %r[%j, %j] : memref<8x16xf32>
+    }
+  }
+  // CHECK: affine.parallel
+  // CHECK:   affine.parallel
+  // CHECK:   }
+  // CHECK:   affine.for
+  // CHECK:   }
+  // CHECK:   affine.parallel
+  // CHECK:   }
+  // CHECK: }
+
+  return
+}
+
+// CHECK-LABEL: @iter_arg_memrefs
+func @iter_arg_memrefs(%in: memref<10xf32>) {
+  %mi = memref.alloc() : memref<f32>
+  // Loop-carried memrefs are treated as serializing the loop.
+  // CHECK: affine.for
+  %mo = affine.for %i = 0 to 10 iter_args(%m_arg = %mi) -> (memref<f32>) {
+    affine.yield %m_arg : memref<f32>
+  }
+  return
+}