[llvm] [GISEL] Fix bugs in G_EXTRACT_SUBVECTOR definition (PR #108848)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland updated https://github.com/llvm/llvm-project/pull/108848

>From fc3b43e6ced2dea4fd3aca63d1a051bd8f726326 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Mon, 16 Sep 2024 08:45:56 -0700
Subject: [PATCH 1/2] [RISCV][GISEL] Fix bug in G_EXTRACT_SUBVECTOR definition

The implementation was missing the fact that G_EXTRACT_SUBVECTOR destination
and source vector can be different types.
---
 llvm/include/llvm/Target/GenericOpcodes.td             | 2 +-
 llvm/lib/CodeGen/MachineVerifier.cpp                   | 5 +++++
 llvm/test/MachineVerifier/test_g_extract_subvector.mir | 3 +++
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Target/GenericOpcodes.td b/llvm/include/llvm/Target/GenericOpcodes.td
index a55d9d3b04e62f..f5e62dda6fd043 100644
--- a/llvm/include/llvm/Target/GenericOpcodes.td
+++ b/llvm/include/llvm/Target/GenericOpcodes.td
@@ -1541,7 +1541,7 @@ def G_INSERT_SUBVECTOR : GenericInstruction {
 // Generic extract subvector.
 def G_EXTRACT_SUBVECTOR : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
-  let InOperandList = (ins type0:$src, untyped_imm_0:$idx);
+  let InOperandList = (ins type1:$src, untyped_imm_0:$idx);
   let hasSideEffects = false;
 }
 
diff --git a/llvm/lib/CodeGen/MachineVerifier.cpp b/llvm/lib/CodeGen/MachineVerifier.cpp
index 759201ed9dadc7..6168027369172a 100644
--- a/llvm/lib/CodeGen/MachineVerifier.cpp
+++ b/llvm/lib/CodeGen/MachineVerifier.cpp
@@ -1778,6 +1778,11 @@ void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
       break;
     }
 
+    if (DstTy.isScalable() && SrcTy.isScalable() && DstTy != SrcTy) {
+      report("Scalable vector types must match", MI);
+      break;
+    }
+
     if (IndexOp.getImm() != 0 &&
         SrcTy.getElementCount().getKnownMinValue() % IndexOp.getImm() != 0) {
       report("Index must be a multiple of the source vector's minimum vector "
diff --git a/llvm/test/MachineVerifier/test_g_extract_subvector.mir b/llvm/test/MachineVerifier/test_g_extract_subvector.mir
index 5a441ff29c1721..938bfe3632000e 100644
--- a/llvm/test/MachineVerifier/test_g_extract_subvector.mir
+++ b/llvm/test/MachineVerifier/test_g_extract_subvector.mir
@@ -29,4 +29,7 @@ body:             |
 
     ; CHECK: Index must be a multiple of the source vector's minimum vector length
     %9:_(<vscale x 4 x s32>) = G_EXTRACT_SUBVECTOR  %1, 3
+
+    ; CHECK: Scalable vector types must match
+    %10:_(<vscale x 4 x s32>) = G_EXTRACT_SUBVECTOR  %1, 2
 ...

>From 158a4ba23f3e021635d93bc64d41b20fb4bae831 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Mon, 16 Sep 2024 08:49:57 -0700
Subject: [PATCH 2/2] [GISel] Fix bug where wrong opcode was created in
 buildExtractSubvector

---
 .../CodeGen/GlobalISel/MachineIRBuilder.cpp   |  2 +-
 .../GlobalISel/MachineIRBuilderTest.cpp       | 24 +++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
index 925a1c7cf6aacc..59f2fc633f5de7 100644
--- a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
@@ -942,7 +942,7 @@ MachineInstrBuilder MachineIRBuilder::buildInsertSubvector(const DstOp &Res,
 MachineInstrBuilder MachineIRBuilder::buildExtractSubvector(const DstOp &Res,
                                                             const SrcOp &Src,
                                                             unsigned Idx) {
-  return buildInstr(TargetOpcode::G_INSERT_SUBVECTOR, Res,
+  return buildInstr(TargetOpcode::G_EXTRACT_SUBVECTOR, Res,
                     {Src, uint64_t(Idx)});
 }
 
diff --git a/llvm/unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp b/llvm/unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp
index 2c068ba7f88fe6..3296b8e64cd5c0 100644
--- a/llvm/unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/MachineIRBuilderTest.cpp
@@ -480,3 +480,27 @@ TEST_F(AArch64GISelMITest, BuildFPEnv) {
 
   EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
 }
+
+TEST_F(AArch64GISelMITest, BuildExtractSubvector) {
+  setUp();
+  if (!TM)
+    GTEST_SKIP();
+
+  LLT VecTy = LLT::fixed_vector(4, 32);
+  LLT SubVecTy = LLT::fixed_vector(2, 32);
+  auto Vec = B.buildUndef(VecTy);
+  B.buildExtractSubvector(SubVecTy, Vec, 0);
+
+  VecTy = LLT::scalable_vector(4, 32);
+  Vec = B.buildUndef(VecTy);
+  B.buildExtractSubvector(VecTy, Vec, 0);
+
+  auto CheckStr = R"(
+  ; CHECK: [[DEF:%[0-9]+]]:_(<4 x s32>) = G_IMPLICIT_DEF
+  ; CHECK-NEXT: [[EXTRACT_SUBVECTOR:%[0-9]+]]:_(<2 x s32>) = G_EXTRACT_SUBVECTOR [[DEF]](<4 x s32>), 0
+  ; CHECK-NEXT: [[DEF1:%[0-9]+]]:_(<vscale x 4 x s32>) = G_IMPLICIT_DEF
+  ; CHECK-NEXT: [[EXTRACT_SUBVECTOR1:%[0-9]+]]:_(<vscale x 4 x s32>) = G_EXTRACT_SUBVECTOR [[DEF1]](<vscale x 4 x s32>), 0
+  )";
+
+  EXPECT_TRUE(CheckMachineFunction(*MF, CheckStr)) << *MF;
+}