[mlir][bufferization] Simplify `getBufferType`

[springerm]

Summary

getBufferType computes the bufferized type of an SSA value without bufferizing any IR. This is useful for predicting the bufferized type of iter_args of a loop.

To avoid endless recursion (e.g., in the case of "scf.for", the type of the iter_arg depends on the type of init_arg and the type of the yielded value; the type of the yielded value depends on the type of the iter_arg again), fixedTypes was used to fall back to "fixed" type. A simpler way is to maintain an "invocation stack". getBufferType implementations can then inspect the invocation stack to detect repetitive computations (typically when computing the bufferized type of a block argument).

[maerhart]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

623

Nit: potentially a more generic type like SmallVectorImpl could be used?

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

742–744

This getBufferType function has the invariant that the passed value is not already pushed to the stack, but the getBufferType method of the OpInterface uses the invariant that it is already pushed to the stack. Couldn't we make this more consistent? Or document it better?

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

385–391

I don't quite understand why both of them can be casted to a ranked memref/tensor. Couldn't they also be unranked?

[springerm]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

623

I've been using SmallVector throughout the bufferization (e.g., AliasingOpOperandList), so keeping that for consistency.

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

742–744

That's so that users do not have to bother with invocationStack in most cases. Added some documentation that the interface method should not be called from user code.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 /// cases. However, when a buffer type should be predicted without modifying any
 /// IR, this function can be used.
 ///
 /// This function is a wrapper around BufferizableOpInterface::getBufferType.
 FailureOr<BaseMemRefType> getBufferType(Value value,
                                         const BufferizationOptions &options);
 
 /// Return the buffer type for a given Value (tensor) after bufferization
-/// without bufferizing any IR. If at any point during the type computation, the
-/// type of a value in `fixedTypes` in required, the mapped type is used.
+/// without bufferizing any IR. This function (and not the other overload
+/// without `invocationStack`) can be used from `getBufferType` implementations
+/// of the `BufferizableOpInterface`.
 ///
 /// Note: It should be sufficient to call `getBuffer()->getType()` in most
 /// cases. However, when a buffer type should be predicted without modifying any
 /// IR, this function can be used.
 ///
-/// This function is a wrapper around BufferizableOpInterface::getBufferType.
-FailureOr<BaseMemRefType>
-getBufferType(Value value, const BufferizationOptions &options,
-              const DenseMap<Value, BaseMemRefType> &fixedTypes);
+/// This function is a wrapper around `BufferizableOpInterface::getBufferType`.
+FailureOr<BaseMemRefType> getBufferType(Value value,
+                                        const BufferizationOptions &options,
+                                        SmallVector<Value> &invocationStack);
 
 /// Replace an op with replacement values. The op is deleted. Tensor OpResults
 /// must be replaced with memref values.
 void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
                                    ValueRange values);
 
 /// Replace an op with a new op. The new op must have the same number of
 /// results as the replaced op. The new op may not return any tensor values.

 AliasingOpOperandList defaultGetAliasingOpOperands(Value value,
                                                    const AnalysisState &state);
 
 /// This is the default implementation of
 /// BufferizableOpInterface::getBufferType. Should not be called from other
 /// places.
 FailureOr<BaseMemRefType>
 defaultGetBufferType(Value value, const BufferizationOptions &options,
-                     const DenseMap<Value, BaseMemRefType> &fixedTypes);
+                     SmallVector<Value> &invocationStack);
 
 /// This is the default implementation of
 /// BufferizableOpInterface::resultBufferizesToMemoryWrite. Should not be called
 /// from other places.
 bool defaultResultBufferizesToMemoryWrite(OpResult opResult,
                                           const AnalysisState &state);
 
 /// This is the default implementation of

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
       InterfaceMethod<
         /*desc=*/[{
           Return the bufferized type of the given tensor value (without
           bufferizing the IR). The value is either a BlockArgument of a block
           that belongs to this op or an OpResult of the given op.
 
           This method is useful when the bufferized type of value must be
           predicted before modifying any IR.
+
+          Implementations may call `bufferization::getBufferType` to compute the
+          bufferized type of another SSA value. The same (unmodified)
+          `invocationStack` must be passed to that function. The stack contains
+          all SSA values for which a buffer type computation is currently in
+          progress. Implementations may inspect the stack to detect repetitive
+          computations for the same SSA value. (E.g., when bufferized types of a
+          loop.)
+
+          Note: This interface method should never be called directly from user
+          code. Always use `bufferization::getBufferType`.
         }],
         /*retType=*/"::mlir::FailureOr<::mlir::BaseMemRefType>",
         /*methodName=*/"getBufferType",
         /*args=*/(ins "::mlir::Value":$value,
                       "const ::mlir::bufferization::BufferizationOptions &":$options,
-                      "const ::mlir::DenseMap<::mlir::Value, ::mlir::BaseMemRefType>":$fixedTypes),
+                      "::llvm::SmallVector<::mlir::Value> &":$invocationStack),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           assert(getOwnerOfValue(value) == $_op.getOperation() &&
                  "expected that value belongs to this op");
