[NVPTX] Use different, convergent MIs for convergent calls.

[jlebar]

Summary

Calls sometimes need to be convergent. This is already handled at the
LLVM IR level, but it also needs to be handled at the MI level.

Ideally we'd propagate convergence from instructions, down through the
selection DAG, and into MIs. But this is Hard, and would affect
optimizations in the SDNs -- right now only SDNs with two operands have
any flags at all.

Instead, here's a much simpler hack: Add new opcodes for NVPTX for
convergent calls, and generate these when lowering convergent LLVM
calls.

[jlebar]

Friendly ping. I know I have some big reviews outstanding from you, but the issue addressed here is a fairly critical correctness fix. I'm also happy to have tra or someone review this, but just wanted to respect your ownership and expertise here.

[mehdi_amini]

If no one else chime in, this LGTM.
(it would be better to have someone more familiar with NVPTX to double check )

[jholewinski]

Looks reasonable to me. Thanks!

include/llvm/Target/TargetLowering.h

struct CallLoweringInfo {
     SDValue Chain;
     Type *RetTy;
     bool RetSExt           : 1;
     bool RetZExt           : 1;
     bool IsVarArg          : 1;
     bool IsInReg           : 1;
     bool DoesNotReturn     : 1;
     bool IsReturnValueUsed : 1;
+    bool IsConvergent      : 1;
 
     // IsTailCall should be modified by implementations of
     // TargetLowering::LowerCall that perform tail call conversions.
     bool IsTailCall;
 
     unsigned NumFixedArgs;
     CallingConv::ID CallConv;
     SDValue Callee;
     ArgListTy Args;
     SelectionDAG &DAG;
     SDLoc DL;
     ImmutableCallSite *CS;
     bool IsPatchPoint;
     SmallVector<ISD::OutputArg, 32> Outs;
     SmallVector<SDValue, 32> OutVals;
     SmallVector<ISD::InputArg, 32> Ins;
 
     CallLoweringInfo(SelectionDAG &DAG)
         : RetTy(nullptr), RetSExt(false), RetZExt(false), IsVarArg(false),
           IsInReg(false), DoesNotReturn(false), IsReturnValueUsed(true),
-        IsTailCall(false), NumFixedArgs(-1), CallConv(CallingConv::C),
-        DAG(DAG), CS(nullptr), IsPatchPoint(false) {}
+          IsConvergent(false), IsTailCall(false), NumFixedArgs(-1),
+          CallConv(CallingConv::C), DAG(DAG), CS(nullptr), IsPatchPoint(false) {
+    }
 
     CallLoweringInfo &setDebugLoc(SDLoc dl) {
       DL = dl;
       return *this;
     }
 
