[flang][cuda] Lower shared global to the correct NVVM address space

[clementval]

Global with the CUDA shared data attribute needs to be lowered to llvm globals with the correct address space (3). Address space is set from the mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace enum from mlir/Dialect/LLVMIR/NVVMDialect.h

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-flang-codegen

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

Global with the CUDA shared data attribute needs to be lowered to llvm globals with the correct address space (3). Address space is set from the mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace enum from mlir/Dialect/LLVMIR/NVVMDialect.h

---

Full diff: https://github.com/llvm/llvm-project/pull/131368.diff

2 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/CodeGen.cpp (+6)
• (modified) flang/test/Fir/CUDA/cuda-code-gen.mlir (+6)

diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 2cb4cea58c2b0..b54b497ee4ba1 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -47,6 +47,7 @@
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/LLVMIR/Transforms/AddComdats.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
@@ -3145,6 +3146,11 @@ struct GlobalOpConversion : public fir::FIROpConversion<fir::GlobalOp> {
         }
       }
     }
+
+    if (global.getDataAttr() &&
+        *global.getDataAttr() == cuf::DataAttribute::Shared)
+      g.setAddrSpace(mlir::NVVM::NVVMMemorySpace::kSharedMemorySpace);
+
     rewriter.eraseOp(global);
     return mlir::success();
   }
diff --git a/flang/test/Fir/CUDA/cuda-code-gen.mlir b/flang/test/Fir/CUDA/cuda-code-gen.mlir
index 063454799502a..fdd9f1ac12b1f 100644
--- a/flang/test/Fir/CUDA/cuda-code-gen.mlir
+++ b/flang/test/Fir/CUDA/cuda-code-gen.mlir
@@ -198,3 +198,9 @@ func.func @_QMm1Psub1(%arg0: !fir.box<!fir.array<?xi32>> {cuf.data_attr = #cuf.c
 
 // CHECK-LABEL: llvm.func @_QMm1Psub1
 // CHECK-COUNT-2: _FortranACUFAllocDescriptor
+
+// -----
+
+fir.global common @_QPshared_static__shared_mem(dense<0> : vector<28xi8>) {alignment = 8 : i64, data_attr = #cuf.cuda<shared>} : !fir.array<28xi8>
+
+// CHECK: llvm.mlir.global common @_QPshared_static__shared_mem(dense<0> : vector<28xi8>) {addr_space = 3 : i32, alignment = 8 : i64} : !llvm.array<28 x i8>