[mlir][bufferization] Drop the assumption for alloc result index (#134503)

Relax the assumption that alloc op always has allocation at
`getResult(0)`, allow to use `optimize-allocation-liveness` pass for
custom ops with >1 results. Ops with multiple allocations are not
handled here yet.

Author: FruitClover

mlir/lib/Dialect/Bufferization/Transforms/OptimizeAllocationLiveness.cpp

@@ -17,7 +17,10 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 
 #define DEBUG_TYPE "optimize-allocation-liveness"
 #define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
@@ -88,6 +91,19 @@ static bool hasMemoryAllocEffect(MemoryEffectOpInterface memEffectOp) {
   return false;
 }
 
+/// Extracts OpResult's with Allocate effects from given op
+static SmallVector<OpResult>
+collectAllocations(MemoryEffectOpInterface allocOp) {
+  SmallVector<MemoryEffects::EffectInstance> effects;
+  allocOp.getEffects(effects);
+  SmallVector<OpResult> allocResults;
+  for (const MemoryEffects::EffectInstance &it : effects)
+    if (isa<MemoryEffects::Allocate>(it.getEffect()))
+      if (auto val = it.getValue(); val && val.getDefiningOp() == allocOp)
+        allocResults.push_back(cast<OpResult>(val));
+  return allocResults;
+}
+
 struct OptimizeAllocationLiveness
     : public bufferization::impl::OptimizeAllocationLivenessPassBase<
           OptimizeAllocationLiveness> {
@@ -109,7 +125,15 @@ struct OptimizeAllocationLiveness
       auto allocOp = memEffectOp;
       LDBG("Checking alloc op: " << allocOp);
 
-      auto deallocOp = findUserWithFreeSideEffect(allocOp->getResult(0));
+      SmallVector<OpResult> allocationResults = collectAllocations(allocOp);
+      // Multiple allocations from a single op are not considered here yet.
+      if (allocationResults.size() != 1)
+        return WalkResult::advance();
+
+      OpResult allocResult = allocationResults[0];
+      LDBG("On allocation result: " << allocResult);
+
+      auto *deallocOp = findUserWithFreeSideEffect(allocResult);
       if (!deallocOp || (deallocOp->getBlock() != allocOp->getBlock())) {
         // The pass handles allocations that have a single dealloc op in the
         // same block. We also should not hoist the dealloc op out of
@@ -119,9 +143,9 @@ struct OptimizeAllocationLiveness
 
       Operation *lastUser = nullptr;
       const BufferViewFlowAnalysis::ValueSetT &deps =
-          analysis.resolve(allocOp->getResult(0));
+          analysis.resolve(allocResult);
       for (auto dep : llvm::make_early_inc_range(deps)) {
-        for (auto user : dep.getUsers()) {
+        for (auto *user : dep.getUsers()) {
           // We are looking for a non dealloc op user.
           // check if user is the dealloc op itself.
           if (user == deallocOp)

mlir/test/Dialect/Bufferization/Transforms/optimize-allocation-liveness.mlir

@@ -209,3 +209,28 @@ func.func private @test_conditional_deallocation() -> memref<32xf32, 1> {
   return %3 : memref<32xf32, 1>
 }
 
+
+// -----
+// CHECK-LABEL:   func.func private @test_alloc_with_multiple_results() {
+// CHECK:           %[[ID1:.+]], %[[ALLOC1:.+]] = test.alloc_with_multiple_results : index, memref<64xf32>
+// CHECK:           memref.expand_shape %[[ALLOC1]]
+// CHECK:           memref.dealloc %[[ALLOC1]] : memref<64xf32>
+// CHECK:           %[[ID2:.+]], %[[ALLOC2:.+]] = test.alloc_with_multiple_results : index, memref<64xf32>
+// CHECK:           memref.expand_shape %[[ALLOC2]]
+// CHECK:           memref.dealloc %[[ALLOC2]] : memref<64xf32>
+// CHECK:           return
+// CHECK:         }
+
+// This test will check that allocations with multiple results and allocated
+// buffer at non-zero position are accepted.
+func.func private @test_alloc_with_multiple_results() -> () {
+  %id1, %alloc1 = test.alloc_with_multiple_results : index, memref<64xf32>
+  %expand_shape1 = memref.expand_shape %alloc1 [[0, 1]] output_shape [8, 8] : memref<64xf32> into memref<8x8xf32>
+
+  %id2, %alloc2 = test.alloc_with_multiple_results : index, memref<64xf32>
+  %expand_shape2 = memref.expand_shape %alloc2 [[0, 1]] output_shape [8, 8] : memref<64xf32> into memref<8x8xf32>
+
+  memref.dealloc %alloc1 : memref<64xf32>
+  memref.dealloc %alloc2 : memref<64xf32>
+  return
+}

mlir/test/lib/Dialect/Test/TestOps.td

@@ -3441,4 +3441,16 @@ def TestMultiSlotAlloca : TEST_Op<"multi_slot_alloca",
   let assemblyFormat = "attr-dict `:` functional-type(operands, results)";
 }
 
+//===----------------------------------------------------------------------===//
+// Test allocation Ops
+//===----------------------------------------------------------------------===//
+
+def TestAllocWithMultipleResults : TEST_Op<"alloc_with_multiple_results"> {
+  let results = (outs Index:$index,
+                      Res<AnyMemRef, "", [MemAlloc]>:$memref);
+  let assemblyFormat = [{
+     attr-dict `:` type($index) `,` type($memref)
+  }];
+}
+
 #endif // TEST_OPS