[RISCV][GlobalISel] Generate shift selections

3 files changed, 181 insertions, 14 deletions

diff --git a/backend/utils/TableGen/Common/CodeGenSelection.cpp b/backend/utils/TableGen/Common/CodeGenSelection.cpp
index 1e65e04..88f85f7 100644
--- a/backend/utils/TableGen/Common/CodeGenSelection.cpp
+++ b/backend/utils/TableGen/Common/CodeGenSelection.cpp
@@ -265,7 +265,8 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
     for (const SourceSelectionOperandDefConstantPredicate &Predicate : Rule.OperandDefConstants) {
       if (!ConstantOperands.insert(Predicate.OperandIndex).second ||
           (MatchInstruction != nullptr &&
-           !isRegisterOperand(*MatchInstruction, Predicate.OperandIndex)) || Predicate.Ranges.empty()) {
+           !isRegisterOperand(*MatchInstruction, Predicate.OperandIndex)) ||
+          Predicate.Ranges.empty()) {
         fail(Predicate.Range,
              "selection operand definition constant predicates require distinct register "
              "operands and at least one range");
@@ -300,6 +301,81 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
       ++Normalized.Specificity;
     }
 
+    std::unordered_map<std::string, std::uint32_t> CaptureIDs;
+    std::unordered_set<std::uint32_t> DefUseRootOperands;
+    for (const SourceSelectionDefUsePredicate &Predicate : Rule.DefUsePredicates) {
+      const auto DefOpcode = resolveOpcode(Result, Predicate.OpcodeSymbol);
+      if (!DefOpcode || !isGenericOpcode(Opcode::fromValue(*DefOpcode)) ||
+          !DefUseRootOperands.insert(Predicate.RootOperandIndex).second ||
+          (MatchInstruction != nullptr &&
+           !isRegisterOperand(*MatchInstruction, Predicate.RootOperandIndex))) {
+        fail(Predicate.Range,
+             "selection def-use predicates require distinct register operands and generic opcodes");
+        return false;
+      }
+
+      NormalizedSelectionDefUsePredicate NormalizedPredicate{
+          .RootOperandIndex = Predicate.RootOperandIndex,
+          .Opcode = *DefOpcode,
+          .Captures = {},
+          .OperandDefConstants = {},
+      };
+      std::unordered_set<std::uint32_t> CapturedOperands;
+      for (const SourceSelectionDefUseCapture &Capture : Predicate.Captures) {
+        if (!isIdentifier(Capture.Symbol) ||
+            !CapturedOperands.insert(Capture.OperandIndex).second ||
+            !CaptureIDs.emplace(Capture.Symbol, CaptureIDs.size() + 1U).second) {
+          fail(Capture.Range,
+               "selection def-use captures require unique identifiers and distinct operands");
+          return false;
+        }
+        NormalizedPredicate.Captures.push_back({
+            .ID = CaptureIDs.at(Capture.Symbol),
+            .OperandIndex = Capture.OperandIndex,
+        });
+      }
+
+      std::unordered_set<std::uint32_t> ConstantOperandsInDef;
+      for (const SourceSelectionOperandDefConstantPredicate &Constant :
+           Predicate.OperandDefConstants) {
+        if (!ConstantOperandsInDef.insert(Constant.OperandIndex).second ||
+            Constant.Ranges.empty()) {
+          fail(Constant.Range, "selection def-use constant predicates require distinct operands "
+                               "and at least one range");
+          return false;
+        }
+        NormalizedSelectionOperandDefConstantPredicate NormalizedConstant{
+            .OperandIndex = Constant.OperandIndex,
+            .Ranges = {},
+        };
+        for (const SourceSelectionOperandDefConstantPredicate::ConstantRange &Range :
+             Constant.Ranges) {
+          if (Range.Minimum > Range.Maximum) {
+            fail(Range.SourceRange, "selection def-use constant range is invalid");
+            return false;
+          }
+          NormalizedConstant.Ranges.push_back({
+              .Minimum = Range.Minimum,
+              .Maximum = Range.Maximum,
+          });
+        }
+        std::ranges::sort(NormalizedConstant.Ranges, {},
+                          &NormalizedSelectionOperandDefConstantPredicate::Range::Minimum);
+        for (std::size_t Index = 1; Index < NormalizedConstant.Ranges.size(); ++Index) {
+          const auto &Previous = NormalizedConstant.Ranges[Index - 1];
+          const auto &Current = NormalizedConstant.Ranges[Index];
+          if (Current.Minimum <= Previous.Maximum) {
+            fail(Constant.Range, "selection def-use constant ranges overlap");
+            return false;
+          }
+        }
+        NormalizedPredicate.OperandDefConstants.push_back(std::move(NormalizedConstant));
+      }
+      Normalized.DefUsePredicates.push_back(std::move(NormalizedPredicate));
+      Normalized.Specificity += 1U + static_cast<std::uint32_t>(Predicate.Captures.size()) +
+                                static_cast<std::uint32_t>(Predicate.OperandDefConstants.size());
+    }
+
     std::unordered_map<std::string, std::uint32_t> TemporaryIDs;
     for (const SourceSelectionTemporary &Temporary : Rule.Temporaries) {
       if (!isIdentifier(Temporary.Symbol) ||
@@ -346,6 +422,22 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
       };
       NormalizedAction.OperandRenderers.reserve(Action.OperandRenderers.size());
       for (const SourceSelectionOperandRenderer &Renderer : Action.OperandRenderers) {
+        if (Renderer.ComplexOperandIndex) {
+          if (Renderer.ComplexRendererSymbol.empty() || !Renderer.Flags.empty() ||
+              Renderer.ImmediateTransform != "identity") {
+            fail(Renderer.Range,
+                 "selection complex operand renderer requires a matcher and no flags or transform");
+            return false;
+          }
+          NormalizedAction.OperandRenderers.push_back({
+              .SourceKind = NormalizedSelectionOperandRenderer::Kind::ComplexOperand,
+              .Value = *Renderer.ComplexOperandIndex,
+              .ComplexRenderer = Renderer.ComplexRendererSymbol,
+              .Flags = 0,
+          });
+          continue;
+        }
+
         if (Renderer.OperandIndex) {
           if (MatchInstruction != nullptr &&
               *Renderer.OperandIndex >= MatchInstruction->Operands.size()) {
@@ -359,6 +451,7 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
           NormalizedAction.OperandRenderers.push_back({
               .SourceKind = NormalizedSelectionOperandRenderer::Kind::Operand,
               .Value = *Renderer.OperandIndex,
+              .ComplexRenderer = {},
               .Flags = 0,
           });
           continue;
@@ -378,6 +471,7 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
           NormalizedAction.OperandRenderers.push_back({
               .SourceKind = NormalizedSelectionOperandRenderer::Kind::PhysicalRegister,
               .Value = Physical->ID,
+              .ComplexRenderer = {},
               .Flags = Flags->rawValue(),
           });
           continue;
@@ -395,14 +489,32 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
           NormalizedAction.OperandRenderers.push_back({
               .SourceKind = NormalizedSelectionOperandRenderer::Kind::Temporary,
               .Value = Temporary->second,
+              .ComplexRenderer = {},
               .Flags = Flags->rawValue(),
           });
           continue;
         }
 
+        if (!Renderer.CapturedOperandSymbol.empty()) {
+          const auto Capture = CaptureIDs.find(Renderer.CapturedOperandSymbol);
+          if (Capture == CaptureIDs.end() || !Renderer.Flags.empty() ||
+              Renderer.ImmediateTransform != "identity") {
+            fail(Renderer.Range,
+                 "selection captured operand renderer requires a matched capture and no flags");
+            return false;
+          }
+          NormalizedAction.OperandRenderers.push_back({
+              .SourceKind = NormalizedSelectionOperandRenderer::Kind::CapturedOperand,
+              .Value = Capture->second,
+              .ComplexRenderer = {},
+              .Flags = 0,
+          });
+          continue;
+        }
+
         if (Renderer.CapturedImmediateOperand) {
-          const auto Transform = [&Renderer]()
-              -> std::optional<NormalizedSelectionImmediateTransform> {
+          const auto Transform =
+              [&Renderer]() -> std::optional<NormalizedSelectionImmediateTransform> {
             if (Renderer.ImmediateTransform == "identity")
               return NormalizedSelectionImmediateTransform::Identity;
             if (Renderer.ImmediateTransform == "negate")
@@ -421,6 +533,7 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
           NormalizedAction.OperandRenderers.push_back({
               .SourceKind = NormalizedSelectionOperandRenderer::Kind::CapturedImmediate,
               .Value = *Renderer.CapturedImmediateOperand,
+              .ComplexRenderer = {},
               .Flags = 0,
               .Transform = *Transform,
           });
@@ -435,6 +548,7 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
           NormalizedAction.OperandRenderers.push_back({
               .SourceKind = NormalizedSelectionOperandRenderer::Kind::Immediate,
               .Value = 0,
+              .ComplexRenderer = {},
               .Flags = 0,
               .Immediate = *Renderer.Immediate,
           });
@@ -458,8 +572,9 @@ bool CodeGenTarget::normalizeSelectionRules(NormalizedTargetModel &Result) {
         fail(Rule.Range, "selection erased definitions require distinct register operands");
         return false;
       }
-      if (!ConstantOperands.contains(OperandIndex)) {
-        fail(Rule.Range, "selection erased definition requires a matched constant operand");
+      if (!ConstantOperands.contains(OperandIndex) && !DefUseRootOperands.contains(OperandIndex)) {
+        fail(Rule.Range,
+             "selection erased definition requires a matched constant or def-use operand");
         return false;
       }
     }
diff --git a/backend/utils/TableGen/Common/CodeGenTarget.h b/backend/utils/TableGen/Common/CodeGenTarget.h
index 42a28c1..47dd598 100644
--- a/backend/utils/TableGen/Common/CodeGenTarget.h
+++ b/backend/utils/TableGen/Common/CodeGenTarget.h
@@ -204,6 +204,18 @@ struct NormalizedSelectionOperandDefConstantPredicate final {
   std::vector<Range> Ranges;
 };
 
+struct NormalizedSelectionDefUseCapture final {
+  std::uint32_t ID = 0;
+  std::uint32_t OperandIndex = 0;
+};
+
+struct NormalizedSelectionDefUsePredicate final {
+  std::uint32_t RootOperandIndex = 0;
+  std::uint32_t Opcode = 0;
+  std::vector<NormalizedSelectionDefUseCapture> Captures;
+  std::vector<NormalizedSelectionOperandDefConstantPredicate> OperandDefConstants;
+};
+
 struct NormalizedSelectionTemporary final {
   std::uint32_t ID = 0;
   std::uint32_t TypeFromOperand = 0;
@@ -218,17 +230,19 @@ enum class NormalizedSelectionImmediateTransform : std::uint8_t {
 struct NormalizedSelectionOperandRenderer final {
   enum class Kind : std::uint8_t {
     Operand,
+    ComplexOperand,
     PhysicalRegister,
     Temporary,
+    CapturedOperand,
     CapturedImmediate,
     Immediate,
   };
 
   Kind SourceKind = Kind::Operand;
   std::uint32_t Value = 0;
+  std::string ComplexRenderer;
   std::uint16_t Flags = 0;
-  NormalizedSelectionImmediateTransform Transform =
-      NormalizedSelectionImmediateTransform::Identity;
+  NormalizedSelectionImmediateTransform Transform = NormalizedSelectionImmediateTransform::Identity;
   std::int64_t Immediate = 0;
 };
 
@@ -249,6 +263,7 @@ struct NormalizedSelectionRule final {
   std::vector<NormalizedSelectionOperandBankPredicate> OperandBanks;
   std::vector<NormalizedSelectionOperandClassPredicate> OperandClasses;
   std::vector<NormalizedSelectionOperandDefConstantPredicate> OperandDefConstants;
+  std::vector<NormalizedSelectionDefUsePredicate> DefUsePredicates;
   std::vector<NormalizedSelectionTemporary> Temporaries;
   std::vector<NormalizedSelectionAction> Actions;
   std::vector<std::uint32_t> EraseOperandDefinitions;
diff --git a/backend/utils/TableGen/GlobalISelEmitter.cpp b/backend/utils/TableGen/GlobalISelEmitter.cpp
index 784d418..562ff7c 100644
--- a/backend/utils/TableGen/GlobalISelEmitter.cpp
+++ b/backend/utils/TableGen/GlobalISelEmitter.cpp
@@ -26,9 +26,11 @@ namespace {
 
 void emitInstruction(std::ostream &Output, std::string_view Opcode, std::uint32_t OperandIndex,
                      std::uint32_t Value = 0, std::string_view Auxiliary = "0",
-                     std::int64_t Immediate = 0) {
+                     std::int64_t Immediate = 0, std::string_view Symbol = {}) {
   Output << "    {GIMatchTableOpcode::" << Opcode << ", " << OperandIndex << ", " << Value << ", "
-         << Auxiliary << ", " << Immediate << "},\n";
+         << Auxiliary << ", " << Immediate << ", ";
+  appendQuoted(Output, Symbol);
+  Output << "},\n";
 }
 
 void emitTypePredicate(std::ostream &Output, std::uint32_t OperandIndex,
@@ -39,11 +41,12 @@ void emitTypePredicate(std::ostream &Output, std::uint32_t OperandIndex,
 }
 
 [[nodiscard]] std::size_t instructionCount(const NormalizedSelectionRule &Rule) noexcept {
-  std::size_t Count = 2U + (Rule.ResultType ? 1U : 0U) +
-                      (Rule.IntegerPredicate ? 1U : 0U) +
-         Rule.OperandTypes.size() + Rule.OperandBanks.size() + Rule.OperandClasses.size() +
-                      Rule.OperandDefConstants.size() + Rule.Temporaries.size() +
-                      Rule.EraseOperandDefinitions.size();
+  std::size_t Count = 2U + (Rule.ResultType ? 1U : 0U) + (Rule.IntegerPredicate ? 1U : 0U) +
+                      Rule.OperandTypes.size() + Rule.OperandBanks.size() +
+                      Rule.OperandClasses.size() + Rule.OperandDefConstants.size() +
+                      Rule.Temporaries.size() + Rule.EraseOperandDefinitions.size();
+  for (const NormalizedSelectionDefUsePredicate &Predicate : Rule.DefUsePredicates)
+    Count += 1U + Predicate.Captures.size() + Predicate.OperandDefConstants.size();
   for (const NormalizedSelectionAction &Action : Rule.Actions)
     Count += 1U + Action.OperandRenderers.size();
   return Count;
@@ -61,6 +64,12 @@ std::string emitGlobalISel(const NormalizedTargetModel &Model, std::string_view
          Rule.OperandDefConstants) {
       ConstantRangeCount += Predicate.Ranges.size();
     }
+    for (const NormalizedSelectionDefUsePredicate &DefUse : Rule.DefUsePredicates) {
+      for (const NormalizedSelectionOperandDefConstantPredicate &Predicate :
+           DefUse.OperandDefConstants) {
+        ConstantRangeCount += Predicate.Ranges.size();
+      }
+    }
   }
 
   std::ostringstream Output;
@@ -77,6 +86,13 @@ std::string emitGlobalISel(const NormalizedTargetModel &Model, std::string_view
       for (const NormalizedSelectionOperandDefConstantPredicate::Range &Range : Predicate.Ranges)
         Output << "    {" << Range.Minimum << ", " << Range.Maximum << "},\n";
     }
+    for (const NormalizedSelectionDefUsePredicate &DefUse : Rule.DefUsePredicates) {
+      for (const NormalizedSelectionOperandDefConstantPredicate &Predicate :
+           DefUse.OperandDefConstants) {
+        for (const NormalizedSelectionOperandDefConstantPredicate::Range &Range : Predicate.Ranges)
+          Output << "    {" << Range.Minimum << ", " << Range.Maximum << "},\n";
+      }
+    }
   }
   Output << "}};\n\n";
 
@@ -103,6 +119,20 @@ std::string emitGlobalISel(const NormalizedTargetModel &Model, std::string_view
                       std::to_string(Predicate.Ranges.size()));
       FirstConstantRange += Predicate.Ranges.size();
     }
+    for (const NormalizedSelectionDefUsePredicate &DefUse : Rule.DefUsePredicates) {
+      emitInstruction(Output, "CheckOperandDefOpcode", DefUse.RootOperandIndex, DefUse.Opcode);
+      for (const NormalizedSelectionOperandDefConstantPredicate &Predicate :
+           DefUse.OperandDefConstants) {
+        emitInstruction(Output, "CheckOperandDefOperandConstantRanges", DefUse.RootOperandIndex,
+                        static_cast<std::uint32_t>(FirstConstantRange),
+                        std::to_string(Predicate.OperandIndex),
+                        static_cast<std::int64_t>(Predicate.Ranges.size()));
+        FirstConstantRange += Predicate.Ranges.size();
+      }
+      for (const NormalizedSelectionDefUseCapture &Capture : DefUse.Captures)
+        emitInstruction(Output, "CaptureDefOperand", DefUse.RootOperandIndex, Capture.ID,
+                        std::to_string(Capture.OperandIndex));
+    }
     for (const NormalizedSelectionTemporary &Temporary : Rule.Temporaries)
       emitInstruction(Output, "CreateTemporary", Temporary.TypeFromOperand, Temporary.ID);
     for (const NormalizedSelectionAction &Action : Rule.Actions) {
@@ -113,6 +143,10 @@ std::string emitGlobalISel(const NormalizedTargetModel &Model, std::string_view
         case NormalizedSelectionOperandRenderer::Kind::Operand:
           emitInstruction(Output, "RenderOperand", Renderer.Value);
           break;
+        case NormalizedSelectionOperandRenderer::Kind::ComplexOperand:
+          emitInstruction(Output, "RenderComplexOperand", Renderer.Value, 0, "0", 0,
+                          Renderer.ComplexRenderer);
+          break;
         case NormalizedSelectionOperandRenderer::Kind::PhysicalRegister:
           emitInstruction(Output, "RenderPhysicalRegister", 0, Renderer.Value,
                           std::to_string(Renderer.Flags));
@@ -121,6 +155,9 @@ std::string emitGlobalISel(const NormalizedTargetModel &Model, std::string_view
           emitInstruction(Output, "RenderTemporary", 0, Renderer.Value,
                           std::to_string(Renderer.Flags));
           break;
+        case NormalizedSelectionOperandRenderer::Kind::CapturedOperand:
+          emitInstruction(Output, "RenderCapturedOperand", 0, Renderer.Value);
+          break;
         case NormalizedSelectionOperandRenderer::Kind::CapturedImmediate:
           emitInstruction(Output, "RenderCapturedImmediate", Renderer.Value,
                           static_cast<std::uint32_t>(Renderer.Transform));