[DAG] scalarizeExtractedVectorLoad - replace getABITypeAlign with allowsMemoryAccess (PR45116)

[RKSimon]

Summary

One of the cases identified in PR45116 - we don't need to limit extracted loads to ABI alignment, we can use allowsMemoryAccess - which tests using getABITypeAlign, but also checks if a target permits (fast) misaligned memory loads by checking allowsMisalignedMemoryAccesses as a fallback.

I've also cleaned up the alignment calculation code - if we have a constant extraction index then the alignment can be based on an offset from the original vector load alignment, but for non-constant indices we should assume the worst (single element alignment only).

[dmgreen]

This sounds OK to me, providing I'm understanding these alignments vs base alignments correctly. LGTM

[RKSimon]

Thanks @dmgreen - I'll keep this open for a few days in case anyone else has any further comments.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::scalarizeExtractedVectorLoad(SDNode *EVE, EVT InVecVT,
   EVT VecEltVT = InVecVT.getVectorElementType();
 
   // If the vector element type is not a multiple of a byte then we are unable
   // to correctly compute an address to load only the extracted element as a
   // scalar.
   if (!VecEltVT.isByteSized())
     return SDValue();
 
-  Align Alignment = OriginalLoad->getAlign();
-  Align NewAlign = DAG.getDataLayout().getABITypeAlign(
-      VecEltVT.getTypeForEVT(*DAG.getContext()));
-
-  if (NewAlign > Alignment ||
-      !TLI.isOperationLegalOrCustom(ISD::LOAD, VecEltVT))
-    return SDValue();
-
-  ISD::LoadExtType ExtTy = ResultVT.bitsGT(VecEltVT) ?
-    ISD::NON_EXTLOAD : ISD::EXTLOAD;
-  if (!TLI.shouldReduceLoadWidth(OriginalLoad, ExtTy, VecEltVT))
+  ISD::LoadExtType ExtTy =
+      ResultVT.bitsGT(VecEltVT) ? ISD::NON_EXTLOAD : ISD::EXTLOAD;
+  if (!TLI.isOperationLegalOrCustom(ISD::LOAD, VecEltVT) ||
+      !TLI.shouldReduceLoadWidth(OriginalLoad, ExtTy, VecEltVT))
     return SDValue();
 
-  Alignment = NewAlign;
-
+  Align Alignment = OriginalLoad->getAlign();
   MachinePointerInfo MPI;
   SDLoc DL(EVE);
   if (auto *ConstEltNo = dyn_cast<ConstantSDNode>(EltNo)) {
     int Elt = ConstEltNo->getZExtValue();
     unsigned PtrOff = VecEltVT.getSizeInBits() * Elt / 8;
     MPI = OriginalLoad->getPointerInfo().getWithOffset(PtrOff);
+    Alignment = commonAlignment(Alignment, PtrOff);
   } else {
     // Discard the pointer info except the address space because the memory
     // operand can't represent this new access since the offset is variable.
     MPI = MachinePointerInfo(OriginalLoad->getPointerInfo().getAddrSpace());
+    Alignment = commonAlignment(Alignment, VecEltVT.getSizeInBits() / 8);
   }
+
+  bool IsFast = false;
+  if (!TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), VecEltVT,
+                              OriginalLoad->getAddressSpace(), Alignment,
+                              OriginalLoad->getMemOperand()->getFlags(),
+                              &IsFast) ||
+      !IsFast)
+    return SDValue();
+
   SDValue NewPtr = TLI.getVectorElementPointer(DAG, OriginalLoad->getBasePtr(),
                                                InVecVT, EltNo);
 
   // The replacement we need to do here is a little tricky: we need to
   // replace an extractelement of a load with a load.
   // Use ReplaceAllUsesOfValuesWith to do the replacement.
   // Note that this replacement assumes that the extractvalue is the only
   // use of the load; that's okay because we don't want to perform this

llvm/test/CodeGen/AArch64/arm64-indexed-vector-ldst.ll

