From f85a126e76b81bd87b421b84697355517d079c62 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland@gmail.com>
Date: Wed, 6 Mar 2024 12:54:05 -0800
Subject: [PATCH 1/4] [GISEL] Add G_VSCALE instruction

---
 llvm/docs/GlobalISel/GenericOpcode.rst           | 11 +++++++++++
 .../llvm/CodeGen/GlobalISel/MachineIRBuilder.h   | 12 ++++++++++++
 llvm/include/llvm/Support/TargetOpcodes.def      |  3 +++
 llvm/include/llvm/Target/GenericOpcodes.td       |  9 +++++++++
 llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp |  5 +++++
 llvm/lib/CodeGen/MachineVerifier.cpp             | 10 ++++++++++
 .../GlobalISel/legalizer-info-validation.mir     |  3 +++
 llvm/test/MachineVerifier/test_g_vscale.mir      | 16 ++++++++++++++++
 8 files changed, 69 insertions(+)
 create mode 100644 llvm/test/MachineVerifier/test_g_vscale.mir

diff --git a/llvm/docs/GlobalISel/GenericOpcode.rst b/llvm/docs/GlobalISel/GenericOpcode.rst
index f9f9e1186460e..bf1b3cb30a52e 100644
--- a/llvm/docs/GlobalISel/GenericOpcode.rst
+++ b/llvm/docs/GlobalISel/GenericOpcode.rst
@@ -607,6 +607,17 @@ See the LLVM LangRef entry on '``llvm.lround.*'`` for details on behaviour.
 Vector Specific Operations
 --------------------------
 
