[mlir] Removed tight coupling of BufferPlacement pass to Alloc and Dealloc operations by using MemoryEffectOpInterface queries.

[dfki-mako]

Summary

The current BufferPlacement implementation tries to find Alloc and Dealloc operations in order to move them. However, this is a tight coupling to standard-dialect ops which has been removed in this CL.

[dfki-mako]

Removed todo comment.

[herhut]

Please wait for @mehdi_amini to take a look.

mlir/lib/Transforms/BufferPlacement.cpp

361

Would just return work?

378

Nit: Error messages start lower case.

[pifon2a]

mlir/lib/Transforms/BufferPlacement.cpp

195

nit: add period after "interface".

361

not sure if return would work, but instead of

  if (allocateEffectInstance == effects.end())
        return WalkResult::advance();
  auto allocResult = allocateEffectInstance->getValue().cast<OpResult>();
  placements.emplace_back(
         allocResult, analysis.computeAllocAndDeallocPositions(allocResult));
  return WalkResult::advance();
});

try doing

if (allocateEffectInstance != effects.end()) {
   auto allocResult = allocateEffectInstance->getValue().cast<OpResult>();
   placements.emplace_back(allocResult, 
     analysis.computeAllocAndDeallocPositions(allocResult));
}});

to get rid of WalkResult::advance() .

394

I am a bit confused by these two lines:

OpBuilder builder(alloc.getOwner());
builder.setInsertionPointAfter(positions.getDeallocPosition());

the first line creates a builder and sets its insertion point, the second line changes insertion point again. Why not do that in "one hop"? I think i am just confused.

[rriddle]

mlir/lib/Transforms/BufferPlacement.cpp

194

This should still be Operation* and not OpOperand?

200

Please make the comments complete sentences.

361

The walk methods are allowed to have void return, so you can strip out the WalkResult::advance. You only need that if you want to interrupt.

[mehdi_amini]

Thanks!
LGTM with the other comments addressed

mlir/lib/Transforms/BufferPlacement.cpp

366

Technically I we could have ops allocating two results, seems like you would only see the first?
We may miscompile in such cases I think, can we not try to move such ops if we see them?

[dfki-mako]

Addressed reviewer comments.

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

366

Makes sense to me.

[herhut]

Thanks, LGTM!

mlir/lib/Transforms/BufferPlacement.cpp

computeAllocAndDeallocPositions(OpResult result) const {
     // Get all possible aliases.
     auto possibleValues = aliases.resolve(result);
     return BufferPlacementPositions(getAllocPosition(result, possibleValues),
                                     getDeallocPosition(result, possibleValues));
   }
 
   /// Finds all associated dealloc nodes for the alloc nodes using alias
   /// information.
-  DeallocSetT findAssociatedDeallocs(AllocOp alloc) const {
+  DeallocSetT findAssociatedDeallocs(OpResult allocResult) const {
     DeallocSetT result;
-    auto possibleValues = aliases.resolve(alloc);
+    auto possibleValues = aliases.resolve(allocResult);
     for (Value alias : possibleValues)
-      for (Operation *user : alias.getUsers()) {
-        if (isa<DeallocOp>(user))
-          result.insert(user);
+      for (OpOperand user : alias.getUsers()) {
+        // Check for an existing memory effect interface
+        auto effectInstance =
+            dyn_cast<MemoryEffectOpInterface>(user.getOwner());
+        if (!effectInstance)
+          continue;
+        // Check whether our value will be freed
+        SmallVector<MemoryEffects::EffectInstance, 2> effects;
+        effectInstance.getEffectsOnValue(alias, effects);
+        if (llvm::any_of(effects, [=](MemoryEffects::EffectInstance &it) {
+              return isa<MemoryEffects::Free>(it.getEffect());
+            }))
+          result.insert(user.getOwner());
       }
     return result;
   }
 
   /// Dumps the buffer placement information to the given stream.
   void print(raw_ostream &os) const {
     os << "// ---- Buffer Placement -----\n";
 

 };
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementPass
 //===----------------------------------------------------------------------===//
 
 /// The actual buffer placement pass that moves alloc and dealloc nodes into
 /// the right positions. It uses the algorithm described at the top of the file.
-// TODO: create a templated version that allows to match dialect-specific
-// alloc/dealloc nodes and to insert dialect-specific dealloc node.
 struct BufferPlacementPass
     : mlir::PassWrapper<BufferPlacementPass, FunctionPass> {
   void runOnFunction() override {
     // Get required analysis information first.
     auto &analysis = getAnalysis<BufferPlacementAnalysis>();
 
     // Compute an initial placement of all nodes.
-    llvm::SmallDenseMap<Value, BufferPlacementPositions, 16> placements;
-    getFunction().walk([&](AllocOp alloc) {
-      placements[alloc] = analysis.computeAllocAndDeallocPositions(
-          alloc.getOperation()->getResult(0));
+    llvm::SmallVector<std::pair<OpResult, BufferPlacementPositions>, 16>
+        placements;
+    getFunction().walk([&](MemoryEffectOpInterface op) {
+      // Try to find an allocation result.
+      SmallVector<MemoryEffects::EffectInstance, 2> effects;
+      op.getEffects(effects);
+      auto allocateEffectInstance =
+          llvm::find_if(effects, [=](MemoryEffects::EffectInstance &it) {
+            Value value = it.getValue();
+            return isa<MemoryEffects::Allocate>(it.getEffect()) && value &&
+                   value.isa<OpResult>();
+          });
+      if (allocateEffectInstance == effects.end())
+        return WalkResult::advance();
+
+      // Insert allocation result.
+      auto allocResult = allocateEffectInstance->getValue().cast<OpResult>();
+      placements.emplace_back(
+          allocResult, analysis.computeAllocAndDeallocPositions(allocResult));
       return WalkResult::advance();
     });
 
-    // Move alloc (and dealloc - if any) nodes into the right places
-    // and insert dealloc nodes if necessary.
-    getFunction().walk([&](AllocOp alloc) {
+    // Move alloc (and dealloc - if any) nodes into the right places and insert
+    // dealloc nodes if necessary.
+    for (auto &entry : placements) {
       // Find already associated dealloc nodes.
+      OpResult alloc = entry.first;
       auto deallocs = analysis.findAssociatedDeallocs(alloc);
       if (deallocs.size() > 1) {
         emitError(alloc.getLoc(),
                   "Not supported number of associated dealloc operations");
-        return WalkResult::interrupt();
+        return;
       }
 
       // Move alloc node to the right place.
-      BufferPlacementPositions &positions = placements[alloc];
-      Operation *allocOperation = alloc.getOperation();
+      BufferPlacementPositions &positions = entry.second;
+      Operation *allocOperation = alloc.getOwner();
       allocOperation->moveBefore(positions.getAllocPosition());
 
       // If there is an existing dealloc, move it to the right place.
       if (deallocs.size()) {
         Operation *nextOp = positions.getDeallocPosition()->getNextNode();
         assert(nextOp && "Invalid Dealloc operation position");
         (*deallocs.begin())->moveBefore(nextOp);
       } else {
         // If there is no dealloc node, insert one in the right place.
-        OpBuilder builder(alloc);
+        OpBuilder builder(alloc.getOwner());
         builder.setInsertionPointAfter(positions.getDeallocPosition());
         builder.create<DeallocOp>(allocOperation->getLoc(), alloc);
       }
-      return WalkResult::advance();
-    });
+    }
   };
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // BufferAssignmentPlacer
 //===----------------------------------------------------------------------===//