[mlir] Add a bufferization pattern for linalg.tiled_loop.

[herhut]

The type converter is a bit of a hack but ok within a test pass to test expected behavior I think.

Otherwise just some nits.

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

269

Should this also check that the block is owned by the TiledLoopOp? You could check whether it is a block arg and then whether the immediate parent op of the blockarg is a TiledLoopOp.

319

Here, too.

449

That change will likely not land. Maybe rephrase that this is needed while dialect conversion does not properly support memrefs with maps.

[pifon2a]

Address the comments.

[silvas]

It feels like this bufferization pattern would miscompile if we were to write integration tests analogous to these:

https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir
https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir

We've tried multiple times to reconcile the desire for in-placeness with the bufferization framework via local changes to patterns, and it never ends well. I feel like we should put more effort into Nicolas' comprehensive bufferize which does the necessary analysis, or implementing a suitable generalization of the current bufferization framework (basically, each pattern would declare if it was going to create aliases / operate in place, and the framework would need to insert copies as needed, similar to how comprehensive bufferize does). Either way, I feel like this patch is a step in the wrong direction. Also, the fact that adding a bufferization pattern for TiledLoop required changing the bufferization pattern for LinalgOp seems like a strong smell to me.

Can we please revert this and work on a more sustainable direction?

[pifon2a]

@silvas Agreed. I am working on a less invasive pattern to bufferize TiledLoopOp, it does not require changes to Extract/InsertSlice patterns and to BufferizeAnyLinalgOp one. If I am not ready to send it out for review today, then I will revert this one.

In D107762#2935671, @silvas wrote:

It feels like this bufferization pattern would miscompile if we were to write integration tests analogous to these:

https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert.mlir
https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/test-subtensor-insert-multiple-uses.mlir

We've tried multiple times to reconcile the desire for in-placeness with the bufferization framework via local changes to patterns, and it never ends well. I feel like we should put more effort into Nicolas' comprehensive bufferize which does the necessary analysis, or implementing a suitable generalization of the current bufferization framework (basically, each pattern would declare if it was going to create aliases / operate in place, and the framework would need to insert copies as needed, similar to how comprehensive bufferize does). Either way, I feel like this patch is a step in the wrong direction. Also, the fact that adding a bufferization pattern for TiledLoop required changing the bufferization pattern for LinalgOp seems like a strong smell to me.

Can we please revert this and work on a more sustainable direction?

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

