[llvm] [AArch64] Roll out zero-register substitution to GPR32/GPR64 (PR #145239)

[via llvm-commits]

https://github.com/AZero13 updated https://github.com/llvm/llvm-project/pull/145239

>From d007e357bc26f35663ea16225398933ac869e023 Mon Sep 17 00:00:00 2001
From: Rose <gfunni234 at gmail.com>
Date: Sun, 22 Jun 2025 11:09:40 -0400
Subject: [PATCH] [AArch64] Roll out zero-register substitution to GPR32/GPR64

No more duplicate StoreUI/StoreUIz multiclasses.
---
 .../lib/Target/AArch64/AArch64InstrFormats.td | 20 +++----------------
 llvm/lib/Target/AArch64/AArch64InstrInfo.td   |  8 ++++----
 .../lib/Target/AArch64/AArch64RegisterInfo.td |  6 ++++--
 3 files changed, 11 insertions(+), 23 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrFormats.td b/llvm/lib/Target/AArch64/AArch64InstrFormats.td
index ba7cbccc0bcd6..0287a2fb48ccd 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrFormats.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrFormats.td
@@ -3761,6 +3761,9 @@ multiclass LoadUI<bits<2> sz, bit V, bits<2> opc, DAGOperand regtype,
                   (!cast<Instruction>(NAME # "ui") regtype:$Rt, GPR64sp:$Rn, 0)>;
 }
 
+// Store with positive immediate. Accepts any DAGOperand type.
+// When RegisterOperand with GIZeroRegister is used, GlobalISel can automatically
+// substitute zero-immediates with the specified zero-register.
 multiclass StoreUI<bits<2> sz, bit V, bits<2> opc, DAGOperand regtype,
              Operand indextype, string asm, list<dag> pattern> {
   let AddedComplexity = 10, mayLoad = 0, mayStore = 1, hasSideEffects = 0 in
@@ -3773,23 +3776,6 @@ multiclass StoreUI<bits<2> sz, bit V, bits<2> opc, DAGOperand regtype,
                   (!cast<Instruction>(NAME # "ui") regtype:$Rt, GPR64sp:$Rn, 0)>;
 }
 
-// Same as StoreUI, but take a RegisterOperand. This is used by GlobalISel to
-// substitute zero-registers automatically.
-//
-// TODO: Roll out zero-register substitution to GPR32/GPR64 and fold this back
-//       into StoreUI.
-multiclass StoreUIz<bits<2> sz, bit V, bits<2> opc, RegisterOperand regtype,
-             Operand indextype, string asm, list<dag> pattern> {
-  let AddedComplexity = 10, mayLoad = 0, mayStore = 1, hasSideEffects = 0 in
-  def ui : BaseLoadStoreUI<sz, V, opc, (outs),
-                           (ins regtype:$Rt, GPR64sp:$Rn, indextype:$offset),
-                           asm, pattern>,
-           Sched<[WriteST]>;
-
-  def : InstAlias<asm # "\t$Rt, [$Rn]",
-                  (!cast<Instruction>(NAME # "ui") regtype:$Rt, GPR64sp:$Rn, 0)>;
-}
-
 def PrefetchOperand : AsmOperandClass {
   let Name = "Prefetch";
   let ParserMethod = "tryParsePrefetch";
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
index 0f3f24f0853c9..7e528611834eb 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td
@@ -4537,10 +4537,10 @@ let AddedComplexity = 19 in {
 
 //---
 // (unsigned immediate)
-defm STRX : StoreUIz<0b11, 0, 0b00, GPR64z, uimm12s8, "str",
+defm STRX : StoreUI<0b11, 0, 0b00, GPR64z, uimm12s8, "str",
                    [(store GPR64z:$Rt,
                             (am_indexed64 GPR64sp:$Rn, uimm12s8:$offset))]>;
-defm STRW : StoreUIz<0b10, 0, 0b00, GPR32z, uimm12s4, "str",
+defm STRW : StoreUI<0b10, 0, 0b00, GPR32z, uimm12s4, "str",
                     [(store GPR32z:$Rt,
                             (am_indexed32 GPR64sp:$Rn, uimm12s4:$offset))]>;
 let Predicates = [HasFPARMv8] in {
@@ -4559,11 +4559,11 @@ defm STRD : StoreUI<0b11, 1, 0b00, FPR64Op, uimm12s8, "str",
 defm STRQ : StoreUI<0b00, 1, 0b10, FPR128Op, uimm12s16, "str", []>;
 }
 
-defm STRHH : StoreUIz<0b01, 0, 0b00, GPR32z, uimm12s2, "strh",
+defm STRHH : StoreUI<0b01, 0, 0b00, GPR32z, uimm12s2, "strh",
                      [(truncstorei16 GPR32z:$Rt,
                                      (am_indexed16 GPR64sp:$Rn,
                                                    uimm12s2:$offset))]>;
-defm STRBB : StoreUIz<0b00, 0, 0b00, GPR32z, uimm12s1,  "strb",
+defm STRBB : StoreUI<0b00, 0, 0b00, GPR32z, uimm12s1,  "strb",
                      [(truncstorei8 GPR32z:$Rt,
                                     (am_indexed8 GPR64sp:$Rn,
                                                  uimm12s1:$offset))]>;
diff --git a/llvm/lib/Target/AArch64/AArch64RegisterInfo.td b/llvm/lib/Target/AArch64/AArch64RegisterInfo.td
index 61bf87fe71441..34f3a79ae6292 100644
--- a/llvm/lib/Target/AArch64/AArch64RegisterInfo.td
+++ b/llvm/lib/Target/AArch64/AArch64RegisterInfo.td
@@ -236,8 +236,10 @@ def GPR64sp0 : RegisterOperand<GPR64sp> {
   let ParserMatchClass = GPR64spPlus0Operand;
 }
 
-// GPR32/GPR64 but with zero-register substitution enabled.
-// TODO: Roll this out to GPR32/GPR64/GPR32all/GPR64all.
+// GPR32/GPR64 with zero-register substitution enabled.
+// The separate StoreUIz multiclass has been removed - StoreUI now works with both
+// DAGOperand and RegisterOperand types, enabling zero-register substitution when RegisterOperand
+// with GIZeroRegister is used while maintaining interface flexibility.
 def GPR32z : RegisterOperand<GPR32> {
   let GIZeroRegister = WZR;
 }