InstCombine: sink loads with invariant.load metadata

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.

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,114 @@ abort:
   call void @abort()
   unreachable
 }
+
+; Loads marked invariant can be sunk past potential memory writes.
+
+define i32 @invariant_load_metadata(ptr %p, i1 %cond) {
+; CHECK-LABEL: @invariant_load_metadata(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BLOCK:%.*]], label [[END:%.*]]
+; CHECK:       block:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[P:%.*]], align 4, !invariant.load [[META0:![0-9]+]]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %v = load i32, ptr %p, !invariant.load !0
+  br i1 %cond, label %block, label %end
+block:
+  call void @fn()
+  br label %end
+end:
+  ret i32 %v
+}
+
+; Loads not marked invariant cannot be sunk past potential memory writes.
+
+define i32 @invariant_load_neg(ptr %p, i1 %cond) {
+; CHECK-LABEL: @invariant_load_neg(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[V:%.*]] = load i32, ptr [[P:%.*]], align 4
+; CHECK-NEXT:    br i1 [[COND:%.*]], label [[BLOCK:%.*]], label [[END:%.*]]
+; CHECK:       block:
+; CHECK-NEXT:    call void @fn()
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret i32 [[V]]
+;
+entry:
+  %v = load i32, ptr %p
+  br i1 %cond, label %block, label %end
+block:
+  call void @fn()
+  br label %end
+end:
+  ret i32 %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 = !{}