[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

   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;
   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(OriginalLoad->getAlign(), 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 = Align();
   }
+
+  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)

llvm/test/CodeGen/X86/vecloadextract.ll

define i32 @const_index(<8 x i32>* %v) {
   ret i32 %b
 }
 
 ; CHECK: name: variable_index
 ; CHECK:  bb.0 (%ir-block.0):
 ; CHECK:    [[INDEX:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.0, 1, $noreg, 0, $noreg :: (load (s32) from %fixed-stack.0)
 ; CHECK:    [[MASKED_INDEX:%[0-9]+]]:gr32_nosp = AND32ri8 [[INDEX]], 7, implicit-def dead $eflags
 ; CHECK:    [[POINTER:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.1, 1, $noreg, 0, $noreg :: (load (s32) from %fixed-stack.1)
-; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[MASKED_INDEX]], 0, $noreg :: (load (s32))
+; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[MASKED_INDEX]], 0, $noreg :: (load (s32), align 1)
 ; CHECK:    $eax = COPY [[LOAD]]
 ; CHECK:    RET 0, $eax
 define i32 @variable_index(<8 x i32>* %v, i32 %i) {
   %a = load <8 x i32>, <8 x i32>* %v
   %b = extractelement <8 x i32> %a, i32 %i
   ret i32 %b
 }
 
 ; CHECK: name: variable_index_with_addrspace
 ; CHECK:  bb.0 (%ir-block.0):
 ; CHECK:    [[INDEX:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.0, 1, $noreg, 0, $noreg :: (load (s32) from %fixed-stack.0)
 ; CHECK:    [[MASKED_INDEX:%[0-9]+]]:gr32_nosp = AND32ri8 [[INDEX]], 7, implicit-def dead $eflags
 ; CHECK:    [[POINTER:%[0-9]+]]:gr32 = MOV32rm %fixed-stack.1, 1, $noreg, 0, $noreg :: (load (s32) from %fixed-stack.1)
-; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[MASKED_INDEX]], 0, $noreg :: (load (s32), addrspace 1)
+; CHECK:    [[LOAD:%[0-9]+]]:gr32 = MOV32rm killed [[POINTER]], 4, killed [[MASKED_INDEX]], 0, $noreg :: (load (s32), align 1, addrspace 1)
 ; CHECK:    $eax = COPY [[LOAD]]
 ; CHECK:    RET 0, $eax
 define i32 @variable_index_with_addrspace(<8 x i32> addrspace(1)* %v, i32 %i) {
   %a = load <8 x i32>, <8 x i32> addrspace(1)* %v
   %b = extractelement <8 x i32> %a, i32 %i
   ret i32 %b
 }