[AVR] Optimize 32-bit shift: move bytes around

[aykevl]

Summary

This patch optimizes 32-bit constant shifts by renaming registers. This
is very effective as the compiler would otherwise need to do a lot of
single bit shift instructions. Instead, the registers are renamed at the
SSA level which means the register allocator will insert the necessary
mov instructions.

Unfortunately, the register allocator will insert some unnecessary movs
with the current code. This will be fixed in a later patch.

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1887

I am a bit confused about here,

std::pair(Out, 0) is first put to the tail, then is dropped by drop_back(1) ?

In the first round of this loop, Regs comes in with size==4, then in the next round, Regs only keeps 3 elements before the moves?

1917

I am also confused by here as your did for the left shift, the std::pair(ShrExtendReg, 0) is put at the head of Regs, and then dropped by drop_front(1) ? Is it necessary to create a temporary copy of Regs ?

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1887

If the last element is dropped, is the above Regs[Regs.size() - 1] = std::pair(Out, 0); redundant ?

1917

If Regs.drop_front(1);, is the above Regs[0] = std::pair(ShrExtendReg, 0); redundant ?

[aykevl]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1887

std::pair(Out, 0) is first put to the tail, then is dropped by drop_back(1) ?

Correct. This is a MutableArrayRef and is modified in-place (it directly modifies the Registers array inside insertWideShift). For example, when shifting left by 8 the least significant register in the Registers array (Registers[3]) is zeroed. The Regs = Regs.drop_back(1) line only modifies the _view_ on the underlying array. That means that in the next iteration of the loop, Regs will only have a length of 3 (while Registers still has a length of four).

See my next comment, where it is hopefully more clear.

1917

If Regs.drop_front(1);, is the above Regs[0] = std::pair(ShrExtendReg, 0); redundant ?

No. When shifting right by 8, Regs[0] = std::pair(ShrExtendReg, 0); directly modifies Registers[0]. After Regs = Regs.drop_front(1), Regs[0] refers to Registers[1], Regs[1] refers to Registers[2], and Regs[2] refers to Registers[3]. This means that if you logically shift right by 16, in the first iteration Regs[0] will zero Registers[0] but in the second iteration Regs[0] will zero Registers[1].

Regs is simply a view on Registers. It is not a separate array.
(This is very similar to slices in Go, which is the language I often use).

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1866–1869

Can this be const DebugLoc dl ?

1917

I see. Your way is good !

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1873

These two variables are only used inside while (ShiftAmt >= 8), so how about moving them down to the head of the loop ?

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1898

I think this if (ShrExtendReg == 0) can be moved out of current loop, something like

Register ShrExtendReg = 0;
if (ShiftAmt >= 8) {
  auto &MSBReg = Regs[0];
  ShrExtendReg = MRI.createVirtualRegister(&AVR::GPR8RegClass);
  if (ArithmeticShift) {
    ...
  } else {
    ...
  }
}
while (ShiftAmt >= 8) {
  Regs[0] = std::pair(ShrExtendReg, 0);
  Regs = Regs.drop_front(1);
  ShiftAmt -= 8;
}

Actually

1. we need not check if (ShrExtendReg == 0) time and time again.
2. we need not expose auto &MSBReg = Regs[0] which is only used locally.

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1883

Do we need to do AVR::COPY for each iteration ? Can it be a unique Register Out shared by all iterations, and even shared with lshr ?

[benshi001]

llvm/test/CodeGen/AVR/shift32.ll

131

How about adding a lshr_i32_9 and see what will happen ?

219

How about adding a ashr_i32_17 and see what will happen ?

[aykevl]

Applied all review feedback and rebased the patch.

llvm/lib/Target/AVR/AVRISelLowering.cpp

1866–1869

This is part of D140569.

1873

Fixed.

1883

Interesting idea! I've tried it, and it seems to work just as well.

1898

Fixed.

[aykevl]

• rebase (somehow the pre-merge checks failed because of an error?)

[benshi001]

llvm/lib/Target/AVR/AVRISelLowering.cpp

1878

This can be for (; ShiftLeft && ShiftAmt >= 8; ShiftAmt -= 8) as you did in the next loop.

You can change that while committing.

llvm/lib/Target/AVR/AVRISelLowering.cpp

MachineBasicBlock *AVRTargetLowering::insertShift(MachineInstr &MI,
   MI.eraseFromParent(); // The pseudo instruction is gone now.
   return RemBB;
 }
 
 // Do a multibyte AVR shift. Insert shift instructions and put the output
 // registers in the Regs array.
 // Because AVR does not have a normal shift instruction (only a single bit shift
 // instruction), we have to emulate this behavior with other instructions.
