Refactor address space handling to use TargetAddressSpaceAttr and update related functionality

Author: RiverDave

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

@@ -10,6 +10,7 @@
 #define LLVM_CLANG_CIR_DIALECT_BUILDER_CIRBASEBUILDER_H
 
 #include "clang/AST/CharUnits.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/CIRTypes.h"
@@ -129,19 +130,29 @@ class CIRBaseBuilderTy : public mlir::OpBuilder {
     return cir::PointerType::get(ty);
   }
 
-  cir::PointerType getPointerTo(mlir::Type ty, cir::AddressSpace as) {
+  cir::PointerType getPointerTo(mlir::Type ty, cir::TargetAddressSpaceAttr as) {
     return cir::PointerType::get(ty, as);
   }
 
   cir::PointerType getPointerTo(mlir::Type ty, clang::LangAS langAS) {
-    return getPointerTo(ty, cir::toCIRAddressSpace(langAS));
+    if (langAS == clang::LangAS::Default) // Default address space.
+      return getPointerTo(ty);
+
+    if (clang::isTargetAddressSpace(langAS)) {
+      unsigned addrSpace = clang::toTargetAddressSpace(langAS);
+      auto asAttr = cir::TargetAddressSpaceAttr::get(
+          getContext(), getUI32IntegerAttr(addrSpace));
+      return getPointerTo(ty, asAttr);
+    }
+
+    llvm_unreachable("language-specific address spaces NYI");
   }
 
   cir::PointerType getVoidPtrTy(clang::LangAS langAS = clang::LangAS::Default) {
     return getPointerTo(cir::VoidType::get(getContext()), langAS);
   }
 
-  cir::PointerType getVoidPtrTy(cir::AddressSpace as) {
+  cir::PointerType getVoidPtrTy(cir::TargetAddressSpaceAttr as) {
     return getPointerTo(cir::VoidType::get(getContext()), as);
   }
 

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

@@ -602,78 +602,25 @@ def CIR_VTableAttr : CIR_Attr<"VTable", "vtable", [TypedAttrInterface]> {
 }
 
 //===----------------------------------------------------------------------===//
-// AddressSpaceAttr
+// TargetAddressSpaceAttr
 //===----------------------------------------------------------------------===//
 
-def CIR_AddressSpaceAttr :  CIR_EnumAttr<CIR_AddressSpace, "address_space"> {
-  let summary = "Attribute representing memory address spaces";
+def CIR_TargetAddressSpaceAttr : CIR_Attr< "TargetAddressSpace",
+                                         "target_address_space"> {
+  let summary = "Attribute representing a target-specific numeric address space";
   let description = [{
-    Represents different memory address spaces for pointer types.
-    Address spaces distinguish between different types of memory regions,
-    such as global, local, or constant memory.
+    Represents a target-specific numeric address space for pointer types.
 
-    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
-    printer and parser for the `value` parameter.
-
-    CIR supports two categories:
-    - Language address spaces: Predefined spaces like `offload_global`,
-      `offload_constant`, `offload_private` for offload programming models
-    - Target address spaces: Numeric spaces `target<N>` for target-specific
-      memory regions, where N is the address space number
-
-    Examples:
+    Example:
     ```mlir
-    // Default address space (implicit)
-    !cir.ptr<!s32i>
-    
-    // Language-defined offload address spaces  
-    !cir.ptr<!s32i, addrspace(offload_global)>
-    !cir.ptr<!s32i, addrspace(offload_constant)>
-    !cir.ptr<!s32i, addrspace(offload_private)>
-    
     // Target-specific numeric address spaces
     !cir.ptr<!s32i, addrspace(target<1>)>
     !cir.ptr<!s32i, addrspace(target<10>)>
     ```
   }];
-  let builders = [
-    AttrBuilder<(ins "clang::LangAS":$langAS), [{
-      return $_get($_ctxt, cir::toCIRAddressSpace(langAS));
-    }]>
-  ];
-
-  let assemblyFormat = [{
-    `(` custom<AddressSpaceValue>($value) `)`
-  }];
-
-  let defaultValue = "cir::AddressSpace::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());
-    }
-
-    bool $cppClass::isTarget() const {
-      return cir::isTargetAddressSpace(getValue());
-    }
-
-    unsigned $cppClass::getTargetValue() const {
-      return cir::getTargetAddressSpaceValue(getValue());
-    }
-  }];
+  let parameters = (ins "mlir::IntegerAttr":$value);
+  let assemblyFormat = "`<` `target` `<` $value `>` `>`";
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -35,23 +35,4 @@ class CIR_DefaultValuedEnumParameter<EnumAttrInfo info, string value = "">
   let defaultValue = value;
 }
 
