[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

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

Is this assert unreachable?

[dnsampaio]

lib/Target/ARM/ARMISelLowering.cpp

13639

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

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

[RKSimon]

lib/Target/ARM/ARMISelLowering.cpp

13639

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

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

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

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

[dnsampaio]

lib/Target/ARM/ARMISelLowering.cpp

13643

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

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

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

13643

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.

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);
+    Known = DAG.computeKnownBits(SrcSV, Depth + 1);
+    EVT VT = Op.getValueType();
+    EVT EltVT = SrcSV.getValueType();
+    const unsigned DstSz = VT.getScalarSizeInBits();
+    const unsigned SrcSz = EltVT.getVectorElementType().getSizeInBits();
+
+    assert(SrcSz == Known.getBitWidth());
+    if (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

test/CodeGen/ARM/dagcombine-anyexttozeroext.ll

define float @f(<4 x i16>* nocapture %in) {
   %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:  vldr  d16, [r0]
+  ; FIXME: We should be able to extend without
+  ; having to generate cross-domain copies.
+; 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
define <4 x i32> @h(<4 x i8> *%in) {
   %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) {
+; CHECK-LABEL: i:
+; CHECK:  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
+}
+
+; CHECK-LABEL: j:
+define float @j(<8 x i8>* nocapture %in) {
+; CHECK:  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
+}
+
+; CHECK-LABEL: k:
+define float @k(<8 x i8>* nocapture %in) {
+; CHECK:  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
+}