[mlir][SCF][bufferize] Support different iter_arg/init_arg types (scf.for)

[springerm]

Summary

Even though iter_arg and init_arg of an scf.for loop may have the same tensor type, their bufferized memref types are not necessarily equal. It is sometimes necessary to insert a cast in case of differing layout maps.

Depends On D132859

[bkramer]

lg

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

 /// 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);
 
+/// 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.
+///
+/// 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);
+
 /// 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.
 template <typename OpTy, typename... Args>

 /// owner of the block. In case of an OpResult that is the defining op.
 Operation *getOwnerOfValue(Value value);
 
 namespace detail {
 /// This is the default implementation of
 /// BufferizableOpInterface::getBufferType. Should not be called from other
 /// places.
 FailureOr<BaseMemRefType>
-defaultGetBufferType(Value value, const BufferizationOptions &options);
+defaultGetBufferType(Value value, const BufferizationOptions &options,
+                     const DenseMap<Value, BaseMemRefType> &fixedTypes);
 } // namespace detail
 
 } // namespace bufferization
 } // namespace mlir
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h.inc"
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

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

let methods = [
           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.
         }],
         /*retType=*/"FailureOr<BaseMemRefType>",
         /*methodName=*/"getBufferType",
         /*args=*/(ins "Value":$value,
-                      "const BufferizationOptions &":$options),
+                      "const BufferizationOptions &":$options,
+                      "const DenseMap<Value, BaseMemRefType>":$fixedTypes),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           assert(getOwnerOfValue(value) == $_op.getOperation() &&
                  "expected that value belongs to this op");
-          return bufferization::detail::defaultGetBufferType(value, options);
+          return bufferization::detail::defaultGetBufferType(
+              value, options, fixedTypes);
         }]
       >,
   ];
 
   let extraClassDeclaration = [{
     /// Resolve out-of-place tensor OpOperands with explicit allocations in the
     /// form of `bufferization.alloc_tensor` ops.
     LogicalResult resolveTensorOpOperandConflicts(

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

let extraClassDeclaration = [{
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
 
     SmallVector<OpResult> getAliasingOpResult(
         OpOperand &opOperand, const AnalysisState &state);
 
     FailureOr<BaseMemRefType> getBufferType(
-        Value value, const BufferizationOptions &options);
+        Value value, const BufferizationOptions &options,
+        const DenseMap<Value, BaseMemRefType> &fixedTypes);
 
     RankedTensorType getType() {
       return getResult().getType().cast<RankedTensorType>();
     }
 
     // Return true if the size of the tensor is dynamic at `idx`
     bool isDynamicDim(unsigned idx) {
       return getType().isDynamicDim(idx);
let extraClassDeclaration = [{
     }
 
     bool isWritable(Value value, const AnalysisState &state) const {
       // It is unknown whether the memref operand is writable or not.
       return false;
     }
 
     FailureOr<BaseMemRefType> getBufferType(
-        Value value, const BufferizationOptions &options) {
+        Value value, const BufferizationOptions &options,
+        const DenseMap<Value, BaseMemRefType> &fixedTypes) {
       return getMemref().getType().cast<BaseMemRefType>();
     }
   }];
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;

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

if (failed(memrefType))
     return failure();
   ensureToMemrefOpIsValid(value, *memrefType);
   return rewriter
       .create<bufferization::ToMemrefOp>(value.getLoc(), *memrefType, value)
       .getResult();
 }
 
 FailureOr<BaseMemRefType> bufferization::detail::defaultGetBufferType(
-    Value value, const BufferizationOptions &options) {
+    Value value, const BufferizationOptions &options,
+    const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value.getType().isa<TensorType>() && "expected tensor type");
 
   // No further analysis is possible for a block argument.
   if (value.isa<BlockArgument>())
     return bufferization::getMemRefType(value, options);
 
   // Value is an OpResult.
   Operation *op = getOwnerOfValue(value);
   auto opResult = value.cast<OpResult>();
   auto bufferizableOp = cast<BufferizableOpInterface>(op);
   AnalysisState state(options);
   auto aliasingOperands = bufferizableOp.getAliasingOpOperand(opResult, state);
   if (!aliasingOperands.empty() &&
       bufferizableOp.bufferRelation(opResult, state) ==
           BufferRelation::Equivalent) {
     // If the OpResult has an equivalent OpOperand, both OpResult and
     // OpOperand bufferize to the exact same buffer type.
     Value equivalentOperand = aliasingOperands.front()->get();
-    return getBufferType(equivalentOperand, options);
+    return getBufferType(equivalentOperand, options, fixedTypes);
   }
 
   // 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=*/{},
                        *options.defaultMemorySpace);
 }
 
 /// 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);
+}
+
+/// 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) {
   assert(value.getType().isa<TensorType>() && "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;
+
+  // Try querying BufferizableOpInterface.
   Operation *op = getOwnerOfValue(value);
   auto bufferizableOp = options.dynCastBufferizableOp(op);
   if (bufferizableOp)
-    return bufferizableOp.getBufferType(value, options);
+    return bufferizableOp.getBufferType(value, options, fixedTypes);
 
   // Op is not bufferizable.
   if (!options.defaultMemorySpace.has_value())
     return op->emitError("could not infer memory space");
 
   return getMemRefType(value, options, /*layout=*/{},
                        *options.defaultMemorySpace);
 }

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

LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
   if (getCopy()) {
     FailureOr<Value> maybeCopyBuffer = getBuffer(rewriter, getCopy(), options);
     if (failed(maybeCopyBuffer))
       return failure();
     copyBuffer = *maybeCopyBuffer;
   }
 
   // Create memory allocation.
-  auto allocType = getBufferType(getResult(), options);
+  auto allocType = bufferization::getBufferType(getResult(), options);
   if (failed(allocType))
     return failure();
   SmallVector<Value> dynamicDims = getDynamicSizes();
   if (getCopy()) {
     assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
     populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
   }
   FailureOr<Value> alloc = options.createAlloc(

 
 SmallVector<OpResult>
 AllocTensorOp::getAliasingOpResult(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) {
+FailureOr<BaseMemRefType> AllocTensorOp::getBufferType(
+    Value value, const BufferizationOptions &options,
+    const DenseMap<Value, BaseMemRefType> &fixedTypes) {
   assert(value == getResult() && "invalid value");
 
   // Compute memory space of this allocation.
   unsigned memorySpace;
   if (getMemorySpace().has_value()) {
     memorySpace = *getMemorySpace();
   } else if (getCopy()) {
-    auto copyBufferType = bufferization::getBufferType(getCopy(), options);
+    auto copyBufferType =
+        bufferization::getBufferType(getCopy(), options, fixedTypes);
     if (failed(copyBufferType))
       return failure();
     memorySpace = copyBufferType->getMemorySpaceAsInt();
   } else if (options.defaultMemorySpace.has_value()) {
     memorySpace = *options.defaultMemorySpace;
   } else {
     return getOperation()->emitError("could not infer memory space");
   }

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

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 {
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
     auto forOp = cast<scf::ForOp>(op);
-    // TODO: Only block arguments supported at the moment.
-    if (value.isa<OpResult>())
+    assert(getOwnerOfValue(value) == op && "invalid value");
+    assert(value.getType().isa<TensorType>() && "expected tensor type");
+
+    // Get result/argument number.
+    unsigned resultNum;
+    if (auto bbArg = value.dyn_cast<BlockArgument>()) {
+      resultNum =
+          forOp.getResultForOpOperand(forOp.getOpOperandForRegionIterArg(bbArg))
+              .getResultNumber();
+    } else {
+      resultNum = value.cast<OpResult>().getResultNumber();
+    }
+
+    // Determine the buffer type of the init_arg.
+    Value initArg = forOp.getInitArgs()[resultNum];
+    auto initArgBufferType =
+        bufferization::getBufferType(initArg, options, fixedTypes);
+    if (failed(initArgBufferType))
       return failure();
-    auto bbArg = value.cast<BlockArgument>();
-    return bufferization::getBufferType(
-        forOp.getOpOperandForRegionIterArg(bbArg).get(), options);
+
+    // 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.
+    //
+    // 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[forOp.getRegionIterArgs()[resultNum]] = *initArgBufferType;
+
+    // Compute the buffer type of the yielded value.
+    auto yieldOp =
+        cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
+    Value yieldedValue = yieldOp.getOperand(resultNum);
+    BaseMemRefType yieldedValueBufferType;
+    if (yieldedValue.getType().isa<BaseMemRefType>()) {
+      // scf.yield was already bufferized.
+      yieldedValueBufferType = yieldedValue.getType().cast<BaseMemRefType>();
+    } else {
+      auto maybeBufferType =
+          bufferization::getBufferType(yieldedValue, options, newFixedTypes);
+      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
+    // buffer type, the buffer type must be promoted to a fully dynamic layout
+    // map.
+    auto yieldedRanked = yieldedValueBufferType.cast<MemRefType>();
+#ifndef NDEBUG
+    auto iterRanked = initArgBufferType->cast<MemRefType>();
+    assert(llvm::equal(yieldedRanked.getShape(), iterRanked.getShape()) &&
+           "expected same shape");
+    assert(yieldedRanked.getMemorySpaceAsInt() ==
+               iterRanked.getMemorySpaceAsInt() &&
+           "expected same memory space");
+#endif // NDEBUG
+    return getMemRefTypeWithFullyDynamicLayout(
+        forOp.getRegionIterArgs()[resultNum].getType().cast<RankedTensorType>(),
+        yieldedRanked.getMemorySpaceAsInt());
   }
 
   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
     // are bufferized.
     DenseSet<int64_t> indices = getTensorIndices(forOp.getInitArgs());
 
     // The new memref init_args of the loop.
     FailureOr<SmallVector<Value>> maybeInitArgs =
         getBuffers(rewriter, forOp.getIterOpOperands(), options);
     if (failed(maybeInitArgs))
       return failure();
     SmallVector<Value> initArgs = *maybeInitArgs;
 
+    // Cast init_args if necessary.
+    SmallVector<Value> castedInitArgs;
+    for (const auto &it : llvm::enumerate(initArgs)) {
+      Value initArg = it.value();
+      auto targetType =
+          bufferization::getBufferType(forOp->getResult(it.index()), options);
+      if (failed(targetType))
+        return failure();
+      castedInitArgs.push_back(castBuffer(rewriter, initArg, *targetType));
+    }
+
     // Construct a new scf.for op with memref instead of tensor values.
     auto newForOp = rewriter.create<scf::ForOp>(
         forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
-        forOp.getStep(), initArgs);
+        forOp.getStep(), castedInitArgs);
     newForOp->setAttrs(forOp->getAttrs());
-    ValueRange initArgsRange(initArgs);
-    TypeRange initArgsTypes(initArgsRange);
     Block *loopBody = &newForOp.getLoopBody().front();
 
     // Set up new iter_args. The loop body uses tensors, so wrap the (memref)
     // iter_args of the new loop in ToTensorOps.
     rewriter.setInsertionPointToStart(loopBody);
     SmallVector<Value> iterArgs =
         getBbArgReplacements(rewriter, newForOp.getRegionIterArgs(), indices);
     iterArgs.insert(iterArgs.begin(), newForOp.getInductionVar());
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     for (const auto &it : llvm::enumerate(yieldOp.getResults())) {
       Value value = it.value();
       if (value.getType().isa<TensorType>()) {
         FailureOr<Value> maybeBuffer = getBuffer(rewriter, value, options);
         if (failed(maybeBuffer))
           return failure();
         Value buffer = *maybeBuffer;
         if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp())) {
-          FailureOr<BaseMemRefType> resultType =
-              cast<BufferizableOpInterface>(forOp.getOperation())
-                  .getBufferType(forOp.getRegionIterArgs()[it.index()],
-                                 options);
+          FailureOr<BaseMemRefType> resultType = bufferization::getBufferType(
+              forOp.getRegionIterArgs()[it.index()], options);
           if (failed(resultType))
             return failure();
           buffer = castBuffer(rewriter, buffer, *resultType);
         }
         newResults.push_back(buffer);
       } else {
         newResults.push_back(value);
       }

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Get source buffer.
     FailureOr<Value> srcMemref =
         getBuffer(rewriter, extractSliceOp.getSource(), options);
     if (failed(srcMemref))
       return failure();
 
     // Take a subview of the source buffer.
     auto resultMemrefType =
-        getBufferType(op, extractSliceOp.getResult(), options);
+        bufferization::getBufferType(extractSliceOp.getResult(), options);
     if (failed(resultMemrefType))
       return failure();
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, resultMemrefType->cast<MemRefType>(), *srcMemref, mixedOffsets,
         mixedSizes, mixedStrides);
 
     replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
 
   FailureOr<BaseMemRefType>
-  getBufferType(Operation *op, Value value,
-                const BufferizationOptions &options) const {
+  getBufferType(Operation *op, Value value, const BufferizationOptions &options,
+                const DenseMap<Value, BaseMemRefType> &fixedTypes) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     assert(value == extractSliceOp.getResult() && "invalid value");
-    auto srcMemrefType =
-        bufferization::getBufferType(extractSliceOp.getSource(), options);
+    auto srcMemrefType = bufferization::getBufferType(
+        extractSliceOp.getSource(), options, fixedTypes);
     if (failed(srcMemrefType))
       return failure();
     SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
     SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
     return memref::SubViewOp::inferRankReducedResultType(
                extractSliceOp.getType().getShape(),
                srcMemrefType->cast<MemRefType>(), mixedOffsets, mixedSizes,

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

%r = scf.for %iv = %c0 to %sz step %c1 iter_args(%t = %0) -> tensor<?xf32> {
     %l = tensor.extract %ll[%c0] : tensor<?xf32>
     %double = arith.mulf %cst, %l : f32
     // CHECK: memref.store %{{.*}}, %[[generate]]
     %s = tensor.insert %double into %t[%iv] : tensor<?xf32>
     scf.yield %s : tensor<?xf32>
   }
   return %r : tensor<?xf32>
 }
+
+// -----
+
+// We just check that this example bufferizes to valid IR.
+
+// CHECK-LABEL: func @scf_for_buffer_type_mismatch
+func.func @scf_for_buffer_type_mismatch(%sz: index, %sz2: index) -> f32 {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %0 = bufferization.alloc_tensor(%sz) : tensor<?xf32>
+  %e2 = tensor.extract_slice %0[1][%sz2][1] : tensor<?xf32> to tensor<?xf32>
+  // init_arg and iter_arg have different buffer types. This must be resolved
+  // with casts.
+  %r = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t = %e2) -> tensor<?xf32> {
+    %s = "test.dummy"() : () -> (index)
+    %e = tensor.extract_slice %t[1][%s][1] : tensor<?xf32> to tensor<?xf32>
+    scf.yield %e : tensor<?xf32>
+  }
+  %x = tensor.extract %r[%c1] : tensor<?xf32>
+  return %x : f32
+}