[CIR] Upstream AddressSpace support for PointerType

Author: RiverDave

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

@@ -129,8 +129,20 @@ class CIRBaseBuilderTy : public mlir::OpBuilder {
     return cir::PointerType::get(ty);
   }
 
-  cir::PointerType getVoidPtrTy() {
-    return getPointerTo(cir::VoidType::get(getContext()));
+  cir::PointerType getPointerTo(mlir::Type ty, cir::AddressSpace as) {
+    return cir::PointerType::get(ty, as);
+  }
+
+  cir::PointerType getPointerTo(mlir::Type ty, clang::LangAS langAS) {
+    return getPointerTo(ty, cir::toCIRAddressSpace(langAS));
+  }
+
+  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::BoolAttr getCIRBoolAttr(bool state) {

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

@@ -15,6 +15,7 @@
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
+#include "clang/Basic/AddressSpaces.h"
 
 #include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 

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

@@ -601,6 +601,47 @@ def CIR_VTableAttr : CIR_Attr<"VTable", "vtable", [TypedAttrInterface]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// AddressSpaceAttr
+//===----------------------------------------------------------------------===//
+
+def CIR_AddressSpaceAttr : CIR_EnumAttr<CIR_AddressSpace, "address_space"> {
+  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());
+    }
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // ConstComplexAttr
 //===----------------------------------------------------------------------===//

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

@@ -35,4 +35,28 @@ 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.
+
+    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.
+  }];
+
+  let genSpecializedAttr = 0;
+}
+
 #endif // CLANG_CIR_DIALECT_IR_CIRENUMATTR_TD

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

@@ -16,6 +16,8 @@
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "clang/Basic/AddressSpaces.h"
+#include "clang/CIR/Dialect/IR/CIROpsEnums.h"
 #include "clang/CIR/Interfaces/CIRTypeInterfaces.h"
 
 namespace cir {
@@ -35,6 +37,42 @@ bool isValidFundamentalIntWidth(unsigned width);
 /// void, or abstract types.
 bool isSized(mlir::Type ty);
 
+//===----------------------------------------------------------------------===//
+// AddressSpace helpers
+//===----------------------------------------------------------------------===//
+
+cir::AddressSpace toCIRAddressSpace(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();
+}
+
+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

@@ -14,10 +14,12 @@
 #define CLANG_CIR_DIALECT_IR_CIRTYPES_TD
 
 include "clang/CIR/Dialect/IR/CIRDialect.td"
+include "clang/CIR/Dialect/IR/CIREnumAttr.td"
 include "clang/CIR/Dialect/IR/CIRTypeConstraints.td"
 include "clang/CIR/Interfaces/CIRTypeInterfaces.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
 
 //===----------------------------------------------------------------------===//
 // CIR Types
@@ -226,31 +228,57 @@ def CIR_PointerType : CIR_Type<"Pointer", "ptr", [
 ]> {
   let summary = "CIR pointer type";
   let description = [{
-    The `!cir.ptr` type represents C and C++ pointer types and C++ reference
-    types, other than pointers-to-members.  The `pointee` type is the type
-    pointed to.
+    The `!cir.ptr` type is a typed pointer type. It is used to represent
+    pointers to objects in C/C++. The type of the pointed-to object is given by
+    the `pointee` parameter. The `addrSpace` parameter is an optional address
+    space attribute that specifies the address space of the pointer. If not
+    specified, the pointer is assumed to be in the default address space.
 
-    TODO(CIR): The address space attribute is not yet implemented.
-  }];
+    The `!cir.ptr` type can point to any type, including fundamental types,
+    records, arrays, vectors, functions, and other pointers. It can also point
+    to incomplete types, such as incomplete records.
 
-  let parameters = (ins "mlir::Type":$pointee);
+    Note: Data-member pointers and method pointers are represented by
+    `!cir.data_member` and `!cir.method` types, respectively not by
+    `!cir.ptr` type.
 
-  let builders = [
-    TypeBuilderWithInferredContext<(ins "mlir::Type":$pointee), [{
-      return $_get(pointee.getContext(), pointee);
-    }]>,
-    TypeBuilder<(ins "mlir::Type":$pointee), [{
-      return $_get($_ctxt, pointee);
-    }]>
-  ];
+    Examples:
 
-  let assemblyFormat = [{
-    `<` $pointee  `>`
+    ```mlir
+    !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>)>
+    ```
   }];
 
-  let genVerifyDecl = 1;
+  let parameters = (ins "mlir::Type":$pointee,
+      CIR_DefaultValuedEnumParameter<CIR_AddressSpace,
+                                     "cir::AddressSpace::Default">:$addrSpace);
 
   let skipDefaultBuilders = 1;
+  let builders = [TypeBuilderWithInferredContext<
+                      (ins "mlir::Type":$pointee,
+                          CArg<"cir::AddressSpace",
+                               "cir::AddressSpace::Default">:$addrSpace),
+                      [{
+        return $_get(pointee.getContext(), pointee, addrSpace);
+    }]>,
+                  TypeBuilder<
+                      (ins "mlir::Type":$pointee,
+                          CArg<"cir::AddressSpace",
+                               "cir::AddressSpace::Default">:$addrSpace),
+                      [{
+        return $_get($_ctxt, pointee, addrSpace);
+    }]>];
+
+  let assemblyFormat = [{
+    `<`
+      $pointee
+      ( `,` `addrspace` `(` custom<AddressSpaceValue>($addrSpace)^ `)` )?
+    `>`
+  }];
 
   let extraClassDeclaration = [{
     template <typename ...Types>

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -2053,7 +2053,8 @@ 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.
   assert(!cir::MissingFeatures::addressSpace());
-  cir::PointerType localVarPtrTy = builder.getPointerTo(ty);
+  cir::PointerType localVarPtrTy =
+      builder.getPointerTo(ty, getCIRAllocaAddressSpace());
   mlir::IntegerAttr alignIntAttr = cgm.getSize(alignment);
 
   mlir::Value addr;

clang/lib/CIR/CodeGen/CIRGenTypeCache.h

@@ -73,13 +73,19 @@ struct CIRGenTypeCache {
   /// The alignment of size_t.
   unsigned char SizeAlignInBytes;
 
+  cir::AddressSpace cirAllocaAddressSpace;
+
   clang::CharUnits getSizeAlign() const {
     return clang::CharUnits::fromQuantity(SizeAlignInBytes);
   }
 
   clang::CharUnits getPointerAlign() const {
     return clang::CharUnits::fromQuantity(PointerAlignInBytes);
   }
+
+  cir::AddressSpace getCIRAllocaAddressSpace() const {
+    return cirAllocaAddressSpace;
+  }
 };
 
 } // namespace clang::CIRGen

clang/lib/CIR/CodeGen/CIRGenTypes.cpp

@@ -417,7 +417,7 @@ mlir::Type CIRGenTypes::convertType(QualType type) {
 
     mlir::Type pointeeType = convertType(elemTy);
 
-    resultType = builder.getPointerTo(pointeeType);
+    resultType = builder.getPointerTo(pointeeType, elemTy.getAddressSpace());
     break;
   }
 

clang/lib/CIR/CodeGen/TargetInfo.h

@@ -32,6 +32,8 @@ bool isEmptyFieldForLayout(const ASTContext &context, const FieldDecl *fd);
 /// if the [[no_unique_address]] attribute would have made them empty.
 bool isEmptyRecordForLayout(const ASTContext &context, QualType t);
 
+class CIRGenFunction;
+
 class TargetCIRGenInfo {
   std::unique_ptr<ABIInfo> info;