[Mlir-commits] [mlir] 4511128 - [mlir][bufferization] Consider defaultMemorySpace when bufferizing FuncOp

[Maya Amrami]

Author: Maya Amrami
Date: 2023-02-22T16:44:16+02:00
New Revision: 4511128fc903378ebafd10cb06dc5ae1753b68fa

URL: https://github.com/llvm/llvm-project/commit/4511128fc903378ebafd10cb06dc5ae1753b68fa
DIFF: https://github.com/llvm/llvm-project/commit/4511128fc903378ebafd10cb06dc5ae1753b68fa.diff

LOG: [mlir][bufferization] Consider defaultMemorySpace when bufferizing FuncOp

The function arguments and results type will have the default memory space.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D144539

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 483d1008b1dc0..1c5767973960c 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -67,11 +67,13 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
   BaseMemRefType memrefType;
   if (options.functionBoundaryTypeConversion ==
       LayoutMapOption::IdentityLayoutMap) {
-    memrefType = getMemRefTypeWithStaticIdentityLayout(tensorType);
+    memrefType = getMemRefTypeWithStaticIdentityLayout(
+        tensorType, *options.defaultMemorySpace);
   } else {
     // Note: Layout maps on function parameters cannot be inferred. The best we
     // can do at the moment is "fully dynamic".
-    memrefType = getMemRefTypeWithFullyDynamicLayout(tensorType);
+    memrefType = getMemRefTypeWithFullyDynamicLayout(
+        tensorType, *options.defaultMemorySpace);
   }
 
   auto layoutAttr = funcOp.getArgAttrOfType<AffineMapAttr>(
@@ -424,10 +426,12 @@ struct FuncOpInterface
       BaseMemRefType resultType;
       if (options.functionBoundaryTypeConversion ==
           LayoutMapOption::IdentityLayoutMap) {
-        resultType = getMemRefTypeWithStaticIdentityLayout(tensorType);
+        resultType = getMemRefTypeWithStaticIdentityLayout(
+            tensorType, *options.defaultMemorySpace);
       } else {
         // Note: If `InferLayoutMap`, cast are later folded away.
-        resultType = getMemRefTypeWithFullyDynamicLayout(tensorType);
+        resultType = getMemRefTypeWithFullyDynamicLayout(
+            tensorType, *options.defaultMemorySpace);
       }
       Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
           loc, resultType, returnVal);