From 5b6c9a5ab092954b5305b51a1f7b02aec48285a5 Mon Sep 17 00:00:00 2001
From: Ben Shi <powerman1st@163.com>
Date: Tue, 17 Aug 2021 16:40:16 +0800
Subject: [PATCH] [RISCV] Optimize add in the zba extension with SH*ADD

Optimize (add x, c) to (SH*ADD (c>>b), x) if c is not simm12
while (c>>b) is simm12 and c has b trailing zeros.

Reviewed By: luismarques

Differential Revision: https://reviews.llvm.org/D108193
---
 llvm/lib/Target/RISCV/RISCVInstrInfoB.td | 33 ++++++++++++++++++++++++
 llvm/test/CodeGen/RISCV/rv32zba.ll       | 20 ++++++--------
 llvm/test/CodeGen/RISCV/rv64zba.ll       | 20 ++++++--------
 3 files changed, 49 insertions(+), 24 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfoB.td b/llvm/lib/Target/RISCV/RISCVInstrInfoB.td
index 71b0f1585aed..a2cd11b1f26e 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfoB.td
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfoB.td
@@ -186,6 +186,32 @@ def C9LeftShift : PatLeaf<(imm), [{
   return C > 9 && ((C % 9) == 0) && isPowerOf2_64(C / 9);
 }]>;
 
+def CSImm12MulBy4 : PatLeaf<(imm), [{
+  if (!N->hasOneUse())
+    return false;
+  int64_t C = N->getSExtValue();
+  // Skip if C is simm12 or can be optimized by the PatLeaf AddiPair.
+  return !isInt<13>(C) && isInt<14>(C) && (C & 3) == 0;
+}]>;
+
+def CSImm12MulBy8 : PatLeaf<(imm), [{
+  if (!N->hasOneUse())
+    return false;
+  int64_t C = N->getSExtValue();
+  // Skip if C is simm12 or can be optimized by the PatLeaf AddiPair.
+  return !isInt<13>(C) && isInt<15>(C) && (C & 7) == 0;
+}]>;
+
+def SimmShiftRightBy2XForm : SDNodeXForm<imm, [{
+  return CurDAG->getTargetConstant(N->getSExtValue() >> 2, SDLoc(N),
+                                   N->getValueType(0));
+}]>;
+
+def SimmShiftRightBy3XForm : SDNodeXForm<imm, [{
+  return CurDAG->getTargetConstant(N->getSExtValue() >> 3, SDLoc(N),
+                                   N->getValueType(0));
+}]>;
+
 //===----------------------------------------------------------------------===//
 // Instruction class templates
 //===----------------------------------------------------------------------===//
@@ -1011,6 +1037,13 @@ def : Pat<(add (mul_oneuse GPR:$rs1, (XLenVT 40)), GPR:$rs2),
 def : Pat<(add (mul_oneuse GPR:$rs1, (XLenVT 72)), GPR:$rs2),
           (SH3ADD (SH3ADD GPR:$rs1, GPR:$rs1), GPR:$rs2)>;
 
+def : Pat<(add GPR:$r, CSImm12MulBy4:$i),
+          (SH2ADD (ADDI X0, (SimmShiftRightBy2XForm CSImm12MulBy4:$i)),
+                  GPR:$r)>;
+def : Pat<(add GPR:$r, CSImm12MulBy8:$i),
+          (SH3ADD (ADDI X0, (SimmShiftRightBy3XForm CSImm12MulBy8:$i)),
+                  GPR:$r)>;
+
 def : Pat<(mul GPR:$r, C3LeftShift:$i),
           (SLLI (SH1ADD GPR:$r, GPR:$r),
                 (TrailingZerosXForm C3LeftShift:$i))>;
