[AArch64] break non-temporal loads over 256 into 256-loads and a smaller load

[zjaffal]

Summary

Currently over 256 non-temporal loads are broken inefficently. For example, v17i32 gets broken into 2 128-bit loads. It is better if we can use
256-bit loads instead.

[fhahn]

Thanks for sharing the patch. As is, it looks like there are multiple test failures that are caused by this patch and need addressing (some comments inline about correctness issues)

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

53

Is this needed?

73

Is this needed?

17820

move all variables that are used after the early exit down after the early exit.

17825

We should also only do this for non-temporal loads.

17829

Could you add a comment here illustrating what kind of DAG nodes we create to replace the original load?

It would also be good to document the motivation for special handling for non-temporal loads here.

17836

I don't think that's correct, you are passing the offset in bits, but I think it should be in bytes.

Looking at some of the test changes, the offsets of the loads are wrong, e.g. ldnp q0, q2, [x0, #256].

17867

Do you need to create NullVector explicitly here? Could you just use DAG.degUNDEF(NewVT) here?

19860

This should be kept I think.

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17829

Would the motivation for the special handling for non-temporal loads be something like:
We have 256-bit non-temporal load so it might be good to utilise them when having large load instructions?

17867

Yes I will do that

[zjaffal]

Fix the issues related to the scalable loads.

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17825

Checks are added now for non-temporal loads

17836

This should be fixed now

[fhahn]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17829

We have 256-bit non-temporal load so it might be good to utilise them when having large load instructions?

Yes something like that, maybe: Try to beak up non-temporal stores into blocks of 256 bits early, so the LDNPQ can be selected.

Might be good to add this as comment for the whole function.

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17829

You can check the comments I added. I will add a comment for the whole function as well

[zjaffal]

Add comments for function

[fhahn]

Thanks for the update, this looks like a nice improvement and should avoid regressions for those cases with D132559. A few more mostly stylistic comments.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17813

It would be good to be more concrete what large means here (>256 bits) will be split into blocks of 256bits.

nit: comments should be full sentences ending with a period (.)

17819

Thinking about it a bit more, I think we should also not handle volatile and atomic loads here to be safe. Could you add tests and the extra conditions for early exit here? I think it would be fine to just add the tests directly in the patch.

17825

LD->isNonTemporal() is already checked above

17830

nit: BasePtr would be more accurate.

17841

nit: used as unsigned, so maybe make this unsigned too?

17851

nit: comments should be full sentences ending with a period (.)

bits in the load operation -> bits of the load operation?

It might also help with readability if you add a newline before the comment.

17854

nit: comments should be full sentences ending with a period (.)

17867

MaskVector not used?

17868

UndefVector would be more accurate?

[zjaffal]

• Add early exit for volatile loads and a test cases to check for it.
• Address some style issues

[zjaffal]

Removing checking if both operands are negative at the same time since it is handled by default.

[fhahn]

It looks like you uploaded the wrong patch here?

[zjaffal]

Fix non-temporal clause

[zjaffal]

readd the volatile test

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17863

Shouldn't the second operand of this call be PtrOffset again?

17870–17872

Is this (and the implementation generally) big-endian correct? I don't know the answer here, I can never remember what's supposed to happen. But someone should definitely try it on an aarch64_be target and at least eyeball the assembly to check the offsets and so on.

17884

The Chain here is the input chain, I think you need to TokenFactor all the loads' output chains together to make sure nothing gets reordered with them.

19928–19929

This looks like it takes precedence over performLOADCombine and disables ldnp formation if the TBI feature is enabled.

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

19928–19929

Removing the checks and calling performLOADCombine caused many tests to fail. Maybe we can check if the load is non-temporal here ?

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17863

Yes it should be, it is worth noting that the offsets on the generated assembly didn't change when I changed from using MemVT.getSizeInBits()) to PtrOffset

[zjaffal]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17870–17872

This is the generated assembly for big-endian
using the following test case

define <17 x float> @test_ldnp_v17f32(<17 x float>* %A) {
  %lv = load<17 x float>, <17 x float>* %A, align 8, !nontemporal !0
  ret <17 x float> %lv
}

!0 = !{i32 1}

test_ldnp_v17f32:                       // @test_ldnp_v17f32
	.cfi_startproc
// %bb.0:
	ldnp	q0, q1, [x0, #32]
	add	x9, x8, #48
	add	x10, x8, #32
	ldnp	q2, q3, [x0]
	add	x11, x8, #16
	ldr	s4, [x0, #64]
	st1	{ v2.4s }, [x8]
	st1	{ v1.4s }, [x9]
	st1	{ v0.4s }, [x10]
	st1	{ v3.4s }, [x11]
	str	s4, [x8, #64]
	ret

Looking at https://godbolt.org I think there are more load instructions before breaking them

[t.p.northover]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17870–17872

It's coming back to me now, and I think that output is incorrect.

Basically, the ld1 and st1 instructions would load each element big-endian but preserve the order (i.e. v0[0] is still loaded from &ptr[0] not &ptr[3]. On the other hand ldr and ldnp load the whole vector-register in big-endian order which loads v0[0] from &ptr[3].

This output mixes the two so it'll rearrange elements.

But I think the real bug is in https://reviews.llvm.org/rG7155ed4289 that you committed earlier, or at least that needs fixing before any bug here is obvious.

A proper fix would be to put an AArch64ISD::REV<n> (where <n> is the element size) after each ldnp to restore what an ld1 would have done, though I'd be OK with disabling the optimization for big-endian instead. I care if we break it, I'm less bothered if it isn't as fast as little-endian.

19928–19929

The check is needed because otherwise we assume TBI even when we shouldn't, but it probably shouldn't cause an early return. Both optimizations should get the opportunity to run if they might be applicable.

[zjaffal]

1. Change Align to use PtrOffset
2. Make sure that performTBISimplification optimisation can run
3. Add TokenFactor node

[zjaffal]

rebase on top of main

[fhahn]

LGTM, thanks for the latest changes. Some small stylistic comments inline.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17813

nit: I think this is easier to read if you keep nontemporal loads together. Also, period at end of sentence.

17839

nit: utilise -> utilize (American spelling)

17856

nit: add newline after here, to separate blocks.

17858

nit: UNDEF vector instead of null value vector?

17880

add newline before the new function/

17880

nit: use ConcatVT to match other names here.

[zjaffal]

Address stylistic comments

[fhahn]

Thanks for the update. 2 more small comments, but I'll fix them before committing the change on your behalf.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

17813

Looks like there's another typo: nontermporal -> nontemporal

17856

The newline was meant to go after the } to separate it from the next block, not before it.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RuntimeLibcalls.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/CodeGen/SelectionDAGNodes.h"
 #include "llvm/CodeGen/TargetCallingConv.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/ValueTypes.h"
 #include "llvm/IR/Attributes.h"
+#include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GetElementPtrTypeIterator.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsAArch64.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/OperandTraits.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/Value.h"
+#include "llvm/IR/VectorBuilder.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/InstructionCost.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MachineValueType.h"
 #include "llvm/Support/MathExtras.h"
+#include "llvm/Support/TypeSize.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include <algorithm>
 #include <bitset>
 #include <cassert>
 #include <cctype>
 #include <cstdint>
#undef LCALLNAME5
   setTargetDAGCombine(ISD::SETCC);
 
   setTargetDAGCombine(ISD::INTRINSIC_WO_CHAIN);
 
   setTargetDAGCombine({ISD::ANY_EXTEND, ISD::ZERO_EXTEND, ISD::SIGN_EXTEND,
                        ISD::VECTOR_SPLICE, ISD::SIGN_EXTEND_INREG,
                        ISD::CONCAT_VECTORS, ISD::EXTRACT_SUBVECTOR,
                        ISD::INSERT_SUBVECTOR, ISD::STORE, ISD::BUILD_VECTOR});
-  if (Subtarget->supportsAddressTopByteIgnored())
-    setTargetDAGCombine(ISD::LOAD);
 
   setTargetDAGCombine(ISD::MSTORE);
 
   setTargetDAGCombine(ISD::MUL);
 
   setTargetDAGCombine({ISD::SELECT, ISD::VSELECT});
 
   setTargetDAGCombine({ISD::INTRINSIC_VOID, ISD::INTRINSIC_W_CHAIN,
if (Store->getMemoryVT() != Orig.getValueType())
       return SDValue();
     return DAG.getStore(Store->getChain(), SDLoc(Store), Orig,
                         Store->getBasePtr(), Store->getMemOperand());
   }
 
   return SDValue();
 }
 
+static SDValue performLOADCombine(SDNode *N,
+                                  TargetLowering::DAGCombinerInfo &DCI,
+                                  SelectionDAG &DAG,
+                                  const AArch64Subtarget *Subtarget) {
+  LoadSDNode *LD = cast<LoadSDNode>(N);
+  SDLoc DL(LD);
+  EVT MemVT = LD->getMemoryVT();
+  SDValue Chain = LD->getChain();
+  SDValue Ptr = LD->getBasePtr();
+  SDNodeFlags Flags = LD->getFlags();
+
+  SmallVector<SDValue, 4> LoadOps;
+  if (MemVT.getSizeInBits() <= 256 || MemVT.getSizeInBits() % 256 == 0 ||
+      256 % MemVT.getScalarSizeInBits() != 0)
+    return SDValue();
+
+  unsigned PtrOffset = 256;
+  MVT NewVT =
+      MVT::getVectorVT(MemVT.getVectorElementType().getSimpleVT(),
+                       256 / MemVT.getVectorElementType().getSizeInBits());
+  for (; PtrOffset <= MemVT.getSizeInBits(); PtrOffset += 256) {
+    // Build 256-bit loads
+    SDValue NewPtr =
+        DAG.getMemBasePlusOffset(Ptr, TypeSize::Fixed(PtrOffset), DL, Flags);
+    Align NewAlign = commonAlignment(LD->getAlign(), PtrOffset);
+    SDValue NewLoad = DAG.getLoad(
+        NewVT, DL, Chain, NewPtr, LD->getPointerInfo().getWithOffset(PtrOffset),
+        NewAlign, LD->getMemOperand()->getFlags(), LD->getAAInfo());
+    LoadOps.push_back(NewLoad);
+  }
+  // Process remaining bits in the load operation
+  // This is done by creating a null value vector to match the size of the
+  // 256-bit loads and inserting the remaining load to it We extract the
+  // original load type at the end
+  unsigned BitsRemaining = MemVT.getSizeInBits() % 256;
+  PtrOffset = MemVT.getScalarSizeInBits() - BitsRemaining;
+  MVT RemainingVT = MVT::getVectorVT(
+      MemVT.getVectorElementType().getSimpleVT(),
+      BitsRemaining / MemVT.getVectorElementType().getSizeInBits());
+  SDValue NewPtr =
+      DAG.getMemBasePlusOffset(Ptr, TypeSize::Fixed(PtrOffset), DL, Flags);
+  Align NewAlign = commonAlignment(LD->getAlign(), MemVT.getSizeInBits());
+  SDValue RemainingLoad =
+      DAG.getLoad(RemainingVT, DL, Chain, NewPtr,
+                  LD->getPointerInfo().getWithOffset(PtrOffset), NewAlign,
+                  LD->getMemOperand()->getFlags(), LD->getAAInfo());
+  SmallVector<SDValue, 4> MaskVector;
+  for (unsigned i = 0; i < NewVT.getVectorNumElements(); i++) {
+    MaskVector.push_back(DAG.getUNDEF(NewVT.getScalarType()));
+  }
+  ArrayRef<SDValue> SV{MaskVector};
+  SDValue NullVector = DAG.getBuildVector(NewVT, DL, SV);
+  SDValue InsertIdx = DAG.getVectorIdxConstant(0, DL);
+  SDValue ExtendedReminingLoad = DAG.getNode(
+      ISD::INSERT_SUBVECTOR, DL, NewVT, {NullVector, RemainingLoad, InsertIdx});
+  LoadOps.push_back(ExtendedReminingLoad);
+
+  EVT ConcatVectorType =
+      EVT::getVectorVT(*DAG.getContext(), MemVT.getScalarType(),
+                       LoadOps.size() * NewVT.getVectorNumElements());
+  SDValue ConcatVectors =
+      DAG.getNode(ISD::CONCAT_VECTORS, DL, ConcatVectorType, LoadOps);
+  SDValue ExtractSubVector =
+      DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MemVT,
+                  {ConcatVectors, DAG.getVectorIdxConstant(0, DL)});
+  return DAG.getMergeValues({ExtractSubVector, Chain}, DL);
+}
 static SDValue performSTORECombine(SDNode *N,
                                    TargetLowering::DAGCombinerInfo &DCI,
                                    SelectionDAG &DAG,
                                    const AArch64Subtarget *Subtarget) {
   StoreSDNode *ST = cast<StoreSDNode>(N);
   SDValue Chain = ST->getChain();
   SDValue Value = ST->getValue();
   SDValue Ptr = ST->getBasePtr();
 
   // If this is an FP_ROUND followed by a store, fold this into a truncating
   // store. We can do this even if this is already a truncstore.
   // We purposefully don't care about legality of the nodes here as we know
   // they can be split down into something legal.
SDValue AArch64TargetLowering::PerformDAGCombine(SDNode *N,
   case ISD::INSERT_SUBVECTOR:
     return performInsertSubvectorCombine(N, DCI, DAG);
   case ISD::SELECT:
     return performSelectCombine(N, DCI);
   case ISD::VSELECT:
     return performVSelectCombine(N, DCI.DAG);
   case ISD::SETCC:
     return performSETCCCombine(N, DCI, DAG);
   case ISD::LOAD:
-    if (performTBISimplification(N->getOperand(1), DCI, DAG))
-      return SDValue(N, 0);
-    break;
+    return performLOADCombine(N, DCI, DAG, Subtarget);
   case ISD::STORE:
     return performSTORECombine(N, DCI, DAG, Subtarget);
   case ISD::MSTORE:
     return performMSTORECombine(N, DCI, DAG, Subtarget);
   case ISD::MGATHER:
   case ISD::MSCATTER:
     return performMaskedGatherScatterCombine(N, DCI, DAG);
   case ISD::VECTOR_SPLICE:

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

 ; 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:    ldp q1, q2, [x0, #32]
-; CHECK-NEXT:    ldp q3, q4, [x0]
-; CHECK-NEXT:    ldr s0, [x0, #64]
-; CHECK-NEXT:    stp q3, q4, [x8]
-; CHECK-NEXT:    stp q1, q2, [x8, #32]
-; CHECK-NEXT:    str s0, [x8, #64]
+; CHECK-NEXT:    ldnp q0, q1, [x0, #512]
+; CHECK-NEXT:    ldnp q2, q3, [x0, #256]
+; CHECK-NEXT:    ldr s4, [x0]
+; CHECK-NEXT:    stp q0, q1, [x8, #32]
+; CHECK-NEXT:    stp q2, q3, [x8]
+; CHECK-NEXT:    str s4, [x8, #64]
 ; CHECK-NEXT:    ret
   %lv = load <17 x float>, <17 x float>* %A, align 8, !nontemporal !0
   ret <17 x float> %lv
 }
 
 define <33 x double> @test_ldnp_v33f64(<33 x double>* %A) {
 ; CHECK-LABEL: test_ldnp_v33f64:
 ; 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:    ldp q16, q17, [x0, #128]
-; CHECK-NEXT:    ldp q18, q19, [x0, #160]
-; CHECK-NEXT:    ldp q21, q22, [x0, #224]
-; CHECK-NEXT:    ldp q23, q24, [x0, #192]
-; CHECK-NEXT:    ldr d20, [x0, #256]
-; CHECK-NEXT:    stp q0, q1, [x8]
-; CHECK-NEXT:    stp q2, q3, [x8, #32]
-; CHECK-NEXT:    stp q4, q5, [x8, #64]
-; CHECK-NEXT:    str d20, [x8, #256]
-; CHECK-NEXT:    stp q6, q7, [x8, #96]
-; CHECK-NEXT:    stp q16, q17, [x8, #128]
-; CHECK-NEXT:    stp q18, q19, [x8, #160]
-; CHECK-NEXT:    stp q23, q24, [x8, #192]
-; CHECK-NEXT:    stp q21, q22, [x8, #224]
+; CHECK-NEXT:    add x9, x0, #1024
+; CHECK-NEXT:    add x10, x0, #1280
+; CHECK-NEXT:    ldnp q4, q5, [x0, #256]
+; CHECK-NEXT:    ldnp q0, q1, [x9]
+; CHECK-NEXT:    add x9, x0, #1792
+; CHECK-NEXT:    ldnp q2, q3, [x10]
+; CHECK-NEXT:    ldnp q6, q7, [x9]
+; CHECK-NEXT:    add x9, x0, #2048
+; CHECK-NEXT:    ldnp q16, q17, [x0, #512]
+; CHECK-NEXT:    ldnp q18, q19, [x9]
+; CHECK-NEXT:    add x9, x0, #1536
+; CHECK-NEXT:    ldnp q20, q21, [x0, #768]
+; CHECK-NEXT:    ldnp q22, q23, [x9]
+; CHECK-NEXT:    ldr d24, [x0]
+; CHECK-NEXT:    stp q4, q5, [x8]
+; CHECK-NEXT:    stp q16, q17, [x8, #32]
+; CHECK-NEXT:    stp q20, q21, [x8, #64]
+; CHECK-NEXT:    stp q0, q1, [x8, #96]
+; CHECK-NEXT:    stp q2, q3, [x8, #128]
+; CHECK-NEXT:    stp q22, q23, [x8, #160]
+; CHECK-NEXT:    stp q6, q7, [x8, #192]
+; CHECK-NEXT:    stp q18, q19, [x8, #224]
+; CHECK-NEXT:    str d24, [x8, #256]
 ; CHECK-NEXT:    ret
   %lv = load <33 x double>, <33 x double>* %A, align 8, !nontemporal !0
   ret <33 x double> %lv
 }
 
 define <33 x i8> @test_ldnp_v33i8(<33 x i8>* %A) {
 ; CHECK-LABEL: test_ldnp_v33i8:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q1, q0, [x0]
-; CHECK-NEXT:    ldrb w9, [x0, #32]
-; CHECK-NEXT:    stp q1, q0, [x8]
-; CHECK-NEXT:    strb w9, [x8, #32]
+; CHECK-NEXT:    ldnp q0, q1, [x0, #256]
+; CHECK-NEXT:    add x9, x8, #32
+; CHECK-NEXT:    ldr b2, [x0]
+; CHECK-NEXT:    stp q0, q1, [x8]
+; CHECK-NEXT:    st1.b { v2 }[0], [x9]
 ; CHECK-NEXT:    ret
   %lv = load<33 x i8>, <33 x i8>* %A, align 8, !nontemporal !0
   ret <33 x i8> %lv
 }
 
 define <4 x i65> @test_ldnp_v4i65(<4 x i65>* %A) {
 ; CHECK-LABEL: test_ldnp_v4i65:
 ; CHECK:       ; %bb.0:

 ; CHECK-NEXT:    ret
   %lv = load <4 x i63>, <4 x i63>* %A, align 8, !nontemporal !0
   ret <4 x i63> %lv
 }
 
 define <5 x double> @test_ldnp_v5f64(<5 x double>* %A) {
 ; CHECK-LABEL: test_ldnp_v5f64:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    ldp q0, q2, [x0]
+; CHECK-NEXT:    ldnp q0, q2, [x0, #256]
+; CHECK-NEXT:    ldr d4, [x0]
 ; CHECK-NEXT:    ext.16b v1, v0, v0, #8
 ; CHECK-NEXT:    ; kill: def $d0 killed $d0 killed $q0
 ; CHECK-NEXT:    ; kill: def $d1 killed $d1 killed $q1
 ; CHECK-NEXT:    ext.16b v3, v2, v2, #8
-; CHECK-NEXT:    ldr d4, [x0, #32]
 ; CHECK-NEXT:    ; kill: def $d2 killed $d2 killed $q2
 ; CHECK-NEXT:    ; kill: def $d3 killed $d3 killed $q3
-; CHECK-NEXT:    ; kill: def $d4 killed $d4 killed $q4
 ; 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: