[AArch64] Add support for 256-bit non temporal loads

[zjaffal]

Summary

Currenlty all temporal loads are mapped to LDP or LDR. This patch will map all the non temporal 256-bit loads into LDNP. Future patches should address other non-temporal loads.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

795

Could you add a comment similar to the one for the stores above?

806

nit: please add a newline before this comment.

20425

Would be good to have a comment here to describe what's lowered here.

20440

leftover commented-out code?

llvm/lib/Target/AArch64/AArch64InstrInfo.td

2588

nit: newline here

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

20426

Why is EC.isKnownEven needed?

20427

This looks like it could drop a set of brackets.

[zjaffal]

Remove even checks and add comments

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

20426

I followed the same code that is used to lower STNP but I think it is not really necessary here.
I think we can remove it.

[dmgreen]

LGTM if @fhahn doesn't know of a reason to add the isKnownEven.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

20426

Yeah - I agree so long as it is checking the size of the MemVT and it's scalar element size.

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

20420–20421

I think this assertion doesn't really make sense any longer, because we now enable the custom lowering for additional types. This probably needs to became just a check (i.e add an early exit here if the cownciditon doesn't hold). At the moment this assertion triggers in multiple tests, causing them to fail.

[fhahn]

In D131773#3725335, @dmgreen wrote:

LGTM if @fhahn doesn't know of a reason to add the isKnownEven.

I can't think of any reason why this should be needed now and also couldn't find anything when looking at the original review. @zjaffal it would be great if you could post a follow-up patch that removes isKnownEven from the checks for STNP or even better move the checks to a helper function used by both STNP and LDNP handling code.

[zjaffal]

Fix for failing tests

[fhahn]

LGTM, thanks!

llvm/lib/Target/AArch64/AArch64ISelLowering.h

enum NodeType : unsigned {
   ST4LANEpost,
 
   STG,
   STZG,
   ST2G,
   STZ2G,
 
   LDP,
+  LDNP,
   STP,
   STNP,
 
   // Memory Operations
   MOPS_MEMSET,
   MOPS_MEMSET_TAGGING,
   MOPS_MEMCOPY,
   MOPS_MEMMOVE,

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

#undef LCALLNAME5
   setOperationAction(ISD::STORE, MVT::v32i8, Custom);
   setOperationAction(ISD::STORE, MVT::v16i16, Custom);
   setOperationAction(ISD::STORE, MVT::v16f16, Custom);
   setOperationAction(ISD::STORE, MVT::v8i32, Custom);
   setOperationAction(ISD::STORE, MVT::v8f32, Custom);
   setOperationAction(ISD::STORE, MVT::v4f64, Custom);
   setOperationAction(ISD::STORE, MVT::v4i64, Custom);
 
+  // 256 bit non-temporal loads can be lowered to LDNP. This is done using
+  // custom lowering, as there are no un-paired non-temporal loads legalization
+  // will break up 256 bit inputs.
+  setOperationAction(ISD::LOAD, MVT::v32i8, Custom);
+  setOperationAction(ISD::LOAD, MVT::v16i16, Custom);
+  setOperationAction(ISD::LOAD, MVT::v16f16, Custom);
+  setOperationAction(ISD::LOAD, MVT::v8i32, Custom);
+  setOperationAction(ISD::LOAD, MVT::v8f32, Custom);
+  setOperationAction(ISD::LOAD, MVT::v4f64, Custom);
+  setOperationAction(ISD::LOAD, MVT::v4i64, Custom);
+
   // Lower READCYCLECOUNTER using an mrs from PMCCNTR_EL0.
   // This requires the Performance Monitors extension.
   if (Subtarget->hasPerfMon())
     setOperationAction(ISD::READCYCLECOUNTER, MVT::i64, Legal);
 
