LLT: Add changeNumElements

[arsenm]

Summary

This is the element analog of changeElementType/changeElementSize

[paquette]

LGTM

[arsenm]

24ab761a60b14cd8824be8d0769d842172ed0334

llvm/include/llvm/Support/LowLevelTypeImpl.h

public:
   /// for pointer types. For pointer types, use changeElementType.
   LLT changeElementSize(unsigned NewEltSize) const {
     assert(!getScalarType().isPointer() &&
            "invalid to directly change element size for pointers");
     return isVector() ? LLT::vector(getNumElements(), NewEltSize)
                       : 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 {
     assert(RawData != 0 && "Invalid Type");
     if (!IsVector) {
       if (!IsPointer)
         return getFieldValue(ScalarSizeFieldInfo);
       else

llvm/unittests/CodeGen/LowLevelTypeTest.cpp

TEST(LowLevelTypeTest, ChangeElementType) {
 
   EXPECT_EQ(P0, S32.changeElementType(P0));
   EXPECT_EQ(S32, P0.changeElementType(S32));
 
   EXPECT_EQ(V2P1, V2P0.changeElementType(P1));
   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
 
 // Invalid to directly change the element size for pointers.
 TEST(LowLevelTypeTest, ChangeElementTypeDeath) {
   const LLT P0 = LLT::pointer(0, 32);
   const LLT V2P0 = LLT::vector(2, P0);