[clang] [CIR][CUDA] Add support for NVVM xchg builtins (PR #211815)

[via cfe-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-clang

Author: Ayokunle Amodu (ayokunle321)

<details>
<summary>Changes</summary>

Adds codegen support for the scoped and unscoped NVVM atomic exchange builtins:
`atom_xchg,` `atom_cta_xchg,` and `atom_sys_xchg.`

These are lowered to the corresponding CIR `cir.atomic.xchg` operations and subsequently lowered to LLVM `atomicrmw xchg` instructions.

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


2 Files Affected:

- (modified) clang/lib/CIR/CodeGen/CIRGenBuiltinNVPTX.cpp (+16-12) 
- (modified) clang/test/CIR/CodeGenCUDA/builtins-nvvm-atomic.cu (+72) 


``````````diff
diff --git a/clang/lib/CIR/CodeGen/CIRGenBuiltinNVPTX.cpp b/clang/lib/CIR/CodeGen/CIRGenBuiltinNVPTX.cpp
index e1e829fc2ab22..a2eac7e80ab06 100644
--- a/clang/lib/CIR/CodeGen/CIRGenBuiltinNVPTX.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenBuiltinNVPTX.cpp
@@ -48,6 +48,19 @@ static mlir::Value makeScopedAtomicRMW(CIRGenFunction &cgf,
   return rmwi->getResult(0);
 }
 
+static mlir::Value makeScopedAtomicXchg(CIRGenFunction &cgf,
+                                        const CallExpr *expr,
+                                        cir::SyncScopeKind scope) {
+  auto &builder = cgf.getBuilder();
+  Address destAddr = cgf.emitPointerWithAlignment(expr->getArg(0));
+  mlir::Value destValue = destAddr.emitRawPointer();
+  mlir::Value val = cgf.emitScalarExpr(expr->getArg(1));
+  auto xchg = cir::AtomicXchgOp::create(
+      builder, cgf.getLoc(expr->getSourceRange()), destValue, val,
+      cir::MemOrder::Relaxed, scope, /*is_volatile=*/false);
+  return xchg.getResult();
+}
+
 std::optional<mlir::Value>
 CIRGenFunction::emitNVPTXBuiltinExpr(unsigned builtinId, const CallExpr *expr) {
   switch (builtinId) {
@@ -89,10 +102,7 @@ CIRGenFunction::emitNVPTXBuiltinExpr(unsigned builtinId, const CallExpr *expr) {
   case NVPTX::BI__nvvm_atom_xchg_gen_i:
   case NVPTX::BI__nvvm_atom_xchg_gen_l:
   case NVPTX::BI__nvvm_atom_xchg_gen_ll:
-    cgm.errorNYI(expr->getSourceRange(),
-                 std::string("unimplemented NVPTX builtin call: ") +
-                     getContext().BuiltinInfo.getName(builtinId));
-    return mlir::Value{};
+    return makeScopedAtomicXchg(*this, expr, cir::SyncScopeKind::System);
   case NVPTX::BI__nvvm_atom_max_gen_i:
   case NVPTX::BI__nvvm_atom_max_gen_l:
   case NVPTX::BI__nvvm_atom_max_gen_ll:
@@ -247,17 +257,11 @@ CIRGenFunction::emitNVPTXBuiltinExpr(unsigned builtinId, const CallExpr *expr) {
   case NVPTX::BI__nvvm_atom_cta_xchg_gen_i:
   case NVPTX::BI__nvvm_atom_cta_xchg_gen_l:
   case NVPTX::BI__nvvm_atom_cta_xchg_gen_ll:
-    cgm.errorNYI(expr->getSourceRange(),
-                 std::string("unimplemented NVPTX builtin call: ") +
-                     getContext().BuiltinInfo.getName(builtinId));
-    return mlir::Value{};
+    return makeScopedAtomicXchg(*this, expr, cir::SyncScopeKind::Workgroup);
   case NVPTX::BI__nvvm_atom_sys_xchg_gen_i:
   case NVPTX::BI__nvvm_atom_sys_xchg_gen_l:
   case NVPTX::BI__nvvm_atom_sys_xchg_gen_ll:
-    cgm.errorNYI(expr->getSourceRange(),
-                 std::string("unimplemented NVPTX builtin call: ") +
-                     getContext().BuiltinInfo.getName(builtinId));
-    return mlir::Value{};
+    return makeScopedAtomicXchg(*this, expr, cir::SyncScopeKind::System);
   case NVPTX::BI__nvvm_atom_cta_max_gen_i:
   case NVPTX::BI__nvvm_atom_cta_max_gen_ui:
   case NVPTX::BI__nvvm_atom_cta_max_gen_l:
diff --git a/clang/test/CIR/CodeGenCUDA/builtins-nvvm-atomic.cu b/clang/test/CIR/CodeGenCUDA/builtins-nvvm-atomic.cu
index a5ed80759d2ad..1ca877017beb2 100644
--- a/clang/test/CIR/CodeGenCUDA/builtins-nvvm-atomic.cu
+++ b/clang/test/CIR/CodeGenCUDA/builtins-nvvm-atomic.cu
@@ -228,6 +228,30 @@ __device__ void test_atom_min_gen_ull(unsigned long long *p, unsigned long long
   __nvvm_atom_min_gen_ull(p, val);
 }
 
+// CIR-LABEL: @_Z20test_atom_xchg_gen_iPii
+// CIR: cir.atomic.xchg relaxed syncscope(system) %{{.*}}, %{{.*}} : (!cir.ptr<!s32i>, !s32i) -> !s32i
+// LLVM-LABEL: @_Z20test_atom_xchg_gen_iPii
+// LLVM: atomicrmw xchg ptr %{{.*}}, i32 %{{.*}} monotonic, align 4
+__device__ void test_atom_xchg_gen_i(int *p, int val) {
+  __nvvm_atom_xchg_gen_i(p, val);
+}
+
+// CIR-LABEL: @_Z20test_atom_xchg_gen_lPll
+// CIR: cir.atomic.xchg relaxed syncscope(system) %{{.*}}, %{{.*}} : (!cir.ptr<!s64i>, !s64i) -> !s64i
+// LLVM-LABEL: @_Z20test_atom_xchg_gen_lPll
+// LLVM: atomicrmw xchg ptr %{{.*}}, i64 %{{.*}} monotonic, align 8
+__device__ void test_atom_xchg_gen_l(long *p, long val) {
+  __nvvm_atom_xchg_gen_l(p, val);
+}
+
+// CIR-LABEL: @_Z21test_atom_xchg_gen_llPxx
+// CIR: cir.atomic.xchg relaxed syncscope(system) %{{.*}}, %{{.*}} : (!cir.ptr<!s64i>, !s64i) -> !s64i
+// LLVM-LABEL: @_Z21test_atom_xchg_gen_llPxx
+// LLVM: atomicrmw xchg ptr %{{.*}}, i64 %{{.*}} monotonic, align 8
+__device__ void test_atom_xchg_gen_ll(long long *p, long long val) {
+  __nvvm_atom_xchg_gen_ll(p, val);
+}
+
 // CIR-LABEL: @_Z23test_atom_cta_add_gen_iPii
 // CIR: cir.atomic.fetch add relaxed syncscope(workgroup) fetch_first %{{.*}}, %{{.*}} : (!cir.ptr<!s32i>, !s32i) -> !s32i
 // LLVM-LABEL: @_Z23test_atom_cta_add_gen_iPii
@@ -675,3 +699,51 @@ __device__ void test_atom_cta_xor_gen_ll(long long *p, long long val) {
 __device__ void test_atom_sys_xor_gen_ll(long long *p, long long val) {
   __nvvm_atom_sys_xor_gen_ll(p, val);
 }
+
+// CIR-LABEL: @_Z24test_atom_cta_xchg_gen_iPii
+// CIR: cir.atomic.xchg relaxed syncscope(workgroup) %{{.*}}, %{{.*}} : (!cir.ptr<!s32i>, !s32i) -> !s32i
+// LLVM-LABEL: @_Z24test_atom_cta_xchg_gen_iPii
+// LLVM: atomicrmw xchg ptr %{{.*}}, i32 %{{.*}} syncscope("block") monotonic, align 4
+__device__ void test_atom_cta_xchg_gen_i(int *p, int val) {
+  __nvvm_atom_cta_xchg_gen_i(p, val);
+}
+
+// CIR-LABEL: @_Z24test_atom_sys_xchg_gen_iPii
+// CIR: cir.atomic.xchg relaxed syncscope(system) %{{.*}}, %{{.*}} : (!cir.ptr<!s32i>, !s32i) -> !s32i
+// LLVM-LABEL: @_Z24test_atom_sys_xchg_gen_iPii
+// LLVM: atomicrmw xchg ptr %{{.*}}, i32 %{{.*}} monotonic, align 4
+__device__ void test_atom_sys_xchg_gen_i(int *p, int val) {
+  __nvvm_atom_sys_xchg_gen_i(p, val);
+}
+
+// CIR-LABEL: @_Z24test_atom_cta_xchg_gen_lPll
+// CIR: cir.atomic.xchg relaxed syncscope(workgroup) %{{.*}}, %{{.*}} : (!cir.ptr<!s64i>, !s64i) -> !s64i
+// LLVM-LABEL: @_Z24test_atom_cta_xchg_gen_lPll
+// LLVM: atomicrmw xchg ptr %{{.*}}, i64 %{{.*}} syncscope("block") monotonic, align 8
+__device__ void test_atom_cta_xchg_gen_l(long *p, long val) {
+  __nvvm_atom_cta_xchg_gen_l(p, val);
+}
+
+// CIR-LABEL: @_Z24test_atom_sys_xchg_gen_lPll
+// CIR: cir.atomic.xchg relaxed syncscope(system) %{{.*}}, %{{.*}} : (!cir.ptr<!s64i>, !s64i) -> !s64i
+// LLVM-LABEL: @_Z24test_atom_sys_xchg_gen_lPll
+// LLVM: atomicrmw xchg ptr %{{.*}}, i64 %{{.*}} monotonic, align 8
+__device__ void test_atom_sys_xchg_gen_l(long *p, long val) {
+  __nvvm_atom_sys_xchg_gen_l(p, val);
+}
+
+// CIR-LABEL: @_Z25test_atom_cta_xchg_gen_llPxx
+// CIR: cir.atomic.xchg relaxed syncscope(workgroup) %{{.*}}, %{{.*}} : (!cir.ptr<!s64i>, !s64i) -> !s64i
+// LLVM-LABEL: @_Z25test_atom_cta_xchg_gen_llPxx
+// LLVM: atomicrmw xchg ptr %{{.*}}, i64 %{{.*}} syncscope("block") monotonic, align 8
+__device__ void test_atom_cta_xchg_gen_ll(long long *p, long long val) {
+  __nvvm_atom_cta_xchg_gen_ll(p, val);
+}
+
+// CIR-LABEL: @_Z25test_atom_sys_xchg_gen_llPxx
+// CIR: cir.atomic.xchg relaxed syncscope(system) %{{.*}}, %{{.*}} : (!cir.ptr<!s64i>, !s64i) -> !s64i
+// LLVM-LABEL: @_Z25test_atom_sys_xchg_gen_llPxx
+// LLVM: atomicrmw xchg ptr %{{.*}}, i64 %{{.*}} monotonic, align 8
+__device__ void test_atom_sys_xchg_gen_ll(long long *p, long long val) {
+  __nvvm_atom_sys_xchg_gen_ll(p, val);
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/211815