[CodeGen][SVE] Don't drop scalable flag in DAGCombiner::visitEXTRACT_SUBVECTOR

[sdesmalen]

Summary

There was a rogue 'assert' in AArch64ISelLowering for the tuple.get intrinsics,
that shouldn't really have been there (I suspect this was a remnant from when
we expected the wider vector always to have come from a vector CONCAT).

When I tried to create a more minimal reproducer, I found a bug in
DAGCombiner where it drops the scalable flag when trying to fold:

extract_subv (bitcast X), Index --> bitcast (extract_subv X, Index')

This patch fixes both issues.

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19221–19222

While we're here, do we need to change these getVectorNumElements() calls to getVectorElementCount()?

19238–19239

Does this math work correctly if we're extracting a fixed vector from a scalable vector?

[sdesmalen]

• Removed uses of getVectorNumElements in favour of getVectorMinNumElements
• Removed another warning from IRTranslator, so the test can have a CHECK line for no warnings (coming from TypeSize)

[sdesmalen]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19238–19239

Yes. To be sure, I tested this with some intrinsic that maps to EXTRACT_SUBVECTOR:

define <2 x i64> @extract_2i64_nxv16i8(<vscale x 16 x i8> %z0) {
  %z0_bc = bitcast <vscale x 16 x i8> %z0 to <vscale x 2 x i64>
  %ext = call <2 x i64> @llvm.experimental.extractsubvec.v2i64.nxv2i64(<vscale x 2 x i64> %z0_bc, i32 2)
  ret <2 x i64> %ext
}
=>
Optimized lowered selection DAG: %bb.0 'extract_2i64_nxv16i8:'
SelectionDAG has 9 nodes:
  t0: ch = EntryToken
        t2: nxv16i8,ch = CopyFromReg t0, Register:nxv16i8 %0
      t10: v16i8 = extract_subvector t2, Constant:i64<16>
    t11: v2i64 = bitcast t10
  t7: ch,glue = CopyToReg t0, Register:v2i64 $q0, t11
  t8: ch = AArch64ISD::RET_FLAG t7, Register:v2i64 $q0, t7:1

and for the other:

define <16 x i8> @extract_16i8_nxv2i64(<vscale x 2 x i64> %z0) {
  %z0_bc = bitcast <vscale x 2 x i64> %z0 to <vscale x 16 x i8>
  %ext = call <16 x i8> @llvm.experimental.extractsubvec.v16i8.nxv16i8(<vscale x 16 x i8> %z0_bc, i32 16)
  ret <16 x i8> %ext
}
=>
Optimized lowered selection DAG: %bb.0 'extract_16i8_nxv2i64:'
SelectionDAG has 9 nodes:
  t0: ch = EntryToken
        t2: nxv2i64,ch = CopyFromReg t0, Register:nxv2i64 %0
      t10: v2i64 = extract_subvector t2, Constant:i64<2>
    t11: v16i8 = bitcast t10
  t7: ch,glue = CopyToReg t0, Register:v16i8 $q0, t11
  t8: ch = AArch64ISD::RET_FLAG t7, Register:v16i8 $q0, t7:1

I wasn't planning to add intrinsic as part of this patch to test the behaviour.

[efriedma]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19238–19239

I'm specifically concerned about cases where the number of lanes in the output fixed vector is greater than the number of lanes in the input scalable vector.

[efriedma]

LGTM

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

19238–19239

Actually, hmm, I think that's fine; the math operations doesn't care about the total number of elements in the output.

llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp

bool IRTranslator::runOnMachineFunction(MachineFunction &CurMF) {
   MF->getFrameInfo().setHasMustTailInVarArgFunc(HasMustTailInVarArgFn);
 
   // Make our arguments/constants entry block fallthrough to the IR entry block.
   EntryBB->addSuccessor(&getMBB(F.front()));
 
   // Lower the actual args into this basic block.
   SmallVector<ArrayRef<Register>, 8> VRegArgs;
   for (const Argument &Arg: F.args()) {
-    if (DL->getTypeStoreSize(Arg.getType()) == 0)
+    if (DL->getTypeStoreSize(Arg.getType()).isZero())
       continue; // Don't handle zero sized types.
     ArrayRef<Register> VRegs = getOrCreateVRegs(Arg);
     VRegArgs.push_back(VRegs);
 
     if (Arg.hasSwiftErrorAttr()) {
       assert(VRegs.size() == 1 && "Too many vregs for Swift error");
       SwiftError.setCurrentVReg(EntryBB, SwiftError.getFunctionArg(), VRegs[0]);
     }

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitEXTRACT_SUBVECTOR(SDNode *N) {
   }
 
   // Try to move vector bitcast after extract_subv by scaling extraction index:
   // extract_subv (bitcast X), Index --> bitcast (extract_subv X, Index')
   if (V.getOpcode() == ISD::BITCAST &&
       V.getOperand(0).getValueType().isVector()) {
     SDValue SrcOp = V.getOperand(0);
     EVT SrcVT = SrcOp.getValueType();
-    unsigned SrcNumElts = SrcVT.getVectorNumElements();
-    unsigned DestNumElts = V.getValueType().getVectorNumElements();
+    unsigned SrcNumElts = SrcVT.getVectorMinNumElements();
+    unsigned DestNumElts = V.getValueType().getVectorMinNumElements();
     if ((SrcNumElts % DestNumElts) == 0) {
       unsigned SrcDestRatio = SrcNumElts / DestNumElts;
-      unsigned NewExtNumElts = NVT.getVectorNumElements() * SrcDestRatio;
+      ElementCount NewExtEC = NVT.getVectorElementCount() * SrcDestRatio;
       EVT NewExtVT = EVT::getVectorVT(*DAG.getContext(), SrcVT.getScalarType(),
-                                      NewExtNumElts);
+                                      NewExtEC);
       if (TLI.isOperationLegalOrCustom(ISD::EXTRACT_SUBVECTOR, NewExtVT)) {
         SDLoc DL(N);
         SDValue NewIndex = DAG.getVectorIdxConstant(ExtIdx * SrcDestRatio, DL);
         SDValue NewExtract = DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, NewExtVT,
                                          V.getOperand(0), NewIndex);
         return DAG.getBitcast(NVT, NewExtract);
       }
     }
     if ((DestNumElts % SrcNumElts) == 0) {
       unsigned DestSrcRatio = DestNumElts / SrcNumElts;
-      if ((NVT.getVectorNumElements() % DestSrcRatio) == 0) {
-        unsigned NewExtNumElts = NVT.getVectorNumElements() / DestSrcRatio;
+      if ((NVT.getVectorMinNumElements() % DestSrcRatio) == 0) {
+        ElementCount NewExtEC = NVT.getVectorElementCount() / DestSrcRatio;
         EVT ScalarVT = SrcVT.getScalarType();
         if ((ExtIdx % DestSrcRatio) == 0) {
           SDLoc DL(N);
           unsigned IndexValScaled = ExtIdx / DestSrcRatio;
           EVT NewExtVT =
-              EVT::getVectorVT(*DAG.getContext(), ScalarVT, NewExtNumElts);
+              EVT::getVectorVT(*DAG.getContext(), ScalarVT, NewExtEC);
           if (TLI.isOperationLegalOrCustom(ISD::EXTRACT_SUBVECTOR, NewExtVT)) {
             SDValue NewIndex = DAG.getVectorIdxConstant(IndexValScaled, DL);
             SDValue NewExtract =
                 DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, NewExtVT,
                             V.getOperand(0), NewIndex);
             return DAG.getBitcast(NVT, NewExtract);
           }
-          if (NewExtNumElts == 1 &&
+          if (NewExtEC == 1 &&
               TLI.isOperationLegalOrCustom(ISD::EXTRACT_VECTOR_ELT, ScalarVT)) {
             SDValue NewIndex = DAG.getVectorIdxConstant(IndexValScaled, DL);
             SDValue NewExtract =
                 DAG.getNode(ISD::EXTRACT_VECTOR_ELT, DL, ScalarVT,
                             V.getOperand(0), NewIndex);
             return DAG.getBitcast(NVT, NewExtract);
           }
         }

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

case Intrinsic::aarch64_sve_st1_scatter_scalar_offset:
       return performScatterStoreCombine(N, DAG, AArch64ISD::SST1_IMM_PRED);
     case Intrinsic::aarch64_sve_tuple_get: {
       SDLoc DL(N);
       SDValue Chain = N->getOperand(0);
       SDValue Src1 = N->getOperand(2);
       SDValue Idx = N->getOperand(3);
 
       uint64_t IdxConst = cast<ConstantSDNode>(Idx)->getZExtValue();
-      if (IdxConst > Src1->getNumOperands() - 1)
-        report_fatal_error("index larger than expected");
-
       EVT ResVT = N->getValueType(0);
       uint64_t NumLanes = ResVT.getVectorElementCount().Min;
       SDValue Val =
           DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, ResVT, Src1,
                       DAG.getConstant(IdxConst * NumLanes, DL, MVT::i32));
       return DAG.getMergeValues({Val, Chain}, DL);
     }
     case Intrinsic::aarch64_sve_tuple_set: {

llvm/test/CodeGen/AArch64/sve-extract-subvector.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s 2>%t | FileCheck %s
+; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
+
+; WARN-NOT: warning
+
+; Test that DAGCombiner doesn't drop the scalable flag when it tries to fold:
+;   extract_subv (bitcast X), Index --> bitcast (extract_subv X, Index')
+define <vscale x 16 x i8> @extract_nxv16i8_nxv4i64(<vscale x 4 x i64> %z0_z1) {
+; CHECK-LABEL: extract_nxv16i8_nxv4i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, z1.d
+; CHECK-NEXT:    ret
+  %z0_z1_bc = bitcast <vscale x 4 x i64> %z0_z1 to <vscale x 32 x i8>
+  %ext = call <vscale x 16 x i8> @llvm.aarch64.sve.tuple.get.nxv32i8(<vscale x 32 x i8> %z0_z1_bc, i32 1)
+  ret <vscale x 16 x i8> %ext
+}
+
+
+define <vscale x 2 x i64> @extract_nxv2i64_nxv32i8(<vscale x 32 x i8> %z0_z1) {
+; CHECK-LABEL: extract_nxv2i64_nxv32i8:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, z1.d
+; CHECK-NEXT:    ret
+  %z0_z1_bc = bitcast <vscale x 32 x i8> %z0_z1 to <vscale x 4 x i64>
+  %ext = call <vscale x 2 x i64> @llvm.aarch64.sve.tuple.get.nxv4i64(<vscale x 4 x i64> %z0_z1_bc, i32 1)
+  ret <vscale x 2 x i64> %ext
+}
+
+declare <vscale x 2 x i64> @llvm.aarch64.sve.tuple.get.nxv4i64(<vscale x 4 x i64>, i32)
+declare <vscale x 16 x i8> @llvm.aarch64.sve.tuple.get.nxv32i8(<vscale x 32 x i8>, i32)