+          assert(invocationStack.back() == value &&
+                 "inconsistant invocation stack");
           return ::mlir::bufferization::detail::defaultGetBufferType(
-              value, options, fixedTypes);
+              value, options, invocationStack);
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the given region of this op is repetitive. By default
           this information is queried from the `RegionBranchOpInterface`. Ops
           that do not implement this inferface can override this method to
           declare regions as repetitive.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

let extraClassDeclaration = [{
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
     AliasingValueList getAliasingValues(
         OpOperand &opOperand, const AnalysisState &state);
 
     FailureOr<BaseMemRefType> getBufferType(
         Value value, const BufferizationOptions &options,
-        const DenseMap<Value, BaseMemRefType> &fixedTypes);
+        SmallVector<Value> &invocationStack);
 
     RankedTensorType getType() {
       return ::llvm::cast<RankedTensorType>(getResult().getType());
     }
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicDim(unsigned idx) {
       return getType().isDynamicDim(idx);
LogicalResult bufferize(RewriterBase &rewriter,
       // to_tensor/to_memref pairs fold away after bufferization.
       return success();
     }
 
     bool isWritable(Value value, const AnalysisState &state);
 
     FailureOr<BaseMemRefType> getBufferType(
         Value value, const BufferizationOptions &options,
-        const DenseMap<Value, BaseMemRefType> &fixedTypes) {
+        SmallVector<Value> &invocationStack) {
       return ::llvm::cast<BaseMemRefType>(getMemref().getType());
     }
   }];
 
   let assemblyFormat = [{
     $memref (`restrict` $restrict^)? (`writable` $writable^)? attr-dict
       `:` type($memref)
   }];

mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     replaceOpWithNewBufferizedOp<arith::SelectOp>(
         rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto selectOp = cast<arith::SelectOp>(op);
     assert(value == selectOp.getResult() && "invalid value");
     auto trueType = bufferization::getBufferType(selectOp.getTrueValue(),
-                                                 options, fixedTypes);
+                                                 options, invocationStack);
     auto falseType = bufferization::getBufferType(selectOp.getFalseValue(),
-                                                  options, fixedTypes);
+                                                  options, invocationStack);
     if (failed(trueType) || failed(falseType))
       return failure();
     if (*trueType == *falseType)
       return *trueType;
     if (trueType->getMemorySpace() != falseType->getMemorySpace())
       return op->emitError("inconsistent memory space on true/false operands");
 
     // If the buffers have different types, they differ only in their layout

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 
 //===----------------------------------------------------------------------===//
 // BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace bufferization {
#endif // NDEBUG
   return rewriter
       .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
 /// Return the buffer type for a given Value (tensor) after bufferization.
 FailureOr<BaseMemRefType>
 bufferization::getBufferType(Value value, const BufferizationOptions &options) {
-  DenseMap<Value, BaseMemRefType> fixedTypes;
-  return getBufferType(value, options, fixedTypes);
+  SmallVector<Value> invocationStack;
+  return getBufferType(value, options, invocationStack);
 }
 
 /// Return the buffer type for a given Value (tensor) after bufferization.