-def CIR_AddressSpace : CIR_I32EnumAttr<
-  "AddressSpace", "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">
-]> {
-  let description = [{
-    The `address_space` attribute is used to represent address spaces for
-    pointer types in CIR. It provides a unified model on top of `clang::LangAS`
-    and simplifies the representation of address spaces.
-  }];
-
-  let genSpecializedAttr = 0;
-}
-
 #endif // CLANG_CIR_DIALECT_IR_CIRENUMATTR_TD

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

@@ -17,6 +17,7 @@
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "clang/Basic/AddressSpaces.h"
+#include "clang/CIR/Dialect/IR/CIRAttrs.h"
 #include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 #include "clang/CIR/Interfaces/CIRTypeInterfaces.h"
 
@@ -41,38 +42,9 @@ bool isSized(mlir::Type ty);
 // AddressSpace helpers
 //===----------------------------------------------------------------------===//
 
-cir::AddressSpace toCIRAddressSpace(clang::LangAS langAS);
+/// Returns the integer value of a CIR address space for LLVM.
+unsigned getTargetAddrSpaceFromAttr(cir::TargetAddressSpaceAttr attr);
 
-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();
-}
-
-constexpr bool isLangAddressSpace(cir::AddressSpace as) {
-  return !isTargetAddressSpace(as);
-}
-
-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);
-}
 } // namespace cir
 
 //===----------------------------------------------------------------------===//

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