     CallLoweringInfo &setChain(SDValue InChain) {
       Chain = InChain;
CallLoweringInfo &setTailCall(bool Value = true) {
       return *this;
     }
 
     CallLoweringInfo &setDiscardResult(bool Value = true) {
       IsReturnValueUsed = !Value;
       return *this;
     }
 
+    CallLoweringInfo &setConvergent(bool Value = true) {
+      IsConvergent = Value;
+      return *this;
+    }
+
     CallLoweringInfo &setSExtResult(bool Value = true) {
       RetSExt = Value;
       return *this;
     }
 
     CallLoweringInfo &setZExtResult(bool Value = true) {
       RetZExt = Value;
       return *this;

lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

void SelectionDAGBuilder::LowerCallTo(ImmutableCallSite CS, SDValue Callee,
   }
 
   // Check if target-independent constraints permit a tail call here.
   // Target-dependent constraints are checked within TLI->LowerCallTo.
   if (isTailCall && !isInTailCallPosition(CS, DAG.getTarget()))
     isTailCall = false;
 
   TargetLowering::CallLoweringInfo CLI(DAG);
-  CLI.setDebugLoc(getCurSDLoc()).setChain(getRoot())
+  CLI.setDebugLoc(getCurSDLoc())
+      .setChain(getRoot())
       .setCallee(RetTy, FTy, Callee, std::move(Args), CS)
-    .setTailCall(isTailCall);
+      .setTailCall(isTailCall)
+      .setConvergent(CS.isConvergent());
   std::pair<SDValue, SDValue> Result = lowerInvokable(CLI, EHPadBB);
 
   if (Result.first.getNode()) {
     const Instruction *Inst = CS.getInstruction();
     Result.first = lowerRangeToAssertZExt(DAG, *Inst, Result.first);
     setValue(Inst, Result.first);
   }
 }

lib/Target/NVPTX/NVPTXISelLowering.h

enum NodeType : unsigned {
   RET_FLAG,
   LOAD_PARAM,
   DeclareParam,
   DeclareScalarParam,
   DeclareRetParam,
   DeclareRet,
   DeclareScalarRet,
   PrintCall,
+  PrintConvergentCall,
   PrintCallUni,
+  PrintConvergentCallUni,
   CallArgBegin,
   CallArg,
   LastCallArg,
   CallArgEnd,
   CallVoid,
   CallVal,
   CallSymbol,
   Prototype,

lib/Target/NVPTX/NVPTXISelLowering.cpp

const char *NVPTXTargetLowering::getTargetNodeName(unsigned Opcode) const {
   case NVPTXISD::DeclareRet:
     return "NVPTXISD::DeclareRet";
   case NVPTXISD::DeclareScalarRet:
     return "NVPTXISD::DeclareScalarRet";
   case NVPTXISD::DeclareRetParam:
     return "NVPTXISD::DeclareRetParam";
   case NVPTXISD::PrintCall:
     return "NVPTXISD::PrintCall";
+  case NVPTXISD::PrintConvergentCall:
+    return "NVPTXISD::PrintConvergentCall";
   case NVPTXISD::PrintCallUni:
     return "NVPTXISD::PrintCallUni";
+  case NVPTXISD::PrintConvergentCallUni:
+    return "NVPTXISD::PrintConvergentCallUni";
   case NVPTXISD::LoadParam:
     return "NVPTXISD::LoadParam";
   case NVPTXISD::LoadParamV2:
     return "NVPTXISD::LoadParamV2";
   case NVPTXISD::LoadParamV4:
     return "NVPTXISD::LoadParamV4";
   case NVPTXISD::StoreParam:
     return "NVPTXISD::StoreParam";
if (!Func) {
     Chain = DAG.getNode(NVPTXISD::CallPrototype, dl, ProtoVTs, ProtoOps);
     InFlag = Chain.getValue(1);
   }
   // Op to just print "call"
   SDVTList PrintCallVTs = DAG.getVTList(MVT::Other, MVT::Glue);
   SDValue PrintCallOps[] = {
     Chain, DAG.getConstant((Ins.size() == 0) ? 0 : 1, dl, MVT::i32), InFlag
   };
-  Chain = DAG.getNode(Func ? (NVPTXISD::PrintCallUni) : (NVPTXISD::PrintCall),
-                      dl, PrintCallVTs, PrintCallOps);
+  // We model convergent calls as separate opcodes.
+  unsigned Opcode = Func ? NVPTXISD::PrintCallUni : NVPTXISD::PrintCall;
+  if (CLI.IsConvergent)
+    Opcode = Opcode == NVPTXISD::PrintCallUni ? NVPTXISD::PrintConvergentCallUni
+                                              : NVPTXISD::PrintConvergentCall;
+  Chain = DAG.getNode(Opcode, dl, PrintCallVTs, PrintCallOps);
   InFlag = Chain.getValue(1);
 
   // Ops to print out the function name
   SDVTList CallVoidVTs = DAG.getVTList(MVT::Other, MVT::Glue);
   SDValue CallVoidOps[] = { Chain, Callee, InFlag };
   Chain = DAG.getNode(NVPTXISD::CallVoid, dl, CallVoidVTs, CallVoidOps);
   InFlag = Chain.getValue(1);
 

lib/Target/NVPTX/NVPTXInstrInfo.td

def LoadParamV2 :
   SDNode<"NVPTXISD::LoadParamV2", SDTLoadParamV2Profile,
          [SDNPHasChain, SDNPMayLoad, SDNPOutGlue, SDNPInGlue]>;
 def LoadParamV4 :
   SDNode<"NVPTXISD::LoadParamV4", SDTLoadParamV4Profile,
          [SDNPHasChain, SDNPMayLoad, SDNPOutGlue, SDNPInGlue]>;
 def PrintCall :
   SDNode<"NVPTXISD::PrintCall", SDTPrintCallProfile,
          [SDNPHasChain, SDNPOutGlue, SDNPInGlue, SDNPSideEffect]>;
+def PrintConvergentCall :
+  SDNode<"NVPTXISD::PrintConvergentCall", SDTPrintCallProfile,
+         [SDNPHasChain, SDNPOutGlue, SDNPInGlue, SDNPSideEffect]>;
 def PrintCallUni :
   SDNode<"NVPTXISD::PrintCallUni", SDTPrintCallUniProfile,
          [SDNPHasChain, SDNPOutGlue, SDNPInGlue, SDNPSideEffect]>;
+def PrintConvergentCallUni :
+  SDNode<"NVPTXISD::PrintConvergentCallUni", SDTPrintCallUniProfile,
+         [SDNPHasChain, SDNPOutGlue, SDNPInGlue, SDNPSideEffect]>;
 def StoreParam :
   SDNode<"NVPTXISD::StoreParam", SDTStoreParamProfile,
          [SDNPHasChain, SDNPOutGlue, SDNPInGlue, SDNPSideEffect]>;
 def StoreParamV2 :
   SDNode<"NVPTXISD::StoreParamV2", SDTStoreParamV2Profile,
          [SDNPHasChain, SDNPOutGlue, SDNPInGlue, SDNPSideEffect]>;
 def StoreParamV4 :
   SDNode<"NVPTXISD::StoreParamV4", SDTStoreParamV4Profile,
class StoreRetvalV4Inst<NVPTXRegClass regclass, string opstr> :
       NVPTXInst<(outs),
                 (ins regclass:$val, regclass:$val2, regclass:$val3,
                      regclass:$val4, i32imm:$a),
                 !strconcat("st.param.v4", opstr,
                            "\t[func_retval0+$a], {{$val, $val2, $val3, $val4}};"),
                 []>;
 
 let isCall=1 in {
+  multiclass CALL<string OpcStr, SDNode OpNode> {
      def PrintCallNoRetInst : NVPTXInst<(outs), (ins),
-   "call ", [(PrintCall (i32 0))]>;
+       !strconcat(OpcStr, " "), [(OpNode (i32 0))]>;
      def PrintCallRetInst1 : NVPTXInst<(outs), (ins),
-   "call (retval0), ", [(PrintCall (i32 1))]>;
+       !strconcat(OpcStr, " (retval0), "), [(OpNode (i32 1))]>;
      def PrintCallRetInst2 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1), ", [(PrintCall (i32 2))]>;
+       !strconcat(OpcStr, " (retval0, retval1), "), [(OpNode (i32 2))]>;
      def PrintCallRetInst3 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1, retval2), ", [(PrintCall (i32 3))]>;
+       !strconcat(OpcStr, " (retval0, retval1, retval2), "), [(OpNode (i32 3))]>;
      def PrintCallRetInst4 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1, retval2, retval3), ", [(PrintCall (i32 4))]>;
+       !strconcat(OpcStr, " (retval0, retval1, retval2, retval3), "),
+       [(OpNode (i32 4))]>;
      def PrintCallRetInst5 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1, retval2, retval3, retval4), ",
-   [(PrintCall (i32 5))]>;
+       !strconcat(OpcStr, " (retval0, retval1, retval2, retval3, retval4), "),
+       [(OpNode (i32 5))]>;
      def PrintCallRetInst6 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1, retval2, retval3, retval4, retval5), ",
-   [(PrintCall (i32 6))]>;
+       !strconcat(OpcStr, " (retval0, retval1, retval2, retval3, retval4, "
+                            "retval5), "),
+       [(OpNode (i32 6))]>;
      def PrintCallRetInst7 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1, retval2, retval3, retval4, retval5, retval6), ",
