[CIR][CUDA] Emit address space casts on lowering

Author: AdUhTkJm

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

@@ -396,9 +396,25 @@ class CIRBaseBuilderTy : public mlir::OpBuilder {
 
   mlir::Value createGetGlobal(mlir::Location loc, cir::GlobalOp global,
                               bool threadLocal = false) {
-    return create<cir::GetGlobalOp>(
+    auto getGlobal = create<cir::GetGlobalOp>(
         loc, getPointerTo(global.getSymType(), global.getAddrSpaceAttr()),
         global.getName(), threadLocal);
+
+    // When AST attribute is not present, the global is a temporary,
+    // and actual & expected address spaces won't mismatch for temporaries.
+    if (global.getAst()) {
+      cir::ASTVarDeclInterface varDecl = global.getAstAttr();
+      auto gpuAS = getAddrSpaceAttr(varDecl.getExpectedAS());
+      if (gpuAS != global.getAddrSpaceAttr()) {
+        auto oldTy = mlir::cast<cir::PointerType>(getGlobal.getType());
+        auto newTy =
+            cir::PointerType::get(oldTy.getPointee(), /*addrspace=*/gpuAS);
+        auto cast = createAddrSpaceCast(loc, getGlobal, newTy);
+        return cast;
+      }
+    }
+
+    return getGlobal;
   }
 
   mlir::Value createGetGlobal(cir::GlobalOp global, bool threadLocal = false) {

clang/include/clang/CIR/Interfaces/ASTAttrInterfaces.td

@@ -68,6 +68,11 @@ let cppNamespace = "::cir" in {
         /*defaultImplementation=*/ [{
           return $_attr.getAst()->getTLSKind();
         }]
+      >,
+      InterfaceMethod<"", "clang::LangAS", "getExpectedAS", (ins), [{}],
+        /*defaultImplementation=*/[{
+          return $_attr.getAst()->getType().getAddressSpace();
+        }]
       >
     ];
   }

clang/lib/CIR/CodeGen/CIRGenDecl.cpp

@@ -470,7 +470,7 @@ CIRGenModule::getOrCreateStaticVarDecl(const VarDecl &D,
     Name = getStaticDeclName(*this, D);
 
   mlir::Type LTy = getTypes().convertTypeForMem(Ty);
-  cir::AddressSpaceAttr AS =
+  cir::AddressSpaceAttr actualAS =
       builder.getAddrSpaceAttr(getGlobalVarAddressSpace(&D));
 
   // OpenCL variables in local address space and CUDA shared
@@ -482,8 +482,9 @@ CIRGenModule::getOrCreateStaticVarDecl(const VarDecl &D,
            !D.hasAttr<CUDASharedAttr>())
     Init = builder.getZeroInitAttr(convertType(Ty));
 
-  cir::GlobalOp GV = builder.createVersionedGlobal(
-      getModule(), getLoc(D.getLocation()), Name, LTy, false, Linkage, AS);
+  cir::GlobalOp GV =
+      builder.createVersionedGlobal(getModule(), getLoc(D.getLocation()), Name,
+                                    LTy, false, Linkage, actualAS);
   // TODO(cir): infer visibility from linkage in global op builder.
   GV.setVisibility(getMLIRVisibilityFromCIRLinkage(Linkage));
   GV.setInitialValueAttr(Init);
@@ -497,14 +498,15 @@ CIRGenModule::getOrCreateStaticVarDecl(const VarDecl &D,
 
   setGVProperties(GV, &D);
 
-  // OG checks if the expected address space, denoted by the type, is the
-  // same as the actual address space indicated by attributes. If they aren't
-  // the same, an addrspacecast is emitted when this variable is accessed.
-  // In CIR however, cir.get_global alreadys carries that information in
-  // !cir.ptr type - if this global is in OpenCL local address space, then its
-  // type would be !cir.ptr<..., addrspace(offload_local)>. Therefore we don't
-  // need an explicit address space cast in CIR: they will get emitted when
-  // lowering to LLVM IR.
+  // OG checks whether the expected address space (AS), denoted by
+  // __attributes__((addrspace(n))), is the same as the actual AS indicated by
+  // other attributes (such as __device__ in CUDA). If they aren't the same, an
+  // addrspacecast is emitted when this variable is accessed, which means we
+  // need it in this function. In CIR however, since we access globals by
+  // `cir.get_global`, we won't emit a cast for GlobalOp here. Instead, we
+  // record the AST, and create a CastOp in
+  // `CIRGenBaseBuilder::createGetGlobal`.
+  GV.setAstAttr(cir::ASTVarDeclAttr::get(&getMLIRContext(), &D));
 
   // Ensure that the static local gets initialized by making sure the parent
   // function gets emitted eventually.
@@ -617,7 +619,10 @@ void CIRGenFunction::emitStaticVarDecl(const VarDecl &D,
   // TODO(cir): we should have a way to represent global ops as values without
   // having to emit a get global op. Sometimes these emissions are not used.
   auto addr = getBuilder().createGetGlobal(globalOp);
-  auto getAddrOp = mlir::cast<cir::GetGlobalOp>(addr.getDefiningOp());
+  auto definingOp = addr.getDefiningOp();
+  bool hasCast = isa<cir::CastOp>(definingOp);
+  auto getAddrOp = mlir::cast<cir::GetGlobalOp>(
+      hasCast ? definingOp->getOperand(0).getDefiningOp() : definingOp);
 
   CharUnits alignment = getContext().getDeclAlign(&D);
 
@@ -633,7 +638,7 @@ void CIRGenFunction::emitStaticVarDecl(const VarDecl &D,
     llvm_unreachable("VLAs are NYI");
 
   // Save the type in case adding the initializer forces a type change.
-  auto expectedType = addr.getType();
+  auto expectedType = cast<cir::PointerType>(addr.getType());
 
   auto var = globalOp;
 
@@ -678,7 +683,25 @@ void CIRGenFunction::emitStaticVarDecl(const VarDecl &D,
   //
   // FIXME: It is really dangerous to store this in the map; if anyone
   // RAUW's the GV uses of this constant will be invalid.
-  auto castedAddr = builder.createBitcast(getAddrOp.getAddr(), expectedType);
+  mlir::Value castedAddr;
+  if (!hasCast)
+    castedAddr = builder.createBitcast(getAddrOp.getAddr(), expectedType);
+  else {
+    // If there is an extra CastOp from createGetGlobal, we need to remove the
+    // existing addrspacecast, then supply a bitcast and a new addrspacecast:
+    //    %1 = cir.get_global @addr
+    //    %2 = cir.cast(addrspacecast, %1) <--- remove
+    //    %2 = cir.cast(bitcast, %1)       <--- insert
+    //    %3 = cir.cast(addrspacecast, %2) <--- insert
+    definingOp->erase();
+
+    auto expectedTypeWithAS = cir::PointerType::get(
+        expectedType.getPointee(), getAddrOp.getType().getAddrSpace());
+    auto converted =
+        builder.createBitcast(getAddrOp.getAddr(), expectedTypeWithAS);
+    castedAddr = builder.createAddrSpaceCast(converted, expectedType);
+  }
+
   LocalDeclMap.find(&D)->second = Address(castedAddr, elemTy, alignment);
   CGM.setStaticLocalDeclAddress(&D, var);
 

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -1267,9 +1267,15 @@ mlir::Value CIRGenModule::getAddrOfGlobalVar(const VarDecl *d, mlir::Type ty,
 
   bool tlsAccess = d->getTLSKind() != VarDecl::TLS_None;
   auto g = getOrCreateCIRGlobal(d, ty, isForDefinition);
-  auto ptrTy = builder.getPointerTo(g.getSymType(), g.getAddrSpaceAttr());
-  return builder.create<cir::GetGlobalOp>(getLoc(d->getSourceRange()), ptrTy,
-                                          g.getSymName(), tlsAccess);
+  mlir::Value globalAddr =
+      builder.createGetGlobal(getLoc(d->getSourceRange()), g);
+  auto definingOp = globalAddr.getDefiningOp();
+  auto getGlobalOp = mlir::cast<cir::GetGlobalOp>(
+      isa<cir::GetGlobalOp>(definingOp)
+          ? definingOp
+          : definingOp->getOperand(0).getDefiningOp());
+  getGlobalOp.setTls(tlsAccess);
+  return globalAddr;
 }
 
 cir::GlobalViewAttr

clang/test/CIR/CodeGen/CUDA/address-spaces.cu

@@ -5,15 +5,35 @@
 // RUN:            %s -o %t.cir
 // RUN: FileCheck --check-prefix=CIR --input-file=%t.cir %s
 
+// RUN: %clang_cc1 -triple nvptx64-nvidia-cuda -fclangir \
+// RUN:            -fcuda-is-device -emit-llvm -target-sdk-version=12.3 \
+// RUN:            %s -o %t.ll
+// RUN: FileCheck --check-prefix=LLVM --input-file=%t.ll %s
+
+__device__ int k;
+
 __global__ void fn() {
   int i = 0;
   __shared__ int j;
   j = i;
+  k = i;
 }
 
 // CIR: cir.global "private" internal dsolocal addrspace(offload_local) @_ZZ2fnvE1j : !s32i
-// CIR: cir.func @_Z2fnv
-// CIR: [[Local:%[0-9]+]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["i", init]
-// CIR: [[Shared:%[0-9]+]] = cir.get_global @_ZZ2fnvE1j : !cir.ptr<!s32i, addrspace(offload_local)>
-// CIR: [[Tmp:%[0-9]+]] = cir.load [[Local]] : !cir.ptr<!s32i>, !s32i
-// CIR: cir.store [[Tmp]], [[Shared]] : !s32i, !cir.ptr<!s32i, addrspace(offload_local)>
+// CIR: cir.func @_Z2fnv() {{.*}} {
+// CIR:   %[[#Local:]] = cir.alloca !s32i, !cir.ptr<!s32i>
+// CIR:   %[[#Shared:]] = cir.get_global @_ZZ2fnvE1j : !cir.ptr<!s32i, addrspace(offload_local)>
+// CIR:   %[[#Converted:]] = cir.cast(address_space, %[[#Shared]] :
+// CIR-SAME: !cir.ptr<!s32i, addrspace(offload_local)>), !cir.ptr<!s32i>
+// CIR:   %[[#Tmp:]] = cir.load %[[#Local]] : !cir.ptr<!s32i>, !s32i
+// CIR:   cir.store %[[#Tmp]], %[[#Converted]] : !s32i, !cir.ptr<!s32i>
+// CIR: }
+
+// LLVM: @_ZZ2fnvE1j = internal addrspace(3) global i32 undef
+// LLVM: define dso_local ptx_kernel void @_Z2fnv() #{{.*}} {
+// LLVM:   %[[#T1:]] = alloca i32, i64 1
+// LLVM:   store i32 0, ptr %[[#T1]]
+// LLVM:   %[[#T2:]] = load i32, ptr %[[#T1]]
+// LLVM:   store i32 %[[#T2]], ptr addrspacecast (ptr addrspace(3) @_ZZ2fnvE1j to ptr)
+// LLVM:   ret void
+// LLVM: }

clang/test/CIR/CodeGen/OpenCL/static-vardecl.cl

@@ -5,11 +5,11 @@
 
 kernel void test_static(int i) {
   static global int b = 15;
-  // CIR-DAG: cir.global "private" internal dsolocal addrspace(offload_global) @test_static.b = #cir.int<15> : !s32i {alignment = 4 : i64}
+  // CIR-DAG: cir.global "private" internal dsolocal addrspace(offload_global) @test_static.b = #cir.int<15> : !s32i {alignment = 4 : i64, ast = {{.*}}}
   // LLVM-DAG: @test_static.b = internal addrspace(1) global i32 15
 
   local int c;
-  // CIR-DAG: cir.global "private" internal dsolocal addrspace(offload_local) @test_static.c : !s32i {alignment = 4 : i64}
+  // CIR-DAG: cir.global "private" internal dsolocal addrspace(offload_local) @test_static.c : !s32i {alignment = 4 : i64, ast = {{.*}}}
   // LLVM-DAG: @test_static.c = internal addrspace(3) global i32 undef
 
   // CIR-DAG: %[[#ADDRB:]] = cir.get_global @test_static.b : !cir.ptr<!s32i, addrspace(offload_global)>

clang/test/CIR/CodeGen/const-array.c

@@ -4,7 +4,7 @@ void bar() {
   const int arr[1] = {1};
 }
 
-// CHECK: cir.global "private" constant internal dsolocal @bar.arr = #cir.const_array<[#cir.int<1> : !s32i]> : !cir.array<!s32i x 1> {alignment = 4 : i64}
+// CHECK: cir.global "private" constant internal dsolocal @bar.arr = #cir.const_array<[#cir.int<1> : !s32i]> : !cir.array<!s32i x 1> {alignment = 4 : i64, ast = {{.*}}}
 // CHECK: cir.func no_proto @bar()
 // CHECK:   {{.*}} = cir.get_global @bar.arr : !cir.ptr<!cir.array<!s32i x 1>>