[Mlir-commits] [mlir] [MLIR][Affine] Fix fusion across ops with unknown memory effects (PR #203231)

[Aditya Pradhan]

================
@@ -38,6 +38,45 @@ using llvm::SmallDenseMap;
 
 using Node = MemRefDependenceGraph::Node;
 
+/// Returns the values that `op` may have a memref effect of type `EffectTys`
+/// on, not considering recursive effects. An op with unknown memory effects
+/// (e.g. a call to an external function without a memory-effect interface) is
+/// conservatively assumed to affect all its memref operands.
+template <typename... EffectTys>
+static void getMayEffectedValues(Operation *op,
+                                 SmallVectorImpl<Value> &values) {
+  auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+  if (!memOp) {
+    if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+      // No effects.
+      return;
+    // Memref operands have to be considered as being affected.
+    for (Value operand : op->getOperands()) {
+      if (isa<MemRefType>(operand.getType()))
+        values.push_back(operand);
+    }
+    return;
+  }
+  SmallVector<SideEffects::EffectInstance<MemoryEffects::Effect>, 4> effects;
+  memOp.getEffects(effects);
+  for (auto &effect : effects) {
+    Value effectVal = effect.getValue();
+    if (isa<EffectTys...>(effect.getEffect()) && effectVal &&
+        isa<MemRefType>(effectVal.getType()))
+      values.push_back(effectVal);
+  };
+}
+
+/// Returns true if `op` may have a memory effect of type `EffectTys` on
+/// `memref`, i.e., whether `memref` is among the values returned by
+/// `getMayEffectedValues` for `op`.
+template <typename... EffectTys>
+static bool mayHaveEffect(Operation *op, Value memref) {
+  SmallVector<Value> values;
+  getMayEffectedValues<EffectTys...>(op, values);
+  return llvm::is_contained(values, memref);
+}
+
----------------
pradhanaditya wrote:

Thanks for your review!

The calculation is cheap in practice -- this branch is only hit for non-affine ops, it runs over just an op's few memref operands and stays in SmallVector's inline storage, no heap allocation.

More importantly, having mayHaveEffect go through getMayEffectedValues is the point of the fix: the bug was that recording accesses and querying them used different logic and drifted apart. Sharing the one helper guarantees they stay in lockstep.

If it ever shows up as a major time contributor in a profile (which is unlikely), caching would be the way to go -- but I would keep that out of this correctness-focused change.

https://github.com/llvm/llvm-project/pull/203231