[IR][Legalizations] Widen illegal deinterleave and interleave vectors

[CarolineConcatto]

Summary

To make legalization easier, the operands and outputs have the same size for
these ISD Nodes. When legalizing the results in WidenVecRes_VECTOR_DEINTERLEAVE
the operands are legalized to the same size as the outputs.
The ISD Node has two output/results, therefore the legalizing functions update
both results/outputs.

[CarolineConcatto]

• Remove old function header

[CarolineConcatto]

Gentle ping!

[MattDevereau]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

5023–5051

I'm not sure if i'm missing something but if i remove this whole block and run check-all tests I don't see anything failing. Is this missing a test case?

5725

DL?

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

11592

This can be defined next to where it's used. Is PackedVT or LegalTY a better name for this? If InVT is nxv3i32 then WideVT = TLI.getTypeToTransformTo(Ctx, InVT); will return nxv4i32 which I found quite surprising

11625

The way i'm reading this if statement is "before we had to widen InVT to WideTV". If we do something like:

bool InVTMustWiden = (TLI.getTypeAction(Ctx, InVT) == TargetLowering::TypeWidenVector);
if (InVTMustWiden) {
  WideVT = TLI.getTypeToTransformTo(Ctx, InVT);
  ...
...
...
if (InVTMustWiden) {
  OutVT = TLI.getValueType(DAG.getDataLayout(), I.getType());
...

Do you think that makes it easier to understand?

[CarolineConcatto]

-Address review comments

[craig.topper]

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

5020

Unrelated change?

5733

GetWidenedVector does a map lookup, but there's no way WidenVec could be in the map. Won't this fail?

llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h

private:
   void WidenVectorResult(SDNode *N, unsigned ResNo);
   SDValue WidenVecRes_MERGE_VALUES(SDNode* N, unsigned ResNo);
   SDValue WidenVecRes_BITCAST(SDNode* N);
   SDValue WidenVecRes_BUILD_VECTOR(SDNode* N);
   SDValue WidenVecRes_CONCAT_VECTORS(SDNode* N);
   SDValue WidenVecRes_EXTEND_VECTOR_INREG(SDNode* N);
   SDValue WidenVecRes_EXTRACT_SUBVECTOR(SDNode* N);
   SDValue WidenVecRes_INSERT_SUBVECTOR(SDNode *N);
+  SDValue WidenVecRes_VECTOR_DEINTERLEAVE(SDNode *N);
   SDValue WidenVecRes_INSERT_VECTOR_ELT(SDNode* N);
   SDValue WidenVecRes_LOAD(SDNode* N);
   SDValue WidenVecRes_VP_LOAD(VPLoadSDNode *N);
   SDValue WidenVecRes_VP_STRIDED_LOAD(VPStridedLoadSDNode *N);
   SDValue WidenVecRes_MLOAD(MaskedLoadSDNode* N);
   SDValue WidenVecRes_MGATHER(MaskedGatherSDNode* N);
   SDValue WidenVecRes_VP_GATHER(VPGatherSDNode* N);
   SDValue WidenVecRes_ScalarOp(SDNode* N);

llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp

#endif
   case ISD::MERGE_VALUES:      Res = WidenVecRes_MERGE_VALUES(N, ResNo); break;
   case ISD::BITCAST:           Res = WidenVecRes_BITCAST(N); break;
   case ISD::BUILD_VECTOR:      Res = WidenVecRes_BUILD_VECTOR(N); break;
   case ISD::CONCAT_VECTORS:    Res = WidenVecRes_CONCAT_VECTORS(N); break;
   case ISD::INSERT_SUBVECTOR:
     Res = WidenVecRes_INSERT_SUBVECTOR(N);
     break;
   case ISD::EXTRACT_SUBVECTOR: Res = WidenVecRes_EXTRACT_SUBVECTOR(N); break;
+  case ISD::VECTOR_DEINTERLEAVE:
+    Res = WidenVecRes_VECTOR_DEINTERLEAVE(N);
+    return;
   case ISD::INSERT_VECTOR_ELT: Res = WidenVecRes_INSERT_VECTOR_ELT(N); break;
   case ISD::LOAD:              Res = WidenVecRes_LOAD(N); break;
   case ISD::STEP_VECTOR:
   case ISD::SPLAT_VECTOR:
   case ISD::SCALAR_TO_VECTOR:
     Res = WidenVecRes_ScalarOp(N);
     break;
   case ISD::SIGN_EXTEND_INREG: Res = WidenVecRes_InregOp(N); break;
if (WidenVT == TLI.getTypeToTransformTo(*DAG.getContext(), InVT)) {
         }
         return DAG.getVectorShuffle(WidenVT, dl,
                                     GetWidenedVector(N->getOperand(0)),
                                     GetWidenedVector(N->getOperand(1)),
                                     MaskOps);
       }
     }
   }
