[mlir][linalg] Use memory effect to detecting allocation

[antiagainst]

Summary

This commit marks AllocLikeOp as MemAlloc in StandardOps.

Also in Linalg dependency analysis use memory effect to detect
allocation. This allows the dependency analysis to be more
general and recognize other allocation-like operations

[rriddle]

Thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

279 ↗

(On Diff #259529)

Can we add to AllocLikeOp instead? Ideally marking the returned value directly.

[mehdi_amini]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

279 ↗

(On Diff #259529)

Should this be rebased on top of https://reviews.llvm.org/D78619 ?

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

41–42

FYI, less operations:

Operation *defOp = v.getDefiningOp();
if (!defOp) return v;

44

The "MemAlloc" trait can be attached to individual results, I believe that the right way here is to look-up which result is v for the operation and query the interface for the MemAlloc trait for this particular result.

[antiagainst]

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

279 ↗

(On Diff #259529)

Oh, thanks for the pointer! I'll wait that CL to land and then rebase. :)

[antiagainst]

Rebase and address comments

[mehdi_amini]

Can this be exercised by a test? (using a test.alloc op?)

[antiagainst]

Address comments

[antiagainst]

In D78705#2006112, @mehdi_amini wrote:

Can this be exercised by a test? (using a test.alloc op?)

Actually we have two allocation-like ops in std. I think leveraging them should suffice for a test over the memory effect here.

[antiagainst]

Rebase

[antiagainst]

@mehdi_amini: Any further comments?

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp

Value Aliases::find(Value v) {
   if (it != aliases.end()) {
     assert(it->getSecond().getType().isa<MemRefType>() && "Memref expected");
     return it->getSecond();
   }
 
   while (true) {
     if (v.isa<BlockArgument>())
       return v;
+
     Operation *defOp = v.getDefiningOp();
-    if (auto alloc = dyn_cast_or_null<AllocOp>(defOp)) {
-      if (isStrided(alloc.getType()))
-        return alloc.getResult();
+    if (!defOp)
+      return v;
+
+    if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(defOp)) {
+      // Collect all memory effects on `v`.
+      SmallVector<MemoryEffects::EffectInstance, 1> effects;
+      memEffect.getEffectsOnValue(v, effects);
+
+      // If we have the 'Allocate' memory effect on `v`, then `v` should be the
+      // original buffer.
+      if (llvm::any_of(
+              effects, [](const MemoryEffects::EffectInstance &instance) {
+                return isa<MemoryEffects::Allocate>(instance.getEffect());
+              }))
+        return v;
     }
-    if (auto viewLikeOp = dyn_cast_or_null<ViewLikeOpInterface>(defOp)) {
+
+    if (auto viewLikeOp = dyn_cast<ViewLikeOpInterface>(defOp)) {
       auto it =
           aliases.insert(std::make_pair(v, find(viewLikeOp.getViewSource())));
       return it.first->second;
     }
 
     llvm::errs() << "View alias analysis reduces to: " << v << "\n";
     llvm_unreachable("unsupported view alias case");
   }