[mlir][Linalg] NFC - More gracefully degrade lookup into failure furing comprehensive bufferization (4/n)

[nicolasvasilache]

Rebase.

[nicolasvasilache]

Rebase.

[nicolasvasilache]

Landed as dd65f420cd2b983ea1e71ed685c811f00110bafb

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

if (auto bbArg = key.dyn_cast<BlockArgument>()) {
         key.getParentBlock()->getParentOp()->dump();
       else
         key.getParentBlock()->getParentOp()->getParentOfType<FuncOp>()->dump();
       bbArg.getOwner()->getParentOp()->dump();
     } else {
       key.getDefiningOp()->getParentOfType<FuncOp>()->dump();
     }
     llvm::errs() << "NO VALUE FOR KEY: " << key << "\n";
-    abort();
+    return Value();
   }
   return bvm.lookup(key);
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific support.
 //===----------------------------------------------------------------------===//
 

 //===----------------------------------------------------------------------===//
 
 /// Helper function for LinalgOp bufferization.
 /// Operate on mixed tensor + buffer Linalg ops for progressive bufferization.
 /// Allocate the output buffers for the remaining tensor output operands of
 /// the Linalg op. If the tensor is an "init" tensor (i.e. its value is
 /// actually used in the payload region), we additionally copy the original
 /// value into the newly allocated buffer.
