diff --git a/llvm/include/llvm/CodeGen/ISDOpcodes.h b/llvm/include/llvm/CodeGen/ISDOpcodes.h --- a/llvm/include/llvm/CodeGen/ISDOpcodes.h +++ b/llvm/include/llvm/CodeGen/ISDOpcodes.h @@ -594,8 +594,8 @@ /// STEP_VECTOR(IMM) - Returns a scalable vector whose lanes are comprised /// of a linear sequence of unsigned values starting from 0 with a step of - /// IMM, where IMM must be a vector index constant positive integer value - /// which must fit in the vector element type. + /// IMM, where IMM must be a vector index constant integer value which must + /// fit in the vector element type. /// Note that IMM may be a smaller type than the vector element type, in /// which case the step is implicitly zero-extended to the vector element /// type. IMM may also be a larger type than the vector element type, in diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp --- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp @@ -3544,6 +3544,14 @@ return DAG.getNode(ISD::ADD, DL, VT, N0, DAG.getVScale(DL, VT, -IntVal)); } + // canonicalize (sub X, step_vector(C)) to (add X, step_vector(-C)) + if (N1.getOpcode() == ISD::STEP_VECTOR) { + SDValue NewStep = DAG.getConstant(-N1.getConstantOperandAPInt(0), DL, + N1.getOperand(0).getValueType()); + return DAG.getNode(ISD::ADD, DL, VT, N0, + DAG.getStepVector(SDLoc(N), VT, NewStep)); + } + // Prefer an add for more folding potential and possibly better codegen: // sub N0, (lshr N10, width-1) --> add N0, (ashr N10, width-1) if (!LegalOperations && N1.getOpcode() == ISD::SRL && N1.hasOneUse()) { diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp --- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp +++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp @@ -4707,10 +4707,9 @@ "STEP_VECTOR can only be used with vectors of integers that are at " "least 8 bits wide"); assert(isa(Operand) && - cast(Operand)->getAPIntValue().isNonNegative() && cast(Operand)->getAPIntValue().isSignedIntN( VT.getScalarSizeInBits()) && - "Expected STEP_VECTOR integer constant to be positive and fit in " + "Expected STEP_VECTOR integer constant to be fit in " "the vector element type"); break; case ISD::FREEZE: diff --git a/llvm/test/CodeGen/AArch64/sve-stepvector.ll b/llvm/test/CodeGen/AArch64/sve-stepvector.ll --- a/llvm/test/CodeGen/AArch64/sve-stepvector.ll +++ b/llvm/test/CodeGen/AArch64/sve-stepvector.ll @@ -259,6 +259,18 @@ ret %3 } +define @sub_stepvector_nxv8i16() { +; CHECK-LABEL: sub_stepvector_nxv8i16: +; CHECK: // %bb.0: // %entry +; CHECK-NEXT: index z0.h, #2, #-1 +; CHECK-NEXT: ret +entry: + %0 = insertelement poison, i16 2, i32 0 + %1 = shufflevector %0, poison, zeroinitializer + %2 = call @llvm.experimental.stepvector.nxv8i16() + %3 = sub %1, %2 + ret %3 +} declare @llvm.experimental.stepvector.nxv2i64() declare @llvm.experimental.stepvector.nxv4i32()