-
   assert(!WidenVT.isScalableVector() &&
          "Cannot use build vectors to widen CONCAT_VECTOR result");
   unsigned WidenNumElts = WidenVT.getVectorNumElements();
   unsigned NumInElts = InVT.getVectorNumElements();
 
   // Fall back to use extracts and build vector.
   EVT EltVT = WidenVT.getVectorElementType();
   SmallVector<SDValue, 16> Ops(WidenNumElts);
   unsigned Idx = 0;
   for (unsigned i=0; i < NumOperands; ++i) {
     SDValue InOp = N->getOperand(i);
     if (InputWidened)
       InOp = GetWidenedVector(InOp);
     for (unsigned j = 0; j < NumInElts; ++j)
       Ops[Idx++] = DAG.getNode(ISD::EXTRACT_VECTOR_ELT, dl, EltVT, InOp,
                                DAG.getVectorIdxConstant(j, dl));
   }
   SDValue UndefVal = DAG.getUNDEF(EltVT);
   for (; Idx < WidenNumElts; ++Idx)
     Ops[Idx] = UndefVal;
   return DAG.getBuildVector(WidenVT, dl, Ops);
 }
 
 SDValue DAGTypeLegalizer::WidenVecRes_INSERT_SUBVECTOR(SDNode *N) {
   EVT VT = N->getValueType(0);
   EVT WidenVT = TLI.getTypeToTransformTo(*DAG.getContext(), VT);
   SDValue InOp1 = GetWidenedVector(N->getOperand(0));
   SDValue InOp2 = N->getOperand(1);
   SDValue Idx = N->getOperand(2);
   SDLoc dl(N);
   return DAG.getNode(ISD::INSERT_SUBVECTOR, dl, WidenVT, InOp1, InOp2, Idx);
 }
 
 SDValue DAGTypeLegalizer::WidenVecRes_EXTRACT_SUBVECTOR(SDNode *N) {
   EVT VT = N->getValueType(0);
   EVT EltVT = VT.getVectorElementType();
   EVT WidenVT = TLI.getTypeToTransformTo(*DAG.getContext(), VT);
SDValue DAGTypeLegalizer::WidenVecRes_VECTOR_REVERSE(SDNode *N) {
   }
   for (unsigned i = VTNumElts; i != WidenNumElts; ++i)
     Mask.push_back(-1);
 
   return DAG.getVectorShuffle(WidenVT, dl, ReverseVal, DAG.getUNDEF(WidenVT),
                               Mask);
 }
 
