[llvm] InstCombine: sink loads with invariant.load metadata (PR #112692)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Ramkumar Ramachandra (artagnon)

<details>
<summary>Changes</summary>



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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+2-1) 
- (modified) llvm/test/Transforms/InstCombine/sink_instruction.ll (+119-4) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 954c4cf19c2077..c8b9f166b16020 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -4822,7 +4822,8 @@ bool InstCombinerImpl::tryToSinkInstruction(Instruction *I,
 
   // We can only sink load instructions if there is nothing between the load and
   // the end of block that could change the value.
-  if (I->mayReadFromMemory()) {
+  if (I->mayReadFromMemory() &&
+      !I->hasMetadata(LLVMContext::MD_invariant_load)) {
     // We don't want to do any sophisticated alias analysis, so we only check
     // the instructions after I in I's parent block if we try to sink to its
     // successor block.
diff --git a/llvm/test/Transforms/InstCombine/sink_instruction.ll b/llvm/test/Transforms/InstCombine/sink_instruction.ll
index c938002788bc28..9341f29c2bdd98 100644
--- a/llvm/test/Transforms/InstCombine/sink_instruction.ll
+++ b/llvm/test/Transforms/InstCombine/sink_instruction.ll
@@ -86,8 +86,8 @@ define i32 @test3(ptr nocapture readonly %P, i32 %i) {
 ; CHECK-LABEL: @test3(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    switch i32 [[I:%.*]], label [[SW_EPILOG:%.*]] [
-; CHECK-NEXT:    i32 5, label [[SW_BB:%.*]]
-; CHECK-NEXT:    i32 2, label [[SW_BB]]
+; CHECK-NEXT:      i32 5, label [[SW_BB:%.*]]
+; CHECK-NEXT:      i32 2, label [[SW_BB]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    [[IDXPROM:%.*]] = sext i32 [[I]] to i64
@@ -190,8 +190,8 @@ define i32 @test6(ptr nocapture readonly %P, i32 %i, i1 %cond) {
 ; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr [[P:%.*]], i64 [[IDXPROM]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[ARRAYIDX]], align 4
 ; CHECK-NEXT:    switch i32 [[I]], label [[SW_BB:%.*]] [
-; CHECK-NEXT:    i32 5, label [[SW_EPILOG:%.*]]
-; CHECK-NEXT:    i32 2, label [[SW_EPILOG]]
+; CHECK-NEXT:      i32 5, label [[SW_EPILOG:%.*]]
+; CHECK-NEXT:      i32 2, label [[SW_EPILOG]]
 ; CHECK-NEXT:    ]
 ; CHECK:       sw.bb:
 ; CHECK-NEXT:    br label [[SW_EPILOG]]
@@ -272,3 +272,118 @@ abort:
   call void @abort()
   unreachable
 }
+
+; Loads marked invariant can be sunk across critical edges.
+
+define <4 x float> @invariant_load_metadata(ptr %p, i32 %cond) {
+; CHECK-LABEL: @invariant_load_metadata(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[COND:%.*]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[BLOCK:%.*]], label [[END:%.*]]
+; CHECK:       block:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    [[V:%.*]] = load <4 x float>, ptr [[P:%.*]], align 16, !invariant.load [[META0:![0-9]+]]
+; CHECK-NEXT:    ret <4 x float> [[V]]
+;
+entry:
+  %v = load <4 x float>, ptr %p, !invariant.load !0
+  %c = icmp eq i32 %cond, 0
+  br i1 %c, label %block, label %end
+block:
+  call void @fn()
+  br label %end
+end:
+  ret <4 x float> %v
+}
+
+; Loads not marked invariant cannot be sunk across critical edges.
+
+define <4 x float> @invariant_load_neg(ptr %p, i32 %cond) {
+; CHECK-LABEL: @invariant_load_neg(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load <4 x float>, ptr [[P:%.*]], align 16
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[COND:%.*]], 0
+; CHECK-NEXT:    br i1 [[C]], label [[BLOCK:%.*]], label [[END:%.*]]
+; CHECK:       block:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret <4 x float> [[V]]
+;
+entry:
+  %v = load <4 x float>, ptr %p
+  %c = icmp eq i32 %cond, 0
+  br i1 %c, label %block, label %end
+block:
+  call void @fn()
+  br label %end
+end:
+  ret <4 x float> %v
+}
+
+; Loads that aren't marked invariant but used in one branch
+; can be sunk to that branch.
+
+define void @invariant_load_use_in_br(ptr %p, i1 %cond) {
+; CHECK-LABEL: @invariant_load_use_in_br(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[TRUE_BR:%.*]], label [[FALSE_BR:%.*]]
+; CHECK:       true.br:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       false.br:
+; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[P:%.*]], align 4
+; CHECK-NEXT:    call void @fn(i32 [[VAL]])
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %val = load i32, ptr %p
+  br i1 %cond, label %true.br, label %false.br
+true.br:
+  call void @fn()
+  br label %exit
+false.br:
+  call void @fn(i32 %val)
+  br label %exit
+exit:
+  ret void
+}
+
+; Invariant loads marked with metadata can be sunk past calls.
+
+define void @invariant_load_metadata_call(ptr %p, i1 %cond) {
+; CHECK-LABEL: @invariant_load_metadata_call(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[TRUE_BR:%.*]], label [[FALSE_BR:%.*]]
+; CHECK:       true.br:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       false.br:
+; CHECK-NEXT:    [[VAL:%.*]] = load i32, ptr [[P:%.*]], align 4, !invariant.load [[META0]]
+; CHECK-NEXT:    call void @fn(i32 [[VAL]])
+; CHECK-NEXT:    br label [[EXIT]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %val = load i32, ptr %p, !invariant.load !0
+  call void @fn()
+  br i1 %cond, label %true.br, label %false.br
+true.br:
+  call void @fn()
+  br label %exit
+false.br:
+  call void @fn(i32 %val)
+  br label %exit
+exit:
+  ret void
+}
+
+declare void @fn()
+
+!0 = !{}

``````````

</details>


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