[llvm] [CodeGen][MIR] Support parsing of scalable vectors in MIR (PR #70893)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland created https://github.com/llvm/llvm-project/pull/70893

This patch builds on the support for vectors by adding ability to parse scalable vectors in MIR.

>From 2ab957c6d5e74912baf97848ec67fc4d76dd8d73 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Tue, 31 Oct 2023 20:39:46 -0700
Subject: [PATCH] [CodeGen][MIR] Support parsing of scalable vectors in MIR

This patch builds on the support for vectors by adding ability to parse
scalable vectors in MIR.
---
 llvm/lib/CodeGen/MIRParser/MIParser.cpp       | 37 +++++++++++++++----
 .../MIR/Generic/scalable-vector-type.mir      | 23 ++++++++++++
 2 files changed, 52 insertions(+), 8 deletions(-)
 create mode 100644 llvm/test/CodeGen/MIR/Generic/scalable-vector-type.mir

diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index c01b34d6f490b0e..0da664e935f5442 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -1946,12 +1946,28 @@ bool MIParser::parseLowLevelType(StringRef::iterator Loc, LLT &Ty) {
 
   // Now we're looking for a vector.
   if (Token.isNot(MIToken::less))
-    return error(Loc,
-                 "expected sN, pA, <M x sN>, or <M x pA> for GlobalISel type");
+    return error(Loc, "expected sN, pA, <M x sN>, <M x pA>, <vscale x M x sN>, "
+                      "or <vscale x M x pA> for GlobalISel type");
   lex();
 
+  bool HasVScale = Token.stringValue() == "vscale";
+  if (HasVScale) {
+    lex();
+    if (Token.stringValue() != "x")
+      return error("expected <vscale x M x sN> or <vscale x M x pA>");
+    lex();
+  }
+
+  auto GetError = [&](bool HasVScale, StringRef::iterator Loc) {
+    if (HasVScale)
+      return error(
+          Loc, "expected <vscale x M x sN> or <vscale M x pA> for vector type");
+    else
+      return error(Loc, "expected <M x sN> or <M x pA> for vector type");
+  };
+
   if (Token.isNot(MIToken::IntegerLiteral))
-    return error(Loc, "expected <M x sN> or <M x pA> for vector type");
+    return GetError(HasVScale, Loc);
   uint64_t NumElements = Token.integerValue().getZExtValue();
   if (!verifyVectorElementCount(NumElements))
     return error("invalid number of vector elements");
@@ -1959,11 +1975,12 @@ bool MIParser::parseLowLevelType(StringRef::iterator Loc, LLT &Ty) {
   lex();
 
   if (Token.isNot(MIToken::Identifier) || Token.stringValue() != "x")
-    return error(Loc, "expected <M x sN> or <M x pA> for vector type");
+    return GetError(HasVScale, Loc);
   lex();
 
   if (Token.range().front() != 's' && Token.range().front() != 'p')
-    return error(Loc, "expected <M x sN> or <M x pA> for vector type");
+    return GetError(HasVScale, Loc);
+
   StringRef SizeStr = Token.range().drop_front();
   if (SizeStr.size() == 0 || !llvm::all_of(SizeStr, isdigit))
     return error("expected integers after 's'/'p' type character");
@@ -1981,14 +1998,18 @@ bool MIParser::parseLowLevelType(StringRef::iterator Loc, LLT &Ty) {
 
     Ty = LLT::pointer(AS, DL.getPointerSizeInBits(AS));
   } else
-    return error(Loc, "expected <M x sN> or <M x pA> for vector type");
+    return GetError(HasVScale, Loc);
   lex();
 
   if (Token.isNot(MIToken::greater))
-    return error(Loc, "expected <M x sN> or <M x pA> for vector type");
+    return GetError(HasVScale, Loc);
+
   lex();
 
-  Ty = LLT::fixed_vector(NumElements, Ty);
+  if (HasVScale)
+    Ty = LLT::scalable_vector(NumElements, Ty);
+  else
+    Ty = LLT::fixed_vector(NumElements, Ty);
   return false;
 }
 
diff --git a/llvm/test/CodeGen/MIR/Generic/scalable-vector-type.mir b/llvm/test/CodeGen/MIR/Generic/scalable-vector-type.mir
new file mode 100644
index 000000000000000..61ac8411b64505b
--- /dev/null
+++ b/llvm/test/CodeGen/MIR/Generic/scalable-vector-type.mir
@@ -0,0 +1,23 @@
+# RUN: llc -run-pass=none -o - %s | FileCheck %s
+
+---
+name:            scalable_vector_type_s
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: scalable_vector_type_s
+    ; CHECK: [[DEF:%[0-9]+]]:_(<vscale x 1 x s8>) = IMPLICIT_DEF
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<vscale x 1 x s8>) = COPY [[DEF]](<vscale x 1 x s8>)
+    %0:_(<vscale x 1 x s8>) = IMPLICIT_DEF
+    %1:_(<vscale x 1 x s8>) = COPY %0
+...
+
+---
+name:            scalable_vector_type_p
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: scalable_vector_type_p
+    ; CHECK: [[DEF:%[0-9]+]]:_(<vscale x 1 x p0>) = IMPLICIT_DEF
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:_(<vscale x 1 x p0>) = COPY [[DEF]](<vscale x 1 x p0>)
+    %0:_(<vscale x 1 x p0>) = IMPLICIT_DEF
+    %1:_(<vscale x 1 x p0>) = COPY %0
+...