[mlir][tensor][bufferize] Bufferize tensor.splat op

[springerm]

Summary

The op bufferizes similarly to tensor.generate: it is lowered to a linalg.map, which may then lower to a loop nest that fills the buffer.

[ingomueller-net]

I wouldn't be able to catch many, even easy flaws, but from what I can judge, all looks good, and I do see the splat op in my own example being bufferized with the patch. Thanks!

[ingomueller-net]

Actually, do we really need the MapOp with an empty body? Can we not just allocate the tensor and that's it?

[springerm]

In D150952#4356162, @ingomueller-net wrote:

Actually, do we really need the MapOp with an empty body? Can we not just allocate the tensor and that's it?

The body of the linalg.map is not empty. It has a linalg.yield terminator. The reason why I used linalg.map is because it can lower to a loop or vectorize. Alternatively, I could've generated a loop nest that fills the allocation, but then the op would not vectorize anymore.

[ingomueller-net]

In D150952#4360295, @springerm wrote:

In D150952#4356162, @ingomueller-net wrote:

Actually, do we really need the MapOp with an empty body? Can we not just allocate the tensor and that's it?

The body of the linalg.map is not empty. It has a linalg.yield terminator. The reason why I used linalg.map is because it can lower to a loop or vectorize. Alternatively, I could've generated a loop nest that fills the allocation, but then the op would not vectorize anymore.

What's the semantic of yielding nothing?

[springerm]

In D150952#4360299, @ingomueller-net wrote:

What's the semantic of yielding nothing?

It does yield something: the operand of the tensor.splat.

[ingomueller-net]

In D150952#4360315, @springerm wrote:

In D150952#4360299, @ingomueller-net wrote:

What's the semantic of yielding nothing?

It does yield something: the operand of the tensor.splat.

Thanks for clarifying and insisting ;) Indeed, it does yield something. I don't know how I got to the impression it did not... All good now :)

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

struct ParallelInsertSliceOpInterface
   bool isNotConflicting(Operation *op, OpOperand *uRead,
                         OpOperand *uConflictingWrite,
                         const AnalysisState &state) const {
     return isNotConflictingInsertSliceLikeOp<tensor::ParallelInsertSliceOp>(
         op, uRead, uConflictingWrite, state);
   }
 };
 
+/// Bufferization of tensor.splat. Bufferizes to a new allocation that is filled
+/// with a linalg.map. Similar to tensor.generate.
+struct SplatOpInterface
+    : public BufferizableOpInterface::ExternalModel<SplatOpInterface,
+                                                    tensor::SplatOp> {
+
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options) const {
+    OpBuilder::InsertionGuard g(rewriter);
+    auto splatOp = cast<tensor::SplatOp>(op);
+
+    // Should the buffer be deallocated?
+    bool dealloc =
+        shouldDeallocateOpResult(cast<OpResult>(splatOp.getResult()), options);
+
+    // TODO: Implement memory space for this op.
+    if (options.defaultMemorySpace != Attribute())
+      return op->emitError("memory space not implemented yet");
+
+    // Allocate memory.
+    Location loc = op->getLoc();
+    FailureOr<Value> tensorAlloc =
+        allocateTensorForShapedValue(rewriter, loc, splatOp.getResult(),
+                                     /*escape=*/!dealloc, options,
+                                     /*copy=*/false);
+    if (failed(tensorAlloc))
+      return failure();
+
+    // Create linalg::MapOp.
+    auto tensorType = cast<RankedTensorType>(tensorAlloc->getType());
+    auto linalgOp =
+        rewriter.create<linalg::MapOp>(loc, tensorType, /*inputs=*/ValueRange(),
+                                       /*init=*/*tensorAlloc);
+    Block &linalgBody = linalgOp.getMapper().emplaceBlock();
+
+    // Create linalg::IndexOps.
+    rewriter.setInsertionPointToStart(&linalgBody);
+    rewriter.create<linalg::YieldOp>(loc, splatOp.getInput());
+    rewriter.replaceOp(splatOp, linalgOp.getResult()[0]);
+
+    return success();
+  }
+};
+
 } // namespace
 } // namespace tensor
 } // namespace mlir
 
 void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
     CastOp::attachInterface<CastOpInterface>(*ctx);
     CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
     DimOp::attachInterface<DimOpInterface>(*ctx);
     EmptyOp::attachInterface<EmptyOpInterface>(*ctx);
     ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
     ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
     ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
     FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
     GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
     InsertOp::attachInterface<InsertOpInterface>(*ctx);
     InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
     PadOp::attachInterface<PadOpInterface>(*ctx);
     ParallelInsertSliceOp::attachInterface<ParallelInsertSliceOpInterface>(
         *ctx);
     RankOp::attachInterface<RankOpInterface>(*ctx);
     ReshapeOp::attachInterface<ReshapeOpInterface>(*ctx);
+    SplatOp::attachInterface<SplatOpInterface>(*ctx);
 
     // Load additional dialects of which ops may get created.
     ctx->loadDialect<arith::ArithDialect, linalg::LinalgDialect>();
   });
 }

mlir/test/Dialect/Tensor/bufferize.mlir

^bb0(%arg0: index, %arg1: index):
     %m = arith.muli %arg0, %arg1 : index
     tensor.yield %m : index
   } : tensor<?x10xindex> to tensor<?x?xindex>
 
   // CHECK:     %[[r:.*]] = bufferization.to_tensor %[[mapped_m]]
   // CHECK:     return %[[r]] : tensor<?x?xindex>
   return %0 : tensor<?x?xindex>
 }
+
+// -----
+
+// CHECK-LABEL:   func @tensor.splat(
+// CHECK-SAME:        %[[F:.*]]: f32)
+// CHECK-DAG:       %[[ALLOC:.*]] = memref.alloc() {{.*}} : memref<10x2x4xf32>
+// CHECK:           %[[ALLOC_T:.*]] = bufferization.to_tensor %[[ALLOC]]
+// CHECK:           %[[MAPPED:.*]] = linalg.map
+// CHECK:                 outs(%[[ALLOC_T]] : tensor<10x2x4xf32>)
+// CHECK:             linalg.yield %[[F]]
+// CHECK:           }
+// CHECK:           return %[[MAPPED]] : tensor<10x2x4xf32>
+// CHECK:         }
+func.func @tensor.splat(%f: f32) -> tensor<10x2x4xf32> {
+  %t = tensor.splat %f : tensor<10x2x4xf32>
+  return %t : tensor<10x2x4xf32>
+}