Implement MemorySpaceInterface for Clang and Target AS attributes

Author: RiverDave

clang/include/clang/CIR/Dialect/Builder/CIRBaseBuilder.h

@@ -19,13 +19,15 @@
 #include "clang/CIR/Dialect/IR/FPEnv.h"
 #include "clang/CIR/MissingFeatures.h"
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Support/LLVM.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/FloatingPointMode.h"
@@ -107,34 +109,30 @@ class CIRBaseBuilderTy : public mlir::OpBuilder {
   }
 
   /// Create a pointer type with an address space attribute.
-  /// Either a cir::ClangAddressSpaceAttr or cir::TargetAddressSpaceAttr is
-  /// expected.
-  cir::PointerType getPointerTo(mlir::Type ty, mlir::Attribute addrSpaceAttr) {
-    if (!addrSpaceAttr)
+  cir::PointerType
+  getPointerTo(mlir::Type ty, mlir::ptr::MemorySpaceAttrInterface memorySpace) {
+    if (!memorySpace)
       return cir::PointerType::get(ty);
-
-    assert((mlir::isa<cir::ClangAddressSpaceAttr>(addrSpaceAttr) ||
-            mlir::isa<cir::TargetAddressSpaceAttr>(addrSpaceAttr)) &&
-           "expected address space attribute");
-
-    return cir::PointerType::get(ty, addrSpaceAttr);
+    return cir::PointerType::get(ty, memorySpace);
   }
 
   cir::PointerType getPointerTo(mlir::Type ty, clang::LangAS langAS) {
     if (langAS == clang::LangAS::Default)
       return getPointerTo(ty);
 
-    mlir::Attribute addrSpaceAttr =
+    mlir::ptr::MemorySpaceAttrInterface addrSpaceAttr =
         cir::toCIRClangAddressSpaceAttr(getContext(), langAS);
+    // Convert to MemorySpaceAttrInterface for the interface overload
     return getPointerTo(ty, addrSpaceAttr);
   }
 
   cir::PointerType getVoidPtrTy(clang::LangAS langAS = clang::LangAS::Default) {
     return getPointerTo(cir::VoidType::get(getContext()), langAS);
   }
 
-  cir::PointerType getVoidPtrTy(mlir::Attribute addrSpaceAttr) {
-    return getPointerTo(cir::VoidType::get(getContext()), addrSpaceAttr);
+  cir::PointerType
+  getVoidPtrTy(mlir::ptr::MemorySpaceAttrInterface memorySpace) {
+    return getPointerTo(cir::VoidType::get(getContext()), memorySpace);
   }
 
   cir::MethodAttr getMethodAttr(cir::MethodType ty, cir::FuncOp methodFuncOp) {
@@ -432,7 +430,8 @@ class CIRBaseBuilderTy : public mlir::OpBuilder {
   mlir::Value createGetGlobal(mlir::Location loc, cir::GlobalOp global,
                               bool threadLocal = false) {
     return cir::GetGlobalOp::create(
-        *this, loc, getPointerTo(global.getSymType(), global.getAddrSpace()),
+        *this, loc,
+        getPointerTo(global.getSymType(), global.getAddrSpaceAttr()),
         global.getName(), threadLocal);
   }
 

clang/include/clang/CIR/Dialect/IR/CIRAttrConstraints.td

@@ -89,24 +89,4 @@ def CIR_AnyTBAAAttr : AnyAttrOf<[
   string cppType = "::mlir::Attribute";
 }
 
-//===----------------------------------------------------------------------===//
-// AddressSpaceAttr constraints
-//===----------------------------------------------------------------------===//
-
-// NOTE: We might end up using this only for GlobalOps, as we cannot apply constraints
-// to types.
-def CIR_AddressSpaceAttrConstraint
-  : CIR_AttrConstraint<"::cir::ClangAddressSpaceAttr", "clang address space attribute">;
-
-def CIR_TargetAddressSpaceAttrConstraint
-  : CIR_AttrConstraint<"::cir::TargetAddressSpaceAttr", "target address space attribute">;
-
-def CIR_AnyAddressSpaceAttr : AnyAttrOf<[
-  CIR_AddressSpaceAttrConstraint,
-  CIR_TargetAddressSpaceAttrConstraint
-]> {
-  string cppType = "::mlir::Attribute";
-  let constBuilderCall = "nullptr";
-}
-
 #endif // CLANG_CIR_DIALECT_IR_CIRATTRCONSTRAINTS_TD

clang/include/clang/CIR/Dialect/IR/CIRAttrs.h

@@ -13,6 +13,7 @@
 #ifndef CLANG_CIR_DIALECT_IR_CIRATTRS_H
 #define CLANG_CIR_DIALECT_IR_CIRATTRS_H
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 

clang/include/clang/CIR/Dialect/IR/CIRAttrs.td

@@ -14,6 +14,7 @@
 #define CLANG_CIR_DIALECT_IR_CIRATTRS_TD
 
 include "mlir/IR/BuiltinAttributeInterfaces.td"
+include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.td"
 include "clang/CIR/Dialect/IR/CIREnumAttr.td"
 
 include "clang/CIR/Dialect/IR/CIRDialect.td"
@@ -963,7 +964,9 @@ def CIR_DynamicCastInfoAttr : CIR_Attr<"DynamicCastInfo", "dyn_cast_info"> {
 // AddressSpaceAttr
 //===----------------------------------------------------------------------===//
 
-def CIR_ClangAddressSpaceAttr :  CIR_EnumAttr<CIR_ClangAddressSpace, "clang_address_space"> {
+def CIR_ClangAddressSpaceAttr :  CIR_EnumAttr<CIR_ClangAddressSpace, "clang_address_space", [
+    DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+  ]> {
 
   let summary = "Represents a language/Clang-level address space";
   let description = [{
@@ -1008,7 +1011,9 @@ def CIR_ClangAddressSpaceAttr :  CIR_EnumAttr<CIR_ClangAddressSpace, "clang_addr
 //===----------------------------------------------------------------------===//
 
 def CIR_TargetAddressSpaceAttr : CIR_Attr< "TargetAddressSpace",
-                                         "target_address_space"> {
+                                         "target_address_space", [
+    DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+  ]> {
   let summary = "Represents a target-specific numeric address space";
   let description = [{
     The TargetAddressSpaceAttr represents a target-specific numeric address space,

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -2616,10 +2616,7 @@ def CIR_GlobalOp : CIR_Op<"global", [
     OptionalAttr<StrAttr>:$sym_visibility,
     TypeAttr:$sym_type,
     CIR_GlobalLinkageKind:$linkage,
-    DefaultValuedAttr<
-      CIR_AnyAddressSpaceAttr,
-      "{}"
-    >:$addr_space,
+    OptionalAttr<MemorySpaceAttrInterface>:$addr_space,
     OptionalAttr<CIR_TLSModel>:$tls_model,
     // Note this can also be a FlatSymbolRefAttr
     OptionalAttr<AnyAttr>:$initial_value,
@@ -2666,7 +2663,7 @@ def CIR_GlobalOp : CIR_Op<"global", [
       // CIR defaults to external linkage.
       CArg<"cir::GlobalLinkageKind",
             "cir::GlobalLinkageKind::ExternalLinkage">:$linkage,
-      CArg<"mlir::Attribute", "{}">:$addrSpace,
+      CArg<"mlir::ptr::MemorySpaceAttrInterface", "{}">:$addrSpace,
       CArg<"llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)>",
            "nullptr">:$ctorBuilder,
       CArg<"llvm::function_ref<void(mlir::OpBuilder &, mlir::Location)>",

clang/include/clang/CIR/Dialect/IR/CIRTypes.h

@@ -13,6 +13,7 @@
 #ifndef CLANG_CIR_DIALECT_IR_CIRTYPES_H
 #define CLANG_CIR_DIALECT_IR_CIRTYPES_H
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
@@ -36,11 +37,13 @@ bool isSized(mlir::Type ty);
 
 cir::ClangAddressSpace toCIRClangAddressSpace(clang::LangAS langAS);
 
-/// Convert a LangAS to the appropriate address space attribute.
-/// Returns ClangAddressSpaceAttr for clang/language-specific address spaces,
-/// or TargetAddressSpaceAttr for target-specific address spaces.
-mlir::Attribute toCIRClangAddressSpaceAttr(mlir::MLIRContext *ctx,
-                                           clang::LangAS langAS);
+/// Convert a LangAS to the appropriate address space attribute interface.
+/// Returns a MemorySpaceAttrInterface.
+mlir::ptr::MemorySpaceAttrInterface
+toCIRClangAddressSpaceAttr(mlir::MLIRContext *ctx, clang::LangAS langAS);
+
+bool isSupportedCIRMemorySpaceAttr(
+    mlir::ptr::MemorySpaceAttrInterface memorySpace);
 
 constexpr unsigned getAsUnsignedValue(cir::ClangAddressSpace as) {
   return static_cast<unsigned>(as);

clang/include/clang/CIR/Dialect/IR/CIRTypes.td

@@ -237,19 +237,19 @@ def CIR_PointerType : CIR_Type<"Pointer", "ptr", [
 
   let parameters = (ins
     "mlir::Type":$pointee,
-    OptionalParameter<"mlir::Attribute">:$addrSpace
+    OptionalParameter<"mlir::ptr::MemorySpaceAttrInterface">:$addrSpace
   );
 
   let skipDefaultBuilders = 1;
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "mlir::Type":$pointee,
-      CArg<"mlir::Attribute", "{}">:$addrSpace), [{
+      CArg<"mlir::ptr::MemorySpaceAttrInterface", "{}">:$addrSpace), [{
         return $_get(pointee.getContext(), pointee, addrSpace);
     }]>,
     TypeBuilder<(ins
       "mlir::Type":$pointee,
-      CArg<"mlir::Attribute", "{}">:$addrSpace), [{
+      CArg<"mlir::ptr::MemorySpaceAttrInterface", "{}">:$addrSpace), [{
         return $_get($_ctxt, pointee, addrSpace);
     }]>
   ];

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -25,6 +25,7 @@
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
 #include "clang/CIR/Dialect/IR/FPEnv.h"
 
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -749,7 +750,7 @@ class CIRGenBuilderTy : public cir::CIRBaseBuilderTy {
   [[nodiscard]] cir::GlobalOp
   createGlobal(mlir::ModuleOp module, mlir::Location loc, mlir::StringRef name,
                mlir::Type type, bool isConst, cir::GlobalLinkageKind linkage,
-               mlir::Attribute addrSpace = {}) {
+               mlir::ptr::MemorySpaceAttrInterface addrSpace = {}) {
     mlir::OpBuilder::InsertionGuard guard(*this);
     setInsertionPointToStart(module.getBody());
     return cir::GlobalOp::create(*this, loc, name, type, isConst, linkage,
@@ -763,7 +764,7 @@ class CIRGenBuilderTy : public cir::CIRBaseBuilderTy {
   createVersionedGlobal(mlir::ModuleOp module, mlir::Location loc,
                         mlir::StringRef name, mlir::Type type, bool isConst,
                         cir::GlobalLinkageKind linkage,
-                        mlir::Attribute addrSpace = {}) {
+                        mlir::ptr::MemorySpaceAttrInterface addrSpace = {}) {
     // Create a unique name if the given name is already taken.
     std::string uniqueName;
     if (unsigned version = GlobalsVersioning[name.str()]++)

clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp

@@ -34,6 +34,7 @@
 #include "clang/Frontend/FrontendDiagnostic.h"
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Value.h"
@@ -1728,8 +1729,8 @@ RValue CIRGenFunction::emitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     // the AST level this is handled within CreateTempAlloca et al., but for the
     // builtin / dynamic alloca we have to handle it here.
     assert(!cir::MissingFeatures::addressSpace());
-    mlir::Attribute AAS = getCIRAllocaAddressSpace();
-    mlir::Attribute EAS = cir::toCIRClangAddressSpaceAttr(
+    mlir::ptr::MemorySpaceAttrInterface AAS = getCIRAllocaAddressSpace();
+    mlir::ptr::MemorySpaceAttrInterface EAS = cir::toCIRClangAddressSpaceAttr(
         &getMLIRContext(), E->getType()->getPointeeType().getAddressSpace());
     if (EAS != AAS) {
       assert(false && "Non-default address space for alloca NYI");

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -15,10 +15,12 @@
 #include "CIRGenFunction.h"
 #include "CIRGenOpenMPRuntime.h"
 #include "EHScopeStack.h"
+#include "mlir/Dialect/Ptr/IR/MemorySpaceInterfaces.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Support/LLVM.h"
 
 #include "clang/AST/Decl.h"
 #include "clang/AST/ExprCXX.h"
@@ -479,7 +481,7 @@ CIRGenModule::getOrCreateStaticVarDecl(const VarDecl &D,
     Name = getStaticDeclName(*this, D);
 
   mlir::Type LTy = getTypes().convertTypeForMem(Ty);
-  mlir::Attribute AS = cir::toCIRClangAddressSpaceAttr(
+  mlir::ptr::MemorySpaceAttrInterface AS = cir::toCIRClangAddressSpaceAttr(
       &getMLIRContext(), getGlobalVarAddressSpace(&D));
 
   // OpenCL variables in local address space and CUDA shared
@@ -595,7 +597,7 @@ cir::GlobalOp CIRGenFunction::addInitializerToStaticVarDecl(
     // Given those constraints, thread in the GetGlobalOp and update it
     // directly.
     GVAddr.getAddr().setType(
-        getBuilder().getPointerTo(Init.getType(), GV.getAddrSpace()));
+        getBuilder().getPointerTo(Init.getType(), GV.getAddrSpaceAttr()));
   }
 
   bool NeedsDtor =