[AArch64TTI] AArch64 supports NT vector stores through STNP.

[fhahn]

Summary

This patch adds a custom implementation of isLegalNTStore to AArch64TTI
that supports vector types that can be directly stored by STNP. Note
that the implementation may not catch all valid cases (e.g. because the
vector is a multiple of 256 and could be broken down to multiple valid 256 bit
stores), but it is good enough for LV to vectorize loops with NT stores,
as LV only passes in a vector with 2 elements to check. LV seems to also
be the only user of isLegalNTStore.

We should also do the same for NT loads, but before that we need to
ensure that we properly lower LDNP of vectors, similar to D72919.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

[fhahn]

Add missing bits of diff.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt

[dmgreen]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

185

Floats too, at least for Neon types?

188–189

Are these because the vectorizer will run this with:

// Arbitrarily try a vector of 2 elements.
Type *VecTy = VectorType::get(T, /*NumElements=*/2);
...
if (!TTI->isLegalNTStore(VecTy, *Alignment))

And decides at that point if it is legal or not? Without knowing real vector widths.

191

isLegalMaskedStore->isLegalNTStore

It seems that the base version of isLegalNTStore just checks:

// By default, assume nontemporal memory stores are available for stores
// that are aligned and have a size that is a power of 2.
unsigned DataSize = DL.getTypeStoreSize(DataType);
return Alignment >= DataSize && isPowerOf2_32(DataSize);

I'm guessing that it is the alignment that is causing problems for these larger vectors?

[fhahn]

Generalize checks, add additional tests.

[fhahn]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

188–189

I've generalised the checks, but yes, that is the reason. I think what LV does is relatively reasonable. LV should only create power of 2 vectors and if <2 x ty> is legal, any vector it creates can be broken down to <2 x ty> I think.

191

I don't think the default implementation is relevant for AArch64. I could not find anything in ArmARM that would indicate any alignment requirements on STNP/LDNP addresses. I think we should probably handle non-vector types here similar to vector types, but probably best done as follow-on change (also, I don't think there are any non-vector users currently in tree)

[dmgreen]

OK. I think this is good. I honestly wouldn't object to just returning isPowerOf2(DataTypeSize), but the code you have looks better. The real answer, as far as I can tell is any scalar/vector that is a multiple of 64. Like you said with any alignment.

But the vectorizer calls this with kind of silly values, so we have to do the best we can. Anything else will just be turned into a normal store anyway, which isn't particularly bad and I don't think should block vectorization (especially if we start with a store size that would not end up being be non-temporal anyway!)

LGTM. Maybe add a comment about how the vectorizer will call the function with a vector width of 2 (in case that changed in the future, we can see where this logic came from).

[fhahn]

Thanks, I'll add a note about the LV use before committing.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

public:
   bool isLegalMaskedLoad(Type *DataType, MaybeAlign Alignment) {
     return isLegalMaskedLoadStore(DataType, Alignment);
   }
 
   bool isLegalMaskedStore(Type *DataType, MaybeAlign Alignment) {
     return isLegalMaskedLoadStore(DataType, Alignment);
   }
 
