[mlir][Bufferize] Fix incorrect bufferization of rank-reducing tensor ops.

[nicolasvasilache]

Summary

This revision fixes SubviewOp, InsertSliceOp, ExtractSliceOp construction during bufferization
where not all offset/size/stride operands were properly specified.

A test that exhibited problematic behaviors related to incorrect memref casts is introduced.
Init tensor optimization is disabled in teh testing func bufferize pass.

[springerm]

mlir/include/mlir/Interfaces/ViewLikeInterface.td

488–489

Maybe worth noting that the missing dims are assumed to be "top" (least significant?) dims, i.e., [offsets.size(); rank). Or is this obvious from the context?

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

81–83 ↗

(On Diff #398456)

Could we fix this by always casting all memrefs to their "most dynamic" variants? I.e., dynamic dims and dynamic offsets + strides in the affine map. Then let canonicalization clean it up.

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-memref-cast-bug.mlir

1 ↗

(On Diff #398456)

Should we merge these into a single file for "regression tests"?

[nicolasvasilache]

Rebase address and cleanup

[nicolasvasilache]

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

81–83 ↗

(On Diff #398456)

The issue was related to rank-reduction due to the init tensor simplification and the missing rank-reducing behavior in the linalgimpl.
We may still want to do what you suggest in the future but this was a legitimate bug.

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-memref-cast-bug.mlir

1 ↗

(On Diff #398456)

yes it was just for the purpose of testing for now, the PR was marked WIP.

[nicolasvasilache]

Reasbe , clean

mlir/include/mlir/Interfaces/ViewLikeInterface.td

let extraClassDeclaration = [{
     static ::llvm::ArrayRef<::llvm::StringRef> getSpecialAttrNames() {
       static ::llvm::SmallVector<::llvm::StringRef, 4> names{
         ::mlir::OffsetSizeAndStrideOpInterface::getStaticOffsetsAttrName(),
         ::mlir::OffsetSizeAndStrideOpInterface::getStaticSizesAttrName(),
         ::mlir::OffsetSizeAndStrideOpInterface::getStaticStridesAttrName(),
         ::mlir::OpTrait::AttrSizedOperandSegments<void>::getOperandSegmentSizeAttr()};
       return names;
    }
+   /// Assume target is a shaped type and offsets/sizes/strides are vectors of
+   /// the same length and lower than target's rank.
+   /// Complete missing dims `i` with offset=0, size=dim(target, i), stride=1
+   /// until all vectors have size rank. The commpletion occurs for the most
+   /// minor dimensions (i.e. fastest varying).
+   /// Take a `createDim` lambda that knows how to build the size of a
+   /// particular dimension of `target` (to avoid dialect dependencies).
+   static void expandToRank(
+     Value target,
+     SmallVector<OpFoldResult> &offsets,
+     SmallVector<OpFoldResult> &sizes,
+     SmallVector<OpFoldResult> &strides,
+     llvm::function_ref<OpFoldResult(Value, int64_t)> createDim);
   }];
 
   let verify = [{
     return ::mlir::detail::verifyOffsetSizeAndStrideOp(
         ::mlir::cast<::mlir::OffsetSizeAndStrideOpInterface>($_op));
   }];
 }
 
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE

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

for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
             aliasInfo.unionAliasSets(opOperand.get(), opResult);
           aliasInfo.markInPlace(opOperand);
         }
     }
     return WalkResult::advance();
   });
 }
 
+// bufferization.to_memref is not allowed to change the rank.
+static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
+  auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>();
+  assert((!rankedTensorType || memrefType.cast<MemRefType>().getRank() ==
+                                   rankedTensorType.getRank()) &&
+         "to_memref would be invalid: mismatching ranks");
+}
+
 static Value lookupBuffer(RewriterBase &rewriter, Value tensor) {
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
 
   // Replace "%t = to_tensor %m" with %m.
   if (auto toTensorOp = tensor.getDefiningOp<bufferization::ToTensorOp>())
     return toTensorOp.memref();
 
   // Insert to_memref op.
   OpBuilder::InsertionGuard g(rewriter);
   setInsertionPointAfter(rewriter, tensor);
   Type memrefType;
   if (auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>()) {
     memrefType = getDynamicMemRefType(rankedTensorType);
   } else {
     memrefType = getUnrankedMemRefType(
         tensor.getType().cast<TensorType>().getElementType());
   }
+  ensureToMemrefOpIsValid(tensor, memrefType);
   return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
                                                     tensor);
 }
 
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 FailureOr<Value>

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