+SDValue DAGTypeLegalizer::WidenVecRes_VECTOR_DEINTERLEAVE(SDNode *N) {
+  SDLoc dl(N);
+  EVT InVT = N->getValueType(0);
+  EVT InConcVT =
+      EVT::getVectorVT(*DAG.getContext(), InVT.getVectorElementType(),
+                       InVT.getVectorElementCount() * 2);
+  SDValue WidenVec = DAG.getNode(ISD::CONCAT_VECTORS, dl, InConcVT,
+                                 N->getOperand(0), N->getOperand(1));
+  if (getTypeAction(WidenVec.getValueType()) == TargetLowering::TypeWidenVector)
+    WidenVec = GetWidenedVector(WidenVec);
+
+  EVT WidenVT = WidenVec.getValueType();
+  EVT SplitVT = EVT::getVectorVT(
+      *DAG.getContext(), WidenVT.getVectorElementType(),
+      WidenVT.getVectorMinNumElements() / 2, InVT.isScalableVector());
+  SDValue Lo = DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, SplitVT, WidenVec,
+                           DAG.getConstant(0, dl, MVT::i64));
+  SDValue Hi = DAG.getNode(
+      ISD::EXTRACT_SUBVECTOR, dl, SplitVT, WidenVec,
+      DAG.getConstant(SplitVT.getVectorMinNumElements(), dl, MVT::i64));
+
+  SDValue Res = DAG.getNode(ISD::VECTOR_DEINTERLEAVE, dl,
+                            DAG.getVTList(SplitVT, SplitVT), Lo, Hi);
+  SetWidenedVector(SDValue(N, 0), Res.getValue(0));
+  SetWidenedVector(SDValue(N, 1), Res.getValue(1));
+  return SDValue();
+}
+
 SDValue DAGTypeLegalizer::WidenVecRes_SETCC(SDNode *N) {
   assert(N->getValueType(0).isVector() &&
          N->getOperand(0).getValueType().isVector() &&
          "Operands must be vectors");
   EVT WidenVT = TLI.getTypeToTransformTo(*DAG.getContext(), N->getValueType(0));
   ElementCount WidenEC = WidenVT.getVectorElementCount();
 
   SDValue InOp1 = N->getOperand(0);

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

SDValue Res = DAG.getNode(ISD::VECTOR_DEINTERLEAVE, DL,
                             DAG.getVTList(OutVT, OutVT), Lo, Hi);
   setValue(&I, Res);
   return;
 }
 
 void SelectionDAGBuilder::visitVectorInterleave(const CallInst &I) {
   auto DL = getCurSDLoc();
   EVT InVT = getValue(I.getOperand(0)).getValueType();
   SDValue InVec0 = getValue(I.getOperand(0));
   SDValue InVec1 = getValue(I.getOperand(1));
   const TargetLowering &TLI = DAG.getTargetLoweringInfo();
   EVT OutVT = TLI.getValueType(DAG.getDataLayout(), I.getType());
 
   // Use VECTOR_SHUFFLE for fixed-length vectors to benefit from existing
   // legalisation and combines.
   if (OutVT.isFixedLengthVector()) {
     unsigned NumElts = InVT.getVectorMinNumElements();
     SDValue V = DAG.getNode(ISD::CONCAT_VECTORS, DL, OutVT, InVec0, InVec1);
     setValue(&I, DAG.getVectorShuffle(OutVT, DL, V, DAG.getUNDEF(OutVT),
                                       createInterleaveMask(NumElts, 2)));
     return;
   }
 
+  // The resulting vector is wrongly concatenated/ordered if it doesn't widen
+  // the vector first if needed.
+  LLVMContext &Ctx = *DAG.getContext();
+  bool IsInVTWiden =
+      (TLI.getTypeAction(Ctx, InVT) == TargetLowering::TypeWidenVector);
+  if (IsInVTWiden) {
+    InVT = TLI.getTypeToTransformTo(Ctx, InVT);
+    InVec0 = widenVectorToPartType(DAG, InVec0, DL, InVT);
+    InVec1 = widenVectorToPartType(DAG, InVec1, DL, InVT);
+    OutVT = InVec0.getValueType().getDoubleNumVectorElementsVT(Ctx);
+  }
+
   SDValue Res = DAG.getNode(ISD::VECTOR_INTERLEAVE, DL,
                             DAG.getVTList(InVT, InVT), InVec0, InVec1);
   Res = DAG.getNode(ISD::CONCAT_VECTORS, DL, OutVT, Res.getValue(0),
                     Res.getValue(1));
+
+  // Return to the original vector size, before widening.
+  if (IsInVTWiden) {
+    OutVT = TLI.getValueType(DAG.getDataLayout(), I.getType());
+    Res = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, OutVT, Res,
+                      DAG.getVectorIdxConstant(0, DL));
+  }
+
   setValue(&I, Res);
   return;
 }
 
 void SelectionDAGBuilder::visitFreeze(const FreezeInst &I) {
   SmallVector<EVT, 4> ValueVTs;
   ComputeValueVTs(DAG.getTargetLoweringInfo(), DAG.getDataLayout(), I.getType(),
                   ValueVTs);

llvm/test/CodeGen/AArch64/sve-vector-deinterleave.ll

 ; CHECK-NEXT:    uunpklo z2.d, z0.s
 ; CHECK-NEXT:    uzp1 z0.d, z2.d, z1.d
 ; CHECK-NEXT:    uzp2 z1.d, z2.d, z1.d
 ; CHECK-NEXT:    ret
 %retval = call {<vscale x 2 x i32>,<vscale x 2 x i32>} @llvm.experimental.vector.deinterleave2.nxv4i32(<vscale x 4 x i32> %vec)
 ret {<vscale x 2 x i32>, <vscale x 2 x i32>} %retval
 }
 
