[AArch64][GlobalISel] Use wzr/xzr for FP stores of zero

[paquette]

Summary

In situations like the example below, GlobalISel produces an unnecessary fmov from wzr.

https://godbolt.org/z/Y45W8T

define void @test(float* %ptr, float %x) {
    %cmp = fcmp une float undef, 0.000000e+00
    br i1 %cmp, label %a, label %b
a:
    %gep = getelementptr inbounds float, float* %ptr, i32 1
    store float 0.000000e+00, float* %gep, align 4
    ret void
b:
    ret void
}

In a lot of cases, matchFConstantToConstant in the pre-legalizer combiner handles this sort of thing.

However, that combine only works when all users of a G_FCONSTANT are stores. In this case, we have a G_FCMP which also uses the G_FCONSTANT.

This patch adds a post-legalization (lowering) combine which looks for G_STORES which store a positive 0. If such a store is found, and we know it can use wzr/xzr, then it's updated to use a G_CONSTANT 0.

During selection, the constant 0 + store will be selected to a store of wzr/xzr.

This is a 0.1% code size improvement at -Os for CTMark/mafft/pairlocalalign. There are minor code size improvements in other CTMark tests as well.

[aemerson]

Why is this in lowering instead of combine?

[paquette]

In D94641#2497021, @aemerson wrote:

Why is this in lowering instead of combine?

Originally I did it in select, since it seemed like something we would want at all optimization levels.

Moving it to lowering gave slightly better code size results on bullet and pairlocalalign. (152 B and 76 B respectively)

I think that it's likely because of regbankselect. Some of the instructions in there check onlyUsesFP to decide on which register bank to use (G_LOAD, G_SELECT, and G_UNMERGE_VALUES). So, we can potentially avoid some extra copies by making sure that instructions like G_STORES of 0.0 don't end up being marked as FPR.

[paquette]

A better, more general solution might be to move some of the G_FCONSTANT selection code into the legalizer or post-legalizer lowering. That might make some of the selection code which relies on seeing a G_FCONSTANT more difficult to get right though.

Or, alternatively, teach regbankselect to check if a G_FCONSTANT is actually producing something that can/should be emitted using fmov. If it's not, then it could actually be assigned a GPR, because we know that's what it will end up on.

[aemerson]

It it possible to extend the existing combine in prelegalizer combiner to handle this too? Or will that cause us to still miss some cases?

[paquette]

In D94641#2498944, @aemerson wrote:

It it possible to extend the existing combine in prelegalizer combiner to handle this too? Or will that cause us to still miss some cases?

The problem with that combine is it looks for G_FCONSTANTS which are only used by stores, so as is, it would miss cases.

I think that we could probably replace that combine with something similar to this though.

[aemerson]

In D94641#2498991, @paquette wrote:

In D94641#2498944, @aemerson wrote:

It it possible to extend the existing combine in prelegalizer combiner to handle this too? Or will that cause us to still miss some cases?

The problem with that combine is it looks for G_FCONSTANTS which are only used by stores, so as is, it would miss cases.

I think that we could probably replace that combine with something similar to this though.

Right. I think it would be nice to keep these int/fp constant of 0 combines somewhat together. If you feel that it's not a net benefit to do that then I'm happy for this to go in as is.

llvm/lib/Target/AArch64/AArch64Combine.td

 def mul_const_matchdata : GIDefMatchData<"std::function<void(MachineIRBuilder&, Register)>">;
 def mul_const : GICombineRule<
   (defs root:$root, mul_const_matchdata:$matchinfo),
   (match (wip_match_opcode G_MUL):$root,
           [{ return matchAArch64MulConstCombine(*${root}, MRI, ${matchinfo}); }]),
   (apply [{ applyAArch64MulConstCombine(*${root}, MRI, B, ${matchinfo}); }])
 >;
 
+def fp_store_with_zero_imm : GICombineRule<
+  (defs root:$root),
+  (match (wip_match_opcode G_STORE):$root,
+         [{ return matchFPStoreWithZeroImm(*${root}, MRI); }]),
+  (apply [{ applyFPStoreWithZeroImm(*${root}, B, Observer); }])>;
+
 // Post-legalization combines which should happen at all optimization levels.
 // (E.g. ones that facilitate matching for the selector) For example, matching
 // pseudos.
 def AArch64PostLegalizerLoweringHelper
     : GICombinerHelper<"AArch64GenPostLegalizerLoweringHelper",
                        [shuffle_vector_pseudos, vashr_vlshr_imm,
-                        icmp_lowering, form_duplane]> {
+                        icmp_lowering, form_duplane, fp_store_with_zero_imm]> {
   let DisableRuleOption = "aarch64postlegalizerlowering-disable-rule";
 }
 
 // Post-legalization combines which are primarily optimizations.
 def AArch64PostLegalizerCombinerHelper
     : GICombinerHelper<"AArch64GenPostLegalizerCombinerHelper",
                        [copy_prop, erase_undef_store, combines_for_extload,
                         sext_trunc_sextload,
                         hoist_logic_op_with_same_opcode_hands,
                         redundant_and, xor_of_and_with_same_reg,
                         extractvecelt_pairwise_add, redundant_or,
                         mul_const]> {
   let DisableRuleOption = "aarch64postlegalizercombiner-disable-rule";
 }

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