-   [(PrintCall (i32 7))]>;
+       !strconcat(OpcStr, " (retval0, retval1, retval2, retval3, retval4, "
+                            "retval5, retval6), "),
+       [(OpNode (i32 7))]>;
      def PrintCallRetInst8 : NVPTXInst<(outs), (ins),
-   "call (retval0, retval1, retval2, retval3, retval4, retval5, retval6, "
-         "retval7), ",
-   [(PrintCall (i32 8))]>;
-
- def PrintCallUniNoRetInst : NVPTXInst<(outs), (ins),
-   "call.uni ", [(PrintCallUni (i32 0))]>;
- def PrintCallUniRetInst1 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0), ", [(PrintCallUni (i32 1))]>;
- def PrintCallUniRetInst2 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1), ", [(PrintCallUni (i32 2))]>;
- def PrintCallUniRetInst3 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1, retval2), ", [(PrintCallUni (i32 3))]>;
- def PrintCallUniRetInst4 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1, retval2, retval3), ", [(PrintCallUni (i32 4))]>;
- def PrintCallUniRetInst5 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1, retval2, retval3, retval4), ",
-   [(PrintCallUni (i32 5))]>;
- def PrintCallUniRetInst6 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1, retval2, retval3, retval4, retval5), ",
-   [(PrintCallUni (i32 6))]>;
- def PrintCallUniRetInst7 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1, retval2, retval3, retval4, retval5, retval6), ",
-   [(PrintCallUni (i32 7))]>;
- def PrintCallUniRetInst8 : NVPTXInst<(outs), (ins),
-   "call.uni (retval0, retval1, retval2, retval3, retval4, retval5, retval6, "
-             "retval7), ",
-   [(PrintCallUni (i32 8))]>;
+       !strconcat(OpcStr, " (retval0, retval1, retval2, retval3, retval4, "
+                            "retval5, retval6, retval7), "),
+       [(OpNode (i32 8))]>;
+  }
+}
+
+defm Call : CALL<"call", PrintCall>;
+defm CallUni : CALL<"call.uni", PrintCallUni>;
+
+// Convergent call instructions.  These are identical to regular calls, except
+// they have the isConvergent bit set.
+let isConvergent=1 in {
+  defm ConvergentCall : CALL<"call", PrintConvergentCall>;
+  defm ConvergentCallUni : CALL<"call.uni", PrintConvergentCallUni>;
 }
 
 def LoadParamMemI64    : LoadParamMemInst<Int64Regs, ".b64">;
 def LoadParamMemI32    : LoadParamMemInst<Int32Regs, ".b32">;
 def LoadParamMemI16    : LoadParamMemInst<Int16Regs, ".b16">;
 def LoadParamMemI8     : LoadParamMemInst<Int16Regs, ".b8">;
 def LoadParamMemV2I64  : LoadParamV2MemInst<Int64Regs, ".b64">;
 def LoadParamMemV2I32  : LoadParamV2MemInst<Int32Regs, ".b32">;

test/CodeGen/NVPTX/convergent-mir-call.ll

+; RUN: llc -mtriple nvptx64-nvidia-cuda -stop-after machine-cp -o - < %s 2>&1 | FileCheck %s
+
+; Check that convergent calls are emitted using convergent MIR instructions,
+; while non-convergent calls are not.
+
+target triple = "nvptx64-nvidia-cuda"
+
+declare void @conv() convergent
+declare void @not_conv()
+
+define void @test(void ()* %f) {
+  ; CHECK: ConvergentCallUniPrintCall
+  ; CHECK-NEXT: @conv
+  call void @conv()
+
+  ; CHECK: CallUniPrintCall
+  ; CHECK-NEXT: @not_conv
+  call void @not_conv()
+
+  ; CHECK: ConvergentCallPrintCall
+  call void %f() convergent
+
+  ; CHECK: CallPrintCall
+  call void %f()
+
+  ret void
+}