[mlir][bufferization] Return BufferLikeType in BufferizableOpInterface (#144867)

Support custom types (2/N): allow value-owning operations (e.g.
allocation ops) to bufferize custom tensors into custom buffers. This
requires BufferizableOpInterface::getBufferType() to return
BufferLikeType instead of BaseMemRefType.

Affected implementors of the interface are updated accordingly.

Relates to ee070d0.

Author: andrey-golubev

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

@@ -480,7 +480,7 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
     FailureOr<BufferLikeType> getBufferType(
         Value value, const BufferizationOptions &options,
         SmallVector<Value> &invocationStack) {
-      return ::llvm::cast<BaseMemRefType>(getBuffer().getType());
+      return getBuffer().getType();
     }
   }];
 

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h

@@ -18,7 +18,6 @@
 
 namespace mlir::bufferization {
 struct BufferizationOptions;
-class BufferizationState;
 class BufferLikeType;
 } // namespace mlir::bufferization
 

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

@@ -953,8 +953,10 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
   auto tensorType = cast<TensorType>(value.getType());
 
   // No further analysis is possible for a block argument.
-  if (llvm::isa<BlockArgument>(value))
-    return bufferization::getMemRefType(tensorType, options);
+  if (llvm::isa<BlockArgument>(value)) {
+    return cast<BufferLikeType>(
+        bufferization::getMemRefType(tensorType, options));
+  }
 
   // Value is an OpResult.
   Operation *op = getOwnerOfValue(value);
@@ -966,8 +968,7 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
     // If the OpResult has an equivalent OpOperand, both OpResult and
     // OpOperand bufferize to the exact same buffer type.
     Value equivalentOperand = aliases.getAliases().front().opOperand->get();
-    return asMemRefType(
-        getBufferType(equivalentOperand, options, invocationStack));
+    return getBufferType(equivalentOperand, options, invocationStack);
   }
 
   // If we do not know the memory space and there is no default memory space,
@@ -977,7 +978,8 @@ FailureOr<BufferLikeType> bufferization::detail::defaultGetBufferType(
   if (!memSpace.has_value())
     return op->emitError("could not infer memory space");
 
-  return getMemRefType(tensorType, options, /*layout=*/{}, *memSpace);
+  return cast<BufferLikeType>(
+      getMemRefType(tensorType, options, /*layout=*/{}, *memSpace));
 }
 
 bool bufferization::detail::defaultIsRepetitiveRegion(

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

@@ -524,8 +524,8 @@ static FailureOr<BufferLikeType> computeLoopRegionIterArgBufferType(
     const BufferizationOptions &options, const BufferizationState &state,
     SmallVector<Value> &invocationStack) {
   // Determine the buffer type of the init_arg.
-  auto initArgBufferType = bufferization::detail::asMemRefType(
-      bufferization::getBufferType(initArg, options, invocationStack));
+  auto initArgBufferType =
+      bufferization::getBufferType(initArg, options, invocationStack);
   if (failed(initArgBufferType))
     return failure();
 
@@ -551,8 +551,8 @@ static FailureOr<BufferLikeType> computeLoopRegionIterArgBufferType(
   } else {
     // Note: This typically triggers a recursive call for the buffer type of
     // the iter_arg.
-    auto maybeBufferType = bufferization::detail::asMemRefType(
-        bufferization::getBufferType(yieldedValue, options, invocationStack));
+    auto maybeBufferType =
+        bufferization::getBufferType(yieldedValue, options, invocationStack);
     if (failed(maybeBufferType))
       return failure();
     yieldedValueBufferType = *maybeBufferType;
@@ -715,12 +715,7 @@ struct ForOpInterface
     if (auto opResult = dyn_cast<OpResult>(value)) {
       // The type of an OpResult must match the corresponding iter_arg type.
       BlockArgument bbArg = forOp.getTiedLoopRegionIterArg(opResult);
-      auto bufferType =
-          bufferization::getBufferType(bbArg, options, invocationStack);
-      if (failed(bufferType))
-        return failure();
-      assert(isa<BaseMemRefType>(*bufferType) && "expected memref type");
-      return cast<BaseMemRefType>(*bufferType);
+      return bufferization::getBufferType(bbArg, options, invocationStack);
     }
 
     // Compute result/argument number.
@@ -1079,8 +1074,8 @@ struct WhileOpInterface
       // scf.condition was already bufferized.
       return cast<BufferLikeType>(conditionYieldedVal.getType());
     }
-    return bufferization::detail::asMemRefType(bufferization::getBufferType(
-        conditionYieldedVal, options, invocationStack));
+    return bufferization::getBufferType(conditionYieldedVal, options,
+                                        invocationStack);
   }
 
   /// Assert that yielded values of an scf.while op are equivalent to their
@@ -1308,14 +1303,14 @@ struct ForallOpInterface
     if (auto bbArg = dyn_cast<BlockArgument>(value))
       // A tensor block argument has the same bufferized type as the
       // corresponding output operand.
-      return bufferization::detail::asMemRefType(bufferization::getBufferType(
-          forallOp.getTiedOpOperand(bbArg)->get(), options, invocationStack));
+      return bufferization::getBufferType(
+          forallOp.getTiedOpOperand(bbArg)->get(), options, invocationStack);
 
     // The bufferized result type is the same as the bufferized type of the
     // corresponding output operand.
-    return bufferization::detail::asMemRefType(bufferization::getBufferType(
+    return bufferization::getBufferType(
         forallOp.getOutputs()[cast<OpResult>(value).getResultNumber()], options,
-        invocationStack));
+        invocationStack);
   }
 
   bool isRepetitiveRegion(Operation *op, unsigned index) const {

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

@@ -82,7 +82,7 @@ struct CastOpInterface
     auto rankedResultType = cast<RankedTensorType>(castOp.getType());
     return cast<BufferLikeType>(MemRefType::get(
         rankedResultType.getShape(), rankedResultType.getElementType(),
-        llvm::cast<MemRefType>(*maybeSrcBufferType).getLayout(), memorySpace);
+        llvm::cast<MemRefType>(*maybeSrcBufferType).getLayout(), memorySpace));
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
@@ -503,8 +503,8 @@ struct FromElementsOpInterface
         /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
