Build token mask caches in worker-local maps

cpp/grammar_compiler.cc

@@ -6,6 +6,7 @@
 #include <xgrammar/compiler.h>
 
 #include <algorithm>
+#include <atomic>
 #include <bitset>
 #include <cctype>
 #include <cstddef>
@@ -22,6 +23,7 @@
 #include "fsm.h"
 #include "grammar_functor.h"
 #include "grammar_impl.h"
+#include "structural_tag.h"
 #include "support/dynamic_bitset.h"
 #include "support/int_set.h"
 #include "support/logging.h"
@@ -1017,6 +1019,7 @@ class GrammarCompilerSub {
   )
       : tokenizer_info_(tokenizer_info),
         max_threads_(max_threads),
+        thread_pool_(max_threads > 1 ? std::make_unique<ThreadPool>(max_threads) : nullptr),
         rule_level_cache_(rule_level_cache) {}
 
   CompiledGrammar CompileBuiltinJSONGrammar();
@@ -1056,14 +1059,17 @@ class GrammarCompilerSub {
   const TokenizerInfo tokenizer_info_;
   /*! \brief The maximum number of threads to use. */
   const int max_threads_;
+  /*! \brief Workers reused across grammar conversion, optimization, and mask generation. */
+  const std::unique_ptr<ThreadPool> thread_pool_;
 
   /*! \brief The manager of the rule level cache.*/
   std::optional<RuleLevelCache> rule_level_cache_;
 };
 
 CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_unoptimized) {
   auto compiled_grammar_impl = std::make_shared<CompiledGrammar::Impl>();
-  compiled_grammar_impl->grammar = GrammarOptimizer::Apply(grammar_unoptimized);
+  compiled_grammar_impl->grammar =
+      GrammarOptimizer::Apply(grammar_unoptimized, max_threads_, thread_pool_.get());
   compiled_grammar_impl->tokenizer_info = tokenizer_info_;
   if (tokenizer_info_.GetVocabSize() == 0) {
     return CompiledGrammar(compiled_grammar_impl);
@@ -1081,42 +1087,21 @@ CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_un
   // 2. All byte strings (with element_in_string=0, 1, 2, ...)
   // since other positions will be expanded to the above positions
 
-  // TODO(Charlie): Figure out how to support ThreadPool and std::mutex in WebAssembly.
-  // Only declare ThreadPool and mutex if max_threads > 1, so when max_threads = 1, we do
-  // not need ThreadPool or std::mutex, which throws error in runtime in WebAssembly.
-  std::optional<ThreadPool> thread_pool;
-  std::optional<std::mutex> adaptive_token_mask_cache_mutex;
-  if (max_threads_ > 1) {
-    thread_pool.emplace(max_threads_);
-    adaptive_token_mask_cache_mutex.emplace();
-  }
+  struct AdaptiveTokenMaskTask {
+    ParserState state;
+    bool is_root_rule;
+  };
+  std::vector<AdaptiveTokenMaskTask> adaptive_token_mask_tasks;
 
-  auto add_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) {
+  auto compute_adaptive_token_mask = [&](const AdaptiveTokenMaskTask& task) {
     auto grammar_matcher = GrammarMatcherForTokenMaskCache(
         compiled_grammar_impl->grammar,
-        state,
+        task.state,
         tag_dispatch_rule_id_to_second_slicing_bitset,
         tokenizer_info_,
         rule_level_cache_
     );
-    auto cur_adaptive_token_mask_cache = grammar_matcher.GetAdaptiveTokenMask(is_root_rule);
-    if (max_threads_ > 1) {
-      std::lock_guard<std::mutex> lock(adaptive_token_mask_cache_mutex.value());
-      compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache;
-    } else {
-      compiled_grammar_impl->adaptive_token_mask_cache[state] = cur_adaptive_token_mask_cache;
-    }
-  };
-
-  auto add_task_adaptive_token_mask = [&](const ParserState& state, bool is_root_rule) {
-    // Execute depending on whether we use thread_pool
-    if (max_threads_ > 1) {
-      thread_pool->Execute([add_adaptive_token_mask, state, is_root_rule]() {
-        add_adaptive_token_mask(state, is_root_rule);
-      });
-    } else {
-      add_adaptive_token_mask(state, is_root_rule);
-    }
+    return grammar_matcher.GetAdaptiveTokenMask(task.is_root_rule);
   };
 
   auto root_rule_id = compiled_grammar_impl->grammar->GetRootRuleId();
@@ -1135,12 +1120,43 @@ CompiledGrammar GrammarCompilerSub::MultiThreadCompileGrammar(Grammar grammar_un
       if (!rule_fsm->GetFsm().IsScanableState(i)) {
         continue;
       }
-      add_task_adaptive_token_mask(cur_stack_element, rule_id == root_rule_id);
+      adaptive_token_mask_tasks.push_back({cur_stack_element, rule_id == root_rule_id});
     }
   }
 
-  if (max_threads_ > 1) {
-    thread_pool->Join();
+  const int32_t num_tasks = static_cast<int32_t>(adaptive_token_mask_tasks.size());
+  if (max_threads_ == 1) {
+    for (const auto& task : adaptive_token_mask_tasks) {
+      compiled_grammar_impl->adaptive_token_mask_cache.emplace(
+          task.state, compute_adaptive_token_mask(task)
+      );
+    }
+  } else if (num_tasks > 0) {
+    const int32_t num_workers = std::min(max_threads_, num_tasks);
+    using AdaptiveTokenMaskMap = decltype(compiled_grammar_impl->adaptive_token_mask_cache);
+    std::vector<AdaptiveTokenMaskMap> worker_results(num_workers);
+    for (auto& worker_result : worker_results) {
+      worker_result.reserve((num_tasks + num_workers - 1) / num_workers);
+    }
+
+    std::atomic<int32_t> next_task_id{0};
+    for (int32_t worker_id = 0; worker_id < num_workers; ++worker_id) {
+      thread_pool_->Execute([&, worker_id]() {
+        auto& worker_result = worker_results[worker_id];
+        int32_t task_id;
+        while ((task_id = next_task_id.fetch_add(1, std::memory_order_relaxed)) < num_tasks) {
+          const auto& task = adaptive_token_mask_tasks[task_id];
+          worker_result.emplace(task.state, compute_adaptive_token_mask(task));
+        }
+      });
+    }
+    thread_pool_->Wait();
+
+    auto& result = compiled_grammar_impl->adaptive_token_mask_cache;
+    result.reserve(num_tasks);
+    for (auto& worker_result : worker_results) {
+      result.merge(worker_result);
+    }
   }
 
   return CompiledGrammar(compiled_grammar_impl);
@@ -1172,7 +1188,9 @@ CompiledGrammar GrammarCompilerSub::CompileJSONSchema(
 }
 
 CompiledGrammar GrammarCompilerSub::CompileStructuralTag(const std::string& structural_tag_json) {
-  auto result = Grammar::FromStructuralTag(structural_tag_json, tokenizer_info_);
+  auto result =
+      StructuralTagToGrammar(structural_tag_json, tokenizer_info_, max_threads_, thread_pool_.get())
+          .ToVariant();
   XGRAMMAR_CHECK(std::holds_alternative<Grammar>(result))
       << GetMessageFromVariantError(std::get<1>(result));
   return MultiThreadCompileGrammar(std::get<0>(result));

cpp/grammar_functor.cc

@@ -8,6 +8,7 @@
 #include <xgrammar/xgrammar.h>
 
 #include <array>
+#include <atomic>
 #include <bitset>
 #include <cstdint>
 #include <map>
@@ -28,6 +29,7 @@
 #include "support/container.h"
 #include "support/encoding.h"
 #include "support/logging.h"
+#include "support/thread_pool.h"
 #include "xgrammar/grammar.h"
 
 namespace xgrammar {
@@ -1306,14 +1308,38 @@ class GrammarFSMBuilderImpl {
   explicit GrammarFSMBuilderImpl(FSM& target_fsm, const std::string* rule_name = nullptr)
       : target_fsm_(target_fsm), rule_name_(rule_name) {}
 
-  static void Apply(Grammar* grammar) {
+  static void Apply(Grammar* grammar, int max_threads, ThreadPool* thread_pool) {
     FSM complete_fsm;
     std::vector<std::optional<FSMWithStartEndWithSize>> per_rule_fsms((*grammar)->NumRules());
     std::vector<int> state_mapping;
 
-    for (int i = 0; i < (*grammar)->NumRules(); ++i) {
-      auto rule_fsm = BuildRuleFSM(*grammar, i);
-      per_rule_fsms[i] = rule_fsm.AddToCompleteFSM(&complete_fsm, &state_mapping);
+    const int num_rules = (*grammar)->NumRules();
+    const int num_threads = std::max(1, std::min(max_threads, num_rules));
+    if (num_threads == 1) {
+      for (int i = 0; i < num_rules; ++i) {
+        auto rule_fsm = BuildRuleFSM(*grammar, i);
+        per_rule_fsms[i] = rule_fsm.AddToCompleteFSM(&complete_fsm, &state_mapping);
+      }
+    } else {
+      std::vector<std::optional<FSMWithStartEnd>> built_rule_fsms(num_rules);
+      std::atomic<int> next_rule_id{0};
+      std::optional<ThreadPool> owned_thread_pool;
+      if (thread_pool == nullptr) {
+        owned_thread_pool.emplace(num_threads);
+        thread_pool = &owned_thread_pool.value();
+      }
+      for (int thread_id = 0; thread_id < num_threads; ++thread_id) {
+        thread_pool->Execute([&]() {
+          int rule_id;
+          while ((rule_id = next_rule_id.fetch_add(1, std::memory_order_relaxed)) < num_rules) {
+            built_rule_fsms[rule_id] = BuildRuleFSM(*grammar, rule_id);
+          }
+        });
+      }
+      thread_pool->Wait();
+      for (int i = 0; i < num_rules; ++i) {
+        per_rule_fsms[i] = built_rule_fsms[i]->AddToCompleteFSM(&complete_fsm, &state_mapping);
+      }
     }
 
     for (int i = 0; i < (*grammar)->NumRules(); ++i) {
@@ -2981,7 +3007,7 @@ class LazyBodyFlattenerImpl : public GrammarMutator {
 
 class GrammarOptimizerImpl {
  public:
-  static Grammar Apply(const Grammar& grammar) {
+  static Grammar Apply(const Grammar& grammar, int max_threads, ThreadPool* thread_pool) {
     // ByteStringFuser and RuleInliner rewrite the grammar in place, so work on a private copy: the
     // input grammar may be shared (e.g. a cached grammar) and must not be mutated. Copy the impl
     // directly (contiguous vector copies) instead of going through GrammarBuilder, which would
@@ -2996,7 +3022,7 @@ class GrammarOptimizerImpl {
     result->allow_empty_rule_ids = AllowEmptyRuleAnalyzer::Apply(result);
     ValidateLazyRules(result);
     RepetitionNormalizer::Apply(&result);
-    GrammarFSMBuilder::Apply(&result);
+    GrammarFSMBuilder::Apply(&result, max_threads, thread_pool);
     result->optimized = true;
     return result;
   }
@@ -3906,7 +3932,17 @@ Grammar StructureNormalizer::Apply(const Grammar& grammar) {
 
 /*************************** Forward grammar optimizers to their impl ***************************/
 
-void GrammarFSMBuilder::Apply(Grammar* grammar) { GrammarFSMBuilderImpl::Apply(grammar); }
+void GrammarFSMBuilder::Apply(Grammar* grammar) {
+  GrammarFSMBuilderImpl::Apply(grammar, 1, nullptr);
+}
+
+void GrammarFSMBuilder::Apply(Grammar* grammar, int max_threads) {
+  GrammarFSMBuilderImpl::Apply(grammar, max_threads, nullptr);
+}
+
+void GrammarFSMBuilder::Apply(Grammar* grammar, int max_threads, ThreadPool* thread_pool) {
+  GrammarFSMBuilderImpl::Apply(grammar, max_threads, thread_pool);
+}
 
 void RepetitionNormalizer::Apply(Grammar* grammar) { RepetitionNormalizerImpl().Apply(grammar); }
 
@@ -3998,7 +4034,15 @@ Grammar RepetitionRangeExpander::Apply(const Grammar& grammar) {
 }
 
 Grammar GrammarOptimizer::Apply(const Grammar& grammar) {
-  return GrammarOptimizerImpl::Apply(grammar);
+  return GrammarOptimizerImpl::Apply(grammar, 1, nullptr);
+}
+
+Grammar GrammarOptimizer::Apply(const Grammar& grammar, int max_threads) {
+  return GrammarOptimizerImpl::Apply(grammar, max_threads, nullptr);
+}
+
+Grammar GrammarOptimizer::Apply(const Grammar& grammar, int max_threads, ThreadPool* thread_pool) {
+  return GrammarOptimizerImpl::Apply(grammar, max_threads, thread_pool);
 }
 
 void ByteStringFuser::Apply(Grammar* grammar) { ByteStringFuserImpl().Apply(grammar); }

cpp/grammar_functor.h

@@ -22,6 +22,8 @@
 
 namespace xgrammar {
 
+class ThreadPool;
+
 /*!
  * \brief Base class for visitors and mutators of the BNF grammar.
  * \tparam T The type of the return value of visitor functions. Typical values:
@@ -383,6 +385,8 @@ class GrammarFSMBuilder {
 
  public:
   static void Apply(Grammar* grammar);
+  static void Apply(Grammar* grammar, int max_threads);
+  static void Apply(Grammar* grammar, int max_threads, ThreadPool* thread_pool);
   static FSMWithStartEnd RuleRef(const GrammarExpr& expr);
   static FSMWithStartEnd CharacterClass(const GrammarExpr& expr);
   static FSMWithStartEnd ByteString(const GrammarExpr& expr);
@@ -440,6 +444,8 @@ class RepetitionRangeExpander {
 class GrammarOptimizer {
  public:
   static Grammar Apply(const Grammar& grammar);
+  static Grammar Apply(const Grammar& grammar, int max_threads);
+  static Grammar Apply(const Grammar& grammar, int max_threads, ThreadPool* thread_pool);
 };
 
 /*!