[GlobalISel] Add G_STEP_VECTOR instruction (#115598)

aka llvm.stepvector Intrinsic

Author: Thorsten Schütt

llvm/docs/GlobalISel/GenericOpcode.rst

@@ -753,6 +753,24 @@ The type of the operand must be equal to or larger than the vector element
 type. If the operand is larger than the vector element type, the scalar is
 implicitly truncated to the vector element type.
 
+G_STEP_VECTOR
+^^^^^^^^^^^^^
+
+Create a scalable vector where all lanes are linear sequences starting at 0
+with a given unsigned step.
+
+The type of the operand must be equal to the vector element type. Arithmetic
+is performed modulo the bitwidth of the element. The step must be > 0.
+Otherwise the vector is zero.
+
+.. code-block::
+
+  %0:_(<vscale x 2 x s64>) = G_STEP_VECTOR i64 4
+
+  %1:_(<vscale x s32>) = G_STEP_VECTOR i32 4
+
+  0, 1*Step, 2*Step, 3*Step, 4*Step, ...
+
 G_VECTOR_COMPRESS
 ^^^^^^^^^^^^^^^^^
 

llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h

@@ -906,6 +906,18 @@ class GVScale : public GenericMachineInstr {
   };
 };
 
+/// Represents a step vector.
+class GStepVector : public GenericMachineInstr {
+public:
+  uint64_t getStep() const {
+    return getOperand(1).getCImm()->getValue().getZExtValue();
+  }
+
+  static bool classof(const MachineInstr *MI) {
+    return MI->getOpcode() == TargetOpcode::G_STEP_VECTOR;
+  };
+};
+
 /// Represents an integer subtraction.
 class GSub : public GIntBinOp {
 public:

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

@@ -1172,6 +1172,17 @@ class MachineIRBuilder {
   MachineInstrBuilder buildInsert(const DstOp &Res, const SrcOp &Src,
                                   const SrcOp &Op, unsigned Index);
 
+  /// Build and insert \p Res = G_STEP_VECTOR \p Step
+  ///
+  /// G_STEP_VECTOR returns a scalable vector of linear sequence of step \p Step
+  /// into \p Res.
+  ///
+  /// \pre setBasicBlock or setMI must have been called.
+  /// \pre \p Res must be a generic virtual register with scalable vector type.
+  ///
+  /// \return a MachineInstrBuilder for the newly created instruction.
+  MachineInstrBuilder buildStepVector(const DstOp &Res, unsigned Step);
+
   /// Build and insert \p Res = G_VSCALE \p MinElts
   ///
   /// G_VSCALE puts the value of the runtime vscale multiplied by \p MinElts

llvm/include/llvm/Support/TargetOpcodes.def

@@ -776,6 +776,9 @@ HANDLE_TARGET_OPCODE(G_SHUFFLE_VECTOR)
 /// Generic splatvector.
 HANDLE_TARGET_OPCODE(G_SPLAT_VECTOR)
 
+/// Generic stepvector.
+HANDLE_TARGET_OPCODE(G_STEP_VECTOR)
+
 /// Generic masked compress.
 HANDLE_TARGET_OPCODE(G_VECTOR_COMPRESS)
 

llvm/include/llvm/Target/GenericOpcodes.td

@@ -1590,6 +1590,13 @@ def G_SPLAT_VECTOR: GenericInstruction {
   let hasSideEffects = false;
 }
 
+// Generic stepvector.
+def G_STEP_VECTOR: GenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins unknown:$step);
+  let hasSideEffects = false;
+}
+
 // Generic masked compress.
 def G_VECTOR_COMPRESS: GenericInstruction {
   let OutOperandList = (outs type0:$dst);

llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

@@ -809,6 +809,17 @@ MachineInstrBuilder MachineIRBuilder::buildInsert(const DstOp &Res,
   return buildInstr(TargetOpcode::G_INSERT, Res, {Src, Op, uint64_t(Index)});
 }
 
+MachineInstrBuilder MachineIRBuilder::buildStepVector(const DstOp &Res,
+                                                      unsigned Step) {
+  ConstantInt *CI =
+      ConstantInt::get(getMF().getFunction().getContext(), APInt(64, Step));
+  auto StepVector = buildInstr(TargetOpcode::G_STEP_VECTOR);
+  StepVector->setDebugLoc(DebugLoc());
+  Res.addDefToMIB(*getMRI(), StepVector);
+  StepVector.addCImm(CI);
+  return StepVector;
+}
+
 MachineInstrBuilder MachineIRBuilder::buildVScale(const DstOp &Res,
                                                   unsigned MinElts) {
 

llvm/lib/CodeGen/MachineVerifier.cpp

@@ -1731,6 +1731,36 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
     }
     break;
   }
