[CSSPGO] MIR target-independent pseudo instruction for pseudo-probe intrinsic

[hoy]

Summary

This change introduces a MIR target-independent pseudo instruction corresponding to the IR intrinsic llvm.pseudoprobe for pseudo-probe block instrumentation. Please refer to https://reviews.llvm.org/D86193 for the whole story.

An llvm.pseudoprobe intrinsic call will be lowered into a target-independent operation named PSEUDO_PROBE. Given the following instrumented IR,

define internal void @foo2(i32 %x, void (i32)* %f) !dbg !4 {
bb0:
   %cmp = icmp eq i32 %x, 0
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 1)
   br i1 %cmp, label %bb1, label %bb2
bb1:                                             
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 2)
   br label %bb3
bb2:                                              
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 3)
   br label %bb3
bb3:                                              
   call void @llvm.pseudoprobe(i64 837061429793323041, i64 4)
   ret void
}

the corresponding MIR is shown below. Note that block bb3 is duplicated into bb1 and bb2 where its probe is duplicated too. This allows for an accurate execution count to be collected for bb3, which is basically the sum of the counts of bb1 and bb2.

bb.0.bb0:
   frame-setup PUSH64r undef $rax, implicit-def $rsp, implicit $rsp
   TEST32rr killed renamable $edi, renamable $edi, implicit-def $eflags
   PSEUDO_PROBE 837061429793323041, 1, 0
   $edi = MOV32ri 1, debug-location !13; test.c:0
   JCC_1 %bb.1, 4, implicit $eflags

bb.2.bb2:
   PSEUDO_PROBE 837061429793323041, 3, 0
   PSEUDO_PROBE 837061429793323041, 4, 0
   $rax = frame-destroy POP64r implicit-def $rsp, implicit $rsp
   RETQ

bb.1.bb1:
   PSEUDO_PROBE 837061429793323041, 2, 0
   PSEUDO_PROBE 837061429793323041, 4, 0
   $rax = frame-destroy POP64r implicit-def $rsp, implicit $rsp
   RETQ

The target op PSEUDO_PROBE will be converted into a piece of binary data by the object emitter with no machine instructions generated. This is done in a different patch.

[hoy]

Updating D86495: [CSSPGO] MIR target-independent pseudo instruction for pseudo-probe intrinsic

[hoy]

@wmi Do you know how to chain this patch with my previous patch https://reviews.llvm.org/D86490 so that the two patches can be tested together? Thanks.

[wmi]

I don't know since I don't use arc.

[wmi]

It is better to have a MIR test.

[hoy]

In D86495#2264868, @wmi wrote:

It is better to have a MIR test.

Yeah, the test is actually included in https://reviews.llvm.org/D86499. It tests both the IR and MIR change.

[wmi]

In D86495#2264877, @hoy wrote:

In D86495#2264868, @wmi wrote:

It is better to have a MIR test.

Yeah, the test is actually included in https://reviews.llvm.org/D86499. It tests both the IR and MIR change.

I mean to have a .mir test which read MIR as input. That is to ensure llc read MIR text containing the new pseudo probe without problem.

[hoy]

In D86495#2264881, @wmi wrote:

In D86495#2264877, @hoy wrote:

In D86495#2264868, @wmi wrote:

It is better to have a MIR test.

Yeah, the test is actually included in https://reviews.llvm.org/D86499. It tests both the IR and MIR change.

I mean to have a .mir test which read MIR as input. That is to ensure llc read MIR text containing the new pseudo probe without problem.

I see. Added a .mir test.

[hoy]

Updating D86495: [CSSPGO] MIR target-independent pseudo instruction for pseudo-probe intrinsic

[hoy]

Updating D86495: [CSSPGO] MIR target-independent pseudo instruction for pseudo-probe intrinsic

Adding an attribute field for use later.

[mceier]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