-    FailureOr<BaseMemRefType> memrefType = bufferization::detail::asMemRefType(
-        bufferization::getBufferType(*tensorAlloc, options));
+    FailureOr<BufferLikeType> memrefType =
+        bufferization::getBufferType(*tensorAlloc, options);
     if (failed(memrefType))
       return failure();
     Value buffer = rewriter.create<bufferization::ToBufferOp>(

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

@@ -272,10 +272,10 @@ func.func @materialize_in_dest_raw(%f: f32, %f2: f32, %idx: index) -> (tensor<5x
 
 // -----
 
-// CHECK-LABEL: func.func @test_dialect_op(
+// CHECK:       func.func @custom_op(
 // CHECK-SAME:    %[[ARG:.*]]: !test.test_tensor<[32, 64], f64>
 // CHECK-SAME:  ) -> !test.test_tensor<[32, 128], f64> {
-func.func @test_dialect_op(%arg: !test.test_tensor<[32, 64], f64>)
+func.func @custom_op(%arg: !test.test_tensor<[32, 64], f64>)
     -> !test.test_tensor<[32, 128], f64> {
   // CHECK: %[[MEMREF:.*]] = bufferization.to_buffer %[[ARG]]
   // CHECK: %[[DUMMY:.*]] = "test.dummy_memref_op"(%[[MEMREF]])
@@ -288,3 +288,22 @@ func.func @test_dialect_op(%arg: !test.test_tensor<[32, 64], f64>)
   // CHECK: return %[[OUT]]
   return %out : !test.test_tensor<[32, 128], f64>
 }
+
+// -----
+
+// CHECK:       func.func @custom_origin_op()
+// CHECK-SAME:  -> !test.test_tensor<[42], f64> {
+func.func @custom_origin_op() -> !test.test_tensor<[42], f64> {
+  // CHECK: %[[MEMREF:.*]] = "test.create_memref_op"() : ()
+  // CHECK-SAME: -> !test.test_memref<[21], f64>
+  // CHECK: %[[DUMMY:.*]] = "test.dummy_memref_op"(%[[MEMREF]])
+  // CHECK-SAME: : (!test.test_memref<[21], f64>)
+  // CHECK-SAME: -> !test.test_memref<[42], f64>
+  %in = "test.create_tensor_op"() : () -> !test.test_tensor<[21], f64>
+  %out = "test.dummy_tensor_op"(%in) : (!test.test_tensor<[21], f64>)
+    -> !test.test_tensor<[42], f64>
+
+  // CHECK: %[[OUT:.*]] = bufferization.to_tensor %[[DUMMY]]
+  // CHECK: return %[[OUT]]
+  return %out : !test.test_tensor<[42], f64>
+}

mlir/test/lib/Dialect/Test/TestOpDefs.cpp

@@ -1418,3 +1418,35 @@ ::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
 
   return mlir::success();
 }
+
+::mlir::LogicalResult test::TestCreateTensorOp::bufferize(
+    ::mlir::RewriterBase &rewriter,
+    const ::mlir::bufferization::BufferizationOptions &options) {
+  // Note: mlir::bufferization::getBufferType() would internally call
+  // TestCreateTensorOp::getBufferType()
+  const auto bufferizedOutType =
+      mlir::bufferization::getBufferType(getOutput(), options);
+  if (mlir::failed(bufferizedOutType))
+    return failure();
+
+  // replace op with memref analogy
+  auto createMemrefOp =
+      rewriter.create<test::TestCreateMemrefOp>(getLoc(), *bufferizedOutType);
+
+  mlir::bufferization::replaceOpWithBufferizedValues(
+      rewriter, getOperation(), createMemrefOp.getResult());
+
+  return mlir::success();
+}
+
+mlir::FailureOr<mlir::bufferization::BufferLikeType>
+test::TestCreateTensorOp::getBufferType(
+    mlir::Value value, const mlir::bufferization::BufferizationOptions &,
+    llvm::SmallVector<::mlir::Value> &) {
+  const auto type = dyn_cast<test::TestTensorType>(value.getType());
+  if (type == nullptr)
+    return failure();
+
+  return cast<mlir::bufferization::BufferLikeType>(test::TestMemrefType::get(
+      getContext(), type.getShape(), type.getElementType(), nullptr));
+}

mlir/test/lib/Dialect/Test/TestOps.td

@@ -3603,28 +3603,16 @@ def TestMultiSlotAlloca : TEST_Op<"multi_slot_alloca",
 // Test Ops bufferization
 //===----------------------------------------------------------------------===//
 
-def TestDummyTensorOp : TEST_Op<"dummy_tensor_op", [BufferizableOpInterface]> {
+def TestDummyTensorOp : TEST_Op<"dummy_tensor_op",
+    [DeclareOpInterfaceMethods<BufferizableOpInterface,
+        ["bufferize", "bufferizesToMemoryRead",
+         "bufferizesToMemoryWrite", "getAliasingValues"]>]> {
   let arguments = (ins
     Arg<TestTensorType>:$input
   );
   let results = (outs
     Arg<TestTensorType>:$output
   );
-  let extraClassDeclaration = [{
-    // BufferizableOpInterface
-    bool bufferizesToMemoryRead(mlir::OpOperand&,
-      const mlir::bufferization::AnalysisState&);
-
-    bool bufferizesToMemoryWrite(mlir::OpOperand&,
-      const mlir::bufferization::AnalysisState&);
-
-    mlir::bufferization::AliasingValueList getAliasingValues(mlir::OpOperand&,
-      const mlir::bufferization::AnalysisState&);
-
-    mlir::LogicalResult bufferize(
-      mlir::RewriterBase& rewriter,
-      const mlir::bufferization::BufferizationOptions& options);
-  }];
 
   let extraClassDefinition = [{
     bool test::TestDummyTensorOp::bufferizesToMemoryRead(::mlir::OpOperand&,
@@ -3652,4 +3640,37 @@ def TestDummyMemrefOp : TEST_Op<"dummy_memref_op", []> {
   );
 }
 
+def TestCreateTensorOp : TEST_Op<"create_tensor_op",
+    [DeclareOpInterfaceMethods<BufferizableOpInterface,
+        ["bufferize", "getBufferType", "bufferizesToMemoryRead",
+         "bufferizesToMemoryWrite", "getAliasingValues",
+         "bufferizesToAllocation"]>]> {
+  let arguments = (ins);
+  let results = (outs Arg<TestTensorType>:$output);
+  let extraClassDefinition = [{
+    bool test::TestCreateTensorOp::bufferizesToMemoryRead(::mlir::OpOperand&,
+        const ::mlir::bufferization::AnalysisState&) {
+      return true;
+    }
+    bool test::TestCreateTensorOp::bufferizesToMemoryWrite(::mlir::OpOperand&,
+        const ::mlir::bufferization::AnalysisState&) {
+      return true;
+    }
+    bool test::TestCreateTensorOp::bufferizesToAllocation(mlir::Value value) {
+      return false;
+    }
+
+    ::mlir::bufferization::AliasingValueList
+    test::TestCreateTensorOp::getAliasingValues(::mlir::OpOperand&,
+        const ::mlir::bufferization::AnalysisState&) {
+      return {};
+    }
+  }];
+}
+
+def TestCreateMemrefOp : TEST_Op<"create_memref_op"> {
+  let arguments = (ins);
+  let results = (outs Arg<TestMemrefType>:$output);
+}
+
 #endif // TEST_OPS