Allow model builder to use address space for MemRef.

PiperOrigin-RevId: 300244466

experimental/ModelBuilder/ModelBuilder.cpp

diff --git a/experimental/ModelBuilder/ModelBuilder.cpp b/experimental/ModelBuilder/ModelBuilder.cpp
index f01aa02..ec3a831 100644
--- a/experimental/ModelBuilder/ModelBuilder.cpp
+++ b/experimental/ModelBuilder/ModelBuilder.cpp
@@ -53,8 +53,9 @@
 }
 
 MemRefType mlir::ModelBuilder::getMemRefType(ArrayRef<int64_t> shape,
-                                             Type elementType) {
-  return MemRefType::get(shape, elementType, {});
+                                             Type elementType,
+                                             unsigned addressSpace) {
+  return MemRefType::get(shape, elementType, {}, addressSpace);
 }
 
 RankedTensorType mlir::ModelBuilder::getRankedTensorType(

experimental/ModelBuilder/ModelBuilder.h

diff --git a/experimental/ModelBuilder/ModelBuilder.h b/experimental/ModelBuilder/ModelBuilder.h
index 5b66d14..ff5513a 100644
--- a/experimental/ModelBuilder/ModelBuilder.h
+++ b/experimental/ModelBuilder/ModelBuilder.h
@@ -106,7 +106,8 @@
   // and contiguous layout.
   // In the future, this can be extended support more advanced layouts, on a
   // per-need basis.
-  MemRefType getMemRefType(ArrayRef<int64_t> shape, Type elementType);
+  MemRefType getMemRefType(ArrayRef<int64_t> shape, Type elementType,
+                           unsigned addressSpace = 0);
 
   // Build an MLIR RankedTensorType with a base `elementType` and a `shape` that
   // can be any mix of static and dynamic values. For now this only supports a