[RISCV][GISel] Instruction selection for G_JUMP_TABLE and G_BRJT.

llvm/lib/Target/RISCV/GISel/RISCVInstructionSelector.cpp

@@ -20,6 +20,7 @@
 #include "llvm/CodeGen/GlobalISel/InstructionSelector.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/MachineIRBuilder.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/IR/IntrinsicsRISCV.h"
 #include "llvm/Support/Debug.h"
 
@@ -64,6 +65,8 @@ class RISCVInstructionSelector : public InstructionSelector {
   bool materializeImm(Register Reg, int64_t Imm, MachineIRBuilder &MIB) const;
   bool selectGlobalValue(MachineInstr &MI, MachineIRBuilder &MIB,
                          MachineRegisterInfo &MRI) const;
+  bool selectJumpTable(MachineInstr &MI, MachineIRBuilder &MIB,
+                       MachineRegisterInfo &MRI) const;
   bool selectSExtInreg(MachineInstr &MI, MachineIRBuilder &MIB) const;
   bool selectSelect(MachineInstr &MI, MachineIRBuilder &MIB,
                     MachineRegisterInfo &MRI) const;
@@ -483,6 +486,8 @@ bool RISCVInstructionSelector::select(MachineInstr &MI) {
   }
   case TargetOpcode::G_GLOBAL_VALUE:
     return selectGlobalValue(MI, MIB, MRI);
+  case TargetOpcode::G_JUMP_TABLE:
+    return selectJumpTable(MI, MIB, MRI);
   case TargetOpcode::G_BRCOND: {
     Register LHS, RHS;
     RISCVCC::CondCode CC;
@@ -493,6 +498,58 @@ bool RISCVInstructionSelector::select(MachineInstr &MI) {
     MI.eraseFromParent();
     return constrainSelectedInstRegOperands(*Bcc, TII, TRI, RBI);
   }
+  case TargetOpcode::G_BRJT: {
+    // FIXME: Move to legalization?
+    const MachineJumpTableInfo *MJTI = MF.getJumpTableInfo();
+    unsigned EntrySize = MJTI->getEntrySize(MF.getDataLayout());
+    assert((EntrySize == 4 || (Subtarget->is64Bit() && EntrySize == 8)) &&
+           "Unsupported jump-table entry size");
+    assert(
+        (MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
+         MJTI->getEntryKind() == MachineJumpTableInfo::EK_Custom32 ||
+         MJTI->getEntryKind() == MachineJumpTableInfo::EK_BlockAddress) &&
+        "Unexpected jump-table entry kind");
+
+    auto SLL =
+        MIB.buildInstr(RISCV::SLLI, {&RISCV::GPRRegClass}, {MI.getOperand(2)})
+            .addImm(Log2_32(EntrySize));
+    if (!SLL.constrainAllUses(TII, TRI, RBI))
+      return false;
+
+    // TODO: Use SHXADD. Moving to legalization would fix this automatically.
+    auto ADD = MIB.buildInstr(RISCV::ADD, {&RISCV::GPRRegClass},
+                              {MI.getOperand(0), SLL.getReg(0)});
+    if (!ADD.constrainAllUses(TII, TRI, RBI))
+      return false;
+
+    unsigned LdOpc = EntrySize == 8 ? RISCV::LD : RISCV::LW;
+    auto Dest =
+        MIB.buildInstr(LdOpc, {&RISCV::GPRRegClass}, {ADD.getReg(0)})
+            .addImm(0)
+            .addMemOperand(MF.getMachineMemOperand(
+                MachinePointerInfo::getJumpTable(MF), MachineMemOperand::MOLoad,
+                EntrySize, Align(MJTI->getEntryAlignment(MF.getDataLayout()))));
+    if (!Dest.constrainAllUses(TII, TRI, RBI))
+      return false;
+
+    // If the Kind is EK_LabelDifference32, the table stores an offset from
+    // the location of the table. Add the table address to get an absolute
+    // address.
+    if (MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32) {
+      Dest = MIB.buildInstr(RISCV::ADD, {&RISCV::GPRRegClass},
+                            {Dest.getReg(0), MI.getOperand(0)});
+      if (!Dest.constrainAllUses(TII, TRI, RBI))
+        return false;
+    }
+
+    auto Branch =
+        MIB.buildInstr(RISCV::PseudoBRIND, {}, {Dest.getReg(0)}).addImm(0);
+    if (!Branch.constrainAllUses(TII, TRI, RBI))
+      return false;
+
+    MI.eraseFromParent();
+    return true;
+  }
   case TargetOpcode::G_SEXT_INREG:
     return selectSExtInreg(MI, MIB);
   case TargetOpcode::G_FRAME_INDEX: {
@@ -784,6 +841,97 @@ bool RISCVInstructionSelector::selectGlobalValue(
   return false;
 }
 
+// FIXME: This is very similar to selectGlobalValue. Merge somehow?
+bool RISCVInstructionSelector::selectJumpTable(MachineInstr &MI,
+                                               MachineIRBuilder &MIB,
+                                               MachineRegisterInfo &MRI) const {
+  assert(MI.getOpcode() == TargetOpcode::G_JUMP_TABLE &&
+         "Expected G_JUMP_TABLE");
+
+  int Idx = MI.getOperand(1).getIndex();
+
+  Register DefReg = MI.getOperand(0).getReg();
+  const LLT DefTy = MRI.getType(DefReg);
+  MachineInstr *Result = nullptr;
+
+  // When HWASAN is used and tagging of global variables is enabled
+  // they should be accessed via the GOT, since the tagged address of a global
+  // is incompatible with existing code models. This also applies to non-pic
+  // mode.
+  if (TM.isPositionIndependent() || Subtarget->allowTaggedGlobals()) {
+    if (!Subtarget->allowTaggedGlobals()) {
+      // Use PC-relative addressing to access the symbol. This generates the
+      // pattern (PseudoLLA sym), which expands to (addi (auipc %pcrel_hi(sym))
+      // %pcrel_lo(auipc)).
+      Result =
+          MIB.buildInstr(RISCV::PseudoLLA, {DefReg}, {}).addJumpTableIndex(Idx);
+    } else {
+      // Use PC-relative addressing to access the GOT for this symbol, then
+      // load the address from the GOT. This generates the pattern (PseudoLGA
+      // sym), which expands to (ld (addi (auipc %got_pcrel_hi(sym))
+      // %pcrel_lo(auipc))).
+      MachineFunction &MF = *MI.getParent()->getParent();
+      MachineMemOperand *MemOp = MF.getMachineMemOperand(
+          MachinePointerInfo::getGOT(MF),
+          MachineMemOperand::MOLoad | MachineMemOperand::MODereferenceable |
+              MachineMemOperand::MOInvariant,
+          DefTy, Align(DefTy.getSizeInBits() / 8));
+
+      Result = MIB.buildInstr(RISCV::PseudoLGA, {DefReg}, {})
+                   .addJumpTableIndex(Idx)
+                   .addMemOperand(MemOp);
+    }
+
+    if (!constrainSelectedInstRegOperands(*Result, TII, TRI, RBI))
+      return false;
+
+    MI.eraseFromParent();
+    return true;
+  }
+
+  switch (TM.getCodeModel()) {
+  default: {
+    reportGISelFailure(const_cast<MachineFunction &>(*MF), *TPC, *MORE,
+                       getName(), "Unsupported code model for lowering", MI);
+    return false;
+  }
+  case CodeModel::Small: {
+    // Must lie within a single 2 GiB address range and must lie between
+    // absolute addresses -2 GiB and +2 GiB. This generates the pattern (addi
+    // (lui %hi(sym)) %lo(sym)).
+    Register AddrHiDest = MRI.createVirtualRegister(&RISCV::GPRRegClass);
+    MachineInstr *AddrHi = MIB.buildInstr(RISCV::LUI, {AddrHiDest}, {})
+                               .addJumpTableIndex(Idx, RISCVII::MO_HI);
+
+    if (!constrainSelectedInstRegOperands(*AddrHi, TII, TRI, RBI))
+      return false;
+
+    Result = MIB.buildInstr(RISCV::ADDI, {DefReg}, {AddrHiDest})
+                 .addJumpTableIndex(Idx, RISCVII::MO_LO);
+
+    if (!constrainSelectedInstRegOperands(*Result, TII, TRI, RBI))
+      return false;
+
+    MI.eraseFromParent();
+    return true;
+  }
+  case CodeModel::Medium: {
+    // Generate a sequence for accessing addresses within any 2GiB range
+    // within the address space. This generates the pattern (PseudoLLA sym),
+    // which expands to (addi (auipc %pcrel_hi(sym)) %pcrel_lo(auipc)).
+    Result =
+        MIB.buildInstr(RISCV::PseudoLLA, {DefReg}, {}).addJumpTableIndex(Idx);
+
+    if (!constrainSelectedInstRegOperands(*Result, TII, TRI, RBI))
+      return false;
+
+    MI.eraseFromParent();
+    return true;
+  }
+  }
+  return false;
+}
+
 bool RISCVInstructionSelector::selectSExtInreg(MachineInstr &MI,
                                                MachineIRBuilder &MIB) const {
   if (!STI.isRV64())

llvm/test/CodeGen/RISCV/GlobalISel/instruction-select/jump-table-brjt-medium-rv64.mir

@@ -0,0 +1,160 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=riscv64 -run-pass=instruction-select %s -o - \
+# RUN:   -code-model=medium | FileCheck %s
+
+--- |
+  define i32 @jt_test(i32 signext %in) {
+  entry:
+    %0 = sext i32 %in to i64
+    switch i64 %0, label %default [
+      i64 1, label %bb1
+      i64 2, label %bb2
+      i64 3, label %bb3
+      i64 4, label %bb4
+      i64 5, label %bb5
+      i64 6, label %bb6
+    ]
+
+  bb1:                                              ; preds = %entry
+    ret i32 4
+
+  bb2:                                              ; preds = %entry
+    ret i32 3
+
+  bb3:                                              ; preds = %entry
+    ret i32 2
+
+  bb4:                                              ; preds = %entry
+    ret i32 1
+
+  bb5:                                              ; preds = %entry
+    ret i32 100
+
+  bb6:                                              ; preds = %entry
+    ret i32 200
+
+  default:                                          ; preds = %entry
+    ret i32 1000
+  }
+
+...
+---
+name:            jt_test
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.2', '%bb.3', '%bb.4', '%bb.5', '%bb.6', '%bb.7' ]
+body:             |
+  ; CHECK-LABEL: name: jt_test
+  ; CHECK: bb.0.entry:
+  ; CHECK-NEXT:   successors: %bb.8(0x40000000), %bb.1(0x40000000)
+  ; CHECK-NEXT:   liveins: $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:gpr = COPY $x10
+  ; CHECK-NEXT:   [[ADDI:%[0-9]+]]:gpr = ADDI $x0, 5
+  ; CHECK-NEXT:   [[ADDIW:%[0-9]+]]:gpr = ADDIW [[COPY]], 0
+  ; CHECK-NEXT:   [[ADDI1:%[0-9]+]]:gpr = ADDI [[ADDIW]], -1
+  ; CHECK-NEXT:   BLTU [[ADDI]], [[ADDI1]], %bb.8
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1.entry:
+  ; CHECK-NEXT:   successors: %bb.2(0x15555555), %bb.3(0x15555555), %bb.4(0x15555555), %bb.5(0x15555555), %bb.6(0x15555555), %bb.7(0x15555555)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[PseudoLLA:%[0-9]+]]:gpr = PseudoLLA %jump-table.0
+  ; CHECK-NEXT:   [[SLLI:%[0-9]+]]:gpr = SLLI [[ADDI1]], 3
+  ; CHECK-NEXT:   [[ADD:%[0-9]+]]:gpr = ADD [[PseudoLLA]], [[SLLI]]
+  ; CHECK-NEXT:   [[LD:%[0-9]+]]:gprjalr = LD [[ADD]], 0 :: (load (s64) from jump-table)
+  ; CHECK-NEXT:   PseudoBRIND [[LD]], 0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2.bb1:
+  ; CHECK-NEXT:   [[ADDI2:%[0-9]+]]:gpr = ADDI $x0, 4
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI2]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3.bb2:
+  ; CHECK-NEXT:   [[ADDI3:%[0-9]+]]:gpr = ADDI $x0, 3
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI3]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.4.bb3:
+  ; CHECK-NEXT:   [[ADDI4:%[0-9]+]]:gpr = ADDI $x0, 2
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI4]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.5.bb4:
+  ; CHECK-NEXT:   [[ADDI5:%[0-9]+]]:gpr = ADDI $x0, 1
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI5]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.6.bb5:
+  ; CHECK-NEXT:   [[ADDI6:%[0-9]+]]:gpr = ADDI $x0, 100
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI6]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.7.bb6:
+  ; CHECK-NEXT:   [[ADDI7:%[0-9]+]]:gpr = ADDI $x0, 200
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI7]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.8.default:
+  ; CHECK-NEXT:   [[ADDI8:%[0-9]+]]:gpr = ADDI $x0, 1000
+  ; CHECK-NEXT:   $x10 = COPY [[ADDI8]]
+  ; CHECK-NEXT:   PseudoRET implicit $x10
+  bb.1.entry:
+    successors: %bb.8, %bb.9
+    liveins: $x10
+
+    %1:gprb(s64) = COPY $x10
+    %2:gprb(s64) = G_ASSERT_SEXT %1, 32
+    %7:gprb(s64) = G_CONSTANT i64 5
+    %3:gprb(s64) = G_SEXT_INREG %2, 32
+    %4:gprb(s64) = G_CONSTANT i64 1
+    %5:gprb(s64) = G_SUB %3, %4
+    %26:gprb(s64) = G_ICMP intpred(ugt), %5(s64), %7
+    G_BRCOND %26(s64), %bb.8
+
+  bb.9.entry:
+    successors: %bb.2, %bb.3, %bb.4, %bb.5, %bb.6, %bb.7
+
+    %10:gprb(p0) = G_JUMP_TABLE %jump-table.0
+    G_BRJT %10(p0), %jump-table.0, %5(s64)
+
+  bb.2.bb1:
+    %22:gprb(s64) = G_CONSTANT i64 4
+    $x10 = COPY %22(s64)
+    PseudoRET implicit $x10
+
+  bb.3.bb2:
+    %20:gprb(s64) = G_CONSTANT i64 3
+    $x10 = COPY %20(s64)
+    PseudoRET implicit $x10
+
+  bb.4.bb3:
+    %18:gprb(s64) = G_CONSTANT i64 2
+    $x10 = COPY %18(s64)
+    PseudoRET implicit $x10
+
+  bb.5.bb4:
+    %16:gprb(s64) = G_CONSTANT i64 1
+    $x10 = COPY %16(s64)
+    PseudoRET implicit $x10
+
+  bb.6.bb5:
+    %14:gprb(s64) = G_CONSTANT i64 100
+    $x10 = COPY %14(s64)
+    PseudoRET implicit $x10
+
+  bb.7.bb6:
+    %12:gprb(s64) = G_CONSTANT i64 200
+    $x10 = COPY %12(s64)
+    PseudoRET implicit $x10
+
+  bb.8.default:
+    %24:gprb(s64) = G_CONSTANT i64 1000
+    $x10 = COPY %24(s64)
+    PseudoRET implicit $x10
+
+...