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

[bondhugula]

Summary

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.

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

 struct Write : public Effect::Base<Write> {};
 } // namespace MemoryEffects
 
 //===----------------------------------------------------------------------===//
 // SideEffect Utilities
 //===----------------------------------------------------------------------===//
 
 /// 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.
 bool isOpTriviallyDead(Operation *op);
 
 /// Return true if the given operation would be dead if unused, and has no side
 /// effects on memory that would prevent erasing. This is equivalent to checking
 /// `isOpTriviallyDead` if `op` was unused.

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

 #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/Func/IR/FuncOps.h"
 #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"
 #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 `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))
+  if (hasSingleEffect<MemoryEffects::Allocate>(defOp, v) &&
+      enclosingOp->isProperAncestor(defOp))
     return true;
 
   // Aliasing ops.
   auto viewOp = dyn_cast<ViewLikeOpInterface>(defOp);
   return viewOp && isLocallyDefined(viewOp.getViewSource(), enclosingOp);
 }
 
 bool mlir::isLoopMemoryParallel(AffineForOp forOp) {
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) &&
-               !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).
       return WalkResult::interrupt();
     }
 
     return WalkResult::advance();
   });

mlir/lib/Interfaces/SideEffectInterfaces.cpp

 bool mlir::hasSingleEffect(Operation *op, Value value) {
   auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
   if (!memOp)
     return false;
   SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
   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)
       return false;
   }
   return hasSingleEffectOnVal;
 }