From 40ae4b890353fcaa4a2c8d84fb0a9a5a9f122bbb Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper@sifive.com>
Date: Sun, 13 Nov 2022 21:54:15 -0800
Subject: [PATCH] [RISCV] Improve PACKH instruction selection

Handle AssertZExt in addition to AND.
---
 llvm/lib/Target/RISCV/RISCVInstrInfo.td   |  1 +
 llvm/lib/Target/RISCV/RISCVInstrInfoZb.td |  6 +++---
 llvm/test/CodeGen/RISCV/rv32zbkb.ll       | 19 +++++++++++++++++++
 llvm/test/CodeGen/RISCV/rv64zbkb.ll       | 18 ++++++++++++++++++
 4 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.td b/llvm/lib/Target/RISCV/RISCVInstrInfo.td
index 38ff929..0089ccb 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfo.td
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfo.td
@@ -1174,6 +1174,7 @@ def assertzexti32 : PatFrag<(ops node:$src), (assertzext node:$src), [{
 }]>;
 def zexti32 : ComplexPattern<i64, 1, "selectZExtBits<32>">;
 def zexti16 : ComplexPattern<XLenVT, 1, "selectZExtBits<16>">;
+def zexti8 : ComplexPattern<XLenVT, 1, "selectZExtBits<8>">;
 
 class binop_oneuse<SDPatternOperator operator>
     : PatFrag<(ops node:$A, node:$B),
diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfoZb.td b/llvm/lib/Target/RISCV/RISCVInstrInfoZb.td
index 44ef576..2928219 100644
--- a/llvm/lib/Target/RISCV/RISCVInstrInfoZb.td
+++ b/llvm/lib/Target/RISCV/RISCVInstrInfoZb.td
@@ -633,10 +633,10 @@ def : Pat<(i64 (bswap GPR:$rs1)), (REV8_RV64 GPR:$rs1)>;
 
 let Predicates = [HasStdExtZbkb] in {
 def : Pat<(or (and (shl GPR:$rs2, (XLenVT 8)), 0xFFFF),
-              (and GPR:$rs1, 0x00FF)),
+              (zexti8 GPR:$rs1)),
           (PACKH GPR:$rs1, GPR:$rs2)>;
-def : Pat<(or (shl (and GPR:$rs2, 0x00FF), (XLenVT 8)),
-              (and GPR:$rs1, 0x00FF)),
+def : Pat<(or (shl (zexti8 GPR:$rs2), (XLenVT 8)),
+              (zexti8 GPR:$rs1)),
           (PACKH GPR:$rs1, GPR:$rs2)>;
 } // Predicates = [HasStdExtZbkb]
 
diff --git a/llvm/test/CodeGen/RISCV/rv32zbkb.ll b/llvm/test/CodeGen/RISCV/rv32zbkb.ll
index 5a48925..7ae4c7b 100644
--- a/llvm/test/CodeGen/RISCV/rv32zbkb.ll
+++ b/llvm/test/CodeGen/RISCV/rv32zbkb.ll
@@ -189,3 +189,22 @@ define i64 @packh_i64_2(i64 %a, i64 %b) nounwind {
   %or = or i64 %shl, %and
   ret i64 %or
 }
+
+
+define zeroext i16 @packh_i16(i8 zeroext %a, i8 zeroext %b) nounwind {
+; RV32I-LABEL: packh_i16:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    slli a1, a1, 8
+; RV32I-NEXT:    or a0, a1, a0
+; RV32I-NEXT:    ret
+;
+; RV32ZBKB-LABEL: packh_i16:
+; RV32ZBKB:       # %bb.0:
+; RV32ZBKB-NEXT:    packh a0, a0, a1
+; RV32ZBKB-NEXT:    ret
+  %zext = zext i8 %a to i16
+  %zext1 = zext i8 %b to i16
+  %shl = shl i16 %zext1, 8
+  %or = or i16 %shl, %zext
+  ret i16 %or
+}
diff --git a/llvm/test/CodeGen/RISCV/rv64zbkb.ll b/llvm/test/CodeGen/RISCV/rv64zbkb.ll
index 6aabf78..757cd78 100644
--- a/llvm/test/CodeGen/RISCV/rv64zbkb.ll
+++ b/llvm/test/CodeGen/RISCV/rv64zbkb.ll
@@ -181,3 +181,21 @@ define i64 @packh_i64_2(i64 %a, i64 %b) nounwind {
   %or = or i64 %shl, %and
   ret i64 %or
 }
+
+define zeroext i16 @packh_i16(i8 zeroext %a, i8 zeroext %b) nounwind {
+; RV64I-LABEL: packh_i16:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    slli a1, a1, 8
+; RV64I-NEXT:    or a0, a1, a0
+; RV64I-NEXT:    ret
+;
+; RV64ZBKB-LABEL: packh_i16:
+; RV64ZBKB:       # %bb.0:
+; RV64ZBKB-NEXT:    packh a0, a0, a1
+; RV64ZBKB-NEXT:    ret
+  %zext = zext i8 %a to i16
+  %zext1 = zext i8 %b to i16
+  %shl = shl i16 %zext1, 8
+  %or = or i16 %shl, %zext
+  ret i16 %or
+}
-- 
2.7.4