GlobalISel: Add LegalizeMutations to help use More/FewerElements

[paquette]

Testcase?

[arsenm]

In D123561#3446114, @paquette wrote:

Testcase?

I don't have an immediate use since the lowering in D123563 shouldn't actually be using MoreElements

[paquette]

Ah. There are no C++ unit tests that can be updated without an immediate use case?

[arsenm]

Add unit test. Found asserts in scalable vectors

[paquette]

LGTM

[arsenm]

507259820aa4a6892585d5435c7459fa13a729c1

llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h

 LegalizeMutation changeTo(unsigned TypeIdx, unsigned FromTypeIdx);
 
 /// Keep the same scalar or element type as the given type index.
 LegalizeMutation changeElementTo(unsigned TypeIdx, unsigned FromTypeIdx);
 
 /// Keep the same scalar or element type as the given type.
 LegalizeMutation changeElementTo(unsigned TypeIdx, LLT Ty);
 
+/// Keep the same scalar or element type as \p TypeIdx, but take the number of
+/// elements from \p FromTypeIdx.
+LegalizeMutation changeElementCountTo(unsigned TypeIdx, unsigned FromTypeIdx);
+
+/// Keep the same scalar or element type as \p TypeIdx, but take the number of
+/// elements from \p Ty.
+LegalizeMutation changeElementCountTo(unsigned TypeIdx, LLT Ty);
+
 /// Change the scalar size or element size to have the same scalar size as type
 /// index \p FromIndex. Unlike changeElementTo, this discards pointer types and
 /// only changes the size.
 LegalizeMutation changeElementSizeTo(unsigned TypeIdx, unsigned FromTypeIdx);
 
 /// Widen the scalar type or vector element type for the given type index to the
 /// next power of 2.
 LegalizeMutation widenScalarOrEltToNextPow2(unsigned TypeIdx, unsigned Min = 0);

llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp

 LegalizeMutation LegalizeMutations::changeElementTo(unsigned TypeIdx,
                                                     LLT NewEltTy) {
   return [=](const LegalityQuery &Query) {
     const LLT OldTy = Query.Types[TypeIdx];
     return std::make_pair(TypeIdx, OldTy.changeElementType(NewEltTy));
   };
 }
 