matchAndRewrite(LinalgOp op, ArrayRef<Value> operands,
     // We abuse the GenericOpAdaptor here.
     // TODO: Manually create an Adaptor that captures inputs and outputs for all
     // linalg::LinalgOp interface ops.
     linalg::GenericOpAdaptor adaptor(operands, op->getAttrDictionary());
 
     Location loc = op.getLoc();
     SmallVector<Value, 2> newOutputBuffers;
 
-    if (failed(allocateBuffersForResults(loc, op, adaptor.outputs(),
+    if (op->getParentOfType<TiledLoopOp>()) {
+      newOutputBuffers = adaptor.outputs();
+    } else if (failed(allocateBuffersForResults(loc, op, adaptor.outputs(),
                                                 newOutputBuffers, rewriter))) {
       return op.emitOpError()
              << "Failed to allocate buffers for tensor results.";
     }
 
     // Delegate to the linalg generic pattern.
     if (auto genericOp = dyn_cast<linalg::GenericOp>(*op)) {
       finalizeBufferAllocationForGenericOp(rewriter, genericOp,
                                            adaptor.inputs(), newOutputBuffers);
       return success();
     }
 
     finalizeBufferAllocation(rewriter, op, adaptor.inputs(), newOutputBuffers);
     return success();
   }
 };
 
+bool IsBlockArgOfTiledLoop(Value tensor) {
+  if (auto tensorLoad = tensor.getDefiningOp<memref::TensorLoadOp>())
+    if (auto blockArgument = tensorLoad.memref().dyn_cast<BlockArgument>())
+      if (isa<TiledLoopOp>(blockArgument.getOwner()->getParentOp()))
+        return true;
+  return false;
+}
+
 /// Convert `extract_slice %t [offsets][sizes][strides] -> %st` to an
 /// alloc + copy pattern.
 /// ```
 ///   %a = alloc(sizes)
 ///   %sv = subview %source [offsets][sizes][strides]
 ///   linalg_copy(%sv, %a)
 /// ```
 ///
 /// This pattern is arguable a std pattern once linalg::CopyOp becomes
 /// std::CopyOp.
 class ExtractSliceOpConverter
     : public OpConversionPattern<tensor::ExtractSliceOp> {
 public:
   using OpConversionPattern<tensor::ExtractSliceOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(tensor::ExtractSliceOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     tensor::ExtractSliceOpAdaptor adaptor(operands, op->getAttrDictionary());
     Value sourceMemref = adaptor.source();
     assert(sourceMemref.getType().isa<MemRefType>());
 
+    // Block arguments of the tiled_loop can be bufferized inplace.
+    if (IsBlockArgOfTiledLoop(op.source())) {
+      Value subView = rewriter.create<memref::SubViewOp>(
+          op.getLoc(), sourceMemref, op.getMixedOffsets(), op.getMixedSizes(),
+          op.getMixedStrides());
+      rewriter.replaceOp(op, subView);
+      return success();
+    }
+
     MemRefType subviewMemRefType =
         getTypeConverter()->convertType(op.getType()).cast<MemRefType>();
     // op.sizes() capture exactly the dynamic alloc operands matching the
     // subviewMemRefType thanks to subview/slice canonicalization and
     // verification.
     Value alloc = rewriter.create<memref::AllocOp>(
         op.getLoc(), subviewMemRefType, op.sizes());
     Value subView = rewriter.create<memref::SubViewOp>(
matchAndRewrite(tensor::InsertSliceOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     tensor::InsertSliceOpAdaptor adaptor(operands, op->getAttrDictionary());
     Value sourceMemRef = adaptor.source();
     assert(sourceMemRef.getType().isa<MemRefType>());
 
     // For now, be conservative and copy the converted input memref.
     // In general, the converted input memref here could be aliased or could
     // point into constant memory, so mutating it would lead to miscompilations.
-    Value destMemRef = cloneMemref(op.getLoc(), adaptor.dest(), rewriter);
+    // Block arguments of the tiled_loop can be bufferized inplace.
+    Value destMemRef;
+    if (IsBlockArgOfTiledLoop(op.dest()))
+      destMemRef = adaptor.dest();
+    else
+      destMemRef = cloneMemref(op.getLoc(), adaptor.dest(), rewriter);
     assert(destMemRef.getType().isa<MemRefType>());
 
     // Take a subview to copy the small memref.
     Value subview = rewriter.create<memref::SubViewOp>(
         op.getLoc(), destMemRef, op.getMixedOffsets(), op.getMixedSizes(),
         op.getMixedStrides());
     // Copy the small memref.
     rewriter.create<linalg::CopyOp>(op.getLoc(), sourceMemRef, subview);
     rewriter.replaceOp(op, destMemRef);
     return success();
   }
 };
 
+class TiledLoopOpConverter : public OpConversionPattern<TiledLoopOp> {
+public:
+  using OpConversionPattern<TiledLoopOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(TiledLoopOp tiledLoop, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    TiledLoopOp::Adaptor adaptor(operands, tiledLoop->getAttrDictionary());
+    Location loc = tiledLoop.getLoc();
+    if (tiledLoop.getNumResults() == 0)
+      return failure();
+    auto newTiledLoop = rewriter.create<TiledLoopOp>(
+        loc, adaptor.lowerBound(), adaptor.upperBound(), adaptor.step(),
+        adaptor.inputs(), adaptor.outputs(), adaptor.iterator_types(),
+        adaptor.distribution_types());
+    // Clone the region.
+    BlockAndValueMapping bvm;
+    bvm.map(tiledLoop.getInductionVars(), newTiledLoop.getInductionVars());
+
+    OpBuilder innerBuilder =
+        OpBuilder::atBlockEnd(newTiledLoop.getBody(), rewriter.getListener());
+
+    // Remap input block arguments.
+    SmallVector<Value, 2> inputs;
+    for (auto en : llvm::zip(newTiledLoop.getRegionInputArgs(),
+                             tiledLoop.getRegionInputArgs())) {
+      auto &newInputArg = std::get<0>(en);
+      if (!newInputArg.getType().isa<ShapedType>()) {
+        inputs.push_back(std::get<0>(en));
+        continue;
+      }
+      inputs.push_back(
+          innerBuilder.create<memref::TensorLoadOp>(loc, newInputArg));
+    }
+    bvm.map(tiledLoop.getRegionInputArgs(), inputs);
+
+    // Remap output block arguments.
+    SmallVector<Value, 2> outputs;
+    for (auto en : llvm::zip(newTiledLoop.getRegionOutputArgs(),
+                             tiledLoop.getRegionOutputArgs())) {
+      auto &newOutputArg = std::get<0>(en);
+      if (!newOutputArg.getType().isa<ShapedType>()) {
+        outputs.push_back(std::get<0>(en));
+        continue;
+      }
+      outputs.push_back(
+          innerBuilder.create<memref::TensorLoadOp>(loc, newOutputArg));
+    }
+    bvm.map(tiledLoop.getRegionOutputArgs(), outputs);
+
+    for (auto &op : tiledLoop.getBody()->without_terminator())
+      innerBuilder.clone(op, bvm);
+    innerBuilder.create<linalg::YieldOp>(loc);
+    rewriter.replaceOp(tiledLoop, newTiledLoop.outputs());
+    return success();
+  }
+};
+
 class VectorTransferReadOpConverter
     : public OpConversionPattern<vector::TransferReadOp> {
 public:
   using OpConversionPattern<vector::TransferReadOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(vector::TransferReadOp readOp, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
rewriter.create<vector::TransferWriteOp>(
         adaptor.permutation_map(),
         adaptor.in_bounds() ? adaptor.in_bounds() : ArrayAttr());
     rewriter.replaceOp(writeOp, adaptor.source());
     return success();
   }
 };
 } // namespace
 
+static Value materializeTensorLoad(OpBuilder &builder, TensorType type,
+                                   ValueRange inputs, Location loc) {
+  assert(inputs.size() == 1);
+  assert(inputs[0].getType().isa<BaseMemRefType>());
+  return builder.create<memref::TensorLoadOp>(loc, type, inputs[0]);
+}
+
 namespace {
+
+/// A helper type converter class that automatically populates the relevant
+/// materializations and type conversions for bufferization.
+//
+// The default BufferizeTypeConverter defined in "Transforms/Bufferize.h" does
+// not properly support memrefs with non-default layout. Whenever a layout of
+// memref changes during bufferization, target materialization call back would
+// assert that the non-matching type is a tensor.
+// There was an attempt to fix this behavior of dialect conversion in a more
+// principal way in https://reviews.llvm.org/D93126 but it had to be reverted
+// due to test failures outside of MLIR Core. It might make sense to revive this
+// PR.
+class CustomBufferizeTypeConverter : public BufferizeTypeConverter {
+public:
+  CustomBufferizeTypeConverter() {
+    // Keep all types unchanged.
+    addConversion([](Type type) { return type; });
+    // Convert RankedTensorType to MemRefType.
+    addConversion([](RankedTensorType type) -> Type {
+      return MemRefType::get(type.getShape(), type.getElementType());
+    });
+    // Convert UnrankedTensorType to UnrankedMemRefType.
+    addConversion([](UnrankedTensorType type) -> Type {
+      return UnrankedMemRefType::get(type.getElementType(), 0);
+    });
+    addArgumentMaterialization(materializeTensorLoad);
+    addSourceMaterialization(materializeTensorLoad);
+    addTargetMaterialization([](OpBuilder &builder, BaseMemRefType type,
+                                ValueRange inputs, Location loc) -> Value {
+      assert(inputs.size() == 1);
+      // Target materialization is invoked if the new operand type does not
+      // match the expected type. A special case is when the new operand type is
+      // a memref with a specified layout, i.e. non-empty affine map.
+      // TODO(pifon) : Change how target materialization is invoked in dialect
+      // conversion.
+      if (auto memrefType = inputs[0].getType().dyn_cast<MemRefType>()) {
+        assert(!memrefType.getAffineMaps().empty());
+        return inputs[0];
+      }
+      assert(inputs[0].getType().isa<TensorType>());
+      return builder.create<memref::BufferCastOp>(loc, type, inputs[0]);
+    });
+  }
+};
+
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
-    BufferizeTypeConverter typeConverter;
+    CustomBufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<AffineDialect, math::MathDialect,
                            memref::MemRefDialect, StandardOpsDialect,
                            tensor::TensorDialect>();
     target.addIllegalOp<InitTensorOp, tensor::ExtractSliceOp,
                         tensor::InsertSliceOp, PadTensorOp>();
 
void mlir::linalg::populateLinalgBufferizePatterns(
   patterns.add<
       BufferizeAnyLinalgOp,
       BufferizeFillOp,
       BufferizeInitTensorOp,
       BufferizeTensorReshapeOp<TensorExpandShapeOp>,
       BufferizeTensorReshapeOp<TensorCollapseShapeOp>,
       ExtractSliceOpConverter,
       InsertSliceOpConverter,
+      TiledLoopOpConverter,
       VectorTransferReadOpConverter,
       VectorTransferWriteOpConverter
     >(typeConverter, patterns.getContext());
   // clang-format on
   patterns.add<GeneralizePadTensorOpPattern>(patterns.getContext());
 }

mlir/test/Dialect/Linalg/bufferize.mlir

%read = vector.transfer_read %in[%c0], %cst {in_bounds = [true]}
       : tensor<4xf32>, vector<4xf32>
   %tanh = math.tanh %read : vector<4xf32>
   %write = vector.transfer_write %tanh, %out[%c0] {in_bounds = [true]}
       : vector<4xf32>, tensor<4xf32>
   return
   // CHECK: vector.transfer_read {{.*}} : memref<4xf32>, vector<4xf32>
   // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, memref<4xf32>
 }
+
+// -----
+
+//      CHECK:  func @tiled_dot
+func @tiled_dot(%A: tensor<10xf32>, %B: tensor<10xf32>,
+                %C: tensor<f32>) -> tensor<f32> {
+  %c0 = constant 0 : index
+  %c2 = constant 2 : index
+  %c10 = constant 10 : index
+
+  %dot = linalg.tiled_loop (%i) = (%c0) to (%c10) step (%c2)
+       ins (%A_ = %A: tensor<10xf32>, %B_ = %B: tensor<10xf32>)
+       outs (%C_ = %C: tensor<f32>)
+       iterators["reduction"] {
+    %A_sub = tensor.extract_slice %A_[%i] [%c2] [1]
+      : tensor<10xf32> to tensor<?xf32>
+    %B_sub = tensor.extract_slice %B_[%i] [%c2] [1]
+      : tensor<10xf32> to tensor<?xf32>
+    %dot_sub = linalg.dot ins(%A_sub, %B_sub : tensor<?xf32>, tensor<?xf32>)
+                          outs(%C_ : tensor<f32>) -> tensor<f32>
+    linalg.yield %dot_sub : tensor<f32>
+  }
+  // CHECK: linalg.tiled_loop
+  // CHECK-SAME: ins (%[[A:.*]] = %{{.*}}: memref<10xf32>, %[[B:.*]] = %{{.*}}: memref<10xf32>)
+  // CHECK-SAME: outs (%[[C:.*]] = %{{.*}}: memref<f32>)
+  //   CHECK-NOT:   alloc
+  //   CHECK:       %[[SV_A:.*]] = memref.subview %[[A]]
+  //   CHECK:       %[[SV_B:.*]] = memref.subview %[[B]]
+  //   CHECK:       linalg.dot ins(%[[SV_A]], %[[SV_B]]
+  //   CHECK-SAME:             outs(%[[C]] : memref<f32>)
+  //   CHECK:   linalg.yield
+  return %dot : tensor<f32>
+}