[llvm] [CodeGen][LLT] Add isFixedVector and isScalableVector (PR #71713)

[Michael Maitland via llvm-commits]

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

The current isScalable function requires a user to call isVector before hand in order to avoid an assertion failure in the case that the LLT is not a vector.

This patch addds helper functions that allow a user to query whether the LLT is fixed or scalable, not wanting an assertion failure in the case that the LLT was never a vector in the first place.

>From a342e9997e7568271a5428d06306c86682ef375d Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Wed, 8 Nov 2023 09:34:42 -0800
Subject: [PATCH] [CodeGen][LLT] Add isFixedVector and isScalableVector

The current isScalable function requires a user to call isVector before
hand in order to avoid an assertion failure in the case that the LLT is
not a vector.

This patch addds helper functions that allow a user to query whether the
LLT is fixed or scalable, not wanting an assertion failure in the case
that the LLT was never a vector in the first place.
---
 llvm/include/llvm/CodeGen/LowLevelType.h    |  8 ++++++++
 llvm/unittests/CodeGen/LowLevelTypeTest.cpp | 14 ++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/LowLevelType.h b/llvm/include/llvm/CodeGen/LowLevelType.h
index b7477834a1ff5e1..d6f1c3b31fbbb5d 100644
--- a/llvm/include/llvm/CodeGen/LowLevelType.h
+++ b/llvm/include/llvm/CodeGen/LowLevelType.h
@@ -163,6 +163,14 @@ class LLT {
                      : getFieldValue(VectorScalableFieldInfo);
   }
 
+  /// Returns true if the LLT is a fixed vector. Returns false otherwise, even
+  /// if the LLT is not a vector type.
+  constexpr bool isFixedVector() const { return isVector() && !isScalable(); }
+
+  /// Returns true if the LLT is a scalable vector. Returns false otherwise,
+  /// even if the LLT is not a vector type.
+  constexpr bool isScalableVector() const { return isVector() && isScalable(); }
+
   constexpr ElementCount getElementCount() const {
     assert(IsVector && "cannot get number of elements on scalar/aggregate");
     return ElementCount::get(IsPointer
diff --git a/llvm/unittests/CodeGen/LowLevelTypeTest.cpp b/llvm/unittests/CodeGen/LowLevelTypeTest.cpp
index 7b13fef23ed00a8..8c909c095cd11ad 100644
--- a/llvm/unittests/CodeGen/LowLevelTypeTest.cpp
+++ b/llvm/unittests/CodeGen/LowLevelTypeTest.cpp
@@ -412,4 +412,18 @@ TEST(LowLevelTypeTest, ConstExpr) {
   EXPECT_EQ(LLT::pointer(0, 32), CEP0);
   EXPECT_EQ(LLT::scalable_vector(2, 32), CESV2S32);
 }
+
+TEST(LowLevelTypeTest, IsFixedVector) {
+  EXPECT_FALSE(LLT::scalar(32).isFixedVector());
+  EXPECT_TRUE(LLT::fixed_vector(2, 32).isFixedVector());
+  EXPECT_FALSE(LLT::scalable_vector(2, 32).isFixedVector());
+  EXPECT_FALSE(LLT::scalable_vector(1, 32).isFixedVector());
+}
+
+TEST(LowLevelTypeTest, IsScalableVector) {
+  EXPECT_FALSE(LLT::scalar(32).isScalableVector());
+  EXPECT_FALSE(LLT::fixed_vector(2, 32).isScalableVector());
+  EXPECT_TRUE(LLT::scalable_vector(2, 32).isScalableVector());
+  EXPECT_TRUE(LLT::scalable_vector(1, 32).isScalableVector());
+}
 }