bool applyDupLane(MachineInstr &MI, MachineRegisterInfo &MRI,
   B.setInstrAndDebugLoc(MI);
   auto Lane = B.buildConstant(LLT::scalar(64), MatchInfo.second);
   B.buildInstr(MatchInfo.first, {MI.getOperand(0).getReg()},
                {MI.getOperand(1).getReg(), Lane});
   MI.eraseFromParent();
   return true;
 }
 
+/// Match a G_STORE which stores a positive 0.0.
+static bool matchFPStoreWithZeroImm(MachineInstr &I, MachineRegisterInfo &MRI) {
+  assert(I.getOpcode() == TargetOpcode::G_STORE);
+  Register Src = I.getOperand(0).getReg();
+  LLT Ty = MRI.getType(Src);
+  if (Ty.isVector())
+    return false;
+  unsigned Size = Ty.getSizeInBits();
+  if (Size != 32 && Size != 64)
+    return false;
+  auto *MaybeZero = getConstantFPVRegVal(Src, MRI);
+  if (!MaybeZero || !MaybeZero->isNullValue())
+    return false;
+  return true;
+}
+
+/// Change a G_STORE which stores a positive 0.0 into a G_STORE which stores
+/// an integer 0.
+static bool applyFPStoreWithZeroImm(MachineInstr &I, MachineIRBuilder &MIB,
+                                    GISelChangeObserver &Observer) {
+  MIB.setInstrAndDebugLoc(I);
+  MachineRegisterInfo &MRI = *MIB.getMRI();
+  auto Zero =
+      MIB.buildConstant(MRI.cloneVirtualRegister(I.getOperand(0).getReg()), 0);
+  Observer.changingInstr(I);
+  I.getOperand(0).setReg(Zero->getOperand(0).getReg());
+  Observer.changedInstr(I);
+  return true;
+}
+
 #define AARCH64POSTLEGALIZERLOWERINGHELPER_GENCOMBINERHELPER_DEPS
 #include "AArch64GenPostLegalizeGILowering.inc"
 #undef AARCH64POSTLEGALIZERLOWERINGHELPER_GENCOMBINERHELPER_DEPS
 
 namespace {
 #define AARCH64POSTLEGALIZERLOWERINGHELPER_GENCOMBINERHELPER_H
 #include "AArch64GenPostLegalizeGILowering.inc"
 #undef AARCH64POSTLEGALIZERLOWERINGHELPER_GENCOMBINERHELPER_H

llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-lowering-zero-store.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple aarch64 -run-pass=aarch64-postlegalizer-lowering -verify-machineinstrs %s -o - | FileCheck %s
+
+# Test that we use G_CONSTANT when we're storing 0, and can use a zero register.
+
+...
+---
+name:            s32
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: s32
+    ; CHECK: liveins: $x0
+    ; CHECK: %ptr:_(p0) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s32) = G_CONSTANT i32 0
+    ; CHECK: G_STORE [[C]](s32), %ptr(p0) :: (store 4)
+    ; CHECK: RET_ReallyLR
+    %ptr:_(p0) = COPY $x0
+    %zero:_(s32) = G_FCONSTANT float 0.000000e+00
+    G_STORE %zero(s32), %ptr(p0) :: (store 4)
+    RET_ReallyLR
+
+...
+---
+name:            s64
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; CHECK-LABEL: name: s64
+    ; CHECK: liveins: $x0
+    ; CHECK: %ptr:_(p0) = COPY $x0
+    ; CHECK: [[C:%[0-9]+]]:_(s64) = G_CONSTANT i64 0
+    ; CHECK: G_STORE [[C]](s64), %ptr(p0) :: (store 8)
+    ; CHECK: RET_ReallyLR
+    %ptr:_(p0) = COPY $x0
+    %zero:_(s64) = G_FCONSTANT double 0.000000e+00
+    G_STORE %zero(s64), %ptr(p0) :: (store 8)
+    RET_ReallyLR
+
+...
+---
+name:            s16
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; We should be able to handle this with full FP16, but right now, we don't.
+
+    ; CHECK-LABEL: name: s16
+    ; CHECK: liveins: $x0
+    ; CHECK: %ptr:_(p0) = COPY $x0
+    ; CHECK: %zero:_(s16) = G_FCONSTANT half 0xH0000
+    ; CHECK: G_STORE %zero(s16), %ptr(p0) :: (store 2)
+    ; CHECK: RET_ReallyLR
+    %ptr:_(p0) = COPY $x0
+    %zero:_(s16) = G_FCONSTANT half 0.000000e+00
+    G_STORE %zero(s16), %ptr(p0) :: (store 2)
+    RET_ReallyLR
+
+...
+---
+name:            negative_zero
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0
+
+    ; Only positive zero should be handled here.
+
+    ; CHECK-LABEL: name: negative_zero
+    ; CHECK: liveins: $x0
+    ; CHECK: %ptr:_(p0) = COPY $x0
+    ; CHECK: %zero:_(s32) = G_FCONSTANT float -0.000000e+00
+    ; CHECK: G_STORE %zero(s32), %ptr(p0) :: (store 4)
+    ; CHECK: RET_ReallyLR
+    %ptr:_(p0) = COPY $x0
+    %zero:_(s32) = G_FCONSTANT float -0.000000e+00
+    G_STORE %zero(s32), %ptr(p0) :: (store 4)
+    RET_ReallyLR