[Mlir-commits] [mlir] [mlir][bufferization] Support custom types at function boundaries (PR #159766)

[Andrei Golubev]

================
@@ -87,6 +87,19 @@ struct BuiltinTensorExternalModel
 
     return mlir::success();
   }
+
+  llvm::FailureOr<BufferLikeType> getBufferTypeAtFunctionBoundary(
+      mlir::Type tensor, mlir::func::FuncOp funcOp,
+      const BufferizationOptions &options,
+      llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
+    auto tensorType = cast<TensorType>(tensor);
+    auto memSpace = options.defaultMemorySpaceFn(tensorType);
+    if (!memSpace.has_value())
+      return emitError() << "could not infer memory space";
+
+    return cast<BufferLikeType>(options.functionArgTypeConverterFn(
+        tensorType, *memSpace, funcOp, options));
+  }
----------------
andrey-golubev wrote:

side-note: there is a new behavior related to mem-space checking. Before, the logic would unconditionally access the value returned by `defaultMemorySpaceFn()`, assuming "no error".

I decided to make the getBufferTypeAtFunctionBoundary() signature aligned with normal getBuffer(). And thus checking for memspace existence seems like a reasonable addition.

Technically, we don't have to return a failure here, but instead can just ignore this case completely (as the original logic goes). Moreover, all current users of this API expect "always success".

https://github.com/llvm/llvm-project/pull/159766