Have a cheaper Phi::create() that takes an ordered list

Summary:
The current way to create a Phi requires building a hashmap of block -> value.
Often times, we already know the order of the (block, value) pairs (or we can
trivially sort and achieve it).  Add a new ctor and setArgs() overload where we
can pass in a `std::span<std::tuple<BasicBlock*, Register*>>` instead.

Reviewed By: yoney

Differential Revision: D114414426

fbshipit-source-id: 4bfa212417539c3e8a2fe93a7cd0d08adc2e8e76

Author: alexmalyshev

cinderx/Jit/hir/hir.cpp

@@ -6,6 +6,7 @@
 #include "cinderx/Jit/threaded_compile.h"
 
 #include <algorithm>
+#include <ranges>
 
 namespace cinderx::jit::hir {
 
@@ -161,7 +162,12 @@ std::string_view CallCFunc::funcName() const {
 }
 
 void Phi::setArgs(const std::unordered_map<BasicBlock*, Register*>& args) {
-  JIT_DCHECK(numOperands() == args.size(), "arg mismatch");
+  // Can't resize a phi, the operands have already been allocated inline.
+  JIT_THROW_IF(
+      numOperands() != args.size(),
+      "Trying to update phi with {} arguments to one with {} arguments",
+      numOperands(),
+      args.size());
 
   basic_blocks_.clear();
   basic_blocks_.reserve(args.size());
@@ -184,6 +190,31 @@ void Phi::setArgs(const std::unordered_map<BasicBlock*, Register*>& args) {
   }
 }
 
+void Phi::setArgs(std::span<std::tuple<BasicBlock*, Register*>> args) {
+  // Can't resize a phi, the operands have already been allocated inline.
+  JIT_THROW_IF(
+      numOperands() != args.size(),
+      "Trying to update phi with {} arguments to one with {} arguments",
+      numOperands(),
+      args.size());
+
+  JIT_DCHECK(
+      std::ranges::is_sorted(
+          args,
+          [](const auto& a, const auto& b) {
+            return std::get<0>(a)->id < std::get<0>(b)->id;
+          }),
+      "Phi must be created with a (block,value) list sorted by block ID");
+
+  basic_blocks_.clear();
+  basic_blocks_.reserve(args.size());
+
+  for (size_t i = 0; i < args.size(); ++i) {
+    basic_blocks_.push_back(std::get<0>(args[i]));
+    operandAt(i) = std::get<1>(args[i]);
+  }
+}
+
 std::size_t Phi::blockIndex(const BasicBlock* block) const {
   auto it = std::lower_bound(
       basic_blocks_.begin(), basic_blocks_.end(), block, [](auto b1, auto b2) {
@@ -1180,32 +1211,6 @@ void BasicBlock::fixupPhis(BasicBlock* old_pred, BasicBlock* new_pred) {
   forEachPhi([&](Phi& phi) { phi.replacePredecessor(old_pred, new_pred); });
 }
 
-void BasicBlock::addPhiPredecessor(BasicBlock* old_pred, BasicBlock* new_pred) {
-  std::vector<Phi*> replacements;
-  forEachPhi([&](Phi& phi) {
-    for (auto block : phi.basicBlocks()) {
-      if (block == old_pred) {
-        replacements.push_back(&phi);
-        break;
-      }
-    }
-  });
-
-  for (auto phi : replacements) {
-    std::unordered_map<BasicBlock*, Register*> args;
-    for (size_t i = 0, n = phi->numOperands(); i < n; ++i) {
-      auto block = phi->basicBlocks()[i];
-      if (block == old_pred) {
-        args[new_pred] = phi->getOperand(i);
-      }
-      args[block] = phi->getOperand(i);
-    }
-
-    phi->replaceWith(*Phi::create(phi->output(), args));
-    delete phi;
-  }
-}
-
 void BasicBlock::removePhiPredecessor(BasicBlock* old_pred) {
   for (auto it = instrs_.begin(); it != instrs_.end();) {
     auto& instr = *it;
@@ -1214,14 +1219,15 @@ void BasicBlock::removePhiPredecessor(BasicBlock* old_pred) {
       break;
     }
 
-    Phi* phi = static_cast<Phi*>(&instr);
-    std::unordered_map<BasicBlock*, Register*> args;
-    for (size_t i = 0, n = phi->numOperands(); i < n; ++i) {
-      auto block = phi->basicBlocks()[i];
-      if (block == old_pred) {
-        continue;
+    auto phi = static_cast<Phi*>(&instr);
+    auto num_operands = phi->numOperands();
+    std::vector<std::tuple<BasicBlock*, Register*>> args;
+    args.reserve(num_operands);
+    for (size_t i = 0; i < num_operands; ++i) {
+      BasicBlock* block = phi->basicBlocks()[i];
+      if (block != old_pred) {
+        args.emplace_back(block, phi->getOperand(i));
       }
-      args[block] = phi->getOperand(i);
     }
     phi->replaceWith(*Phi::create(phi->output(), args));
     delete phi;

cinderx/Jit/hir/hir.h

@@ -18,6 +18,7 @@
 #include <span>
 #include <string>
 #include <string_view>
+#include <tuple>
 #include <type_traits>
 #include <unordered_map>
 #include <unordered_set>
@@ -1144,12 +1145,11 @@ class INSTR_CLASS(Phi, (TTop), HasOutput, Operands<>) {
  public:
   explicit Phi(Register* dst) : InstrT(dst) {}
 
-  static Phi* create(
-      Register* dst,
-      const std::unordered_map<BasicBlock*, Register*>& args) {
+  template <typename PhiArgs>
+  static Phi* create(Register* dst, PhiArgs&& args) {
     void* ptr = Instr::allocate(sizeof(Phi), args.size());
     auto phi = new (ptr) Phi(dst);
-    phi->setArgs(args);
+    phi->setArgs(std::forward<PhiArgs>(args));
     return phi;
   }
 
@@ -1172,15 +1172,22 @@ class INSTR_CLASS(Phi, (TTop), HasOutput, Operands<>) {
   // Return the index of the given predecessor in basic_blocks.
   std::size_t blockIndex(const BasicBlock* block) const;
 
+  // Replace a predecessor block with a new block in the phi, keeping the same
+  // associated value.
   void replacePredecessor(BasicBlock* old_pred, BasicBlock* new_pred);
 
   const std::vector<BasicBlock*> basicBlocks() const {
     return basic_blocks_;
   }
 
+ private:
+  // Set a phi's arguments through a block -> value map.
   void setArgs(const std::unordered_map<BasicBlock*, Register*>& args);
 
- private:
+  // Set a phi's arguments through a (block, value) list.  The list must already
+  // be sorted by the block's ID.
+  void setArgs(std::span<std::tuple<BasicBlock*, Register*>> args);
+
   // List of incoming blocks, sorted by ascending block ID.
   std::vector<BasicBlock*> basic_blocks_;
 };
@@ -4116,9 +4123,8 @@ class BasicBlock : public IntrusiveListNode<BasicBlock> {
 
   // Replace any references to old_pred in this block's Phis with new_pred.
   void fixupPhis(BasicBlock* old_pred, BasicBlock* new_pred);
-  // Adds a new predecessor to the phi that follows from the old predecessor
-  void addPhiPredecessor(BasicBlock* old_pred, BasicBlock* new_pred);
-  // Removes any references to old_pred in this block's Phis
+
+  // Removes any references to old_pred in this block's Phis.
   void removePhiPredecessor(BasicBlock* old_pred);
 
   // Call f with each Phi instruction at the beginning of this block.

cinderx/Jit/hir/parser.cpp

@@ -805,6 +805,10 @@ HIRParser::parseInstr(std::string_view opcode, Register* dst, int bb_index) {
       for (auto& input : info.inputs) {
         input.value = parseRegister();
       }
+      std::sort(
+          info.inputs.begin(),
+          info.inputs.end(),
+          [](const PhiInput& a, const PhiInput& b) { return a.bb < b.bb; });
       phis_[bb_index].emplace_back(std::move(info));
       break;
     }
@@ -1316,9 +1320,10 @@ void HIRParser::realizePhis() {
     auto& front = block->front();
 
     for (auto& phi : pair.second) {
-      std::unordered_map<BasicBlock*, Register*> inputs;
+      std::vector<std::tuple<BasicBlock*, Register*>> inputs;
+      inputs.reserve(phi.inputs.size());
       for (auto& info : phi.inputs) {
-        inputs.emplace(index_to_bb_[info.bb], info.value);
+        inputs.emplace_back(index_to_bb_[info.bb], info.value);
       }
       (Phi::create(phi.dst, inputs))->insertBefore(front);
     }