[CodeGen] Add NoMerge MIFlag to avoid MIR branch folding

[zequanwu]

duplicate existing diff

llvm/include/llvm/CodeGen/MachineInstr.h

 //===- llvm/CodeGen/MachineInstr.h - MachineInstr class ---------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the declaration of the MachineInstr class, which is the
enum MIFlag {
     NoUWrap      = 1 << 11,             // Instruction supports binary operator
                                         // no unsigned wrap.
     NoSWrap      = 1 << 12,             // Instruction supports binary operator
                                         // no signed wrap.
     IsExact      = 1 << 13,             // Instruction supports division is
                                         // known to be exact.
     NoFPExcept   = 1 << 14,             // Instruction does not raise
                                         // floatint-point exceptions.
+    NoMerge      = 1 << 15,             // Passes that drop source location info
+                                        // (e.g. branch folding) should skip
+                                        // this instruction.
   };
 
 private:
   const MCInstrDesc *MCID;              // Instruction descriptor.
   MachineBasicBlock *Parent = nullptr;  // Pointer to the owning basic block.
 
   // Operands are allocated by an ArrayRecycler.
   MachineOperand *Operands = nullptr;   // Pointer to the first operand.

llvm/include/llvm/CodeGen/SelectionDAG.h

class SelectionDAG {
   SDDbgInfo *DbgInfo;
 
   using CallSiteInfo = MachineFunction::CallSiteInfo;
   using CallSiteInfoImpl = MachineFunction::CallSiteInfoImpl;
 
   struct CallSiteDbgInfo {
     CallSiteInfo CSInfo;
     MDNode *HeapAllocSite = nullptr;
+    bool NoMerge = false;
   };
 
   DenseMap<const SDNode *, CallSiteDbgInfo> SDCallSiteDbgInfo;
 
   uint16_t NextPersistentId = 0;
 
 public:
   /// Clients of various APIs that cause global effects on
public:
   /// Return the HeapAllocSite type associated with the SDNode, if it exists.
   MDNode *getHeapAllocSite(const SDNode *Node) {
     auto It = SDCallSiteDbgInfo.find(Node);
     if (It == SDCallSiteDbgInfo.end())
       return nullptr;
     return It->second.HeapAllocSite;
   }
 
+  void addNoMergeSiteInfo(const SDNode *Node, bool NoMerge) {
+    if (NoMerge)
+      SDCallSiteDbgInfo[Node].NoMerge = NoMerge;
+  }
+
+  bool getNoMergeSiteInfo(const SDNode *Node) {
+    auto I = SDCallSiteDbgInfo.find(Node);
+    if (I == SDCallSiteDbgInfo.end())
+      return false;
+    return I->second.NoMerge;
+  }
+
   /// Return the current function's default denormal handling kind for the given
   /// floating point type.
   DenormalMode getDenormalMode(EVT VT) const {
     return MF->getDenormalMode(EVTToAPFloatSemantics(VT));
   }
 
   bool shouldOptForSize() const;
 

llvm/include/llvm/CodeGen/TargetLowering.h

 //===- llvm/CodeGen/TargetLowering.h - Target Lowering Info -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 ///
 /// \file
struct CallLoweringInfo {
     bool RetZExt           : 1;
     bool IsVarArg          : 1;
     bool IsInReg           : 1;
     bool DoesNotReturn     : 1;
     bool IsReturnValueUsed : 1;
     bool IsConvergent      : 1;
     bool IsPatchPoint      : 1;
     bool IsPreallocated : 1;
+    bool NoMerge           : 1;
 
     // IsTailCall should be modified by implementations of
     // TargetLowering::LowerCall that perform tail call conversions.
     bool IsTailCall = false;
 
     // Is Call lowering done post SelectionDAG type legalization.
     bool IsPostTypeLegalization = false;
 
     unsigned NumFixedArgs = -1;
     CallingConv::ID CallConv = CallingConv::C;
     SDValue Callee;
     ArgListTy Args;
     SelectionDAG &DAG;
     SDLoc DL;
     const CallBase *CB = nullptr;
     SmallVector<ISD::OutputArg, 32> Outs;
     SmallVector<SDValue, 32> OutVals;
     SmallVector<ISD::InputArg, 32> Ins;
     SmallVector<SDValue, 4> InVals;
 
