[BPF] fix sub-register handling for bpf_fastcall

bpf_fastcall induced spill/fill pairs should be generated for
sub-register as well as for sub-registers. At the moment this is not
the case, e.g.:

    $ cat t.c
    extern int foo(void) __attribute__((bpf_fastcall));

    int bar(int a) {
      foo();
      return a;
    }

    $ clang --target=bpf -mcpu=v3 -O2 -S t.c -o -
    ...
    call foo
    w0 = w1
    exit

Modify BPFMIPeephole.cpp:collectBPFFastCalls() to check sub-registers
liveness and thus produce correct code for example above:

    *(u64 *)(r10 - 8) = r1
    call foo
    r1 = *(u64 *)(r10 - 8)
    w0 = w1
    exit

Author: eddyz87

llvm/lib/Target/BPF/BPFMIPeephole.cpp

@@ -614,11 +614,20 @@ static void collectBPFFastCalls(const TargetRegisterInfo *TRI,
   LiveRegs.init(*TRI);
   LiveRegs.addLiveOuts(BB);
   Calls.clear();
+
+  auto IsLiveBeforeInsn = [&](MCRegister R, MachineInstr &MI) {
+    return !MI.definesRegister(R, TRI) && LiveRegs.contains(R);
+  };
+
   for (MachineInstr &MI : llvm::reverse(BB)) {
     if (MI.isCall()) {
       unsigned LiveCallerSavedRegs = 0;
       for (MCRegister R : CallerSavedRegs) {
-        bool DoSpillFill = !MI.definesRegister(R, TRI) && LiveRegs.contains(R);
+        bool DoSpillFill = IsLiveBeforeInsn(R, MI);
+        for (MCSubRegIndexIterator SRI(R, TRI); SRI.isValid(); ++SRI) {
+          MCRegister SR = SRI.getSubReg();
+          DoSpillFill |= IsLiveBeforeInsn(SR, MI);
+        }
         if (!DoSpillFill)
           continue;
         LiveCallerSavedRegs |= 1 << R;

llvm/test/CodeGen/BPF/bpf-fastcall-4.ll

@@ -0,0 +1,35 @@
+; RUN: llc -O2 --march=bpfel %s -o - | FileCheck %s
+
+; Generated from the following C code:
+;
+;   extern int foo(void) __attribute__((bpf_fastcall));
+;
+;   int bar(int a) {
+;     foo();
+;     return a;
+;   }
+;
+; Using the following command:
+;
+;   clang --target=bpf -emit-llvm -O2 -S -o - t.c
+;
+; (unnecessary attrs removed maually)
+
+; Check that function marked with bpf_fastcall does not clobber W1-W5.
+
+define dso_local i32 @bar(i32 %a) {
+entry:
+  %call = tail call i32 @foo() #0
+  ret i32 %a
+}
+
+; CHECK:      # %bb.0:
+; CHECK-NEXT:   *(u64 *)(r10 - 8) = r1
+; CHECK-NEXT:   call foo
+; CHECK-NEXT:   r1 = *(u64 *)(r10 - 8)
+; CHECK-NEXT:   w0 = w1
+; CHECK-NEXT:   exit
+
+declare dso_local i32 @foo() #0
+
+attributes #0 = { "bpf_fastcall" }

llvm/test/CodeGen/BPF/bpf-fastcall-5.ll

@@ -0,0 +1,35 @@
+; RUN: llc -O2 --march=bpfel %s -o - | FileCheck %s
+
+; Generated from the following C code:
+;
+;   extern int foo(void) __attribute__((bpf_fastcall));
+;
+;   int bar(int a, int b, int c, int d, int e) {
+;     foo();
+;     return e;
+;   }
+;
+; Using the following command:
+;
+;   clang --target=bpf -emit-llvm -O2 -S -o - t.c
+;
+; (unnecessary attrs removed maually)
+
+; Check that function marked with bpf_fastcall does not clobber W1-W5.
+
+define dso_local i32 @bar(i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) {
+entry:
+  %call = tail call i32 @foo() #0
+  ret i32 %e
+}
+
+; CHECK:      # %bb.0:
+; CHECK-NEXT:   *(u64 *)(r10 - 8) = r5
+; CHECK-NEXT:   call foo
+; CHECK-NEXT:   r5 = *(u64 *)(r10 - 8)
+; CHECK-NEXT:   w0 = w5
+; CHECK-NEXT:   exit
+
+declare dso_local i32 @foo() #0
+
+attributes #0 = { "bpf_fastcall" }