[CIR] Backport TargetAddressSpaceAttr and Support both existing Lang AS and target AS attributes in pointer types and Global Ops (#1986)

Bringing the changes we performed at:
llvm/llvm-project#161028 down to the incubator
isn't as straight-forward as I though. Therefore — this PR might be a
bit long but hopefully not too complicated, bare with me :)

The main purpose of this is to bring the recent upstreamed
`TargetAddressSpaceAttr` and couple it with `PointerType` and
`GlobalOp`. The main challenge is that this collides with the already
implemented infrastructure related to AS that revolves around our
unified enum approach (That handles numeric and lang AS). My main
rationale here is that we let our new attribute to coexist with the
already existing `cir::AddressSpaceAttr` (renamed now). so that we don't
break any tests related to offload-type languages.

Considering the above what I'm essentially doing is:

1. Letting `TargetAddressSpaceAttr` handle numeric target AS as it is
done upstream.
2. Rename our previous `AddressSpaceAttr` to `LangAdddressSpaceAttr`.
3. Implement `MemorySpaceAttrInterface`for lang and target AS.
4. The current implementation of `LangAddressSpaceAttr`(which handles an
enum) will handle language/clang specific AS (CUDA, OpenCL, etc..) it
will no longer handle target as it used to!
5. `PointerType` and globalOps will hold a `MemorySpaceAttrInterface`.
Similar to the upstream ptr type in the ptr dialect.
6. Adjusted the assembly format for all tests containing any AS
attributes to conform with our new format:
`lang_address_space(n)`/`clang_address_space(x)`.
7. For the relevant Ops: Default address spaces will now be represented
as an empty `MemorySpaceAttrInterface` attribute.
8. Remove any comments that make reference to our old approach regarding
AS handling.
9. Updated description where relevant.
10. Bring any extra upstream changes that deviate with what we have
here.

Author: RiverDave

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

@@ -11,20 +11,23 @@
 
 #include "clang/AST/Decl.h"
 #include "clang/AST/Type.h"
+#include "clang/Basic/AddressSpaces.h"
 #include "clang/CIR/Dialect/IR/CIRAttrs.h"
 #include "clang/CIR/Dialect/IR/CIRDialect.h"
 #include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
 #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"
@@ -105,20 +108,30 @@ class CIRBaseBuilderTy : public mlir::OpBuilder {
     return cir::PointerType::get(ty);
   }
 
