[llvm] 24ab761 - LLT: Add changeNumElements

[Matt Arsenault via llvm-commits]

Author: Matt Arsenault
Date: 2020-01-29T07:32:07-08:00
New Revision: 24ab761a60b14cd8824be8d0769d842172ed0334

URL: https://github.com/llvm/llvm-project/commit/24ab761a60b14cd8824be8d0769d842172ed0334
DIFF: https://github.com/llvm/llvm-project/commit/24ab761a60b14cd8824be8d0769d842172ed0334.diff

LOG: LLT: Add changeNumElements

This is the element analog of changeElementType/changeElementSize

Added: 
    

Modified: 
    llvm/include/llvm/Support/LowLevelTypeImpl.h
    llvm/unittests/CodeGen/LowLevelTypeTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Support/LowLevelTypeImpl.h b/llvm/include/llvm/Support/LowLevelTypeImpl.h
index 6ef7c298bc28..6b0e8d3898e6 100644
--- a/llvm/include/llvm/Support/LowLevelTypeImpl.h
+++ b/llvm/include/llvm/Support/LowLevelTypeImpl.h
@@ -137,6 +137,12 @@ class LLT {
                       : LLT::scalar(NewEltSize);
   }
 
+  /// Return a vector or scalar with the same element type and the new number of
+  /// elements.
+  LLT changeNumElements(unsigned NewNumElts) const {
+    return LLT::scalarOrVector(NewNumElts, getScalarType());
+  }
+
   bool isByteSized() const { return (getSizeInBits() & 7) == 0; }
 
   unsigned getScalarSizeInBits() const {

diff  --git a/llvm/unittests/CodeGen/LowLevelTypeTest.cpp b/llvm/unittests/CodeGen/LowLevelTypeTest.cpp
index bf4277d82fd2..2094be181f8d 100644
--- a/llvm/unittests/CodeGen/LowLevelTypeTest.cpp
+++ b/llvm/unittests/CodeGen/LowLevelTypeTest.cpp
@@ -138,6 +138,29 @@ TEST(LowLevelTypeTest, ChangeElementType) {
   EXPECT_EQ(V2S32, V2P0.changeElementType(S32));
 }
 
+TEST(LowLevelTypeTest, ChangeNumElements) {
+  const LLT P0 = LLT::pointer(0, 32);
+  const LLT V2P0 = LLT::vector(2, P0);
+  const LLT V3P0 = LLT::vector(3, P0);
+
+  const LLT S64 = LLT::scalar(64);
+  const LLT V2S64 = LLT::vector(2, 64);
+  const LLT V3S64 = LLT::vector(3, 64);
+
+  // Vector to scalar
+  EXPECT_EQ(S64, V2S64.changeNumElements(1));
+
+  // Vector to vector
+  EXPECT_EQ(V3S64, V2S64.changeNumElements(3));
+
+  // Scalar to vector
+  EXPECT_EQ(V2S64, S64.changeNumElements(2));
+
+  EXPECT_EQ(P0, V2P0.changeNumElements(1));
+  EXPECT_EQ(V3P0, V2P0.changeNumElements(3));
+  EXPECT_EQ(V2P0, P0.changeNumElements(2));
+}
+
 #ifdef GTEST_HAS_DEATH_TEST
 #ifndef NDEBUG