i32/i64 support after code restructure for vectors.

shami · shami · commit f5dbf6615a27 · 2024-04-10T19:36:25.000+05:30
diff --git a/llvm/lib/Target/X86/X86ISelLowering.cpp b/llvm/lib/Target/X86/X86ISelLowering.cpp
@@ -1496,6 +1496,11 @@ X86TargetLowering::X86TargetLowering(const X86TargetMachine &TM,
     setOperationAction(ISD::TRUNCATE,          MVT::v32i32, Custom);
     setOperationAction(ISD::TRUNCATE,          MVT::v32i64, Custom);
 
+    if (Subtarget.hasGFNI()) {
+      setOperationAction(ISD::BITREVERSE,      MVT::i32, Custom);
+      setOperationAction(ISD::BITREVERSE,      MVT::i64, Custom);
+    }
+
     for (auto VT : { MVT::v32i8, MVT::v16i16, MVT::v8i32, MVT::v4i64 }) {
       setOperationAction(ISD::SETCC,           VT, Custom);
       setOperationAction(ISD::CTPOP,           VT, Custom);
@@ -31312,53 +31317,18 @@ static SDValue LowerBITREVERSE_XOP(SDValue Op, SelectionDAG &DAG) {
   return DAG.getBitcast(VT, Res);
 }
 
-static void createBSWAPShuffleMask(EVT VT, SmallVector<int, 16>& ShuffleMask) {
-  int ScalarSizeInBytes = VT.getScalarSizeInBits() / 8;
-  for (int I = 0, E = VT.getVectorNumElements(); I != E; ++I)
-    for (int J = ScalarSizeInBytes - 1; J >= 0; --J)
-      ShuffleMask.push_back((I * ScalarSizeInBytes) + J);
-}
-
 static SDValue LowerBITREVERSE(SDValue Op, const X86Subtarget &Subtarget,
                                SelectionDAG &DAG) {
   MVT VT = Op.getSimpleValueType();
-  SDValue In = Op.getOperand(0);
-  SDLoc DL(Op);
-
-  auto HasGFNI = Subtarget.hasGFNI();
-  auto ScalarType = VT.getScalarType();
-
-  if (HasGFNI && ((ScalarType == MVT::i32) || (ScalarType == MVT::i64))) {
-    if (VT.isVector()) {
-      SmallVector<int, 16> BSWAPMask = createBSWAPShuffleMask(VT);
-      MVT ByteVT =
-          EVT::getVectorVT(*DAG.getContext(), MVT::i8, BSWAPMask.size());
-      SDValue VecShuffle = DAG.getVectorShuffle(
-          ByteVT, DL, DAG.getNode(ISD::BITCAST, DL, ByteVT, In),
-          DAG.getUNDEF(ByteVT), BSWAPMask);
-      SDValue BitReverse = DAG.getNode(ISD::BITREVERSE, DL, ByteVT, VecShuffle);
-      return DAG.getBitcast(VT, BitReverse);
-    } else {
-      auto CastTo = ScalarType == MVT::i32 ? MVT::v4i32 : MVT::v2i64;
-      SDValue ScalarToVector =
-          DAG.getNode(ISD::SCALAR_TO_VECTOR, DL, CastTo, In);
-      SDValue BitReverse =
-          DAG.getNode(ISD::BITREVERSE, DL, MVT::v16i8,
-                      DAG.getBitcast(MVT::v16i8, ScalarToVector));
-      SDValue ExtractElementZero = DAG.getNode(
-          ISD::EXTRACT_VECTOR_ELT, DL, ScalarType,
-          DAG.getBitcast(CastTo, BitReverse), DAG.getIntPtrConstant(0, DL));
-      return DAG.getNode(ISD::BSWAP, DL, ScalarType, ExtractElementZero);
-    }
-  }
-
-  // Split v64i8 without BWI so that we can still use the PSHUFB lowering.
 
   if (Subtarget.hasXOP() && !VT.is512BitVector())
     return LowerBITREVERSE_XOP(Op, DAG);
 
   assert(Subtarget.hasSSSE3() && "SSSE3 required for BITREVERSE");
 
+  SDValue In = Op.getOperand(0);
+  SDLoc DL(Op);
+
   // Split 512-bit ops without BWI so that we can still use the PSHUFB lowering.
   if (VT.is512BitVector() && !Subtarget.hasBWI())
     return splitVectorIntUnary(Op, DAG, DL);
@@ -31367,6 +31337,21 @@ static SDValue LowerBITREVERSE(SDValue Op, const X86Subtarget &Subtarget,
   if (VT.is256BitVector() && !Subtarget.hasInt256())
     return splitVectorIntUnary(Op, DAG, DL);
 
+  // Lower i32/i64 to GFNI as i32/i64 -> Convert to vector (V = v16i32/v8i64) -> vXi8 BITREVERSE -> V[0] -> BSWAP
+  if (Subtarget.hasGFNI() && !VT.isVector()) {
+
+    assert ((VT.getScalarType() == MVT::i32) || (VT.getScalarType() == MVT::i64));
+
+    auto ScalarType = VT.getScalarType();
+    auto CastTo = ScalarType == MVT::i32 ? MVT::v4i32 : MVT::v2i64;
+    SDValue Res = DAG.getNode(ISD::SCALAR_TO_VECTOR, DL, CastTo, In);
+    Res = DAG.getNode(ISD::BITREVERSE, DL, MVT::v16i8, DAG.getBitcast(MVT::v16i8, Res));
+    Res = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ScalarType, DAG.getBitcast(CastTo, Res), DAG.getIntPtrConstant(0, DL));
+    return DAG.getNode(ISD::BSWAP, DL, ScalarType, Res);
+  }
+
+  assert (VT.isVector() && VT.getSizeInBits() >= 128);
+
   // Lower vXi16/vXi32/vXi64 as BSWAP + vXi8 BITREVERSE.
   if (VT.getScalarType() != MVT::i8) {
     MVT ByteVT = MVT::getVectorVT(MVT::i8, VT.getSizeInBits() / 8);
@@ -31381,7 +31366,7 @@ static SDValue LowerBITREVERSE(SDValue Op, const X86Subtarget &Subtarget,
   unsigned NumElts = VT.getVectorNumElements();
 
   // If we have GFNI, we can use GF2P8AFFINEQB to reverse the bits.
-  if (HasGFNI) {
+  if (Subtarget.hasGFNI()) {
     MVT MatrixVT = MVT::getVectorVT(MVT::i64, NumElts / 8);
     SDValue Matrix = DAG.getConstant(0x8040201008040201ULL, DL, MatrixVT);
     Matrix = DAG.getBitcast(VT, Matrix);
diff --git a/llvm/test/CodeGen/X86/vector-bitreverse.ll b/llvm/test/CodeGen/X86/vector-bitreverse.ll
@@ -274,57 +274,13 @@ define i32 @test_bitreverse_i32(i32 %a) nounwind {
 ; GFNISSE-NEXT:    leal (%rax,%rcx,2), %eax
 ; GFNISSE-NEXT:    retq
 ;
-; GFNIAVX1-LABEL: test_bitreverse_i32:
-; GFNIAVX1:       # %bb.0:
-; GFNIAVX1-NEXT:    # kill: def $edi killed $edi def $rdi
-; GFNIAVX1-NEXT:    bswapl %edi
-; GFNIAVX1-NEXT:    movl %edi, %eax
-; GFNIAVX1-NEXT:    andl $252645135, %eax # imm = 0xF0F0F0F
-; GFNIAVX1-NEXT:    shll $4, %eax
-; GFNIAVX1-NEXT:    shrl $4, %edi
-; GFNIAVX1-NEXT:    andl $252645135, %edi # imm = 0xF0F0F0F
-; GFNIAVX1-NEXT:    orl %eax, %edi
-; GFNIAVX1-NEXT:    movl %edi, %eax
-; GFNIAVX1-NEXT:    andl $858993459, %eax # imm = 0x33333333
-; GFNIAVX1-NEXT:    shrl $2, %edi
-; GFNIAVX1-NEXT:    andl $858993459, %edi # imm = 0x33333333
-; GFNIAVX1-NEXT:    leal (%rdi,%rax,4), %eax
-; GFNIAVX1-NEXT:    movl %eax, %ecx
-; GFNIAVX1-NEXT:    andl $1431655765, %ecx # imm = 0x55555555
-; GFNIAVX1-NEXT:    shrl %eax
-; GFNIAVX1-NEXT:    andl $1431655765, %eax # imm = 0x55555555
-; GFNIAVX1-NEXT:    leal (%rax,%rcx,2), %eax
-; GFNIAVX1-NEXT:    retq
-;
-; GFNIAVX2-LABEL: test_bitreverse_i32:
-; GFNIAVX2:       # %bb.0:
-; GFNIAVX2-NEXT:    # kill: def $edi killed $edi def $rdi
-; GFNIAVX2-NEXT:    bswapl %edi
-; GFNIAVX2-NEXT:    movl %edi, %eax
-; GFNIAVX2-NEXT:    andl $252645135, %eax # imm = 0xF0F0F0F
-; GFNIAVX2-NEXT:    shll $4, %eax
-; GFNIAVX2-NEXT:    shrl $4, %edi
-; GFNIAVX2-NEXT:    andl $252645135, %edi # imm = 0xF0F0F0F
-; GFNIAVX2-NEXT:    orl %eax, %edi
-; GFNIAVX2-NEXT:    movl %edi, %eax
-; GFNIAVX2-NEXT:    andl $858993459, %eax # imm = 0x33333333
-; GFNIAVX2-NEXT:    shrl $2, %edi
-; GFNIAVX2-NEXT:    andl $858993459, %edi # imm = 0x33333333
-; GFNIAVX2-NEXT:    leal (%rdi,%rax,4), %eax
-; GFNIAVX2-NEXT:    movl %eax, %ecx
-; GFNIAVX2-NEXT:    andl $1431655765, %ecx # imm = 0x55555555
-; GFNIAVX2-NEXT:    shrl %eax
-; GFNIAVX2-NEXT:    andl $1431655765, %eax # imm = 0x55555555
-; GFNIAVX2-NEXT:    leal (%rax,%rcx,2), %eax
-; GFNIAVX2-NEXT:    retq
-;
-; GFNIAVX512-LABEL: test_bitreverse_i32:
-; GFNIAVX512:       # %bb.0:
-; GFNIAVX512-NEXT:    vmovd %edi, %xmm0
-; GFNIAVX512-NEXT:    vgf2p8affineqb $0, {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
-; GFNIAVX512-NEXT:    vmovd %xmm0, %eax
-; GFNIAVX512-NEXT:    bswapl %eax
-; GFNIAVX512-NEXT:    retq
+; GFNIAVX-LABEL: test_bitreverse_i32:
+; GFNIAVX:       # %bb.0:
+; GFNIAVX-NEXT:    vmovd %edi, %xmm0
+; GFNIAVX-NEXT:    vgf2p8affineqb $0, {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; GFNIAVX-NEXT:    vmovd %xmm0, %eax
+; GFNIAVX-NEXT:    bswapl %eax
+; GFNIAVX-NEXT:    retq
   %b = call i32 @llvm.bitreverse.i32(i32 %a)
   ret i32 %b
 }
@@ -409,61 +365,13 @@ define i64 @test_bitreverse_i64(i64 %a) nounwind {
 ; GFNISSE-NEXT:    leaq (%rax,%rdx,2), %rax
 ; GFNISSE-NEXT:    retq
 ;
-; GFNIAVX1-LABEL: test_bitreverse_i64:
-; GFNIAVX1:       # %bb.0:
-; GFNIAVX1-NEXT:    bswapq %rdi
-; GFNIAVX1-NEXT:    movq %rdi, %rax
-; GFNIAVX1-NEXT:    shrq $4, %rax
-; GFNIAVX1-NEXT:    movabsq $1085102592571150095, %rcx # imm = 0xF0F0F0F0F0F0F0F
-; GFNIAVX1-NEXT:    andq %rcx, %rax
-; GFNIAVX1-NEXT:    andq %rcx, %rdi
-; GFNIAVX1-NEXT:    shlq $4, %rdi
-; GFNIAVX1-NEXT:    orq %rax, %rdi
-; GFNIAVX1-NEXT:    movabsq $3689348814741910323, %rax # imm = 0x3333333333333333
-; GFNIAVX1-NEXT:    movq %rdi, %rcx
-; GFNIAVX1-NEXT:    andq %rax, %rcx
-; GFNIAVX1-NEXT:    shrq $2, %rdi
-; GFNIAVX1-NEXT:    andq %rax, %rdi
-; GFNIAVX1-NEXT:    leaq (%rdi,%rcx,4), %rax
-; GFNIAVX1-NEXT:    movabsq $6148914691236517205, %rcx # imm = 0x5555555555555555
-; GFNIAVX1-NEXT:    movq %rax, %rdx
-; GFNIAVX1-NEXT:    andq %rcx, %rdx
-; GFNIAVX1-NEXT:    shrq %rax
-; GFNIAVX1-NEXT:    andq %rcx, %rax
-; GFNIAVX1-NEXT:    leaq (%rax,%rdx,2), %rax
-; GFNIAVX1-NEXT:    retq
-;
-; GFNIAVX2-LABEL: test_bitreverse_i64:
-; GFNIAVX2:       # %bb.0:
-; GFNIAVX2-NEXT:    bswapq %rdi
-; GFNIAVX2-NEXT:    movq %rdi, %rax
-; GFNIAVX2-NEXT:    shrq $4, %rax
-; GFNIAVX2-NEXT:    movabsq $1085102592571150095, %rcx # imm = 0xF0F0F0F0F0F0F0F
-; GFNIAVX2-NEXT:    andq %rcx, %rax
-; GFNIAVX2-NEXT:    andq %rcx, %rdi
-; GFNIAVX2-NEXT:    shlq $4, %rdi
-; GFNIAVX2-NEXT:    orq %rax, %rdi
-; GFNIAVX2-NEXT:    movabsq $3689348814741910323, %rax # imm = 0x3333333333333333
-; GFNIAVX2-NEXT:    movq %rdi, %rcx
-; GFNIAVX2-NEXT:    andq %rax, %rcx
-; GFNIAVX2-NEXT:    shrq $2, %rdi
-; GFNIAVX2-NEXT:    andq %rax, %rdi
-; GFNIAVX2-NEXT:    leaq (%rdi,%rcx,4), %rax
-; GFNIAVX2-NEXT:    movabsq $6148914691236517205, %rcx # imm = 0x5555555555555555
-; GFNIAVX2-NEXT:    movq %rax, %rdx
-; GFNIAVX2-NEXT:    andq %rcx, %rdx
-; GFNIAVX2-NEXT:    shrq %rax
-; GFNIAVX2-NEXT:    andq %rcx, %rax
-; GFNIAVX2-NEXT:    leaq (%rax,%rdx,2), %rax
-; GFNIAVX2-NEXT:    retq
-;
-; GFNIAVX512-LABEL: test_bitreverse_i64:
-; GFNIAVX512:       # %bb.0:
-; GFNIAVX512-NEXT:    vmovq %rdi, %xmm0
-; GFNIAVX512-NEXT:    vgf2p8affineqb $0, {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
-; GFNIAVX512-NEXT:    vmovq %xmm0, %rax
-; GFNIAVX512-NEXT:    bswapq %rax
-; GFNIAVX512-NEXT:    retq
+; GFNIAVX-LABEL: test_bitreverse_i64:
+; GFNIAVX:       # %bb.0:
+; GFNIAVX-NEXT:    vmovq %rdi, %xmm0
+; GFNIAVX-NEXT:    vgf2p8affineqb $0, {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0, %xmm0
+; GFNIAVX-NEXT:    vmovq %xmm0, %rax
+; GFNIAVX-NEXT:    bswapq %rax
+; GFNIAVX-NEXT:    retq
   %b = call i64 @llvm.bitreverse.i64(i64 %a)
   ret i64 %b
 }
