stackmaps.diff

diff --git a/llvm/include/llvm/CodeGen/ISDOpcodes.h b/llvm/include/llvm/CodeGen/ISDOpcodes.h
index ea1d3170acba..531edff27c70 100644
--- a/llvm/include/llvm/CodeGen/ISDOpcodes.h
+++ b/llvm/include/llvm/CodeGen/ISDOpcodes.h
@@ -1269,6 +1269,9 @@ enum NodeType {
 #define BEGIN_REGISTER_VP_SDNODE(VPSDID, ...) VPSDID,
 #include "llvm/IR/VPIntrinsics.def"
 
+  // The `llvm.experimental.stackmap` intrinsic.
+  STACKMAP,
+
   /// BUILTIN_OP_END - This must be the last enum value in this list.
   /// The target-specific pre-isel opcode values start here.
   BUILTIN_OP_END
diff --git a/llvm/include/llvm/CodeGen/SelectionDAGISel.h b/llvm/include/llvm/CodeGen/SelectionDAGISel.h
index 35fb0bc80593..87df6d1b1604 100644
--- a/llvm/include/llvm/CodeGen/SelectionDAGISel.h
+++ b/llvm/include/llvm/CodeGen/SelectionDAGISel.h
@@ -321,6 +321,7 @@ private:
 
   void Select_FREEZE(SDNode *N);
   void Select_ARITH_FENCE(SDNode *N);
+  void Select_STACKMAP(SDNode *N);
 
 private:
   void DoInstructionSelection();
diff --git a/llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp b/llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp
index 1f4256662031..c2cad0dc3c30 100644
--- a/llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp
@@ -1150,6 +1150,7 @@ EmitSpecialNode(SDNode *Node, bool IsClone, bool IsCloned,
 #ifndef NDEBUG
     Node->dump();
 #endif
+    Node->dump();
     llvm_unreachable("This target-independent node should have been selected!");
   case ISD::EntryToken:
     llvm_unreachable("EntryToken should have been excluded from the schedule!");
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
index a1ddb02563e3..1c929f7dee2c 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeIntegerTypes.cpp
@@ -24,6 +24,7 @@
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
+#include <vector>
 using namespace llvm;
 
 #define DEBUG_TYPE "legalize-types"
@@ -1683,6 +1684,9 @@ bool DAGTypeLegalizer::PromoteIntegerOperand(SDNode *N, unsigned OpNo) {
     break;
 
   case ISD::SET_ROUNDING: Res = PromoteIntOp_SET_ROUNDING(N); break;
+  case ISD::STACKMAP:
+    Res = PromoteIntOp_STACKMAP(N, OpNo);
+    break;
   }
 
   // If the result is null, the sub-method took care of registering results etc.
@@ -2246,6 +2250,17 @@ SDValue DAGTypeLegalizer::PromoteIntOp_SET_ROUNDING(SDNode *N) {
   return SDValue(DAG.UpdateNodeOperands(N, N->getOperand(0), Op), 0);
 }
 
+SDValue DAGTypeLegalizer::PromoteIntOp_STACKMAP(SDNode *N, unsigned OpNo) {
+  assert(OpNo > 1); // Because the first two arguments are guaranteed legal.
+
+  // XXX: This legalises all operands at once. Is that OK?
+  std::vector<SDValue> NewOps = {N->getOperand(0), N->getOperand(1)};
+  for (size_t i = 2; i < N->getNumOperands(); i++)
+      NewOps.push_back(ZExtPromotedInteger(N->getOperand(i)));
+
+  return SDValue(DAG.UpdateNodeOperands(N, NewOps), 0);
+}
+
 //===----------------------------------------------------------------------===//
 //  Integer Result Expansion
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index 5642908961bf..57c4cf3c5741 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
@@ -403,6 +403,8 @@ private:
   SDValue PromoteIntOp_VECREDUCE(SDNode *N);
   SDValue PromoteIntOp_VP_REDUCE(SDNode *N, unsigned OpNo);
   SDValue PromoteIntOp_SET_ROUNDING(SDNode *N);
+  SDValue PromoteIntOp_STACKMAP(SDNode *N, unsigned OpNo);
+
 
   void PromoteSetCCOperands(SDValue &LHS,SDValue &RHS, ISD::CondCode Code);
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
index 90e4b5d71be6..3835955ce625 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp
@@ -9935,10 +9935,13 @@ void SelectionDAG::ReplaceAllUsesWith(SDValue FromN, SDValue To) {
 ///
 void SelectionDAG::ReplaceAllUsesWith(SDNode *From, SDNode *To) {
 #ifndef NDEBUG
-  for (unsigned i = 0, e = From->getNumValues(); i != e; ++i)
+  for (unsigned i = 0, e = From->getNumValues(); i != e; ++i) {
+    errs() << "has use: " << From->hasAnyUseOfValue(i) << "\n";
+    errs() << "type ok: " << (From->getValueType(i) == To->getValueType(i)) << "\n";
     assert((!From->hasAnyUseOfValue(i) ||
             From->getValueType(i) == To->getValueType(i)) &&
            "Cannot use this version of ReplaceAllUsesWith!");
+  }
 #endif
 
   // Handle the trivial case.
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 170fa25a2960..ab31d2d719c0 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -9376,64 +9376,22 @@ static void addStackMapLiveVars(const CallBase &Call, unsigned StartIdx,
   }
 }
 
-/// Lower llvm.experimental.stackmap directly to its target opcode.
+/// Lower llvm.experimental.stackmap.
 void SelectionDAGBuilder::visitStackmap(const CallInst &CI) {
   // void @llvm.experimental.stackmap(i32 <id>, i32 <numShadowBytes>,
   //                                  [live variables...])
 
   assert(CI.getType()->isVoidTy() && "Stackmap cannot return a value.");
 
-  SDValue Chain, InFlag, Callee, NullPtr;
-  SmallVector<SDValue, 32> Ops;
+  // Create nodes for the operands.
+  std::vector<SDValue> Operands;
+  for (unsigned I = 0; I < CI.arg_size(); I++)
+    Operands.push_back(getValue(CI.getArgOperand(I)));
 
+  // Emit the STACKMAP node.
   SDLoc DL = getCurSDLoc();
-  Callee = getValue(CI.getCalledOperand());
-  NullPtr = DAG.getIntPtrConstant(0, DL, true);
-
-  // The stackmap intrinsic only records the live variables (the arguments
-  // passed to it) and emits NOPS (if requested). Unlike the patchpoint
-  // intrinsic, this won't be lowered to a function call. This means we don't
-  // have to worry about calling conventions and target specific lowering code.
-  // Instead we perform the call lowering right here.
-  //
-  // chain, flag = CALLSEQ_START(chain, 0, 0)
-  // chain, flag = STACKMAP(id, nbytes, ..., chain, flag)
-  // chain, flag = CALLSEQ_END(chain, 0, 0, flag)
-  //
-  Chain = DAG.getCALLSEQ_START(getRoot(), 0, 0, DL);
-  InFlag = Chain.getValue(1);
-
-  // Add the <id> and <numBytes> constants.
-  SDValue IDVal = getValue(CI.getOperand(PatchPointOpers::IDPos));
-  Ops.push_back(DAG.getTargetConstant(
-                  cast<ConstantSDNode>(IDVal)->getZExtValue(), DL, MVT::i64));
-  SDValue NBytesVal = getValue(CI.getOperand(PatchPointOpers::NBytesPos));
-  Ops.push_back(DAG.getTargetConstant(
-                  cast<ConstantSDNode>(NBytesVal)->getZExtValue(), DL,
-                  MVT::i32));
-
-  // Push live variables for the stack map.
-  addStackMapLiveVars(CI, 2, DL, Ops, *this);
-
-  // We are not pushing any register mask info here on the operands list,
-  // because the stackmap doesn't clobber anything.
-
-  // Push the chain and the glue flag.
-  Ops.push_back(Chain);
-  Ops.push_back(InFlag);
-
-  // Create the STACKMAP node.
-  SDVTList NodeTys = DAG.getVTList(MVT::Other, MVT::Glue);
-  SDNode *SM = DAG.getMachineNode(TargetOpcode::STACKMAP, DL, NodeTys, Ops);
-  Chain = SDValue(SM, 0);
-  InFlag = Chain.getValue(1);
-
-  Chain = DAG.getCALLSEQ_END(Chain, NullPtr, NullPtr, InFlag, DL);
-
-  // Stackmaps don't generate values, so nothing goes into the NodeMap.
-
-  // Set the root to the target-lowered call chain.
-  DAG.setRoot(Chain);
+  SDValue Res = DAG.getNode(ISD::STACKMAP, DL, MVT::Other, Operands);
+  DAG.setRoot(Res);
 
   // Inform the Frame Information that we have a stackmap in this function.
   FuncInfo.MF->getFrameInfo().setHasStackMap();
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
index 60a3c02e38cf..6effa2987bda 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
@@ -490,6 +490,8 @@ std::string SDNode::getOperationName(const SelectionDAG *G) const {
   case ISD::SDID:                                                              \
     return #NAME;
 #include "llvm/IR/VPIntrinsics.def"
+
+  case ISD::STACKMAP:                   return "stackmap";
   }
 }
 
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
index 2b6524223546..6704e54da728 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
@@ -51,6 +51,7 @@
 #include "llvm/CodeGen/SchedulerRegistry.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/SelectionDAGNodes.h"
+#include "llvm/CodeGen/StackMaps.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/SwiftErrorValueTracking.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
@@ -2238,6 +2239,58 @@ void SelectionDAGISel::Select_ARITH_FENCE(SDNode *N) {
                        N->getOperand(0));
 }
 
+void SelectionDAGISel::Select_STACKMAP(SDNode *N) {
+  std::vector<SDValue> Ops;
+  std::vector<EVT> VTs;
+  auto *It = N->op_begin();
+  SDLoc DL(N);
+
+  // Lower the "stackmap ID" operand.
+  SDValue ID = It->get();
+  assert(ID.getOpcode() == ISD::Constant);
+  assert(ID.getValueType() == MVT::i64);
+  SDValue IDConst = CurDAG->getTargetConstant(cast<ConstantSDNode>(ID)->getZExtValue(),
+              DL, ID.getValueType());
+  Ops.push_back(IDConst);
+  VTs.push_back(ID.getValueType());
+  It++;
+
+  // Lower the "num shadow bytes" operand.
+  SDValue Shad = It->get();
+  assert(Shad.getOpcode() == ISD::Constant);
+  assert(Shad.getValueType() == MVT::i32);
+  SDValue ShadConst = CurDAG->getTargetConstant(cast<ConstantSDNode>(Shad)->getZExtValue(),
+          DL, Shad.getValueType());
+  Ops.push_back(ShadConst);
+  VTs.push_back(Shad.getValueType());
+  It++;
+
+  // Lower the live variable operands.
+  for (;It != N->op_end(); It++)
+    Ops.push_back(*It);
+
+  SDValue New = CurDAG->getNode(TargetOpcode::STACKMAP, DL, VTs, Ops);
+  New->setNodeId(-1);
+  // The following line crashes:
+  // clang-15:
+  // /home/vext01/research/ykllvm/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp:9941:
+  // void llvm::SelectionDAG::ReplaceAllUsesWith(llvm::SDNode*, llvm::SDNode*):
+  // Assertion `(!From->hasAnyUseOfValue(i) || From->getValueType(i) ==
+  // To->getValueType(i)) && "Cannot use this version of ReplaceAllUsesWith!"'
+  // failed.
+  ReplaceNode(N, New.getNode());
+  CurDAG->RemoveDeadNode(N);
+
+  // This would have been shorter, but sadly it selects the first two operands
+  // into registers, later crashing the lowering code, which expects
+  // immediates.
+  //std::vector<SDValue> Ops2;
+  //for (auto &N: N->ops()) {
+  //    Ops2.push_back(N);
+  //}
+  //CurDAG->SelectNodeTo(N, TargetOpcode::STACKMAP, N->getValueType(0), Ops2);
+}
+
 /// GetVBR - decode a vbr encoding whose top bit is set.
 LLVM_ATTRIBUTE_ALWAYS_INLINE static uint64_t
 GetVBR(uint64_t Val, const unsigned char *MatcherTable, unsigned &Idx) {
@@ -2792,6 +2845,9 @@ void SelectionDAGISel::SelectCodeCommon(SDNode *NodeToMatch,
   case ISD::ARITH_FENCE:
     Select_ARITH_FENCE(NodeToMatch);
     return;
+  case ISD::STACKMAP:
+    Select_STACKMAP(NodeToMatch);
+    return;
   }
 
   assert(!NodeToMatch->isMachineOpcode() && "Node already selected!");
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
index 14e8f2d3f934..e3bfc0544781 100644
--- a/llvm/lib/Target/X86/X86ISelLowering.cpp
+++ b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -2209,6 +2209,8 @@ X86TargetLowering::X86TargetLowering(const X86TargetMachine &TM,
       if (isOperationExpand(Op, MVT::f32))
         setOperationAction(Op, MVT::f32, Promote);
 
+  setOperationAction(ISD::STACKMAP, MVT::Any, Custom);
+
   // We have target-specific dag combine patterns for the following nodes:
   setTargetDAGCombine({ISD::VECTOR_SHUFFLE,
                        ISD::SCALAR_TO_VECTOR,