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 constant positive integer value. The operation - /// does not support returning fixed-width vectors or non-constant operands. + /// IMM, where IMM must be a constant integer value. The operation does not + /// support returning fixed-width vectors or non-constant operands. /// If the sequence value exceeds the limit allowed for the element type then /// the values for those lanes are undefined. STEP_VECTOR, 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 @@ -4709,8 +4709,7 @@ assert(Operand.getValueType().bitsGE(VT.getScalarType()) && "Operand type should be at least as large as the element type"); assert(isa(Operand) && - cast(Operand)->getAPIntValue().isNonNegative() && - "Expected positive integer constant for STEP_VECTOR"); + "Expected integer constant for STEP_VECTOR"); break; case ISD::FREEZE: assert(VT == Operand.getValueType() && "Unexpected VT!"); 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()