-FailureOr<BaseMemRefType> bufferization::getBufferType(
-    Value value, const BufferizationOptions &options,
-    const DenseMap<Value, BaseMemRefType> &fixedTypes) {
+FailureOr<BaseMemRefType>
+bufferization::getBufferType(Value value, const BufferizationOptions &options,
+                             SmallVector<Value> &invocationStack) {
   assert(llvm::isa<TensorType>(value.getType()) &&
          "unexpected non-tensor type");
-
-  // If the `value` is in `fixedTypes`, return the mapped type.
-  const auto &it = fixedTypes.find(value);
-  if (it != fixedTypes.end())
-    return it->second;
+  invocationStack.push_back(value);
+  auto popFromStack =
+      llvm::make_scope_exit([&]() { invocationStack.pop_back(); });
 
   // Try querying BufferizableOpInterface.
   Operation *op = getOwnerOfValue(value);
   auto bufferizableOp = options.dynCastBufferizableOp(op);
   if (bufferizableOp)
-    return bufferizableOp.getBufferType(value, options, fixedTypes);
+    return bufferizableOp.getBufferType(value, options, invocationStack);
 
   // Op is not bufferizable.
   if (!options.defaultMemorySpace.has_value())
     return op->emitError("could not infer memory space");
 
   return getMemRefType(value, options, /*layout=*/{},
                        *options.defaultMemorySpace);
 }
for (const auto &it : aliasingValues)
       if (it.value == value)
         result.emplace_back(&opOperand, it.relation, it.isDefinite);
   }
   return AliasingOpOperandList(std::move(result));
 }
 
 FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
     Value value, const BufferizationOptions &options,
