[wasm][compiler] Keep call counts in a side table

This patch adds a side table to the MachineGraph that stores the
previously observed call count for the Call nodes used for Wasm
direct calls. This replaces a more convoluted system that accessed
processed feedback during compilation, keyed on source position.

Bug: v8:12166
Change-Id: I06109918030b8f256c5f170da5853394c1a69cc2
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3644803Reviewed-by: Tobias Tebbi <tebbi@chromium.org>
Commit-Queue: Jakob Kummerow <jkummerow@chromium.org>
Cr-Commit-Position: refs/heads/main@{#80558}

src/compiler/machine-graph.h

@@ -11,6 +11,7 @@
 #include "src/compiler/common-operator.h"
 #include "src/compiler/graph.h"
 #include "src/compiler/machine-operator.h"
+#include "src/compiler/node-aux-data.h"
 #include "src/runtime/runtime.h"
 
 namespace v8 {
@@ -24,7 +25,11 @@ class V8_EXPORT_PRIVATE MachineGraph : public NON_EXPORTED_BASE(ZoneObject) {
  public:
   MachineGraph(Graph* graph, CommonOperatorBuilder* common,
                MachineOperatorBuilder* machine)
-      : graph_(graph), common_(common), machine_(machine), cache_(zone()) {}
+      : graph_(graph),
+        common_(common),
+        machine_(machine),
+        cache_(zone()),
+        call_counts_(zone()) {}
   MachineGraph(const MachineGraph&) = delete;
   MachineGraph& operator=(const MachineGraph&) = delete;
 
@@ -75,6 +80,16 @@ class V8_EXPORT_PRIVATE MachineGraph : public NON_EXPORTED_BASE(ZoneObject) {
     return Dead_ ? Dead_ : Dead_ = graph_->NewNode(common_->Dead());
   }
 
+  // Store and retrieve call count information.
+  void StoreCallCount(NodeId call_id, int count) {
+    call_counts_.Put(call_id, count);
+  }
+  int GetCallCount(NodeId call_id) { return call_counts_.Get(call_id); }
+  // Use this to keep the number of map rehashings to a minimum.
+  void ReserveCallCounts(size_t num_call_instructions) {
+    call_counts_.Reserve(num_call_instructions);
+  }
+
   CommonOperatorBuilder* common() const { return common_; }
   MachineOperatorBuilder* machine() const { return machine_; }
   Graph* graph() const { return graph_; }
@@ -85,6 +100,7 @@ class V8_EXPORT_PRIVATE MachineGraph : public NON_EXPORTED_BASE(ZoneObject) {
   CommonOperatorBuilder* common_;
   MachineOperatorBuilder* machine_;
   CommonNodeCache cache_;
+  NodeAuxDataMap<int, -1> call_counts_;
   Node* Dead_ = nullptr;
 };
 

src/compiler/node-aux-data.h

@@ -105,6 +105,28 @@ typename NodeAuxData<T, def>::const_iterator NodeAuxData<T, def>::end() const {
                                                       aux_data_.size());
 }
 
+template <class T, T kNonExistent>
+class NodeAuxDataMap {
+ public:
+  explicit NodeAuxDataMap(Zone* zone) : map_(zone) {}
+
+  void Put(NodeId key, T value) { map_[key] = value; }
+
+  T Get(NodeId key) const {
+    auto entry = map_.find(key);
+    if (entry == map_.end()) return kNonExistent;
+    return entry->second;
+  }
+
+  void Reserve(size_t count) {
+    size_t new_capacity = map_.size() + count;
+    map_.reserve(new_capacity);
+  }
+
+ private:
+  ZoneUnorderedMap<NodeId, T> map_;
+};
+
 }  // namespace compiler
 }  // namespace internal
 }  // namespace v8

src/compiler/wasm-compiler.cc

@@ -8402,6 +8402,14 @@ const wasm::FunctionSig* WasmGraphBuilder::Int64LoweredSig(
              : sig;
 }
 
