[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?

5755

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

11624

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?

5763

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);
       }
     }
   }
+  if (WidenVT.isScalableVector()) {
+    // Break down the inputs such that we can concatenate individual parts
+    // into one wider vector, e.g.
+    //
+    // nxv6i32 = concat(nxv3i32 t0, nxv3i32 t1)
+    // =>
+    // nxv8i32 = concat(nxv1i32 (extract.. t0),
+    //                  nxv1i32 (extract.. t0),
+    //                  nxv1i32 (extract.. t0),
+    //                  nxv1i32 (extract.. t1),
+    //                  nxv1i32 (extract.. t1),
+    //                  nxv1i32 (extract.. t1),
+    //                  undef, undef)
+    unsigned GCD = std::gcd(InVT.getVectorMinNumElements(),
+                            WidenVT.getVectorMinNumElements());
+    EVT PartVT =
+        EVT::getVectorVT(*DAG.getContext(), InVT.getVectorElementType(),
+                         ElementCount::getScalable(GCD));
+    SmallVector<SDValue> Parts(WidenVT.getVectorMinNumElements() / GCD,
+                               DAG.getUNDEF(PartVT));
+    unsigned NumSubParts = InVT.getVectorMinNumElements() / GCD;
+    for (unsigned I = 0; I < NumOperands; ++I)
+      for (unsigned SI = 0; SI < NumSubParts; ++SI)
+        Parts[I * NumSubParts + SI] =
+            DAG.getNode(ISD::EXTRACT_SUBVECTOR, dl, PartVT,
+                        GetWidenedVector(N->getOperand(I)),
+                        DAG.getVectorIdxConstant(SI * GCD, dl));
+    return DAG.getNode(ISD::CONCAT_VECTORS, dl, WidenVT, Parts);
+  }
 
   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);
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();
+  EVT WideVT = InVT;
   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();
+  if (TLI.getTypeAction(Ctx, InVT) == TargetLowering::TypeWidenVector) {
+    WideVT = TLI.getTypeToTransformTo(Ctx, InVT);
+    InVec0 = widenVectorToPartType(DAG, InVec0, DL, WideVT);
+    InVec1 = widenVectorToPartType(DAG, InVec1, DL, WideVT);
+    OutVT = InVec0.getValueType().getDoubleNumVectorElementsVT(Ctx);
+  }
+
   SDValue Res = DAG.getNode(ISD::VECTOR_INTERLEAVE, DL,
-                            DAG.getVTList(InVT, InVT), InVec0, InVec1);
+                            DAG.getVTList(WideVT, WideVT), 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 (InVT.getVectorElementCount() != WideVT.getVectorElementCount()) {
+    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 i64>, <vscale x 6 x i64>} @vector_deinterleave_nxv6i64_nxv12i64(<vscale x 12 x i64> %vec) {
+; CHECK-LABEL: vector_deinterleave_nxv6i64_nxv12i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    uzp1 z6.d, z4.d, z5.d
+; CHECK-NEXT:    uzp1 z7.d, z2.d, z3.d
+; CHECK-NEXT:    uzp1 z24.d, z0.d, z1.d
+; CHECK-NEXT:    uzp2 z25.d, z0.d, z1.d
+; CHECK-NEXT:    uzp2 z26.d, z2.d, z3.d
+; CHECK-NEXT:    uzp2 z5.d, z4.d, z5.d
+; CHECK-NEXT:    mov z0.d, z24.d
+; CHECK-NEXT:    mov z1.d, z7.d
+; CHECK-NEXT:    mov z2.d, z6.d
+; CHECK-NEXT:    mov z3.d, z25.d
+; CHECK-NEXT:    mov z4.d, z26.d
+; CHECK-NEXT:    ret
+%retval = call {<vscale x 6 x i64>, <vscale x 6 x i64>} @llvm.experimental.vector.deinterleave2.nxv12i64(<vscale x 12 x i64> %vec)
+ret {<vscale x 6 x i64>, <vscale x 6 x i64>} %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 i64>, <vscale x 6 x i64>} @llvm.experimental.vector.deinterleave2.nxv12i64(<vscale x 12 x i64>)

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>)