[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

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

[mehdi_amini]

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

279

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

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

40

FYI, less operations:

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

43

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

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/include/mlir/Dialect/StandardOps/IR/Ops.td

def AddIOp : IntArithmeticOp<"addi", [Commutative]> {
   }];
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
-def AllocOp : AllocLikeOp<"alloc"> {
+def AllocOp : AllocLikeOp<"alloc", [MemoryEffects<[MemAlloc]>]> {
   let summary = "memory allocation operation";
   let description = [{
     The `alloc` operation allocates a region of memory, as specified by its
     memref type.
 
     Example:
 
     ```mlir

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 (auto allocEffect = dyn_cast_or_null<MemoryEffectOpInterface>(defOp)) {
+      if (allocEffect.hasEffect<MemoryEffects::Allocate>() &&
+          defOp->getNumResults() == 1 &&
+          isStrided(v.getType().cast<MemRefType>()))
+        return v;
     }
     if (auto viewLikeOp = dyn_cast_or_null<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";