diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -67,11 +67,13 @@
   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 @@
       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);