-    const DenseMap<Value, BaseMemRefType> &fixedTypes) {
+    SmallVector<Value> &invocationStack) {
   assert(llvm::isa<TensorType>(value.getType()) && "expected tensor type");
 
   // No further analysis is possible for a block argument.
   if (llvm::isa<BlockArgument>(value))
     return bufferization::getMemRefType(value, options);
 
   // Value is an OpResult.
   Operation *op = getOwnerOfValue(value);
   auto opResult = llvm::cast<OpResult>(value);
   AnalysisState state(options);
   AliasingOpOperandList aliases = state.getAliasingOpOperands(opResult);
   if (aliases.getNumAliases() > 0 &&
       aliases.getAliases()[0].relation == BufferRelation::Equivalent) {
     // If the OpResult has an equivalent OpOperand, both OpResult and
     // OpOperand bufferize to the exact same buffer type.
     Value equivalentOperand = aliases.getAliases().front().opOperand->get();
-    return getBufferType(equivalentOperand, options, fixedTypes);
+    return getBufferType(equivalentOperand, options, invocationStack);
   }
 
   // If we do not know the memory space and there is no default memory space,
   // report a failure.
   if (!options.defaultMemorySpace.has_value())
     return op->emitError("could not infer memory space");
 
   return getMemRefType(value, options, /*layout=*/{},

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

 }
 
 AliasingValueList AllocTensorOp::getAliasingValues(OpOperand &opOperand,
                                                    const AnalysisState &state) {
   // This is a new allocation. It does not alias with any other buffer.
   return {};
 }
 
-FailureOr<BaseMemRefType> AllocTensorOp::getBufferType(
-    Value value, const BufferizationOptions &options,
-    const DenseMap<Value, BaseMemRefType> &fixedTypes) {
+FailureOr<BaseMemRefType>
+AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options,
+                             SmallVector<Value> &invocationStack) {
   assert(value == getResult() && "invalid value");
 
   // Compute memory space of this allocation.
   Attribute memorySpace;
   if (getMemorySpace().has_value()) {
     memorySpace = *getMemorySpace();
   } else if (getCopy()) {
     auto copyBufferType =
-        bufferization::getBufferType(getCopy(), options, fixedTypes);
+        bufferization::getBufferType(getCopy(), options, invocationStack);
     if (failed(copyBufferType))
       return failure();
     memorySpace = copyBufferType->getMemorySpace();
   } else if (options.defaultMemorySpace.has_value()) {
     memorySpace = *options.defaultMemorySpace;
   } else {
     return getOperation()->emitError("could not infer memory space");
   }

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

for (int64_t resultIdx : aliasingReturnVals)
                        equivalent.has_value() ? BufferRelation::Equivalent
                                               : BufferRelation::Unknown,
                        /*isDefinite=*/equivalent.has_value()});
     return result;
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
 
     // The callee was already bufferized, so we can directly take the type from
     // its signature.
     FunctionType funcType = funcOp.getFunctionType();
     return cast<BaseMemRefType>(
#endif // NDEBUG
     return success();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto funcOp = cast<FuncOp>(op);
     auto bbArg = cast<BlockArgument>(value);
     // Unstructured control flow is not supported.
     assert(bbArg.getOwner() == &funcOp.getBody().front() &&
            "expected that block argument belongs to first block");
     return getBufferizedFunctionArgType(funcOp, bbArg.getArgNumber(), options);
   }
 

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Replace op results.
     replaceOpWithBufferizedValues(rewriter, op, newIfOp->getResults());
 
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto ifOp = cast<scf::IfOp>(op);
     auto thenYieldOp = cast<scf::YieldOp>(ifOp.thenBlock()->getTerminator());
     auto elseYieldOp = cast<scf::YieldOp>(ifOp.elseBlock()->getTerminator());
     assert(value.getDefiningOp() == op && "invalid valid");
 
     // Determine buffer types of the true/false branches.
     auto opResult = cast<OpResult>(value);
     auto thenValue = thenYieldOp.getOperand(opResult.getResultNumber());
     auto elseValue = elseYieldOp.getOperand(opResult.getResultNumber());
     BaseMemRefType thenBufferType, elseBufferType;
     if (isa<BaseMemRefType>(thenValue.getType())) {
       // True branch was already bufferized.
       thenBufferType = cast<BaseMemRefType>(thenValue.getType());
     } else {
       auto maybeBufferType =
-          bufferization::getBufferType(thenValue, options, fixedTypes);
+          bufferization::getBufferType(thenValue, options, invocationStack);
       if (failed(maybeBufferType))
         return failure();
       thenBufferType = *maybeBufferType;
     }
     if (isa<BaseMemRefType>(elseValue.getType())) {
       // False branch was already bufferized.
       elseBufferType = cast<BaseMemRefType>(elseValue.getType());
     } else {
       auto maybeBufferType =
-          bufferization::getBufferType(elseValue, options, fixedTypes);
+          bufferization::getBufferType(elseValue, options, invocationStack);
       if (failed(maybeBufferType))
         return failure();
       elseBufferType = *maybeBufferType;
     }
 
     // Best case: Both branches have the exact same buffer type.
     if (thenBufferType == elseBufferType)
       return thenBufferType;

 }
 
 /// Compute the bufferized type of a loop iter_arg. This type must be equal to
 /// the bufferized type of the corresponding init_arg and the bufferized type
 /// of the corresponding yielded value.
 ///
 /// This function uses bufferization::getBufferType to compute the bufferized
 /// type of the init_arg and of the yielded value. (The computation of the
