Add NoMerge MIFlag to avoid MIR branch folding

[zequanwu]

Summary

Related to: D78659

[zequanwu]

add testcase

[zequanwu]

Handle nomerge flag in LowerCall at targets AArch64, ARM, PowerPC, RISCV, and add test cases for each target.

[rnk]

Overall, this looks good to me, but I'd like to ask @craig.topper for an opinion. His area of expertise is the x86 backend.

llvm/include/llvm/CodeGen/SelectionDAG.h

281

Please zero initialize this, otherwise the default constructor will not initialize it, and if entries in this hash table are added for other reasons, it will be uninitialized. You could observe this by combining the nomerge attribute with __declspec(allocator), but IMO a test is not necessary.

1916

If NoMerge is false, you can skip this hash table insertion. Otherwise, we will do this hash insertion for each and ever call instruction, but we only need to do it for this feature.

[zequanwu]

Zero Initialize NoMerge in CallSiteDbgInfo.
Put Check in addNoMergeSiteInfo to avoid unnecessary hash table insertion.

llvm/include/llvm/CodeGen/SelectionDAG.h

1916

Previously, I put the check before each call site of this function.
I agree that put the check inside the function is better.

[rnk]

After this lands, can you follow this up by migrating the sanitizer passes to using the attribute instead of using empty inline asm blobs? Send the patch to @vitalybuka. This is the code I'm thinking of:
https://github.com/llvm/llvm-project/blob/028bfdd8913616f7a3e57e8ef5c2a9990e528ff0/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp#L1626
https://github.com/llvm/llvm-project/blob/d6695e18763a05b30cb336c18157175277da8f4b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp#L864
... and any similar uses.

Run check-asan, check-msan, etc to run the sanitizer integration tests, and confirm that they pass. You can run check-all to run everything and be sure as well.

[arsenm]

llvm/include/llvm/CodeGen/MachineInstr.h

109

It's not obvious to me what this means. These are also supposed to be droppable optimization flags, but it seems like you're using this for something semantically required?

[rnk]

llvm/include/llvm/CodeGen/MachineInstr.h

109

I don't think BundledPred or FrameSetup are droppable, IMO they are semantically significant.

The only alternative I see to this would be to encode this information in the MCInstrDesc, which would require mirroring the entire family of X86 CALL instructions, of which there are many.

[zequanwu]

In D79537#2035041, @rnk wrote:

After this lands, can you follow this up by migrating the sanitizer passes to using the attribute instead of using empty inline asm blobs? Send the patch to @vitalybuka. This is the code I'm thinking of:
https://github.com/llvm/llvm-project/blob/028bfdd8913616f7a3e57e8ef5c2a9990e528ff0/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp#L1626
https://github.com/llvm/llvm-project/blob/d6695e18763a05b30cb336c18157175277da8f4b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp#L864
... and any similar uses.

Run check-asan, check-msan, etc to run the sanitizer integration tests, and confirm that they pass. You can run check-all to run everything and be sure as well.

Sure, I will do after this lands.

llvm/include/llvm/CodeGen/MachineInstr.h

109

This NoMerge flag is used to disable branch-fold optimization when present. What do you mean something semantically required?

[arsenm]

llvm/include/llvm/CodeGen/MachineInstr.h

109

I forgot those were also here. Maybe this should be sorted up to the other mandatory flags. I'm pretty sure FrameSetup only changes debug info, so I guess that's not exactly mandatory?

I would still like a better comment here. This one just restates the name of the flag and doesn't tell me what it means by merge

[rnk]

I think this looks good, but I want to find another reviewer for CodeGen. @craig.topper or @arsenm, are you OK reviewing this, or can you recommend another reviewer?

llvm/include/llvm/CodeGen/MachineInstr.h

109

I second Matt's comments on reordering the flags to group this flag away from the math flags, which are generally ignorable, and you can expand on this comment here. The notion of "merging" instructions isn't self-explanatory. Perhaps "Passes that drop source location (e.g. tail merging) info should skip this instruction." could work.

[zequanwu]

Update comment for NoMerge flag

[MaskRay]

llvm/test/CodeGen/PowerPC/nomerge.ll

31 ↗

(On Diff #263842)

Change the three CHECK: bl bar to CHECK-NEXT: bl bar

It makes the reader clear they are first instructions in their basic blocks and there will be a friendly error message if things go wrong.

; CHECK:      # %if.then2
; CHECK-NEXT:  bl bar

llvm/test/CodeGen/RISCV/nomerge.ll

31 ↗

(On Diff #263842)

ditto

[zequanwu]

Change some CHECK to CHECK-NEXT where the call to bar is first instruction in blocks.

[craig.topper]

llvm/lib/Target/X86/X86ISelLowering.cpp

4272

Does this need to be in this else or can it apply to both paths?

[zequanwu]

rebase and address comment.

llvm/lib/Target/X86/X86ISelLowering.cpp

4272

Yes, it can be applied to both paths.

[craig.topper]

LGTM

llvm/include/llvm/CodeGen/MachineInstr.h

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,             // Instruction should not be merged
   };
 
 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;
   };
 
   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) {
+    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

struct CallLoweringInfo {
     bool RetSExt           : 1;
     bool RetZExt           : 1;
     bool IsVarArg          : 1;
     bool IsInReg           : 1;
     bool DoesNotReturn     : 1;
     bool IsReturnValueUsed : 1;
     bool IsConvergent      : 1;
     bool IsPatchPoint      : 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), DAG(DAG) {}
+          IsPatchPoint(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/X86/X86ISelLowering.cpp

if (isTailCall) {
     DAG.addCallSiteInfo(Ret.getNode(), std::move(CSInfo));
     return Ret;
   }
 
   if (HasNoCfCheck && IsCFProtectionSupported) {
     Chain = DAG.getNode(X86ISD::NT_CALL, dl, NodeTys, Ops);
   } else {
     Chain = DAG.getNode(X86ISD::CALL, dl, NodeTys, Ops);
+    // set nomerge.
+    DAG.addNoMergeSiteInfo(Chain.getNode(), CLI.NoMerge);
   }
   InFlag = Chain.getValue(1);
   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);

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: callq bar
+; CHECK: jmp bar # TAILCALL
+; CHECK: .LBB0_3: # %if.then2
+; CHECK: callq bar
+; CHECK: .LBB0_4: # %if.end3
+; CHECK: jmp bar # TAILCALL