draft rule composition

Author: tgymnich

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

@@ -14,6 +14,7 @@
 #ifndef LLVM_CODEGEN_GLOBALISEL_LEGALIZERINFO_H
 #define LLVM_CODEGEN_GLOBALISEL_LEGALIZERINFO_H
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/GlobalISel/LegacyLegalizerInfo.h"
@@ -365,6 +366,8 @@ LegalizeMutation changeTo(unsigned TypeIdx, LLT Ty);
 /// Keep the same type as the given type index.
 LegalizeMutation changeTo(unsigned TypeIdx, unsigned FromTypeIdx);
 
+LegalizeMutation changeToInteger(unsigned TypeIdx);
+
 /// Keep the same scalar or element type as the given type index.
 LegalizeMutation changeElementTo(unsigned TypeIdx, unsigned FromTypeIdx);
 
@@ -695,6 +698,57 @@ class LegalizeRuleSet {
     return actionIf(LegalizeAction::Bitcast, Predicate, Mutation);
   }
 
+  LegalizeRuleSet &bitcast(unsigned TypeIdx, LLT From, LLT To) {
+    assert(From.getScalarSizeInBits() == To.getScalarSizeInBits());
+    return actionIf(LegalizeActions::Bitcast,
+                    LegalityPredicates::typeIs(TypeIdx, From),
+                    LegalizeMutations::changeTo(TypeIdx, To));
+  }
+
+  LegalizeRuleSet &bitcast(unsigned TypeIdx, LLT To) {
+    return actionIf(
+        LegalizeActions::Bitcast,
+        LegalityPredicates::sizeIs(TypeIdx, To.getScalarSizeInBits()),
+        LegalizeMutations::changeTo(TypeIdx, To));
+  }
+
+  LegalizeRuleSet &bitcast(unsigned TypeIdx,
+                           std::initializer_list<std::pair<LLT, LLT>> TypeMap) {
+    assert(all_of(TypeMap,
+                  [](auto &P) {
+                    return P.first.getScalarSizeInBits() ==
+                           P.second.getScalarSizeInBits();
+                  }) &&
+           "All type pairs must have the same scalar size");
+
+    SmallVector<std::pair<LLT, LLT>, 4> Types = TypeMap;
+    auto Predicate = [=](const LegalityQuery &Query) {
+      LLT Ty = Query.Types[TypeIdx];
+      return llvm::any_of(Types, [&](auto P) { return P.first == Ty; });
+    };
+
+    auto Mutation = [=](const LegalityQuery &Query) {
+      LLT Ty = Query.Types[TypeIdx];
+      auto Map = llvm::find_if(Types, [&](auto P) { return P.first == Ty; });
+      return std::make_pair(TypeIdx, Map->second);
+    };
+
+    return actionIf(LegalizeActions::Bitcast, Predicate, Mutation);
+  }
+
+  LegalizeRuleSet &bitcastToInteger(unsigned TypeIdx) {
+    return actionIf(LegalizeActions::Bitcast,
+                    LegalityPredicates::isScalar(TypeIdx),
+                    LegalizeMutations::changeToInteger(TypeIdx));
+  }
+
+  LegalizeRuleSet &bitcastToInteger(unsigned TypeIdx,
+                                    std::initializer_list<LLT> TypesInit) {
+    return actionIf(LegalizeActions::Bitcast,
+                    LegalityPredicates::typeInSet(TypeIdx, TypesInit),
+                    LegalizeMutations::changeToInteger(TypeIdx));
+  }
+
   /// The instruction is lowered.
   LegalizeRuleSet &lower() {
     using namespace LegalizeMutations;
@@ -1117,6 +1171,35 @@ class LegalizeRuleSet {
     return clampScalar(TypeIdx, MinTy, MaxTy);
   }
 
+  template <typename Func> LegalizeRuleSet &compose(Func F) {
+    LegalizeRuleSet RuleSet;
+    F(RuleSet);
+
+    for (auto &&[Idx, Rule] : enumerate(RuleSet.Rules)) {
+      LegalityPredicate All = [=](const LegalityQuery &Q) {
+        SmallVector<LLT> Tys(Q.Types);
+        for (auto &Rule : drop_begin(RuleSet.Rules, Idx)) {
+          LegalityQuery NewQ(Q.Opcode, Tys, Q.MMODescrs);
+          if (!Rule.match(NewQ))
+            return false;
+
+          auto &&[Idx, Ty] = Rule.determineMutation(Q);
+          Tys[Idx] = Ty;
+        }
+
+        return true;
+      };
+
+      LegalizeMutation Mutation = [=](const LegalityQuery &Q) {
+        return Rule.determineMutation(Q);
+      };
+
+      add({All, Rule.getAction(), Mutation});
+    }
+
+    return *this;
+  }
+
   /// Limit the range of scalar sizes to MinTy and MaxTy.
   LegalizeRuleSet &clampScalarOrElt(unsigned TypeIdx, const LLT MinTy,
                                     const LLT MaxTy) {

llvm/lib/CodeGen/GlobalISel/LegalizeMutations.cpp

@@ -26,6 +26,12 @@ LegalizeMutation LegalizeMutations::changeTo(unsigned TypeIdx,
   };
 }
 
+LegalizeMutation LegalizeMutations::changeToInteger(unsigned TypeIdx) {
+  return [=](const LegalityQuery &Query) {
+    return std::make_pair(TypeIdx, Query.Types[TypeIdx].dropType());
+  };
+}
+
 LegalizeMutation LegalizeMutations::changeElementTo(unsigned TypeIdx,
                                                     unsigned FromTypeIdx) {
   return [=](const LegalityQuery &Query) {

llvm/lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

@@ -1804,6 +1804,9 @@ AMDGPULegalizerInfo::AMDGPULegalizerInfo(const GCNSubtarget &ST_,
                                  LLT::fixed_vector(2, LocalPtr),
                                  LLT::fixed_vector(2, PrivatePtr)},
                                 {I1, I32})
+      .compose([](LegalizeRuleSet &RuleSet) {
+        RuleSet.bitcastToInteger(0).clampScalar(0, I16, I64);
+      })
       .clampScalar(0, I16, I64)
       .scalarize(1)
       .moreElementsIf(isSmallOddVector(0), oneMoreElement(0))