+; Widen illegal type size
+
+define {<vscale x 6 x i32>, <vscale x 6 x i32>} @vector_deinterleave_nxv6i64_nxv12i64(<vscale x 12 x i32> %vec) {
+; CHECK-LABEL: vector_deinterleave_nxv6i64_nxv12i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    str x29, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w29, -16
+; CHECK-NEXT:    addvl sp, sp, #-4
+; CHECK-NEXT:    .cfi_escape 0x0f, 0x0c, 0x8f, 0x00, 0x11, 0x10, 0x22, 0x11, 0x20, 0x92, 0x2e, 0x00, 0x1e, 0x22 // sp + 16 + 32 * VG
+; CHECK-NEXT:    uzp2 z3.s, z2.s, z0.s
+; CHECK-NEXT:    uzp2 z4.s, z0.s, z1.s
+; CHECK-NEXT:    uzp1 z0.s, z0.s, z1.s
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    uzp1 z1.s, z2.s, z0.s
+; CHECK-NEXT:    st1w { z0.s }, p0, [sp, #2, mul vl]
+; CHECK-NEXT:    uunpklo z0.d, z1.s
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    st1w { z0.d }, p1, [sp, #6, mul vl]
+; CHECK-NEXT:    uunpklo z0.d, z3.s
+; CHECK-NEXT:    st1w { z4.s }, p0, [sp]
+; CHECK-NEXT:    st1w { z0.d }, p1, [sp, #2, mul vl]
+; CHECK-NEXT:    ld1w { z0.s }, p0/z, [sp, #2, mul vl]
+; CHECK-NEXT:    ld1w { z1.s }, p0/z, [sp, #3, mul vl]
+; CHECK-NEXT:    ld1w { z2.s }, p0/z, [sp]
+; CHECK-NEXT:    ld1w { z3.s }, p0/z, [sp, #1, mul vl]
+; CHECK-NEXT:    addvl sp, sp, #4
+; CHECK-NEXT:    ldr x29, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+%retval = call {<vscale x 6 x i32>, <vscale x 6 x i32>} @llvm.experimental.vector.deinterleave2.nxv12i64(<vscale x 12 x i32> %vec)
+ret {<vscale x 6 x i32>, <vscale x 6 x i32>} %retval
+}
+
 
 ; Floating declarations
 declare {<vscale x 2 x half>,<vscale x 2 x half>} @llvm.experimental.vector.deinterleave2.nxv4f16(<vscale x 4 x half>)
 declare {<vscale x 4 x half>, <vscale x 4 x half>} @llvm.experimental.vector.deinterleave2.nxv8f16(<vscale x 8 x half>)
 declare {<vscale x 2 x float>, <vscale x 2 x float>} @llvm.experimental.vector.deinterleave2.nxv4f32(<vscale x 4 x float>)
 declare {<vscale x 8 x half>, <vscale x 8 x half>} @llvm.experimental.vector.deinterleave2.nxv16f16(<vscale x 16 x half>)
 declare {<vscale x 4 x float>, <vscale x 4 x float>} @llvm.experimental.vector.deinterleave2.nxv8f32(<vscale x 8 x float>)
 declare {<vscale x 2 x double>, <vscale x 2 x double>} @llvm.experimental.vector.deinterleave2.nxv4f64(<vscale x 4 x double>)

 
 ; Illegal size type
 declare {<vscale x 4 x i64>, <vscale x 4 x i64>} @llvm.experimental.vector.deinterleave2.nxv8i64(<vscale x 8 x i64>)
 declare {<vscale x 8 x i64>, <vscale x 8 x i64>} @llvm.experimental.vector.deinterleave2.nxv16i64(<vscale x 16 x i64>)
 
 declare {<vscale x 8 x i8>, <vscale x 8 x i8>} @llvm.experimental.vector.deinterleave2.nxv16i8(<vscale x 16 x i8>)
 declare {<vscale x 4 x i16>, <vscale x 4 x i16>} @llvm.experimental.vector.deinterleave2.nxv8i16(<vscale x 8 x i16>)
 declare {<vscale x 2 x i32>, <vscale x 2 x i32>} @llvm.experimental.vector.deinterleave2.nxv4i32(<vscale x 4 x i32>)
+
+declare {<vscale x 6 x i32>, <vscale x 6 x i32>} @llvm.experimental.vector.deinterleave2.nxv12i64(<vscale x 12 x i32>)

llvm/test/CodeGen/AArch64/sve-vector-interleave.ll

 ; CHECK-NEXT:    zip2 z2.d, z0.d, z1.d
 ; CHECK-NEXT:    zip1 z0.d, z0.d, z1.d
 ; CHECK-NEXT:    uzp1 z0.s, z0.s, z2.s
 ; CHECK-NEXT:    ret
   %retval = call <vscale x 4 x i32> @llvm.experimental.vector.interleave2.nxv4i32(<vscale x 2 x i32> %vec0, <vscale x 2 x i32> %vec1)
   ret <vscale x 4 x i32> %retval
 }
 
+; Widen illegal type size
+
+define <vscale x 6 x i32> @interleave2_nxv6i32(<vscale x 3 x i32> %vec0, <vscale x 3 x i32> %vec1) nounwind {
+; CHECK-LABEL: interleave2_nxv6i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    zip1 z2.s, z0.s, z1.s
+; CHECK-NEXT:    zip2 z1.s, z0.s, z1.s
+; CHECK-NEXT:    mov z0.d, z2.d
+; CHECK-NEXT:    ret
+  %retval = call <vscale x 6 x i32> @llvm.experimental.vector.interleave2.nxv6i32(<vscale x 3 x i32> %vec0, <vscale x 3 x i32> %vec1)
+  ret <vscale x 6 x i32> %retval
+}
+
+
 ; Float declarations
 declare <vscale x 4 x half> @llvm.experimental.vector.interleave2.nxv4f16(<vscale x 2 x half>, <vscale x 2 x half>)
 declare <vscale x 8 x half> @llvm.experimental.vector.interleave2.nxv8f16(<vscale x 4 x half>, <vscale x 4 x half>)
 declare <vscale x 16 x half> @llvm.experimental.vector.interleave2.nxv16f16(<vscale x 8 x half>, <vscale x 8 x half>)
 declare <vscale x 4 x float> @llvm.experimental.vector.interleave2.nxv4f32(<vscale x 2 x float>, <vscale x 2 x float>)
 declare <vscale x 8 x float> @llvm.experimental.vector.interleave2.nxv8f32(<vscale x 4 x float>, <vscale x 4 x float>)
 declare <vscale x 4 x double> @llvm.experimental.vector.interleave2.nxv4f64(<vscale x 2 x double>, <vscale x 2 x double>)
 

 
 ; Illegal type size
 declare <vscale x 16 x i32> @llvm.experimental.vector.interleave2.nxv16i32(<vscale x 8 x i32>, <vscale x 8 x i32>)
 declare <vscale x 8 x i64> @llvm.experimental.vector.interleave2.nxv8i64(<vscale x 4 x i64>, <vscale x 4 x i64>)
 
 declare <vscale x 16 x i8> @llvm.experimental.vector.interleave2.nxv16i8(<vscale x 8 x i8>, <vscale x 8 x i8>)
 declare <vscale x 8 x i16> @llvm.experimental.vector.interleave2.nxv8i16(<vscale x 4 x i16>, <vscale x 4 x i16>)
 declare <vscale x 4 x i32> @llvm.experimental.vector.interleave2.nxv4i32(<vscale x 2 x i32>, <vscale x 2 x i32>)
+
+declare <vscale x 6 x i32> @llvm.experimental.vector.interleave2.nxv6i32(<vscale x 3 x i32>, <vscale x 3 x i32>)