From b6d2a38769da7f4637c380cf4e601b8f0d54050d Mon Sep 17 00:00:00 2001
From: Jessica Paquette <jpaquette@apple.com>
Date: Fri, 24 Jul 2020 16:57:37 -0700
Subject: [PATCH] [AArch64][GlobalISel] Use wzr/xzr for 16 and 32 bit stores of
 zero

We weren't performing this optimization on 16 and 32 bit stores. SDAG happily
does this though.

e.g. https://godbolt.org/z/cWocKr

This saves about 0.2% in code size on CTMark at -O3.

Differential Revision: https://reviews.llvm.org/D84568
---
 .../GISel/AArch64InstructionSelector.cpp      | 15 ++++++---
 .../AArch64/GlobalISel/select-store.mir       | 33 +++++++++++++++++++
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp b/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
index eb6a4aa3d82..0b0de09fff2 100644
--- a/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
+++ b/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp
@@ -2306,10 +2306,17 @@ bool AArch64InstructionSelector::select(MachineInstr &I) {
     // If we're storing a 0, use WZR/XZR.
     if (auto CVal = getConstantVRegVal(ValReg, MRI)) {
       if (*CVal == 0 && Opcode == TargetOpcode::G_STORE) {
-        if (I.getOpcode() == AArch64::STRWui)
-          I.getOperand(0).setReg(AArch64::WZR);
-        else if (I.getOpcode() == AArch64::STRXui)
-          I.getOperand(0).setReg(AArch64::XZR);
+        unsigned Opc = I.getOpcode();
+        switch(Opc) {
+          case AArch64::STRWui:
+          case AArch64::STRHHui:
+          case AArch64::STRBBui:
+            I.getOperand(0).setReg(AArch64::WZR);
+            break;
+          case AArch64::STRXui:
+            I.getOperand(0).setReg(AArch64::XZR);
+            break;
+        }
       }
     }
 
diff --git a/test/CodeGen/AArch64/GlobalISel/select-store.mir b/test/CodeGen/AArch64/GlobalISel/select-store.mir
index b72b8e5d0d7..d60a34ef663 100644
--- a/test/CodeGen/AArch64/GlobalISel/select-store.mir
+++ b/test/CodeGen/AArch64/GlobalISel/select-store.mir
@@ -11,6 +11,8 @@
 
   define void @store_zero_s64_gpr(i64* %addr) { ret void }
   define void @store_zero_s32_gpr(i32* %addr) { ret void }
+  define void @store_zero_s16(i32* %addr) { ret void }
+  define void @store_zero_s8(i32* %addr) { ret void }
 
   define void @store_fi_s64_gpr() {
     %ptr0 = alloca i64
@@ -176,6 +178,37 @@ body:             |
 
 ...
 
+---
+name:            store_zero_s16
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: store_zero_s16
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: STRHHui $wzr, [[COPY]], 0 :: (store 2)
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s16) = G_CONSTANT i16 0
+    G_STORE %1(s16), %0(p0) :: (store 2)
+
+...
+
+---
+name:            store_zero_s8
+legalized:       true
+regBankSelected: true
+body:             |
+  bb.0:
+    liveins: $x0
+    ; CHECK-LABEL: name: store_zero_s8
+    ; CHECK: [[COPY:%[0-9]+]]:gpr64sp = COPY $x0
+    ; CHECK: STRBBui $wzr, [[COPY]], 0 :: (store 1)
+    %0:gpr(p0) = COPY $x0
+    %1:gpr(s8) = G_CONSTANT i8 0
+    G_STORE %1(s8), %0(p0) :: (store 1)
+...
+
 ---
 name:            store_fi_s64_gpr
 legalized:       true