[mlir] Add linalg.fill bufferization conversion

[ezhulenev]

Summary

BufferizeAnyLinalgOp fails because FillOp is not a LinalgGenericOp and it fails while reading operand sizes attribute.

[nicolasvasilache]

FillOp abides by the structured op interface, what is annoying here is there is no adaptor for an interface.
Here is what I am currently iterating on, which does not use the conversion infrastructure: https://github.com/google/iree-llvm-sandbox/blob/main/runners/LinalgComprehensiveBufferizePass.cpp#L1043.

I guess this is fine in the meantime but I would be very cautious on depending on anything related to core linalg bufferization: it will likely break in a bunch of ways before it gets better (see e.g. https://llvm.discourse.group/t/properly-using-bufferization-related-passes/2913).

[antiagainst]

Note that tensor_to_memref is now memref.buffer_cast and tensor_load should be memref.tesnor_load.

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

static Value cloneMemref(Location loc, Value memref, OpBuilder &b) {
   auto memrefType = memref.getType().cast<MemRefType>();
   auto alloc =
       b.create<AllocOp>(loc, memrefType, getDynOperands(loc, memref, b));
   b.create<linalg::CopyOp>(loc, memref, alloc);
   return alloc;
 }
 
 static LogicalResult
-allocateBuffersForResults(Location loc, LinalgOp linalgOp,
-                          linalg::GenericOpAdaptor &adaptor,
+allocateBuffersForResults(Location loc, LinalgOp linalgOp, ValueRange outputs,
                           SmallVectorImpl<Value> &resultBuffers, OpBuilder &b) {
   // Lazily compute loopRanges.
   SmallVector<Range, 4> loopRanges;
 
   // Allocate a buffer for every tensor result.
   assert(linalgOp.getNumOutputs() == linalgOp->getNumResults());
   for (auto en : llvm::enumerate(linalgOp->getResultTypes())) {
     size_t resultIndex = en.index();
     Type resultType = en.value();
 
     auto tensorType = resultType.dyn_cast<RankedTensorType>();
     if (tensorType == nullptr) {
       linalgOp.emitOpError()
           << "tensor to buffer conversion expects ranked tensor results";
       return failure();
     }
     auto tensorShape = tensorType.getShape();
     auto memrefType = MemRefType::get(tensorShape, tensorType.getElementType());
-    Value resultTensor = adaptor.outputs()[resultIndex];
+    Value resultTensor = outputs[resultIndex];
 
     // Clone output buffers whose value is actually used.
     if (linalgOp.payloadUsesValueFromOutputOperandIndex(resultIndex)) {
       resultBuffers.push_back(cloneMemref(loc, resultTensor, b));
       continue;
     }
 
     if (auto alloc = resultTensor.getDefiningOp<AllocOp>()) {

 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization patterns.
 //===----------------------------------------------------------------------===//
 
 namespace {
 
-/// Generic conversion pattern that matches any LinalgOp. This avoids template
-/// instantiating one pattern for each LinalgOp.
+/// Conversion pattern that replaces `linalg.init_tensor` with allocation.
 class BufferizeInitTensorOp : public OpConversionPattern<InitTensorOp> {
 public:
   using OpConversionPattern<InitTensorOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(InitTensorOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
     linalg::InitTensorOpAdaptor adaptor(operands, op->getAttrDictionary());
     rewriter.replaceOpWithNewOp<AllocOp>(
         op, getTypeConverter()->convertType(op.getType()).cast<MemRefType>(),
         adaptor.sizes());
     return success();
   }
 };
 
+/// Conversion pattern that bufferizes `linalg.fill` operation.
+class BufferizeFillOp : public OpConversionPattern<FillOp> {
+public:
+  using OpConversionPattern<FillOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(FillOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const final {
+    linalg::FillOpAdaptor adaptor(operands, op->getAttrDictionary());
+    if (!op.output().getType().isa<TensorType>())
+      return rewriter.notifyMatchFailure(op,
+                                         "operand must be of a tensor type");
+
+    rewriter.create<FillOp>(op.getLoc(), adaptor.output(), adaptor.value());
+    rewriter.replaceOp(op, adaptor.output());
+
+    return success();
+  }
+};
+
 /// Generic conversion pattern that matches any LinalgOp. This avoids template
 /// instantiating one pattern for each LinalgOp.
 class BufferizeAnyLinalgOp : public ConversionPattern {
 public:
   BufferizeAnyLinalgOp(TypeConverter &typeConverter)
       : ConversionPattern(/*benefit=*/1, typeConverter, MatchAnyOpTypeTag()) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const final {
 
     LinalgOp linalgOp = dyn_cast<linalg::LinalgOp>(op);
     if (!linalgOp)
       return failure();
 
     // 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 = linalgOp.getLoc();
     SmallVector<Value, 2> newOutputBuffers;
 
-    if (failed(allocateBuffersForResults(loc, linalgOp, adaptor,
+    if (failed(allocateBuffersForResults(loc, linalgOp, adaptor.outputs(),
                                          newOutputBuffers, rewriter))) {
       linalgOp.emitOpError()
           << "Failed to allocate buffers for tensor results.";
       return failure();
     }
 
     // Delegate to the linalg generic pattern.
     if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {

 
 void mlir::linalg::populateLinalgBufferizePatterns(
     MLIRContext *context, BufferizeTypeConverter &typeConverter,
     OwningRewritePatternList &patterns) {
   patterns.insert<BufferizeAnyLinalgOp>(typeConverter);
   // TODO: Drop this once tensor constants work in standard.
   // clang-format off
   patterns.insert<
+      BufferizeFillOp,
       BufferizeInitTensorOp,
       SubTensorOpConverter,
       SubTensorInsertOpConverter
     >(typeConverter, context);
   // clang-format on
 }

mlir/test/Dialect/Linalg/bufferize.mlir

func @bufferize_subtensor_insert(%t : tensor<?x?xf32>, %st0 : tensor<2x3xf32>, %st1 : tensor<2x?xf32>) ->
   // CHECK-NEXT: linalg.copy(%[[SM1]], %[[SUBVIEW1]]) : memref<2x?xf32>, memref<2x?xf32, #[[$MAP1]]>
   // CHECK-NEXT: %[[RT1:.*]] = tensor_load %[[M1_COPY]] : memref<?x?xf32>
   %t1 = subtensor_insert %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor<?x?xf32>
 
   //     CHECK: return %[[RT0]], %[[RT1]]
   return %t0, %t1: tensor<?x?xf32>, tensor<?x?xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @bufferize_fill(
+// CHECK-SAME:    %[[IN:.*]]: tensor<?xf32>
+func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  %c0 = constant 0.0 : f32
+  // CHECK: %[[MEMREF:.*]] = tensor_to_memref %[[IN]] : memref<?xf32>
+  // CHECK: linalg.fill(%[[MEMREF]], %cst) : memref<?xf32>, f32
+  // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF]] : memref<?xf32>
+  // CHECK: return %[[TENSOR]]
+  %0 = linalg.fill(%arg0, %c0) : tensor<?xf32>, f32 -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}