+  bool isLegalNTStore(Type *DataType, Align Alignment) {
+    // NOTE: The logic below is mostly geared towards LV, which calls it with
+    //       vectors with 2 elements. We might want to improve that, if other
+    //       users show up.
+    // Nontemporal vector stores can be directly lowered to STNP, if the vector
+    // can be halved so that each half fits into a register. That's the case if
+    // the element type fits into a register and the number of elements is a
+    // power of 2 > 1.
+    if (isa<VectorType>(DataType)) {
+      unsigned NumElements = DataType->getVectorNumElements();
+      unsigned EltSize =
+          DataType->getVectorElementType()->getScalarSizeInBits();
+      return NumElements > 1 && isPowerOf2_64(NumElements) && EltSize >= 8 &&
+             EltSize <= 128 && isPowerOf2_64(EltSize);
+    }
+    return BaseT::isLegalNTStore(DataType, Alignment);
+  }
+
   int getInterleavedMemoryOpCost(unsigned Opcode, Type *VecTy, unsigned Factor,
                                  ArrayRef<unsigned> Indices, unsigned Alignment,
                                  unsigned AddressSpace,
                                  bool UseMaskForCond = false,
                                  bool UseMaskForGaps = false);
 
   bool
   shouldConsiderAddressTypePromotion(const Instruction &I,

llvm/test/Transforms/LoopVectorize/AArch64/nontemporal-load-store.ll

+; RUN: opt -loop-vectorize -mtriple=arm64-apple-iphones -force-vector-width=4 -force-vector-interleave=1 %s -S | FileCheck %s
+
+; Vectors with i4 elements may not legal with nontemporal stores.
+define void @test_i4_store(i4* %ddst) {
+; CHECK-LABEL: define void @test_i4_store(
+; CHECK-NOT:   vector.body:
+; CHECK:        ret void
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i4* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i4, i4* %ddst.addr, i64 1
+  store i4 10, i4* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i8_store(i8* %ddst) {
+; CHECK-LABEL: define void @test_i8_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <4 x i8> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i8* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i8, i8* %ddst.addr, i64 1
+  store i8 10, i8* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_half_store(half* %ddst) {
+; CHECK-LABEL: define void @test_half_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <4 x half> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi half* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds half, half* %ddst.addr, i64 1
+  store half 10.0, half* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i16_store(i16* %ddst) {
+; CHECK-LABEL: define void @test_i16_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <4 x i16> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i16* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i16, i16* %ddst.addr, i64 1
+  store i16 10, i16* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i32_store(i32* nocapture %ddst) {
+; CHECK-LABEL: define void @test_i32_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <16 x i32> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i32* [ %ddst, %entry ], [ %incdec.ptr3, %for.body ]
+  %incdec.ptr = getelementptr inbounds i32, i32* %ddst.addr, i64 1
+  store i32 10, i32* %ddst.addr, align 4, !nontemporal !8
+  %incdec.ptr1 = getelementptr inbounds i32, i32* %ddst.addr, i64 2
+  store i32 20, i32* %incdec.ptr, align 4, !nontemporal !8
+  %incdec.ptr2 = getelementptr inbounds i32, i32* %ddst.addr, i64 3
+  store i32 30, i32* %incdec.ptr1, align 4, !nontemporal !8
+  %incdec.ptr3 = getelementptr inbounds i32, i32* %ddst.addr, i64 4
+  store i32 40, i32* %incdec.ptr2, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i33_store(i33* nocapture %ddst) {
+; CHECK-LABEL: define void @test_i33_store(
+; CHECK-NOT:   vector.body:
+; CHECK:         ret
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i33* [ %ddst, %entry ], [ %incdec.ptr3, %for.body ]
+  %incdec.ptr = getelementptr inbounds i33, i33* %ddst.addr, i64 1
+  store i33 10, i33* %ddst.addr, align 4, !nontemporal !8
+  %incdec.ptr1 = getelementptr inbounds i33, i33* %ddst.addr, i64 2
+  store i33 20, i33* %incdec.ptr, align 4, !nontemporal !8
+  %incdec.ptr2 = getelementptr inbounds i33, i33* %ddst.addr, i64 3
+  store i33 30, i33* %incdec.ptr1, align 4, !nontemporal !8
+  %incdec.ptr3 = getelementptr inbounds i33, i33* %ddst.addr, i64 4
+  store i33 40, i33* %incdec.ptr2, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 3
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i40_store(i40* nocapture %ddst) {
+; CHECK-LABEL: define void @test_i40_store(
+; CHECK-NOT:   vector.body:
+; CHECK:         ret
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i40* [ %ddst, %entry ], [ %incdec.ptr3, %for.body ]
+  %incdec.ptr = getelementptr inbounds i40, i40* %ddst.addr, i64 1
+  store i40 10, i40* %ddst.addr, align 4, !nontemporal !8
+  %incdec.ptr1 = getelementptr inbounds i40, i40* %ddst.addr, i64 2
+  store i40 20, i40* %incdec.ptr, align 4, !nontemporal !8
+  %incdec.ptr2 = getelementptr inbounds i40, i40* %ddst.addr, i64 3
+  store i40 30, i40* %incdec.ptr1, align 4, !nontemporal !8
+  %incdec.ptr3 = getelementptr inbounds i40, i40* %ddst.addr, i64 4
+  store i40 40, i40* %incdec.ptr2, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 3
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+define void @test_i64_store(i64* nocapture %ddst) local_unnamed_addr #0 {
+; CHECK-LABEL: define void @test_i64_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <4 x i64> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i64* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i64, i64* %ddst.addr, i64 1
+  store i64 10, i64* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_double_store(double* %ddst) {
+; CHECK-LABEL: define void @test_double_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <4 x double> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi double* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds double, double* %ddst.addr, i64 1
+  store double 10.0, double* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i128_store(i128* %ddst) {
+; CHECK-LABEL: define void @test_i128_store(
+; CHECK-LABEL: vector.body:
+; CHECK:         store <4 x i128> {{.*}} !nontemporal !0
+; CHECK:         br
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i128* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i128, i128* %ddst.addr, i64 1
+  store i128 10, i128* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+define void @test_i256_store(i256* %ddst) {
+; CHECK-LABEL: define void @test_i256_store(
+; CHECK-NOT:   vector.body:
+; CHECK:        ret void
+;
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %entry, %for.body
+  %i = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %ddst.addr = phi i256* [ %ddst, %entry ], [ %incdec.ptr, %for.body ]
+  %incdec.ptr = getelementptr inbounds i256, i256* %ddst.addr, i64 1
+  store i256 10, i256* %ddst.addr, align 4, !nontemporal !8
+  %add = add nuw nsw i32 %i, 4
+  %cmp = icmp ult i32 %i, 4092
+  br i1 %cmp, label %for.body, label %for.cond.cleanup
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+}
+
+!8 = !{i32 1}