[Mlir-commits] [mlir] 8d7f270 - [MLIR] NFC. Introduce mlir::hasEffect and refactor usages dialect util

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2022-09-07T12:21:57+05:30
New Revision: 8d7f2701860ee587057770652a58e55eb7a001b9

URL: https://github.com/llvm/llvm-project/commit/8d7f2701860ee587057770652a58e55eb7a001b9
DIFF: https://github.com/llvm/llvm-project/commit/8d7f2701860ee587057770652a58e55eb7a001b9.diff

LOG: [MLIR] NFC. Introduce mlir::hasEffect and refactor usages dialect util

Introduce mlir::hasEffect and refactor existing usage to use utility.
NFC.

Reviewed By: rriddle, mehdi_amini

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

Added: 
    

Modified: 
    mlir/include/mlir/Interfaces/SideEffectInterfaces.h
    mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
    mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
    mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
    mlir/lib/Interfaces/SideEffectInterfaces.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index f9cd96e9762ed..4b219e71aa756 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -256,6 +256,12 @@ struct Write : public Effect::Base<Write> {};
 template <typename EffectTy>
 bool hasSingleEffect(Operation *op, Value value = nullptr);
 
+/// Returns true if `op` has an effect of type `EffectTy` on `value`. If no
+/// `value` is provided, simply check if effects of the given type(s) are
+/// present.
+template <typename... EffectTys>
+bool hasEffect(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);

diff  --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index efa1c96a06933..edffbe5e815b5 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -803,18 +803,11 @@ bool MemRefDependenceGraph::init(func::FuncOp f) {
         Node node(nextNodeId++, &op);
         nodes.insert({node.id, node});
       }
-    } else if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+    } else if (hasEffect<MemoryEffects::Write, MemoryEffects::Free>(&op)) {
       // Create graph node for top-level op, which could have a memory write
       // side effect.
-      SmallVector<MemoryEffects::EffectInstance, 1> effects;
-      effectInterface.getEffects(effects);
-      if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &it) {
-            return isa<MemoryEffects::Write, MemoryEffects::Free>(
-                it.getEffect());
-          })) {
-        Node node(nextNodeId++, &op);
-        nodes.insert({node.id, node});
-      }
+      Node node(nextNodeId++, &op);
+      nodes.insert({node.id, node});
     }
   }
 

diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
index 2d5a63561282c..955871b31be5d 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefDialect.cpp
@@ -42,22 +42,11 @@ void mlir::memref::MemRefDialect::initialize() {
   addInterfaces<MemRefInlinerInterface>();
 }
 
-/// Finds a single dealloc operation for the given allocated value.
+/// Finds the unique dealloc operation (if one exists) for `allocValue`.
 llvm::Optional<Operation *> mlir::memref::findDealloc(Value allocValue) {
   Operation *dealloc = nullptr;
   for (Operation *user : allocValue.getUsers()) {
-    auto effectInterface = dyn_cast<MemoryEffectOpInterface>(user);
-    if (!effectInterface)
-      continue;
-    // Try to find a free effect that is applied to one of our values
-    // that will be automatically freed by our pass.
-    SmallVector<MemoryEffects::EffectInstance, 2> effects;
-    effectInterface.getEffectsOnValue(allocValue, effects);
-    const bool isFree =
-        llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) {
-          return isa<MemoryEffects::Free>(it.getEffect());
-        });
-    if (!isFree)
+    if (!hasEffect<MemoryEffects::Free>(user, allocValue))
       continue;
     // If we found > 1 dealloc, return None.
     if (dealloc)

diff  --git a/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp b/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
index 265797bc61bf2..9ddaa088e0af1 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformInterfaces.cpp
@@ -374,8 +374,8 @@ bool transform::isHandleConsumed(Value handle,
   auto iface = cast<MemoryEffectOpInterface>(transform.getOperation());
   SmallVector<MemoryEffects::EffectInstance> effects;
   iface.getEffectsOnValue(handle, effects);
-  return hasEffect<MemoryEffects::Read, TransformMappingResource>(effects) &&
-         hasEffect<MemoryEffects::Free, TransformMappingResource>(effects);
+  return ::hasEffect<MemoryEffects::Read, TransformMappingResource>(effects) &&
+         ::hasEffect<MemoryEffects::Free, TransformMappingResource>(effects);
 }
 
 void transform::producesHandle(

diff  --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 8685ba270a63c..b6884f076a151 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -128,6 +128,26 @@ template bool mlir::hasSingleEffect<MemoryEffects::Free>(Operation *, Value);
 template bool mlir::hasSingleEffect<MemoryEffects::Read>(Operation *, Value);
 template bool mlir::hasSingleEffect<MemoryEffects::Write>(Operation *, Value);
 
+template <typename... EffectTys>
+bool mlir::hasEffect(Operation *op, Value value) {
+  auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+  if (!memOp)
+    return false;
+  SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+  memOp.getEffects(effects);
+  return llvm::any_of(effects, [&](MemoryEffects::EffectInstance &effect) {
+    if (value && effect.getValue() != value)
+      return false;
+    return isa<EffectTys...>(effect.getEffect());
+  });
+}
+template bool mlir::hasEffect<MemoryEffects::Allocate>(Operation *, Value);
+template bool mlir::hasEffect<MemoryEffects::Free>(Operation *, Value);
+template bool mlir::hasEffect<MemoryEffects::Read>(Operation *, Value);
+template bool mlir::hasEffect<MemoryEffects::Write>(Operation *, Value);
+template bool
+mlir::hasEffect<MemoryEffects::Write, MemoryEffects::Free>(Operation *, Value);
+
 bool mlir::wouldOpBeTriviallyDead(Operation *op) {
   if (op->mightHaveTrait<OpTrait::IsTerminator>())
     return false;