[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

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

20605

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.isScalableVector() == ConcatSrcVT.isScalableVector() && ConcatSrcNumElts == ExtNumElts)
       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-scalable-vector.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve,+bf16 < %s | FileCheck %s --check-prefixes=CHECK
 
 ; Extracting illegal subvectors
 
+define <4 x i32> @extract_v4i32_nxv16i32(<vscale x 16 x i32> %arg) {
+; CHECK-LABEL: extract_v4i32_nxv16i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    mov z0.d, z3.d
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT:    ret
+  %ext = call <4 x i32> @llvm.experimental.vector.extract.v4i32.nxv16i32(<vscale x 16 x i32> %arg, i64 12)
+  ret <4 x i32> %ext
+}
+
 define <vscale x 1 x i32> @extract_nxv1i32_nxv4i32(<vscale x 4 x i32> %vec) nounwind {
 ; CHECK-LABEL: extract_nxv1i32_nxv4i32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ret
   %retval = call <vscale x 1 x i32> @llvm.experimental.vector.extract.nxv1i32.nxv4i32(<vscale x 4 x i32> %vec, i64 0)
   ret <vscale x 1 x i32> %retval
 }
 

 ; CHECK-NEXT:    uunpkhi z0.s, z1.h
 ; CHECK-NEXT:    ret
   %res = call <vscale x 4 x bfloat> @llvm.experimental.vector.extract.nxv4bf16.nxv16bf16(<vscale x 16 x bfloat> %in, i64 12)
   ret <vscale x 4 x bfloat> %res
 }
 
 declare <vscale x 4 x bfloat> @llvm.experimental.vector.extract.nxv4bf16.nxv16bf16(<vscale x 16 x bfloat>, i64)
 
+declare <4 x i32> @llvm.experimental.vector.extract.v4i32.nxv16i32(<vscale x 16 x i32>, i64)