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

[Andrei Golubev]

================
@@ -49,29 +49,47 @@ void FuncAnalysisState::startFunctionAnalysis(FuncOp funcOp) {
 #endif // NDEBUG
 }
 
+// Note: this is a local adaptor to unify TensorType and TensorLikeType code
+// paths that both work with BufferizationOptions.
+static mlir::Attribute
+getDefaultMemorySpace(const BufferizationOptions &options,
+                      TensorLikeType type) {
+  if (auto tensorType = dyn_cast<TensorType>(type)) {
+    return *options.defaultMemorySpaceFn(tensorType);
+  }
+  return nullptr;
+}
----------------
andrey-golubev wrote:

note: i guess we could also make `defaultMemorySpaceFn` work with tensor-like type. but i'd prefer to do it separately to simplify the updates for the users.

otoh, TensorLikeType::getBufferType() already has access to the tensor type and can infer memory space from it. I am not sure why we need it explicitly at all.

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