Fix

Author: rayngun

mlir/test/lib/Dialect/Test/TestDialect.cpp

@@ -482,11 +482,11 @@ mlir::LogicalResult test::TestDummyTensorOp::bufferize(
     const mlir::bufferization::BufferizationOptions &options) {
   const auto inputType = getInput().getType();
   const auto bufferizedInputType = test::TestMemrefType::get(
-      getContext(), inputType.getShape(), inputType.getElementType(), nullptr);
+      getContext(), inputType.getShape(), inputType.getElementType(), nullptr, 0);
   const auto outputType = getOutput().getType();
   const auto bufferizedOutputType =
       test::TestMemrefType::get(getContext(), outputType.getShape(),
-                                outputType.getElementType(), nullptr);
+                                outputType.getElementType(), nullptr, 0);
 
   // replace op with memref analogy, preserve correct types at the boundaries
   auto toMemref = rewriter.create<mlir::bufferization::ToMemrefOp>(

mlir/test/lib/Dialect/Test/TestTypeDefs.td

@@ -421,23 +421,27 @@ def TestMemrefType : Test_Type<"TestMemref", [BaseMemRefTypeInterface]> {
   let parameters = (ins
     ArrayRefParameter<"int64_t">:$shape,
     "mlir::Type":$elementType,
-    DefaultValuedParameter<"mlir::Attribute", "nullptr">:$memSpace
+    DefaultValuedParameter<"mlir::Attribute", "nullptr">:$memSpace,
+    DefaultValuedParameter<"unsigned", "0">:$memSpaceAsInt
   );
-  let assemblyFormat = "`<` `[` $shape `]` `,` $elementType (`,` $memSpace^)? `>`";
+  let assemblyFormat = "`<` `[` $shape `]` `,` $elementType (`,` $memSpace^)? (`,` $memSpaceAsInt^)? `>`";
 
   let extraClassDeclaration = [{
     // ShapedTypeInterface:
     bool hasRank() const {
       return true;
     }
     mlir::ShapedType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape, mlir::Type elementType) const {
-      return test::TestMemrefType::get(getContext(), shape.value_or(getShape()), elementType, getMemSpace());
+      return test::TestMemrefType::get(getContext(), shape.value_or(getShape()), elementType, getMemSpace(), getMemSpaceAsInt());
     }
 
     // BaseMemRefTypeInterface:
     mlir::Attribute getMemorySpace() const {
       return getMemSpace();
     }
+    unsigned getMemorySpaceAsInt() const {
+      return getMemSpaceAsInt();
+    }
   }];
 }