-static void allocateBuffersForResults(OpBuilder &b, Location loc, LinalgOp op,
-                                      SmallVectorImpl<Value> &resultBuffers,
-                                      BlockAndValueMapping &bvm) {
+static LogicalResult
+allocateBuffersForResults(OpBuilder &b, Location loc, LinalgOp op,
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
   // Lazily compute loopRanges.
   SmallVector<Range, 4> loopRanges;
 
   // Linalg invariant: output tensors and result match 1-1.
   assert(op.getNumOutputTensors() == op->getNumResults());
   for (auto &opOperand : op.getOutputOpOperands()) {
     Value output = opOperand.get();
     if (output.getType().isa<MemRefType>()) {
       resultBuffers.push_back(output);
       continue;
     }
 
     // If output tensor is marked inPlace, just use the buffer.
     // The following uses internal knowledge of the position of tied operand /
     // results. A proper TieOperandInterface would be much better.
     auto inplaceOp = cast<InplaceOpInterface>(op.getOperation());
     OpResult tiedResult = inplaceOp.getTiedOpResult(opOperand);
     if (getInPlace(tiedResult) == InPlaceSpec::True) {
-      resultBuffers.push_back(lookup(bvm, output));
+      Value v = lookup(bvm, output);
+      if (!v)
+        return failure();
+      resultBuffers.push_back(v);
       continue;
     }
 
     Value dimTensor = bvm.lookupOrDefault(output);
     Value alloc = createNewAllocDeallocPairForShapedValue(b, loc, dimTensor);
     b.setInsertionPointAfter(alloc.getDefiningOp());
     resultBuffers.push_back(alloc);
 
     // Additionally, if the output buffer is used, clone its value for now.
-    if (op.payloadUsesValueFromOpOperand(&opOperand))
-      b.create<CopyOp>(loc, lookup(bvm, output), alloc);
+    if (op.payloadUsesValueFromOpOperand(&opOperand)) {
+      Value v = lookup(bvm, output);
+      if (!v)
+        return failure();
+      b.create<CopyOp>(loc, v, alloc);
+    }
   }
   if (op->getNumResults())
     map(bvm, op->getResults(), resultBuffers);
+
+  return success();
 }
 
 static void finalizeBufferAllocation(OpBuilder &b, LinalgOp op,
                                      ValueRange inputs, ValueRange outputs,
                                      BlockAndValueMapping &bvm) {
   SmallVector<Value, 8> newOperands = inputs;
   newOperands.append(outputs.begin(), outputs.end());
   auto otherOperands = op.getAssumedNonShapedOperands();
if (op.hasBufferSemantics())
     return failure();
 
   LLVM_DEBUG(DBGS() << "convert: " << *op << "\n");
 
   b.setInsertionPoint(op);
   Location loc = op.getLoc();
   SmallVector<Value, 2> newInputBuffers;
   newInputBuffers.reserve(op.getNumInputs());
-  for (Value v : op.getInputs())
-    newInputBuffers.push_back(lookup(bvm, v));
+  for (Value in : op.getInputs()) {
+    Value v = lookup(bvm, in);
+    if (!v)
+      return failure();
+    newInputBuffers.push_back(v);
+  }
   SmallVector<Value, 2> newOutputBuffers;
-  allocateBuffersForResults(b, loc, op, newOutputBuffers, bvm);
+  if (failed(allocateBuffersForResults(b, loc, op, newOutputBuffers, bvm)))
+    return failure();
   finalizeBufferAllocation(b, op, newInputBuffers, newOutputBuffers, bvm);
   return success();
 }
 
 /// DimOp tensor operand is modified inplace. This allows leaving dead tensors
 /// behind that will get DCE'd.
 static LogicalResult bufferize(OpBuilder &b, memref::DimOp dimOp,
                                BlockAndValueMapping &bvm) {
-  if (dimOp.memrefOrTensor().getType().isa<RankedTensorType>())
-    dimOp.memrefOrTensorMutable().assign(lookup(bvm, dimOp.memrefOrTensor()));
+  if (dimOp.memrefOrTensor().getType().isa<RankedTensorType>()) {
+    Value v = lookup(bvm, dimOp.memrefOrTensor());
+    if (!v)
+      return failure();
+    dimOp.memrefOrTensorMutable().assign(v);
+  }
   return success();
 }
 
 /// FuncOp always creates TensorToMemRef ops.
 static LogicalResult bufferize(OpBuilder &b, FuncOp funcOp,
                                BlockAndValueMapping &bvm) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
static LogicalResult bufferize(OpBuilder &b, ReturnOp returnOp,
   b.setInsertionPoint(returnOp);
 
   FuncOp funcOp = cast<FuncOp>(returnOp->getParentOp());
   assert(funcOp && "only support FuncOp parent for ReturnOp");
   for (OpOperand &operand : returnOp->getOpOperands()) {
     auto tensorType = operand.get().getType().dyn_cast<TensorType>();
     if (!tensorType)
       continue;
-    operand.set(b.create<memref::TensorLoadOp>(returnOp.getLoc(),
-                                               lookup(bvm, operand.get())));
+    Value v = lookup(bvm, operand.get());
+    if (!v)
+      return failure();
+    operand.set(b.create<memref::TensorLoadOp>(returnOp.getLoc(), v));
   }
   return success();
 }
 
 static LogicalResult bufferize(OpBuilder &b,
                                SubTensorInsertOp subTensorInsertOp,
                                BlockAndValueMapping &bvm) {
   LLVM_DEBUG(DBGS() << "convert: " << *subTensorInsertOp << "\n");
 
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(subTensorInsertOp);
   Location loc = subTensorInsertOp.getLoc();
 
   Value dstMemref;
   auto inPlace = getInPlace(subTensorInsertOp->getResult(0));
   if (inPlace != InPlaceSpec::True) {
     dstMemref = createNewAllocDeallocPairForShapedValue(
         b, loc, subTensorInsertOp.result());
     b.setInsertionPointAfter(dstMemref.getDefiningOp());
     map(bvm, subTensorInsertOp.result(), dstMemref);
   } else {
     dstMemref = lookup(bvm, subTensorInsertOp.dest());
+    if (!dstMemref)
+      return failure();
   }
   auto dstMemrefType = dstMemref.getType().cast<MemRefType>();
 
   Value srcMemref = lookup(bvm, subTensorInsertOp.source());
+  if (!srcMemref)
+    return failure();
   auto subviewMemRefType =
       memref::SubViewOp::inferRankReducedResultType(
           subTensorInsertOp.getSourceType().getRank(), dstMemrefType,
           subTensorInsertOp.getMixedOffsets(),
           subTensorInsertOp.getMixedSizes(),
           subTensorInsertOp.getMixedStrides())
           .cast<MemRefType>();
 
static LogicalResult bufferize(OpBuilder &b, VectorTransferOpInterface op,
 
   if (op.getShapedType().isa<MemRefType>())
     return failure();
 
   LLVM_DEBUG(DBGS() << "convert: " << *op << "\n");
 
   /// transfer_read from buffer
   if (auto readOp = dyn_cast<vector::TransferReadOp>(op.getOperation())) {
-    readOp.sourceMutable().assign(lookup(bvm, op.source()));
+    Value v = lookup(bvm, op.source());
+    if (!v)
+      return failure();
+    readOp.sourceMutable().assign(lookup(bvm, v));
     return success();
   }
 
   auto inPlace = getInPlace(op->getResult(0));
   auto writeOp = cast<vector::TransferWriteOp>(op.getOperation());
 
   // If transfer_write is not inPlace, allocate a new buffer.
   Value newInputBuffer;
   if (inPlace != InPlaceSpec::True) {
     newInputBuffer =
         createNewAllocDeallocPairForShapedValue(b, loc, writeOp.result());
     b.setInsertionPointAfter(newInputBuffer.getDefiningOp());
     map(bvm, writeOp.result(), newInputBuffer);
   } else {
     // InPlace write will result in memref.tensor_load(x) which must
     // canonicalize away with one of it uses.
     newInputBuffer = lookup(bvm, writeOp.source());
+    if (!newInputBuffer)
+      return failure();
   }
 
   // Create a new transfer_write on buffer that doesn't have a return value.
   // Leave the previous transfer_write to dead code as it still has uses at
   // this point.
   b.create<vector::TransferWriteOp>(
       loc, writeOp.vector(), newInputBuffer, writeOp.indices(),
       writeOp.permutation_map(),
static LogicalResult bufferizeFuncOpInternals(
   if (failed(bufferize(b, funcOp, bvm)))
     return failure();
   WalkResult result = funcOp.walk<WalkOrder::PreOrder>([&](Operation *op) {
     LogicalResult status =
         llvm::TypeSwitch<Operation *, LogicalResult>(op)
             // Skip BufferCast and TensorLoad ops.
             .Case<memref::BufferCastOp, memref::TensorLoadOp>(
                 [&](auto) { return success(); })
-            .Case<memref::DimOp, LinalgOp, ReturnOp,
-                  SubTensorInsertOp, VectorTransferOpInterface>(
+            .Case<memref::DimOp, LinalgOp, ReturnOp, SubTensorInsertOp,
+                  VectorTransferOpInterface>(
                 [&](auto op) { return bufferize(b, op, bvm); })
             .Default([&](Operation *op) {
               auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
               if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
                   llvm::any_of(op->getResultTypes(), isaTensor))
                 return failure();
               return success();
             });