+G_VSCALE
+^^^^^^^^
+
+Puts the value of the runtime ``vscale`` multiplied by the value in the source
+operand into the destination register. This can be useful in determining the
+actual runtime number of elements in a vector.
+
+.. code-block::
+
+  %0:_(s32) = G_VSCALE 4
+
 G_INSERT_SUBVECTOR
 ^^^^^^^^^^^^^^^^^^
 
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h b/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
index 4732eaf4ee27c..85affb9d784f8 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
@@ -1143,6 +1143,18 @@ class MachineIRBuilder {
   MachineInstrBuilder buildInsert(const DstOp &Res, const SrcOp &Src,
                                   const SrcOp &Op, unsigned Index);
 
+  /// Build and insert \p Res = G_VSCALE \p MinElts
+  ///
+  /// G_VSCALE puts the value of the runtime vscale multiplied by \p MinElts
+  /// into \p Res.
+  ///
+  /// \pre setBasicBlock or setMI must have been called.
+  /// \pre \p Res must be a generic virtual register with scalar type.
+  /// \pre \p Src must be a generic virtual register with scalar type.
+  ///
+  /// \return a MachineInstrBuilder for the newly created instruction.
+  MachineInstrBuilder buildVScale(const DstOp &Res, unsigned MinElts);
+
   /// Build and insert a G_INTRINSIC instruction.
   ///
   /// There are four different opcodes based on combinations of whether the
diff --git a/llvm/include/llvm/Support/TargetOpcodes.def b/llvm/include/llvm/Support/TargetOpcodes.def
index 3dade14f043b6..899eaad5842ae 100644
--- a/llvm/include/llvm/Support/TargetOpcodes.def
+++ b/llvm/include/llvm/Support/TargetOpcodes.def
@@ -727,6 +727,9 @@ HANDLE_TARGET_OPCODE(G_BR)
 /// Generic branch to jump table entry.
 HANDLE_TARGET_OPCODE(G_BRJT)
 
+/// Generic vscale.
+HANDLE_TARGET_OPCODE(G_VSCALE)
+
 /// Generic insert subvector.
 HANDLE_TARGET_OPCODE(G_INSERT_SUBVECTOR)
 
diff --git a/llvm/include/llvm/Target/GenericOpcodes.td b/llvm/include/llvm/Target/GenericOpcodes.td
index 8dc84fb0ba052..b044b8efd4205 100644
--- a/llvm/include/llvm/Target/GenericOpcodes.td
+++ b/llvm/include/llvm/Target/GenericOpcodes.td
@@ -1289,6 +1289,15 @@ def G_MERGE_VALUES : GenericInstruction {
   let variadicOpsType = type1;
 }
 
+// Generic vscale.
+// Puts the value of the runtime vscale multiplied by the value in the source
+// operand into the destination register.
+def G_VSCALE : GenericInstruction {
+  let OutOperandList = (outs type0:$dst);
+  let InOperandList = (ins untyped_imm_0:$src);
+  let hasSideEffects = false;
+}
+
 /// Create a vector from multiple scalar registers. No implicit
 /// conversion is performed (i.e. the result element type must be the
 /// same as all source operands)
diff --git a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
index 9b12d443c96e9..ce6a2f9ca087e 100644
--- a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
@@ -793,6 +793,11 @@ MachineInstrBuilder MachineIRBuilder::buildInsert(const DstOp &Res,
   return buildInstr(TargetOpcode::G_INSERT, Res, {Src, Op, uint64_t(Index)});
 }
 
+MachineInstrBuilder MachineIRBuilder::buildVScale(const DstOp &Res,
+                                                  unsigned MinElts) {
+  return buildInstr(TargetOpcode::G_VSCALE, {Res}, {uint64_t(MinElts)});
+}
+
 static unsigned getIntrinsicOpcode(bool HasSideEffects, bool IsConvergent) {
   if (HasSideEffects && IsConvergent)
     return TargetOpcode::G_INTRINSIC_CONVERGENT_W_SIDE_EFFECTS;
diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index 90cbf097370de..cd0b795e7e39e 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -1613,6 +1613,16 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
       report("G_BSWAP size must be a multiple of 16 bits", MI);
     break;
   }
+  case TargetOpcode::G_VSCALE: {
+    if (!MI->getOperand(1).isImm()) {
+      report("G_VSCALE operand #1 must be an immediate", MI);
+      break;
+    }
+    if (MI->getOperand(1).getImm() == 0) {
+      report("G_VSCALE immediate cannot be zero", MI);
+      break;
+    }
+
   case TargetOpcode::G_INSERT_SUBVECTOR: {
     const MachineOperand &Src0Op = MI->getOperand(1);
     if (!Src0Op.isReg()) {
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
index ac330918b430a..c9ebbf3e738d4 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
@@ -616,6 +616,9 @@
 # DEBUG-NEXT: G_BRJT (opcode {{[0-9]+}}): 2 type indices
 # DEBUG-NEXT: .. the first uncovered type index: 2, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
+# DEBUG-NEXT: G_VSCALE (opcode {{[0-9]+}}): 1 type index, 1 imm index
+# DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
+# DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: G_INSERT_SUBVECTOR (opcode {{[0-9]+}}): 2 type indices, 1 imm index
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 # DEBUG-NEXT: .. imm index coverage check SKIPPED: no rules defined
diff --git a/llvm/test/MachineVerifier/test_g_vscale.mir b/llvm/test/MachineVerifier/test_g_vscale.mir
new file mode 100644
index 0000000000000..74c147a42bc6e
--- /dev/null
+++ b/llvm/test/MachineVerifier/test_g_vscale.mir
@@ -0,0 +1,16 @@
+# RUN: not --crash llc -verify-machineinstrs -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
+# REQUIRES: aarch64-registered-target
+
+---
+name:            g_vscale
+body: |
+  bb.0:
+
+  %1:_(s32) = G_CONSTANT 4
+
+  ; CHECK: G_VSCALE operand #1 must be an immediate
+  %2:_(s32) = G_VSCALE %1
+
+  ; CHECK: G_VSCALE immediate cannot be zero
+  %3:_(s32) = G_VSCALE 0
+...

From 8770ba0541d081397c4954d41646a66cf2d3af48 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland@gmail.com>
Date: Sat, 9 Mar 2024 08:24:47 -0800
Subject: [PATCH 2/4] fixup! remove requires

---
 llvm/test/MachineVerifier/test_g_vscale.mir | 1 -
 1 file changed, 1 deletion(-)

diff --git a/llvm/test/MachineVerifier/test_g_vscale.mir b/llvm/test/MachineVerifier/test_g_vscale.mir
index 74c147a42bc6e..e068c2e1fc464 100644
--- a/llvm/test/MachineVerifier/test_g_vscale.mir
+++ b/llvm/test/MachineVerifier/test_g_vscale.mir
@@ -1,5 +1,4 @@
 # RUN: not --crash llc -verify-machineinstrs -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
-# REQUIRES: aarch64-registered-target
 
 ---
 name:            g_vscale

From e13e848a00e74c327f3f3f38295421be3dc90344 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland@gmail.com>
Date: Mon, 11 Mar 2024 06:37:05 -0700
Subject: [PATCH 3/4] fixup! use ConstantInt to represent G_VSCALE operand

---
 .../llvm/CodeGen/GlobalISel/MachineIRBuilder.h    | 12 +++++++++++-
 llvm/include/llvm/Target/GenericOpcodes.td        |  2 +-
 llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp  | 15 ++++++++++++++-
 llvm/lib/CodeGen/MachineVerifier.cpp              |  6 +++---
 .../GlobalISel/legalizer-info-validation.mir      |  2 +-
 llvm/test/MachineVerifier/test_g_vscale.mir       |  2 +-
 6 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h b/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
index 85affb9d784f8..aaa81342845bf 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
@@ -1150,11 +1150,21 @@ class MachineIRBuilder {
   ///
   /// \pre setBasicBlock or setMI must have been called.
   /// \pre \p Res must be a generic virtual register with scalar type.
-  /// \pre \p Src must be a generic virtual register with scalar type.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildVScale(const DstOp &Res, unsigned MinElts);
 
+  /// Build and insert \p Res = G_VSCALE \p MinElts
+  ///
+  /// G_VSCALE puts the value of the runtime vscale multiplied by \p MinElts
+  /// into \p Res.
+  ///
+  /// \pre setBasicBlock or setMI must have been called.
+  /// \pre \p Res must be a generic virtual register with scalar type.
+  ///
+  /// \return a MachineInstrBuilder for the newly created instruction.
+  MachineInstrBuilder buildVScale(const DstOp &Res, const ConstantInt &MinElts);
+
   /// Build and insert a G_INTRINSIC instruction.
   ///
   /// There are four different opcodes based on combinations of whether the
diff --git a/llvm/include/llvm/Target/GenericOpcodes.td b/llvm/include/llvm/Target/GenericOpcodes.td
index b044b8efd4205..67d405ba96fa1 100644
--- a/llvm/include/llvm/Target/GenericOpcodes.td
+++ b/llvm/include/llvm/Target/GenericOpcodes.td
@@ -1294,7 +1294,7 @@ def G_MERGE_VALUES : GenericInstruction {
 // operand into the destination register.
 def G_VSCALE : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
-  let InOperandList = (ins untyped_imm_0:$src);
+  let InOperandList = (ins unknown:$src);
   let hasSideEffects = false;
 }
 
diff --git a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
index ce6a2f9ca087e..f7aaa0f02efcb 100644
--- a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
@@ -795,7 +795,20 @@ MachineInstrBuilder MachineIRBuilder::buildInsert(const DstOp &Res,
 
 MachineInstrBuilder MachineIRBuilder::buildVScale(const DstOp &Res,
                                                   unsigned MinElts) {
-  return buildInstr(TargetOpcode::G_VSCALE, {Res}, {uint64_t(MinElts)});
+
+  auto IntN = IntegerType::get(getMF().getFunction().getContext(),
+                               Res.getLLTTy(*getMRI()).getScalarSizeInBits());
+  ConstantInt *CI = ConstantInt::get(IntN, MinElts);
+  return buildVScale(Res, *CI);
+}
+
+MachineInstrBuilder MachineIRBuilder::buildVScale(const DstOp &Res,
+                                                  const ConstantInt &MinElts) {
+  auto VScale = buildInstr(TargetOpcode::G_VSCALE);
+  VScale->setDebugLoc(DebugLoc());
+  Res.addDefToMIB(*getMRI(), VScale);
+  VScale.addCImm(&MinElts);
+  return VScale;
 }
 
 static unsigned getIntrinsicOpcode(bool HasSideEffects, bool IsConvergent) {
diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index cd0b795e7e39e..c7098d0b7b912 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -1614,11 +1614,11 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
     break;
   }
   case TargetOpcode::G_VSCALE: {
-    if (!MI->getOperand(1).isImm()) {
-      report("G_VSCALE operand #1 must be an immediate", MI);
+    if (!MI->getOperand(1).isCImm()) {
+      report("G_VSCALE operand must be cimm", MI);
       break;
     }
-    if (MI->getOperand(1).getImm() == 0) {
+    if (MI->getOperand(1).getCImm()->isZero()) {
       report("G_VSCALE immediate cannot be zero", MI);
       break;
     }
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir b/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
index c9ebbf3e738d4..c9e5f8924f8a0 100644
--- a/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir
@@ -616,7 +616,7 @@
 # DEBUG-NEXT: G_BRJT (opcode {{[0-9]+}}): 2 type indices
 # DEBUG-NEXT: .. the first uncovered type index: 2, OK
 # DEBUG-NEXT: .. the first uncovered imm index: 0, OK
-# DEBUG-NEXT: G_VSCALE (opcode {{[0-9]+}}): 1 type index, 1 imm index
+# DEBUG-NEXT: G_VSCALE (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_INSERT_SUBVECTOR (opcode {{[0-9]+}}): 2 type indices, 1 imm index
diff --git a/llvm/test/MachineVerifier/test_g_vscale.mir b/llvm/test/MachineVerifier/test_g_vscale.mir
index e068c2e1fc464..009658a14e436 100644
--- a/llvm/test/MachineVerifier/test_g_vscale.mir
+++ b/llvm/test/MachineVerifier/test_g_vscale.mir
@@ -7,7 +7,7 @@ body: |
 
   %1:_(s32) = G_CONSTANT 4
 
-  ; CHECK: G_VSCALE operand #1 must be an immediate
+  ; CHECK: G_VSCALE operand must be cimm
   %2:_(s32) = G_VSCALE %1
 
   ; CHECK: G_VSCALE immediate cannot be zero

From 36306fddad98eaebc75093d07c48fa57fa0588d9 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland@gmail.com>
Date: Mon, 11 Mar 2024 08:40:02 -0700
Subject: [PATCH 4/4] fixup! fix test case

---
 llvm/lib/CodeGen/MachineVerifier.cpp        | 3 ++-
 llvm/test/MachineVerifier/test_g_vscale.mir | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index c7098d0b7b912..c2d6dd35e1cb2 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -1622,7 +1622,8 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
       report("G_VSCALE immediate cannot be zero", MI);
       break;
     }
-
+    break;
+  }
   case TargetOpcode::G_INSERT_SUBVECTOR: {
     const MachineOperand &Src0Op = MI->getOperand(1);
     if (!Src0Op.isReg()) {
diff --git a/llvm/test/MachineVerifier/test_g_vscale.mir b/llvm/test/MachineVerifier/test_g_vscale.mir
index 009658a14e436..78854620913a1 100644
--- a/llvm/test/MachineVerifier/test_g_vscale.mir
+++ b/llvm/test/MachineVerifier/test_g_vscale.mir
@@ -5,11 +5,11 @@ name:            g_vscale
 body: |
   bb.0:
 
-  %1:_(s32) = G_CONSTANT 4
+  %1:_(s32) = G_CONSTANT i32 4
 
   ; CHECK: G_VSCALE operand must be cimm
   %2:_(s32) = G_VSCALE %1
 
   ; CHECK: G_VSCALE immediate cannot be zero
-  %3:_(s32) = G_VSCALE 0
+  %3:_(s32) = G_VSCALE i32 0
 ...