[llvm] [llvm] Add support for atomicrmw and cmpxchg in AssumeBundleBuilder (PR #194630)

[via llvm-commits]

https://github.com/SekaiArendelle created https://github.com/llvm/llvm-project/pull/194630

The assume builder currently only preserves dereferenceable, nonnull, and alignment knowledge for regular load/store instructions and calls. Atomic memory accessing instructions (atomicrmw and cmpxchg) also dereference their pointer operands, but were previously skipped, causing useful knowledge to be lost across these operations.

Add handling for AtomicRMWInst and AtomicCmpXchgInst in AssumeBuilderState::addInstruction(), using the same addAccessedPtr() path as loads and stores. The accessed type is taken from the value operand (atomicrmw) or compare operand (cmpxchg), which corresponds to the in-memory element type, and the alignment is taken from the instruction's explicit alignment.

Add a test to verify that assume bundles are correctly generated before atomicrmw and cmpxchg instructions.

>From 9229c38dc9d3035f49e4578c03938a097f063245 Mon Sep 17 00:00:00 2001
From: Arendelle <Arendelle_ at outlook.com>
Date: Tue, 28 Apr 2026 16:35:03 +0800
Subject: [PATCH] [llvm] Add support for atomicrmw and cmpxchg

The assume builder currently only preserves dereferenceable, nonnull, and
alignment knowledge for regular load/store instructions and calls. Atomic
memory accessing instructions (atomicrmw and cmpxchg) also dereference
their pointer operands, but were previously skipped, causing useful
knowledge to be lost across these operations.

Add handling for AtomicRMWInst and AtomicCmpXchgInst in
AssumeBuilderState::addInstruction(), using the same addAccessedPtr()
path as loads and stores. The accessed type is taken from the value
operand (atomicrmw) or compare operand (cmpxchg), which corresponds to
the in-memory element type, and the alignment is taken from the
instruction's explicit alignment.

Add a test to verify that assume bundles are correctly generated before
atomicrmw and cmpxchg instructions.
---
 .../Transforms/Utils/AssumeBundleBuilder.cpp    |  8 +++++++-
 .../Transforms/Util/assume-builder-atomics.ll   | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/Transforms/Util/assume-builder-atomics.ll

diff --git a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
index 4aebeeda7f262..26c9aa1ff1fcc 100644
--- a/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
+++ b/llvm/lib/Transforms/Utils/AssumeBundleBuilder.cpp
@@ -274,7 +274,13 @@ struct AssumeBuilderState {
       return addAccessedPtr(I, Store->getPointerOperand(),
                             Store->getValueOperand()->getType(),
                             Store->getAlign());
-    // TODO: Add support for the other Instructions.
+    if (auto *RMW = dyn_cast<AtomicRMWInst>(I))
+      return addAccessedPtr(I, RMW->getPointerOperand(),
+                            RMW->getValOperand()->getType(), RMW->getAlign());
+    if (auto *CmpXchg = dyn_cast<AtomicCmpXchgInst>(I))
+      return addAccessedPtr(I, CmpXchg->getPointerOperand(),
+                            CmpXchg->getCompareOperand()->getType(),
+                            CmpXchg->getAlign());
     // TODO: Maybe we should look around and merge with other llvm.assume.
   }
 };
diff --git a/llvm/test/Transforms/Util/assume-builder-atomics.ll b/llvm/test/Transforms/Util/assume-builder-atomics.ll
new file mode 100644
index 0000000000000..0136555289fb8
--- /dev/null
+++ b/llvm/test/Transforms/Util/assume-builder-atomics.ll
@@ -0,0 +1,17 @@
+; RUN: opt -passes='assume-builder,verify' --enable-knowledge-retention -S %s | FileCheck %s
+
+define void @test_atomic_accesses(ptr %p, ptr %q) {
+; CHECK-LABEL: define {{[^@]+}}@test_atomic_accesses
+; CHECK-SAME: (ptr [[P:%.*]], ptr [[Q:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[P]], i64 4), "nonnull"(ptr [[P]]), "align"(ptr [[P]], i64 8) ]
+; CHECK-NEXT:    [[OLD:%.*]] = atomicrmw add ptr [[P]], i32 1 monotonic, align 8
+; CHECK-NEXT:    call void @llvm.assume(i1 true) [ "dereferenceable"(ptr [[Q]], i64 8), "nonnull"(ptr [[Q]]), "align"(ptr [[Q]], i64 16) ]
+; CHECK-NEXT:    [[PAIR:%.*]] = cmpxchg ptr [[Q]], i64 0, i64 1 monotonic monotonic, align 16
+; CHECK-NEXT:    ret void
+;
+entry:
+  %old = atomicrmw add ptr %p, i32 1 monotonic, align 8
+  %pair = cmpxchg ptr %q, i64 0, i64 1 monotonic monotonic, align 16
+  ret void
+}