diff --git a/llvm/test/CodeGen/RISCV/rv32zba.ll b/llvm/test/CodeGen/RISCV/rv32zba.ll
index f0a030045176..6ee346d2029c 100644
--- a/llvm/test/CodeGen/RISCV/rv32zba.ll
+++ b/llvm/test/CodeGen/RISCV/rv32zba.ll
@@ -750,16 +750,14 @@ define i32 @add4104(i32 %a) {
 ;
 ; RV32B-LABEL: add4104:
 ; RV32B:       # %bb.0:
-; RV32B-NEXT:    lui a1, 1
-; RV32B-NEXT:    addi a1, a1, 8
-; RV32B-NEXT:    add a0, a0, a1
+; RV32B-NEXT:    addi a1, zero, 1026
+; RV32B-NEXT:    sh2add a0, a1, a0
 ; RV32B-NEXT:    ret
 ;
 ; RV32ZBA-LABEL: add4104:
 ; RV32ZBA:       # %bb.0:
-; RV32ZBA-NEXT:    lui a1, 1
-; RV32ZBA-NEXT:    addi a1, a1, 8
-; RV32ZBA-NEXT:    add a0, a0, a1
+; RV32ZBA-NEXT:    addi a1, zero, 1026
+; RV32ZBA-NEXT:    sh2add a0, a1, a0
 ; RV32ZBA-NEXT:    ret
   %c = add i32 %a, 4104
   ret i32 %c
@@ -775,16 +773,14 @@ define i32 @add8208(i32 %a) {
 ;
 ; RV32B-LABEL: add8208:
 ; RV32B:       # %bb.0:
-; RV32B-NEXT:    lui a1, 2
-; RV32B-NEXT:    addi a1, a1, 16
-; RV32B-NEXT:    add a0, a0, a1
+; RV32B-NEXT:    addi a1, zero, 1026
+; RV32B-NEXT:    sh3add a0, a1, a0
 ; RV32B-NEXT:    ret
 ;
 ; RV32ZBA-LABEL: add8208:
 ; RV32ZBA:       # %bb.0:
-; RV32ZBA-NEXT:    lui a1, 2
-; RV32ZBA-NEXT:    addi a1, a1, 16
-; RV32ZBA-NEXT:    add a0, a0, a1
+; RV32ZBA-NEXT:    addi a1, zero, 1026
+; RV32ZBA-NEXT:    sh3add a0, a1, a0
 ; RV32ZBA-NEXT:    ret
   %c = add i32 %a, 8208
   ret i32 %c
diff --git a/llvm/test/CodeGen/RISCV/rv64zba.ll b/llvm/test/CodeGen/RISCV/rv64zba.ll
index 16ce6196d836..0c6c2f23b896 100644
--- a/llvm/test/CodeGen/RISCV/rv64zba.ll
+++ b/llvm/test/CodeGen/RISCV/rv64zba.ll
@@ -1321,16 +1321,14 @@ define i64 @add4104(i64 %a) {
 ;
 ; RV64B-LABEL: add4104:
 ; RV64B:       # %bb.0:
-; RV64B-NEXT:    lui a1, 1
-; RV64B-NEXT:    addiw a1, a1, 8
-; RV64B-NEXT:    add a0, a0, a1
+; RV64B-NEXT:    addi a1, zero, 1026
+; RV64B-NEXT:    sh2add a0, a1, a0
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: add4104:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    lui a1, 1
-; RV64ZBA-NEXT:    addiw a1, a1, 8
-; RV64ZBA-NEXT:    add a0, a0, a1
+; RV64ZBA-NEXT:    addi a1, zero, 1026
+; RV64ZBA-NEXT:    sh2add a0, a1, a0
 ; RV64ZBA-NEXT:    ret
   %c = add i64 %a, 4104
   ret i64 %c
@@ -1346,16 +1344,14 @@ define i64 @add8208(i64 %a) {
 ;
 ; RV64B-LABEL: add8208:
 ; RV64B:       # %bb.0:
-; RV64B-NEXT:    lui a1, 2
-; RV64B-NEXT:    addiw a1, a1, 16
-; RV64B-NEXT:    add a0, a0, a1
+; RV64B-NEXT:    addi a1, zero, 1026
+; RV64B-NEXT:    sh3add a0, a1, a0
 ; RV64B-NEXT:    ret
 ;
 ; RV64ZBA-LABEL: add8208:
 ; RV64ZBA:       # %bb.0:
-; RV64ZBA-NEXT:    lui a1, 2
-; RV64ZBA-NEXT:    addiw a1, a1, 16
-; RV64ZBA-NEXT:    add a0, a0, a1
+; RV64ZBA-NEXT:    addi a1, zero, 1026
+; RV64ZBA-NEXT:    sh3add a0, a1, a0
 ; RV64ZBA-NEXT:    ret
   %c = add i64 %a, 8208
   ret i64 %c