[llvm] [AArch64][NFC] Extract a `LowerNTStore()` method (PR #176872)

[Tomer Shafir via llvm-commits]

https://github.com/tomershafir created https://github.com/llvm/llvm-project/pull/176872

Extracts nontermporal store lowering from `LowerSTORE` into its own private method `LowerNTStore`. The refactoring improves overall code structure (and prepares for later scaling of custom STNP lowering logic).

>From cfff768a25b6873394ac5094f808391859cabdba Mon Sep 17 00:00:00 2001
From: tomershafir <tomer.shafir8 at gmail.com>
Date: Sun, 18 Jan 2026 15:06:59 +0200
Subject: [PATCH] [AArch64][NFC] Extract a `LowerNTStore()` method

Extracts nontermporal store lowering from `LowerSTORE` into its own private method `LowerNTStore`. The refactoring improves overall code structure (and prepares for later scaling of custom STNP lowering logic).
---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 76 ++++++++++++-------
 llvm/lib/Target/AArch64/AArch64ISelLowering.h |  2 +
 2 files changed, 49 insertions(+), 29 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index c81f457898d7c..73b5474152592 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -7335,6 +7335,13 @@ SDValue AArch64TargetLowering::LowerSTORE(SDValue Op,
   EVT VT = Value.getValueType();
   EVT MemVT = StoreNode->getMemoryVT();
 
+  // Currently, STNP lowering can only either keep or increase code size, thus
+  // we predicate it to not apply when optimizing for code size.
+  if (StoreNode->isNonTemporal() && !DAG.shouldOptForSize()) {
+    if (auto MaybeSTNP = LowerNTStore(StoreNode, VT, MemVT, Dl, DAG))
+      return MaybeSTNP;
+  }
+
   if (VT.isVector()) {
     if (useSVEForFixedLengthVectorVT(
             VT,
@@ -7354,55 +7361,66 @@ SDValue AArch64TargetLowering::LowerSTORE(SDValue Op,
         MemVT == MVT::v4i8) {
       return LowerTruncateVectorStore(Dl, StoreNode, VT, MemVT, DAG);
     }
+  } else if (MemVT == MVT::i128 && StoreNode->isVolatile()) {
+    return LowerStore128(Op, DAG);
+  } else if (MemVT == MVT::i64x8) {
+    SDValue Value = StoreNode->getValue();
+    assert(Value->getValueType(0) == MVT::i64x8);
+    SDValue Chain = StoreNode->getChain();
+    SDValue Base = StoreNode->getBasePtr();
+    EVT PtrVT = Base.getValueType();
+    for (unsigned i = 0; i < 8; i++) {
+      SDValue Part = DAG.getNode(AArch64ISD::LS64_EXTRACT, Dl, MVT::i64,
+                                 Value, DAG.getConstant(i, Dl, MVT::i32));
+      SDValue Ptr = DAG.getNode(ISD::ADD, Dl, PtrVT, Base,
+                                DAG.getConstant(i * 8, Dl, PtrVT));
+      Chain = DAG.getStore(Chain, Dl, Part, Ptr, StoreNode->getPointerInfo(),
+                           StoreNode->getBaseAlign());
+    }
+    return Chain;
+  }
+
+  return SDValue();
+}
+
+// Lower non-temporal stores that would otherwise be broken by legalization.
+SDValue AArch64TargetLowering::LowerNTStore(StoreSDNode *StoreNode, EVT VT,
+                                            EVT MemVT, const SDLoc &DL,
+                                            SelectionDAG &DAG) const {
+  assert(StoreNode && "Expected a store operation");
+  assert(StoreNode->isNonTemporal() && "Expected a non-temporal store");
+
+  // Currently we only support NT stores lowering for little-endian targets.
+  if (!Subtarget->isLittleEndian())
+    return SDValue();
+
+  if (VT.isVector()) {
     // 256 bit non-temporal stores can be lowered to STNP. Do this as part of
     // the custom lowering, as there are no un-paired non-temporal stores and
     // legalization will break up 256 bit inputs.
-    //
-    // Currently, STNP lowering can only either keep or increase code size, thus
-    // we predicate it to not apply when optimizing for code size.
     ElementCount EC = MemVT.getVectorElementCount();
-    if (StoreNode->isNonTemporal() && MemVT.getSizeInBits() == 256u &&
-        EC.isKnownEven() && DAG.getDataLayout().isLittleEndian() &&
-        !DAG.shouldOptForSize() &&
+    if (VT.isVector() && MemVT.getSizeInBits() == 256u && EC.isKnownEven() &&
         (MemVT.getScalarSizeInBits() == 8u ||
          MemVT.getScalarSizeInBits() == 16u ||
          MemVT.getScalarSizeInBits() == 32u ||
          MemVT.getScalarSizeInBits() == 64u)) {
       SDValue Lo =
-          DAG.getNode(ISD::EXTRACT_SUBVECTOR, Dl,
+          DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL,
                       MemVT.getHalfNumVectorElementsVT(*DAG.getContext()),
-                      StoreNode->getValue(), DAG.getConstant(0, Dl, MVT::i64));
+                      StoreNode->getValue(), DAG.getConstant(0, DL, MVT::i64));
       SDValue Hi =
-          DAG.getNode(ISD::EXTRACT_SUBVECTOR, Dl,
+          DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL,
                       MemVT.getHalfNumVectorElementsVT(*DAG.getContext()),
                       StoreNode->getValue(),
-                      DAG.getConstant(EC.getKnownMinValue() / 2, Dl, MVT::i64));
+                      DAG.getConstant(EC.getKnownMinValue() / 2, DL, MVT::i64));
       SDValue Result = DAG.getMemIntrinsicNode(
-          AArch64ISD::STNP, Dl, DAG.getVTList(MVT::Other),
+          AArch64ISD::STNP, DL, DAG.getVTList(MVT::Other),
           {StoreNode->getChain(), DAG.getBitcast(MVT::v2i64, Lo),
            DAG.getBitcast(MVT::v2i64, Hi), StoreNode->getBasePtr()},
           StoreNode->getMemoryVT(), StoreNode->getMemOperand());
       return Result;
     }
-  } else if (MemVT == MVT::i128 && StoreNode->isVolatile()) {
-    return LowerStore128(Op, DAG);
-  } else if (MemVT == MVT::i64x8) {
-    SDValue Value = StoreNode->getValue();
-    assert(Value->getValueType(0) == MVT::i64x8);
-    SDValue Chain = StoreNode->getChain();
-    SDValue Base = StoreNode->getBasePtr();
-    EVT PtrVT = Base.getValueType();
-    for (unsigned i = 0; i < 8; i++) {
-      SDValue Part = DAG.getNode(AArch64ISD::LS64_EXTRACT, Dl, MVT::i64,
-                                 Value, DAG.getConstant(i, Dl, MVT::i32));
-      SDValue Ptr = DAG.getNode(ISD::ADD, Dl, PtrVT, Base,
-                                DAG.getConstant(i * 8, Dl, PtrVT));
-      Chain = DAG.getStore(Chain, Dl, Part, Ptr, StoreNode->getPointerInfo(),
-                           StoreNode->getBaseAlign());
-    }
-    return Chain;
   }
-
   return SDValue();
 }
 
diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.h b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
index 8dacd08ddf8a7..dec34b2923fcf 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.h
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.h
@@ -614,6 +614,8 @@ class AArch64TargetLowering : public TargetLowering {
 
   SDValue LowerLOAD(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerSTORE(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerNTStore(StoreSDNode *, EVT VT, EVT MemVT, const SDLoc &DL,
+                       SelectionDAG &DAG) const;
   SDValue LowerStore128(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerABS(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerFMUL(SDValue Op, SelectionDAG &DAG) const;