[mlir][Linalg] Add comprehensive bufferization support for linalg::InitTensor and tensor::CastOp (11/n)

[nicolasvasilache]

Summary

Also add an integration test that connects all the dots end to end, including with cast to unranked tensor for external library calls.

[ftynse]

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

1589–1598

Looks shorter

[nicolasvasilache]

Rebase.

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

 
 /// Return `true` if the op is explicitly supported by bufferization or if it
 /// has no result tensors.
 /// Other cases must be conservative.
 static bool hasKnownBufferizationAliasingBehavior(Operation *op) {
   return
       // clang-format off
       isa<CallOpInterface,
+          tensor::CastOp,
           scf::ForOp,
+          InitTensorOp,
           LinalgOp,
           ReturnOp,
           ExtractSliceOp,
           InsertSliceOp,
           VectorTransferOpInterface,
           scf::YieldOp>(op)
       // clang-format on
       || (none_of(op->getResultTypes(), isaTensor) &&

 static OpResult getInplaceableOpResult(InsertSliceOp op, OpOperand &opOperand) {
   if (opOperand.get() != op.dest())
     return OpResult();
   return op->getResult(0);
 }
 
 /// Return the OpResult that may bufferize into the same buffer as `opOperand`
 /// when the op is bufferized inplace.
+/// Return null if no such result exists.
+static OpResult getInplaceableOpResult(tensor::CastOp op,
+                                       OpOperand &opOperand) {
+  return op->getResult(0);
+}
+
+/// Return the OpResult that may bufferize into the same buffer as `opOperand`
+/// when the op is bufferized inplace.
 /// The inplace analysis uses this information along with interfering read
 /// analysis to determine which op results reuse the same buffer as some
 /// operand.
 static OpResult getInplaceableOpResult(OpOperand &opOperand) {
   return TypeSwitch<Operation *, OpResult>(opOperand.getOwner())
       // clang-format off
         // Ops that perform destructive updates on operand(s) to produce
         // result(s).
-        .Case<scf::ForOp,
+        .Case<tensor::CastOp,
+              scf::ForOp,
               LinalgOp,
               InsertSliceOp,
               VectorTransferOpInterface>(
             [&](auto op) { return getInplaceableOpResult(op, opOperand); })
         // ExtractSliceOp is special, when bufferized inplace it just returns an
         // alias to its operand. Its result is never inplaceable on its operand.
         .Case([&](ExtractSliceOp op) { return OpResult(); })
         // CallOpInterface is special, it needs to wait for the callee to be

 /// Return None if the owner of `opOperand` does not have known
 /// bufferization aliasing behavior, which indicates that the op must allocate
 /// all of its tensor results.
 /// TODO: in the future this may need to evolve towards a list of OpOperand*.
 static Optional<OpOperand *> getAliasingOpOperand(OpResult result) {
   if (!hasKnownBufferizationAliasingBehavior(result.getDefiningOp()))
     return None;
   return TypeSwitch<Operation *, OpOperand *>(result.getDefiningOp())
+      .Case([&](tensor::CastOp op) { return &op->getOpOperand(0); })
       .Case([&](LinalgOp op) {
         return op.getOutputTensorOperands()[result.getResultNumber()];
       })
       .Case([&](ExtractSliceOp op) { return &op->getOpOperand(0); })
       .Case([&](InsertSliceOp op) { return &op->getOpOperand(1); })
       .Case([&](vector::TransferWriteOp op) { return &op->getOpOperand(1); })
       // In the case of scf::ForOp, this currently assumes the iter_args / yield
       // are 1-1. This may fail and is verified at the end.
bufferize(OpBuilder &b, CallOpInterface callOp, BlockAndValueMapping &bvm,
 
   // 4. Create the new CallOp.
   Operation *newCallOp = b.create<CallOp>(callOp.getLoc(), funcOp.sym_name(),
                                           resultTypes, newOperands);
   newCallOp->setAttrs(callOp->getAttrs());
   return success();
 }
 
+/// tensor::CastOp bufferizes to memref::CastOp.
+static LogicalResult bufferize(OpBuilder &b, tensor::CastOp castOp,
+                               BlockAndValueMapping &bvm,
+                               BufferizationAliasInfo &aliasInfo) {
+  // Take a guard before anything else.
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(castOp);
+
+  Type sourceType = lookup(bvm, castOp.source()).getType();
+  auto rankedMemRefType = sourceType.dyn_cast<MemRefType>();
+  auto unrankedMemRefType = sourceType.dyn_cast<UnrankedMemRefType>();
+  assert(rankedMemRefType || unrankedMemRefType);
+  unsigned memorySpace = rankedMemRefType
+                             ? rankedMemRefType.getMemorySpaceAsInt()
+                             : unrankedMemRefType.getMemorySpaceAsInt();
+  TensorType tensorType = castOp.getResult().getType().cast<TensorType>();
+  ArrayRef<AffineMap> affineMaps =
+      rankedMemRefType && tensorType.isa<RankedTensorType>()
+          ? rankedMemRefType.getAffineMaps()
+          : ArrayRef<AffineMap>{};
+  Type memRefType = getContiguousOrUnrankedMemRefType(
+      castOp.getResult().getType(), {}, memorySpace);
+  Value res = b.create<memref::CastOp>(castOp.getLoc(), memRefType,
+                                       lookup(bvm, castOp.source()));
+  aliasInfo.insertNewBufferEquivalence(res, castOp.getResult());
+  map(bvm, castOp.getResult(), res);
+  return success();
+}
+
 /// DimOp tensor operand is modified inplace. This allows leaving dead
 /// tensors behind that will get DCE'd.
 static LogicalResult bufferize(OpBuilder &b, tensor::DimOp dimOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   if (dimOp.source().getType().isa<RankedTensorType>()) {
     Value v = lookup(bvm, dimOp.source());
     if (!v)
for (auto bbArg : funcOp.getArguments()) {
     Value bufferCast =
         b.create<memref::BufferCastOp>(funcOp.getLoc(), memRefType, bbArg);
     aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
     map(bvm, bbArg, bufferCast);
   }
   return success();
 }
 
+/// InitTensor always allocates.
+/// TODO: consider hoisting across function boundaries prior to bufferization.
+static LogicalResult bufferize(OpBuilder &b, InitTensorOp initTensorOp,
+                               BlockAndValueMapping &bvm,
+                               BufferizationAliasInfo &aliasInfo) {
+  // Take a guard before anything else.
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(initTensorOp);
+
+  Value alloc = createNewAllocDeallocPairForShapedValue(
+      b, initTensorOp->getLoc(), initTensorOp.result(), aliasInfo);
+  map(bvm, initTensorOp.result(), alloc);
+  return success();
+}
+
 /// ReturnOp always creates memref::TensorLoadOp.
 static LogicalResult bufferize(OpBuilder &b, ReturnOp returnOp,
                                BlockAndValueMapping &bvm,
                                BufferizationAliasInfo &aliasInfo) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(returnOp);
 
static LogicalResult bufferizeFuncOpInternals(
   OpBuilder b(funcOp->getContext());
   /// Start by bufferizing `funcOp` arguments.
   if (failed(bufferize(b, funcOp, bvm, aliasInfo)))
     return failure();
   // Walk in PreOrder to ensure ops with regions are handled before their body.
   // Since walk has to be PreOrder, we need to erase ops that require it
   // separately: this is the case for CallOp
   SmallVector<Operation *> toErase;
-  WalkResult result = funcOp.walk<WalkOrder::PreOrder>([&](Operation *op)
-                                                           -> WalkResult {
+  WalkResult result =
+      funcOp.walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
         // clang-format off
     WalkResult result =
       TypeSwitch<Operation *, LogicalResult>(op)
       // Skip BufferCast and TensorLoad ops.
       .Case<memref::BufferCastOp,
             memref::TensorLoadOp>([&](auto) { return success(); })
-      .Case<tensor::DimOp,
+      .Case<tensor::CastOp,
+            tensor::DimOp,
             scf::ForOp,
+            InitTensorOp,
             LinalgOp,
             ReturnOp,
             ExtractSliceOp,
             InsertSliceOp,
             VectorTransferOpInterface,
             scf::YieldOp>([&](auto op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo);
       })
       .Case([&](CallOpInterface op) {
         LDBG("Begin bufferize:\n" << op << '\n');
         return bufferize(b, op, bvm, aliasInfo, bufferizedFunctionTypes);
       })
       .Default([&](Operation *op) {
         auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
         if (any_of(op->getOperandTypes(), isaTensor) ||
             any_of(op->getResultTypes(), isaTensor))
           return failure();
         return success();
       });
         // clang-format on
 
         // Register post-walk erasure, if necessary.
         if (isa<CallOpInterface>(op))
           if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
               llvm::any_of(op->getResultTypes(), isaTensor))
             toErase.push_back(op);
 
         return result;
       });
   LDBG("End BufferizeFuncOpInternals:\n" << funcOp << '\n');
 
   for (Operation *op : toErase)
     op->erase();
 
   return failure(result.wasInterrupted());
 }
 

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

// CHECK-NEXT:   call @scf_for_with_tensor_insert_slice(%[[A]], %[[B]], %[[C]]
 
   // %r0#0 is actually %B after inplaceable results are swapped in the callee.
 // CHECK-NEXT:   call @some_external_func(%[[B]]) : (memref<?xf32, #[[$DYN_1D_MAP]]>) -> ()
   call @some_external_func(%r0#0) : (tensor<?xf32>) -> ()
 
 // CHECK-NEXT:   return
   return %r0#0, %r0#1: tensor<?xf32>, tensor<?xf32>
 }
+
+// -----
+
+//  CHECK-DAG: #[[$DYN_0D_MAP:.*]] = affine_map<()[s0] -> (s0)>
+//  CHECK-DAG: #[[$DYN_1D_MAP:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+
+//      CHECK:  func @init_and_dot(
+// CHECK-SAME:    %[[A:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[B:[a-zA-Z0-9]*]]: memref<64xf32, #[[$DYN_1D_MAP]]>
+// CHECK-SAME:    %[[C:[a-zA-Z0-9]*]]: memref<f32, #[[$DYN_0D_MAP]]>
+func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> {
+  // CHECK-NEXT:   %[[C0:.*]] = constant 0{{.*}} : f32
+  %v0 = constant 0.0 : f32
+
+  // CHECK-NEXT:   linalg.fill(%[[C0]], %[[C]]) : f32, memref<f32, #[[$DYN_0D_MAP]]>
+  %d = linalg.fill(%v0, %c) : f32, tensor<f32> -> tensor<f32>
+
+  // CHECK-NEXT:   linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, #[[$DYN_1D_MAP]]>, memref<64xf32, #[[$DYN_1D_MAP]]>) outs(%[[C]] : memref<f32, #[[$DYN_0D_MAP]]>)
+  %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>)
+    outs(%d: tensor<f32>) -> tensor<f32>
+
+  // CHECK-NEXT:   return
+  return %e : tensor<f32>
+}
+
+//      CHECK:  func @main()
+func @main() {
+  //  CHECK-DAG:   %[[C0:.*]] = constant 0{{.*}} : f32
+  //  CHECK-DAG:   %[[C1:.*]] = constant 1{{.*}} : f32
+  //  CHECK-DAG:   %[[C2:.*]] = constant 2{{.*}} : f32
+  %v0 = constant 0.0 : f32
+  %v1 = constant 1.0 : f32
+  %v2 = constant 2.0 : f32
+
+  // CHECK-NEXT:   %[[A:.*]] = memref.alloc() : memref<64xf32>
+  // CHECK-NEXT:   %[[B:.*]] = memref.alloc() : memref<64xf32>
+  // CHECK-NEXT:   %[[C:.*]] = memref.alloc() : memref<f32>
+  %A = linalg.init_tensor [64] : tensor<64xf32>
+  %B = linalg.init_tensor [64] : tensor<64xf32>
+  %C = linalg.init_tensor [] : tensor<f32>
+
+  // CHECK-NEXT:   linalg.fill(%[[C1]], %[[A]]) : f32, memref<64xf32>
+  // CHECK-NEXT:   linalg.fill(%[[C2]], %[[B]]) : f32, memref<64xf32>
+  // CHECK-NEXT:   linalg.fill(%[[C0]], %[[C]]) : f32, memref<f32>
+  %AA = linalg.fill(%v1, %A) : f32, tensor<64xf32> -> tensor<64xf32>
+  %BB = linalg.fill(%v2, %B) : f32, tensor<64xf32> -> tensor<64xf32>
+  %CC = linalg.fill(%v0, %C) : f32, tensor<f32> -> tensor<f32>
+
+  // CHECK-NEXT:   %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
+  // CHECK-NEXT:   %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, #[[$DYN_1D_MAP]]>
+  // CHECK-NEXT:   %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, #[[$DYN_0D_MAP]]>
+  // CHECK-NEXT:   call @init_and_dot(%[[cA]], %[[cB]], %[[cC]])
+  %res = call @init_and_dot(%AA, %BB, %CC) :
+    (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
+
+  // CHECK-NEXT:   %[[dC:.*]] = memref.cast %[[C]] : memref<f32> to memref<*xf32>
+  %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
+
+  // CHECK-NEXT:   call @print_memref_f32(%[[dC]]) : (memref<*xf32>) -> ()
+  call @print_memref_f32(%res2) : (tensor<*xf32>) -> ()
+
+  // CHECK-DAG:   memref.dealloc %[[A]] : memref<64xf32>
+  // CHECK-DAG:   memref.dealloc %[[B]] : memref<64xf32>
+  // CHECK-DAG:   memref.dealloc %[[C]] : memref<f32>
+  // CHECK-NEXT:   return
+  return
+}
+
+//     CHECK:   func private @print_memref_f32(memref<*xf32>)
+func private @print_memref_f32(tensor<*xf32>)

mlir/test/Integration/Dialect/Linalg/CPU/test-comprehensive-bufferize.mlir

+// RUN: mlir-opt %s -canonicalize -cse -linalg-comprehensive-module-bufferize |\
+// RUN: mlir-opt -convert-vector-to-scf -lower-affine -convert-linalg-to-loops |\
+// RUN: mlir-opt -canonicalize -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
+
+// RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext |\
+// RUN: FileCheck %s
+
+func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> {
+  %v0 = constant 0.0 : f32
+
+  %d = linalg.fill(%v0, %c) : f32, tensor<f32> -> tensor<f32>
+
+  %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>)
+    outs(%d: tensor<f32>) -> tensor<f32>
+
+  return %e : tensor<f32>
+}
+
+func @main() {
+  %v0 = constant 0.0 : f32
+  %v1 = constant 1.0 : f32
+  %v2 = constant 2.0 : f32
+
+  %A = linalg.init_tensor [64] : tensor<64xf32>
+  %B = linalg.init_tensor [64] : tensor<64xf32>
+  %C = linalg.init_tensor [] : tensor<f32>
+  %AA = linalg.fill(%v1, %A) : f32, tensor<64xf32> -> tensor<64xf32>
+  %BB = linalg.fill(%v2, %B) : f32, tensor<64xf32> -> tensor<64xf32>
+  %CC = linalg.fill(%v0, %C) : f32, tensor<f32> -> tensor<f32>
+
+  %res = call @init_and_dot(%AA, %BB, %CC) :
+    (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
+
+  %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
+
+//      CHECK: Unranked Memref base@ = {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data =
+// CHECK-NEXT: [128]
+  call @print_memref_f32(%res2) : (tensor<*xf32>) -> ()
+
+  return
+}
+
+func private @print_memref_f32(tensor<*xf32>) attributes { llvm.emit_c_interface }