return eliminateInitTensors(
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
         if (!aliasInfo.isInPlace(insertSliceOp->getOpOperand(1) /*dest*/))
           return false;
         return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
       },
       /*rewriteFunc=*/
       [](OpBuilder &b, Location loc, OpOperand &operand) {
-        auto insertSliceOp = cast<tensor::InsertSliceOp>(operand.getOwner());
+        auto insertOp = cast<tensor::InsertSliceOp>(operand.getOwner());
+        // Expand offsets, sizes and strides to the full rank to handle the
+        // rank-reducing case.
+        SmallVector<OpFoldResult> mixedOffsets = insertOp.getMixedOffsets();
+        SmallVector<OpFoldResult> mixedSizes = insertOp.getMixedSizes();
+        SmallVector<OpFoldResult> mixedStrides = insertOp.getMixedStrides();
+        OffsetSizeAndStrideOpInterface::expandToRank(
+            insertOp.dest(), mixedOffsets, mixedSizes, mixedStrides,
+            [&](Value target, int64_t dim) -> OpFoldResult {
+              auto shapedType = target.getType().cast<ShapedType>();
+              if (shapedType.isDynamicDim(dim))
+                return b.create<tensor::DimOp>(loc, target, dim).result();
+              return b.getIndexAttr(shapedType.getDimSize(dim));
+            });
+        auto t = tensor::ExtractSliceOp::inferRankReducedResultType(
+            insertOp.getSourceType().getRank(),
+            insertOp.dest().getType().cast<RankedTensorType>(), mixedOffsets,
+            mixedSizes, mixedStrides);
         auto extractOp = b.create<tensor::ExtractSliceOp>(
-            loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
-            insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
+            loc, t, insertOp.dest(), mixedOffsets, mixedSizes, mixedStrides);
         return extractOp.result();
       },
       newOps);
 }
 
 void mlir::linalg::comprehensive_bufferize::linalg_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<linalg::InitTensorOp, InitTensorOpInterface>();

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

 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
+// bufferization.to_memref is not allowed to change the rank.
+static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
+  auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>();
+  assert(!rankedTensorType ||
+         memrefType.cast<MemRefType>().getRank() ==
+                 rankedTensorType.getRank() &&
+             "to_memref would be invalid: mismatching ranks");
+}
+
 /// Bufferization of scf.execute_region. Can be analyzed, but bufferization not
 /// fully implemented at the moment.
 struct ExecuteRegionOpInterface
     : public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
                                                     scf::ExecuteRegionOp> {
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
                        const BufferizationState &state) const {
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     rewriter.mergeBlocks(ifOp.elseBlock(), newIfOp.elseBlock());
 
     // Update scf.yield of new then-block.
     auto thenYieldOp = cast<scf::YieldOp>(newIfOp.thenBlock()->getTerminator());
     rewriter.setInsertionPoint(thenYieldOp);
     SmallVector<Value> thenYieldValues;
     for (OpOperand &operand : thenYieldOp->getOpOperands()) {
       if (operand.get().getType().isa<TensorType>()) {
+        ensureToMemrefOpIsValid(operand.get(),
+                                newTypes[operand.getOperandNumber()]);
         Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
             operand.get().getLoc(), newTypes[operand.getOperandNumber()],
             operand.get());
         operand.set(toMemrefOp);
       }
     }
 
     // Update scf.yield of new else-block.
     auto elseYieldOp = cast<scf::YieldOp>(newIfOp.elseBlock()->getTerminator());
     rewriter.setInsertionPoint(elseYieldOp);
     SmallVector<Value> elseYieldValues;
     for (OpOperand &operand : elseYieldOp->getOpOperands()) {
       if (operand.get().getType().isa<TensorType>()) {
+        ensureToMemrefOpIsValid(operand.get(),
+                                newTypes[operand.getOperandNumber()]);
         Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
             operand.get().getLoc(), newTypes[operand.getOperandNumber()],
             operand.get());
         operand.set(toMemrefOp);
       }
     }
 
     // Replace op results.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Move loop body to new loop.
     rewriter.mergeBlocks(oldLoopBody, loopBody, iterArgs);
 
     // Update scf.yield of new loop.
     auto yieldOp = cast<scf::YieldOp>(loopBody->getTerminator());
     rewriter.setInsertionPoint(yieldOp);
     SmallVector<Value> yieldValues =
         convert(yieldOp.getResults(), [&](Value val, int64_t index) {
+          ensureToMemrefOpIsValid(val, initArgs[index].getType());
           return rewriter.create<bufferization::ToMemrefOp>(
               val.getLoc(), initArgs[index].getType(), val);
         });
     yieldOp.getResultsMutable().assign(yieldValues);
 
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newForOp->getResults());
 

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     MemRefLayoutAttrInterface layout;
 
     if (auto rankedMemRefType = sourceMemRefType.dyn_cast<MemRefType>())
       if (resultTensorType.isa<RankedTensorType>())
         layout = rankedMemRefType.getLayout();
 
     // Compute the new memref type.
     Type resultMemRefType;