+// It first tries large steps (moving registers around) and then smaller steps
+// like single bit shifts.
+// Large shifts actually reduce the number of shifted registers, so the below
+// algorithms have to work independently of the number of registers that are
+// shifted.
+// For more information and background, see this blogpost:
+// https://aykevl.nl/2021/02/avr-bitshift
 static void insertMultibyteShift(MachineInstr &MI, MachineBasicBlock *BB,
                                  MutableArrayRef<std::pair<Register, int>> Regs,
                                  ISD::NodeType Opc, int64_t ShiftAmt) {
   const TargetInstrInfo &TII = *BB->getParent()->getSubtarget().getInstrInfo();
+  const AVRSubtarget &STI = BB->getParent()->getSubtarget<AVRSubtarget>();
   MachineRegisterInfo &MRI = BB->getParent()->getRegInfo();
   const DebugLoc &dl = MI.getDebugLoc();
 
   const bool ShiftLeft = Opc == ISD::SHL;
   const bool ArithmeticShift = Opc == ISD::SRA;
 
+  // Zero a register, for use in later operations.
+  Register ZeroReg = MRI.createVirtualRegister(&AVR::GPR8RegClass);
+  BuildMI(*BB, MI, dl, TII.get(AVR::COPY), ZeroReg)
+      .addReg(STI.getZeroRegister());
+
+  // For shift amounts of at least one register, simply rename the registers and
+  // zero the bottom registers.
+  while (ShiftLeft && ShiftAmt >= 8) {
+    // Move all registers one to the left.
+    for (size_t I = 0; I < Regs.size() - 1; I++) {
+      Regs[I] = Regs[I + 1];
+    }
+
+    // Zero the least significant register.
+    Regs[Regs.size() - 1] = std::pair(ZeroReg, 0);
+
+    // Continue shifts with the leftover registers.
+    Regs = Regs.drop_back(1);
+
+    ShiftAmt -= 8;
+  }
+
+  // And again, the same for right shifts.
+  Register ShrExtendReg = 0;
+  if (!ShiftLeft && ShiftAmt >= 8) {
+    if (ArithmeticShift) {
+      // Sign extend the most significant register into ShrExtendReg.
+      ShrExtendReg = MRI.createVirtualRegister(&AVR::GPR8RegClass);
+      Register Tmp = MRI.createVirtualRegister(&AVR::GPR8RegClass);
+      BuildMI(*BB, MI, dl, TII.get(AVR::ADDRdRr), Tmp)
+          .addReg(Regs[0].first, 0, Regs[0].second)
+          .addReg(Regs[0].first, 0, Regs[0].second);
+      BuildMI(*BB, MI, dl, TII.get(AVR::SBCRdRr), ShrExtendReg)
+          .addReg(Tmp)
+          .addReg(Tmp);
+    } else {
+      ShrExtendReg = ZeroReg;
+    }
+    for (; ShiftAmt >= 8; ShiftAmt -= 8) {
+      // Move all registers one to the right.
+      for (size_t I = Regs.size() - 1; I != 0; I--) {
+        Regs[I] = Regs[I - 1];
+      }
+
+      // Zero or sign extend the most significant register.
+      Regs[0] = std::pair(ShrExtendReg, 0);
+
+      // Continue shifts with the leftover registers.
+      Regs = Regs.drop_front(1);
+    }
+  }
+
+  // The bigger shifts are already handled above.
+  assert((ShiftAmt < 8) && "Unexpect shift amount");
+
   // Shift by one. This is the fallback that always works, and the shift
   // operation that is used for 1, 2, and 3 bit shifts.
   while (ShiftLeft && ShiftAmt) {
     // Shift one to the left.
     for (ssize_t I = Regs.size() - 1; I >= 0; I--) {
       Register Out = MRI.createVirtualRegister(&AVR::GPR8RegClass);
       Register In = Regs[I].first;
       Register InSubreg = Regs[I].second;

llvm/test/CodeGen/AVR/shift32.ll

 ; CHECK-NEXT:    rol r23
 ; CHECK-NEXT:    rol r24
 ; CHECK-NEXT:    rol r25
 ; CHECK-NEXT:    ret
   %res = shl i32 %a, 2
   ret i32 %res
 }
 
+define i32 @shl_i32_8(i32 %a) {
+; CHECK-LABEL: shl_i32_8:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov r25, r24
+; CHECK-NEXT:    mov r24, r23
+; CHECK-NEXT:    mov r23, r22
+; CHECK-NEXT:    mov r22, r1
+; CHECK-NEXT:    ret
+  %res = shl i32 %a, 8
+  ret i32 %res
+}
+
+define i32 @shl_i32_9(i32 %a) {
+; CHECK-LABEL: shl_i32_9:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    lsl r22
+; CHECK-NEXT:    rol r23
+; CHECK-NEXT:    rol r24
+; CHECK-NEXT:    mov r25, r24
+; CHECK-NEXT:    mov r24, r23
+; CHECK-NEXT:    mov r23, r22
+; CHECK-NEXT:    mov r22, r1
+; CHECK-NEXT:    ret
+  %res = shl i32 %a, 9
+  ret i32 %res
+}
+
 ; This is a special case: this shift is performed directly inside SelectionDAG
 ; instead of as a custom lowering like the other shift operations.
 define i32 @shl_i32_16(i32 %a) {
 ; CHECK-LABEL: shl_i32_16:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    movw r24, r22
 ; CHECK-NEXT:    ldi r22, 0
 ; CHECK-NEXT:    ldi r23, 0

 ; CHECK-NEXT:    ror r24
 ; CHECK-NEXT:    ror r23
 ; CHECK-NEXT:    ror r22
 ; CHECK-NEXT:    ret
   %res = lshr i32 %a, 2
   ret i32 %res
 }
 
+define i32 @lshr_i32_8(i32 %a) {
+; CHECK-LABEL: lshr_i32_8:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov r19, r1
+; CHECK-NEXT:    mov r18, r25
+; CHECK-NEXT:    mov r25, r24
+; CHECK-NEXT:    mov r24, r23
+; CHECK-NEXT:    movw r22, r24
+; CHECK-NEXT:    movw r24, r18
+; CHECK-NEXT:    ret
+  %res = lshr i32 %a, 8
+  ret i32 %res
+}
+
+define i32 @lshr_i32_9(i32 %a) {
+; CHECK-LABEL: lshr_i32_9:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    lsr r25
+; CHECK-NEXT:    ror r24
+; CHECK-NEXT:    ror r23
+; CHECK-NEXT:    mov r19, r1
+; CHECK-NEXT:    mov r18, r25
+; CHECK-NEXT:    mov r25, r24
+; CHECK-NEXT:    mov r24, r23
+; CHECK-NEXT:    movw r22, r24
+; CHECK-NEXT:    movw r24, r18
+; CHECK-NEXT:    ret
+  %res = lshr i32 %a, 9
+  ret i32 %res
+}
+
 define i32 @lshr_i32_16(i32 %a) {
 ; CHECK-LABEL: lshr_i32_16:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    movw r22, r24
 ; CHECK-NEXT:    ldi r24, 0
 ; CHECK-NEXT:    ldi r25, 0
 ; CHECK-NEXT:    ret
   %res = lshr i32 %a, 16
   ret i32 %res
 }
 
+define i32 @lshr_i32_24(i32 %a) {
+; CHECK-LABEL: lshr_i32_24:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov r19, r1
+; CHECK-NEXT:    mov r18, r1
+; CHECK-NEXT:    mov r23, r1
+; CHECK-NEXT:    mov r22, r25
+; CHECK-NEXT:    movw r24, r18
+; CHECK-NEXT:    ret
+  %res = lshr i32 %a, 24
+  ret i32 %res
+}
+
 define i32 @ashr_i32_1(i32 %a) {
 ; CHECK-LABEL: ashr_i32_1:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    asr r25
 ; CHECK-NEXT:    ror r24
 ; CHECK-NEXT:    ror r23
 ; CHECK-NEXT:    ror r22
 ; CHECK-NEXT:    ret

 ; CHECK-NEXT:    asr r25
 ; CHECK-NEXT:    ror r24
 ; CHECK-NEXT:    ror r23
 ; CHECK-NEXT:    ror r22
 ; CHECK-NEXT:    ret
   %res = ashr i32 %a, 2
   ret i32 %res
 }
+
+; TODO: this could be optimized to 4 movs, instead of 6.
+define i32 @ashr_i32_8(i32 %a) {
+; CHECK-LABEL: ashr_i32_8:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    mov r19, r25
+; CHECK-NEXT:    lsl r19
+; CHECK-NEXT:    sbc r19, r19
+; CHECK-NEXT:    mov r18, r25
+; CHECK-NEXT:    mov r25, r24
+; CHECK-NEXT:    mov r24, r23
+; CHECK-NEXT:    movw r22, r24
+; CHECK-NEXT:    movw r24, r18
+; CHECK-NEXT:    ret
+  %res = ashr i32 %a, 8
+  ret i32 %res
+}
+
+define i32 @ashr_i32_16(i32 %a) {
+; CHECK-LABEL: ashr_i32_16:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    movw r22, r24
+; CHECK-NEXT:    lsl r25
+; CHECK-NEXT:    sbc r25, r25
+; CHECK-NEXT:    mov r24, r25
+; CHECK-NEXT:    ret
+  %res = ashr i32 %a, 16
+  ret i32 %res
+}
+
+define i32 @ashr_i32_17(i32 %a) {
+; CHECK-LABEL: ashr_i32_17:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    movw r22, r24
+; CHECK-NEXT:    lsl r25
+; CHECK-NEXT:    sbc r25, r25
+; CHECK-NEXT:    asr r23
+; CHECK-NEXT:    ror r22
+; CHECK-NEXT:    mov r24, r25
+; CHECK-NEXT:    ret
+  %res = ashr i32 %a, 17
+  ret i32 %res
+}