[SVE] Fix incorrect DAG combines when extracting fixed-width from scalable vectors

[david-arm]

Summary

We were previously silently generating incorrect code when extracting a
fixed-width vector from a scalable vector. This is worse than crashing,
since the user will have no indication that this is currently unsupported
behaviour. I have fixed the code to only perform DAG combines when safe
to do so, i.e. the input and output vectors are both fixed-width or
both scalable.

Test added here:

CodeGen/AArch64/sve-extract-scalable-vector.ll

[david-arm]

llvm/test/CodeGen/AArch64/sve-extract-scalable-vector.ll

6 ↗

(On Diff #375554)

Hi all, my apologies. I just realised I've uploaded the wrong test. I will fix asap!

[david-arm]

• Actually upload the correct test. :)

[CarolineConcatto]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

20613

Can we replace :

if (NVT.isScalableVector() == ConcatSrcVT.isScalableVector() && ConcatSrcNumElts == ExtNumElts)

by:

if (NVT.getElementCount() == ConcatSrcVT.getElementCount() && NVT.isScalableVector())

[david-arm]

Hi @CarolineConcatto, thanks for the suggestion! I couldn't add the && NVT.isScalableVector() bit because I also want that particular DAG combine to trigger for fixed-width vectors too.

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

assert((ExtIdx % ExtNumElts) == 0 &&
            "Extract index is not a multiple of the input vector length.");
 
     unsigned ConcatSrcNumElts = ConcatSrcVT.getVectorMinNumElements();
     unsigned ConcatOpIdx = ExtIdx / ConcatSrcNumElts;
 
     // If the concatenated source types match this extract, it's a direct
     // simplification:
     // extract_subvec (concat V1, V2, ...), i --> Vi
-    if (ConcatSrcNumElts == ExtNumElts)
+    if (NVT.getVectorElementCount() == ConcatSrcVT.getVectorElementCount())
       return V.getOperand(ConcatOpIdx);
 
     // If the concatenated source vectors are a multiple length of this extract,
     // then extract a fraction of one of those source vectors directly from a
     // concat operand. Example:
     //   v2i8 extract_subvec (v16i8 concat (v8i8 X), (v8i8 Y), 14 -->
     //   v2i8 extract_subvec v8i8 Y, 6
-    if (NVT.isFixedLengthVector() && ConcatSrcNumElts % ExtNumElts == 0) {
+    if (NVT.isFixedLengthVector() && ConcatSrcVT.isFixedLengthVector() &&
+        ConcatSrcNumElts % ExtNumElts == 0) {
       SDLoc DL(N);
       unsigned NewExtIdx = ExtIdx - ConcatOpIdx * ConcatSrcNumElts;
       assert(NewExtIdx + ExtNumElts <= ConcatSrcNumElts &&
              "Trying to extract from >1 concat operand?");
       assert(NewExtIdx % ExtNumElts == 0 &&
              "Extract index is not a multiple of the input vector length.");
       SDValue NewIndexC = DAG.getVectorIdxConstant(NewExtIdx, DL);
       return DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, NVT,

llvm/test/CodeGen/AArch64/sve-extract-fixed-from-scalable-vector.ll

+; RUN: not --crash llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s -o - 2>&1 | FileCheck %s --check-prefix=CHECK-ERROR
+
+; Extracting a fixed-length vector from an illegal subvector
+
+; CHECK-ERROR: ERROR: Extracting a fixed-length vector from an illegal scalable vector is not yet supported
+define <4 x i32> @extract_v4i32_nxv16i32_12(<vscale x 16 x i32> %arg) {
+  %ext = call <4 x i32> @llvm.experimental.vector.extract.v4i32.nxv16i32(<vscale x 16 x i32> %arg, i64 12)
+  ret <4 x i32> %ext
+}
+
+declare <4 x i32> @llvm.experimental.vector.extract.v4i32.nxv16i32(<vscale x 16 x i32>, i64)