[llvm] [llubi] Update the implementation of assume align operand bundle (PR #203775)

[Zhige Chen via llvm-commits]

https://github.com/nofe1248 created https://github.com/llvm/llvm-project/pull/203775

This PR updates the implementation for the `"align"(ptr %p, i64 %align, i64 %offset)` operand bundle of `@llvm.assume`: The old implementation reduced align/offset to the largest shared power of two. But LangRef requires checking %alloc - %offset in this case.

>From 4b3520b36b64e2ff8a340763c2cd595bdb19b3d6 Mon Sep 17 00:00:00 2001
From: Zhige Chen <zhigec_cpp at outlook.com>
Date: Sun, 14 Jun 2026 23:55:35 +0800
Subject: [PATCH] [llubi] Update the implementation of assume align operand
 bundle

---
 llvm/test/tools/llubi/assume_align_offset.ll    | 16 ++++++++++++++++
 llvm/test/tools/llubi/assume_operand_bundles.ll |  4 ++--
 llvm/tools/llubi/lib/Interpreter.cpp            | 16 ++++++----------
 3 files changed, 24 insertions(+), 12 deletions(-)
 create mode 100644 llvm/test/tools/llubi/assume_align_offset.ll

diff --git a/llvm/test/tools/llubi/assume_align_offset.ll b/llvm/test/tools/llubi/assume_align_offset.ll
new file mode 100644
index 0000000000000..4c13b6f959db3
--- /dev/null
+++ b/llvm/test/tools/llubi/assume_align_offset.ll
@@ -0,0 +1,16 @@
+; NOTE: Assertions have been autogenerated by utils/update_llubi_test_checks.py UTC_ARGS: --version 6
+; RUN: not llubi --verbose < %s 2>&1 | FileCheck %s
+
+define void @main() {
+  %alloc = alloca i32
+  call void @llvm.assume(i1 true) ["align"(ptr %alloc, i32 17, i32 8)]
+  call void @llvm.assume(i1 true) ["align"(ptr %alloc, i32 17, i32 4)]
+  ret void
+}
+; CHECK: Entering function: main
+; CHECK-NEXT:   %alloc = alloca i32, align 4 => ptr 0x8 [alloc]
+; CHECK-NEXT:   call void @llvm.assume(i1 true) [ "align"(ptr %alloc, i32 17, i32 8) ]
+; CHECK-NEXT: Stacktrace:
+; CHECK-NEXT: #0   call void @llvm.assume(i1 true) [ "align"(ptr %alloc, i32 17, i32 4) ] at @main <stdin>:10
+; CHECK-NEXT: Immediate UB detected: Assume on pointer ptr 0x8 [alloc] with a nonzero adjusted address and a non-power-of-two alignment 17.
+; CHECK-NEXT: error: Execution of function 'main' failed.
diff --git a/llvm/test/tools/llubi/assume_operand_bundles.ll b/llvm/test/tools/llubi/assume_operand_bundles.ll
index e9abb574e0a1b..0c13993b08cf6 100644
--- a/llvm/test/tools/llubi/assume_operand_bundles.ll
+++ b/llvm/test/tools/llubi/assume_operand_bundles.ll
@@ -4,7 +4,7 @@
 define void @assume_align_dynamic(ptr %p, i32 %align) {
   call void @llvm.assume(i1 true) ["align"(ptr %p, i32 4)]
   call void @llvm.assume(i1 true) ["align"(ptr %p, i32 %align)]
-  call void @llvm.assume(i1 true) ["align"(ptr %p, i32 %align, i32 20)]
+  call void @llvm.assume(i1 true) ["align"(ptr %p, i32 %align, i32 0)]
   ret void
 }
 
@@ -31,7 +31,7 @@ define void @main() {
 ; CHECK-NEXT:   i32 %align = i32 8
 ; CHECK-NEXT:   call void @llvm.assume(i1 true) [ "align"(ptr %p, i32 4) ]
 ; CHECK-NEXT:   call void @llvm.assume(i1 true) [ "align"(ptr %p, i32 %align) ]
-; CHECK-NEXT:   call void @llvm.assume(i1 true) [ "align"(ptr %p, i32 %align, i32 20) ]
+; CHECK-NEXT:   call void @llvm.assume(i1 true) [ "align"(ptr %p, i32 %align, i32 0) ]
 ; CHECK-NEXT:   ret void
 ; CHECK-NEXT: Exiting function: assume_align_dynamic
 ; CHECK-NEXT:   call void @assume_align_dynamic(ptr %alloc, i32 8)
diff --git a/llvm/tools/llubi/lib/Interpreter.cpp b/llvm/tools/llubi/lib/Interpreter.cpp
index 5769d0a4a9c4d..a1460f20e6c58 100644
--- a/llvm/tools/llubi/lib/Interpreter.cpp
+++ b/llvm/tools/llubi/lib/Interpreter.cpp
@@ -830,25 +830,21 @@ class InstExecutor : public InstVisitor<InstExecutor, void>,
           switch (Kind) {
           case Attribute::Alignment: {
             // Alignment assumptions should have 2 or 3 arguments.
-            // If there are two integer arguments, use the largest power of 2
-            // that divides them as the alignment.
             APInt Alignment = getIntNonPoison(getValue(GetBundleArg(1)));
+            APInt CheckedAddr = WasOnPtr.address();
             if (OBU.Inputs.size() == 3) {
               APInt Offset = getIntNonPoison(getValue(GetBundleArg(2)));
-              if (!Alignment.isZero() || !Offset.isZero())
-                Alignment = APInt::getOneBitSet(
-                    std::max(Alignment.getBitWidth(), Offset.getBitWidth()),
-                    std::min(Alignment.countr_zero(), Offset.countr_zero()));
+              CheckedAddr -= Offset.zextOrTrunc(CheckedAddr.getBitWidth());
             }
             if (!Alignment.isPowerOf2()) {
-              if (!WasOnPtr.address().isZero())
-                reportImmediateUB() << "Assume on nonzero pointer " << WasOn
-                                    << " with a "
+              if (!CheckedAddr.isZero())
+                reportImmediateUB() << "Assume on pointer " << WasOn
+                                    << " with a nonzero adjusted address and a "
                                        "non-power-of-two alignment "
                                     << Alignment << '.';
               break;
             }
-            if (WasOnPtr.address().countr_zero() < Alignment.logBase2())
+            if (CheckedAddr.countr_zero() < Alignment.logBase2())
               reportImmediateUB()
                   << "The pointer " << WasOn << " violates align(" << Alignment
                   << ") assumption.";