Extend SimpleGVN to support load-load forwarding (#2625)

* Initial plan

* Implement load-load forwarding in SimpleGVN

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Apply clang-format to SimpleGVN.cpp

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Remove unused next_argument label

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Update SimpleGVN README to document load-load forwarding

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Combine dominatesAndCovers functions to avoid redundancy

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Refactor collectMemoryOps per review feedback

- Use pointer parameter instead of bool for Calls to indicate if nocapture calls should be collected
- Iterate over uses instead of users to get argument index directly
- Use isNoCapture utility function from Utils.h
- Avoid recollecting in second phase by tracking eliminated loads

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Change Calls parameter from pointer to reference with sentinel pattern

- Changed Calls from pointer to reference parameter in collectMemoryOps
- Use sentinel value (nullptr CallInst) to distinguish between modes
- Empty Calls on entry = reject nocapture calls (store-load forwarding)
- Non-empty Calls on entry = collect nocapture calls (load-load forwarding)
- Eliminates need for separate boolean parameter

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Remove unnecessary sentinel pattern for nocapture call collection

Instead of using sentinel values and calling collectMemoryOps again,
directly traverse uses to collect nocapture calls in the second phase.
This simplifies the code and avoids the complexity of the sentinel pattern.

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* Refactor aliasing check into common helper and improve tests

- Extract hasAliasingWriteBetween() helper to check for stores/calls between instructions
- Reuse this helper in both store-load and load-load forwarding
- Update load_load_no_nocapture_call.ll: place both loads after call to properly test no forwarding
- Update load_load_with_store_between.ll: demonstrate both load-load forwarding before and after store

Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>

* fix

* fix

* fix

* fix

* fix

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: wsmoses <1260124+wsmoses@users.noreply.github.com>
Co-authored-by: William S. Moses <gh@wsmoses.com>

Author: Copilot

enzyme/Enzyme/SimpleGVN.cpp

@@ -0,0 +1 @@
+./enzyme

enzyme/enzyme

@@ -8,24 +8,41 @@ SimpleGVN is a GVN-like (Global Value Numbering) optimization pass that forwards
 
 ## How It Works
 
-The pass:
-1. Identifies function arguments with both `noalias` and `nocapture` attributes
-2. Verifies all uses are exclusively loads, stores, or GEP instructions
+The pass operates in two phases:
+
+### Phase 1: Store-to-Load Forwarding
+1. Identifies function arguments with both `noalias` and `nocapture` attributes, and allocas
+2. Verifies all uses are exclusively loads, stores, GEP instructions, or casts
 3. For each load, finds dominating stores that cover the load's memory range
 4. Replaces the load with the stored value if no aliasing store exists in between
 
+### Phase 2: Load-to-Load Forwarding
+1. Re-collects loads and stores, this time allowing nocapture function calls
+2. For each load, finds dominating loads that cover the same memory range
+3. Replaces the load with the value from the dominating load if:
+   - No aliasing store exists between the two loads
+   - No nocapture function call exists between the two loads (as they may modify memory)
+
 ## Test Cases
 
+### Store-to-Load Forwarding Tests
 - **basic.ll** - Simple store-to-load forwarding
 - **offset.ll** - Forwarding with GEP offsets
 - **dominance.ll** - Verifies dominance requirements
 - **intermediate_store.ll** - Handles intermediate stores correctly
 - **no_noalias.ll** - Rejects optimization when noalias is missing
-- **call_use.ll** - Rejects when argument has non-memory uses
+- **call_use.ll** - Rejects when argument has non-memory uses (non-nocapture calls)
 - **struct_field.ll** - Handles struct field accesses
 - **type_conversion.ll** - Tests byte-level extraction
 - **comprehensive.ll** - Multiple loads/stores at different offsets
 
+### Load-to-Load Forwarding Tests
+- **load_load_basic.ll** - Simple load-to-load forwarding
+- **load_load_offset.ll** - Load-to-load forwarding with GEP offsets
+- **load_load_nocapture_call.ll** - No forwarding when nocapture call exists between loads
+- **load_load_no_nocapture_call.ll** - Optimization disabled when call lacks nocapture attribute
+- **load_load_with_store_between.ll** - No load-to-load forwarding when store exists between loads
+
 ## Running the Tests
 
 Using opt with the new pass manager:
@@ -38,21 +55,44 @@ Using opt with the legacy pass manager (LLVM < 16):
 opt -load LLVMEnzyme-18.so -simple-gvn -S < test.ll
 ```
 
-## Example
+## Examples
+
+### Store-to-Load Forwarding Example
 
 Input:
 ```llvm
-define i32 @foo(i32* noalias nocapture %ptr) {
-  store i32 42, i32* %ptr
-  %v = load i32, i32* %ptr
+define i32 @foo(ptr noalias nocapture %ptr) {
+  store i32 42, ptr %ptr
+  %v = load i32, ptr %ptr
   ret i32 %v
 }
 ```
 
 Output after SimpleGVN:
 ```llvm
-define i32 @foo(i32* noalias nocapture %ptr) {
-  store i32 42, i32* %ptr
+define i32 @foo(ptr noalias nocapture %ptr) {
+  store i32 42, ptr %ptr
   ret i32 42
 }
 ```
+
+### Load-to-Load Forwarding Example
+
+Input:
+```llvm
+define i32 @bar(ptr noalias nocapture %ptr) {
+  %v1 = load i32, ptr %ptr
+  %v2 = load i32, ptr %ptr
+  %sum = add i32 %v1, %v2
+  ret i32 %sum
+}
+```
+
+Output after SimpleGVN:
+```llvm
+define i32 @bar(ptr noalias nocapture %ptr) {
+  %v1 = load i32, ptr %ptr
+  %sum = add i32 %v1, %v1
+  ret i32 %sum
+}
+```

enzyme/test/Enzyme/SimpleGVN/README.md

@@ -0,0 +1,31 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -passes="simple-gvn" -S | FileCheck %s
+
+declare  void @julia___conv_filter__271_37475({ { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* noalias nocapture nofree noundef nonnull writeonly sret({ { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }) align 8 dereferenceable(72) %0);
+
+; Function Attrs: noinline
+define private {} addrspace(10)* @julia__conv_filter__37469([1 x {} addrspace(10)*]* %return_roots) {
+top:
+  %0 = alloca { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }, align 8
+  call fastcc void @julia___conv_filter__271_37475({ { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* noalias nocapture nofree noundef nonnull writeonly sret({ { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }) align 8 dereferenceable(72) %0)
+  %a5 = getelementptr inbounds { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }, { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* %0, i64 0, i32 0, i32 0
+  %a6 = load {} addrspace(10)*, {} addrspace(10)** %a5, align 8
+  %a7 = getelementptr inbounds [1 x {} addrspace(10)*], [1 x {} addrspace(10)*]* %return_roots, i64 0, i64 0
+  store {} addrspace(10)* %a6, {} addrspace(10)** %a7, align 8
+  %srcloccs2 = getelementptr inbounds { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }, { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* %0, i64 0, i32 0, i32 0
+  %a8 = load {} addrspace(10)*, {} addrspace(10)** %srcloccs2, align 8
+  ret {} addrspace(10)* %a8
+}
+
+; CHECK: define private {} addrspace(10)* @julia__conv_filter__37469([1 x {} addrspace(10)*]* %return_roots)
+; CHECK-NEXT: top:
+; CHECK-NEXT:   %0 = alloca { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }, align 8
+; CHECK-NEXT:   call fastcc void @julia___conv_filter__271_37475({ { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* noalias nocapture nofree noundef nonnull writeonly sret({ { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }) align 8 dereferenceable(72) %0)
+; CHECK-NEXT:   %a5 = getelementptr inbounds { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }, { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* %0, i64 0, i32 0, i32 0
+; CHECK-NEXT:   %a6 = load {} addrspace(10)*, {} addrspace(10)** %a5, align 8
+; CHECK-NEXT:   %a7 = getelementptr inbounds [1 x {} addrspace(10)*], [1 x {} addrspace(10)*]* %return_roots, i64 0, i64 0
+; CHECK-NEXT:   store {} addrspace(10)* %a6, {} addrspace(10)** %a7, align 8
+; CHECK-NEXT:   %srcloccs2 = getelementptr inbounds { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }, { { {} addrspace(10)*, [1 x [2 x i64]], i64, i64 }, [4 x i64] }* %0, i64 0, i32 0, i32 0
+; CHECK-NEXT:   ret {} addrspace(10)* %a6
+; CHECK-NEXT: }
+

enzyme/test/Enzyme/SimpleGVN/jlloadload.ll

@@ -0,0 +1,18 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -opaque-pointers -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -opaque-pointers -passes="simple-gvn" -S | FileCheck %s
+
+; Test basic load-to-load forwarding with noalias nocapture argument
+
+define i32 @test_load_load_basic(ptr noalias nocapture %ptr) {
+entry:
+  %val1 = load i32, ptr %ptr, align 4
+  %val2 = load i32, ptr %ptr, align 4
+  %sum = add i32 %val1, %val2
+  ret i32 %sum
+}
+
+; CHECK: define i32 @test_load_load_basic(ptr noalias nocapture %ptr)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %val1 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   %sum = add i32 %val1, %val1
+; CHECK-NEXT:   ret i32 %sum

enzyme/test/Enzyme/SimpleGVN/load_load_basic.ll

@@ -0,0 +1,24 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -opaque-pointers -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -opaque-pointers -passes="simple-gvn" -S | FileCheck %s
+
+; Test that optimization is NOT applied when call does not have nocapture attribute
+; and that even with nocapture, we don't forward if there's an intermediate call
+
+declare void @external_func(ptr)
+
+define i32 @test_load_load_no_nocapture_call(ptr noalias nocapture %ptr) {
+entry:
+  call void @external_func(ptr %ptr)
+  %val1 = load i32, ptr %ptr, align 4
+  %val2 = load i32, ptr %ptr, align 4
+  %sum = add i32 %val1, %val2
+  ret i32 %sum
+}
+
+; CHECK: define i32 @test_load_load_no_nocapture_call(ptr noalias nocapture %ptr)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   call void @external_func(ptr %ptr)
+; CHECK-NEXT:   %val1 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   %val2 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   %sum = add i32 %val1, %val2
+; CHECK-NEXT:   ret i32 %sum

enzyme/test/Enzyme/SimpleGVN/load_load_no_nocapture_call.ll

@@ -0,0 +1,24 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -opaque-pointers -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -opaque-pointers -passes="simple-gvn" -S | FileCheck %s
+
+; Test load-to-load forwarding with nocapture function call between loads
+; The nocapture call should prevent forwarding
+
+declare void @nocapture_func(ptr nocapture)
+
+define i32 @test_load_load_nocapture_call(ptr noalias nocapture %ptr) {
+entry:
+  %val1 = load i32, ptr %ptr, align 4
+  call void @nocapture_func(ptr nocapture %ptr)
+  %val2 = load i32, ptr %ptr, align 4
+  %sum = add i32 %val1, %val2
+  ret i32 %sum
+}
+
+; CHECK: define i32 @test_load_load_nocapture_call(ptr noalias nocapture %ptr)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %val1 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   call void @nocapture_func(ptr nocapture %ptr)
+; CHECK-NEXT:   %val2 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   %sum = add i32 %val1, %val2
+; CHECK-NEXT:   ret i32 %sum

enzyme/test/Enzyme/SimpleGVN/load_load_nocapture_call.ll

@@ -0,0 +1,20 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -opaque-pointers -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -opaque-pointers -passes="simple-gvn" -S | FileCheck %s
+
+; Test load-to-load forwarding with GEP offsets
+
+define i32 @test_load_load_offset(ptr noalias nocapture %ptr) {
+entry:
+  %gep = getelementptr i32, ptr %ptr, i64 1
+  %val1 = load i32, ptr %gep, align 4
+  %val2 = load i32, ptr %gep, align 4
+  %sum = add i32 %val1, %val2
+  ret i32 %sum
+}
+
+; CHECK: define i32 @test_load_load_offset(ptr noalias nocapture %ptr)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %gep = getelementptr i32, ptr %ptr, i64 1
+; CHECK-NEXT:   %val1 = load i32, ptr %gep, align 4
+; CHECK-NEXT:   %sum = add i32 %val1, %val1
+; CHECK-NEXT:   ret i32 %sum

enzyme/test/Enzyme/SimpleGVN/load_load_offset.ll

@@ -0,0 +1,29 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -opaque-pointers -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -opaque-pointers -passes="simple-gvn" -S | FileCheck %s
+
+; Test that load-to-load forwarding does not happen when there's a store between loads
+; But does happen after the store when there's no intervening write
+
+define i32 @test_load_load_with_store_between(ptr noalias nocapture %ptr) {
+entry:
+  %val1 = load i32, ptr %ptr, align 4
+  %val2 = load i32, ptr %ptr, align 4
+  store i32 %val1, ptr %ptr, align 4
+  %val3 = load i32, ptr %ptr, align 4
+  %val4 = load i32, ptr %ptr, align 4
+  %sum1 = add i32 %val1, %val2
+  %sum2 = add i32 %sum1, %val3
+  %sum3 = add i32 %sum2, %val4
+  ret i32 %sum3
+}
+
+; val1 and val2 should be forwarded (load-load before store)
+; val3 and val4 should both be forwarded from the store (store-load forwarding)
+; CHECK: define i32 @test_load_load_with_store_between(ptr noalias nocapture %ptr)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %val1 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   store i32 %val1, ptr %ptr, align 4
+; CHECK-NEXT:   %sum1 = add i32 %val1, %val1
+; CHECK-NEXT:   %sum2 = add i32 %sum1, %val1
+; CHECK-NEXT:   %sum3 = add i32 %sum2, %val1
+; CHECK-NEXT:   ret i32 %sum3

enzyme/test/Enzyme/SimpleGVN/load_load_with_store_between.ll

@@ -0,0 +1,29 @@
+; RUN: if [ %llvmver -lt 16 ]; then %opt < %s %loadEnzyme -opaque-pointers -simple-gvn -S | FileCheck %s; fi
+; RUN: %opt < %s %newLoadEnzyme -opaque-pointers -passes="simple-gvn" -S | FileCheck %s
+
+; Test load-to-load forwarding with nocapture function call between loads
+; The nocapture call should prevent forwarding
+
+declare void @nocapture_func(ptr nocapture)
+
+define i32 @test_load_load_nocapture_call(ptr noalias nocapture %ptr) {
+entry:
+  store i32 0, ptr %ptr, align 4
+  br label %next
+
+next:
+  call void @nocapture_func(ptr nocapture %ptr)
+  %val2 = load i32, ptr %ptr, align 4
+  ret i32 %val2
+}
+
+; CHECK: define i32 @test_load_load_nocapture_call(ptr noalias nocapture %ptr)
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   store i32 0, ptr %ptr, align 4
+; CHECK-NEXT:   br label %next
+
+; CHECK: next:                                             ; preds = %entry
+; CHECK-NEXT:   call void @nocapture_func(ptr nocapture %ptr)
+; CHECK-NEXT:   %val2 = load i32, ptr %ptr, align 4
+; CHECK-NEXT:   ret i32 %val2
+; CHECK-NEXT: }