[llvm] [InstCombine] Support offsets in `memset` to load forwarding (PR #151924)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Pedro Lobo (pedroclobo)

<details>
<summary>Changes</summary>

Adds support for load offsets when performing `memset` load forwarding.

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


3 Files Affected:

- (modified) llvm/lib/Analysis/Loads.cpp (+13-4) 
- (modified) llvm/test/Analysis/GlobalsModRef/memset-escape.ll (+1-4) 
- (modified) llvm/test/Transforms/InstCombine/load-store-forward.ll (+38-4) 


``````````diff
diff --git a/llvm/lib/Analysis/Loads.cpp b/llvm/lib/Analysis/Loads.cpp
index 6fc81d787c220..6bc5124cc02f0 100644
--- a/llvm/lib/Analysis/Loads.cpp
+++ b/llvm/lib/Analysis/Loads.cpp
@@ -631,9 +631,14 @@ static Value *getAvailableLoadStore(Instruction *Inst, const Value *Ptr,
     if (!Val || !Len)
       return nullptr;
 
-    // TODO: Handle offsets.
-    Value *Dst = MSI->getDest();
-    if (!AreEquivalentAddressValues(Dst, Ptr))
+    // Handle offsets.
+    int64_t StoreOffset = 0, LoadOffset = 0;
+    const Value *StoreBase =
+        GetPointerBaseWithConstantOffset(MSI->getDest(), StoreOffset, DL);
+    const Value *LoadBase =
+        GetPointerBaseWithConstantOffset(Ptr, LoadOffset, DL);
+    int64_t Offset = LoadOffset - StoreOffset;
+    if (StoreBase != LoadBase || Offset < 0)
       return nullptr;
 
     if (IsLoadCSE)
@@ -645,7 +650,11 @@ static Value *getAvailableLoadStore(Instruction *Inst, const Value *Ptr,
 
     // Make sure the read bytes are contained in the memset.
     uint64_t LoadSize = LoadTypeSize.getFixedValue();
-    if ((Len->getValue() * 8).ult(LoadSize))
+    if ((Len->getValue() * 8).ult(LoadSize + Offset * 8))
+      return nullptr;
+
+    // If there is an offset, make sure the load size is a multiple of 8.
+    if (Offset && LoadSize % 8 != 0)
       return nullptr;
 
     APInt Splat = LoadSize >= 8 ? APInt::getSplat(LoadSize, Val->getValue())
diff --git a/llvm/test/Analysis/GlobalsModRef/memset-escape.ll b/llvm/test/Analysis/GlobalsModRef/memset-escape.ll
index a84987c3bc695..faf2fa5a95953 100644
--- a/llvm/test/Analysis/GlobalsModRef/memset-escape.ll
+++ b/llvm/test/Analysis/GlobalsModRef/memset-escape.ll
@@ -10,11 +10,8 @@ target triple = "x86_64-apple-macosx10.10.0"
 ; @a after the memset.
 
 ; CHECK-LABEL: @main
-; CHECK: call void @llvm.memset.p0.i64{{.*}} @a
 ; CHECK: store i32 3
-; CHECK: load i32, ptr getelementptr {{.*}} @a
-; CHECK: icmp eq i32
-; CHECK: br i1
+; CHECK: ret i32 0
 
 define i32 @main() {
 entry:
diff --git a/llvm/test/Transforms/InstCombine/load-store-forward.ll b/llvm/test/Transforms/InstCombine/load-store-forward.ll
index 9a5db318df5e7..e3f5b81e87140 100644
--- a/llvm/test/Transforms/InstCombine/load-store-forward.ll
+++ b/llvm/test/Transforms/InstCombine/load-store-forward.ll
@@ -365,13 +365,10 @@ define i32 @load_after_memset_unknown(ptr %a, i8 %byte) {
   ret i32 %v
 }
 
-; TODO: Handle load at offset.
 define i32 @load_after_memset_0_offset(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0_offset(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)
-; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 4
-; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[GEP]], align 4
-; CHECK-NEXT:    ret i32 [[V]]
+; CHECK-NEXT:    ret i32 0
 ;
   call void @llvm.memset.p0.i64(ptr %a, i8 0, i64 16, i1 false)
   %gep = getelementptr i8, ptr %a, i64 4
@@ -379,6 +376,43 @@ define i32 @load_after_memset_0_offset(ptr %a) {
   ret i32 %v
 }
 
+define i32 @load_after_memset_1_offset(ptr %a) {
+; CHECK-LABEL: @load_after_memset_1_offset(
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 1, i64 16, i1 false)
+; CHECK-NEXT:    ret i32 16843009
+;
+  call void @llvm.memset.p0.i64(ptr %a, i8 1, i64 16, i1 false)
+  %gep = getelementptr i8, ptr %a, i64 4
+  %v = load i32, ptr %gep
+  ret i32 %v
+}
+
+define i1 @neg_load_after_memset_0_offset_i1(ptr %a) {
+; CHECK-LABEL: @neg_load_after_memset_0_offset_i1(
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 15, i64 16, i1 false)
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A]], i64 12
+; CHECK-NEXT:    [[V:%.*]] = load i1, ptr [[GEP]], align 1
+; CHECK-NEXT:    ret i1 [[V]]
+;
+  call void @llvm.memset.p0.i64(ptr %a, i8 15, i64 16, i1 false)
+  %gep = getelementptr i1, ptr %a, i64 12
+  %v = load i1, ptr %gep
+  ret i1 %v
+}
+
+define i8 @neg_load_after_memset_0_neg_offset(ptr %a) {
+; CHECK-LABEL: @neg_load_after_memset_0_neg_offset(
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr [[A:%.*]], i64 2
+; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[GEP]], i8 15, i64 16, i1 false)
+; CHECK-NEXT:    [[V:%.*]] = load i8, ptr [[A]], align 1
+; CHECK-NEXT:    ret i8 [[V]]
+;
+  %gep = getelementptr i8, ptr %a, i64 2
+  call void @llvm.memset.p0.i64(ptr %gep, i8 15, i64 16, i1 false)
+  %v = load i8, ptr %a
+  ret i8 %v
+}
+
 define i32 @load_after_memset_0_offset_too_large(ptr %a) {
 ; CHECK-LABEL: @load_after_memset_0_offset_too_large(
 ; CHECK-NEXT:    call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(16) [[A:%.*]], i8 0, i64 16, i1 false)

``````````

</details>


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