Preserve important metadata in JumpThreadingPass::unfoldSelectInstr

[mmendell]

Summary

Useful metadata like llvm.loop can be lost when the new branch instruction is created. Grab that metadata from the old branch.
Add a correct LIT test to check

[nikic]

LGTM, but I do wonder whether we aren't working around a different issue here. JumpThreading generally does not perform threading across loop headers (it's behind a default-disabled option). As far as I understand, the select unfolding is just done to enable future threading. In this case, we unfold the select, but then can't thread it because that would require threading across a header.

We're left with what appear to me an unprofitable transform: We now have two loop latches, which is non-canonical. LoopSimplify will have to convert this back into a single latch.

So my thinking is that we should probably prevent select unfolding from happening here in the first place. I see that tryToUnfoldSelectInCurrBB() already guards against this, but tryToUnfoldSelec() does not.

llvm/lib/Transforms/Scalar/JumpThreading.cpp

BasicBlock *NewBB = BasicBlock::Create(BB->getContext(), "select.unfold",
                                          BB->getParent(), BB);
   // Move the unconditional branch to NewBB.
   PredTerm->removeFromParent();
   PredTerm->insertInto(NewBB, NewBB->end());
   // Create a conditional branch and update PHI nodes.
   auto *BI = BranchInst::Create(NewBB, BB, SI->getCondition(), Pred);
   BI->applyMergedLocation(PredTerm->getDebugLoc(), SI->getDebugLoc());
   BI->copyMetadata(*SI, {LLVMContext::MD_prof});
+  // Also preserve the metadata from the branch
+  BI->copyMetadata(*PredTerm,
+                   {LLVMContext::MD_loop, LLVMContext::MD_annotation});
   SIUse->setIncomingValue(Idx, SI->getFalseValue());
   SIUse->addIncoming(SI->getTrueValue(), NewBB);
 
   uint64_t TrueWeight = 1;
   uint64_t FalseWeight = 1;
   // Copy probabilities from 'SI' to created conditional branch in 'Pred'.
   if (extractBranchWeights(*SI, TrueWeight, FalseWeight) &&
       (TrueWeight + FalseWeight) != 0) {

llvm/test/Transforms/JumpThreading/preserve-select-unfold.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals
+; RUN: opt -passes=jump-threading -S < %s | FileCheck %s
+; Ensure that we preserve llvm.loop information
+define spir_kernel void @test(i1 %Input) {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[A:%.*]]
+; CHECK:       select.unfold:
+; CHECK-NEXT:    br label [[A]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK:       A:
+; CHECK-NEXT:    [[CURRENT_STATE_0_I_I:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ 0, [[M:%.*]] ], [ 2, [[SELECT_UNFOLD:%.*]] ]
+; CHECK-NEXT:    [[B:%.*]] = icmp eq i32 [[CURRENT_STATE_0_I_I]], 0
+; CHECK-NEXT:    br i1 [[B]], label [[M]], label [[B2:%.*]]
+; CHECK:       B2:
+; CHECK-NEXT:    ret void
+; CHECK:       M:
+; CHECK-NEXT:    br i1 [[INPUT:%.*]], label [[SELECT_UNFOLD]], label [[A]], !llvm.loop [[LOOP0]]
+;
+entry:
+  br label %A
+
+A:
+  %current_state.0.i.i = phi i32 [ 0, %entry ], [ %X, %M ]
+  %B = icmp eq i32 %current_state.0.i.i, 0
+  br i1 %B, label %M, label %B2
+
+B2:
+  ret void
+
+M:
+  %X = select i1 %Input, i32 2, i32 0
+  br label %A, !llvm.loop !0
+}
+
+!0 = distinct !{!0, !1}
+!1 = !{!"llvm.loop.unroll.disable"}
+;.
+; CHECK: [[LOOP0]] = distinct !{!0, !1}
+; CHECK: [[META1:![0-9]+]] = !{!"llvm.loop.unroll.disable"}
+;.