[clang] [CIR] More atomic load and store (PR #155168)

[Sirui Mu via cfe-commits]

https://github.com/Lancern updated https://github.com/llvm/llvm-project/pull/155168

>From 0b6f6f7897f0bd6d22252ffd05b3aa7bc7740f4e Mon Sep 17 00:00:00 2001
From: Sirui Mu <msrlancern at gmail.com>
Date: Sun, 24 Aug 2025 21:33:31 +0800
Subject: [PATCH] [CIR] More atomic load and store

This patch adds support for `__atomic_load_n` and `__atomic_store_n` that were
missed by #153814.
---
 clang/lib/CIR/CodeGen/CIRGenAtomic.cpp | 21 ++++++++--
 clang/test/CIR/CodeGen/atomic.c        | 53 ++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 4 deletions(-)

diff --git a/clang/lib/CIR/CodeGen/CIRGenAtomic.cpp b/clang/lib/CIR/CodeGen/CIRGenAtomic.cpp
index 80aacb946e7c7..d8981c8c7ca56 100644
--- a/clang/lib/CIR/CodeGen/CIRGenAtomic.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenAtomic.cpp
@@ -265,6 +265,7 @@ static void emitAtomicOp(CIRGenFunction &cgf, AtomicExpr *expr, Address dest,
     llvm_unreachable("already handled!");
 
   case AtomicExpr::AO__c11_atomic_load:
+  case AtomicExpr::AO__atomic_load_n:
   case AtomicExpr::AO__atomic_load: {
     cir::LoadOp load =
         builder.createLoad(loc, ptr, /*isVolatile=*/expr->isVolatile());
@@ -278,6 +279,7 @@ static void emitAtomicOp(CIRGenFunction &cgf, AtomicExpr *expr, Address dest,
   }
 
   case AtomicExpr::AO__c11_atomic_store:
+  case AtomicExpr::AO__atomic_store_n:
   case AtomicExpr::AO__atomic_store: {
     cir::LoadOp loadVal1 = builder.createLoad(loc, val1);
 
@@ -305,13 +307,11 @@ static void emitAtomicOp(CIRGenFunction &cgf, AtomicExpr *expr, Address dest,
 
   case AtomicExpr::AO__opencl_atomic_load:
   case AtomicExpr::AO__hip_atomic_load:
-  case AtomicExpr::AO__atomic_load_n:
   case AtomicExpr::AO__scoped_atomic_load_n:
   case AtomicExpr::AO__scoped_atomic_load:
 
   case AtomicExpr::AO__opencl_atomic_store:
   case AtomicExpr::AO__hip_atomic_store:
-  case AtomicExpr::AO__atomic_store_n:
   case AtomicExpr::AO__scoped_atomic_store:
   case AtomicExpr::AO__scoped_atomic_store_n:
 
@@ -450,6 +450,7 @@ RValue CIRGenFunction::emitAtomicExpr(AtomicExpr *e) {
   case AtomicExpr::AO__c11_atomic_init:
     llvm_unreachable("already handled above with emitAtomicInit");
 
+  case AtomicExpr::AO__atomic_load_n:
   case AtomicExpr::AO__c11_atomic_load:
     break;
 
@@ -461,6 +462,7 @@ RValue CIRGenFunction::emitAtomicExpr(AtomicExpr *e) {
     val1 = emitPointerWithAlignment(e->getVal1());
     break;
 
+  case AtomicExpr::AO__atomic_store_n:
   case AtomicExpr::AO__c11_atomic_store:
     val1 = emitValToTemp(*this, e->getVal1());
     break;
@@ -507,9 +509,20 @@ RValue CIRGenFunction::emitAtomicExpr(AtomicExpr *e) {
   }
 
   bool isStore = e->getOp() == AtomicExpr::AO__c11_atomic_store ||
-                 e->getOp() == AtomicExpr::AO__atomic_store;
+                 e->getOp() == AtomicExpr::AO__opencl_atomic_store ||
+                 e->getOp() == AtomicExpr::AO__hip_atomic_store ||
+                 e->getOp() == AtomicExpr::AO__atomic_store ||
+                 e->getOp() == AtomicExpr::AO__atomic_store_n ||
+                 e->getOp() == AtomicExpr::AO__scoped_atomic_store ||
+                 e->getOp() == AtomicExpr::AO__scoped_atomic_store_n ||
+                 e->getOp() == AtomicExpr::AO__atomic_clear;
   bool isLoad = e->getOp() == AtomicExpr::AO__c11_atomic_load ||
-                e->getOp() == AtomicExpr::AO__atomic_load;
+                e->getOp() == AtomicExpr::AO__opencl_atomic_load ||
+                e->getOp() == AtomicExpr::AO__hip_atomic_load ||
+                e->getOp() == AtomicExpr::AO__atomic_load ||
+                e->getOp() == AtomicExpr::AO__atomic_load_n ||
+                e->getOp() == AtomicExpr::AO__scoped_atomic_load ||
+                e->getOp() == AtomicExpr::AO__scoped_atomic_load_n;
 
   if (!order) {
     // We have evaluated the memory order as an integer constant in orderConst.
diff --git a/clang/test/CIR/CodeGen/atomic.c b/clang/test/CIR/CodeGen/atomic.c
index 6c37fb7cd432d..8b947f795d1d4 100644
--- a/clang/test/CIR/CodeGen/atomic.c
+++ b/clang/test/CIR/CodeGen/atomic.c
@@ -75,6 +75,35 @@ void load(int *ptr) {
 // OGCG:   %{{.+}} = load atomic i32, ptr %{{.+}} seq_cst, align 4
 // OGCG: }
 
+void load_n(int *ptr) {
+  int a;
+  a = __atomic_load_n(ptr, __ATOMIC_RELAXED);
+  a = __atomic_load_n(ptr, __ATOMIC_CONSUME);
+  a = __atomic_load_n(ptr, __ATOMIC_ACQUIRE);
+  a = __atomic_load_n(ptr, __ATOMIC_SEQ_CST);
+}
+
+// CIR-LABEL: @load_n
+// CIR:   %{{.+}} = cir.load align(4) atomic(relaxed) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR:   %{{.+}} = cir.load align(4) atomic(consume) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR:   %{{.+}} = cir.load align(4) atomic(acquire) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR:   %{{.+}} = cir.load align(4) atomic(seq_cst) %{{.+}} : !cir.ptr<!s32i>, !s32i
+// CIR: }
+
+// LLVM-LABEL: @load_n
+// LLVM:   %{{.+}} = load atomic i32, ptr %{{.+}} monotonic, align 4
+// LLVM:   %{{.+}} = load atomic i32, ptr %{{.+}} acquire, align 4
+// LLVM:   %{{.+}} = load atomic i32, ptr %{{.+}} acquire, align 4
+// LLVM:   %{{.+}} = load atomic i32, ptr %{{.+}} seq_cst, align 4
+// LLVM: }
+
+// OGCG-LABEL: @load_n
+// OGCG:   %{{.+}} = load atomic i32, ptr %{{.+}} monotonic, align 4
+// OGCG:   %{{.+}} = load atomic i32, ptr %{{.+}} acquire, align 4
+// OGCG:   %{{.+}} = load atomic i32, ptr %{{.+}} acquire, align 4
+// OGCG:   %{{.+}} = load atomic i32, ptr %{{.+}} seq_cst, align 4
+// OGCG: }
+
 void c11_load(_Atomic(int) *ptr) {
   __c11_atomic_load(ptr, __ATOMIC_RELAXED);
   __c11_atomic_load(ptr, __ATOMIC_CONSUME);
@@ -127,6 +156,30 @@ void store(int *ptr, int x) {
 // OGCG:   store atomic i32 %{{.+}}, ptr %{{.+}} seq_cst, align 4
 // OGCG: }
 
+void store_n(int *ptr, int x) {
+  __atomic_store_n(ptr, x, __ATOMIC_RELAXED);
+  __atomic_store_n(ptr, x, __ATOMIC_RELEASE);
+  __atomic_store_n(ptr, x, __ATOMIC_SEQ_CST);
+}
+
+// CIR-LABEL: @store_n
+// CIR:   cir.store align(4) atomic(relaxed) %{{.+}}, %{{.+}} : !s32i, !cir.ptr<!s32i>
+// CIR:   cir.store align(4) atomic(release) %{{.+}}, %{{.+}} : !s32i, !cir.ptr<!s32i>
+// CIR:   cir.store align(4) atomic(seq_cst) %{{.+}}, %{{.+}} : !s32i, !cir.ptr<!s32i>
+// CIR: }
+
+// LLVM-LABEL: @store_n
+// LLVM:   store atomic i32 %{{.+}}, ptr %{{.+}} monotonic, align 4
+// LLVM:   store atomic i32 %{{.+}}, ptr %{{.+}} release, align 4
+// LLVM:   store atomic i32 %{{.+}}, ptr %{{.+}} seq_cst, align 4
+// LLVM: }
+
+// OGCG-LABEL: @store_n
+// OGCG:   store atomic i32 %{{.+}}, ptr %{{.+}} monotonic, align 4
+// OGCG:   store atomic i32 %{{.+}}, ptr %{{.+}} release, align 4
+// OGCG:   store atomic i32 %{{.+}}, ptr %{{.+}} seq_cst, align 4
+// OGCG: }
+
 void c11_store(_Atomic(int) *ptr, int x) {
   __c11_atomic_store(ptr, x, __ATOMIC_RELAXED);
   __c11_atomic_store(ptr, x, __ATOMIC_RELEASE);