[CIR][CodeGen] Set address space for OpenCL globals (#788)

This PR sets proper address space when creating `cir.global` and
`cir.get_global`.

Different languages use different ways to encode the address space in
AST constructs (i.e. VarDecl *). OpenCL and SYCL use an address space
qualifier on the type of `VarDecl`, while CUDA uses separate AST
attributes like `CUDASharedAttr`. Similarily, some targets may want to
use special address space for global variables. So a per-language +
per-target hook is needed to provide this customization point. In the
LLVM CodeGen, it's the helper method `getGlobalVarAddressSpace` that
takes on the role.

For OpenCL C + SPIR-V combination, OpenCL C converts the address space
qualifier to corresponding LangAS, but SPIR-V does not require any
action.

This PR implements `global` qualifier in OpenCL C, but does not include
`constant` qualifier. Although the modified part works for `constant`,
CIRGen is not yet able to set constant attribute for global ops (there
is a TODO line).

Static variable decl and `local` qualifier work in a similar way and
come in later patches.

Author: seven-mile

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -674,9 +674,10 @@ class CIRGenBuilderTy : public CIRBaseBuilderTy {
 
   mlir::Value createGetGlobal(mlir::cir::GlobalOp global,
                               bool threadLocal = false) {
-    return create<mlir::cir::GetGlobalOp>(global.getLoc(),
-                                          getPointerTo(global.getSymType()),
-                                          global.getName(), threadLocal);
+    return create<mlir::cir::GetGlobalOp>(
+        global.getLoc(),
+        getPointerTo(global.getSymType(), global.getAddrSpaceAttr()),
+        global.getName(), threadLocal);
   }
 
   mlir::Value createGetBitfield(mlir::Location loc, mlir::Type resultType,

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -895,7 +895,9 @@ static LValue buildGlobalVarDeclLValue(CIRGenFunction &CGF, const Expr *E,
   auto V = CGF.CGM.getAddrOfGlobalVar(VD);
 
   auto RealVarTy = CGF.getTypes().convertTypeForMem(VD->getType());
-  auto realPtrTy = CGF.getBuilder().getPointerTo(RealVarTy);
+  mlir::cir::PointerType realPtrTy = CGF.getBuilder().getPointerTo(
+      RealVarTy, cast_if_present<mlir::cir::AddressSpaceAttr>(
+                     cast<mlir::cir::PointerType>(V.getType()).getAddrSpace()));
   if (realPtrTy != V.getType())
     V = CGF.getBuilder().createBitcast(V.getLoc(), V, realPtrTy);
 

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -637,11 +637,11 @@ mlir::Value CIRGenModule::getGlobalValue(const Decl *D) {
   return CurCGF->symbolTable.lookup(D);
 }
 
-mlir::cir::GlobalOp CIRGenModule::createGlobalOp(CIRGenModule &CGM,
-                                                 mlir::Location loc,
-                                                 StringRef name, mlir::Type t,
-                                                 bool isCst,
-                                                 mlir::Operation *insertPoint) {
+mlir::cir::GlobalOp
+CIRGenModule::createGlobalOp(CIRGenModule &CGM, mlir::Location loc,
+                             StringRef name, mlir::Type t, bool isCst,
+                             mlir::cir::AddressSpaceAttr addrSpace,
+                             mlir::Operation *insertPoint) {
   mlir::cir::GlobalOp g;
   auto &builder = CGM.getBuilder();
   {
@@ -655,7 +655,8 @@ mlir::cir::GlobalOp CIRGenModule::createGlobalOp(CIRGenModule &CGM,
     if (curCGF)
       builder.setInsertionPoint(curCGF->CurFn);
 
-    g = builder.create<mlir::cir::GlobalOp>(loc, name, t, isCst);
+    g = builder.create<mlir::cir::GlobalOp>(
+        loc, name, t, isCst, GlobalLinkageKind::ExternalLinkage, addrSpace);
     if (!curCGF) {
       if (insertPoint)
         CGM.getModule().insert(insertPoint, g);
@@ -742,6 +743,12 @@ void CIRGenModule::replaceGlobal(mlir::cir::GlobalOp Old,
   // If the types does not match, update all references to Old to the new type.
   auto OldTy = Old.getSymType();
   auto NewTy = New.getSymType();
+  mlir::cir::AddressSpaceAttr oldAS = Old.getAddrSpaceAttr();
+  mlir::cir::AddressSpaceAttr newAS = New.getAddrSpaceAttr();
+  // TODO(cir): If the AS differs, we should also update all references.
+  if (oldAS != newAS) {
+    llvm_unreachable("NYI");
+  }
   if (OldTy != NewTy) {
     auto OldSymUses = Old.getSymbolUses(theModule.getOperation());
     if (OldSymUses.has_value()) {
@@ -809,7 +816,7 @@ void CIRGenModule::setTLSMode(mlir::Operation *Op, const VarDecl &D) const {
 /// mangled name but some other type.
 mlir::cir::GlobalOp
 CIRGenModule::getOrCreateCIRGlobal(StringRef MangledName, mlir::Type Ty,
-                                   LangAS AddrSpace, const VarDecl *D,
+                                   LangAS langAS, const VarDecl *D,
                                    ForDefinition_t IsForDefinition) {
   // Lookup the entry, lazily creating it if necessary.
   mlir::cir::GlobalOp Entry;
@@ -818,8 +825,9 @@ CIRGenModule::getOrCreateCIRGlobal(StringRef MangledName, mlir::Type Ty,
     Entry = dyn_cast_or_null<mlir::cir::GlobalOp>(V);
   }
 
-  // unsigned TargetAS = astCtx.getTargetAddressSpace(AddrSpace);
+  mlir::cir::AddressSpaceAttr cirAS = builder.getAddrSpaceAttr(langAS);
   if (Entry) {
+    auto entryCIRAS = Entry.getAddrSpaceAttr();
     if (WeakRefReferences.erase(Entry)) {
       if (D && !D->hasAttr<WeakAttr>()) {
         auto LT = mlir::cir::GlobalLinkageKind::ExternalLinkage;
@@ -837,8 +845,7 @@ CIRGenModule::getOrCreateCIRGlobal(StringRef MangledName, mlir::Type Ty,
     if (langOpts.OpenMP && !langOpts.OpenMPSimd && D)
       getOpenMPRuntime().registerTargetGlobalVariable(D, Entry);
 
-    // TODO(cir): check TargetAS matches Entry address space
-    if (Entry.getSymType() == Ty && !MissingFeatures::addressSpaceInGlobalVar())
+    if (Entry.getSymType() == Ty && entryCIRAS == cirAS)
       return Entry;
 
     // If there are two attempts to define the same mangled name, issue an
@@ -867,14 +874,16 @@ CIRGenModule::getOrCreateCIRGlobal(StringRef MangledName, mlir::Type Ty,
 
     // TODO(cir): LLVM codegen makes sure the result is of the correct type
     // by issuing a address space cast.
+    if (entryCIRAS != cirAS)
+      llvm_unreachable("NYI");
 
     // (If global is requested for a definition, we always need to create a new
     // global, not just return a bitcast.)
     if (!IsForDefinition)
       return Entry;
   }
 
-  // TODO(cir): auto DAddrSpace = GetGlobalVarAddressSpace(D);
+  auto declCIRAS = builder.getAddrSpaceAttr(getGlobalVarAddressSpace(D));
   // TODO(cir): do we need to strip pointer casts for Entry?
 
   auto loc = getLoc(D->getSourceRange());
@@ -883,6 +892,7 @@ CIRGenModule::getOrCreateCIRGlobal(StringRef MangledName, mlir::Type Ty,
   // mark it as such.
   auto GV = CIRGenModule::createGlobalOp(*this, loc, MangledName, Ty,
                                          /*isConstant=*/false,
+                                         /*addrSpace=*/declCIRAS,
                                          /*insertPoint=*/Entry.getOperation());
 
   // If we already created a global with the same mangled name (but different
@@ -992,8 +1002,7 @@ mlir::Value CIRGenModule::getAddrOfGlobalVar(const VarDecl *D, mlir::Type Ty,
 
   bool tlsAccess = D->getTLSKind() != VarDecl::TLS_None;
   auto g = buildGlobal(D, Ty, IsForDefinition);
-  auto ptrTy =
-      mlir::cir::PointerType::get(builder.getContext(), g.getSymType());
+  auto ptrTy = builder.getPointerTo(g.getSymType(), g.getAddrSpaceAttr());
   return builder.create<mlir::cir::GetGlobalOp>(
       getLoc(D->getSourceRange()), ptrTy, g.getSymName(), tlsAccess);
 }
@@ -1076,7 +1085,8 @@ void CIRGenModule::buildGlobalVarDefinition(const clang::VarDecl *D,
   // If this is OpenMP device, check if it is legal to emit this global
   // normally.
   QualType ASTTy = D->getType();
-  if (getLangOpts().OpenCL || getLangOpts().OpenMPIsTargetDevice)
+  if ((getLangOpts().OpenCL && ASTTy->isSamplerT()) ||
+      getLangOpts().OpenMPIsTargetDevice)
     llvm_unreachable("not implemented");
 
   // TODO(cir): LLVM's codegen uses a llvm::TrackingVH here. Is that
@@ -1409,7 +1419,7 @@ LangAS CIRGenModule::getLangTempAllocaAddressSpace() const {
   if (getLangOpts().OpenCL)
     return LangAS::opencl_private;
   if (getLangOpts().SYCLIsDevice || getLangOpts().CUDAIsDevice ||
-    (getLangOpts().OpenMP && getLangOpts().OpenMPIsTargetDevice))
+      (getLangOpts().OpenMP && getLangOpts().OpenMPIsTargetDevice))
     llvm_unreachable("NYI");
   return LangAS::Default;
 }
@@ -3100,3 +3110,25 @@ mlir::cir::SourceLanguage CIRGenModule::getCIRSourceLanguage() {
   // TODO(cir): support remaining source languages.
   llvm_unreachable("CIR does not yet support the given source language");
 }
+
+LangAS CIRGenModule::getGlobalVarAddressSpace(const VarDecl *D) {
+  if (langOpts.OpenCL) {
+    LangAS AS = D ? D->getType().getAddressSpace() : LangAS::opencl_global;
+    assert(AS == LangAS::opencl_global || AS == LangAS::opencl_global_device ||
+           AS == LangAS::opencl_global_host || AS == LangAS::opencl_constant ||
+           AS == LangAS::opencl_local || AS >= LangAS::FirstTargetAddressSpace);
+    return AS;
+  }
+
+  if (langOpts.SYCLIsDevice &&
+      (!D || D->getType().getAddressSpace() == LangAS::Default))
+    llvm_unreachable("NYI");
+
+  if (langOpts.CUDA && langOpts.CUDAIsDevice)
+    llvm_unreachable("NYI");
+
+  if (langOpts.OpenMP)
+    llvm_unreachable("NYI");
+
+  return getTargetCIRGenInfo().getGlobalVarAddressSpace(*this, D);
+}

clang/lib/CIR/CodeGen/CIRGenModule.h

@@ -226,6 +226,7 @@ class CIRGenModule : public CIRGenTypeCache {
   static mlir::cir::GlobalOp
   createGlobalOp(CIRGenModule &CGM, mlir::Location loc, StringRef name,
                  mlir::Type t, bool isCst = false,
+                 mlir::cir::AddressSpaceAttr addrSpace = {},
                  mlir::Operation *insertPoint = nullptr);
 
   // FIXME: Hardcoding priority here is gross.
@@ -328,6 +329,16 @@ class CIRGenModule : public CIRGenTypeCache {
     return (Twine(".compoundLiteral.") + Twine(CompoundLitaralCnt++)).str();
   }
 
+  /// Return the AST address space of the underlying global variable for D, as
+  /// determined by its declaration. Normally this is the same as the address
+  /// space of D's type, but in CUDA, address spaces are associated with
+  /// declarations, not types. If D is nullptr, return the default address
+  /// space for global variable.
+  ///
+  /// For languages without explicit address spaces, if D has default address
+  /// space, target-specific global or constant address space may be returned.
+  LangAS getGlobalVarAddressSpace(const VarDecl *D);
+
   /// Return the AST address space of constant literal, which is used to emit
   /// the constant literal as global variable in LLVM IR.
   /// Note: This is not necessarily the address space of the constant literal

clang/lib/CIR/CodeGen/TargetInfo.cpp

@@ -552,6 +552,15 @@ ABIArgInfo X86_64ABIInfo::classifyReturnType(QualType RetTy) const {
   return ABIArgInfo::getDirect(ResType);
 }
 
+clang::LangAS
+TargetCIRGenInfo::getGlobalVarAddressSpace(cir::CIRGenModule &CGM,
+                                           const clang::VarDecl *D) const {
+  assert(!CGM.getLangOpts().OpenCL &&
+         !(CGM.getLangOpts().CUDA && CGM.getLangOpts().CUDAIsDevice) &&
+         "Address space agnostic languages only");
+  return D ? D->getType().getAddressSpace() : LangAS::Default;
+}
+
 mlir::Value TargetCIRGenInfo::performAddrSpaceCast(
     CIRGenFunction &CGF, mlir::Value Src, mlir::cir::AddressSpaceAttr SrcAddr,
     mlir::cir::AddressSpaceAttr DestAddr, mlir::Type DestTy,

clang/lib/CIR/CodeGen/TargetInfo.h

@@ -62,6 +62,13 @@ class TargetCIRGenInfo {
                            std::vector<LValue> &ResultRegDests,
                            std::string &AsmString, unsigned NumOutputs) const {}
 
+  /// Get target favored AST address space of a global variable for languages
+  /// other than OpenCL and CUDA.
+  /// If \p D is nullptr, returns the default target favored address space
+  /// for global variable.
+  virtual clang::LangAS getGlobalVarAddressSpace(CIRGenModule &CGM,
+                                                 const clang::VarDecl *D) const;
+
   /// Get the CIR address space for alloca.
   virtual mlir::cir::AddressSpaceAttr getCIRAllocaAddressSpace() const {
     // Return the null attribute, which means the target does not care about the

clang/test/CIR/CodeGen/OpenCL/global.cl

@@ -0,0 +1,23 @@
+// RUN: %clang_cc1 -cl-std=CL3.0 -O0 -fclangir -emit-cir -triple spirv64-unknown-unknown %s -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s --check-prefix=CIR
+// RUN: %clang_cc1 -cl-std=CL3.0 -O0 -fclangir -emit-llvm -triple spirv64-unknown-unknown %s -o %t.ll
+// RUN: FileCheck --input-file=%t.ll %s --check-prefix=LLVM
+
+global int a = 13;
+// CIR-DAG: cir.global external addrspace(offload_global) @a = #cir.int<13> : !s32i
+// LLVM-DAG: @a = addrspace(1) global i32 13
+
+global int b = 15;
+// CIR-DAG: cir.global external addrspace(offload_global) @b = #cir.int<15> : !s32i
+// LLVM-DAG: @b = addrspace(1) global i32 15
+
+kernel void test_get_global() {
+  a = b;
+  // CIR:      %[[#ADDRB:]] = cir.get_global @b : !cir.ptr<!s32i, addrspace(offload_global)>
+  // CIR-NEXT: %[[#LOADB:]] = cir.load %[[#ADDRB]] : !cir.ptr<!s32i, addrspace(offload_global)>, !s32i
+  // CIR-NEXT: %[[#ADDRA:]] = cir.get_global @a : !cir.ptr<!s32i, addrspace(offload_global)>
+  // CIR-NEXT: cir.store %[[#LOADB]], %[[#ADDRA]] : !s32i, !cir.ptr<!s32i, addrspace(offload_global)>
+
+  // LLVM:      %[[#LOADB:]] = load i32, ptr addrspace(1) @b, align 4
+  // LLVM-NEXT: store i32 %[[#LOADB]], ptr addrspace(1) @a, align 4
+}