+LegalizeMutation LegalizeMutations::changeElementCountTo(unsigned TypeIdx,
+                                                         unsigned FromTypeIdx) {
+  return [=](const LegalityQuery &Query) {
+    const LLT OldTy = Query.Types[TypeIdx];
+    const LLT NewTy = Query.Types[FromTypeIdx];
+    ElementCount NewEltCount =
+        NewTy.isVector() ? NewTy.getElementCount() : ElementCount::getFixed(1);
+    return std::make_pair(TypeIdx, OldTy.changeElementCount(NewEltCount));
+  };
+}
+
+LegalizeMutation LegalizeMutations::changeElementCountTo(unsigned TypeIdx,
+                                                         LLT NewEltTy) {
+  return [=](const LegalityQuery &Query) {
+    const LLT OldTy = Query.Types[TypeIdx];
+    ElementCount NewEltCount = NewEltTy.isVector() ? NewEltTy.getElementCount()
+                                                   : ElementCount::getFixed(1);
+    return std::make_pair(TypeIdx, OldTy.changeElementCount(NewEltCount));
+  };
+}
+
 LegalizeMutation LegalizeMutations::changeElementSizeTo(unsigned TypeIdx,
                                                         unsigned FromTypeIdx) {
   return [=](const LegalityQuery &Query) {
     const LLT OldTy = Query.Types[TypeIdx];
     const LLT NewTy = Query.Types[FromTypeIdx];
     const LLT NewEltTy = LLT::scalar(NewTy.getScalarSizeInBits());
     return std::make_pair(TypeIdx, OldTy.changeElementType(NewEltTy));
   };

llvm/unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp

 //===- llvm/unittest/CodeGen/GlobalISel/LegalizerInfoTest.cpp -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/LegalizerInfo.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "GISelMITest.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 using namespace LegalizeActions;
+using namespace LegalityPredicates;
+using namespace LegalizeMutations;
 
 // Define a couple of pretty printers to help debugging when things go wrong.
 namespace llvm {
 std::ostream &
 operator<<(std::ostream &OS, const LegalizeAction Act) {
   switch (Act) {
   case Lower: OS << "Lower"; break;
   case Legal: OS << "Legal"; break;
TEST(LegalizerInfoTest, RuleSets) {
   const LLT s64 = LLT::scalar(64);
 
   const LLT v2s5 = LLT::fixed_vector(2, 5);
   const LLT v2s8 = LLT::fixed_vector(2, 8);
   const LLT v2s16 = LLT::fixed_vector(2, 16);
   const LLT v2s32 = LLT::fixed_vector(2, 32);
   const LLT v3s32 = LLT::fixed_vector(3, 32);
   const LLT v4s32 = LLT::fixed_vector(4, 32);
+  const LLT v8s32 = LLT::fixed_vector(8, 32);
   const LLT v2s33 = LLT::fixed_vector(2, 33);
   const LLT v2s64 = LLT::fixed_vector(2, 64);
 
   const LLT p0 = LLT::pointer(0, 32);
   const LLT v3p0 = LLT::fixed_vector(3, p0);
   const LLT v4p0 = LLT::fixed_vector(4, p0);
 
+  const LLT s1 = LLT::scalar(1);
+  const LLT v2s1 = LLT::fixed_vector(2, 1);
+  const LLT v4s1 = LLT::fixed_vector(4, 1);
+
   {
     LegalizerInfo LI;
     auto &LegacyInfo = LI.getLegacyLegalizerInfo();
 
     LI.getActionDefinitionsBuilder(G_IMPLICIT_DEF)
       .legalFor({v4s32, v4p0})
       .moreElementsToNextPow2(0);
     LegacyInfo.computeTables();
// Test widenScalarToNextPow2
 
     EXPECT_ACTION(WidenScalar, 0, s32, LegalityQuery(G_AND, {s5}));
     EXPECT_ACTION(WidenScalar, 0, s64, LegalityQuery(G_AND, {s33}));
 
     // Do nothing for vectors.
     EXPECT_ACTION(Unsupported, 0, LLT(), LegalityQuery(G_AND, {v2s5}));
     EXPECT_ACTION(Unsupported, 0, LLT(), LegalityQuery(G_AND, {v2s33}));
   }
+
+  // Test changeElementCountTo
+  {
+    LegalizerInfo LI;
+    auto &LegacyInfo = LI.getLegacyLegalizerInfo();
+
+    // Type index form
+    LI.getActionDefinitionsBuilder(G_SELECT)
+      .moreElementsIf(isScalar(1), changeElementCountTo(1, 0));
+
+    // Raw type form
+    LI.getActionDefinitionsBuilder(G_ADD)
+      .fewerElementsIf(typeIs(0, v4s32), changeElementCountTo(0, v2s32))
+      .fewerElementsIf(typeIs(0, v8s32), changeElementCountTo(0, s32))
+      .fewerElementsIf(typeIs(0, LLT::scalable_vector(4, 16)),
+                       changeElementCountTo(0, LLT::scalable_vector(2, 16)))
+      .fewerElementsIf(typeIs(0, LLT::scalable_vector(8, 16)),
+                       changeElementCountTo(0, s16));
+
+    LegacyInfo.computeTables();
+
+    EXPECT_ACTION(MoreElements, 1, v4s1, LegalityQuery(G_SELECT, {v4s32, s1}));
+    EXPECT_ACTION(MoreElements, 1, v2s1, LegalityQuery(G_SELECT, {v2s32, s1}));
+    EXPECT_ACTION(MoreElements, 1, v2s1, LegalityQuery(G_SELECT, {v2s32, s1}));
+    EXPECT_ACTION(MoreElements, 1, v4s1, LegalityQuery(G_SELECT, {v4p0, s1}));
+
+    EXPECT_ACTION(MoreElements, 1, LLT::scalable_vector(2, 1),
+                  LegalityQuery(G_SELECT, {LLT::scalable_vector(2, 32), s1}));
+    EXPECT_ACTION(MoreElements, 1, LLT::scalable_vector(4, 1),
+                  LegalityQuery(G_SELECT, {LLT::scalable_vector(4, 32), s1}));
+    EXPECT_ACTION(MoreElements, 1, LLT::scalable_vector(2, s1),
+                  LegalityQuery(G_SELECT, {LLT::scalable_vector(2, p0), s1}));
+
+    EXPECT_ACTION(FewerElements, 0, v2s32, LegalityQuery(G_ADD, {v4s32}));
+    EXPECT_ACTION(FewerElements, 0, s32, LegalityQuery(G_ADD, {v8s32}));
+
+    EXPECT_ACTION(FewerElements, 0, LLT::scalable_vector(2, 16),
+                  LegalityQuery(G_ADD, {LLT::scalable_vector(4, 16)}));
+    EXPECT_ACTION(FewerElements, 0, s16,
+                  LegalityQuery(G_ADD, {LLT::scalable_vector(8, 16)}));
+  }
 }
 
 TEST(LegalizerInfoTest, MMOAlignment) {
   using namespace TargetOpcode;
 
   const LLT s32 = LLT::scalar(32);
   const LLT p0 = LLT::pointer(0, 64);