[WIP][BPF] Add sext load instructions

[yonghong-song]

Summary

Add sign-extension load instructions.

For kernel bpf selftest pyperf180, the number of
instructions without this patch is 91472. With this
patch, the number of instructions is 80652. Most improvements
come from

r1 = *(u32 *)(r8 + 12) 
r1 <<= 32
r1 s>>= 32

to

r1 = *(s32 *)(r8 + 12)

Such a new instruction not only improves bpf code, but
can also improve jitted code since most architecture
has explicit sign extension insntructions.

To-Do list:

• put the new functionality into e.g, -mcpu=v4
• implement the linux kernel part to support the new insn

[ast]

llvm/lib/Target/BPF/BPFInstrInfo.td

469

I believe x86 and arm64 have instructions that sign extends one register into another,
but not insns that sign extend during the load.
Such insns can be composed better to in case s32 came as a return value from a call.
Have you considered going that route?

[anakryiko]

llvm/lib/Target/BPF/BPFInstrInfo.td

469

one of the motivations for this instruction was to allow libbpf to automatically change load instruction to a proper size of the field (https://github.com/libbpf/libbpf/blob/master/src/relo_core.c#L930-L940). This works currently only for unsigned fields. If we add some instruction which would have to be called after load, that will not help in this case. So maybe let's add both? Register-to-register would help with calls, load with sign extension will help CO-RE and generally signed loads?

[anakryiko]

llvm/lib/Target/BPF/BPFInstrInfo.td

469

for call optimization, maybe having a family of calls that would perform desired size adjustment (32-bit to 64-bit) and optionally sign extension is another improvement?

E.g., what we have as "call" today would be basically 64-bit unsigned call, but we can have call{8,16,32,64}{s,u}, where call16s means that callee function is returning short in r0 and we request that r0 is sign-extended so it can be used as s64.

Just a thought.

[ast]

llvm/lib/Target/BPF/BPFInstrInfo.td

469

calls don't have such attrs. that's no go.
I prefer to avoid new loads, since they don't map to insns directly.
libbpf should be able to insert an insn. It can adjust offsets already. Not that hard.

[anakryiko]

llvm/lib/Target/BPF/BPFInstrInfo.td

469

Understood about call, np. And while you are technically correct about inserting instructions, that's a major PITA, so I'm not so keen about that.

But then re: x86 not having sign-extending loads. Isn't movsx exactly that (https://www.felixcloutier.com/x86/movsx:movsxd). Not an assembly expert, but seems like source can be a memory location, so that's exactly 1:1 mapping?

[anakryiko]

llvm/lib/Target/BPF/BPFInstrInfo.td

469

arm64 seems to have something like this as well? E.g., LDRSB (https://developer.arm.com/documentation/102374/0100/Loads-and-stores---zero-and-sign-extension)

[yonghong-song]

All major architectures support sign extension load (x64, arm64, sparc, webassembly, lanai, etc.)

Yes, x64/arm64 support register->register sign extension (https://www.felixcloutier.com/x86/movsx:movsxd, https://developer.arm.com/documentation/dui0801/g/A64-General-Instructions/SXTW) let me take a look.

[yonghong-song]

I hacked the kernel x64 jit part, with the following jit for sext load, it seems working okay from jit perspective.

+/* LDX: dst_reg = *(s8*)(src_reg + off) */
+static void emit_ldx_sext(u8 **pprog, u32 size, u32 dst_reg, u32 src_reg, int off)
+{
+       u8 *prog = *pprog;
+
+       switch (size) {
+       case BPF_B:
+               /* Emit 'movsx rax, byte ptr [rax + off]' */
+               EMIT3(add_2mod(0x48, src_reg, dst_reg), 0x0F, 0xBE);
+               break;
+       case BPF_H:
+               /* Emit 'movsx rax, word ptr [rax + off]' */
+               EMIT3(add_2mod(0x48, src_reg, dst_reg), 0x0F, 0xBF);
+               break;
+       case BPF_W:
+               /* Emit 'movsx rax, dword ptr [rax+0x14]' */
+               EMIT2(add_2mod(0x48, src_reg, dst_reg), 0x63);
+               break;
+       }
+       emit_insn_suffix(&prog, src_reg, dst_reg, off);
+       *pprog = prog;
+}

My hacking is incomplete so I only showed the partial x86 jit in the above.

llvm/lib/Target/BPF/BPFISelLowering.cpp

BPFTargetLowering::BPFTargetLowering(const TargetMachine &TM,
   setOperationAction(ISD::SIGN_EXTEND_INREG, MVT::i16, Expand);
   setOperationAction(ISD::SIGN_EXTEND_INREG, MVT::i32, Expand);
 
   // Extended load operations for i1 types must be promoted
   for (MVT VT : MVT::integer_valuetypes()) {
     setLoadExtAction(ISD::EXTLOAD, VT, MVT::i1, Promote);
     setLoadExtAction(ISD::ZEXTLOAD, VT, MVT::i1, Promote);
     setLoadExtAction(ISD::SEXTLOAD, VT, MVT::i1, Promote);
-
-    setLoadExtAction(ISD::SEXTLOAD, VT, MVT::i8, Expand);
-    setLoadExtAction(ISD::SEXTLOAD, VT, MVT::i16, Expand);
-    setLoadExtAction(ISD::SEXTLOAD, VT, MVT::i32, Expand);
   }
 
   setBooleanContents(ZeroOrOneBooleanContent);
 
   // Function alignments
   setMinFunctionAlignment(Align(8));
   setPrefFunctionAlignment(Align(8));
 

llvm/lib/Target/BPF/BPFInstrFormats.td

 class BPFModeModifer<bits<3> val> {
   bits<3> Value = val;
 }
 
 def BPF_IMM  : BPFModeModifer<0x0>;
 def BPF_ABS  : BPFModeModifer<0x1>;
 def BPF_IND  : BPFModeModifer<0x2>;
 def BPF_MEM  : BPFModeModifer<0x3>;
+def BPF_MEM_SEXT : BPFModeModifer<0x4>;
 def BPF_ATOMIC : BPFModeModifer<0x6>;
 
 class BPFAtomicFlag<bits<4> val> {
   bits<4> Value = val;
 }
 
 def BPF_FETCH : BPFAtomicFlag<0x1>;
 

llvm/lib/Target/BPF/BPFInstrInfo.td

 let Predicates = [BPFNoALU32] in {
   def STW : STOREi64<BPF_W, "u32", truncstorei32>;
   def STH : STOREi64<BPF_H, "u16", truncstorei16>;
   def STB : STOREi64<BPF_B, "u8", truncstorei8>;
 }
 def STD : STOREi64<BPF_DW, "u64", store>;
 
 // LOAD instructions
-class LOAD<BPFWidthModifer SizeOp, string OpcodeStr, list<dag> Pattern>
-    : TYPE_LD_ST<BPF_MEM.Value, SizeOp.Value,
+class LOAD<BPFWidthModifer SizeOp, BPFModeModifer ModOp, string OpcodeStr, list<dag> Pattern>
+    : TYPE_LD_ST<ModOp.Value, SizeOp.Value,
                  (outs GPR:$dst),
                  (ins MEMri:$addr),
                  "$dst = *("#OpcodeStr#" *)($addr)",
                  Pattern> {
   bits<4> dst;
   bits<20> addr;
 
   let Inst{51-48} = dst;
   let Inst{55-52} = addr{19-16};
   let Inst{47-32} = addr{15-0};
   let BPFClass = BPF_LDX;
 }
 
-class LOADi64<BPFWidthModifer SizeOp, string OpcodeStr, PatFrag OpNode>
-    : LOAD<SizeOp, OpcodeStr, [(set i64:$dst, (OpNode ADDRri:$addr))]>;
+class LOADi64<BPFWidthModifer SizeOp, BPFModeModifer ModOp, string OpcodeStr, PatFrag OpNode>
+    : LOAD<SizeOp, ModOp, OpcodeStr, [(set i64:$dst, (OpNode ADDRri:$addr))]>;
 
 let isCodeGenOnly = 1 in {
   def CORE_MEM : TYPE_LD_ST<BPF_MEM.Value, BPF_W.Value,
                             (outs GPR:$dst),
                             (ins u64imm:$opcode, GPR:$src, u64imm:$offset),
                             "$dst = core_mem($opcode, $src, $offset)",
                             []>;
   def CORE_ALU32_MEM : TYPE_LD_ST<BPF_MEM.Value, BPF_W.Value,
                                   (outs GPR32:$dst),
                                   (ins u64imm:$opcode, GPR:$src, u64imm:$offset),
                                   "$dst = core_alu32_mem($opcode, $src, $offset)",
                                   []>;
   let Constraints = "$dst = $src" in {
     def CORE_SHIFT : ALU_RR<BPF_ALU64, BPF_LSH,
                              (outs GPR:$dst),
                              (ins u64imm:$opcode, GPR:$src, u64imm:$offset),
                              "$dst = core_shift($opcode, $src, $offset)",
                              []>;
   }
 }
 
 let Predicates = [BPFNoALU32] in {
-  def LDW : LOADi64<BPF_W, "u32", zextloadi32>;
-  def LDH : LOADi64<BPF_H, "u16", zextloadi16>;
-  def LDB : LOADi64<BPF_B, "u8", zextloadi8>;
+  def LDW : LOADi64<BPF_W, BPF_MEM, "u32", zextloadi32>;
+  def LDH : LOADi64<BPF_H, BPF_MEM, "u16", zextloadi16>;
+  def LDB : LOADi64<BPF_B, BPF_MEM, "u8", zextloadi8>;
+
+  def LDWS : LOADi64<BPF_W, BPF_MEM_SEXT, "s32", sextloadi32>;
+  def LDHS : LOADi64<BPF_H, BPF_MEM_SEXT, "s16", sextloadi16>;
+  def LDBS : LOADi64<BPF_B, BPF_MEM_SEXT, "s8",  sextloadi8>;
 }
 
-def LDD : LOADi64<BPF_DW, "u64", load>;
+def LDD : LOADi64<BPF_DW, BPF_MEM, "u64", load>;
 
 class BRANCH<BPFJumpOp Opc, string OpcodeStr, list<dag> Pattern>
     : TYPE_ALU_JMP<Opc.Value, BPF_K.Value,
                    (outs),
                    (ins brtarget:$BrDst),
                    !strconcat(OpcodeStr, " $BrDst"),
                    Pattern> {
   bits<16> BrDst;
class STOREi32<BPFWidthModifer Opc, string OpcodeStr, PatFrag OpNode>
     : STORE32<Opc, OpcodeStr, [(OpNode i32:$src, ADDRri:$addr)]>;
 
 let Predicates = [BPFHasALU32], DecoderNamespace = "BPFALU32" in {
   def STW32 : STOREi32<BPF_W, "u32", store>;
   def STH32 : STOREi32<BPF_H, "u16", truncstorei16>;
   def STB32 : STOREi32<BPF_B, "u8", truncstorei8>;
 }
 
-class LOAD32<BPFWidthModifer SizeOp, string OpcodeStr, list<dag> Pattern>
-    : TYPE_LD_ST<BPF_MEM.Value, SizeOp.Value,
+class LOAD32<BPFWidthModifer SizeOp, BPFModeModifer ModOp, string OpcodeStr, list<dag> Pattern>
+    : TYPE_LD_ST<ModOp.Value, SizeOp.Value,
                 (outs GPR32:$dst),
                 (ins MEMri:$addr),
                 "$dst = *("#OpcodeStr#" *)($addr)",
                 Pattern> {
   bits<4> dst;
   bits<20> addr;
 
   let Inst{51-48} = dst;
   let Inst{55-52} = addr{19-16};
   let Inst{47-32} = addr{15-0};
   let BPFClass = BPF_LDX;
 }
 
-class LOADi32<BPFWidthModifer SizeOp, string OpcodeStr, PatFrag OpNode>
-    : LOAD32<SizeOp, OpcodeStr, [(set i32:$dst, (OpNode ADDRri:$addr))]>;
+class LOADi32<BPFWidthModifer SizeOp, BPFModeModifer ModOp, string OpcodeStr, PatFrag OpNode>
+    : LOAD32<SizeOp, ModOp, OpcodeStr, [(set i32:$dst, (OpNode ADDRri:$addr))]>;
 
 let Predicates = [BPFHasALU32], DecoderNamespace = "BPFALU32" in {
-  def LDW32 : LOADi32<BPF_W, "u32", load>;
-  def LDH32 : LOADi32<BPF_H, "u16", zextloadi16>;
-  def LDB32 : LOADi32<BPF_B, "u8", zextloadi8>;
+  def LDW32 : LOADi32<BPF_W, BPF_MEM, "u32", load>;
+  def LDH32 : LOADi32<BPF_H, BPF_MEM, "u16", zextloadi16>;
+  def LDB32 : LOADi32<BPF_B, BPF_MEM, "u8", zextloadi8>;
+
+  def LDH32S : LOADi32<BPF_H, BPF_MEM_SEXT, "s16", sextloadi16>;
+  def LDB32S : LOADi32<BPF_B, BPF_MEM_SEXT, "s8", sextloadi8>;
 }
 
 let Predicates = [BPFHasALU32] in {
   def : Pat<(truncstorei8 GPR:$src, ADDRri:$dst),
             (STB32 (EXTRACT_SUBREG GPR:$src, sub_32), ADDRri:$dst)>;
   def : Pat<(truncstorei16 GPR:$src, ADDRri:$dst),
             (STH32 (EXTRACT_SUBREG GPR:$src, sub_32), ADDRri:$dst)>;
   def : Pat<(truncstorei32 GPR:$src, ADDRri:$dst),
             (STW32 (EXTRACT_SUBREG GPR:$src, sub_32), ADDRri:$dst)>;
   def : Pat<(i32 (extloadi8 ADDRri:$src)), (i32 (LDB32 ADDRri:$src))>;
   def : Pat<(i32 (extloadi16 ADDRri:$src)), (i32 (LDH32 ADDRri:$src))>;
+  def : Pat<(i32 (sextloadi8 ADDRri:$src)), (i32 (LDB32S ADDRri:$src))>;
+  def : Pat<(i32 (sextloadi16 ADDRri:$src)), (i32 (LDH32S ADDRri:$src))>;
   def : Pat<(i64 (zextloadi8  ADDRri:$src)),
             (SUBREG_TO_REG (i64 0), (LDB32 ADDRri:$src), sub_32)>;
   def : Pat<(i64 (zextloadi16 ADDRri:$src)),
             (SUBREG_TO_REG (i64 0), (LDH32 ADDRri:$src), sub_32)>;
   def : Pat<(i64 (zextloadi32 ADDRri:$src)),
             (SUBREG_TO_REG (i64 0), (LDW32 ADDRri:$src), sub_32)>;
   def : Pat<(i64 (extloadi8  ADDRri:$src)),
             (SUBREG_TO_REG (i64 0), (LDB32 ADDRri:$src), sub_32)>;
   def : Pat<(i64 (extloadi16 ADDRri:$src)),
             (SUBREG_TO_REG (i64 0), (LDH32 ADDRri:$src), sub_32)>;
   def : Pat<(i64 (extloadi32 ADDRri:$src)),
             (SUBREG_TO_REG (i64 0), (LDW32 ADDRri:$src), sub_32)>;
+  def : Pat<(i64 (sextloadi8 ADDRri:$src)), (i64 (LDBS ADDRri:$src))>;
+  def : Pat<(i64 (sextloadi16 ADDRri:$src)), (i64 (LDHS ADDRri:$src))>;
+  def : Pat<(i64 (sextloadi32 ADDRri:$src)), (i64 (LDWS ADDRri:$src))>;
 }
 
 let usesCustomInserter = 1, isCodeGenOnly = 1 in {
     def MEMCPY : Pseudo<
       (outs),
       (ins GPR:$dst, GPR:$src, i64imm:$len, i64imm:$align, variable_ops),
       "#memcpy dst: $dst, src: $src, len: $len, align: $align",
       [(BPFmemcpy GPR:$dst, GPR:$src, imm:$len, imm:$align)]>;
 }

llvm/lib/Target/BPF/Disassembler/BPFDisassembler.cpp

enum BPF_SIZE {
     BPF_DW = 0x3
   };
 
   enum BPF_MODE {
     BPF_IMM = 0x0,
     BPF_ABS = 0x1,
     BPF_IND = 0x2,
     BPF_MEM = 0x3,
+    BPF_MEM_SEXT = 0x4,
     BPF_LEN = 0x4,
     BPF_MSH = 0x5,
     BPF_ATOMIC = 0x6
   };
 
   BPFDisassembler(const MCSubtargetInfo &STI, MCContext &Ctx)
       : MCDisassembler(STI, Ctx) {}
   ~BPFDisassembler() override = default;

llvm/test/CodeGen/BPF/sext.ll

+; RUN: llc -O2 -march=bpfel -mcpu=v1 < %s | FileCheck -check-prefixes=CHECK -check-prefixes=CHECKV1 %s
+; RUN: llc -O2 -march=bpfel -mcpu=v3 < %s | FileCheck -check-prefixes=CHECK -check-prefixes=CHECKV3 %s
+; Source:
+;  int f1(char *p) {
+;    return *p;
+;  }
+;  int f2(short *p) {
+;    return *p;
+;  }
+;  int f3(int *p) {
+;    return *p;
+;  }
+;  long f4(char *p) {
+;    return *p;
+;  }
+;  long f5(short *p) {
+;    return *p;
+;  }
+;  long f6(int *p) {
+;    return *p;
+;  }
+;  long f7(long *p) {
+;    return *p;
+;  }
+; Compilation flags:
+;   clang -target bpf -O2 -S -emit-llvm -Xclang -disable-llvm-passes t.c
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i32 @f1(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i8, ptr %p, align 1, !tbaa !3
+  %conv = sext i8 %0 to i32
+; CHECKV1:  r0 = *(s8 *)(r1 + 0)
+; CHECKV3:  w0 = *(s8 *)(r1 + 0)
+  ret i32 %conv
+}
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i32 @f2(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i16, ptr %p, align 2, !tbaa !6
+  %conv = sext i16 %0 to i32
+; CHECKV1:  r0 = *(s16 *)(r1 + 0)
+; CHECKV3:  w0 = *(s16 *)(r1 + 0)
+  ret i32 %conv
+}
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i32 @f3(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i32, ptr %p, align 4, !tbaa !8
+; CHECKV1:  r0 = *(u32 *)(r1 + 0)
+; CHECKV3:  w0 = *(u32 *)(r1 + 0)
+  ret i32 %0
+}
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i64 @f4(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i8, ptr %p, align 1, !tbaa !3
+  %conv = sext i8 %0 to i64
+  ret i64 %conv
+; CHECK:      r0 = *(s8 *)(r1 + 0)
+; CHECK-NEXT: exit
+}
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i64 @f5(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i16, ptr %p, align 2, !tbaa !6
+  %conv = sext i16 %0 to i64
+  ret i64 %conv
+; CHECK:      r0 = *(s16 *)(r1 + 0)
+; CHECK-NEXT: exit
+}
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i64 @f6(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i32, ptr %p, align 4, !tbaa !8
+  %conv = sext i32 %0 to i64
+  ret i64 %conv
+; CHECK:      r0 = *(s32 *)(r1 + 0)
+; CHECK-NEXT: exit
+}
+
+; Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn
+define dso_local i64 @f7(ptr nocapture noundef readonly %p) local_unnamed_addr #0 {
+entry:
+  %0 = load i64, ptr %p, align 8, !tbaa !10
+  ret i64 %0
+; CHECK:      r0 = *(u64 *)(r1 + 0)
+; CHECK-NEXT: exit
+}
+
+attributes #0 = { argmemonly mustprogress nofree norecurse nosync nounwind readonly willreturn "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }
+
+!llvm.module.flags = !{!0, !1}
+!llvm.ident = !{!2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"frame-pointer", i32 2}
+!2 = !{!"clang version 16.0.0 (https://github.com/llvm/llvm-project.git 68665544c7d59735e9c0bb32b08829c006c7c594)"}
+!3 = !{!4, !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C/C++ TBAA"}
+!6 = !{!7, !7, i64 0}
+!7 = !{!"short", !4, i64 0}
+!8 = !{!9, !9, i64 0}
+!9 = !{!"int", !4, i64 0}
+!10 = !{!11, !11, i64 0}
+!11 = !{!"long", !4, i64 0}