[Mlir-commits] [mlir] [mlir][bufferization] Alternate op bufferization approach (PR #199898)

[Andrei Golubev]

================
@@ -229,25 +229,32 @@ AllocTensorOp::getBufferType(Value value, const BufferizationOptions &options,
                              SmallVector<Value> &invocationStack) {
   assert(value == getResult() && "invalid value");
 
-  // Compute memory space of this allocation.
-  Attribute memorySpace;
-  if (getMemorySpace().has_value()) {
-    memorySpace = *getMemorySpace();
-  } else if (getCopy()) {
-    auto copyBufferType =
-        bufferization::detail::asMemRefType(bufferization::getBufferType(
-            getCopy(), options, state, invocationStack));
-    if (failed(copyBufferType))
-      return failure();
-    memorySpace = copyBufferType->getMemorySpace();
-  } else if (auto ms = options.defaultMemorySpaceFn(getType())) {
-    memorySpace = *ms;
-  } else {
-    return getOperation()->emitError("could not infer memory space");
-  }
+  const auto defaultGetBufferType =
+      [&](TensorLikeType tensorLikeType) -> FailureOr<BufferLikeType> {
+    auto tensorType = cast<TensorType>(tensorLikeType);
+    // Compute memory space of this allocation.
+    Attribute memorySpace;
+    if (getMemorySpace().has_value()) {
+      memorySpace = *getMemorySpace();
+    } else if (getCopy()) {
+      auto copyBufferType =
+          bufferization::detail::asMemRefType(bufferization::getBufferType(
+              getCopy(), options, state, invocationStack));
+      if (failed(copyBufferType))
+        return failure();
+      memorySpace = copyBufferType->getMemorySpace();
+    } else if (auto ms = options.defaultMemorySpaceFn(tensorType)) {
+      memorySpace = *ms;
+    } else {
+      return getOperation()->emitError("could not infer memory space");
+    }
+
+    return cast<BufferLikeType>(
+        getMemRefTypeWithStaticIdentityLayout(tensorType, memorySpace));
+  };
 
-  return cast<BufferLikeType>(
-      getMemRefTypeWithStaticIdentityLayout(getType(), memorySpace));
+  return cast<TensorLikeType>(getType()).getBufferType(
+      options, [&]() { return emitError(); }, defaultGetBufferType);
----------------
andrey-golubev wrote:

@matthias-springer this would be the way I'd love to see bufferization done (or somehow similarly) for "owners" of values. what we have from this is 2 things:
a) custom type bufferization (e.g. for non-upstream-tensors) is now "natively" supported - i.e. we don't need any special dispatch, it just works due to an interface
b) "unknown" encoding causes unknown type conversion behind the scenes for upstream tensor; and users can define what "unknown encoding" means also.

at the same time, the "fallback" bufferization also exists and can still be used if desired.

I am not 100% sure about the new API, but this feels to be the direction.

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