[llvm] [Attributor] Check range size before constant fold load (PR #151359)

[Shilei Tian via llvm-commits]

https://github.com/shiltian updated https://github.com/llvm/llvm-project/pull/151359

>From c6a3ff254aca84b45042d87898a07ec0c30c00d9 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Wed, 15 Oct 2025 14:35:49 -0400
Subject: [PATCH 1/2] [Attributor] Check range size before constant fold load

If the range size doesn't match the type size, it might read wrong data.
---
 llvm/lib/Transforms/IPO/Attributor.cpp        |  3 ++
 .../Attributor/range-and-constant-fold.ll     | 34 +++++++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 llvm/test/Transforms/Attributor/range-and-constant-fold.ll

diff --git a/llvm/lib/Transforms/IPO/Attributor.cpp b/llvm/lib/Transforms/IPO/Attributor.cpp
index 077d29f7499a4..3b59ebbbb9322 100644
--- a/llvm/lib/Transforms/IPO/Attributor.cpp
+++ b/llvm/lib/Transforms/IPO/Attributor.cpp
@@ -272,6 +272,9 @@ AA::getInitialValueForObj(Attributor &A, const AbstractAttribute &QueryingAA,
   }
 
   if (RangePtr && !RangePtr->offsetOrSizeAreUnknown()) {
+    int64_t StorageSize = DL.getTypeStoreSize(&Ty);
+    if (StorageSize != RangePtr->Size)
+      return nullptr;
     APInt Offset = APInt(64, RangePtr->Offset);
     return ConstantFoldLoadFromConst(Initializer, &Ty, Offset, DL);
   }
diff --git a/llvm/test/Transforms/Attributor/range-and-constant-fold.ll b/llvm/test/Transforms/Attributor/range-and-constant-fold.ll
new file mode 100644
index 0000000000000..fcc5e6e0fd52f
--- /dev/null
+++ b/llvm/test/Transforms/Attributor/range-and-constant-fold.ll
@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes=attributor %s -o - | FileCheck %s
+
+ at g = internal unnamed_addr addrspace(4) constant [3 x i8] c"12\00", align 16
+
+define void @foo(i32 %a, i32 %b, ptr %p) {
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]], ptr nofree nonnull writeonly captures(none) dereferenceable(1) [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A]], [[B]]
+; CHECK-NEXT:    br i1 [[CMP]], label %[[L1:.*]], label %[[L2:.*]]
+; CHECK:       [[L1]]:
+; CHECK-NEXT:    br label %[[L3:.*]]
+; CHECK:       [[L2]]:
+; CHECK-NEXT:    br label %[[L3]]
+; CHECK:       [[L3]]:
+; CHECK-NEXT:    [[PHI:%.*]] = phi ptr addrspace(4) [ @g, %[[L1]] ], [ getelementptr inbounds nuw (i8, ptr addrspace(4) @g, i64 1), %[[L2]] ]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(4) [[PHI]], align 1
+; CHECK-NEXT:    store i8 [[LOAD]], ptr [[P]], align 1
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp = icmp ne i32 %a, %b
+  br i1 %cmp, label %l1, label %l2
+l1:
+  br label %l3
+l2:
+  br label %l3
+l3:
+  %phi = phi ptr addrspace(4) [ @g, %l1 ], [ getelementptr inbounds nuw (i8, ptr addrspace(4) @g, i64 1), %l2 ]
+  %load = load i8, ptr addrspace(4) %phi
+  store i8 %load, ptr %p
+  ret void
+}

>From 003c4a92e9f207ca2eac39d235e4310bee68f2d7 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Thu, 23 Oct 2025 23:15:17 -0400
Subject: [PATCH 2/2] add extra test

---
 .../Attributor/range-and-constant-fold.ll        | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/llvm/test/Transforms/Attributor/range-and-constant-fold.ll b/llvm/test/Transforms/Attributor/range-and-constant-fold.ll
index fcc5e6e0fd52f..a8f33092d0d0b 100644
--- a/llvm/test/Transforms/Attributor/range-and-constant-fold.ll
+++ b/llvm/test/Transforms/Attributor/range-and-constant-fold.ll
@@ -3,9 +3,9 @@
 
 @g = internal unnamed_addr addrspace(4) constant [3 x i8] c"12\00", align 16
 
-define void @foo(i32 %a, i32 %b, ptr %p) {
+define void @foo(i32 %a, i32 %b, ptr %p.0, ptr %p.1) {
 ; CHECK-LABEL: define void @foo(
-; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]], ptr nofree nonnull writeonly captures(none) dereferenceable(1) [[P:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]], ptr nofree nonnull writeonly captures(none) dereferenceable(1) [[P_0:%.*]], ptr nofree nonnull writeonly align 4 captures(none) dereferenceable(8) [[P_1:%.*]]) #[[ATTR0:[0-9]+]] {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[A]], [[B]]
 ; CHECK-NEXT:    br i1 [[CMP]], label %[[L1:.*]], label %[[L2:.*]]
@@ -15,8 +15,10 @@ define void @foo(i32 %a, i32 %b, ptr %p) {
 ; CHECK-NEXT:    br label %[[L3]]
 ; CHECK:       [[L3]]:
 ; CHECK-NEXT:    [[PHI:%.*]] = phi ptr addrspace(4) [ @g, %[[L1]] ], [ getelementptr inbounds nuw (i8, ptr addrspace(4) @g, i64 1), %[[L2]] ]
-; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(4) [[PHI]], align 1
-; CHECK-NEXT:    store i8 [[LOAD]], ptr [[P]], align 1
+; CHECK-NEXT:    [[LOAD_SMALL:%.*]] = load i8, ptr addrspace(4) [[PHI]], align 4
+; CHECK-NEXT:    store i8 [[LOAD_SMALL]], ptr [[P_0]], align 1
+; CHECK-NEXT:    [[LOAD_LARGE:%.*]] = load i64, ptr addrspace(4) [[PHI]], align 4
+; CHECK-NEXT:    store i64 [[LOAD_LARGE]], ptr [[P_1]], align 4
 ; CHECK-NEXT:    ret void
 ;
 entry:
@@ -28,7 +30,9 @@ l2:
   br label %l3
 l3:
   %phi = phi ptr addrspace(4) [ @g, %l1 ], [ getelementptr inbounds nuw (i8, ptr addrspace(4) @g, i64 1), %l2 ]
-  %load = load i8, ptr addrspace(4) %phi
-  store i8 %load, ptr %p
+  %load.small = load i8, ptr addrspace(4) %phi
+  store i8 %load.small, ptr %p.0
+  %load.large = load i64, ptr addrspace(4) %phi
+  store i64 %load.large, ptr %p.1
   ret void
 }