From f4ef79306cee2b5866aff681174f16b816810c4a Mon Sep 17 00:00:00 2001
From: Ben Shi <ben.shi@streamcomputing.com>
Date: Tue, 4 Jan 2022 03:20:29 +0000
Subject: [PATCH] [AVR] Optimize int8 arithmetic right shift 6 bits

Reviewed By: aykevl

Differential Revision: https://reviews.llvm.org/D115593
---
 llvm/lib/Target/AVR/AVRExpandPseudoInsts.cpp | 42 ++++++++++++++++++++
 llvm/lib/Target/AVR/AVRISelLowering.cpp      |  5 +++
 llvm/test/CodeGen/AVR/shift.ll               | 12 +++++-
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AVR/AVRExpandPseudoInsts.cpp b/llvm/lib/Target/AVR/AVRExpandPseudoInsts.cpp
index cb85d73772c55..eaff501ebe3b0 100644
--- a/llvm/lib/Target/AVR/AVRExpandPseudoInsts.cpp
+++ b/llvm/lib/Target/AVR/AVRExpandPseudoInsts.cpp
@@ -92,6 +92,7 @@ class AVRExpandPseudo : public MachineFunctionPass {
   /// Specific shift implementation.
   bool expandLSLB7Rd(Block &MBB, BlockIt MBBI);
   bool expandLSRB7Rd(Block &MBB, BlockIt MBBI);
+  bool expandASRB6Rd(Block &MBB, BlockIt MBBI);
   bool expandASRB7Rd(Block &MBB, BlockIt MBBI);
   bool expandLSLW4Rd(Block &MBB, BlockIt MBBI);
   bool expandLSRW4Rd(Block &MBB, BlockIt MBBI);
@@ -1921,6 +1922,45 @@ bool AVRExpandPseudo::expand<AVR::LSRBNRd>(Block &MBB, BlockIt MBBI) {
   }
 }
 
+bool AVRExpandPseudo::expandASRB6Rd(Block &MBB, BlockIt MBBI) {
+  MachineInstr &MI = *MBBI;
+  Register DstReg = MI.getOperand(0).getReg();
+  bool DstIsDead = MI.getOperand(0).isDead();
+  bool DstIsKill = MI.getOperand(1).isKill();
+  bool ImpIsDead = MI.getOperand(3).isDead();
+
+  // bst r24, 6
+  // lsl r24
+  // sbc r24, r24
+  // bld r24, 0
+
+  buildMI(MBB, MBBI, AVR::BST)
+      .addReg(DstReg)
+      .addImm(6)
+      ->getOperand(2)
+      .setIsUndef(true);
+
+  buildMI(MBB, MBBI, AVR::ADDRdRr) // LSL Rd <==> ADD Rd, Rd
+      .addReg(DstReg, RegState::Define | getDeadRegState(DstIsDead))
+      .addReg(DstReg, getKillRegState(DstIsKill))
+      .addReg(DstReg, getKillRegState(DstIsKill));
+
+  buildMI(MBB, MBBI, AVR::SBCRdRr)
+      .addReg(DstReg, RegState::Define | getDeadRegState(DstIsDead))
+      .addReg(DstReg, getKillRegState(DstIsKill))
+      .addReg(DstReg, getKillRegState(DstIsKill));
+
+  buildMI(MBB, MBBI, AVR::BLD)
+      .addReg(DstReg, RegState::Define | getDeadRegState(DstIsDead))
+      .addReg(DstReg, getKillRegState(DstIsKill))
+      .addImm(0)
+      ->getOperand(3)
+      .setIsKill();
+
+  MI.eraseFromParent();
+  return true;
+}
+
 bool AVRExpandPseudo::expandASRB7Rd(Block &MBB, BlockIt MBBI) {
   MachineInstr &MI = *MBBI;
   Register DstReg = MI.getOperand(0).getReg();
@@ -1957,6 +1997,8 @@ bool AVRExpandPseudo::expand<AVR::ASRBNRd>(Block &MBB, BlockIt MBBI) {
   MachineInstr &MI = *MBBI;
   unsigned Imm = MI.getOperand(2).getImm();
   switch (Imm) {
+  case 6:
+    return expandASRB6Rd(MBB, MBBI);
   case 7:
     return expandASRB7Rd(MBB, MBBI);
   default:
diff --git a/llvm/lib/Target/AVR/AVRISelLowering.cpp b/llvm/lib/Target/AVR/AVRISelLowering.cpp
index 39fba74a1ec78..f3e74e8436956 100644
--- a/llvm/lib/Target/AVR/AVRISelLowering.cpp
+++ b/llvm/lib/Target/AVR/AVRISelLowering.cpp
@@ -359,6 +359,11 @@ SDValue AVRTargetLowering::LowerShifts(SDValue Op, SelectionDAG &DAG) const {
       Victim = DAG.getNode(AVRISD::LSRBN, dl, VT, Victim,
                            DAG.getConstant(7, dl, VT));
       ShiftAmount = 0;
+    } else if (Op.getOpcode() == ISD::SRA && ShiftAmount == 6) {
+      // Optimize ASR when ShiftAmount == 6.
+      Victim = DAG.getNode(AVRISD::ASRBN, dl, VT, Victim,
+                           DAG.getConstant(6, dl, VT));
+      ShiftAmount = 0;
     } else if (Op.getOpcode() == ISD::SRA && ShiftAmount == 7) {
       // Optimize ASR when ShiftAmount == 7.
       Victim = DAG.getNode(AVRISD::ASRBN, dl, VT, Victim,
diff --git a/llvm/test/CodeGen/AVR/shift.ll b/llvm/test/CodeGen/AVR/shift.ll
index 24bc369cf6143..90e1b25bd762a 100644
--- a/llvm/test/CodeGen/AVR/shift.ll
+++ b/llvm/test/CodeGen/AVR/shift.ll
@@ -1,4 +1,4 @@
-; RUN: llc < %s -march=avr | FileCheck %s
+; RUN: llc < %s -march=avr -verify-machineinstrs | FileCheck %s
 
 ; Optimize for speed.
 ; CHECK-LABEL: shift_i8_i8_speed
@@ -171,6 +171,16 @@ define i8 @lsr_i8_7(i8 %a) {
   ret i8 %result
 }
 
+define i8 @asr_i8_6(i8 %a) {
+; CHECK-LABEL: asr_i8_6
+; CHECK:      bst r24, 6
+; CHECK-NEXT: lsl r24
+; CHECK-NEXT: sbc r24, r24
+; CHECK-NEXT: bld r24, 0
+  %result = ashr i8 %a, 6
+  ret i8 %result
+}
+
 define i8 @asr_i8_7(i8 %a) {
 ; CHECK-LABEL: asr_i8_7
 ; CHECK:      lsl r24