[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

261

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.

314

Here, too.

443

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);
public:
 
   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 (op->getParentOfType<TiledLoopOp>()) {
+      auto tensorLoad = op.source().getDefiningOp<memref::TensorLoadOp>();
+      if (tensorLoad != nullptr && tensorLoad.memref().isa<BlockArgument>()) {
+        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 = op->getParentOfType<TiledLoopOp>() &&
+                               adaptor.dest().isa<BlockArgument>()
+                           ? adaptor.dest()
+                           : 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.
+//
+// TODO(pifon): Remove as soon as https://reviews.llvm.org/D93126 is landed.
+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 memref_type = inputs[0].getType().dyn_cast<MemRefType>()) {
+        assert(!memref_type.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>
+}