[llvm] [AArch64] Avoid creating a new generic constant in SelectSMETileSlice (PR #203344)

[Benjamin Maxwell via llvm-commits]

https://github.com/MacDue updated https://github.com/llvm/llvm-project/pull/203344

>From 79695468c7f17294981e4b64058052f8c20b0675 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Thu, 11 Jun 2026 16:42:58 +0000
Subject: [PATCH 1/2] [AArch64] Avoid creating a new generic constant in
 SelectSMETileSlice

This was creating a new ISD::Constant node during instruction selection,
which may also need lowering (e.g., to a mov/movimm). The issue with
this is the new constant may not end up in the instruction selection
worklist, as this happens after it has been sorted and started
processing nodes. This results it it being directly lowered to an
immediate in some cases.

This patch works around this by directly lowering the constant to a
`CopyFromReg WZR` within SelectSMETileSlice.

Fixes #203295
---
 .../Target/AArch64/AArch64ISelDAGToDAG.cpp    | 25 ++++++++-----------
 .../AArch64/sme-intrinsics-mova-insert.ll     | 11 ++++++++
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 499bb2325186d..886fb89365d5c 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -5032,6 +5032,13 @@ bool AArch64DAGToDAGISel::trySelectXAR(SDNode *N) {
   return true;
 }
 
+static SDValue getZeroRegister(SelectionDAG &DAG, SDLoc DL, EVT VT) {
+  if (VT == MVT::i32)
+    return DAG.getCopyFromReg(DAG.getEntryNode(), DL, AArch64::WZR, MVT::i32);
+  assert(VT == MVT::i64);
+  return DAG.getCopyFromReg(DAG.getEntryNode(), DL, AArch64::XZR, MVT::i64);
+}
+
 void AArch64DAGToDAGISel::Select(SDNode *Node) {
   // If we have a custom node, we already have selected!
   if (Node->isMachineOpcode()) {
@@ -5115,19 +5122,9 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
     // Materialize zero constants as copies from WZR/XZR.  This allows
     // the coalescer to propagate these into other instructions.
     ConstantSDNode *ConstNode = cast<ConstantSDNode>(Node);
-    if (ConstNode->isZero()) {
-      if (VT == MVT::i32) {
-        SDValue New = CurDAG->getCopyFromReg(
-            CurDAG->getEntryNode(), SDLoc(Node), AArch64::WZR, MVT::i32);
-        ReplaceNode(Node, New.getNode());
-        return;
-      } else if (VT == MVT::i64) {
-        SDValue New = CurDAG->getCopyFromReg(
-            CurDAG->getEntryNode(), SDLoc(Node), AArch64::XZR, MVT::i64);
-        ReplaceNode(Node, New.getNode());
-        return;
-      }
-    }
+    if (ConstNode->isZero())
+      return ReplaceNode(Node,
+                         getZeroRegister(*CurDAG, SDLoc(Node), VT).getNode());
     break;
   }
 
@@ -7972,7 +7969,7 @@ bool AArch64DAGToDAGISel::SelectSMETileSlice(SDValue N, unsigned MaxSize,
   };
 
   if (SDValue C = MatchConstantOffset(N)) {
-    Base = CurDAG->getConstant(0, SDLoc(N), MVT::i32);
+    Base = getZeroRegister(*CurDAG, SDLoc(N), MVT::i32);
     Offset = C;
     return true;
   }
diff --git a/llvm/test/CodeGen/AArch64/sme-intrinsics-mova-insert.ll b/llvm/test/CodeGen/AArch64/sme-intrinsics-mova-insert.ll
index 9c5becf6ffbf2..3d9065bda9c26 100644
--- a/llvm/test/CodeGen/AArch64/sme-intrinsics-mova-insert.ll
+++ b/llvm/test/CodeGen/AArch64/sme-intrinsics-mova-insert.ll
@@ -486,6 +486,17 @@ define void @test_add_with_disjoint_or(i64 %idx, <vscale x 4 x i1> %pg) {
   ret void
 }
 
+; Test we materialize the zero register base within a GPR when we only have an immediate slice offset.
+define void @test_imm_offset_with_zero_register_base(<vscale x 4 x i1> %pg, <vscale x 4 x i32> %vec) {
+; CHECK-LABEL: test_imm_offset_with_zero_register_base:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov w12, wzr
+; CHECK-NEXT:    mov za0h.s[w12, 1], p0/m, z0.s
+; CHECK-NEXT:    ret
+  call void @llvm.aarch64.sme.write.horiz.nxv4i32(i32 0, i32 1, <vscale x 4 x i1> %pg, <vscale x 4 x i32> %vec)
+  ret void
+}
+
 declare void @llvm.aarch64.sme.write.horiz.nxv16i8(i32, i32, <vscale x 16 x i1>, <vscale x 16 x i8>)
 declare void @llvm.aarch64.sme.write.horiz.nxv8i16(i32, i32, <vscale x 8 x i1>, <vscale x 8 x i16>)
 declare void @llvm.aarch64.sme.write.horiz.nxv8f16(i32, i32, <vscale x 8 x i1>, <vscale x 8 x half>)

>From 34b8cd21e10eb95ab09e7a057ba3b6c55ece6167 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Thu, 11 Jun 2026 17:06:26 +0000
Subject: [PATCH 2/2] Preserve check

---
 llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 886fb89365d5c..4f4b626ab807b 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -5122,7 +5122,7 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
     // Materialize zero constants as copies from WZR/XZR.  This allows
     // the coalescer to propagate these into other instructions.
     ConstantSDNode *ConstNode = cast<ConstantSDNode>(Node);
-    if (ConstNode->isZero())
+    if (ConstNode->isZero() && (VT == MVT::i32 || VT == MVT::i64))
       return ReplaceNode(Node,
                          getZeroRegister(*CurDAG, SDLoc(Node), VT).getNode());
     break;