-    if (auto rankedTensorType = resultTensorType.isa<RankedTensorType>()) {
+    if (resultTensorType.isa<RankedTensorType>()) {
       resultMemRefType =
           getContiguousMemRefType(resultTensorType, layout, memorySpace);
     } else {
       resultMemRefType =
           getUnrankedMemRefType(resultTensorType.getElementType(), memorySpace);
     }
 
     // Replace the op with a memref.cast.
if (!inplace) {
       FailureOr<Value> allocOrFailure =
           state.createAlloc(rewriter, loc, extractSliceOp.result(),
                             state.getOptions().createDeallocs);
       if (failed(allocOrFailure))
         return failure();
       alloc = *allocOrFailure;
     }
 
+    // Expand offsets, sizes and strides to the full rank to handle the
+    // rank-reducing case.
+    SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = extractSliceOp.getMixedStrides();
+    OffsetSizeAndStrideOpInterface::expandToRank(
+        srcMemref, mixedOffsets, mixedSizes, mixedStrides,
+        [&](Value target, int64_t dim) -> OpFoldResult {
+          auto shapedType = target.getType().cast<ShapedType>();
+          if (shapedType.isDynamicDim(dim))
+            return rewriter.create<memref::DimOp>(loc, target, dim).result();
+          return rewriter.getIndexAttr(shapedType.getDimSize(dim));
+        });
     // Bufferize to subview.
-    auto subviewMemRefType =
-        memref::SubViewOp::inferRankReducedResultType(
+    auto subviewMemRefType = memref::SubViewOp::inferRankReducedResultType(
                                  dstTensorType.getRank(), srcMemrefType,
-            extractSliceOp.getMixedOffsets(), extractSliceOp.getMixedSizes(),
-            extractSliceOp.getMixedStrides())
+                                 mixedOffsets, mixedSizes, mixedStrides)
                                  .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
-        loc, subviewMemRefType, srcMemref, extractSliceOp.getMixedOffsets(),
-        extractSliceOp.getMixedSizes(), extractSliceOp.getMixedStrides());
+        loc, subviewMemRefType, srcMemref, mixedOffsets, mixedSizes,
+        mixedStrides);
 
     // If not inplaceable, copy.
     if (!inplace) {
       // Do not copy if the copied data is never read.
       if (state.isValueRead(extractSliceOp.result()))
         state.createMemCpy(rewriter, extractSliceOp.getLoc(), subView, alloc);
       subView = alloc;
     }
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     Location loc = insertSliceOp.getLoc();
 
     // When bufferizing out-of-place, `getResultBuffer` allocates.
     FailureOr<Value> dstMemref =
         state.getBuffer(rewriter, insertSliceOp->getOpOperand(1) /*dest*/);
     if (failed(dstMemref))
       return failure();
 
+    // Expand offsets, sizes and strides to the full rank to handle the
+    // rank-reducing case.
+    SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
+    SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
+    SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
+    OffsetSizeAndStrideOpInterface::expandToRank(
+        *dstMemref, mixedOffsets, mixedSizes, mixedStrides,
+        [&](Value target, int64_t dim) -> OpFoldResult {
+          auto shapedType = target.getType().cast<ShapedType>();
+          if (shapedType.isDynamicDim(dim))
+            return rewriter.create<memref::DimOp>(loc, target, dim).result();
+          return rewriter.getIndexAttr(shapedType.getDimSize(dim));
+        });
     // Take a subview of the dst.
     auto dstMemrefType = dstMemref->getType().cast<MemRefType>();
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
             insertSliceOp.getSourceType().getRank(), dstMemrefType,