The size of this class (66 if I computed it correctly) exceeds the size of LargestSDNode (64 bytes) in llvm/include/llvm/CodeGen/SelectionDAGNodes.h causing static_assert to fail in Recycler.h when compiling SelectionDAG.cpp (at least) on 32-bit system:

In file included from llvm/include/llvm/CodeGen/MachineFunction.h:35,
                 from llvm/include/llvm/CodeGen/SelectionDAG.h:31,
                 from llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp:13:
llvm/include/llvm/Support/Recycler.h: In instantiation of ‘SubClass* llvm::Recycler<T, Size, Align>::Allocate(AllocatorType&) [with SubClass = llvm::PseudoProbeSDNode; AllocatorType = llvm::BumpPtrAllocatorImpl<>; T = llvm::SDNode; unsigned int Size = 64; unsigned int Align = 4]’:
llvm/include/llvm/Support/RecyclingAllocator.h:43:65:   required from ‘SubClass* llvm::RecyclingAllocator<AllocatorType, T, Size, Align>::Allocate() [with SubClass = llvm::PseudoProbeSDNode; AllocatorType = llvm::BumpPtrAllocatorImpl<>; T = llvm::SDNode; unsigned int Size = 64; unsigned int Align = 4]’
llvm/include/llvm/CodeGen/SelectionDAG.h:380:57:   required from ‘SDNodeT* llvm::SelectionDAG::newSDNode(ArgTypes&& ...) [with SDNodeT = llvm::PseudoProbeSDNode; ArgTypes = {const unsigned int&, unsigned int, const llvm::DebugLoc&, const llvm::SDVTList&, long long unsigned int&, long long unsigned int&, unsigned int&}]’
llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp:6906:72:   required from here
llvm/include/llvm/Support/Recycler.h:86:36: error: static assertion failed: Recycler allocation size is less than object size!
   86 |     static_assert(sizeof(SubClass) <= Size,
      |                   ~~~~~~~~~~~~~~~~~^~~~~~~

Relevant bug: https://bugs.llvm.org/show_bug.cgi?id=48259

I'm guessing this code in llvm/include/llvm/CodeGen/SelectionDAGNodes.h:

using LargestSDNode = AlignedCharArrayUnion<AtomicSDNode, TargetIndexSDNode,
                                            BlockAddressSDNode,
                                            GlobalAddressSDNode>;

should be changed to:

using LargestSDNode = AlignedCharArrayUnion<AtomicSDNode, TargetIndexSDNode,
                                            BlockAddressSDNode,
                                            GlobalAddressSDNode,
                                            PseudoProbeSDNode>;

[hoy]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

Thanks for reporting the bug! I'm sorry for the breakage. Will send out a fix.

[hoy]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

@mceier Could you do a favor by giving your suggested fix a try? I just realized that my system doesn't have a 32-bit toolset. Thanks!

[mceier]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

Yes. SelectionDAG.cpp compiled with that change, I'm still waiting for whole build to finish.

[hoy]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

Thanks for the update! Please let me know once the whole build is fine and I'll push the fix in.

[mceier]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

Just finished building llvm - both 32-bit and 64-bit versions build with this change.

[hoy]

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

1752–1777

Thanks for the confirmation! I just pushed the change in.

llvm/include/llvm/CodeGen/ISDOpcodes.h

   GC_TRANSITION_END,
 
   /// GET_DYNAMIC_AREA_OFFSET - get offset from native SP to the address of
   /// the most recent dynamic alloca. For most targets that would be 0, but
   /// for some others (e.g. PowerPC, PowerPC64) that would be compile-time
   /// known nonzero constant. The only operand here is the chain.
   GET_DYNAMIC_AREA_OFFSET,
 
+  /// Pseudo probe for AutoFDO, as a place holder in a basic block to improve
+  /// the sample counts quality.
+  PSEUDO_PROBE,
+
   /// VSCALE(IMM) - Returns the runtime scaling factor used to calculate the
   /// number of elements within a scalable vector. IMM is a constant integer
   /// multiplier that is applied to the runtime value.
   VSCALE,
 
   /// Generic reduction nodes. These nodes represent horizontal vector
   /// reduction operations, producing a scalar result.
   /// The SEQ variants perform reductions in sequential order. The first

llvm/include/llvm/CodeGen/MachineInstr.h

bool isMetaInstruction() const {
     case TargetOpcode::CFI_INSTRUCTION:
     case TargetOpcode::EH_LABEL:
     case TargetOpcode::GC_LABEL:
     case TargetOpcode::DBG_VALUE:
     case TargetOpcode::DBG_INSTR_REF:
     case TargetOpcode::DBG_LABEL:
     case TargetOpcode::LIFETIME_START:
     case TargetOpcode::LIFETIME_END:
+    case TargetOpcode::PSEUDO_PROBE:
       return true;
     }
   }
 
   /// Return true if this is a transient instruction that is either very likely
   /// to be eliminated during register allocation (such as copy-like
   /// instructions), or if this instruction doesn't have an execution-time cost.
   bool isTransient() const {

llvm/include/llvm/CodeGen/SelectionDAG.h

SDValue getMemIntrinsicNode(unsigned Opcode, const SDLoc &dl, SDVTList VTList,
                               MachineMemOperand *MMO);
 
   /// Creates a LifetimeSDNode that starts (`IsStart==true`) or ends
   /// (`IsStart==false`) the lifetime of the portion of `FrameIndex` between
   /// offsets `Offset` and `Offset + Size`.
   SDValue getLifetimeNode(bool IsStart, const SDLoc &dl, SDValue Chain,
                           int FrameIndex, int64_t Size, int64_t Offset = -1);
 
+  /// Creates a PseudoProbeSDNode with function GUID `Guid` and
+  /// the index of the block `Index` it is probing, as well as the attributes
+  /// `attr` of the probe.
+  SDValue getPseudoProbeNode(const SDLoc &Dl, SDValue Chain, uint64_t Guid,
+                             uint64_t Index, uint32_t Attr);
+
   /// Create a MERGE_VALUES node from the given operands.
   SDValue getMergeValues(ArrayRef<SDValue> Ops, const SDLoc &dl);
 
   /// Loads are not normal binary operators: their result type is not
   /// determined by their operands, and they produce a value AND a token chain.
   ///
   /// This function will set the MOLoad flag on MMOFlags, but you can set it if
   /// you want.  The MOStore flag must not be set.

llvm/include/llvm/CodeGen/SelectionDAGNodes.h

public:
 
   // Methods to support isa and dyn_cast
   static bool classof(const SDNode *N) {
     return N->getOpcode() == ISD::LIFETIME_START ||
            N->getOpcode() == ISD::LIFETIME_END;
   }
 };
 
