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

Benjamin Maxwell via llvm-commits llvm-commits at lists.llvm.org
Fri Jun 12 09:17:56 PDT 2026


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/5] [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/5] 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;

>From 7a5aa74666817af67cfacffd815dc365f10850d7 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 12 Jun 2026 09:28:12 +0000
Subject: [PATCH 3/5] Fixups

---
 llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 4f4b626ab807b..8a42da97db799 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -5032,6 +5032,8 @@ bool AArch64DAGToDAGISel::trySelectXAR(SDNode *N) {
   return true;
 }
 
+/// Returns a copy from WZR or XZR. This can be used during instruction
+/// selection (it does not require any further selection/legalization).
 static SDValue getZeroRegister(SelectionDAG &DAG, SDLoc DL, EVT VT) {
   if (VT == MVT::i32)
     return DAG.getCopyFromReg(DAG.getEntryNode(), DL, AArch64::WZR, MVT::i32);
@@ -5123,8 +5125,7 @@ void AArch64DAGToDAGISel::Select(SDNode *Node) {
     // the coalescer to propagate these into other instructions.
     ConstantSDNode *ConstNode = cast<ConstantSDNode>(Node);
     if (ConstNode->isZero() && (VT == MVT::i32 || VT == MVT::i64))
-      return ReplaceNode(Node,
-                         getZeroRegister(*CurDAG, SDLoc(Node), VT).getNode());
+      ReplaceNode(Node, getZeroRegister(*CurDAG, SDLoc(Node), VT).getNode());
     break;
   }
 

>From 4cfba5f645ebaaeb71033f3cf1875082b70a4efe Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 12 Jun 2026 10:12:11 +0000
Subject: [PATCH 4/5] Add return back...

---
 llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index 8a42da97db799..e38d10ed7d8db 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -5124,8 +5124,10 @@ 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() && (VT == MVT::i32 || VT == MVT::i64))
+    if (ConstNode->isZero() && (VT == MVT::i32 || VT == MVT::i64)) {
       ReplaceNode(Node, getZeroRegister(*CurDAG, SDLoc(Node), VT).getNode());
+      return;
+    }
     break;
   }
 

>From 5f601aaaa7c4b65c6da3cb39421e6086aa4b8988 Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 12 Jun 2026 16:17:38 +0000
Subject: [PATCH 5/5] Fixups

---
 llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
index e38d10ed7d8db..97a8bfeedca78 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelDAGToDAG.cpp
@@ -5035,10 +5035,9 @@ bool AArch64DAGToDAGISel::trySelectXAR(SDNode *N) {
 /// Returns a copy from WZR or XZR. This can be used during instruction
 /// selection (it does not require any further selection/legalization).
 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);
+  assert(VT == MVT::i32 || VT == MVT::i64);
+  return DAG.getCopyFromReg(DAG.getEntryNode(), DL,
+                            VT == MVT::i32 ? AArch64::WZR : AArch64::XZR, VT);
 }
 
 void AArch64DAGToDAGISel::Select(SDNode *Node) {



More information about the llvm-commits mailing list