@@ -248,37 +248,36 @@ def CIR_PointerType : CIR_Type<"Pointer", "ptr", [
     !cir.ptr<!cir.int<u, 8>>
     !cir.ptr<!cir.float>
     !cir.ptr<!cir.record<struct "MyStruct">>
-    !cir.ptr<!cir.record<struct "MyStruct">, addrspace(offload_private)>
-    !cir.ptr<!cir.int<u, 8>, addrspace(target<1>)>
+    !cir.ptr<!cir.int<u, 8>, target_address_space(1)>
+    !cir.ptr<!cir.record<struct "MyStruct">, target_address_space(5)>
     ```
   }];
 
   let parameters = (ins
     "mlir::Type":$pointee,
-    CIR_DefaultValuedEnumParameter<
-      CIR_AddressSpace,
-      "cir::AddressSpace::Default"
+    OptionalParameter<
+    "cir::TargetAddressSpaceAttr"
     >:$addrSpace
   );
 
   let skipDefaultBuilders = 1;
   let builders = [
     TypeBuilderWithInferredContext<(ins
       "mlir::Type":$pointee,
-      CArg<"cir::AddressSpace", "cir::AddressSpace::Default">:$addrSpace), [{
+       CArg<"cir::TargetAddressSpaceAttr", "nullptr">:$addrSpace), [{
         return $_get(pointee.getContext(), pointee, addrSpace);
     }]>,
     TypeBuilder<(ins
       "mlir::Type":$pointee,
-      CArg<"cir::AddressSpace", "cir::AddressSpace::Default">:$addrSpace), [{
+      CArg<"cir::TargetAddressSpaceAttr", "nullptr">:$addrSpace), [{
         return $_get($_ctxt, pointee, addrSpace);
     }]>
   ];
 
   let assemblyFormat = [{
     `<`
       $pointee
-      ( `,` `addrspace` `(` `` custom<AddressSpaceValue>($addrSpace)^ `)` )?
+      ( `,` ` ` custom<TargetAddressSpace>($addrSpace)^ )?
     `>`
   }];
 

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -2053,7 +2053,7 @@ mlir::Value CIRGenFunction::emitAlloca(StringRef name, mlir::Type ty,
   // layout like original CodeGen. The data layout awareness should be done in
   // the lowering pass instead.
   cir::PointerType localVarPtrTy =
-      builder.getPointerTo(ty, getCIRAllocaAddressSpace());
+      builder.getPointerTo(ty, getASTAllocaAddressSpace());
   mlir::IntegerAttr alignIntAttr = cgm.getSize(alignment);
 
   mlir::Value addr;

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -76,7 +76,7 @@ CIRGenModule::CIRGenModule(mlir::MLIRContext &mlirContext,
   SInt128Ty = cir::IntType::get(&getMLIRContext(), 128, /*isSigned=*/true);
   UInt8Ty = cir::IntType::get(&getMLIRContext(), 8, /*isSigned=*/false);
   UInt8PtrTy = cir::PointerType::get(UInt8Ty);
-  cirAllocaAddressSpace = getTargetCIRGenInfo().getCIRAllocaAddressSpace();
+  ASTAllocaAddressSpace = getTargetCIRGenInfo().getASTAllocaAddressSpace();
   UInt16Ty = cir::IntType::get(&getMLIRContext(), 16, /*isSigned=*/false);
   UInt32Ty = cir::IntType::get(&getMLIRContext(), 32, /*isSigned=*/false);
   UInt64Ty = cir::IntType::get(&getMLIRContext(), 64, /*isSigned=*/false);

clang/lib/CIR/CodeGen/CIRGenTypeCache.h

@@ -14,6 +14,7 @@
 #define LLVM_CLANG_LIB_CIR_CIRGENTYPECACHE_H
 
 #include "clang/AST/CharUnits.h"
+#include "clang/Basic/AddressSpaces.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
 
 namespace clang::CIRGen {
@@ -73,7 +74,7 @@ struct CIRGenTypeCache {
   /// The alignment of size_t.
   unsigned char SizeAlignInBytes;
 
-  cir::AddressSpace cirAllocaAddressSpace;
+  LangAS ASTAllocaAddressSpace;
 
   clang::CharUnits getSizeAlign() const {
     return clang::CharUnits::fromQuantity(SizeAlignInBytes);
@@ -83,9 +84,7 @@ struct CIRGenTypeCache {
     return clang::CharUnits::fromQuantity(PointerAlignInBytes);
   }
 
-  cir::AddressSpace getCIRAllocaAddressSpace() const {
-    return cirAllocaAddressSpace;
-  }
+  LangAS getASTAllocaAddressSpace() const { return ASTAllocaAddressSpace; }
 };
 
 } // namespace clang::CIRGen

clang/lib/CIR/CodeGen/TargetInfo.h

@@ -16,6 +16,7 @@
 
 #include "ABIInfo.h"
 #include "CIRGenTypes.h"
+#include "clang/Basic/AddressSpaces.h"
 
 #include <memory>
 #include <utility>
@@ -43,10 +44,8 @@ class TargetCIRGenInfo {
   /// Returns ABI info helper for the target.
   const ABIInfo &getABIInfo() const { return *info; }
 
-  /// Get the CIR address space for alloca.
-  virtual cir::AddressSpace getCIRAllocaAddressSpace() const {
-    return cir::AddressSpace::Default;
-  }
+  /// Get the AST address space for alloca.
+  virtual LangAS getASTAllocaAddressSpace() const { return LangAS::Default; }
 
   /// Determine whether a call to an unprototyped functions under
   /// the given calling convention should use the variadic

clang/lib/CIR/Dialect/IR/CIRAttrs.cpp

@@ -47,10 +47,11 @@ parseFloatLiteral(mlir::AsmParser &parser,
 // AddressSpaceAttr
 //===----------------------------------------------------------------------===//
 
-mlir::ParseResult parseAddressSpaceValue(mlir::AsmParser &p,
-                                         cir::AddressSpace &addrSpace);
+mlir::ParseResult
+parseTargetAddressSpace(mlir::AsmParser &p, cir::TargetAddressSpaceAttr &attr);
 
-void printAddressSpaceValue(mlir::AsmPrinter &p, cir::AddressSpace addrSpace);
+void printTargetAddressSpace(mlir::AsmPrinter &p,
+                                    cir::TargetAddressSpaceAttr attr);
 
 static mlir::ParseResult parseConstPtr(mlir::AsmParser &parser,
                                        mlir::IntegerAttr &value);