Handle jump table change in machine-sink pass

Yonghong Song · Yonghong Song · commit e709a509df8c · 2025-05-07T19:16:48.000-07:00
The implementation is similar to func getJumpTableIndex() in X86InstrInfo.cpp.

For the following example:
    struct simple_ctx {
	int x;
	int y;
	int z;
    };

    int ret_user, ret_user2;
    void bar(void);

    int foo(struct simple_ctx *ctx, struct simple_ctx *ctx2)
    {
        switch (ctx-&gt;x) {
        case 1:
                ret_user = 8;
                break;
        case 6:
                ret_user = 3;
                break;
        case 2:
                ret_user = 4;
                break;
        case 31:
                ret_user = 5;
                break;
        default:
                ret_user = 19;
                break;
        }

	bar();

        switch (ctx2-&gt;x) {
        case 0:
                ret_user2 = 8;
                break;
        case 7:
                ret_user2 = 3;
                break;
        case 9:
                ret_user2 = 4;
                break;
        case 31:
                ret_user2 = 5;
                break;
        default:
                ret_user2 = 29;
                break;
        }

        return 0;
    }

Before machine-sink pass,
Jump Tables:
%jump-table.0: %bb.5 %bb.2 %bb.4 %bb.4 %bb.4 %bb.1 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.3
%jump-table.1: %bb.10 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.6 %bb.9 %bb.7 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.8

Machine-level IR:
bb.0.entry:
  successors: %bb.4(0x0ccccccb), %bb.11(0x73333335); %bb.4(10.00%), %bb.11(90.00%)
  liveins: $r1, $r2
  %3:gpr = COPY $r2
  %2:gpr = COPY $r1
  %4:gpr32 = MOV_ri_32 8
  %6:gpr32 = LDW32 %2:gpr, 0 :: (load (s32) from %ir.ctx, !tbaa !3)
  %7:gpr32 = ADD_ri_32 %6:gpr32(tied-def 0), -1
  %5:gpr = MOV_32_64 %7:gpr32
  JUGT_ri_32 %7:gpr32, 30, %bb.4

bb.11.entry:
; predecessors: %bb.0
  successors: %bb.5(0x1c71c71c), %bb.2(0x1c71c71c), %bb.4(0x0e38e38e), %bb.1(0x1c71c71c), %bb.3(0x1c71c71c); %bb.5(22.22%), %bb.2(22.22%), %bb.4(11.11%), %bb.1(22.22%), %bb.3(22.22%)

  %8:gpr = SLL_ri %5:gpr(tied-def 0), 3
  %9:gpr = LD_imm64 %jump-table.0
  %10:gpr = ADD_rr %9:gpr(tied-def 0), killed %8:gpr
  %11:gpr = LDD killed %10:gpr, 0 :: (load (s64) from jump-table)
  JX killed %11:gpr

bb.1.sw.bb1:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %14:gpr32 = MOV_ri_32 3
  JMP %bb.5

bb.2.sw.bb2:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %13:gpr32 = MOV_ri_32 4
  JMP %bb.5

bb.3.sw.bb3:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %12:gpr32 = MOV_ri_32 5
  JMP %bb.5

bb.4.sw.default:
; predecessors: %bb.11, %bb.0
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %15:gpr32 = MOV_ri_32 19

bb.5.sw.epilog:
...

After machine-sink pass:
Jump Tables:
%jump-table.0: %bb.13 %bb.2 %bb.4 %bb.4 %bb.4 %bb.1 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.4 %bb.3
%jump-table.1: %bb.14 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.6 %bb.9 %bb.7 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.9 %bb.8

Machine-level IR:
bb.0.entry:
  successors: %bb.4(0x0ccccccb), %bb.11(0x73333335); %bb.4(10.00%), %bb.11(90.00%)
  liveins: $r1, $r2
  %3:gpr = COPY $r2
  %2:gpr = COPY $r1
  %6:gpr32 = LDW32 %2:gpr, 0 :: (load (s32) from %ir.ctx, !tbaa !3)
  %7:gpr32 = ADD_ri_32 %6:gpr32(tied-def 0), -1
  JUGT_ri_32 %7:gpr32, 30, %bb.4