-            insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
-            insertSliceOp.getMixedStrides())
+            mixedOffsets, mixedSizes, mixedStrides)
             .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
-        loc, subviewMemRefType, *dstMemref, insertSliceOp.getMixedOffsets(),
-        insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
+        loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
+        mixedStrides);
 
     // Copy tensor. If this tensor.insert_slice has a matching
     // tensor.extract_slice, the copy operation will eventually fold away.
     Value srcMemref =
         *state.getBuffer(rewriter, insertSliceOp->getOpOperand(0) /*source*/);
     state.createMemCpy(rewriter, loc, srcMemref, subView);
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);

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

void LinalgComprehensiveModuleBufferize::runOnOperation() {
   options->testAnalysisOnly = testAnalysisOnly;
   options->printConflicts = printConflicts;
 
   // Enable InitTensorOp elimination.
   if (initTensorElimination) {
     options->addPostAnalysisStep<
         linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   }
+
   if (!allowReturnMemref)
     options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, std::move(options)))) {
     signalPassFailure();

mlir/lib/Interfaces/ViewLikeInterface.cpp

bool mlir::detail::sameOffsetsSizesAndStrides(
   for (auto it : llvm::zip(a.getMixedSizes(), b.getMixedSizes()))
     if (!cmp(std::get<0>(it), std::get<1>(it)))
       return false;
   for (auto it : llvm::zip(a.getMixedStrides(), b.getMixedStrides()))
     if (!cmp(std::get<0>(it), std::get<1>(it)))
       return false;
   return true;
 }
+
+void OffsetSizeAndStrideOpInterface::expandToRank(
+    Value target, SmallVector<OpFoldResult> &offsets,
+    SmallVector<OpFoldResult> &sizes, SmallVector<OpFoldResult> &strides,
+    llvm::function_ref<OpFoldResult(Value, int64_t)> createOrFoldDim) {
+  auto shapedType = target.getType().cast<ShapedType>();
+  unsigned rank = shapedType.getRank();
+  assert(offsets.size() == sizes.size() && "mismatched lengths");
+  assert(offsets.size() == strides.size() && "mismatched lengths");
+  assert(offsets.size() <= rank && "rank overflow");
+  MLIRContext *ctx = target.getContext();
+  Attribute zero = IntegerAttr::get(IndexType::get(ctx), APInt(64, 0));
+  Attribute one = IntegerAttr::get(IndexType::get(ctx), APInt(64, 1));
+  for (unsigned i = offsets.size(); i < rank; ++i) {
+    offsets.push_back(zero);
+    sizes.push_back(createOrFoldDim(target, i));
+    strides.push_back(one);
+  }
+}

mlir/test/Dialect/Linalg/comprehensive-function-bufferize.mlir

 
 // CHECK-LABEL: func private @private_func
 func private @private_func(tensor<?xf32>) -> ()
 
 // CHECK-LABEL: func @empty_func()
 func @empty_func() -> () {
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @rank_reducing
+func @rank_reducing(
+    %i: index, %j: index,
+    %arg0: tensor<8x18x32xf32>) 
+      -> tensor<?x1x6x8xf32> {
+  %c1 = arith.constant 1 : index
+  %c6 = arith.constant 6 : index
+  %c8 = arith.constant 8 : index
+  %c32 = arith.constant 32 : index
+  %c0 = arith.constant 0 : index
+  %0 = linalg.init_tensor [4, 1, 6, 8] : tensor<4x1x6x8xf32>
+  %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>
+  %2 = linalg.init_tensor [1, 6, 8] : tensor<1x6x8xf32>
+  %5 = scf.for %arg7 = %c0 to %c32 step %c8 iter_args(%arg8 = %1) -> (tensor<?x1x6x8xf32>) {
+    %7 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg7)
+    %8 = tensor.extract_slice %arg0[%i, %j, %arg7] [1, 6, 8] [1, 1, 1] : tensor<8x18x32xf32> to tensor<1x6x8xf32>
+    %9 = scf.for %arg9 = %c0 to %c6 step %c1 iter_args(%arg10 = %2) -> (tensor<1x6x8xf32>) {
+      %11 = tensor.extract_slice %8[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x6x8xf32> to tensor<1x1x8xf32>
+      %12 = tensor.insert_slice %11 into %arg10[0, %arg9, 0] [1, 1, 8] [1, 1, 1] : tensor<1x1x8xf32> into tensor<1x6x8xf32>
+      scf.yield %12 : tensor<1x6x8xf32>
+    }
+    %10 = tensor.insert_slice %9 into %arg8[%7, 0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] : tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
+    scf.yield %10 : tensor<?x1x6x8xf32>
+  }
+  return %5: tensor<?x1x6x8xf32>
+}
 No newline at end of file

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir

func @equivalent_func_arg_2(%c0: index, %c10: index, %c1: index, %t0: tensor<?xf32>) -> tensor<?xf32> {
   // This test does not check IR. It just asserts there is no failure due to
   // non-equivalent scf.for yield values.
   %1 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%t1 = %t0) -> (tensor<?xf32>) {
     %3 = call @inner_func_2(%t1) : (tensor<?xf32>) -> tensor<?xf32>
     scf.yield %3 : tensor<?xf32>
   }
   return %1: tensor<?xf32>
 }
-
-// -----
-
-//===----------------------------------------------------------------------===//
-// InitTensorOp elimination would produce SSA violations for the example below.
-//===----------------------------------------------------------------------===//
-
-func @depthwise_conv_1d_nwc_wc(%arg0: index, %arg1: index, %arg2: tensor<8x18x32xf32>) 
-    -> tensor<?x1x6x8xf32> {
-  %c0 = arith.constant 0 : index
-  %c32 = arith.constant 32 : index
-  %c8 = arith.constant 8 : index
-  %0 = linalg.init_tensor [4, 1, 6, 8] : tensor<4x1x6x8xf32>
-  %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>
-  %2 = linalg.init_tensor [1, 6, 8] : tensor<1x6x8xf32>
-  %3 = scf.for %arg3 = %c0 to %c32 step %c8 iter_args(%arg4 = %1) -> (tensor<?x1x6x8xf32>) {
-    %4 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg3)
-    %5 = tensor.insert_slice %2 into %arg4[%4,0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] :
-      tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
-    scf.yield %5 : tensor<?x1x6x8xf32>
-  }
-  return %3 : tensor<?x1x6x8xf32>
-}
 

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @op_is_reading_but_following_ops_are_not(
   %r1 = linalg.generic #trait outs (%r0 : tensor<?xf32>) {
       ^bb(%0: f32) :
         linalg.yield %cst : f32
     } -> (tensor<?xf32>)
 
   // CHECK: return %[[ALLOC]]
   return %r1 : tensor<?xf32>
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// InitTensorOp elimination would produce SSA violations for the example below.
+//===----------------------------------------------------------------------===//
+
+func @depthwise_conv_1d_nwc_wc(%arg0: index, %arg1: index, %arg2: tensor<8x18x32xf32>) 
+    -> tensor<?x1x6x8xf32> {
+  %c0 = arith.constant 0 : index
+  %c32 = arith.constant 32 : index
+  %c8 = arith.constant 8 : index
+  %0 = linalg.init_tensor [4, 1, 6, 8] : tensor<4x1x6x8xf32>
+  %1 = tensor.cast %0 : tensor<4x1x6x8xf32> to tensor<?x1x6x8xf32>
+  %2 = linalg.init_tensor [1, 6, 8] : tensor<1x6x8xf32>
+  %3 = scf.for %arg3 = %c0 to %c32 step %c8 iter_args(%arg4 = %1) -> (tensor<?x1x6x8xf32>) {
+    %4 = affine.apply affine_map<(d0) -> (d0 ceildiv 8)>(%arg3)
+    %5 = tensor.insert_slice %2 into %arg4[%4,0, 0, 0] [1, 1, 6, 8] [1, 1, 1, 1] :
+      tensor<1x6x8xf32> into tensor<?x1x6x8xf32>
+    scf.yield %5 : tensor<?x1x6x8xf32>
+  }
+  return %3 : tensor<?x1x6x8xf32>
+}

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

ListOption<std::string> dialectFilter{
       llvm::cl::desc("Bufferize only ops from the specified dialects"),
       llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnFunction() {
   auto options = std::make_unique<BufferizationOptions>();
 
-  // Enable InitTensorOp elimination.
-  options->addPostAnalysisStep<
-      linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   if (!allowReturnMemref)
     options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;