[RISCV] Generate 32 bits jumptable entries when code model is small

[• pcwang-thead]

Summary

The code can only address the whole RV32 address space or the lower 2 GiB
of the RV64 address space in small code model, so 32 bits entry is enough.
Cache hit ratio and code size have some improvements.

[asb]

LGTM. It doesn't make much practical difference, but I wonder if gating the getJumpTableEncoding case on RV64 (as well as your other preconditions) might make sense - as the default EK_BlockAddress lowering is equivalent in that case.

[jrtc27]

AArch64 does even better and has a PC-relative jump table lowering whose element size depends on how big your switch is, not the code model in use. Ideally this would be made target-independent, or at least copied to RISC-V. That'd let you have small entries even for medany RV64.

[• pcwang-thead]

In D116435#3225279, @jrtc27 wrote:

AArch64 does even better and has a PC-relative jump table lowering whose element size depends on how big your switch is, not the code model in use. Ideally this would be made target-independent, or at least copied to RISC-V. That'd let you have small entries even for medany RV64.

Thanks for your advice!

It seems that AArch64 has a pass called AArch64CompressJumpTables to do this, I may add it to RISCV later (or make it a target-independent pass).

[• pcwang-thead]

Address comment.

[• pcwang-thead]

In D116435#3226574, @pcwang-thead wrote:

In D116435#3225279, @jrtc27 wrote:

AArch64 does even better and has a PC-relative jump table lowering whose element size depends on how big your switch is, not the code model in use. Ideally this would be made target-independent, or at least copied to RISC-V. That'd let you have small entries even for medany RV64.

Thanks for your advice!

It seems that AArch64 has a pass called AArch64CompressJumpTables to do this, I may add it to RISCV later (or make it a target-independent pass).

It may take some time to do these changes, so I will land this patch first.

llvm/lib/Target/RISCV/RISCVISelLowering.h

public:
   MVT getContainerForFixedLengthVector(MVT VT) const;
 
   bool shouldRemoveExtendFromGSIndex(EVT VT) const override;
 
   bool isLegalElementTypeForRVV(Type *ScalarTy) const;
 
   bool shouldConvertFpToSat(unsigned Op, EVT FPVT, EVT VT) const override;
 
+  unsigned getJumpTableEncoding() const override;
+
+  const MCExpr *LowerCustomJumpTableEntry(const MachineJumpTableInfo *MJTI,
+                                          const MachineBasicBlock *MBB,
+                                          unsigned uid,
+                                          MCContext &Ctx) const override;
+
 private:
   /// RISCVCCAssignFn - This target-specific function extends the default
   /// CCValAssign with additional information used to lower RISC-V calling
   /// conventions.
   typedef bool RISCVCCAssignFn(const DataLayout &DL, RISCVABI::ABI,
                                unsigned ValNo, MVT ValVT, MVT LocVT,
                                CCValAssign::LocInfo LocInfo,
                                ISD::ArgFlagsTy ArgFlags, CCState &State,

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

 #include "RISCVSubtarget.h"
 #include "RISCVTargetMachine.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/DiagnosticPrinter.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicsRISCV.h"
 #include "llvm/IR/PatternMatch.h"
case MVT::f32:
     return Subtarget.hasStdExtF();
   case MVT::f64:
     return Subtarget.hasStdExtD();
   default:
     return false;
   }
 }
 
+unsigned RISCVTargetLowering::getJumpTableEncoding() const {
+  // If we are using the small code model, we can reduce size of jump table
+  // entry to 4 bytes.
+  if (!isPositionIndependent() &&
+      getTargetMachine().getCodeModel() == CodeModel::Small) {
+    return MachineJumpTableInfo::EK_Custom32;
+  }
+  return TargetLowering::getJumpTableEncoding();
+}
+
+const MCExpr *RISCVTargetLowering::LowerCustomJumpTableEntry(
+    const MachineJumpTableInfo *MJTI, const MachineBasicBlock *MBB,
+    unsigned uid, MCContext &Ctx) const {
+  assert(!isPositionIndependent() &&
+         getTargetMachine().getCodeModel() == CodeModel::Small);
+  return MCSymbolRefExpr::create(MBB->getSymbol(), Ctx);
+}
+
 bool RISCVTargetLowering::isFMAFasterThanFMulAndFAdd(const MachineFunction &MF,
                                                      EVT VT) const {
   VT = VT.getScalarType();
 
   if (!VT.isSimple())
     return false;
 
   switch (VT.getSimpleVT().SimpleTy) {

llvm/test/CodeGen/RISCV/jumptable.ll

 ;
 ; RV64I-SMALL-LABEL: above_threshold:
 ; RV64I-SMALL:       # %bb.0: # %entry
 ; RV64I-SMALL-NEXT:    sext.w a0, a0
 ; RV64I-SMALL-NEXT:    addi a0, a0, -1
 ; RV64I-SMALL-NEXT:    li a2, 5
 ; RV64I-SMALL-NEXT:    bltu a2, a0, .LBB1_9
 ; RV64I-SMALL-NEXT:  # %bb.1: # %entry
-; RV64I-SMALL-NEXT:    slli a0, a0, 3
+; RV64I-SMALL-NEXT:    slli a0, a0, 2
 ; RV64I-SMALL-NEXT:    lui a2, %hi(.LJTI1_0)
 ; RV64I-SMALL-NEXT:    addi a2, a2, %lo(.LJTI1_0)
 ; RV64I-SMALL-NEXT:    add a0, a0, a2
-; RV64I-SMALL-NEXT:    ld a0, 0(a0)
+; RV64I-SMALL-NEXT:    lw a0, 0(a0)
 ; RV64I-SMALL-NEXT:    jr a0
 ; RV64I-SMALL-NEXT:  .LBB1_2: # %bb1
 ; RV64I-SMALL-NEXT:    li a0, 4
 ; RV64I-SMALL-NEXT:    j .LBB1_8
 ; RV64I-SMALL-NEXT:  .LBB1_3: # %bb2
 ; RV64I-SMALL-NEXT:    li a0, 3
 ; RV64I-SMALL-NEXT:    j .LBB1_8
 ; RV64I-SMALL-NEXT:  .LBB1_4: # %bb3