     CallLoweringInfo(SelectionDAG &DAG)
         : RetSExt(false), RetZExt(false), IsVarArg(false), IsInReg(false),
           DoesNotReturn(false), IsReturnValueUsed(true), IsConvergent(false),
-          IsPatchPoint(false), IsPreallocated(false), DAG(DAG) {}
+          IsPatchPoint(false), IsPreallocated(false), NoMerge(false),
+          DAG(DAG) {}
 
     CallLoweringInfo &setDebugLoc(const SDLoc &dl) {
       DL = dl;
       return *this;
     }
 
     CallLoweringInfo &setChain(SDValue InChain) {
       Chain = InChain;
CallLoweringInfo &setCallee(Type *ResultType, FunctionType *FTy,
       IsInReg = Call.hasRetAttr(Attribute::InReg);
       DoesNotReturn =
           Call.doesNotReturn() ||
           (!isa<InvokeInst>(Call) && isa<UnreachableInst>(Call.getNextNode()));
       IsVarArg = FTy->isVarArg();
       IsReturnValueUsed = !Call.use_empty();
       RetSExt = Call.hasRetAttr(Attribute::SExt);
       RetZExt = Call.hasRetAttr(Attribute::ZExt);
+      NoMerge = Call.hasFnAttr(Attribute::NoMerge);
       
       Callee = Target;
 
       CallConv = Call.getCallingConv();
       NumFixedArgs = FTy->getNumParams();
       Args = std::move(ArgsList);
 
       CB = &Call;
 

llvm/lib/CodeGen/BranchFolding.cpp

if (!MBBI1->isIdenticalTo(*MBBI2) ||
         // FIXME: This check is dubious. It's used to get around a problem where
         // people incorrectly expect inline asm directives to remain in the same
         // relative order. This is untenable because normal compiler
         // optimizations (like this one) may reorder and/or merge these
         // directives.
         MBBI1->isInlineAsm()) {
       break;
     }
+    if (MBBI1->getFlag(MachineInstr::NoMerge) ||
+        MBBI2->getFlag(MachineInstr::NoMerge))
+      break;
     ++TailLen;
     I1 = MBBI1;
     I2 = MBBI2;
   }
 
   return TailLen;
 }
 

llvm/lib/CodeGen/MIRPrinter.cpp

void MIPrinter::print(const MachineInstr &MI) {
   if (MI.getFlag(MachineInstr::NoUWrap))
     OS << "nuw ";
   if (MI.getFlag(MachineInstr::NoSWrap))
     OS << "nsw ";
   if (MI.getFlag(MachineInstr::IsExact))
     OS << "exact ";
   if (MI.getFlag(MachineInstr::NoFPExcept))
     OS << "nofpexcept ";
+  if (MI.getFlag(MachineInstr::NoMerge))
+    OS << "nomerge ";
 
   OS << TII->getName(MI.getOpcode());
   if (I < E)
     OS << ' ';
 
   bool NeedComma = false;
   for (; I < E; ++I) {
     if (NeedComma)

llvm/lib/CodeGen/MachineInstr.cpp

void MachineInstr::print(raw_ostream &OS, ModuleSlotTracker &MST,
   if (getFlag(MachineInstr::NoUWrap))
     OS << "nuw ";
   if (getFlag(MachineInstr::NoSWrap))
     OS << "nsw ";
   if (getFlag(MachineInstr::IsExact))
     OS << "exact ";
   if (getFlag(MachineInstr::NoFPExcept))
     OS << "nofpexcept ";
+  if (getFlag(MachineInstr::NoMerge))
+    OS << "nomerge ";
 
   // Print the opcode name.
   if (TII)
     OS << TII->getName(getOpcode());
   else
     OS << "UNKNOWN";
 
   if (SkipOpers)

llvm/lib/CodeGen/SelectionDAG/ScheduleDAGSDNodes.cpp

if (Before == BB->end()) {
       // Return first instruction after the pre-existing instructions.
       MI = &*std::next(Before);
     }
 
     if (MI->isCandidateForCallSiteEntry() &&
         DAG->getTarget().Options.EmitCallSiteInfo)
       MF.addCallArgsForwardingRegs(MI, DAG->getSDCallSiteInfo(Node));
 
+    if (DAG->getNoMergeSiteInfo(Node)) {
+      MI->setFlag(MachineInstr::MIFlag::NoMerge);
+    }
+
     return MI;
   };
 
   // If this is the first BB, emit byval parameter dbg_value's.
   if (HasDbg && BB->getParent()->begin() == MachineFunction::iterator(BB)) {
     SDDbgInfo::DbgIterator PDI = DAG->ByvalParmDbgBegin();
     SDDbgInfo::DbgIterator PDE = DAG->ByvalParmDbgEnd();
     for (; PDI != PDE; ++PDI) {

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (IsTailCall) {
     MF.getFrameInfo().setHasTailCall();
     SDValue Ret = DAG.getNode(AArch64ISD::TC_RETURN, DL, NodeTys, Ops);
     DAG.addCallSiteInfo(Ret.getNode(), std::move(CSInfo));
     return Ret;
   }
 
   // Returns a chain and a flag for retval copy to use.
   Chain = DAG.getNode(AArch64ISD::CALL, DL, NodeTys, Ops);
+  DAG.addNoMergeSiteInfo(Chain.getNode(), CLI.NoMerge);
   InFlag = Chain.getValue(1);
   DAG.addCallSiteInfo(Chain.getNode(), std::move(CSInfo));
 
   uint64_t CalleePopBytes =
       DoesCalleeRestoreStack(CallConv, TailCallOpt) ? alignTo(NumBytes, 16) : 0;
 
   Chain = DAG.getCALLSEQ_END(Chain, DAG.getIntPtrConstant(NumBytes, DL, true),
                              DAG.getIntPtrConstant(CalleePopBytes, DL, true),

llvm/lib/Target/ARM/ARMISelLowering.cpp

if (isTailCall) {
     MF.getFrameInfo().setHasTailCall();
     SDValue Ret = DAG.getNode(ARMISD::TC_RETURN, dl, NodeTys, Ops);
     DAG.addCallSiteInfo(Ret.getNode(), std::move(CSInfo));
     return Ret;
   }
 
   // Returns a chain and a flag for retval copy to use.
   Chain = DAG.getNode(CallOpc, dl, NodeTys, Ops);
+  DAG.addNoMergeSiteInfo(Chain.getNode(), CLI.NoMerge);
   InFlag = Chain.getValue(1);
   DAG.addCallSiteInfo(Chain.getNode(), std::move(CSInfo));
 
   Chain = DAG.getCALLSEQ_END(Chain, DAG.getIntPtrConstant(NumBytes, dl, true),
                              DAG.getIntPtrConstant(0, dl, true), InFlag, dl);
   if (!Ins.empty())
     InFlag = Chain.getValue(1);
 

llvm/lib/Target/PowerPC/PPCISelLowering.h

public:
     /// between the functions for call lowering.
     struct CallFlags {
       const CallingConv::ID CallConv;
       const bool IsTailCall : 1;
       const bool IsVarArg : 1;
       const bool IsPatchPoint : 1;
       const bool IsIndirect : 1;
       const bool HasNest : 1;
+      const bool NoMerge : 1;
 
       CallFlags(CallingConv::ID CC, bool IsTailCall, bool IsVarArg,
-                bool IsPatchPoint, bool IsIndirect, bool HasNest)
+                bool IsPatchPoint, bool IsIndirect, bool HasNest, bool NoMerge)
           : CallConv(CC), IsTailCall(IsTailCall), IsVarArg(IsVarArg),
             IsPatchPoint(IsPatchPoint), IsIndirect(IsIndirect),
-            HasNest(HasNest) {}
+            HasNest(HasNest), NoMerge(NoMerge) {}
     };
 
   private:
     struct ReuseLoadInfo {
       SDValue Ptr;
       SDValue Chain;
       SDValue ResChain;
       MachinePointerInfo MPI;

llvm/lib/Target/PowerPC/PPCISelLowering.cpp

if (CFlags.IsTailCall) {
     assert(CallOpc == PPCISD::TC_RETURN &&
            "Unexpected call opcode for a tail call.");
     DAG.getMachineFunction().getFrameInfo().setHasTailCall();
     return DAG.getNode(CallOpc, dl, MVT::Other, Ops);
   }
 
   std::array<EVT, 2> ReturnTypes = {{MVT::Other, MVT::Glue}};
   Chain = DAG.getNode(CallOpc, dl, ReturnTypes, Ops);
+  DAG.addNoMergeSiteInfo(Chain.getNode(), CFlags.NoMerge);
   Glue = Chain.getValue(1);
 
   // When performing tail call optimization the callee pops its arguments off
   // the stack. Account for this here so these bytes can be pushed back on in
   // PPCFrameLowering::eliminateCallFramePseudoInstr.
   int BytesCalleePops = (CFlags.CallConv == CallingConv::Fast &&
                          getTargetMachine().Options.GuaranteedTailCallOpt)
                             ? NumBytes
if (Subtarget.useLongCalls() && isa<GlobalAddressSDNode>(Callee) &&
       !isTailCall)
     Callee = LowerGlobalAddress(Callee, DAG);
 
   CallFlags CFlags(
       CallConv, isTailCall, isVarArg, isPatchPoint,
       isIndirectCall(Callee, DAG, Subtarget, isPatchPoint),
       // hasNest
       Subtarget.is64BitELFABI() &&
-          any_of(Outs, [](ISD::OutputArg Arg) { return Arg.Flags.isNest(); }));
+          any_of(Outs, [](ISD::OutputArg Arg) { return Arg.Flags.isNest(); }),
+      CLI.NoMerge);
 
   if (Subtarget.isSVR4ABI() && Subtarget.isPPC64())
     return LowerCall_64SVR4(Chain, Callee, CFlags, Outs, OutVals, Ins, dl, DAG,
                             InVals, CB);
 
   if (Subtarget.isSVR4ABI())
     return LowerCall_32SVR4(Chain, Callee, CFlags, Outs, OutVals, Ins, dl, DAG,
                             InVals, CB);

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

SDValue RISCVTargetLowering::LowerCall(CallLoweringInfo &CLI,
   SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
 
   if (IsTailCall) {
     MF.getFrameInfo().setHasTailCall();
     return DAG.getNode(RISCVISD::TAIL, DL, NodeTys, Ops);
   }
 
   Chain = DAG.getNode(RISCVISD::CALL, DL, NodeTys, Ops);
+  DAG.addNoMergeSiteInfo(Chain.getNode(), CLI.NoMerge);
   Glue = Chain.getValue(1);
 
   // Mark the end of the call, which is glued to the call itself.
   Chain = DAG.getCALLSEQ_END(Chain,
                              DAG.getConstant(NumBytes, DL, PtrVT, true),
                              DAG.getConstant(0, DL, PtrVT, true),
                              Glue, DL);
   Glue = Chain.getValue(1);

llvm/lib/Target/X86/X86ISelLowering.cpp

X86TargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
   }
 
   if (HasNoCfCheck && IsCFProtectionSupported) {
     Chain = DAG.getNode(X86ISD::NT_CALL, dl, NodeTys, Ops);
   } else {
     Chain = DAG.getNode(X86ISD::CALL, dl, NodeTys, Ops);
   }
   InFlag = Chain.getValue(1);
+  DAG.addNoMergeSiteInfo(Chain.getNode(), CLI.NoMerge);
   DAG.addCallSiteInfo(Chain.getNode(), std::move(CSInfo));
 
   // Save heapallocsite metadata.
   if (CLI.CB)
     if (MDNode *HeapAlloc = CLI.CB->getMetadata("heapallocsite"))
       DAG.addHeapAllocSite(Chain.getNode(), HeapAlloc);
 
   // Create the CALLSEQ_END node.

llvm/test/CodeGen/AArch64/nomerge.ll

+; RUN: llc < %s -mtriple=aarch64 -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+entry:
+  switch i32 %i, label %if.end3 [
+    i32 5, label %if.then
+    i32 7, label %if.then2
+  ]
+
+if.then:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.then2:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.end3:
+  tail call void @bar() #0
+  ret void
+}
+
+declare void @bar()
+
+attributes #0 = { nomerge }
+
+; CHECK-LABEL: foo:
+; CHECK: // %bb.0: // %entry
+; CHECK: // %bb.1: // %entry
+; CHECK: // %bb.2: // %if.then
+; CHECK-NEXT: bl bar
+; CHECK: b bar
+; CHECK: .LBB0_3: // %if.then2
+; CHECK-NEXT: bl bar
+; CHECK: .LBB0_4: // %if.end3
+; CHECK: b bar

llvm/test/CodeGen/ARM/nomerge.ll

+; RUN: llc < %s -mtriple=arm -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+entry:
+  switch i32 %i, label %if.end3 [
+    i32 5, label %if.then
+    i32 7, label %if.then2
+  ]
+
+if.then:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.then2:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.end3:
+  tail call void @bar() #0
+  ret void
+}
+
+declare void @bar()
+
+attributes #0 = { nomerge }
+
+; CHECK-LABEL: foo:
+; CHECK: @ %bb.0: @ %entry
+; CHECK: @ %bb.1: @ %entry
+; CHECK: @ %bb.2: @ %if.then
+; CHECK-NEXT: bl bar
+; CHECK: b bar
+; CHECK: .LBB0_3: @ %if.then2
+; CHECK-NEXT: bl bar
+; CHECK: .LBB0_4: @ %if.end3
+; CHECK: b bar

llvm/test/CodeGen/PowerPC/nomerge.ll

+; RUN: llc < %s -mtriple=powerpc -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+entry:
+  switch i32 %i, label %if.end3 [
+    i32 5, label %if.then
+    i32 7, label %if.then2
+  ]
+
+if.then:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.then2:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.end3:
+  tail call void @bar() #0
+  ret void
+}
+
+declare void @bar()
+
+attributes #0 = { nomerge }
+
+; CHECK-LABEL: foo:
+; CHECK: # %bb.0: # %entry
+; CHECK: # %bb.1: # %entry
+; CHECK: # %bb.2: # %if.then
+; CHECK-NEXT: bl bar
+; CHECK: .LBB0_3: # %if.then2
+; CHECK-NEXT: bl bar
+; CHECK: .LBB0_4: # %if.end3
+; CHECK-NEXT: bl bar

llvm/test/CodeGen/RISCV/nomerge.ll

+; RUN: llc < %s -mtriple=riscv64 -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+entry:
+  switch i32 %i, label %if.end3 [
+    i32 5, label %if.then
+    i32 7, label %if.then2
+  ]
+
+if.then:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.then2:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.end3:
+  tail call void @bar() #0
+  ret void
+}
+
+declare void @bar()
+
+attributes #0 = { nomerge }
+
+; CHECK-LABEL: foo:
+; CHECK: # %bb.0: # %entry
+; CHECK: # %bb.1: # %entry
+; CHECK: # %bb.2: # %if.then
+; CHECK-NEXT: call bar
+; CHECK: .LBB0_3: # %if.then2
+; CHECK-NEXT: call bar
+; CHECK: .LBB0_4: # %if.end3
+; CHECK: tail bar

llvm/test/CodeGen/X86/nomerge.ll

+; RUN: llc < %s -mtriple=x86_64 -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+entry:
+  switch i32 %i, label %if.end3 [
+    i32 5, label %if.then
+    i32 7, label %if.then2
+  ]
+
+if.then:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.then2:
+  tail call void @bar() #0
+  br label %if.end3
+
+if.end3:
+  tail call void @bar() #0
+  ret void
+}
+
+declare void @bar()
+
+attributes #0 = { nomerge }
+
+; CHECK-LABEL: foo:
+; CHECK: # %bb.0: # %entry
+; CHECK: # %bb.1: # %entry
+; CHECK: # %bb.2: # %if.then
+; CHECK-NEXT: callq bar
+; CHECK: jmp bar # TAILCALL
+; CHECK: .LBB0_3: # %if.then2
+; CHECK: callq bar
+; CHECK: .LBB0_4: # %if.end3
+; CHECK: jmp bar # TAILCALL