bb.11.entry:
; predecessors: %bb.0
  successors: %bb.13(0x1c71c71c), %bb.2(0x1c71c71c), %bb.4(0x0e38e38e), %bb.1(0x1c71c71c), %bb.3(0x1c71c71c); %bb.13(22.22%), %bb.2(22.22%), %bb.4(11.11%), %bb.1(22.22%), %bb.3(22.22%)

  %5:gpr = MOV_32_64 %7:gpr32
  %8:gpr = SLL_ri %5:gpr(tied-def 0), 3
  %9:gpr = LD_imm64 %jump-table.0
  %10:gpr = ADD_rr %9:gpr(tied-def 0), killed %8:gpr
  %11:gpr = LDD killed %10:gpr, 0 :: (load (s64) from jump-table)
  JX killed %11:gpr

bb.13:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %4:gpr32 = MOV_ri_32 8
  JMP %bb.5
bb.1.sw.bb1:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %14:gpr32 = MOV_ri_32 3
  JMP %bb.5

bb.2.sw.bb2:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %13:gpr32 = MOV_ri_32 4
  JMP %bb.5

bb.3.sw.bb3:
; predecessors: %bb.11
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %12:gpr32 = MOV_ri_32 5
  JMP %bb.5

bb.4.sw.default:
; predecessors: %bb.11, %bb.0
  successors: %bb.5(0x80000000); %bb.5(100.00%)

  %15:gpr32 = MOV_ri_32 19

bb.5.sw.epilog:

Before machine-sink pass, '%4:gpr32 = MOV_ri_32 8' is in the entry block so
there is no switch-branch for value 8. But machine-sink pass later removed
'%4:gpr32 = MOV_ri_32 8' and add back the switch-branch with value 8. Such
transformation requires adjust the jump table. This commit implemented
backend callback function getJumpTableIndex() so jump table can be properly
updated.
diff --git a/llvm/lib/Target/BPF/BPFInstrInfo.cpp b/llvm/lib/Target/BPF/BPFInstrInfo.cpp
@@ -263,3 +263,40 @@ unsigned BPFInstrInfo::removeBranch(MachineBasicBlock &MBB,
 
   return Count;
 }
+
+int BPFInstrInfo::getJumpTableIndex(const MachineInstr &MI) const {
+  // The pattern looks like:
+  // %0 = LD_imm64 %jump-table.0   ; load jump-table address
+  // %1 = ADD_rr %0, $another_reg  ; address + offset
+  // %2 = LDD %1, 0                ; load the actual label
+  // JX %2
+  const MachineFunction &MF = *MI.getParent()->getParent();
+  const MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  Register Reg = MI.getOperand(0).getReg();
+  if (!Reg.isVirtual())
+    return -1;
+  MachineInstr *Ldd = MRI.getUniqueVRegDef(Reg);
+  if (Ldd == nullptr || Ldd->getOpcode() != BPF::LDD)
+    return -1;
+
+  Reg = Ldd->getOperand(1).getReg();
+  if (!Reg.isVirtual())
+    return -1;
+  MachineInstr *Add = MRI.getUniqueVRegDef(Reg);
+  if (Add == nullptr || Add->getOpcode() != BPF::ADD_rr)
+    return -1;
+
+  Reg = Add->getOperand(1).getReg();
+  if (!Reg.isVirtual())
+    return -1;
+  MachineInstr *LDimm64 = MRI.getUniqueVRegDef(Reg);
+  if (LDimm64 == nullptr || LDimm64->getOpcode() != BPF::LD_imm64)
+    return -1;
+
+  const MachineOperand &MO = LDimm64->getOperand(1);
+  if (!MO.isJTI())
+    return -1;
+
+  return MO.getIndex();
+}
diff --git a/llvm/lib/Target/BPF/BPFInstrInfo.h b/llvm/lib/Target/BPF/BPFInstrInfo.h
@@ -58,6 +58,9 @@ class BPFInstrInfo : public BPFGenInstrInfo {
                         MachineBasicBlock *FBB, ArrayRef<MachineOperand> Cond,
                         const DebugLoc &DL,
                         int *BytesAdded = nullptr) const override;
+
+  int getJumpTableIndex(const MachineInstr &MI) const override;
+
 private:
   void expandMEMCPY(MachineBasicBlock::iterator) const;
 
