[ARM] Teach ComputeKnownBitsTarget to handle extract vectors

[dnsampaio]

Summary

Teaching the ARM backend to set upper zero bits when performing a sign or zero extend allows to remove redundant masking.

[sbaranga]

lib/Target/ARM/ARMISelLowering.cpp

13639 ↗

(On Diff #179545)

IIUC we're computing here the known bits on the input vector operand. However, when computing the result of the VGETLANE operation we're not taking into account the lane operand, which seems wrong to me.

We should have additional tests to cover both the sext and zext cases and cases where the lane operand is not zero.

13656 ↗

(On Diff #179545)

Is this assert unreachable?

[dnsampaio]

lib/Target/ARM/ARMISelLowering.cpp

13639 ↗

(On Diff #179545)

Perhaps I've got it wrong, but from what I've seen is that computeKnownBits tells bits that are known to be zero or one in all vector elements (lanes). So it really does not matter which lane is being extracted. Even the width is the number of bits of a single element, not the entire vector.

I'll add a signed extension test, to show that the behavior is the same.

13656 ↗

(On Diff #179545)

It is. Compiling fore release does not test it. Thanks.

[RKSimon]

lib/Target/ARM/ARMISelLowering.cpp

13639 ↗

(On Diff #179545)

computeKnownBits can takes DemandedElts mask to indicate the vector elements that you care about - see the ISD::EXTRACT_VECTOR_ELT code in SelectionDAG::computeKnownBits for an example of how it works.

[dnsampaio]

Track a single element if the index is a constant.

[dnsampaio]

lib/Target/ARM/ARMISelLowering.cpp

13639 ↗

(On Diff #179545)

Hi, thanks for the suggestion. I just couldn't find a straightforward test to capture such behavior, as most simplifications are done before extractelement gets replaced by VGETLANE. It doesn't seems to have an intrinsic as well.

[RKSimon]

lib/Target/ARM/ARMISelLowering.cpp

13639 ↗

(On Diff #179545)

Something that make use of the implicit sext/zext should be testable? IIRC we do something similar on X86 for PEXTRB/PEXTRW

[sbaranga]

lib/Target/ARM/ARMISelLowering.cpp

13643 ↗

(On Diff #179634)

I think this test should always evaluate to true. Maybe it should be an assert?

[dnsampaio]

lib/Target/ARM/ARMISelLowering.cpp

13643 ↗

(On Diff #179634)

Don't know. Can't Pos be nullptr? That is, the position is not a constant value? (vector[variableX] ?)
if (Pos >= NumSrcElts), perhaps I could just consider it a poison and return as a value zero.

[dnsampaio]

Removed treatment for out of bounds constant, as it is simplified to poison before being transformed to VGETLANE.
Also removed treatment for variable index, as it is not accepted by VGETLANE.
Added extra tests.

lib/Target/ARM/ARMISelLowering.cpp

13643 ↗

(On Diff #179634)

Oh, I see that when creating these operands we only allow constant values. And trying llvm-mc it confirms it so. Sorry, did not get the suggestion. :)

[sbaranga]

LGTM! (with one nit)

lib/Target/ARM/ARMISelLowering.cpp

13642 ↗

(On Diff #180253)

Nit: you can use a cast<> instead of dyn_cast here and remove the assert.

13643 ↗

(On Diff #179634)

Yeah, sorry, I should have given a justification for that.

These nodes are also used to only select instructions that have the lane encoded as an immediate.

llvm/trunk/lib/Target/ARM/ARMISelLowering.cpp

case ARMISD::BFI: {
     // The operand to BFI is already a mask suitable for removing the bits it
     // sets.
     ConstantSDNode *CI = cast<ConstantSDNode>(Op.getOperand(2));
     const APInt &Mask = CI->getAPIntValue();
     Known.Zero &= Mask;
     Known.One &= Mask;
     return;
   }
+  case ARMISD::VGETLANEs:
+  case ARMISD::VGETLANEu: {
+    const SDValue &SrcSV = Op.getOperand(0);
+    EVT VecVT = SrcSV.getValueType();
+    assert(VecVT.isVector() && "VGETLANE expected a vector type");
+    const unsigned NumSrcElts = VecVT.getVectorNumElements();
+    ConstantSDNode *Pos = cast<ConstantSDNode>(Op.getOperand(1).getNode());
+    assert(Pos->getAPIntValue().ult(NumSrcElts) &&
+           "VGETLANE index out of bounds");
+    unsigned Idx = Pos->getZExtValue();
+    APInt DemandedElt = APInt::getOneBitSet(NumSrcElts, Idx);
+    Known = DAG.computeKnownBits(SrcSV, DemandedElt, Depth + 1);
+
+    EVT VT = Op.getValueType();
+    const unsigned DstSz = VT.getScalarSizeInBits();
+    const unsigned SrcSz = VecVT.getVectorElementType().getSizeInBits();
+    assert(SrcSz == Known.getBitWidth());
+    assert(DstSz > SrcSz);
+    if (Op.getOpcode() == ARMISD::VGETLANEs)
+      Known = Known.sext(DstSz);
+    else {
+      Known = Known.zext(DstSz);
+      Known.Zero.setBitsFrom(SrcSz);
+    }
+    assert(DstSz == Known.getBitWidth());
+    break;
+  }
   }
 }
 
 bool
 ARMTargetLowering::targetShrinkDemandedConstant(SDValue Op,
                                                 const APInt &DemandedAPInt,
                                                 TargetLoweringOpt &TLO) const {
   // Delay optimization, so we don't have to deal with illegal types, or block

llvm/trunk/test/CodeGen/ARM/dagcombine-anyexttozeroext.ll

 ; RUN: llc -mtriple armv7 %s -o - | FileCheck %s
 
-; CHECK-LABEL: f:
 define float @f(<4 x i16>* nocapture %in) {
-  ; CHECK: vld1
-  ; CHECK: vmovl.u16
-  ; CHECK-NOT: vand
+; CHECK-LABEL: f:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vld1.16 {d16}, [r0:64]
+; CHECK-NEXT:    vmovl.u16 q8, d16
+; CHECK-NEXT:    vcvt.f32.u32 q0, q8
+; CHECK-NEXT:    vadd.f32 s4, s0, s1
+; CHECK-NEXT:    vadd.f32 s0, s4, s2
+; CHECK-NEXT:    vmov r0, s0
+; CHECK-NEXT:    bx lr
   %1 = load <4 x i16>, <4 x i16>* %in
-  ; CHECK: vcvt.f32.u32
   %2 = uitofp <4 x i16> %1 to <4 x float>
   %3 = extractelement <4 x float> %2, i32 0
   %4 = extractelement <4 x float> %2, i32 1
   %5 = extractelement <4 x float> %2, i32 2
 
-  ; CHECK: vadd.f32
   %6 = fadd float %3, %4
   %7 = fadd float %6, %5
 
   ret float %7
 }
 
-; CHECK-LABEL: g:
 define float @g(<4 x i16>* nocapture %in) {
-  ; CHECK: vldr
+; CHECK-LABEL: g:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vldr d16, [r0]
+; CHECK-NEXT:    vmov.u16 r0, d16[0]
+; CHECK-NEXT:    vmov s0, r0
+; CHECK-NEXT:    vcvt.f32.u32 s0, s0
+; CHECK-NEXT:    vmov r0, s0
+; CHECK-NEXT:    bx lr
   %1 = load <4 x i16>, <4 x i16>* %in
-
-  ; For now we're generating a vmov.16 and a uxth instruction.
-  ; The uxth is redundant, and we should be able to extend without
-  ; having to generate cross-domain copies. Once we can do this
-  ; we should modify the checks below.
-
-  ; CHECK: uxth
   %2 = extractelement <4 x i16> %1, i32 0
-  ; CHECK: vcvt.f32.u32
   %3 = uitofp i16 %2 to float
   ret float %3
 }
 
 ; Make sure we generate zext from <4 x i8> to <4 x 32>.
-
-; CHECK-LABEL: h:
-; CHECK: vld1.32
-; CHECK: vmovl.u8 q8, d16
-; CHECK: vmovl.u16 q8, d16
-; CHECK: vmov r0, r1, d16
-; CHECK: vmov r2, r3, d17
 define <4 x i32> @h(<4 x i8> *%in) {
+; CHECK-LABEL: h:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vld1.32 {d16[0]}, [r0:32]
+; CHECK-NEXT:    vmovl.u8 q8, d16
+; CHECK-NEXT:    vmovl.u16 q8, d16
+; CHECK-NEXT:    vmov r0, r1, d16
+; CHECK-NEXT:    vmov r2, r3, d17
+; CHECK-NEXT:    bx lr
   %1 = load <4 x i8>, <4 x i8>* %in, align 4
   %2 = extractelement <4 x i8> %1, i32 0
   %3 = zext i8 %2 to i32
   %4 = insertelement <4 x i32> undef, i32 %3, i32 0
   %5 = extractelement <4 x i8> %1, i32 1
   %6 = zext i8 %5 to i32
   %7 = insertelement <4 x i32> %4, i32 %6, i32 1
   %8 = extractelement <4 x i8> %1, i32 2
   %9 = zext i8 %8 to i32
   %10 = insertelement <4 x i32> %7, i32 %9, i32 2
   %11 = extractelement <4 x i8> %1, i32 3
   %12 = zext i8 %11 to i32
   %13 = insertelement <4 x i32> %10, i32 %12, i32 3
   ret <4 x i32> %13
 }
+
+define float @i(<4 x i16>* nocapture %in) {
+  ; FIXME: The vmov.u + sxt can convert to a vmov.s
+; CHECK-LABEL: i:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vldr d16, [r0]
+; CHECK-NEXT:    vmov.u16 r0, d16[0]
+; CHECK-NEXT:    sxth r0, r0
+; CHECK-NEXT:    vmov s0, r0
+; CHECK-NEXT:    vcvt.f32.s32 s0, s0
+; CHECK-NEXT:    vmov r0, s0
+; CHECK-NEXT:    bx lr
+  %1 = load <4 x i16>, <4 x i16>* %in
+  %2 = extractelement <4 x i16> %1, i32 0
+  %3 = sitofp i16 %2 to float
+  ret float %3
+}
+
+define float @j(<8 x i8>* nocapture %in) {
+; CHECK-LABEL: j:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vldr d16, [r0]
+; CHECK-NEXT:    vmov.u8 r0, d16[7]
+; CHECK-NEXT:    vmov s0, r0
+; CHECK-NEXT:    vcvt.f32.u32 s0, s0
+; CHECK-NEXT:    vmov r0, s0
+; CHECK-NEXT:    bx lr
+  %1 = load <8 x i8>, <8 x i8>* %in
+  %2 = extractelement <8 x i8> %1, i32 7
+  %3 = uitofp i8 %2 to float
+  ret float %3
+}
+
+define float @k(<8 x i8>* nocapture %in) {
+; FIXME: The vmov.u + sxt can convert to a vmov.s
+; CHECK-LABEL: k:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vldr d16, [r0]
+; CHECK-NEXT:    vmov.u8 r0, d16[7]
+; CHECK-NEXT:    sxtb r0, r0
+; CHECK-NEXT:    vmov s0, r0
+; CHECK-NEXT:    vcvt.f32.s32 s0, s0
+; CHECK-NEXT:    vmov r0, s0
+; CHECK-NEXT:    bx lr
+  %1 = load <8 x i8>, <8 x i8>* %in
+  %2 = extractelement <8 x i8> %1, i32 7
+  %3 = sitofp i8 %2 to float
+  ret float %3
+}
+
+define float @KnownUpperZero(<4 x i16> %v) {
+; FIXME: uxtb are not required
+; CHECK-LABEL: KnownUpperZero:
+; CHECK:       @ %bb.0:
+; CHECK-NEXT:    vmov.i16 d16, #0x3
+; CHECK-NEXT:    vmov d17, r0, r1
+; CHECK-NEXT:    vand d16, d17, d16
+; CHECK-NEXT:    vmov.u16 r0, d16[0]
+; CHECK-NEXT:    vmov.u16 r1, d16[3]
+; CHECK-NEXT:    uxtb r0, r0
+; CHECK-NEXT:    vmov s0, r0
+; CHECK-NEXT:    uxtb r0, r1
+; CHECK-NEXT:    vmov s2, r0
+; CHECK-NEXT:    vcvt.f32.s32 s0, s0
+; CHECK-NEXT:    vcvt.f32.s32 s2, s2
+; CHECK-NEXT:    vadd.f32 s0, s2, s0
+; CHECK-NEXT:    vmov r0, s0
+; CHECK-NEXT:    bx lr
+  %1 = and <4 x i16> %v, <i16 3,i16 3,i16 3,i16 3>
+  %2 = extractelement <4 x i16> %1, i32 3
+  %3 = extractelement <4 x i16> %1, i32 0
+  %sinf1 = sitofp i16 %2 to float
+  %sinf2 = sitofp i16 %3 to float
+  %sum =   fadd float %sinf1, %sinf2
+  ret float %sum
+}