[MLIR] Clean up checks for alloc-like ops in analysis

Clean up checks for alloc-like ops in analysis. Use the analysis
utility to properly check for the desired kind of effects. The previous
locality utility worked for all practical purposes but wasn't sound and
was locally duplicate code. Instead, use mlir::hasSingleEffect.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D129439

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -251,9 +251,10 @@ struct Write : public Effect::Base<Write> {};
 //===----------------------------------------------------------------------===//
 
 /// Returns true if `op` has only an effect of type `EffectTy` (and of no other
-/// type) on `value`.
+/// type) on `value`. If no value is provided, simply check if effects of that
+/// type and only of that type are present.
 template <typename EffectTy>
-bool hasSingleEffect(Operation *op, Value value);
+bool hasSingleEffect(Operation *op, Value value = nullptr);
 
 /// Return true if the given operation is unused, and has no side effects on
 /// memory that prevent erasing.

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

@@ -21,6 +21,7 @@
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IntegerSet.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
@@ -114,12 +115,6 @@ bool mlir::isLoopParallel(AffineForOp forOp,
   return isLoopMemoryParallel(forOp);
 }
 
-/// 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) {
@@ -127,7 +122,8 @@ static bool isLocallyDefined(Value v, Operation *enclosingOp) {
   if (!defOp)
     return false;
 
-  if (isAllocLikeOp(defOp) && enclosingOp->isProperAncestor(defOp))
+  if (hasSingleEffect<MemoryEffects::Allocate>(defOp, v) &&
+      enclosingOp->isProperAncestor(defOp))
     return true;
 
   // Aliasing ops.
@@ -153,7 +149,7 @@ bool mlir::isLoopMemoryParallel(AffineForOp forOp) {
       if (!isLocallyDefined(writeOp.getMemRef(), forOp))
         loadAndStoreOps.push_back(op);
     } else if (!isa<AffineForOp, AffineYieldOp, AffineIfOp>(op) &&
-               !isAllocLikeOp(op) &&
+               !hasSingleEffect<MemoryEffects::Allocate>(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).

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -99,9 +99,10 @@ bool mlir::hasSingleEffect(Operation *op, Value value) {
   memOp.getEffects(effects);
   bool hasSingleEffectOnVal = false;
   // Iterate through `effects` and check if an effect of type `EffectTy` and
-  // only of that type is present.
+  // only of that type is present. A `value` to check the effect on may or may
+  // not have been provided.
   for (auto &effect : effects) {
-    if (effect.getValue() != value)
+    if (value && effect.getValue() != value)
       continue;
     hasSingleEffectOnVal = isa<EffectTy>(effect.getEffect());
     if (!hasSingleEffectOnVal)