-/// usually requires computing the bufferized type of the corresponding
-/// iter_arg; the implementation of getBufferType traces back the use-def chain
-/// of the given value and computes a buffer type along the way.) If both buffer
-/// types are equal, no casts are needed the computed buffer type can be used
-/// directly. Otherwise, the buffer types can only differ in their layout map
-/// and a cast must be inserted.
+/// bufferized yielded value type usually requires computing the bufferized type
+/// of the iter_arg again; the implementation of getBufferType traces back the
+/// use-def chain of the given value and computes a buffer type along the way.)
+/// If both buffer types are equal, no casts are needed the computed buffer type
+/// can be used directly. Otherwise, the buffer types can only differ in their
+/// layout map and a cast must be inserted.
 static FailureOr<BaseMemRefType> computeLoopRegionIterArgBufferType(
-    BlockArgument iterArg, Value initArg, Value yieldedValue,
-    const BufferizationOptions &options,
-    const DenseMap<Value, BaseMemRefType> &fixedTypes) {
+    Operation *loopOp, BlockArgument iterArg, Value initArg, Value yieldedValue,
+    const BufferizationOptions &options, SmallVector<Value> &invocationStack) {
   // Determine the buffer type of the init_arg.
   auto initArgBufferType =
-      bufferization::getBufferType(initArg, options, fixedTypes);
+      bufferization::getBufferType(initArg, options, invocationStack);
   if (failed(initArgBufferType))
     return failure();
 
-  // Fix the iter_arg type, so that recursive lookups return the buffer type
-  // of the init_arg. This is to avoid infinite loops when calculating the
-  // buffer type of the yielded value.
-  //
+  if (llvm::count(invocationStack, iterArg) >= 2) {
+    // If the iter_arg is already twice on the invocation stack, just take the
+    // type of the init_arg. This is to avoid infinite loops when calculating
+    // the buffer type. This will most likely result in computing a memref type
+    // with a fully dynamic layout map.
+
     // Note: For more precise layout map computation, a fixpoint iteration could
     // be done (i.e., re-computing the yielded buffer type until the bufferized
     // iter_arg type no longer changes). This current implementation immediately
     // switches to a fully dynamic layout map when a mismatch between bufferized
     // init_arg type and bufferized yield value type is detected.
-  DenseMap<Value, BaseMemRefType> newFixedTypes(fixedTypes);
-  newFixedTypes[iterArg] = *initArgBufferType;
+    return *initArgBufferType;
+  }
 
   // Compute the buffer type of the yielded value.
   BaseMemRefType yieldedValueBufferType;
   if (isa<BaseMemRefType>(yieldedValue.getType())) {
     // scf.yield was already bufferized.
     yieldedValueBufferType = cast<BaseMemRefType>(yieldedValue.getType());
   } else {
+    // Note: This typically triggers a recursive call for the buffer type of
+    // the iter_arg.
     auto maybeBufferType =
-        bufferization::getBufferType(yieldedValue, options, newFixedTypes);
+        bufferization::getBufferType(yieldedValue, options, invocationStack);
     if (failed(maybeBufferType))
       return failure();
     yieldedValueBufferType = *maybeBufferType;
   }
 
   // If yielded type and init_arg type are the same, use that type directly.
   if (*initArgBufferType == yieldedValueBufferType)
     return yieldedValueBufferType;
 
-  // If there is a mismatch between the yielded buffer type and the iter_arg
+  // If there is a mismatch between the yielded buffer type and the init_arg
   // buffer type, the buffer type must be promoted to a fully dynamic layout
   // map.
-  auto yieldedRanked = cast<MemRefType>(yieldedValueBufferType);
+  auto yieldedBufferType = cast<BaseMemRefType>(yieldedValueBufferType);
+  auto iterTensorType = cast<TensorType>(iterArg.getType());
+  auto initBufferType = llvm::cast<BaseMemRefType>(*initArgBufferType);
+  if (initBufferType.getMemorySpace() != yieldedBufferType.getMemorySpace())
+    return loopOp->emitOpError(
+        "init_arg and yielded value bufferize to inconsistent memory spaces");
 #ifndef NDEBUG
-  auto iterRanked = llvm::cast<MemRefType>(*initArgBufferType);
-  assert(llvm::equal(yieldedRanked.getShape(), iterRanked.getShape()) &&
-         "expected same shape");
-  assert(yieldedRanked.getMemorySpace() == iterRanked.getMemorySpace() &&
-         "expected same memory space");
+  if (auto yieldedRankedBufferType = dyn_cast<MemRefType>(yieldedBufferType)) {
+    assert(
+        llvm::all_equal({yieldedRankedBufferType.getShape(),
+                         cast<MemRefType>(initBufferType).getShape(),
+                         cast<RankedTensorType>(iterTensorType).getShape()}) &&
+  }
 #endif // NDEBUG
   return getMemRefTypeWithFullyDynamicLayout(
-      cast<RankedTensorType>(iterArg.getType()),
-      yieldedRanked.getMemorySpace());
+      iterTensorType, yieldedBufferType.getMemorySpace());
 }
 
 /// Return `true` if the given loop may have 0 iterations.
 bool mayHaveZeroIterations(scf::ForOp forOp) {
   std::optional<int64_t> lb = getConstantIntValue(forOp.getLowerBound());
   std::optional<int64_t> ub = getConstantIntValue(forOp.getUpperBound());
   if (!lb.has_value() || !ub.has_value())
     return true;
LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
 
     rewriter.updateRootInPlace(
         yieldOp, [&]() { yieldOp.getResultsMutable().assign(yieldValues); });
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto forOp = cast<scf::ForOp>(op);
     assert(getOwnerOfValue(value) == op && "invalid value");
     assert(isa<TensorType>(value.getType()) && "expected tensor type");
 
-    // Get result/argument number.
-    unsigned resultNum;
-    if (auto bbArg = dyn_cast<BlockArgument>(value)) {
-      resultNum =
-          forOp.getResultForOpOperand(forOp.getOpOperandForRegionIterArg(bbArg))
-              .getResultNumber();
-    } else {
-      resultNum = cast<OpResult>(value).getResultNumber();
-    }
+    if (auto opResult = dyn_cast<OpResult>(value)) {
+      // The type of an OpResult must match the corresponding iter_arg type.
+      BlockArgument bbArg = forOp.getRegionIterArgForOpOperand(
+          forOp.getOpOperandForResult(opResult));
+      return bufferization::getBufferType(bbArg, options, invocationStack);
+    }
+
+    // Compute result/argument number.
+    BlockArgument bbArg = cast<BlockArgument>(value);
+    unsigned resultNum =
 
     // Compute the bufferized type.
     auto yieldOp =
         cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
     Value yieldedValue = yieldOp.getOperand(resultNum);
     BlockArgument iterArg = forOp.getRegionIterArgs()[resultNum];
     Value initArg = forOp.getInitArgs()[resultNum];
-    return computeLoopRegionIterArgBufferType(iterArg, initArg, yieldedValue,
-                                              options, fixedTypes);
+    return computeLoopRegionIterArgBufferType(
+        op, iterArg, initArg, yieldedValue, options, invocationStack);
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto forOp = cast<scf::ForOp>(op);
     Block *oldLoopBody = &forOp.getLoopBody().front();
 
     // Indices of all iter_args that have tensor type. These are the ones that
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newWhileOp->getResults());
 
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto whileOp = cast<scf::WhileOp>(op);
     assert(getOwnerOfValue(value) == op && "invalid value");
     assert(isa<TensorType>(value.getType()) && "expected tensor type");
 
     // Case 1: Block argument of the "before" region.
     if (auto bbArg = dyn_cast<BlockArgument>(value)) {
       if (bbArg.getOwner()->getParent() == &whileOp.getBefore()) {
         Value initArg = whileOp.getInits()[bbArg.getArgNumber()];
         auto yieldOp = whileOp.getYieldOp();
         Value yieldedValue = yieldOp.getOperand(bbArg.getArgNumber());
-        return computeLoopRegionIterArgBufferType(bbArg, initArg, yieldedValue,
-                                                  options, fixedTypes);
+        return computeLoopRegionIterArgBufferType(
+            op, bbArg, initArg, yieldedValue, options, invocationStack);
       }
     }
 
     // Case 2: OpResult of the loop or block argument of the "after" region.
     // The bufferized "after" bbArg type can be directly computed from the
     // bufferized "before" bbArg type.
     unsigned resultNum;
     if (auto opResult = dyn_cast<OpResult>(value)) {
       resultNum = opResult.getResultNumber();
     } else if (cast<BlockArgument>(value).getOwner()->getParent() ==
                &whileOp.getAfter()) {
       resultNum = cast<BlockArgument>(value).getArgNumber();
     } else {
       llvm_unreachable("invalid value");
     }
     Value conditionYieldedVal = whileOp.getConditionOp().getArgs()[resultNum];
     if (!isa<TensorType>(conditionYieldedVal.getType())) {
       // scf.condition was already bufferized.
       return cast<BaseMemRefType>(conditionYieldedVal.getType());
     }
     return bufferization::getBufferType(conditionYieldedVal, options,
-                                        fixedTypes);
+                                        invocationStack);
   }
 
   /// Assert that yielded values of an scf.while op are equivalent to their
   /// corresponding bbArgs. In that case, the buffer relations of the
   /// corresponding OpResults are "Equivalent".
   ///
   /// If this is not the case, allocs+copies are inserted and yielded from
   /// the loop. This could be a performance problem, so it must be explicitly
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Remove the old op and replace all of its uses.
     replaceOpWithBufferizedValues(rewriter, op, buffers);
 
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto forallOp = cast<ForallOp>(op);
 
     if (auto bbArg = dyn_cast<BlockArgument>(value))
       // A tensor block argument has the same bufferized type as the
       // corresponding output operand.
       return bufferization::getBufferType(
-          forallOp.getTiedOpOperand(bbArg)->get(), options, fixedTypes);
+          forallOp.getTiedOpOperand(bbArg)->get(), options, invocationStack);
 
     // The bufferized result type is the same as the bufferized type of the
     // corresponding output operand.
     return bufferization::getBufferType(
         forallOp.getOutputs()[cast<OpResult>(value).getResultNumber()], options,
-        fixedTypes);
+        invocationStack);
   }
 
   bool isRepetitiveRegion(Operation *op, unsigned index) const {
     auto forallOp = cast<ForallOp>(op);
 
     // This op is repetitive if it has 1 or more steps.
     // If the control variables are dynamic, it is also considered so.
     for (auto [lb, ub, step] :

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

struct CastOpInterface
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {{op->getResult(0), BufferRelation::Equivalent}};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto castOp = cast<tensor::CastOp>(op);
-    auto maybeSrcBufferType =
-        bufferization::getBufferType(castOp.getSource(), options, fixedTypes);
+    auto maybeSrcBufferType = bufferization::getBufferType(
+        castOp.getSource(), options, invocationStack);
     if (failed(maybeSrcBufferType))
       return failure();
     Attribute memorySpace = maybeSrcBufferType->getMemorySpace();
 
     // Note: `getMemRefTypeWithFullyDynamicLayout` returns an unranked memref
     // type in case the input is an unranked tensor type.
 
     // Case 1: Casting an unranked tensor
struct CollapseShapeOpInterface
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     // TODO: CollapseShapeOp may allocate at runtime.
     return {{op->getOpResult(0), BufferRelation::Equivalent}};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     auto maybeSrcBufferType = bufferization::getBufferType(
-        collapseShapeOp.getSrc(), options, fixedTypes);
+        collapseShapeOp.getSrc(), options, invocationStack);
     if (failed(maybeSrcBufferType))
       return failure();
     auto srcBufferType = llvm::cast<MemRefType>(*maybeSrcBufferType);
     bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
         srcBufferType, collapseShapeOp.getReassociationIndices());
 
     if (!canBeCollapsed) {
       // If dims cannot be collapsed, this op bufferizes to a new allocation.
struct ExpandShapeOpInterface
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {{op->getOpResult(0), BufferRelation::Equivalent}};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
     auto maybeSrcBufferType = bufferization::getBufferType(
-        expandShapeOp.getSrc(), options, fixedTypes);
+        expandShapeOp.getSrc(), options, invocationStack);
     if (failed(maybeSrcBufferType))
       return failure();
     auto srcBufferType = llvm::cast<MemRefType>(*maybeSrcBufferType);
     auto maybeResultType = memref::ExpandShapeOp::computeExpandedType(
         srcBufferType, expandShapeOp.getResultType().getShape(),
         expandShapeOp.getReassociationIndices());
     if (failed(maybeResultType))
       return failure();