+  case TargetOpcode::G_STEP_VECTOR: {
+    if (!MI->getOperand(1).isCImm()) {
+      report("operand must be cimm", MI);
+      break;
+    }
+
+    if (!MI->getOperand(1).getCImm()->getValue().isStrictlyPositive()) {
+      report("step must be > 0", MI);
+      break;
+    }
+
+    LLT DstTy = MRI->getType(MI->getOperand(0).getReg());
+    if (!DstTy.isScalableVector()) {
+      report("Destination type must be a scalable vector", MI);
+      break;
+    }
+
+    // <vscale x 2 x p0>
+    if (!DstTy.getElementType().isScalar()) {
+      report("Destination element type must be scalar", MI);
+      break;
+    }
+
+    if (MI->getOperand(1).getCImm()->getBitWidth() !=
+        DstTy.getElementType().getScalarSizeInBits()) {
+      report("step bitwidth differs from result type element bitwidth", MI);
+      break;
+    }
+    break;
+  }
   case TargetOpcode::G_INSERT_SUBVECTOR: {
     const MachineOperand &Src0Op = MI->getOperand(1);
     if (!Src0Op.isReg()) {

llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

@@ -665,6 +665,9 @@
 # DEBUG-NEXT: G_SPLAT_VECTOR (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
+# DEBUG-NEXT: G_STEP_VECTOR (opcode {{[0-9]+}}): 1 type index, 0 imm indices
+# DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_VECTOR_COMPRESS (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected

llvm/test/CodeGen/RISCV/GlobalISel/legalizer-info-validation.mir

@@ -646,6 +646,9 @@
 # DEBUG-NEXT: G_SPLAT_VECTOR (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: user-defined predicate detected
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: user-defined predicate detected
+# DEBUG-NEXT: G_STEP_VECTOR (opcode {{[0-9]+}}): 1 type index, 0 imm indices
+# DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_VECTOR_COMPRESS (opcode {{[0-9]+}}): 2 type indices, 0 imm indices
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined

llvm/test/MachineVerifier/test_step-vector.mir

@@ -0,0 +1,29 @@
+# RUN: not --crash llc -verify-machineinstrs -mtriple=arm64 -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# REQUIRES: aarch64-registered-target
+
+---
+name:            g_step_vector
+body: |
+  bb.0:
+
+  %0:_(s32) = G_CONSTANT i32 4
+
+  ; CHECK: operand must be cimm
+  %1:_(s32) = G_STEP_VECTOR %0
+
+  ; CHECK: step must be > 0
+  %2:_(s32) = G_STEP_VECTOR i32 -1
+
+  ; CHECK: Destination type must be a scalable vector
+  %3:_(<4 x s64>) = G_STEP_VECTOR i32 5
+
+  ; CHECK: Destination element type must be scalar
+  %4:_(<vscale x 4 x p0>) = G_STEP_VECTOR i32 9
+
+  ; CHECK: step bitwidth differs from result type element bitwidth
+  %6:_(<vscale x 2 x s33>) = G_STEP_VECTOR i32 56
+
+  %7:_(<vscale x 2 x s128>) = G_STEP_VECTOR i128 79
+
+...
+