   if (getLibcallName(RTLIB::SINCOS_STRET_F32) != nullptr &&
       getLibcallName(RTLIB::SINCOS_STRET_F64) != nullptr) {
     // Issue __sincos_stret if available.
     setOperationAction(ISD::FSINCOS, MVT::f64, Custom);
case AArch64ISD::FIRST_NUMBER:
     MAKE_CASE(AArch64ISD::SST1_SXTW_PRED)
     MAKE_CASE(AArch64ISD::SST1_UXTW_PRED)
     MAKE_CASE(AArch64ISD::SST1_SXTW_SCALED_PRED)
     MAKE_CASE(AArch64ISD::SST1_UXTW_SCALED_PRED)
     MAKE_CASE(AArch64ISD::SST1_IMM_PRED)
     MAKE_CASE(AArch64ISD::SSTNT1_PRED)
     MAKE_CASE(AArch64ISD::SSTNT1_INDEX_PRED)
     MAKE_CASE(AArch64ISD::LDP)
+    MAKE_CASE(AArch64ISD::LDNP)
     MAKE_CASE(AArch64ISD::STP)
     MAKE_CASE(AArch64ISD::STNP)
     MAKE_CASE(AArch64ISD::BITREVERSE_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::BSWAP_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::REVH_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::REVW_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::REVD_MERGE_PASSTHRU)
     MAKE_CASE(AArch64ISD::CTLZ_MERGE_PASSTHRU)
void AArch64TargetLowering::ReplaceNodeResults(
   case ISD::STRICT_FP_TO_UINT:
     assert(N->getValueType(0) == MVT::i128 && "unexpected illegal conversion");
     // Let normal code take care of it by not adding anything to Results.
     return;
   case ISD::ATOMIC_CMP_SWAP:
     ReplaceCMP_SWAP_128Results(N, Results, DAG, Subtarget);
     return;
   case ISD::ATOMIC_LOAD:
   case ISD::LOAD: {
-    assert(SDValue(N, 0).getValueType() == MVT::i128 &&
-           "unexpected load's value type");
     MemSDNode *LoadNode = cast<MemSDNode>(N);
+    EVT MemVT = LoadNode->getMemoryVT();
+    // Handle lowering 256 bit non temporal loads into LDNP.
+    if (LoadNode->isNonTemporal() && MemVT.getSizeInBits() == 256u &&
+        (MemVT.getScalarSizeInBits() == 8u ||
+         MemVT.getScalarSizeInBits() == 16u ||
+         MemVT.getScalarSizeInBits() == 32u ||
+         MemVT.getScalarSizeInBits() == 64u)) {
+
+      SDValue Result = DAG.getMemIntrinsicNode(
+          AArch64ISD::LDNP, SDLoc(N),
+          DAG.getVTList({MemVT.getHalfNumVectorElementsVT(*DAG.getContext()),
+                         MemVT.getHalfNumVectorElementsVT(*DAG.getContext()),
+                         MVT::Other}),
+          {LoadNode->getChain(), LoadNode->getBasePtr()},
+          LoadNode->getMemoryVT(), LoadNode->getMemOperand());
+
+      SDValue Pair = DAG.getNode(ISD::CONCAT_VECTORS, SDLoc(N), MemVT,
+                                 Result.getValue(0), Result.getValue(1));
+      Results.append({Pair, Result.getValue(2) /* Chain */});
+      return;
+    }
+
     if ((!LoadNode->isVolatile() && !LoadNode->isAtomic()) ||
         LoadNode->getMemoryVT() != MVT::i128) {
       // Non-volatile or atomic loads are optimized later in AArch64's load/store
       // optimizer.
       return;
     }
 
+    if (SDValue(N, 0).getValueType() == MVT::i128) {
       SDValue Result = DAG.getMemIntrinsicNode(
           AArch64ISD::LDP, SDLoc(N),
           DAG.getVTList({MVT::i64, MVT::i64, MVT::Other}),
-        {LoadNode->getChain(), LoadNode->getBasePtr()}, LoadNode->getMemoryVT(),
-        LoadNode->getMemOperand());
+          {LoadNode->getChain(), LoadNode->getBasePtr()},
+          LoadNode->getMemoryVT(), LoadNode->getMemOperand());
 
       SDValue Pair = DAG.getNode(ISD::BUILD_PAIR, SDLoc(N), MVT::i128,
                                  Result.getValue(0), Result.getValue(1));
       Results.append({Pair, Result.getValue(2) /* Chain */});
+    }
     return;
   }
   case ISD::EXTRACT_SUBVECTOR:
     ReplaceExtractSubVectorResults(N, Results, DAG);
     return;
   case ISD::INSERT_SUBVECTOR:
   case ISD::CONCAT_VECTORS:
     // Custom lowering has been requested for INSERT_SUBVECTOR and

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 def SDT_AArch64ITOF  : SDTypeProfile<1, 1, [SDTCisFP<0>, SDTCisSameAs<0,1>]>;
 
 def SDT_AArch64TLSDescCall : SDTypeProfile<0, -2, [SDTCisPtrTy<0>,
                                                  SDTCisPtrTy<1>]>;
 
 def SDT_AArch64uaddlp : SDTypeProfile<1, 1, [SDTCisVec<0>, SDTCisVec<1>]>;
 
 def SDT_AArch64ldp : SDTypeProfile<2, 1, [SDTCisVT<0, i64>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
+def SDT_AArch64ldnp : SDTypeProfile<2, 1, [SDTCisVT<0, v4i32>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
 def SDT_AArch64stp : SDTypeProfile<0, 3, [SDTCisVT<0, i64>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
 def SDT_AArch64stnp : SDTypeProfile<0, 3, [SDTCisVT<0, v4i32>, SDTCisSameAs<0, 1>, SDTCisPtrTy<2>]>;
 
 // Generates the general dynamic sequences, i.e.
 //  adrp  x0, :tlsdesc:var
 //  ldr   x1, [x0, #:tlsdesc_lo12:var]
 //  add   x0, x0, #:tlsdesc_lo12:var
 //  .tlsdesccall var
def SDT_AArch64unpk : SDTypeProfile<1, 1, [
     SDTCisInt<0>, SDTCisInt<1>, SDTCisOpSmallerThanOp<1, 0>
 ]>;
 def AArch64sunpkhi : SDNode<"AArch64ISD::SUNPKHI", SDT_AArch64unpk>;
 def AArch64sunpklo : SDNode<"AArch64ISD::SUNPKLO", SDT_AArch64unpk>;
 def AArch64uunpkhi : SDNode<"AArch64ISD::UUNPKHI", SDT_AArch64unpk>;
 def AArch64uunpklo : SDNode<"AArch64ISD::UUNPKLO", SDT_AArch64unpk>;
 
 def AArch64ldp : SDNode<"AArch64ISD::LDP", SDT_AArch64ldp, [SDNPHasChain, SDNPMayLoad, SDNPMemOperand]>;
+def AArch64ldnp : SDNode<"AArch64ISD::LDNP", SDT_AArch64ldnp, [SDNPHasChain, SDNPMayLoad, SDNPMemOperand]>;
 def AArch64stp : SDNode<"AArch64ISD::STP", SDT_AArch64stp, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;
 def AArch64stnp : SDNode<"AArch64ISD::STNP", SDT_AArch64stnp, [SDNPHasChain, SDNPMayStore, SDNPMemOperand]>;
 
 def AArch64tbl : SDNode<"AArch64ISD::TBL", SDT_AArch64TBL>;
 def AArch64mrs : SDNode<"AArch64ISD::MRS",
                         SDTypeProfile<1, 1, [SDTCisVT<0, i64>, SDTCisVT<1, i32>]>,
                         [SDNPHasChain, SDNPOutGlue]>;
 

 defm LDNPX : LoadPairNoAlloc<0b10, 0, GPR64z, simm7s8, "ldnp">;
 defm LDNPS : LoadPairNoAlloc<0b00, 1, FPR32Op, simm7s4, "ldnp">;
 defm LDNPD : LoadPairNoAlloc<0b01, 1, FPR64Op, simm7s8, "ldnp">;
 defm LDNPQ : LoadPairNoAlloc<0b10, 1, FPR128Op, simm7s16, "ldnp">;
 
 def : Pat<(AArch64ldp (am_indexed7s64 GPR64sp:$Rn, simm7s8:$offset)),
           (LDPXi GPR64sp:$Rn, simm7s8:$offset)>;
 
+def : Pat<(AArch64ldnp (am_indexed7s128 GPR64sp:$Rn, simm7s16:$offset)),
+          (LDNPQi GPR64sp:$Rn, simm7s16:$offset)>;
 //---
 // (register offset)
 //---
 
 // Integer
 defm LDRBB : Load8RO<0b00,  0, 0b01, GPR32, "ldrb", i32, zextloadi8>;
 defm LDRHH : Load16RO<0b01, 0, 0b01, GPR32, "ldrh", i32, zextloadi16>;
 defm LDRW  : Load32RO<0b10, 0, 0b01, GPR32, "ldr", i32, load>;
 defm LDRX  : Load64RO<0b11, 0, 0b01, GPR64, "ldr", i64, load>;

llvm/test/CodeGen/AArch64/nontemporal-load.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple aarch64-apple-darwin | FileCheck %s
 
 define <4 x double> @test_ldnp_v4f64(<4 x double>* %A) {
 ; CHECK-LABEL: test_ldnp_v4f64:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <4 x double>, <4 x double>* %A, align 8, !nontemporal !0
   ret <4 x double> %lv
 }
 
 define <4 x i64> @test_ldnp_v4i64(<4 x i64>* %A) {
 ; CHECK-LABEL: test_ldnp_v4i64:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <4 x i64>, <4 x i64>* %A, align 8, !nontemporal !0
   ret <4 x i64> %lv
 }
 
 define <8 x i32> @test_ldnp_v8i32(<8 x i32>* %A) {
 ; CHECK-LABEL: test_ldnp_v8i32:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <8 x i32>, <8 x i32>* %A, align 8, !nontemporal !0
   ret <8 x i32> %lv
 }
 
 define <8 x float> @test_ldnp_v8f32(<8 x float>* %A) {
 ; CHECK-LABEL: test_ldnp_v8f32:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <8 x float>, <8 x float>* %A, align 8, !nontemporal !0
   ret <8 x float> %lv
 }
 
 define <16 x i16> @test_ldnp_v16i16(<16 x i16>* %A) {
 ; CHECK-LABEL: test_ldnp_v16i16:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <16 x i16>, <16 x i16>* %A, align 8, !nontemporal !0
   ret <16 x i16> %lv
 }
 
 define <16 x half> @test_ldnp_v16f16(<16 x half>* %A) {
 ; CHECK-LABEL: test_ldnp_v16f16:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <16 x half>, <16 x half>* %A, align 8, !nontemporal !0
   ret <16 x half> %lv
 }
 
 define <32 x i8> @test_ldnp_v32i8(<32 x i8>* %A) {
 ; CHECK-LABEL: test_ldnp_v32i8:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
 ; CHECK-NEXT:    ret
   %lv = load <32 x i8>, <32 x i8>* %A, align 8, !nontemporal !0
   ret <32 x i8> %lv
 }
 
 define <4 x i32> @test_ldnp_v4i32(<4 x i32>* %A) {
 ; CHECK-LABEL: test_ldnp_v4i32:
 ; CHECK:       ; %bb.0:

 ; CHECK-NEXT:    ret
   %lv = load <1 x i64>, <1 x i64>* %A, align 8, !nontemporal !0
   ret <1 x i64> %lv
 }
 
 define <32 x i16> @test_ldnp_v32i16(<32 x i16>* %A) {
 ; CHECK-LABEL: test_ldnp_v32i16:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
-; CHECK-NEXT:    ldp q2, q3, [x0, #32]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #32]
 ; CHECK-NEXT:    ret
   %lv = load <32 x i16>, <32 x i16>* %A, align 8, !nontemporal !0
   ret <32 x i16> %lv
 }
 
 define <32 x half> @test_ldnp_v32f16(<32 x half>* %A) {
 ; CHECK-LABEL: test_ldnp_v32f16:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
-; CHECK-NEXT:    ldp q2, q3, [x0, #32]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #32]
 ; CHECK-NEXT:    ret
   %lv = load <32 x half>, <32 x half>* %A, align 8, !nontemporal !0
   ret <32 x half> %lv
 }
 
 define <16 x i32> @test_ldnp_v16i32(<16 x i32>* %A) {
 ; CHECK-LABEL: test_ldnp_v16i32:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
-; CHECK-NEXT:    ldp q2, q3, [x0, #32]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #32]
 ; CHECK-NEXT:    ret
   %lv = load <16 x i32>, <16 x i32>* %A, align 8, !nontemporal !0
   ret <16 x i32> %lv
 }
 
 define <16 x float> @test_ldnp_v16f32(<16 x float>* %A) {
 ; CHECK-LABEL: test_ldnp_v16f32:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
-; CHECK-NEXT:    ldp q2, q3, [x0, #32]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #32]
 ; CHECK-NEXT:    ret
   %lv = load <16 x float>, <16 x float>* %A, align 8, !nontemporal !0
   ret <16 x float> %lv
 }
 
 define <17 x float> @test_ldnp_v17f32(<17 x float>* %A) {
 ; CHECK-LABEL: test_ldnp_v17f32:
 ; CHECK:       ; %bb.0:

 ; CHECK-NEXT:    ret
   %lv = load<5 x double>, <5 x double>* %A, align 8, !nontemporal !0
   ret <5 x double> %lv
 }
 
 define <16 x i64> @test_ldnp_v16i64(<16 x i64>* %A) {
 ; CHECK-LABEL: test_ldnp_v16i64:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
-; CHECK-NEXT:    ldp q2, q3, [x0, #32]
-; CHECK-NEXT:    ldp q4, q5, [x0, #64]
-; CHECK-NEXT:    ldp q6, q7, [x0, #96]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #32]
+; CHECK-NEXT:    ldnp q4, q5, [x0, #64]
+; CHECK-NEXT:    ldnp q6, q7, [x0, #96]
 ; CHECK-NEXT:    ret
   %lv = load <16 x i64>, <16 x i64>* %A, align 8, !nontemporal !0
   ret <16 x i64> %lv
 }
 
 define <16 x double> @test_ldnp_v16f64(<16 x double>* %A) {
 ; CHECK-LABEL: test_ldnp_v16f64:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q1, [x0]
-; CHECK-NEXT:    ldp q2, q3, [x0, #32]
-; CHECK-NEXT:    ldp q4, q5, [x0, #64]
-; CHECK-NEXT:    ldp q6, q7, [x0, #96]
+; CHECK-NEXT:    ldnp q0, q1, [x0]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #32]
+; CHECK-NEXT:    ldnp q4, q5, [x0, #64]
+; CHECK-NEXT:    ldnp q6, q7, [x0, #96]
 ; CHECK-NEXT:    ret
   %lv = load <16 x double>, <16 x double>* %A, align 8, !nontemporal !0
   ret <16 x double> %lv
 }
 
 
 !0 = !{i32 1}