Value subView = rewriter.create<memref::SubViewOp>(
         mixedSizes, mixedStrides);
 
     replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     assert(value == extractSliceOp.getResult() && "invalid value");
     auto srcMemrefType = bufferization::getBufferType(
-        extractSliceOp.getSource(), options, fixedTypes);
+        extractSliceOp.getSource(), options, invocationStack);
     if (failed(srcMemrefType))
       return failure();
     SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
     SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
     return cast<BaseMemRefType>(memref::SubViewOp::inferRankReducedResultType(
         extractSliceOp.getType().getShape(), llvm::cast<MemRefType>(*srcMemrefType),
         mixedOffsets, mixedSizes, mixedStrides));
struct PadOpInterface
 
   AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                       const AnalysisState &state) const {
     return {};
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     // Infer memory space from the source tensor.
     auto padOp = cast<tensor::PadOp>(op);
-    auto maybeSrcBufferType =
-        bufferization::getBufferType(padOp.getSource(), options, fixedTypes);
+    auto maybeSrcBufferType = bufferization::getBufferType(
+        padOp.getSource(), options, invocationStack);
     if (failed(maybeSrcBufferType))
       return failure();
     MemRefLayoutAttrInterface layout;
     return MemRefType::get(padOp.getResultType().getShape(),
                            padOp.getResultType().getElementType(), layout,
                            maybeSrcBufferType->getMemorySpace());
   }
 
