[flang][CodeGen] Fix address space mismatch for CUF globals in AddrOfOpConversion (llvm#190408)

AddrOfOpConversion in CodeGen.cpp only handled `LLVM::GlobalOp` when
determining the address space for `llvm.mlir.addressof`. When the global
was still a `fir::GlobalOp` (not yet converted), it fell back to address
space 0, breaking CUF constant globals (addr_space 4) and AMDGPU targets
(global addr_space 1).

This extends the upstream fix (llvm#192111, which only covered Constant) to
also handle Shared and Managed CUF data attributes, and returns
`std::nullopt` instead of 0 for non-CUF globals so the target's default
address space is preserved.

Author: khaki3

flang/lib/Optimizer/CodeGen/CodeGen.cpp

@@ -196,26 +196,38 @@ mlir::Value replaceWithAddrOfOrASCast(mlir::ConversionPatternRewriter &rewriter,
   return mlir::LLVM::AddressOfOp::create(rewriter, loc, type, symName);
 }
 
-static std::uint64_t getAddressSpace(fir::AddrOfOp addr,
-                                     mlir::ConversionPatternRewriter &rewriter,
-                                     std::uint64_t defaultAS) {
-  auto global = addr->getParentOfType<mlir::ModuleOp>()
-                    .lookupSymbol<mlir::LLVM::GlobalOp>(addr.getSymbol());
-  if (global)
-    return global.getAddrSpace();
-  auto firGlobal =
-      addr->getParentOfType<mlir::ModuleOp>().lookupSymbol<fir::GlobalOp>(
-          addr.getSymbol());
-  if (firGlobal && firGlobal.getDataAttr() &&
-      *firGlobal.getDataAttr() == cuf::DataAttribute::Constant)
-    return static_cast<unsigned>(mlir::NVVM::NVVMMemorySpace::Constant);
-  return defaultAS;
+/// Return the NVVM address space implied by a CUF data attribute on a
+/// fir::GlobalOp that has not yet been converted to llvm.mlir.global.
+/// Returns std::nullopt if no CUF-specific address space applies.
+static std::optional<unsigned> getCUFAddrSpace(fir::GlobalOp global) {
+  if (auto dataAttr = global.getDataAttr()) {
+    if (*dataAttr == cuf::DataAttribute::Constant)
+      return static_cast<unsigned>(mlir::NVVM::NVVMMemorySpace::Constant);
+    if (*dataAttr == cuf::DataAttribute::Shared)
+      return static_cast<unsigned>(mlir::NVVM::NVVMMemorySpace::Shared);
+    if (*dataAttr == cuf::DataAttribute::Managed)
+      return static_cast<unsigned>(mlir::NVVM::NVVMMemorySpace::Global);
+  }
+  return std::nullopt;
 }
 
 /// Lower `fir.address_of` operation to `llvm.address_of` operation.
 struct AddrOfOpConversion : public fir::FIROpConversion<fir::AddrOfOp> {
   using FIROpConversion::FIROpConversion;
 
+  /// Look up the address space for a symbol in \p mod, handling both
+  /// already-converted llvm.mlir.global and not-yet-converted fir.global.
+  template <typename ModOp>
+  unsigned getAddrSpaceForGlobal(ModOp mod, mlir::SymbolRefAttr sym,
+                                 unsigned fallback) const {
+    if (auto g = mod.template lookupSymbol<mlir::LLVM::GlobalOp>(sym))
+      return g.getAddrSpace();
+    if (auto g = mod.template lookupSymbol<fir::GlobalOp>(sym))
+      if (auto as = getCUFAddrSpace(g))
+        return *as;
+    return fallback;
+  }
+
   llvm::LogicalResult
   matchAndRewrite(fir::AddrOfOp addr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
@@ -224,20 +236,22 @@ struct AddrOfOpConversion : public fir::FIROpConversion<fir::AddrOfOp> {
       auto global = gpuMod.lookupSymbol<mlir::LLVM::GlobalOp>(addr.getSymbol());
       replaceWithAddrOfOrASCast(
           rewriter, addr->getLoc(),
-          global ? global.getAddrSpace() : getGlobalAddressSpace(rewriter),
+          getAddrSpaceForGlobal(gpuMod, addr.getSymbol(),
+                                getGlobalAddressSpace(rewriter)),
           getProgramAddressSpace(rewriter),
           global ? global.getSymName()
                  : addr.getSymbol().getRootReference().getValue(),
           convertType(addr.getType()), addr);
       return mlir::success();
     }
 
-    std::uint64_t globalAS =
-        getAddressSpace(addr, rewriter, getGlobalAddressSpace(rewriter));
-    auto global = addr->getParentOfType<mlir::ModuleOp>()
-                      .lookupSymbol<mlir::LLVM::GlobalOp>(addr.getSymbol());
+    auto mod = addr->getParentOfType<mlir::ModuleOp>();
+    auto global = mod.lookupSymbol<mlir::LLVM::GlobalOp>(addr.getSymbol());
     replaceWithAddrOfOrASCast(
-        rewriter, addr->getLoc(), globalAS, getProgramAddressSpace(rewriter),
+        rewriter, addr->getLoc(),
+        getAddrSpaceForGlobal(mod, addr.getSymbol(),
+                              getGlobalAddressSpace(rewriter)),
+        getProgramAddressSpace(rewriter),
         global ? global.getSymName()
                : addr.getSymbol().getRootReference().getValue(),
         convertType(addr.getType()), addr);

flang/test/Fir/CUDA/cuda-code-gen.mlir

@@ -349,7 +349,7 @@ func.func @sub16_(%arg0: !fir.ref<f32> {fir.bindc_name = "h16"}) attributes {fir
   return
 }
 fir.global linkonce @_QQclXc8657e47c19bb9e89730387c9d99c2da constant : !fir.char<1,38> {
-  %0 = fir.string_lit "/local/home/vclement/lorado/dummy.cuf\00"(38) : !fir.char<1,38>
+  %0 = fir.string_lit "/path/to/source/test/dummy_module.cuf\00"(38) : !fir.char<1,38>
   fir.has_value %0 : !fir.char<1,38>
 }
 fir.global @_QMdevice_dataEd16 {data_attr = #cuf.cuda<constant>} : f32 {
@@ -360,3 +360,45 @@ func.func private @_FortranACUFGetDeviceAddress(!fir.llvm_ptr<i8>, !fir.ref<i8>,
 
 // CHECK-LABEL: llvm.func @sub16_
 // CHECK: llvm.mlir.addressof @_QMdevice_dataEd16 : !llvm.ptr<4>
+
+// -----
+
+// Test that a host-side fir.address_of referencing a fir.global with CUF
+// shared data_attr produces an addrspacecast from ptr<3> to ptr.
+
+fir.global @_QMmodEsval {data_attr = #cuf.cuda<shared>} : i32 {
+  %0 = fir.zero_bits i32
+  fir.has_value %0 : i32
+}
+func.func @_QQhost_shared() {
+  %0 = fir.address_of(@_QMmodEsval) : !fir.ref<i32>
+  return
+}
+
+// CHECK: llvm.mlir.global external @_QMmodEsval() {addr_space = 3 : i32} : i32
+// CHECK-LABEL: llvm.func @_QQhost_shared()
+// CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @_QMmodEsval : !llvm.ptr<3>
+// CHECK: %{{.*}} = llvm.addrspacecast %[[ADDR]] : !llvm.ptr<3> to !llvm.ptr
+
+// -----
+
+// Test that fir.address_of inside gpu.module referencing a managed fir.global
+// produces an addressof with ptr<1> and an addrspacecast.
+
+module attributes {gpu.container_module} {
+  gpu.module @cuda_device_mod {
+    fir.global @_QMmodEmval {data_attr = #cuf.cuda<managed>} : i32 {
+      %0 = fir.zero_bits i32
+      fir.has_value %0 : i32
+    }
+    gpu.func @_QMkernelsPuse_managed() kernel {
+      %0 = fir.address_of(@_QMmodEmval) : !fir.ref<i32>
+      gpu.return
+    }
+  }
+}
+
+// CHECK: llvm.mlir.global external @_QMmodEmval() {addr_space = 1 : i32, nvvm.managed} : i32
+// CHECK-LABEL: gpu.func @_QMkernelsPuse_managed()
+// CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @_QMmodEmval : !llvm.ptr<1>
+// CHECK: %{{.*}} = llvm.addrspacecast %[[ADDR]] : !llvm.ptr<1> to !llvm.ptr