; CHECK-NEXT:    ret
   %lv = load <4 x i32>, <4 x i32>* %A
   %e = extractelement <4 x i32> %lv, i32 %idx
   ret i32 %e
 }
 
 define i32 @load_single_extract_variable_index_v3i32_small_align(<3 x i32>* %A, i32 %idx) {
 ; CHECK-LABEL: load_single_extract_variable_index_v3i32_small_align:
 ; CHECK:       ; %bb.0:
-; CHECK-NEXT:    sub sp, sp, #16
-; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    ldr d0, [x0]
-; CHECK-NEXT:    add x8, x0, #8
 ; CHECK-NEXT:    ; kill: def $w1 killed $w1 def $x1
-; CHECK-NEXT:    mov x9, sp
-; CHECK-NEXT:    ld1.s { v0 }[2], [x8]
-; CHECK-NEXT:    and x8, x1, #0x3
-; CHECK-NEXT:    bfi x9, x8, #2, #2
-; CHECK-NEXT:    str q0, [sp]
-; CHECK-NEXT:    ldr w0, [x9]
-; CHECK-NEXT:    add sp, sp, #16
+; CHECK-NEXT:    sxtw x8, w1
+; CHECK-NEXT:    cmp x8, #2
+; CHECK-NEXT:    mov w9, #2
+; CHECK-NEXT:    csel x8, x8, x9, lo
+; CHECK-NEXT:    ldr w0, [x0, x8, lsl #2]
 ; CHECK-NEXT:    ret
   %lv = load <3 x i32>, <3 x i32>* %A, align 2
   %e = extractelement <3 x i32> %lv, i32 %idx
   ret i32 %e
 }
 
 define i32 @load_single_extract_variable_index_v3i32_default_align(<3 x i32>* %A, i32 %idx) {
 ; CHECK-LABEL: load_single_extract_variable_index_v3i32_default_align:

llvm/test/CodeGen/X86/merge-consecutive-stores-nt.ll

 ; X86-SSE2-NEXT:    movd %xmm0, %ecx
 ; X86-SSE2-NEXT:    movntil %ecx, 20(%eax)
 ; X86-SSE2-NEXT:    retl
 ;
 ; X86-SSE4A-LABEL: merge_2_v4f32_align1_ntstore:
 ; X86-SSE4A:       # %bb.0:
 ; X86-SSE4A-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-SSE4A-NEXT:    movl {{[0-9]+}}(%esp), %ecx
-; X86-SSE4A-NEXT:    movups (%ecx), %xmm0
-; X86-SSE4A-NEXT:    movups 16(%ecx), %xmm1
+; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm0 = mem[0],zero
+; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm1 = mem[0],zero
 ; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm2 = mem[0],zero
 ; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm3 = mem[0],zero
-; X86-SSE4A-NEXT:    movntsd %xmm2, 8(%eax)
 ; X86-SSE4A-NEXT:    movntsd %xmm0, (%eax)
+; X86-SSE4A-NEXT:    movntsd %xmm1, 8(%eax)
 ; X86-SSE4A-NEXT:    movntsd %xmm3, 24(%eax)
-; X86-SSE4A-NEXT:    movntsd %xmm1, 16(%eax)
+; X86-SSE4A-NEXT:    movntsd %xmm2, 16(%eax)
 ; X86-SSE4A-NEXT:    retl
 ;
 ; X64-SSE2-LABEL: merge_2_v4f32_align1_ntstore:
 ; X64-SSE2:       # %bb.0:
 ; X64-SSE2-NEXT:    movdqu (%rdi), %xmm0
 ; X64-SSE2-NEXT:    movdqu 16(%rdi), %xmm1
 ; X64-SSE2-NEXT:    movq %xmm0, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, (%rsi)
 ; X64-SSE2-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[2,3,2,3]
 ; X64-SSE2-NEXT:    movq %xmm0, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, 8(%rsi)
 ; X64-SSE2-NEXT:    movq %xmm1, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, 16(%rsi)
 ; X64-SSE2-NEXT:    pshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
 ; X64-SSE2-NEXT:    movq %xmm0, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, 24(%rsi)
 ; X64-SSE2-NEXT:    retq
 ;
 ; X64-SSE4A-LABEL: merge_2_v4f32_align1_ntstore:
 ; X64-SSE4A:       # %bb.0:
-; X64-SSE4A-NEXT:    movups (%rdi), %xmm0
-; X64-SSE4A-NEXT:    movups 16(%rdi), %xmm1
+; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm0 = mem[0],zero
+; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm1 = mem[0],zero
 ; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm2 = mem[0],zero
 ; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm3 = mem[0],zero
-; X64-SSE4A-NEXT:    movntsd %xmm2, 8(%rsi)
 ; X64-SSE4A-NEXT:    movntsd %xmm0, (%rsi)
+; X64-SSE4A-NEXT:    movntsd %xmm1, 8(%rsi)
 ; X64-SSE4A-NEXT:    movntsd %xmm3, 24(%rsi)
-; X64-SSE4A-NEXT:    movntsd %xmm1, 16(%rsi)
+; X64-SSE4A-NEXT:    movntsd %xmm2, 16(%rsi)
 ; X64-SSE4A-NEXT:    retq
 ;
 ; X64-SSE41-LABEL: merge_2_v4f32_align1_ntstore:
 ; X64-SSE41:       # %bb.0:
 ; X64-SSE41-NEXT:    movdqu (%rdi), %xmm0
 ; X64-SSE41-NEXT:    movdqu 16(%rdi), %xmm1
 ; X64-SSE41-NEXT:    pextrq $1, %xmm0, %rax
 ; X64-SSE41-NEXT:    movntiq %rax, 8(%rsi)

 ; X86-SSE2-NEXT:    movd %xmm0, %ecx
 ; X86-SSE2-NEXT:    movntil %ecx, 20(%eax)
 ; X86-SSE2-NEXT:    retl
 ;
 ; X86-SSE4A-LABEL: merge_2_v4f32_align1:
 ; X86-SSE4A:       # %bb.0:
 ; X86-SSE4A-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; X86-SSE4A-NEXT:    movl {{[0-9]+}}(%esp), %ecx
-; X86-SSE4A-NEXT:    movups (%ecx), %xmm0
-; X86-SSE4A-NEXT:    movups 16(%ecx), %xmm1
+; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm0 = mem[0],zero
+; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm1 = mem[0],zero
 ; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm2 = mem[0],zero
 ; X86-SSE4A-NEXT:    movsd {{.*#+}} xmm3 = mem[0],zero
-; X86-SSE4A-NEXT:    movntsd %xmm2, 8(%eax)
 ; X86-SSE4A-NEXT:    movntsd %xmm0, (%eax)
+; X86-SSE4A-NEXT:    movntsd %xmm1, 8(%eax)
 ; X86-SSE4A-NEXT:    movntsd %xmm3, 24(%eax)
-; X86-SSE4A-NEXT:    movntsd %xmm1, 16(%eax)
+; X86-SSE4A-NEXT:    movntsd %xmm2, 16(%eax)
 ; X86-SSE4A-NEXT:    retl
 ;
 ; X64-SSE2-LABEL: merge_2_v4f32_align1:
 ; X64-SSE2:       # %bb.0:
 ; X64-SSE2-NEXT:    movdqu (%rdi), %xmm0
 ; X64-SSE2-NEXT:    movdqu 16(%rdi), %xmm1
 ; X64-SSE2-NEXT:    movq %xmm0, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, (%rsi)
 ; X64-SSE2-NEXT:    pshufd {{.*#+}} xmm0 = xmm0[2,3,2,3]
 ; X64-SSE2-NEXT:    movq %xmm0, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, 8(%rsi)
 ; X64-SSE2-NEXT:    movq %xmm1, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, 16(%rsi)
 ; X64-SSE2-NEXT:    pshufd {{.*#+}} xmm0 = xmm1[2,3,2,3]
 ; X64-SSE2-NEXT:    movq %xmm0, %rax
 ; X64-SSE2-NEXT:    movntiq %rax, 24(%rsi)
 ; X64-SSE2-NEXT:    retq
 ;
 ; X64-SSE4A-LABEL: merge_2_v4f32_align1:
 ; X64-SSE4A:       # %bb.0:
-; X64-SSE4A-NEXT:    movups (%rdi), %xmm0
-; X64-SSE4A-NEXT:    movups 16(%rdi), %xmm1
+; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm0 = mem[0],zero
+; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm1 = mem[0],zero
 ; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm2 = mem[0],zero
 ; X64-SSE4A-NEXT:    movsd {{.*#+}} xmm3 = mem[0],zero
-; X64-SSE4A-NEXT:    movntsd %xmm2, 8(%rsi)
 ; X64-SSE4A-NEXT:    movntsd %xmm0, (%rsi)
+; X64-SSE4A-NEXT:    movntsd %xmm1, 8(%rsi)
 ; X64-SSE4A-NEXT:    movntsd %xmm3, 24(%rsi)
-; X64-SSE4A-NEXT:    movntsd %xmm1, 16(%rsi)
+; X64-SSE4A-NEXT:    movntsd %xmm2, 16(%rsi)
 ; X64-SSE4A-NEXT:    retq
 ;
 ; X64-SSE41-LABEL: merge_2_v4f32_align1:
 ; X64-SSE41:       # %bb.0:
 ; X64-SSE41-NEXT:    movdqu (%rdi), %xmm0
 ; X64-SSE41-NEXT:    movdqu 16(%rdi), %xmm1
 ; X64-SSE41-NEXT:    pextrq $1, %xmm0, %rax
 ; X64-SSE41-NEXT:    movntiq %rax, 8(%rsi)