if (failed(maybeResultMemRefType))
       return failure();
     replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
         rewriter, op, maybeResultMemRefType.value(), *srcBuffer, *shapeBuffer);
     return success();
   }
 
   FailureOr<BaseMemRefType>
   getBufferType(Operation *op, Value value, const BufferizationOptions &options,
-                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
+                SmallVector<Value> &invocationStack) const {
     auto reshapeOp = cast<tensor::ReshapeOp>(op);
     assert(value == reshapeOp.getResult() && "unexpected value provided");
     auto maybeSourceBufferType = bufferization::getBufferType(
-        reshapeOp.getSource(), options, fixedTypes);
+        reshapeOp.getSource(), options, invocationStack);
     if (failed(maybeSourceBufferType))
       return failure();
     return getMemRefTypeWithStaticIdentityLayout(
         reshapeOp.getResult().getType(),
         cast<BaseMemRefType>(maybeSourceBufferType.value()).getMemorySpace());
   }
 };
 

mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir

func.func @execute_region_multiple_blocks(%t: tensor<5xf32>) -> tensor<5xf32> {
   // expected-error @below{{op or BufferizableOpInterface implementation does not support unstructured control flow, but at least one region has multiple blocks}}
   %0 = scf.execute_region -> tensor<5xf32> {
     cf.br ^bb1(%t : tensor<5xf32>)
   ^bb1(%arg1 : tensor<5xf32>):
     scf.yield %arg1 : tensor<5xf32>
   }
   func.return %0 : tensor<5xf32>
 }
+
+// -----
+
+func.func @inconsistent_memory_space_scf_for(%lb: index, %ub: index, %step: index) -> tensor<10xf32> {
+  %0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
+  %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
+  // expected-error @below{{init_arg and yielded value bufferize to inconsistent memory spaces}}
+  %2 = scf.for %iv = %lb to %ub step %step iter_args(%arg = %0) -> tensor<10xf32> {
+    // expected-error @below {{failed to bufferize op}}
+    scf.yield %1 : tensor<10xf32>
+  }
+  return %2 : tensor<10xf32>
+}