[CodeGen] Fix incorrect uses of getVectorNumElements()

[david-arm]

Summary

I have fixed up some places in SelectionDAG::getNode() where we
used to assert that the number of vector elements for two types
are the same. I have changed such cases to assert that the
element counts are the same instead. I've added new tests that
exercise the code paths for all the truncations. All the extend
operations are covered by this existing test:

CodeGen/AArch64/sve-sext-zext.ll

For the ISD::SETCC case I fixed this code path is exercised by
these existing tests:

CodeGen/AArch64/sve-fcmp.ll
CodeGen/AArch64/sve-intrinsics-int-compares-with-imm.ll

[sdesmalen]

llvm/test/CodeGen/AArch64/sve-ext-trunc.ll

1 ↗

(On Diff #262025)

You can drop REQUIRES: asserts since this test is equally valid without asserts.

6 ↗

(On Diff #262025)

missing CHECK-NEXT: ret

For a patch like this, you can also generate the CHECK lines with the update_llc_test_checks.py script.

102 ↗

(On Diff #262025)

If this is a no-op, please add a CHECK-NEXT: ret

[david-arm]

llvm/test/CodeGen/AArch64/sve-ext-trunc.ll

1 ↗

(On Diff #262025)

OK, I thought that since I was trying to test the actual code change (within an assert) that the test would only exercise that code path with asserts on. You are right though that it's equally valid without the assert, which is then just testing the codegen is correct. I assume buildbot builds and tests with asserts anyway?

[david-arm]

Hi Sander, I think I've addressed your review comments now, thanks!

[sdesmalen]

LGTM, cheers @david-arm

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

SDValue SelectionDAG::getNode(unsigned Opcode, const SDLoc &DL, EVT VT,
   case ISD::SIGN_EXTEND:
     assert(VT.isInteger() && Operand.getValueType().isInteger() &&
            "Invalid SIGN_EXTEND!");
     assert(VT.isVector() == Operand.getValueType().isVector() &&
            "SIGN_EXTEND result type type should be vector iff the operand "
            "type is vector!");
     if (Operand.getValueType() == VT) return Operand;   // noop extension
     assert((!VT.isVector() ||
-            VT.getVectorNumElements() ==
-            Operand.getValueType().getVectorNumElements()) &&
+            VT.getVectorElementCount() ==
+                Operand.getValueType().getVectorElementCount()) &&
            "Vector element count mismatch!");
     assert(Operand.getValueType().bitsLT(VT) &&
            "Invalid sext node, dst < src!");
     if (OpOpcode == ISD::SIGN_EXTEND || OpOpcode == ISD::ZERO_EXTEND)
       return getNode(OpOpcode, DL, VT, Operand.getOperand(0));
     else if (OpOpcode == ISD::UNDEF)
       // sext(undef) = 0, because the top bits will all be the same.
       return getConstant(0, DL, VT);
     break;
   case ISD::ZERO_EXTEND:
     assert(VT.isInteger() && Operand.getValueType().isInteger() &&
            "Invalid ZERO_EXTEND!");
     assert(VT.isVector() == Operand.getValueType().isVector() &&
            "ZERO_EXTEND result type type should be vector iff the operand "
            "type is vector!");
     if (Operand.getValueType() == VT) return Operand;   // noop extension
     assert((!VT.isVector() ||
-            VT.getVectorNumElements() ==
-            Operand.getValueType().getVectorNumElements()) &&
+            VT.getVectorElementCount() ==
+                Operand.getValueType().getVectorElementCount()) &&
            "Vector element count mismatch!");
     assert(Operand.getValueType().bitsLT(VT) &&
            "Invalid zext node, dst < src!");
     if (OpOpcode == ISD::ZERO_EXTEND)   // (zext (zext x)) -> (zext x)
       return getNode(ISD::ZERO_EXTEND, DL, VT, Operand.getOperand(0));
     else if (OpOpcode == ISD::UNDEF)
       // zext(undef) = 0, because the top bits will be zero.
       return getConstant(0, DL, VT);
     break;
   case ISD::ANY_EXTEND:
     assert(VT.isInteger() && Operand.getValueType().isInteger() &&
            "Invalid ANY_EXTEND!");
     assert(VT.isVector() == Operand.getValueType().isVector() &&
            "ANY_EXTEND result type type should be vector iff the operand "
            "type is vector!");
     if (Operand.getValueType() == VT) return Operand;   // noop extension
     assert((!VT.isVector() ||
-            VT.getVectorNumElements() ==
-            Operand.getValueType().getVectorNumElements()) &&
+            VT.getVectorElementCount() ==
+                Operand.getValueType().getVectorElementCount()) &&
            "Vector element count mismatch!");
     assert(Operand.getValueType().bitsLT(VT) &&
            "Invalid anyext node, dst < src!");
 
     if (OpOpcode == ISD::ZERO_EXTEND || OpOpcode == ISD::SIGN_EXTEND ||
         OpOpcode == ISD::ANY_EXTEND)
       // (ext (zext x)) -> (zext x)  and  (ext (sext x)) -> (sext x)
       return getNode(OpOpcode, DL, VT, Operand.getOperand(0));
SDValue SelectionDAG::getNode(unsigned Opcode, const SDLoc &DL, EVT VT,
   case ISD::TRUNCATE:
     assert(VT.isInteger() && Operand.getValueType().isInteger() &&
            "Invalid TRUNCATE!");
     assert(VT.isVector() == Operand.getValueType().isVector() &&
            "TRUNCATE result type type should be vector iff the operand "
            "type is vector!");
     if (Operand.getValueType() == VT) return Operand;   // noop truncate
     assert((!VT.isVector() ||
-            VT.getVectorNumElements() ==
-            Operand.getValueType().getVectorNumElements()) &&
+            VT.getVectorElementCount() ==
+                Operand.getValueType().getVectorElementCount()) &&
            "Vector element count mismatch!");
     assert(Operand.getValueType().bitsGT(VT) &&
            "Invalid truncate node, src < dst!");
     if (OpOpcode == ISD::TRUNCATE)
       return getNode(ISD::TRUNCATE, DL, VT, Operand.getOperand(0));
     if (OpOpcode == ISD::ZERO_EXTEND || OpOpcode == ISD::SIGN_EXTEND ||
         OpOpcode == ISD::ANY_EXTEND) {
       // If the source is smaller than the dest, we still need an extend.
case ISD::SIGN_EXTEND_INREG: {
     EVT EVT = cast<VTSDNode>(N2)->getVT();
     assert(VT == N1.getValueType() && "Not an inreg extend!");
     assert(VT.isInteger() && EVT.isInteger() &&
            "Cannot *_EXTEND_INREG FP types");
     assert(EVT.isVector() == VT.isVector() &&
            "SIGN_EXTEND_INREG type should be vector iff the operand "
            "type is vector!");
     assert((!EVT.isVector() ||
-            EVT.getVectorNumElements() == VT.getVectorNumElements()) &&
+            EVT.getVectorElementCount() == VT.getVectorElementCount()) &&
            "Vector element counts must match in SIGN_EXTEND_INREG");
     assert(EVT.bitsLE(VT) && "Not extending!");
     if (EVT == VT) return N1;  // Not actually extending
 
     auto SignExtendInReg = [&](APInt Val, llvm::EVT ConstantVT) {
       unsigned FromBits = EVT.getScalarSizeInBits();
       Val <<= Val.getBitWidth() - FromBits;
       Val.ashrInPlace(Val.getBitWidth() - FromBits);
case ISD::CONCAT_VECTORS: {
     break;
   }
   case ISD::SETCC: {
     assert(VT.isInteger() && "SETCC result type must be an integer!");
     assert(N1.getValueType() == N2.getValueType() &&
            "SETCC operands must have the same type!");
     assert(VT.isVector() == N1.getValueType().isVector() &&
            "SETCC type should be vector iff the operand type is vector!");
-    assert((!VT.isVector() ||
-            VT.getVectorNumElements() == N1.getValueType().getVectorNumElements()) &&
+    assert((!VT.isVector() || VT.getVectorElementCount() ==
+                                  N1.getValueType().getVectorElementCount()) &&
            "SETCC vector element counts must match!");
     // Use FoldSetCC to simplify SETCC's.
     if (SDValue V = FoldSetCC(VT, N1, N2, cast<CondCodeSDNode>(N3)->get(), DL))
       return V;
     // Vector constant folding.
     SDValue Ops[] = {N1, N2, N3};
     if (SDValue V = FoldConstantVectorArithmetic(Opcode, DL, VT, Ops)) {
       NewSDValueDbgMsg(V, "New node vector constant folding: ", this);

llvm/test/CodeGen/AArch64/sve-trunc.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+; For all the functions below should the operation is a nop
+define <vscale x 8 x i8> @trunc_i16toi8(<vscale x 8 x i16> %in) {
+; CHECK-LABEL: trunc_i16toi8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 8 x i16> %in to <vscale x 8 x i8>
+  ret <vscale x 8 x i8> %out
+}
+
+define <vscale x 4 x i8> @trunc_i32toi8(<vscale x 4 x i32> %in) {
+; CHECK-LABEL: trunc_i32toi8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 4 x i32> %in to <vscale x 4 x i8>
+  ret <vscale x 4 x i8> %out
+}
+
+define <vscale x 2 x i8> @trunc_i64toi8(<vscale x 2 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi8:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 2 x i64> %in to <vscale x 2 x i8>
+  ret <vscale x 2 x i8> %out
+}
+
+define <vscale x 4 x i16> @trunc_i32toi16(<vscale x 4 x i32> %in) {
+; CHECK-LABEL: trunc_i32toi16:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 4 x i32> %in to <vscale x 4 x i16>
+  ret <vscale x 4 x i16> %out
+}
+
+define <vscale x 2 x i16> @trunc_i64toi16(<vscale x 2 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi16:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 2 x i64> %in to <vscale x 2 x i16>
+  ret <vscale x 2 x i16> %out
+}
+
+define <vscale x 2 x i32> @trunc_i64toi32(<vscale x 2 x i64> %in) {
+; CHECK-LABEL: trunc_i64toi32:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    ret
+entry:
+  %out = trunc <vscale x 2 x i64> %in to <vscale x 2 x i32>
+  ret <vscale x 2 x i32> %out
+}