[X86][SSE] Improve extraction of elements from v16i8 (pre-SSE41)

[RKSimon]

Summary

Without SSE41 (pextrb) we currently extract byte elements from a vector by spilling to stack and reloading the byte.

This patch is an initial attempt at using MOVD/PEXTRW to extract the relevant DWORD/WORD from the vector and then shift+truncate to collect the correct byte.

Extraction of multiple bytes this way results in code bloat, but as explained in the patch we could probably afford to be more aggressive with the supported extractions before again falling back on spilling - possibly through counting the number of extracts and which DWORD/WORD they originate?

[RKSimon]

ping?

[spatel]

LGTM.

llvm/trunk/lib/Target/X86/X86ISelLowering.cpp

SDValue Assert  = DAG.getNode(ISD::AssertZext, dl, MVT::i32, Extract,
                                   DAG.getValueType(VT));
     return DAG.getNode(ISD::TRUNCATE, dl, VT, Assert);
   }
 
   if (Subtarget.hasSSE41())
     if (SDValue Res = LowerEXTRACT_VECTOR_ELT_SSE4(Op, DAG))
       return Res;
 
-  // TODO: handle v16i8.
+  // TODO: We only extract a single element from v16i8, we can probably afford
+  // to be more aggressive here before using the default approach of spilling to
+  // stack.
+  if (VT.getSizeInBits() == 8 && Op->isOnlyUserOf(Vec.getNode())) {
+    // Extract either the lowest i32 or any i16, and extract the sub-byte.
+    int DWordIdx = IdxVal / 4;
+    if (DWordIdx == 0) {
+      SDValue Res = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, MVT::i32,
+                                DAG.getBitcast(MVT::v4i32, Vec),
+                                DAG.getIntPtrConstant(DWordIdx, dl));
+      int ShiftVal = (IdxVal % 4) * 8;
+      if (ShiftVal != 0)
+        Res = DAG.getNode(ISD::SRL, dl, MVT::i32, Res,
+                          DAG.getConstant(ShiftVal, dl, MVT::i32));
+      return DAG.getNode(ISD::TRUNCATE, dl, VT, Res);
+    }
+
+    int WordIdx = IdxVal / 2;
+    SDValue Res = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, MVT::i16,
+                              DAG.getBitcast(MVT::v8i16, Vec),
+                              DAG.getIntPtrConstant(WordIdx, dl));
+    int ShiftVal = (IdxVal % 2) * 8;
+    if (ShiftVal != 0)
+      Res = DAG.getNode(ISD::SRL, dl, MVT::i16, Res,
+                        DAG.getConstant(ShiftVal, dl, MVT::i16));
+    return DAG.getNode(ISD::TRUNCATE, dl, VT, Res);
+  }
 
   if (VT.getSizeInBits() == 32) {
     if (IdxVal == 0)
       return Op;
 
     // SHUFPS the element to the lowest double word, then movss.
     int Mask[4] = { static_cast<int>(IdxVal), -1, -1, -1 };
     Vec = DAG.getVectorShuffle(VecVT, dl, Vec, DAG.getUNDEF(VecVT), Mask);

llvm/trunk/test/CodeGen/X86/extract-store.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=i686-unknown -mattr=+sse2     | FileCheck %s --check-prefix=X32 --check-prefix=SSE-X32 --check-prefix=SSE2-X32
 ; RUN: llc < %s -mtriple=x86_64-unknown -mattr=+sse2   | FileCheck %s --check-prefix=X64 --check-prefix=SSE-X64 --check-prefix=SSE2-X64
 ; RUN: llc < %s -mtriple=i686-unknown -mattr=+sse4.1   | FileCheck %s --check-prefix=X32 --check-prefix=SSE-X32 --check-prefix=SSE41-X32
 ; RUN: llc < %s -mtriple=x86_64-unknown -mattr=+sse4.1 | FileCheck %s --check-prefix=X64 --check-prefix=SSE-X64 --check-prefix=SSE41-X64
 ; RUN: llc < %s -mtriple=i686-unknown -mattr=+avx      | FileCheck %s --check-prefix=X32 --check-prefix=AVX-X32
 ; RUN: llc < %s -mtriple=x86_64-unknown -mattr=+avx    | FileCheck %s --check-prefix=X64 --check-prefix=AVX-X64
 
 define void @extract_i8_0(i8* nocapture %dst, <16 x i8> %foo) nounwind {
 ; SSE2-X32-LABEL: extract_i8_0:
 ; SSE2-X32:       # BB#0:
-; SSE2-X32-NEXT:    pushl %ebp
-; SSE2-X32-NEXT:    movl %esp, %ebp
-; SSE2-X32-NEXT:    andl $-16, %esp
-; SSE2-X32-NEXT:    subl $32, %esp
-; SSE2-X32-NEXT:    movl 8(%ebp), %eax
-; SSE2-X32-NEXT:    movaps %xmm0, (%esp)
-; SSE2-X32-NEXT:    movb (%esp), %cl
+; SSE2-X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; SSE2-X32-NEXT:    movd %xmm0, %ecx
 ; SSE2-X32-NEXT:    movb %cl, (%eax)
-; SSE2-X32-NEXT:    movl %ebp, %esp
-; SSE2-X32-NEXT:    popl %ebp
 ; SSE2-X32-NEXT:    retl
 ;
 ; SSE2-X64-LABEL: extract_i8_0:
 ; SSE2-X64:       # BB#0:
-; SSE2-X64-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-X64-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-X64-NEXT:    movd %xmm0, %eax
 ; SSE2-X64-NEXT:    movb %al, (%rdi)
 ; SSE2-X64-NEXT:    retq
 ;
 ; SSE41-X32-LABEL: extract_i8_0:
 ; SSE41-X32:       # BB#0:
 ; SSE41-X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; SSE41-X32-NEXT:    pextrb $0, %xmm0, (%eax)
 ; SSE41-X32-NEXT:    retl
; AVX-X64-NEXT:    retq
   %vecext = extractelement <16 x i8> %foo, i32 0
   store i8 %vecext, i8* %dst, align 1
   ret void
 }
 
 define void @extract_i8_3(i8* nocapture %dst, <16 x i8> %foo) nounwind {
 ; SSE2-X32-LABEL: extract_i8_3:
 ; SSE2-X32:       # BB#0:
-; SSE2-X32-NEXT:    pushl %ebp
-; SSE2-X32-NEXT:    movl %esp, %ebp
-; SSE2-X32-NEXT:    andl $-16, %esp
-; SSE2-X32-NEXT:    subl $32, %esp
-; SSE2-X32-NEXT:    movl 8(%ebp), %eax
-; SSE2-X32-NEXT:    movaps %xmm0, (%esp)
-; SSE2-X32-NEXT:    movb {{[0-9]+}}(%esp), %cl
+; SSE2-X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; SSE2-X32-NEXT:    movd %xmm0, %ecx
+; SSE2-X32-NEXT:    shrl $24, %ecx
 ; SSE2-X32-NEXT:    movb %cl, (%eax)
-; SSE2-X32-NEXT:    movl %ebp, %esp
-; SSE2-X32-NEXT:    popl %ebp
 ; SSE2-X32-NEXT:    retl
 ;
 ; SSE2-X64-LABEL: extract_i8_3:
 ; SSE2-X64:       # BB#0:
-; SSE2-X64-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-X64-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-X64-NEXT:    movd %xmm0, %eax
+; SSE2-X64-NEXT:    shrl $24, %eax
 ; SSE2-X64-NEXT:    movb %al, (%rdi)
 ; SSE2-X64-NEXT:    retq
 ;
 ; SSE41-X32-LABEL: extract_i8_3:
 ; SSE41-X32:       # BB#0:
 ; SSE41-X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; SSE41-X32-NEXT:    pextrb $3, %xmm0, (%eax)
 ; SSE41-X32-NEXT:    retl
; AVX-X64-NEXT:    retq
   %vecext = extractelement <16 x i8> %foo, i32 3
   store i8 %vecext, i8* %dst, align 1
   ret void
 }
 
 define void @extract_i8_15(i8* nocapture %dst, <16 x i8> %foo) nounwind {
 ; SSE2-X32-LABEL: extract_i8_15:
 ; SSE2-X32:       # BB#0:
-; SSE2-X32-NEXT:    pushl %ebp
-; SSE2-X32-NEXT:    movl %esp, %ebp
-; SSE2-X32-NEXT:    andl $-16, %esp
-; SSE2-X32-NEXT:    subl $32, %esp
-; SSE2-X32-NEXT:    movl 8(%ebp), %eax
-; SSE2-X32-NEXT:    movaps %xmm0, (%esp)
-; SSE2-X32-NEXT:    movb {{[0-9]+}}(%esp), %cl
-; SSE2-X32-NEXT:    movb %cl, (%eax)
-; SSE2-X32-NEXT:    movl %ebp, %esp
-; SSE2-X32-NEXT:    popl %ebp
+; SSE2-X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
+; SSE2-X32-NEXT:    pextrw $7, %xmm0, %ecx
+; SSE2-X32-NEXT:    movb %ch, (%eax)
 ; SSE2-X32-NEXT:    retl
 ;
 ; SSE2-X64-LABEL: extract_i8_15:
 ; SSE2-X64:       # BB#0:
-; SSE2-X64-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-X64-NEXT:    movb -{{[0-9]+}}(%rsp), %al
-; SSE2-X64-NEXT:    movb %al, (%rdi)
+; SSE2-X64-NEXT:    pextrw $7, %xmm0, %eax
+; SSE2-X64-NEXT:    movb %ah, (%rdi) # NOREX
 ; SSE2-X64-NEXT:    retq
 ;
 ; SSE41-X32-LABEL: extract_i8_15:
 ; SSE41-X32:       # BB#0:
 ; SSE41-X32-NEXT:    movl {{[0-9]+}}(%esp), %eax
 ; SSE41-X32-NEXT:    pextrb $15, %xmm0, (%eax)
 ; SSE41-X32-NEXT:    retl
 ;

llvm/trunk/test/CodeGen/X86/extractelement-index.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -mattr=+sse2   | FileCheck %s --check-prefix=SSE --check-prefix=SSE2
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -mattr=+sse4.1 | FileCheck %s --check-prefix=SSE --check-prefix=SSE41
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -mattr=+avx    | FileCheck %s --check-prefix=AVX --check-prefix=AVX1
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -mattr=+avx2   | FileCheck %s --check-prefix=AVX --check-prefix=AVX2
 
 ;
 ; ExtractElement - Constant Index
 ;
 
 define i8 @extractelement_v16i8_1(<16 x i8> %a) nounwind {
 ; SSE2-LABEL: extractelement_v16i8_1:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-NEXT:    movd %xmm0, %eax
+; SSE2-NEXT:    shrl $8, %eax
+; SSE2-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE2-NEXT:    retq
 ;
 ; SSE41-LABEL: extractelement_v16i8_1:
 ; SSE41:       # BB#0:
 ; SSE41-NEXT:    pextrb $1, %xmm0, %eax
 ; SSE41-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE41-NEXT:    retq
 ;
 ; AVX-LABEL: extractelement_v16i8_1:
 ; AVX:       # BB#0:
 ; AVX-NEXT:    vpextrb $1, %xmm0, %eax
 ; AVX-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; AVX-NEXT:    retq
   %b = extractelement <16 x i8> %a, i256 1
   ret i8 %b
 }
 
 define i8 @extractelement_v16i8_11(<16 x i8> %a) nounwind {
 ; SSE2-LABEL: extractelement_v16i8_11:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-NEXT:    pextrw $5, %xmm0, %eax
+; SSE2-NEXT:    shrl $8, %eax
+; SSE2-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE2-NEXT:    retq
 ;
 ; SSE41-LABEL: extractelement_v16i8_11:
 ; SSE41:       # BB#0:
 ; SSE41-NEXT:    pextrb $11, %xmm0, %eax
 ; SSE41-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE41-NEXT:    retq
 ;
 ; AVX-LABEL: extractelement_v16i8_11:
 ; AVX:       # BB#0:
 ; AVX-NEXT:    vpextrb $11, %xmm0, %eax
 ; AVX-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; AVX-NEXT:    retq
   %b = extractelement <16 x i8> %a, i256 11
   ret i8 %b
 }
 
 define i8 @extractelement_v16i8_14(<16 x i8> %a) nounwind {
 ; SSE2-LABEL: extractelement_v16i8_14:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-NEXT:    pextrw $7, %xmm0, %eax
+; SSE2-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE2-NEXT:    retq
 ;
 ; SSE41-LABEL: extractelement_v16i8_14:
 ; SSE41:       # BB#0:
 ; SSE41-NEXT:    pextrb $14, %xmm0, %eax
 ; SSE41-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE41-NEXT:    retq
 ;
 ; AVX-LABEL: extractelement_v16i8_14:
 ; AVX:       # BB#0:
 ; AVX-NEXT:    vpextrb $14, %xmm0, %eax
 ; AVX-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; AVX-NEXT:    retq
   %b = extractelement <16 x i8> %a, i256 14
   ret i8 %b
 }
 
 define i8 @extractelement_v32i8_1(<32 x i8> %a) nounwind {
 ; SSE2-LABEL: extractelement_v32i8_1:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    movaps %xmm0, -{{[0-9]+}}(%rsp)
-; SSE2-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-NEXT:    movd %xmm0, %eax
+; SSE2-NEXT:    shrl $8, %eax
+; SSE2-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE2-NEXT:    retq
 ;
 ; SSE41-LABEL: extractelement_v32i8_1:
 ; SSE41:       # BB#0:
 ; SSE41-NEXT:    pextrb $1, %xmm0, %eax
 ; SSE41-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE41-NEXT:    retq
 ;
 ; AVX-LABEL: extractelement_v32i8_1:
 ; AVX:       # BB#0:
 ; AVX-NEXT:    vpextrb $1, %xmm0, %eax
 ; AVX-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; AVX-NEXT:    vzeroupper
 ; AVX-NEXT:    retq
   %b = extractelement <32 x i8> %a, i256 1
   ret i8 %b
 }
 
 define i8 @extractelement_v32i8_17(<32 x i8> %a) nounwind {
 ; SSE2-LABEL: extractelement_v32i8_17:
 ; SSE2:       # BB#0:
-; SSE2-NEXT:    movaps %xmm1, -{{[0-9]+}}(%rsp)
-; SSE2-NEXT:    movb -{{[0-9]+}}(%rsp), %al
+; SSE2-NEXT:    movd %xmm1, %eax
+; SSE2-NEXT:    shrl $8, %eax
+; SSE2-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE2-NEXT:    retq
 ;
 ; SSE41-LABEL: extractelement_v32i8_17:
 ; SSE41:       # BB#0:
 ; SSE41-NEXT:    pextrb $1, %xmm1, %eax
 ; SSE41-NEXT:    # kill: %AL<def> %AL<kill> %EAX<kill>
 ; SSE41-NEXT:    retq
 ;