-  cir::PointerType getPointerTo(mlir::Type ty, cir::AddressSpace as) {
-    return cir::PointerType::get(ty, as);
+  /// Create a pointer type with an address space attribute.
+  cir::PointerType
+  getPointerTo(mlir::Type ty, mlir::ptr::MemorySpaceAttrInterface memorySpace) {
+    if (!memorySpace)
+      return cir::PointerType::get(ty);
+    return cir::PointerType::get(ty, memorySpace);
   }
 
   cir::PointerType getPointerTo(mlir::Type ty, clang::LangAS langAS) {
-    return getPointerTo(ty, cir::toCIRAddressSpace(langAS));
+    if (langAS == clang::LangAS::Default)
+      return getPointerTo(ty);
+
+    mlir::ptr::MemorySpaceAttrInterface addrSpaceAttr =
+        cir::toCIRLangAddressSpaceAttr(getContext(), langAS);
+    return getPointerTo(ty, addrSpaceAttr);
   }
 
   cir::PointerType getVoidPtrTy(clang::LangAS langAS = clang::LangAS::Default) {
     return getPointerTo(cir::VoidType::get(getContext()), langAS);
   }
 
-  cir::PointerType getVoidPtrTy(cir::AddressSpace as) {
-    return getPointerTo(cir::VoidType::get(getContext()), as);
+  cir::PointerType
+  getVoidPtrTy(mlir::ptr::MemorySpaceAttrInterface memorySpace) {
+    return getPointerTo(cir::VoidType::get(getContext()), memorySpace);
   }
 
   cir::MethodAttr getMethodAttr(cir::MethodType ty, cir::FuncOp methodFuncOp) {
@@ -419,7 +432,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/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,43 +964,76 @@ def CIR_DynamicCastInfoAttr : CIR_Attr<"DynamicCastInfo", "dyn_cast_info"> {
 // AddressSpaceAttr
 //===----------------------------------------------------------------------===//
 
-def CIR_AddressSpaceAttr :  CIR_EnumAttr<CIR_AddressSpace, "address_space"> {
+def CIR_LangAddressSpaceAttr :
+              CIR_EnumAttr<CIR_LangAddressSpace, "lang_address_space", [
+    DeclareAttrInterfaceMethods<MemorySpaceAttrInterface>
+  ]> {
+
+  let summary = "Represents a language address space";
+  let description = [{
+    Encodes the semantic address spaces defined by the front-end language
+    (e.g. `__shared__`, `__constant__`, `__local__`). Values are stored using the
+    `cir::LangAddressSpace` enum, keeping the representation compact while and
+     preserves the qualifier until it is mapped onto target/LLVM address-space numbers.
+
+     Example: 
+     ``` mlir
+     !cir.ptr<!s32i, lang_address_space(offload_local)>
+     cir.global constant external lang_address_space(offload_constant)
+
+     ```
+  }];
+  
   let builders = [
     AttrBuilder<(ins "clang::LangAS":$langAS), [{
-      return $_get($_ctxt, cir::toCIRAddressSpace(langAS));
+      return $_get($_ctxt, cir::toCIRLangAddressSpace(langAS));
     }]>
   ];
 
   let assemblyFormat = [{
     `(` custom<AddressSpaceValue>($value) `)`
   }];
 
-  let defaultValue = "cir::AddressSpace::Default";
+  let defaultValue = "cir::LangAddressSpace::Default";
 
   let extraClassDeclaration = [{
-    bool isLang() const;
-    bool isTarget() const;
-    unsigned getTargetValue() const;
     unsigned getAsUnsignedValue() const;
   }];
 
   let extraClassDefinition = [{
     unsigned $cppClass::getAsUnsignedValue() const {
       return static_cast<unsigned>(getValue());
     }
+  }];
+}
 
-    bool $cppClass::isLang() const {
-      return cir::isLangAddressSpace(getValue());
-    }
+//===----------------------------------------------------------------------===//
+// TargetAddressSpaceAttr
+//===----------------------------------------------------------------------===//
 
-    bool $cppClass::isTarget() const {
-      return cir::isTargetAddressSpace(getValue());
-    }
+def CIR_TargetAddressSpaceAttr : CIR_Attr< "TargetAddressSpace",
+                                         "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,
+    corresponding to the LLVM IR `addressspace` qualifier and the clang
+     `address_space` attribute.
+    
+    A value of zero represents the default address space. The semantics of non-zero
+    address spaces are target-specific.
 
-    unsigned $cppClass::getTargetValue() const {
-      return cir::getTargetAddressSpaceValue(getValue());
-    }
+    Example:
+    ```mlir
+    // Target-specific numeric address spaces
+    !cir.ptr<!s32i, target_address_space(1)>
+    !cir.ptr<!s32i, target_address_space(1)>
+    ```
   }];
+
+  let parameters = (ins "unsigned":$value);
+  let assemblyFormat = "`<` `target` `<` $value `>` `>`";
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -41,15 +41,14 @@ class CIR_DefaultValuedEnumParameter<EnumAttrInfo info, string value = "">
   let defaultValue = value;
 }
 
-def CIR_AddressSpace : CIR_I32EnumAttr<
-  "AddressSpace", "address space kind", [
+def CIR_LangAddressSpace : CIR_I32EnumAttr<
+  "LangAddressSpace", "language address space kind", [
   I32EnumAttrCase<"Default", 0, "default">,
   I32EnumAttrCase<"OffloadPrivate", 1, "offload_private">,
   I32EnumAttrCase<"OffloadLocal", 2, "offload_local">,
   I32EnumAttrCase<"OffloadGlobal", 3, "offload_global">,
   I32EnumAttrCase<"OffloadConstant", 4, "offload_constant">,
-  I32EnumAttrCase<"OffloadGeneric", 5, "offload_generic">,
-  I32EnumAttrCase<"Target", 6, "target">
+  I32EnumAttrCase<"OffloadGeneric", 5, "offload_generic">
 ]> {
   let description = [{
     The `address_space` attribute is used to represent address spaces for
@@ -58,7 +57,7 @@ def CIR_AddressSpace : CIR_I32EnumAttr<
 
     The `value` parameter is an extensible enum, which encodes target address
     space as an offset to the last language address space. For that reason, the
-    attribute is implemented as custom AddressSpaceAttr, which provides custom
+    attribute is implemented as custom LangAddressSpaceAttr, which provides custom
     printer and parser for the `value` parameter.
   }];
 

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

@@ -2618,10 +2618,7 @@ def CIR_GlobalOp : CIR_Op<"global", [
     OptionalAttr<StrAttr>:$sym_visibility,
     TypeAttr:$sym_type,
     CIR_GlobalLinkageKind:$linkage,
-    DefaultValuedAttr<
-      CIR_AddressSpaceAttr,
-      "AddressSpace::Default"
-    >:$addr_space,
+    OptionalAttr<MemorySpaceAttrInterface>:$addr_space,
     OptionalAttr<CIR_TLSModel>:$tls_model,
     // Note this can also be a FlatSymbolRefAttr
     OptionalAttr<AnyAttr>:$initial_value,
@@ -2644,7 +2641,7 @@ def CIR_GlobalOp : CIR_Op<"global", [
        (`comdat` $comdat^)?
        ($tls_model^)?
        (`dso_local` $dso_local^)?
-       (`addrspace` `` $addr_space^)?
+       (` ` custom<GlobalAddressSpaceValue>($addr_space)^ )?
        $sym_name
        custom<GlobalOpTypeAndInitialValue>($sym_type, $initial_value, $ctorRegion, $dtorRegion)
        ($annotations^)?
@@ -2668,7 +2665,7 @@ def CIR_GlobalOp : CIR_Op<"global", [
       // CIR defaults to external linkage.
       CArg<"cir::GlobalLinkageKind",
             "cir::GlobalLinkageKind::ExternalLinkage">:$linkage,
-      CArg<"cir::AddressSpace", "cir::AddressSpace::Default">:$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"
@@ -34,37 +35,18 @@ bool isSized(mlir::Type ty);
 // AddressSpace helpers
 //===----------------------------------------------------------------------===//
 
-cir::AddressSpace toCIRAddressSpace(clang::LangAS langAS);
+cir::LangAddressSpace toCIRLangAddressSpace(clang::LangAS langAS);
 
-constexpr unsigned getAsUnsignedValue(cir::AddressSpace as) {
-  return static_cast<unsigned>(as);
-}
-
-inline constexpr unsigned TargetAddressSpaceOffset =
-    cir::getMaxEnumValForAddressSpace();
-
-// Target address space is used for target-specific address spaces that are not
-// part of the enum. Its value is represented as an offset from the maximum
-// value of the enum. Make sure that it is always the last enum value.
-static_assert(getAsUnsignedValue(cir::AddressSpace::Target) ==
-                  cir::getMaxEnumValForAddressSpace(),
-              "Target address space must be the last enum value");
-
-constexpr bool isTargetAddressSpace(cir::AddressSpace as) {
-  return getAsUnsignedValue(as) >= cir::getMaxEnumValForAddressSpace();
-}
+/// Convert a LangAS to the appropriate address space attribute interface.
+/// Returns a MemorySpaceAttrInterface.
+mlir::ptr::MemorySpaceAttrInterface
+toCIRLangAddressSpaceAttr(mlir::MLIRContext *ctx, clang::LangAS langAS);
 
-constexpr bool isLangAddressSpace(cir::AddressSpace as) {
-  return !isTargetAddressSpace(as);
-}
+bool isSupportedCIRMemorySpaceAttr(
+    mlir::ptr::MemorySpaceAttrInterface memorySpace);
 
-constexpr unsigned getTargetAddressSpaceValue(cir::AddressSpace as) {
-  assert(isTargetAddressSpace(as) && "expected target address space");
-  return getAsUnsignedValue(as) - TargetAddressSpaceOffset;
-}
-
-constexpr cir::AddressSpace computeTargetAddressSpace(unsigned v) {
-  return static_cast<cir::AddressSpace>(v + TargetAddressSpaceOffset);
+constexpr unsigned getAsUnsignedValue(cir::LangAddressSpace as) {
+  return static_cast<unsigned>(as);
 }
 
 } // namespace cir

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

@@ -16,11 +16,13 @@
 include "clang/CIR/Dialect/IR/CIRDialect.td"
 include "clang/CIR/Dialect/IR/CIREnumAttr.td"
 include "clang/CIR/Dialect/IR/CIRTypeConstraints.td"
+include "clang/CIR/Dialect/IR/CIRAttrConstraints.td"
 include "clang/CIR/Interfaces/ASTAttrInterfaces.td"
 include "clang/CIR/Interfaces/CIRTypeInterfaces.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/IR/CommonAttrConstraints.td"
 
 //===----------------------------------------------------------------------===//
 // CIR Types
@@ -231,32 +233,31 @@ def CIR_PointerType : CIR_Type<"Pointer", "ptr", [
     ```
   }];
 
+  let genVerifyDecl = 1;
+
   let parameters = (ins
     "mlir::Type":$pointee,
-    CIR_DefaultValuedEnumParameter<
-      CIR_AddressSpace,
-      "cir::AddressSpace::Default"
-    >:$addrSpace
+    OptionalParameter<"mlir::ptr::MemorySpaceAttrInterface">:$addrSpace
   );
 
   let skipDefaultBuilders = 1;
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "mlir::Type":$pointee,
-      CArg<"cir::AddressSpace", "cir::AddressSpace::Default">:$addrSpace), [{
+      CArg<"mlir::ptr::MemorySpaceAttrInterface", "{}">:$addrSpace), [{
         return $_get(pointee.getContext(), pointee, addrSpace);
     }]>,
     TypeBuilder<(ins
       "mlir::Type":$pointee,
-      CArg<"cir::AddressSpace", "cir::AddressSpace::Default">:$addrSpace), [{
+      CArg<"mlir::ptr::MemorySpaceAttrInterface", "{}">:$addrSpace), [{
         return $_get($_ctxt, pointee, addrSpace);
     }]>
   ];
 
   let assemblyFormat = [{
     `<`
       $pointee
-      ( `,` `addrspace` `(` `` custom<AddressSpaceValue>($addrSpace)^ `)` )?
+      ( `,` ` ` custom<AddressSpaceValue>($addrSpace)^ )?
     `>`
   }];
 

clang/include/clang/CIR/MissingFeatures.h

@@ -480,6 +480,8 @@ struct MissingFeatures {
   static bool mustProgress() { return false; }
 
   static bool skipTempCopy() { return false; }
+
+  static bool dataLayoutPtrHandlingBasedOnLangAS() { return false; }
 };
 
 } // namespace cir

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,
-               cir::AddressSpace addrSpace = cir::AddressSpace::Default) {
+               mlir::ptr::MemorySpaceAttrInterface addrSpace = {}) {
     mlir::OpBuilder::InsertionGuard guard(*this);
     setInsertionPointToStart(module.getBody());
     return cir::GlobalOp::create(*this, loc, name, type, isConst, linkage,
@@ -759,10 +760,11 @@ class CIRGenBuilderTy : public cir::CIRBaseBuilderTy {
   /// Creates a versioned global variable. If the symbol is already taken, an ID
   /// will be appended to the symbol. The returned global must always be queried
   /// for its name so it can be referenced correctly.
-  [[nodiscard]] cir::GlobalOp createVersionedGlobal(
-      mlir::ModuleOp module, mlir::Location loc, mlir::StringRef name,
-      mlir::Type type, bool isConst, cir::GlobalLinkageKind linkage,
-      cir::AddressSpace addrSpace = cir::AddressSpace::Default) {
+  [[nodiscard]] cir::GlobalOp
+  createVersionedGlobal(mlir::ModuleOp module, mlir::Location loc,
+                        mlir::StringRef name, mlir::Type type, bool isConst,
+                        cir::GlobalLinkageKind linkage,
+                        mlir::ptr::MemorySpaceAttrInterface addrSpace = {}) {
     // Create a unique name if the given name is already taken.
     std::string uniqueName;
     if (unsigned version = GlobalsVersioning[name.str()]++)