[mlir][Ptr] Add the MemorySpaceAttrInterface interface and dependencies. (#86870)

This patch introduces the `MemorySpaceAttrInterface` interface. This
interface is responsible for handling the semantics of `ptr` operations.

For example, this interface can be used to create read-only memory
spaces, making any other operation other than a load a verification
error, see `TestConstMemorySpaceAttr` for a possible implementation of
this concept.

This patch also introduces Enum dependencies `AtomicOrdering`, and
`AtomicBinOp`, both enumerations are clones of the Enums with the same
name in the LLVM Dialect.

Also, see:
- [[RFC] `ptr` dialect & modularizing ptr ops in the LLVM
dialect](https://discourse.llvm.org/t/rfc-ptr-dialect-modularizing-ptr-ops-in-the-llvm-dialect/75142)
for rationale.
- https://github.com/llvm/llvm-project/pull/73057 for a prototype
implementation of the full change.

**Note: Ignore the first commit, that's being reviewed in
https://github.com/llvm/llvm-project/pull/86860 .**

mlir/include/mlir/Dialect/Ptr/IR/CMakeLists.txt

@@ -5,3 +5,15 @@ set(LLVM_TARGET_DEFINITIONS PtrOps.td)
 mlir_tablegen(PtrOpsAttrs.h.inc -gen-attrdef-decls -attrdefs-dialect=ptr)
 mlir_tablegen(PtrOpsAttrs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=ptr)
 add_public_tablegen_target(MLIRPtrOpsAttributesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS MemorySpaceInterfaces.td)
+mlir_tablegen(MemorySpaceInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(MemorySpaceInterfaces.cpp.inc -gen-op-interface-defs)
+mlir_tablegen(MemorySpaceAttrInterfaces.h.inc -gen-attr-interface-decls)
+mlir_tablegen(MemorySpaceAttrInterfaces.cpp.inc -gen-attr-interface-defs)
+add_public_tablegen_target(MLIRPtrMemorySpaceInterfacesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS PtrOps.td)
+mlir_tablegen(PtrOpsEnums.h.inc -gen-enum-decls)
+mlir_tablegen(PtrOpsEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRPtrOpsEnumsGen)

mlir/include/mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h

@@ -0,0 +1,32 @@
+//===-- MemorySpaceInterfaces.h - ptr memory space interfaces ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the ptr dialect memory space interfaces.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_PTR_IR_MEMORYSPACEINTERFACES_H
+#define MLIR_DIALECT_PTR_IR_MEMORYSPACEINTERFACES_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+class Operation;
+namespace ptr {
+enum class AtomicBinOp : uint64_t;
+enum class AtomicOrdering : uint64_t;
+} // namespace ptr
+} // namespace mlir
+
+#include "mlir/Dialect/Ptr/IR/MemorySpaceAttrInterfaces.h.inc"
+
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h.inc"
+
+#endif // MLIR_DIALECT_PTR_IR_MEMORYSPACEINTERFACES_H

mlir/include/mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td

@@ -0,0 +1,117 @@
+//===-- MemorySpaceInterfaces.td - Memory space interfaces ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines memory space attribute interfaces.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PTR_MEMORYSPACEINTERFACES
+#define PTR_MEMORYSPACEINTERFACES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// Memory space attribute interface.
+//===----------------------------------------------------------------------===//
+
+def MemorySpaceAttrInterface : AttrInterface<"MemorySpaceAttrInterface"> {
+  let description = [{
+    This interface defines a common API for interacting with the memory model of
+    a memory space and the operations in the pointer dialect.
+
+    Furthermore, this interface allows concepts such as read-only memory to be
+    adequately modeled and enforced.
+  }];
+  let cppNamespace = "::mlir::ptr";
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/        [{
+        This method checks if it's valid to load a value from the memory space
+        with a specific type, alignment, and atomic ordering.
+        If `emitError` is non-null then the method is allowed to emit errors.
+      }],
+      /*returnType=*/  "::mlir::LogicalResult",
+      /*methodName=*/  "isValidLoad",
+      /*args=*/        (ins "::mlir::Type":$type,
+                            "::mlir::ptr::AtomicOrdering":$ordering,
+                            "::mlir::IntegerAttr":$alignment,
+                            "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >,
+    InterfaceMethod<
+      /*desc=*/        [{
+        This method checks if it's valid to store a value in the memory space
+        with a specific type, alignment, and atomic ordering.
+        If `emitError` is non-null then the method is allowed to emit errors.
+      }],
+      /*returnType=*/  "::mlir::LogicalResult",
+      /*methodName=*/  "isValidStore",
+      /*args=*/        (ins "::mlir::Type":$type,
+                            "::mlir::ptr::AtomicOrdering":$ordering,
+                            "::mlir::IntegerAttr":$alignment,
+                            "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >,
+    InterfaceMethod<
+      /*desc=*/        [{
+        This method checks if it's valid to perform an atomic operation in the
+        memory space with a specific type, alignment, and atomic ordering.
+        If `emitError` is non-null then the method is allowed to emit errors.
+      }],
+      /*returnType=*/  "::mlir::LogicalResult",
+      /*methodName=*/  "isValidAtomicOp",
+      /*args=*/        (ins "::mlir::ptr::AtomicBinOp":$op,
+                            "::mlir::Type":$type,
+                            "::mlir::ptr::AtomicOrdering":$ordering,
+                            "::mlir::IntegerAttr":$alignment,
+                            "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >,
+    InterfaceMethod<
+      /*desc=*/        [{
+        This method checks if it's valid to perform an atomic exchange operation
+        in the memory space with a specific type, alignment, and atomic
+        orderings.
+        If `emitError` is non-null then the method is allowed to emit errors.
+      }],
+      /*returnType=*/  "::mlir::LogicalResult",
+      /*methodName=*/  "isValidAtomicXchg",
+      /*args=*/        (ins "::mlir::Type":$type,
+                            "::mlir::ptr::AtomicOrdering":$successOrdering,
+                            "::mlir::ptr::AtomicOrdering":$failureOrdering,
+                            "::mlir::IntegerAttr":$alignment,
+                            "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >,
+    InterfaceMethod<
+      /*desc=*/        [{
+        This method checks if it's valid to perform an `addrspacecast` op
+        in the memory space.
+        If `emitError` is non-null then the method is allowed to emit errors.
+      }],
+      /*returnType=*/  "::mlir::LogicalResult",
+      /*methodName=*/  "isValidAddrSpaceCast",
+      /*args=*/        (ins "::mlir::Type":$tgt,
+                            "::mlir::Type":$src,
+                            "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >,
+    InterfaceMethod<
+      /*desc=*/        [{
+        This method checks if it's valid to perform a `ptrtoint` or `inttoptr`
+        op in the memory space.
+        The first type is expected to be integer-like, while the second must be a
+        ptr-like type.
+        If `emitError` is non-null then the method is allowed to emit errors.
+      }],
+      /*returnType=*/  "::mlir::LogicalResult",
+      /*methodName=*/  "isValidPtrIntCast",
+      /*args=*/        (ins "::mlir::Type":$intLikeTy,
+                            "::mlir::Type":$ptrLikeTy,
+                            "::llvm::function_ref<::mlir::InFlightDiagnostic()>":$emitError)
+    >,
+  ];
+}
+
+#endif // PTR_MEMORYSPACEINTERFACES

mlir/include/mlir/Dialect/Ptr/IR/PtrAttrs.h

@@ -18,4 +18,6 @@
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOpsAttrs.h.inc"
 
+#include "mlir/Dialect/Ptr/IR/PtrOpsEnums.h.inc"
+
 #endif // MLIR_DIALECT_PTR_IR_PTRATTRS_H

mlir/include/mlir/Dialect/Ptr/IR/PtrDialect.td

@@ -45,7 +45,7 @@ def Ptr_PtrType : Ptr_Type<"Ptr", "ptr", [
   let description = [{
     The `ptr` type is an opaque pointer type. This type typically represents a
     handle to an object in memory or target-dependent values like `nullptr`.
-    Pointers are optionally parameterized by a memory space.
+    Pointers are parameterized by a memory space.
 
     Syntax:
 
@@ -54,14 +54,14 @@ def Ptr_PtrType : Ptr_Type<"Ptr", "ptr", [
     memory-space ::= attribute-value
     ```
   }];
-  let parameters = (ins OptionalParameter<"Attribute">:$memorySpace);
-  let assemblyFormat = "(`<` $memorySpace^ `>`)?";
+  let parameters = (ins "MemorySpaceAttrInterface":$memorySpace);
+  let assemblyFormat = "`<` $memorySpace `>`";
   let builders = [
-    TypeBuilder<(ins CArg<"Attribute", "nullptr">:$memorySpace), [{
-      return $_get($_ctxt, memorySpace);
+    TypeBuilderWithInferredContext<(ins
+      "MemorySpaceAttrInterface":$memorySpace), [{
+      return $_get(memorySpace.getContext(), memorySpace);
     }]>
   ];
-  let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/Ptr/IR/PtrEnums.td

@@ -0,0 +1,69 @@
+//===-- PtrEnums.td - Ptr dialect enumerations -------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PTR_ENUMS
+#define PTR_ENUMS
+
+include "mlir/IR/EnumAttr.td"
+
+//===----------------------------------------------------------------------===//
+// Atomic binary op enum attribute.
+//===----------------------------------------------------------------------===//
+
+def AtomicBinOpXchg : I64EnumAttrCase<"xchg", 0, "xchg">;
+def AtomicBinOpAdd  : I64EnumAttrCase<"add", 1, "add">;
+def AtomicBinOpSub  : I64EnumAttrCase<"sub", 2, "sub">;
+def AtomicBinOpAnd  : I64EnumAttrCase<"_and", 3, "_and">;
+def AtomicBinOpNand : I64EnumAttrCase<"nand", 4, "nand">;
+def AtomicBinOpOr   : I64EnumAttrCase<"_or", 5, "_or">;
+def AtomicBinOpXor  : I64EnumAttrCase<"_xor", 6, "_xor">;
+def AtomicBinOpMax  : I64EnumAttrCase<"max", 7, "max">;
+def AtomicBinOpMin  : I64EnumAttrCase<"min", 8, "min">;
+def AtomicBinOpUMax : I64EnumAttrCase<"umax", 9, "umax">;
+def AtomicBinOpUMin : I64EnumAttrCase<"umin", 10, "umin">;
+def AtomicBinOpFAdd : I64EnumAttrCase<"fadd", 11, "fadd">;
+def AtomicBinOpFSub : I64EnumAttrCase<"fsub", 12, "fsub">;
+def AtomicBinOpFMax : I64EnumAttrCase<"fmax", 13, "fmax">;
+def AtomicBinOpFMin : I64EnumAttrCase<"fmin", 14, "fmin">;
+def AtomicBinOpUIncWrap : I64EnumAttrCase<"uinc_wrap", 15, "uinc_wrap">;
+def AtomicBinOpUDecWrap : I64EnumAttrCase<"udec_wrap", 16, "udec_wrap">;
+
+def AtomicBinOp : I64EnumAttr<
+    "AtomicBinOp",
+    "ptr.atomicrmw binary operations",
+    [AtomicBinOpXchg, AtomicBinOpAdd, AtomicBinOpSub, AtomicBinOpAnd,
+     AtomicBinOpNand, AtomicBinOpOr, AtomicBinOpXor, AtomicBinOpMax,
+     AtomicBinOpMin, AtomicBinOpUMax, AtomicBinOpUMin, AtomicBinOpFAdd,
+     AtomicBinOpFSub, AtomicBinOpFMax, AtomicBinOpFMin, AtomicBinOpUIncWrap,
+     AtomicBinOpUDecWrap]> {
+  let cppNamespace = "::mlir::ptr";
+}
+
+//===----------------------------------------------------------------------===//
+// Atomic ordering enum attribute.
+//===----------------------------------------------------------------------===//
+
+def AtomicOrderingNotAtomic : I64EnumAttrCase<"not_atomic", 0, "not_atomic">;
+def AtomicOrderingUnordered : I64EnumAttrCase<"unordered", 1, "unordered">;
+def AtomicOrderingMonotonic : I64EnumAttrCase<"monotonic", 2, "monotonic">;
+def AtomicOrderingAcquire   : I64EnumAttrCase<"acquire", 3, "acquire">;
+def AtomicOrderingRelease   : I64EnumAttrCase<"release", 4, "release">;
+def AtomicOrderingAcqRel    : I64EnumAttrCase<"acq_rel", 5, "acq_rel">;
+def AtomicOrderingSeqCst    : I64EnumAttrCase<"seq_cst", 6, "seq_cst">;
+
+def AtomicOrdering : I64EnumAttr<
+    "AtomicOrdering",
+    "Atomic ordering for LLVM's memory model",
+    [AtomicOrderingNotAtomic, AtomicOrderingUnordered, AtomicOrderingMonotonic,
+     AtomicOrderingAcquire, AtomicOrderingRelease, AtomicOrderingAcqRel,
+     AtomicOrderingSeqCst
+    ]> {
+  let cppNamespace = "::mlir::ptr";
+}
+
+#endif // PTR_ENUMS

mlir/include/mlir/Dialect/Ptr/IR/PtrOps.td

@@ -11,6 +11,7 @@
 
 include "mlir/Dialect/Ptr/IR/PtrDialect.td"
 include "mlir/Dialect/Ptr/IR/PtrAttrDefs.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 
 #endif // PTR_OPS

mlir/include/mlir/Dialect/Ptr/IR/PtrTypes.h

@@ -13,6 +13,7 @@
 #ifndef MLIR_DIALECT_PTR_IR_PTRTYPES_H
 #define MLIR_DIALECT_PTR_IR_PTRTYPES_H
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 

mlir/lib/Dialect/Ptr/IR/CMakeLists.txt

@@ -7,7 +7,8 @@ add_mlir_dialect_library(
   DEPENDS
   MLIRPtrOpsAttributesIncGen
   MLIRPtrOpsIncGen
-
+  MLIRPtrOpsEnumsGen
+  MLIRPtrMemorySpaceInterfacesIncGen
   LINK_LIBS
   PUBLIC
   MLIRIR

mlir/lib/Dialect/Ptr/IR/PtrDialect.cpp

@@ -48,6 +48,12 @@ void PtrDialect::initialize() {
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOpsAttrs.cpp.inc"
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.cpp.inc"
+
+#include "mlir/Dialect/Ptr/IR/MemorySpaceAttrInterfaces.cpp.inc"
+
+#include "mlir/Dialect/Ptr/IR/PtrOpsEnums.cpp.inc"
+
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/Ptr/IR/PtrOpsTypes.cpp.inc"
 

mlir/lib/Dialect/Ptr/IR/PtrTypes.cpp

@@ -28,7 +28,7 @@ constexpr const static unsigned kDefaultPointerAlignmentBits = 8;
 /// Searches the data layout for the pointer spec, returns nullptr if it is not
 /// found.
 static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type,
-                               Attribute defaultMemorySpace) {
+                               MemorySpaceAttrInterface defaultMemorySpace) {
   for (DataLayoutEntryInterface entry : params) {
     if (!entry.isTypeEntry())
       continue;
@@ -38,9 +38,11 @@ static SpecAttr getPointerSpec(DataLayoutEntryListRef params, PtrType type,
         return spec;
     }
   }
-  // If not found, and this is the pointer to the default memory space, assume
-  // 64-bit pointers.
-  if (type.getMemorySpace() == defaultMemorySpace)
+  // If not found, and this is the pointer to the default memory space or if
+  // `defaultMemorySpace` is null, assume 64-bit pointers. `defaultMemorySpace`
+  // might be null if the data layout doesn't define the default memory space.
+  if (type.getMemorySpace() == defaultMemorySpace ||
+      defaultMemorySpace == nullptr)
     return SpecAttr::get(type.getContext(), kDefaultPointerSizeBits,
                          kDefaultPointerAlignmentBits,
                          kDefaultPointerAlignmentBits, kDefaultPointerSizeBits);
@@ -93,44 +95,47 @@ bool PtrType::areCompatible(DataLayoutEntryListRef oldLayout,
 
 uint64_t PtrType::getABIAlignment(const DataLayout &dataLayout,
                                   DataLayoutEntryListRef params) const {
-  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  auto defaultMemorySpace = llvm::cast_if_present<MemorySpaceAttrInterface>(
+      dataLayout.getDefaultMemorySpace());
   if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
     return spec.getAbi() / kBitsInByte;
 
-  return dataLayout.getTypeABIAlignment(get(getContext(), defaultMemorySpace));
+  return dataLayout.getTypeABIAlignment(get(defaultMemorySpace));
 }
 
 std::optional<uint64_t>
 PtrType::getIndexBitwidth(const DataLayout &dataLayout,
                           DataLayoutEntryListRef params) const {
-  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  auto defaultMemorySpace = llvm::cast_if_present<MemorySpaceAttrInterface>(
+      dataLayout.getDefaultMemorySpace());
   if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace)) {
     return spec.getIndex() == SpecAttr::kOptionalSpecValue ? spec.getSize()
                                                            : spec.getIndex();
   }
 
-  return dataLayout.getTypeIndexBitwidth(get(getContext(), defaultMemorySpace));
+  return dataLayout.getTypeIndexBitwidth(get(defaultMemorySpace));
 }
 
 llvm::TypeSize PtrType::getTypeSizeInBits(const DataLayout &dataLayout,
                                           DataLayoutEntryListRef params) const {
-  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  auto defaultMemorySpace = llvm::cast_if_present<MemorySpaceAttrInterface>(
+      dataLayout.getDefaultMemorySpace());
   if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
     return llvm::TypeSize::getFixed(spec.getSize());
 
   // For other memory spaces, use the size of the pointer to the default memory
   // space.
-  return dataLayout.getTypeSizeInBits(get(getContext(), defaultMemorySpace));
+  return dataLayout.getTypeSizeInBits(get(defaultMemorySpace));
 }
 
 uint64_t PtrType::getPreferredAlignment(const DataLayout &dataLayout,
                                         DataLayoutEntryListRef params) const {
-  Attribute defaultMemorySpace = dataLayout.getDefaultMemorySpace();
+  auto defaultMemorySpace = llvm::cast_if_present<MemorySpaceAttrInterface>(
+      dataLayout.getDefaultMemorySpace());
   if (SpecAttr spec = getPointerSpec(params, *this, defaultMemorySpace))
     return spec.getPreferred() / kBitsInByte;
 
-  return dataLayout.getTypePreferredAlignment(
-      get(getContext(), defaultMemorySpace));
+  return dataLayout.getTypePreferredAlignment(get(defaultMemorySpace));
 }
 
 LogicalResult PtrType::verifyEntries(DataLayoutEntryListRef entries,

mlir/test/Dialect/Ptr/layout.mlir

@@ -1,61 +1,61 @@
 // RUN: mlir-opt --test-data-layout-query --split-input-file --verify-diagnostics %s | FileCheck %s
 
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 32, abi = 32, preferred = 64>>,
-  #dlti.dl_entry<!ptr.ptr<5>,#ptr.spec<size = 64, abi = 64, preferred = 64>>,
-  #dlti.dl_entry<!ptr.ptr<4>, #ptr.spec<size = 32, abi = 64, preferred = 64, index = 24>>,
-  #dlti.dl_entry<"dlti.default_memory_space", 7 : ui64>,
-  #dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui64>,
-  #dlti.dl_entry<"dlti.global_memory_space", 2 : ui64>,
-  #dlti.dl_entry<"dlti.program_memory_space", 3 : ui64>,
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 32, abi = 32, preferred = 64>>,
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space<5>>,#ptr.spec<size = 64, abi = 64, preferred = 64>>,
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space<4>>, #ptr.spec<size = 32, abi = 64, preferred = 64, index = 24>>,
+  #dlti.dl_entry<"dlti.default_memory_space", #test.const_memory_space<7>>,
+  #dlti.dl_entry<"dlti.alloca_memory_space", #test.const_memory_space<5>>,
+  #dlti.dl_entry<"dlti.global_memory_space", #test.const_memory_space<2>>,
+  #dlti.dl_entry<"dlti.program_memory_space", #test.const_memory_space<3>>,
   #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>
 >} {
   // CHECK-LABEL: @spec
   func.func @spec() {
     // CHECK: alignment = 4
-    // CHECK: alloca_memory_space = 5
+    // CHECK: alloca_memory_space = #test.const_memory_space<5>
     // CHECK: bitsize = 32
-    // CHECK: default_memory_space = 7
-    // CHECK: global_memory_space = 2
+    // CHECK: default_memory_space = #test.const_memory_space<7>
+    // CHECK: global_memory_space = #test.const_memory_space<2>
     // CHECK: index = 32
     // CHECK: preferred = 8
-    // CHECK: program_memory_space = 3
+    // CHECK: program_memory_space = #test.const_memory_space<3>
     // CHECK: size = 4
     // CHECK: stack_alignment = 128
-    "test.data_layout_query"() : () -> !ptr.ptr
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space>
     // CHECK: alignment = 1
-    // CHECK: alloca_memory_space = 5
+    // CHECK: alloca_memory_space = #test.const_memory_space<5>
     // CHECK: bitsize = 64
-    // CHECK: default_memory_space = 7
-    // CHECK: global_memory_space = 2
+    // CHECK: default_memory_space = #test.const_memory_space<7>
+    // CHECK: global_memory_space = #test.const_memory_space<2>
     // CHECK: index = 64
     // CHECK: preferred = 1
-    // CHECK: program_memory_space = 3
+    // CHECK: program_memory_space = #test.const_memory_space<3>
     // CHECK: size = 8
     // CHECK: stack_alignment = 128
-    "test.data_layout_query"() : () -> !ptr.ptr<3>
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space<3>>
     // CHECK: alignment = 8
-    // CHECK: alloca_memory_space = 5
+    // CHECK: alloca_memory_space = #test.const_memory_space<5>
     // CHECK: bitsize = 64
-    // CHECK: default_memory_space = 7
-    // CHECK: global_memory_space = 2
+    // CHECK: default_memory_space = #test.const_memory_space<7>
+    // CHECK: global_memory_space = #test.const_memory_space<2>
     // CHECK: index = 64
     // CHECK: preferred = 8
-    // CHECK: program_memory_space = 3
+    // CHECK: program_memory_space = #test.const_memory_space<3>
     // CHECK: size = 8
     // CHECK: stack_alignment = 128
-    "test.data_layout_query"() : () -> !ptr.ptr<5>
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space<5>>
     // CHECK: alignment = 8
-    // CHECK: alloca_memory_space = 5
+    // CHECK: alloca_memory_space = #test.const_memory_space<5>
     // CHECK: bitsize = 32
-    // CHECK: default_memory_space = 7
-    // CHECK: global_memory_space = 2
+    // CHECK: default_memory_space = #test.const_memory_space<7>
+    // CHECK: global_memory_space = #test.const_memory_space<2>
     // CHECK: index = 24
     // CHECK: preferred = 8
-    // CHECK: program_memory_space = 3
+    // CHECK: program_memory_space = #test.const_memory_space<3>
     // CHECK: size = 4
     // CHECK: stack_alignment = 128
-    "test.data_layout_query"() : () -> !ptr.ptr<4>
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space<4>>
     return
   }
 }
@@ -63,8 +63,8 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 // -----
 
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr<1 : ui64>, #ptr.spec<size = 32, abi = 32, preferred = 32>>,
-  #dlti.dl_entry<"dlti.default_memory_space", 1 : ui64>
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 32, abi = 32, preferred = 32>>,
+  #dlti.dl_entry<"dlti.default_memory_space", #test.const_memory_space>
 >} {
   // CHECK-LABEL: @default_memory_space
   func.func @default_memory_space() {
@@ -73,19 +73,19 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
     // CHECK: index = 32
     // CHECK: preferred = 4
     // CHECK: size = 4
-    "test.data_layout_query"() : () -> !ptr.ptr
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space>
     // CHECK: alignment = 4
     // CHECK: bitsize = 32
     // CHECK: index = 32
     // CHECK: preferred = 4
     // CHECK: size = 4
-    "test.data_layout_query"() : () -> !ptr.ptr<1>
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space<1>>
     // CHECK: alignment = 4
     // CHECK: bitsize = 32
     // CHECK: index = 32
     // CHECK: preferred = 4
     // CHECK: size = 4
-    "test.data_layout_query"() : () -> !ptr.ptr<2>
+    "test.data_layout_query"() : () -> !ptr.ptr<#test.const_memory_space<2>>
     return
   }
 }
@@ -94,7 +94,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 
 // expected-error@+2 {{preferred alignment is expected to be at least as large as ABI alignment}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 64, abi = 64, preferred = 32>>
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 64, abi = 64, preferred = 32>>
 >} {
   func.func @pointer() {
     return
@@ -105,7 +105,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 
 // expected-error@+2 {{size entry must be divisible by 8}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 33, abi = 32, preferred = 32>>
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 33, abi = 32, preferred = 32>>
 >} {
   func.func @pointer() {
     return
@@ -117,7 +117,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 
 // expected-error@+2 {{abi entry must be divisible by 8}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 32, abi = 33, preferred = 64>>
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 32, abi = 33, preferred = 64>>
 >} {
   func.func @pointer() {
     return
@@ -129,7 +129,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 
 // expected-error@+2 {{preferred entry must be divisible by 8}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 32, abi = 32, preferred = 33>>
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 32, abi = 32, preferred = 33>>
 >} {
   func.func @pointer() {
     return
@@ -141,7 +141,7 @@ module attributes { dlti.dl_spec = #dlti.dl_spec<
 
 // expected-error@+2 {{index entry must be divisible by 8}}
 module attributes { dlti.dl_spec = #dlti.dl_spec<
-  #dlti.dl_entry<!ptr.ptr, #ptr.spec<size = 32, abi = 32, preferred = 32, index = 33>>
+  #dlti.dl_entry<!ptr.ptr<#test.const_memory_space>, #ptr.spec<size = 32, abi = 32, preferred = 32, index = 33>>
 >} {
   func.func @pointer() {
     return

mlir/test/Dialect/Ptr/types.mlir

@@ -1,17 +1,24 @@
 // RUN: mlir-opt %s -split-input-file | mlir-opt | FileCheck %s
 
 // CHECK-LABEL: func @ptr_test
-// CHECK: (%[[ARG0:.*]]: !ptr.ptr, %[[ARG1:.*]]: !ptr.ptr<1 : i32>)
-// CHECK: -> (!ptr.ptr<1 : i32>, !ptr.ptr)
-func.func @ptr_test(%arg0: !ptr.ptr, %arg1: !ptr.ptr<1 : i32>) -> (!ptr.ptr<1 : i32>, !ptr.ptr) {
-  // CHECK: return %[[ARG1]], %[[ARG0]] : !ptr.ptr<1 : i32>, !ptr.ptr
-  return %arg1, %arg0 : !ptr.ptr<1 : i32>, !ptr.ptr
+// CHECK: (%[[ARG0:.*]]: !ptr.ptr<#test.const_memory_space>, %[[ARG1:.*]]: !ptr.ptr<#test.const_memory_space<1>>)
+// CHECK: -> (!ptr.ptr<#test.const_memory_space<1>>, !ptr.ptr<#test.const_memory_space>)
+func.func @ptr_test(%arg0: !ptr.ptr<#test.const_memory_space>, %arg1: !ptr.ptr<#test.const_memory_space<1>>) -> (!ptr.ptr<#test.const_memory_space<1>>, !ptr.ptr<#test.const_memory_space>) {
+  // CHECK: return %[[ARG1]], %[[ARG0]] : !ptr.ptr<#test.const_memory_space<1>>, !ptr.ptr<#test.const_memory_space>
+  return %arg1, %arg0 : !ptr.ptr<#test.const_memory_space<1>>, !ptr.ptr<#test.const_memory_space>
 }
 
 // -----
 
 // CHECK-LABEL: func @ptr_test
-// CHECK: %[[ARG:.*]]: memref<!ptr.ptr>
-func.func @ptr_test(%arg0: memref<!ptr.ptr>) {
+// CHECK: %[[ARG:.*]]: memref<!ptr.ptr<#test.const_memory_space>>
+func.func @ptr_test(%arg0: memref<!ptr.ptr<#test.const_memory_space>>) {
+  return
+}
+
+// CHECK-LABEL: func @ptr_test_1
+// CHECK: (%[[ARG0:.*]]: !ptr.ptr<#test.const_memory_space>, %[[ARG1:.*]]: !ptr.ptr<#test.const_memory_space<3>>)
+func.func @ptr_test_1(%arg0: !ptr.ptr<#test.const_memory_space>,
+                      %arg1: !ptr.ptr<#test.const_memory_space<3>>) {
   return
 }

mlir/test/lib/Dialect/Test/CMakeLists.txt

@@ -84,6 +84,7 @@ mlir_target_link_libraries(MLIRTestDialect PUBLIC
   MLIRInferTypeOpInterface
   MLIRLinalgDialect
   MLIRLinalgTransforms
+  MLIRPtrDialect
   MLIRLLVMDialect
   MLIRPass
   MLIRPolynomialDialect

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

@@ -17,6 +17,7 @@
 include "TestDialect.td"
 include "TestEnumDefs.td"
 include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
 include "mlir/Dialect/Utils/StructuredOpsUtils.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
@@ -384,6 +385,14 @@ def NestedPolynomialAttr2 : Test_Attr<"NestedPolynomialAttr2"> {
   }];
 }
 
+// Test a ptr constant memory space.
+def TestConstMemorySpaceAttr : Test_Attr<"TestConstMemorySpace", [
+    DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+  ]> {
+  let mnemonic = "const_memory_space";
+  let parameters = (ins DefaultValuedParameter<"unsigned", "0">:$addressSpace);
+  let assemblyFormat = "(`<` $addressSpace^ `>`)?";
+}
 
 // Test custom location handling.
 def TestCustomLocationAttr : Test_LocAttr<"TestCustomLocation"> {

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

@@ -327,6 +327,49 @@ TestOpAsmAttrInterfaceAttr::getAlias(::llvm::raw_ostream &os) const {
   return ::mlir::OpAsmDialectInterface::AliasResult::FinalAlias;
 }
 
+//===----------------------------------------------------------------------===//
+// TestConstMemorySpaceAttr
+//===----------------------------------------------------------------------===//
+
+LogicalResult TestConstMemorySpaceAttr::isValidLoad(
+    Type type, mlir::ptr::AtomicOrdering ordering, IntegerAttr alignment,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return success();
+}
+
+LogicalResult TestConstMemorySpaceAttr::isValidStore(
+    Type type, mlir::ptr::AtomicOrdering ordering, IntegerAttr alignment,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return emitError ? (emitError() << "memory space is read-only") : failure();
+}
+
+LogicalResult TestConstMemorySpaceAttr::isValidAtomicOp(
+    mlir::ptr::AtomicBinOp binOp, Type type, mlir::ptr::AtomicOrdering ordering,
+    IntegerAttr alignment, function_ref<InFlightDiagnostic()> emitError) const {
+  return emitError ? (emitError() << "memory space is read-only") : failure();
+}
+
+LogicalResult TestConstMemorySpaceAttr::isValidAtomicXchg(
+    Type type, mlir::ptr::AtomicOrdering successOrdering,
+    mlir::ptr::AtomicOrdering failureOrdering, IntegerAttr alignment,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return emitError ? (emitError() << "memory space is read-only") : failure();
+}
+
+LogicalResult TestConstMemorySpaceAttr::isValidAddrSpaceCast(
+    Type tgt, Type src, function_ref<InFlightDiagnostic()> emitError) const {
+  return emitError
+             ? (emitError() << "memory space doesn't allow addrspace casts")
+             : failure();
+}
+
+LogicalResult TestConstMemorySpaceAttr::isValidPtrIntCast(
+    Type intLikeTy, Type ptrLikeTy,
+    function_ref<InFlightDiagnostic()> emitError) const {
+  return emitError ? (emitError() << "memory space doesn't allow int-ptr casts")
+                   : failure();
+}
+
 //===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//

mlir/test/lib/Dialect/Test/TestAttributes.h

@@ -18,6 +18,7 @@
 
 #include "TestTraits.h"
 #include "mlir/Dialect/Polynomial/IR/PolynomialAttributes.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"