+/// This SDNode is used for PSEUDO_PROBE values, which are the function guid and
+/// the index of the basic block being probed. A pseudo probe serves as a place
+/// holder and will be removed at the end of compilation. It does not have any
+/// operand because we do not want the instruction selection to deal with any.
+class PseudoProbeSDNode : public SDNode {
+  friend class SelectionDAG;
+  uint64_t Guid;
+  uint64_t Index;
+  uint32_t Attributes;
+
+  PseudoProbeSDNode(unsigned Opcode, unsigned Order, const DebugLoc &Dl,
+                    SDVTList VTs, uint64_t Guid, uint64_t Index, uint32_t Attr)
+      : SDNode(Opcode, Order, Dl, VTs), Guid(Guid), Index(Index),
+        Attributes(Attr) {}
+
+public:
+  uint64_t getGuid() const { return Guid; }
+  uint64_t getIndex() const { return Index; }
+  uint32_t getAttributes() const { return Attributes; }
+
+  // Methods to support isa and dyn_cast
+  static bool classof(const SDNode *N) {
+    return N->getOpcode() == ISD::PSEUDO_PROBE;
+  }
+};
+
 class JumpTableSDNode : public SDNode {
   friend class SelectionDAG;
 
   int JTI;
   unsigned TargetFlags;
 
   JumpTableSDNode(int jti, EVT VT, bool isTarg, unsigned TF)
     : SDNode(isTarg ? ISD::TargetJumpTable : ISD::JumpTable,

llvm/include/llvm/Support/TargetOpcodes.def

 /// which immediately follow a BUNDLE instruction which are marked with
 /// 'InsideBundle' flag are inside the bundle.
 HANDLE_TARGET_OPCODE(BUNDLE)
 
 /// Lifetime markers.
 HANDLE_TARGET_OPCODE(LIFETIME_START)
 HANDLE_TARGET_OPCODE(LIFETIME_END)
 
+/// Pseudo probe
+HANDLE_TARGET_OPCODE(PSEUDO_PROBE)
+
 /// A Stackmap instruction captures the location of live variables at its
 /// position in the instruction stream. It is followed by a shadow of bytes
 /// that must lie within the function and not contain another stackmap.
 HANDLE_TARGET_OPCODE(STACKMAP)
 
 /// FEntry all - This is a marker instruction which gets translated into a raw fentry call.
 HANDLE_TARGET_OPCODE(FENTRY_CALL)
 

llvm/include/llvm/Target/Target.td

def LIFETIME_START : StandardPseudoInstruction {
   let hasSideEffects = false;
 }
 def LIFETIME_END : StandardPseudoInstruction {
   let OutOperandList = (outs);
   let InOperandList = (ins i32imm:$id);
   let AsmString = "LIFETIME_END";
   let hasSideEffects = false;
 }
+def PSEUDO_PROBE : StandardPseudoInstruction {
+  let OutOperandList = (outs);
+  let InOperandList = (ins i64imm:$guid, i64imm:$index, i8imm:$type, i32imm:$attr);
+  let AsmString = "PSEUDO_PROBE";
+  let hasSideEffects = 1;
+}
+
 def STACKMAP : StandardPseudoInstruction {
   let OutOperandList = (outs);
   let InOperandList = (ins i64imm:$id, i32imm:$nbytes, variable_ops);
   let hasSideEffects = true;
   let isCall = true;
   let mayLoad = true;
   let usesCustomInserter = true;
 }

llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp

case ISD::LIFETIME_END: {
     TargetOpcode::LIFETIME_START : TargetOpcode::LIFETIME_END;
 
     FrameIndexSDNode *FI = dyn_cast<FrameIndexSDNode>(Node->getOperand(1));
     BuildMI(*MBB, InsertPos, Node->getDebugLoc(), TII->get(TarOp))
     .addFrameIndex(FI->getIndex());
     break;
   }
 
+  case ISD::PSEUDO_PROBE: {
+    unsigned TarOp = TargetOpcode::PSEUDO_PROBE;
+    auto Guid = cast<PseudoProbeSDNode>(Node)->getGuid();
+    auto Index = cast<PseudoProbeSDNode>(Node)->getIndex();
+    auto Attr = cast<PseudoProbeSDNode>(Node)->getAttributes();
+
+    BuildMI(*MBB, InsertPos, Node->getDebugLoc(), TII->get(TarOp))
+        .addImm(Guid)
+        .addImm(Index)
+        .addImm(0) // 0 for block probes
+        .addImm(Attr);
+    break;
+  }
+
   case ISD::INLINEASM:
   case ISD::INLINEASM_BR: {
     unsigned NumOps = Node->getNumOperands();
     if (Node->getOperand(NumOps-1).getValueType() == MVT::Glue)
       --NumOps;  // Ignore the glue operand.
 
     // Create the inline asm machine instruction.
     unsigned TgtOpc = Node->getOpcode() == ISD::INLINEASM_BR

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::TargetFrameIndex:
     break;
   case ISD::LIFETIME_START:
   case ISD::LIFETIME_END:
     if (cast<LifetimeSDNode>(N)->hasOffset()) {
       ID.AddInteger(cast<LifetimeSDNode>(N)->getSize());
       ID.AddInteger(cast<LifetimeSDNode>(N)->getOffset());
     }
     break;
+  case ISD::PSEUDO_PROBE:
+    ID.AddInteger(cast<PseudoProbeSDNode>(N)->getGuid());
+    ID.AddInteger(cast<PseudoProbeSDNode>(N)->getIndex());
+    ID.AddInteger(cast<PseudoProbeSDNode>(N)->getAttributes());
+    break;
   case ISD::JumpTable:
   case ISD::TargetJumpTable:
     ID.AddInteger(cast<JumpTableSDNode>(N)->getIndex());
     ID.AddInteger(cast<JumpTableSDNode>(N)->getTargetFlags());
     break;
   case ISD::ConstantPool:
   case ISD::TargetConstantPool: {
     const ConstantPoolSDNode *CP = cast<ConstantPoolSDNode>(N);
SDValue SelectionDAG::getLifetimeNode(bool IsStart, const SDLoc &dl,
   createOperands(N, Ops);
   CSEMap.InsertNode(N, IP);
   InsertNode(N);
   SDValue V(N, 0);
   NewSDValueDbgMsg(V, "Creating new node: ", this);
   return V;
 }
 
+SDValue SelectionDAG::getPseudoProbeNode(const SDLoc &Dl, SDValue Chain,
+                                         uint64_t Guid, uint64_t Index,
+                                         uint32_t Attr) {
+  const unsigned Opcode = ISD::PSEUDO_PROBE;
+  const auto VTs = getVTList(MVT::Other);
+  SDValue Ops[] = {Chain};
+  FoldingSetNodeID ID;
+  AddNodeIDNode(ID, Opcode, VTs, Ops);
+  ID.AddInteger(Guid);
+  ID.AddInteger(Index);
+  void *IP = nullptr;
+  if (SDNode *E = FindNodeOrInsertPos(ID, Dl, IP))
+    return SDValue(E, 0);
+
+  auto *N = newSDNode<PseudoProbeSDNode>(
+      Opcode, Dl.getIROrder(), Dl.getDebugLoc(), VTs, Guid, Index, Attr);
+  createOperands(N, Ops);
+  CSEMap.InsertNode(N, IP);
+  InsertNode(N);
+  SDValue V(N, 0);
+  NewSDValueDbgMsg(V, "Creating new node: ", this);
+  return V;
+}
+
 /// InferPointerInfo - If the specified ptr/offset is a frame index, infer a
 /// MachinePointerInfo record from it.  This is particularly useful because the
 /// code generator has many cases where it doesn't bother passing in a
 /// MachinePointerInfo to getLoad or getStore when it has "FI+Cst".
 static MachinePointerInfo InferPointerInfo(const MachinePointerInfo &Info,
                                            SelectionDAG &DAG, SDValue Ptr,
                                            int64_t Offset = 0) {
   // If this is FI+Offset, we can model it.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

for (SmallVectorImpl<const Value*>::iterator Object = Allocas.begin(),
               ObjectPtr, Offset, DAG.getDataLayout()) != LifetimeObject)
         Offset = -1; // Cannot determine offset from alloca to lifetime object.
       Res = DAG.getLifetimeNode(IsStart, sdl, getRoot(), FrameIndex, ObjectSize,
                                 Offset);
       DAG.setRoot(Res);
     }
     return;
   }
+  case Intrinsic::pseudoprobe: {
+    auto Guid = cast<ConstantInt>(I.getArgOperand(0))->getZExtValue();
+    auto Index = cast<ConstantInt>(I.getArgOperand(1))->getZExtValue();
+    auto Attr = cast<ConstantInt>(I.getArgOperand(2))->getZExtValue();
+    Res = DAG.getPseudoProbeNode(sdl, getRoot(), Guid, Index, Attr);
+    DAG.setRoot(Res);
+    return;
+  }
   case Intrinsic::invariant_start:
     // Discard region information.
     setValue(&I, DAG.getUNDEF(TLI.getPointerTy(DAG.getDataLayout())));
     return;
   case Intrinsic::invariant_end:
     // Discard region information.
     return;
   case Intrinsic::clear_cache:

llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp

#endif
   case ISD::EXTRACT_ELEMENT:            return "extract_element";
   case ISD::BUILD_PAIR:                 return "build_pair";
   case ISD::STACKSAVE:                  return "stacksave";
   case ISD::STACKRESTORE:               return "stackrestore";
   case ISD::TRAP:                       return "trap";
   case ISD::DEBUGTRAP:                  return "debugtrap";
   case ISD::LIFETIME_START:             return "lifetime.start";
   case ISD::LIFETIME_END:               return "lifetime.end";
+  case ISD::PSEUDO_PROBE:
+    return "pseudoprobe";
   case ISD::GC_TRANSITION_START:        return "gc_transition.start";
   case ISD::GC_TRANSITION_END:          return "gc_transition.end";
   case ISD::GET_DYNAMIC_AREA_OFFSET:    return "get.dynamic.area.offset";
   case ISD::FREEZE:                     return "freeze";
   case ISD::PREALLOCATED_SETUP:
     return "call_setup";
   case ISD::PREALLOCATED_ARG:
     return "call_alloc";

llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp

void SelectionDAGISel::SelectCodeCommon(SDNode *NodeToMatch,
   case ISD::TargetGlobalAddress:
   case ISD::TokenFactor:
   case ISD::CopyFromReg:
   case ISD::CopyToReg:
   case ISD::EH_LABEL:
   case ISD::ANNOTATION_LABEL:
   case ISD::LIFETIME_START:
   case ISD::LIFETIME_END:
+  case ISD::PSEUDO_PROBE:
     NodeToMatch->setNodeId(-1); // Mark selected.
     return;
   case ISD::AssertSext:
   case ISD::AssertZext:
   case ISD::AssertAlign:
     ReplaceUses(SDValue(NodeToMatch, 0), NodeToMatch->getOperand(0));
     CurDAG->RemoveDeadNode(NodeToMatch);
     return;

llvm/test/Transforms/SampleProfile/pseudo-probe.mir

+
+# REQUIRES: x86-registered-target
+# Ensure llc can read and parse MIR pseudo probe operations.
+# RUN: llc -O0 -mtriple x86_64-- -run-pass none %s -o - | FileCheck %s
+
+# CHECK: PSEUDO_PROBE 6699318081062747564, 1, 0, 0
+# CHECK: PSEUDO_PROBE 6699318081062747564, 3, 0, 0
+# CHECK: PSEUDO_PROBE 6699318081062747564, 4, 0, 0
+# CHECK: PSEUDO_PROBE 6699318081062747564, 2, 0, 0
+# CHECK: PSEUDO_PROBE 6699318081062747564, 4, 0, 0
+
+name:            foo
+body:             |
+  bb.0:
+    TEST32rr killed renamable $edi, renamable $edi, implicit-def $eflags
+    PSEUDO_PROBE 6699318081062747564, 1, 0, 0
+    JCC_1 %bb.1, 4, implicit $eflags
+  
+  bb.2:
+    PSEUDO_PROBE 6699318081062747564, 3, 0, 0
+    PSEUDO_PROBE 6699318081062747564, 4, 0, 0
+    RETQ
+  
+  bb.1:
+    PSEUDO_PROBE 6699318081062747564, 2, 0, 0
+    PSEUDO_PROBE 6699318081062747564, 4, 0, 0
+    RETQ
+
+...