[AArch64][GlobalISel] Use wzr/xzr for 16 and 32 bit stores of zero

[paquette]

Summary

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.

[aemerson]

LGTM.

llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp

#endif
     if (!Offset && PtrMI->getOpcode() == TargetOpcode::G_FRAME_INDEX)
       I.getOperand(1).ChangeToFrameIndex(PtrMI->getOperand(1).getIndex());
 
     I.addOperand(MachineOperand::CreateImm(Offset));
 
     // 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)
+        unsigned Opc = I.getOpcode();
+        switch(Opc) {
+          case AArch64::STRWui:
+          case AArch64::STRHHui:
+          case AArch64::STRBBui:
             I.getOperand(0).setReg(AArch64::WZR);
-        else if (I.getOpcode() == AArch64::STRXui)
+            break;
+          case AArch64::STRXui:
             I.getOperand(0).setReg(AArch64::XZR);
+            break;
+        }
       }
     }
 
     if (IsZExtLoad) {
       // The zextload from a smaller type to i32 should be handled by the importer.
       if (MRI.getType(ValReg).getSizeInBits() != 64)
         return false;
       // If we have a ZEXTLOAD then change the load's type to be a narrower reg

llvm/test/CodeGen/AArch64/GlobalISel/select-store.mir

 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -mtriple=aarch64-- -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s
 
 --- |
   target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
 
   define void @store_s64_gpr(i64* %addr) { ret void }
   define void @store_s32_gpr(i32* %addr) { ret void }
   define void @store_s16_gpr(i16* %addr) { ret void }
   define void @store_s8_gpr(i8* %addr) { ret void }
 
   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
     ret void
   }
 
   define void @store_gep_128_s64_gpr(i64* %addr) { ret void }
   define void @store_gep_512_s32_gpr(i32* %addr) { ret void }
bb.0:
     ; CHECK: STRWui $wzr, [[COPY]], 0 :: (store 4 into %ir.addr)
     %0(p0) = COPY $x0
     %1(s32) = G_CONSTANT i32 0
     G_STORE  %1, %0 :: (store 4 into %ir.addr)
 
 ...
 
 ---
+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
 regBankSelected: true
 
 registers:
   - { id: 0, class: gpr }
   - { id: 1, class: gpr }