[llvm] 5072598 - GlobalISel: Add LegalizeMutations to help use More/FewerElements

[Matt Arsenault via llvm-commits]

Author: Matt Arsenault
Date: 2022-04-19T21:04:32-04:00
New Revision: 507259820aa4a6892585d5435c7459fa13a729c1

URL: https://github.com/llvm/llvm-project/commit/507259820aa4a6892585d5435c7459fa13a729c1
DIFF: https://github.com/llvm/llvm-project/commit/507259820aa4a6892585d5435c7459fa13a729c1.diff

LOG: GlobalISel: Add LegalizeMutations to help use More/FewerElements

Added: 
    

Modified: 
    llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
    llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp
    llvm/unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
index 8607db1beadd3..c0cad8ff675dc 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/LegalizerInfo.h
@@ -357,6 +357,14 @@ 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.

diff  --git a/llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp b/llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp
index 75b7fcb5663ab..25c1db91b05d8 100644
--- a/llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp
@@ -43,6 +43,27 @@ LegalizeMutation LegalizeMutations::changeElementTo(unsigned TypeIdx,
   };
 }
 
+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) {

diff  --git a/llvm/unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp b/llvm/unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp
index 7ca4678e2c3ce..0cce67411e220 100644
--- a/llvm/unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp
+++ b/llvm/unittests/CodeGen/GlobalISel/LegalizerInfoTest.cpp
@@ -13,6 +13,8 @@
 
 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 {
@@ -234,6 +236,7 @@ TEST(LegalizerInfoTest, RuleSets) {
   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);
 
@@ -241,6 +244,10 @@ TEST(LegalizerInfoTest, RuleSets) {
   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();
@@ -379,6 +386,47 @@ TEST(LegalizerInfoTest, RuleSets) {
     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) {