[mlir][memref]: Fix Bug in GlobalOp Verifier

When reconstructing the corresponding tensor type of a memref ensure we
include the memory space of the tensor if it exists.

Signed-off-by: Jack Frankland <[email protected]>

Author: FranklandJack

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

@@ -59,7 +59,8 @@ LogicalResult mlir::memref::foldMemRefCast(Operation *op, Value inner) {
 /// type.
 Type mlir::memref::getTensorTypeFromMemRefType(Type type) {
   if (auto memref = llvm::dyn_cast<MemRefType>(type))
-    return RankedTensorType::get(memref.getShape(), memref.getElementType());
+    return RankedTensorType::get(memref.getShape(), memref.getElementType(),
+                                 memref.getMemorySpace());
   if (auto memref = llvm::dyn_cast<UnrankedMemRefType>(type))
     return UnrankedTensorType::get(memref.getElementType());
   return NoneType::get(type.getContext());

mlir/test/Dialect/MemRef/invalid.mlir

@@ -342,6 +342,11 @@ memref.global "priate" constant @memref5 : memref<2xf32>  = uninitialized
 
 // -----
 
+// expected-error @+1 {{op initial value expected to be of type 'tensor<1xf16>', but was of type 'tensor<1xf16, 1 : i32>'}}
+"memref.global"() <{constant, initial_value = dense<1.000000e+00> : tensor<1xf16, 1 : i32>, sym_name = "memref6", sym_visibility = "private", type = memref<1xf16>}> : () -> ()
+
+// -----
+
 func.func @nonexistent_global_memref() {
   // expected-error @+1 {{'gv' does not reference a valid global memref}}
   %0 = memref.get_global @gv : memref<3xf32>

mlir/test/Dialect/MemRef/ops.mlir

@@ -174,6 +174,9 @@ memref.global "private" @memref3 : memref<2xf32>  = uninitialized
 // CHECK-LABEL: memref.global "private" constant @memref4 : memref<2xf32> = uninitialized
 memref.global "private" constant @memref4 : memref<2xf32>  = uninitialized
 
+// CHECK-LABEL: memref.global "private" constant @memref5 : memref<1xf16, 1 : i32> = dense<1.000000e+00>
+"memref.global"() <{constant, initial_value = dense<1.000000e+00> : tensor<1xf16, 1 : i32>, sym_name = "memref5", sym_visibility = "private", type = memref<1xf16, 1 : i32>}> : () -> ()
+
 // CHECK-LABEL: func @read_global_memref
 func.func @read_global_memref() {
   %0 = memref.get_global @memref0 : memref<2xf32>