+void WasmGraphBuilder::StoreCallCount(Node* call, int count) {
+  mcgraph()->StoreCallCount(call->id(), count);
+}
+
+void WasmGraphBuilder::ReserveCallCounts(size_t num_call_instructions) {
+  mcgraph()->ReserveCallCounts(num_call_instructions);
+}
+
 CallDescriptor* GetI32WasmCallDescriptor(
     Zone* zone, const CallDescriptor* call_descriptor) {
   return ReplaceTypeInCallDescriptorWith(zone, call_descriptor, 2,

src/compiler/wasm-compiler.h

@@ -560,6 +560,9 @@ class WasmGraphBuilder {
   static const wasm::FunctionSig* Int64LoweredSig(Zone* zone,
                                                   const wasm::FunctionSig* sig);
 
+  void StoreCallCount(Node* call, int count);
+  void ReserveCallCounts(size_t num_call_instructions);
+
  protected:
   V8_EXPORT_PRIVATE WasmGraphBuilder(wasm::CompilationEnv* env, Zone* zone,
                                      MachineGraph* mcgraph,

src/compiler/wasm-inlining.cc

@@ -36,39 +36,9 @@ void WasmInliner::Trace(Node* call, int inlinee, const char* decision) {
         call->id(), inlinee, decision);
 }
 
-uint32_t WasmInliner::FindOriginatingFunction(Node* call) {
-  DCHECK_EQ(inlined_functions_.size(), first_node_id_.size());
-  NodeId id = call->id();
-  if (inlined_functions_.size() == 0 || id < first_node_id_[0]) {
-    return function_index_;
-  }
-  for (size_t i = 1; i < first_node_id_.size(); i++) {
-    if (id < first_node_id_[i]) return inlined_functions_[i - 1];
-  }
-  DCHECK_GE(id, first_node_id_.back());
-  return inlined_functions_.back();
-}
-
 int WasmInliner::GetCallCount(Node* call) {
   if (!FLAG_wasm_speculative_inlining) return 0;
-  base::MutexGuard guard(&module()->type_feedback.mutex);
-  wasm::WasmCodePosition position =
-      source_positions_->GetSourcePosition(call).ScriptOffset();
-  uint32_t func = FindOriginatingFunction(call);
-  auto maybe_feedback =
-      module()->type_feedback.feedback_for_function.find(func);
-  if (maybe_feedback == module()->type_feedback.feedback_for_function.end()) {
-    return 0;
-  }
-  wasm::FunctionTypeFeedback feedback = maybe_feedback->second;
-  // It's possible that we haven't processed the feedback yet. Currently,
-  // this can happen for targets of call_direct that haven't gotten hot yet,
-  // and for functions where Liftoff bailed out.
-  if (feedback.feedback_vector.size() == 0) return 0;
-  auto index_in_vector = feedback.positions.find(position);
-  if (index_in_vector == feedback.positions.end()) return 0;
-  return feedback.feedback_vector[index_in_vector->second]
-      .absolute_call_frequency;
+  return mcgraph()->GetCallCount(call->id());
 }
 
 // TODO(12166): Save inlined frames for trap/--trace-wasm purposes. Consider
@@ -232,9 +202,6 @@ void WasmInliner::Finalize() {
     size_t additional_nodes = graph()->NodeCount() - subgraph_min_node_id;
     Trace(candidate, "inlining!");
     current_graph_size_ += additional_nodes;
-    inlined_functions_.push_back(candidate.inlinee_index);
-    static_assert(std::is_same_v<NodeId, uint32_t>);
-    first_node_id_.push_back(static_cast<uint32_t>(subgraph_min_node_id));
 
     if (call->opcode() == IrOpcode::kCall) {
       InlineCall(call, inlinee_start, inlinee_end, inlinee->sig,

src/compiler/wasm-inlining.h

@@ -79,8 +79,6 @@ class WasmInliner final : public AdvancedReducer {
     }
   };
 
-  uint32_t FindOriginatingFunction(Node* call);
-
   Zone* zone() const { return mcgraph_->zone(); }
   CommonOperatorBuilder* common() const { return mcgraph_->common(); }
   Graph* graph() const { return mcgraph_->graph(); }
@@ -113,11 +111,6 @@ class WasmInliner final : public AdvancedReducer {
                       LexicographicOrdering>
       inlining_candidates_;
   std::unordered_set<Node*> seen_;
-  std::vector<uint32_t> inlined_functions_;
-  // Stores the graph size before an inlining was performed, to make it
-  // possible to map back from nodes to the function they came from.
-  // Guaranteed to have the same length as {inlined_functions_}.
-  std::vector<uint32_t> first_node_id_;
 };
 
 }  // namespace compiler

src/wasm/baseline/liftoff-compiler.cc

@@ -6118,8 +6118,7 @@ class LiftoffCompiler {
     if (FLAG_wasm_speculative_inlining) {
       base::MutexGuard mutex_guard(&decoder->module_->type_feedback.mutex);
       decoder->module_->type_feedback.feedback_for_function[func_index_]
-          .positions[decoder->position()] =
-          static_cast<int>(num_call_instructions_);
+          .call_targets.push_back(imm.index);
       num_call_instructions_++;
     }
 
@@ -6359,8 +6358,7 @@ class LiftoffCompiler {
       {
         base::MutexGuard mutex_guard(&decoder->module_->type_feedback.mutex);
         decoder->module_->type_feedback.feedback_for_function[func_index_]
-            .positions[decoder->position()] =
-            static_cast<int>(num_call_instructions_);
+            .call_targets.push_back(FunctionTypeFeedback::kNonDirectCall);
       }
       num_call_instructions_++;
       __ LoadConstant(index, WasmValue::ForUintPtr(vector_slot));

src/wasm/graph-builder-interface.cc

@@ -128,6 +128,7 @@ class WasmGraphBuildingInterface {
       auto feedback = feedbacks.feedback_for_function.find(func_index_);
       if (feedback != feedbacks.feedback_for_function.end()) {
         type_feedback_ = feedback->second.feedback_vector;
+        builder_->ReserveCallCounts(type_feedback_.size());
         // We need to keep the feedback in the module to inline later. However,
         // this means we are stuck with it forever.
         // TODO(jkummerow): Reconsider our options here.
@@ -640,21 +641,25 @@ class WasmGraphBuildingInterface {
   void CallDirect(FullDecoder* decoder,
                   const CallFunctionImmediate<validate>& imm,
                   const Value args[], Value returns[]) {
+    int maybe_call_count = -1;
     if (FLAG_wasm_speculative_inlining && type_feedback_.size() > 0) {
-      DCHECK_LT(feedback_instruction_index_, type_feedback_.size());
-      feedback_instruction_index_++;
+      const CallSiteFeedback& feedback = next_call_feedback();
+      maybe_call_count = feedback.absolute_call_frequency;
     }
-    DoCall(decoder, CallInfo::CallDirect(imm.index), imm.sig, args, returns);
+    DoCall(decoder, CallInfo::CallDirect(imm.index, maybe_call_count), imm.sig,
+           args, returns);
   }
 
   void ReturnCall(FullDecoder* decoder,
                   const CallFunctionImmediate<validate>& imm,
                   const Value args[]) {
+    int maybe_call_count = -1;
     if (FLAG_wasm_speculative_inlining && type_feedback_.size() > 0) {
-      DCHECK_LT(feedback_instruction_index_, type_feedback_.size());
-      feedback_instruction_index_++;
+      const CallSiteFeedback& feedback = next_call_feedback();
+      maybe_call_count = feedback.absolute_call_frequency;
     }
-    DoReturnCall(decoder, CallInfo::CallDirect(imm.index), imm.sig, args);
+    DoReturnCall(decoder, CallInfo::CallDirect(imm.index, maybe_call_count),
+                 imm.sig, args);
   }
 
   void CallIndirect(FullDecoder* decoder, const Value& index,
@@ -679,11 +684,11 @@ class WasmGraphBuildingInterface {
                const FunctionSig* sig, uint32_t sig_index, const Value args[],
                Value returns[]) {
     int maybe_feedback = -1;
+    int maybe_call_count = -1;
     if (FLAG_wasm_speculative_inlining && type_feedback_.size() > 0) {
-      DCHECK_LT(feedback_instruction_index_, type_feedback_.size());
-      maybe_feedback =
-          type_feedback_[feedback_instruction_index_].function_index;
-      feedback_instruction_index_++;
+      const CallSiteFeedback& feedback = next_call_feedback();
+      maybe_feedback = feedback.function_index;
+      maybe_call_count = feedback.absolute_call_frequency;
     }
     if (maybe_feedback == -1) {
       DoCall(decoder, CallInfo::CallRef(func_ref, NullCheckFor(func_ref.type)),
@@ -713,7 +718,8 @@ class WasmGraphBuildingInterface {
     ssa_env_->control = success_control;
     Value* returns_direct =
         decoder->zone()->NewArray<Value>(sig->return_count());
-    DoCall(decoder, CallInfo::CallDirect(expected_function_index),
+    DoCall(decoder,
+           CallInfo::CallDirect(expected_function_index, maybe_call_count),
            decoder->module_->signature(sig_index), args, returns_direct);
     TFNode* control_direct = control();
     TFNode* effect_direct = effect();
@@ -749,11 +755,11 @@ class WasmGraphBuildingInterface {
                      const FunctionSig* sig, uint32_t sig_index,
                      const Value args[]) {
     int maybe_feedback = -1;
+    int maybe_call_count = -1;
     if (FLAG_wasm_speculative_inlining && type_feedback_.size() > 0) {
-      DCHECK_LT(feedback_instruction_index_, type_feedback_.size());
-      maybe_feedback =
-          type_feedback_[feedback_instruction_index_].function_index;
-      feedback_instruction_index_++;
+      const CallSiteFeedback& feedback = next_call_feedback();
+      maybe_feedback = feedback.function_index;
+      maybe_call_count = feedback.absolute_call_frequency;
     }
     if (maybe_feedback == -1) {
       DoReturnCall(decoder,
@@ -782,8 +788,10 @@ class WasmGraphBuildingInterface {
 
     builder_->SetControl(success_control);
     ssa_env_->control = success_control;
-    DoReturnCall(decoder, CallInfo::CallDirect(expected_function_index), sig,
-                 args);
+    DoReturnCall(
+        decoder,
+        CallInfo::CallDirect(expected_function_index, maybe_call_count), sig,
+        args);
 
     builder_->SetEffectControl(initial_effect, failure_control);
     ssa_env_->effect = initial_effect;
@@ -1758,8 +1766,9 @@ class WasmGraphBuildingInterface {
    public:
     enum CallMode { kCallDirect, kCallIndirect, kCallRef };
 
-    static CallInfo CallDirect(uint32_t callee_index) {
-      return {kCallDirect, callee_index, nullptr, 0,
+    static CallInfo CallDirect(uint32_t callee_index, int call_count) {
+      return {kCallDirect, callee_index, nullptr,
+              static_cast<uint32_t>(call_count),
               CheckForNull::kWithoutNullCheck};
     }
 
@@ -1786,6 +1795,11 @@ class WasmGraphBuildingInterface {
       return callee_or_sig_index_;
     }
 
+    int call_count() {
+      DCHECK_EQ(call_mode_, kCallDirect);
+      return static_cast<int>(table_index_or_call_count_);
+    }
+
     CheckForNull null_check() {
       DCHECK_EQ(call_mode_, kCallRef);
       return null_check_;
@@ -1798,22 +1812,22 @@ class WasmGraphBuildingInterface {
 
     uint32_t table_index() {
       DCHECK_EQ(call_mode_, kCallIndirect);
-      return table_index_;
+      return table_index_or_call_count_;
     }
 
    private:
     CallInfo(CallMode call_mode, uint32_t callee_or_sig_index,
-             const Value* index_or_callee_value, uint32_t table_index,
-             CheckForNull null_check)
+             const Value* index_or_callee_value,
+             uint32_t table_index_or_call_count, CheckForNull null_check)
         : call_mode_(call_mode),
           callee_or_sig_index_(callee_or_sig_index),
           index_or_callee_value_(index_or_callee_value),
-          table_index_(table_index),
+          table_index_or_call_count_(table_index_or_call_count),
           null_check_(null_check) {}
     CallMode call_mode_;
     uint32_t callee_or_sig_index_;
     const Value* index_or_callee_value_;
-    uint32_t table_index_;
+    uint32_t table_index_or_call_count_;
     CheckForNull null_check_;
   };
 
@@ -1850,13 +1864,14 @@ class WasmGraphBuildingInterface {
                          real_sig, base::VectorOf(arg_nodes),
                          base::VectorOf(return_nodes), decoder->position()));
         break;
-      case CallInfo::kCallDirect:
-        CheckForException(
-            decoder, builder_->CallDirect(call_info.callee_index(), real_sig,
-                                          base::VectorOf(arg_nodes),
-                                          base::VectorOf(return_nodes),
-                                          decoder->position()));
+      case CallInfo::kCallDirect: {
+        TFNode* call = builder_->CallDirect(
+            call_info.callee_index(), real_sig, base::VectorOf(arg_nodes),
+            base::VectorOf(return_nodes), decoder->position());
+        builder_->StoreCallCount(call, call_info.call_count());
+        CheckForException(decoder, call);
         break;
+      }
       case CallInfo::kCallRef:
         CheckForException(
             decoder,
@@ -1916,10 +1931,13 @@ class WasmGraphBuildingInterface {
             call_info.table_index(), call_info.sig_index(), real_sig,
             base::VectorOf(arg_nodes), decoder->position());
         break;
-      case CallInfo::kCallDirect:
-        builder_->ReturnCall(call_info.callee_index(), real_sig,
-                             base::VectorOf(arg_nodes), decoder->position());
+      case CallInfo::kCallDirect: {
+        TFNode* call = builder_->ReturnCall(call_info.callee_index(), real_sig,
+                                            base::VectorOf(arg_nodes),
+                                            decoder->position());
+        builder_->StoreCallCount(call, call_info.call_count());
         break;
+      }
       case CallInfo::kCallRef:
         builder_->ReturnCallRef(real_sig, base::VectorOf(arg_nodes),
                                 call_info.null_check(), decoder->position());
@@ -1927,6 +1945,11 @@ class WasmGraphBuildingInterface {
     }
   }
 
+  const CallSiteFeedback& next_call_feedback() {
+    DCHECK_LT(feedback_instruction_index_, type_feedback_.size());
+    return type_feedback_[feedback_instruction_index_++];
+  }
+
   void BuildLoopExits(FullDecoder* decoder, Control* loop) {
     builder_->LoopExit(loop->loop_node);
     ssa_env_->control = control();

src/wasm/module-compiler.cc

@@ -1282,9 +1282,7 @@ class TransitiveTypeFeedbackProcessor {
   }
 
  private:
-  static constexpr uint32_t kNonDirectCall = 0xFFFFFFFF;
   void Process(int func_index);
-  std::vector<uint32_t> GetCallDirectTargets(int func_index);
 
   void EnqueueCallees(std::vector<CallSiteFeedback> feedback) {
     for (size_t i = 0; i < feedback.size(); i++) {
@@ -1308,44 +1306,6 @@ class TransitiveTypeFeedbackProcessor {
   std::unordered_set<int> queue_;
 };
 
-// For every recorded position of a call instruction in the given function,
-// extracts the bytes that constitute the target function index if the call
-// is a call_direct. We don't know the type of call, so this will produce
-// bogus bytes for other types of calls.
-std::vector<uint32_t> TransitiveTypeFeedbackProcessor::GetCallDirectTargets(
-    int func_index) {
-  WasmModuleObject module_object = instance_->module_object();
-  const NativeModule* native_module = module_object.native_module();
-  base::Vector<const uint8_t> wire_bytes = native_module->wire_bytes();
-  const WasmModule* module = native_module->module();
-  const WasmFunction* func = &module->functions[func_index];
-  const byte* func_start = wire_bytes.begin() + func->code.offset();
-  Decoder decoder(wire_bytes);
-
-  std::map<WasmCodePosition, int> positions =
-      module->type_feedback.feedback_for_function[func_index].positions;
-  size_t num_calls = positions.size();
-  std::vector<uint32_t> result(num_calls);
-
-  for (auto entry : positions) {
-    int position = entry.first;
-    int call_index = entry.second;
-    const byte* pc = func_start + position;
-    uint8_t call = decoder.read_u8<Decoder::kNoValidation>(pc);
-    if (call != kExprCallFunction) {
-      result[call_index] = kNonDirectCall;
-      continue;
-    }
-    static constexpr int kCallDirectInstructionLength = 1;
-    pc += kCallDirectInstructionLength;
-    uint32_t length_dummy;
-    uint32_t immediate =
-        decoder.read_u32v<Decoder::kNoValidation>(pc, &length_dummy);
-    result[call_index] = immediate;
-  }
-  return result;
-}
-
 void TransitiveTypeFeedbackProcessor::Process(int func_index) {
   int which_vector = declared_function_index(instance_->module(), func_index);
   Object maybe_feedback = instance_->feedback_vectors().get(which_vector);
@@ -1354,7 +1314,11 @@ void TransitiveTypeFeedbackProcessor::Process(int func_index) {
   std::vector<CallSiteFeedback> result(feedback.length() / 2);
   int imported_functions =
       static_cast<int>(instance_->module()->num_imported_functions);
-  std::vector<uint32_t> call_direct_targets = GetCallDirectTargets(func_index);
+  WasmModuleObject module_object = instance_->module_object();
+  const NativeModule* native_module = module_object.native_module();
+  const WasmModule* module = native_module->module();
+  const std::vector<uint32_t>& call_direct_targets(
+      module->type_feedback.feedback_for_function[func_index].call_targets);
   for (int i = 0; i < feedback.length(); i += 2) {
     Object value = feedback.get(i);
     if (value.IsWasmInternalFunction() &&
@@ -1429,7 +1393,7 @@ void TransitiveTypeFeedbackProcessor::Process(int func_index) {
     } else if (value.IsSmi()) {
       // Uninitialized, or a direct call collecting call count.
       uint32_t target = call_direct_targets[i / 2];
-      if (target != kNonDirectCall) {
+      if (target != FunctionTypeFeedback::kNonDirectCall) {
         int count = Smi::cast(value).value();
         if (FLAG_trace_wasm_speculative_inlining) {
           PrintF("[Function #%d call_direct #%d: frequency %d]\n", func_index,

src/wasm/wasm-module.h

@@ -405,8 +405,10 @@ struct CallSiteFeedback {
 };
 struct FunctionTypeFeedback {
   std::vector<CallSiteFeedback> feedback_vector;
-  std::map<WasmCodePosition, int> positions;
+  std::vector<uint32_t> call_targets;
   int tierup_priority = 0;
+
+  static constexpr uint32_t kNonDirectCall = 0xFFFFFFFF;
 };
 struct TypeFeedbackStorage {
   std::map<uint32_t, FunctionTypeFeedback> feedback_for_function;