Upload ChatIPC.cpp

ChatIPC.cpp

@@ -286,16 +286,79 @@ static std::vector<std::string> tokenize_non_alnum(const std::string &s){
 
 // --------------------------- String interning (short methods) --------------
 
+using StrPtr = const std::string*;
+using TokenId = std::uint32_t;
+static constexpr TokenId TOKEN_ID_INVALID = 0xFFFFFFFFu;
+
+static inline std::size_t popcount_u64(std::uint64_t x){
+#if defined(_MSC_VER)
+    return static_cast<std::size_t>(__popcnt64(x));
+#else
+    return static_cast<std::size_t>(__builtin_popcountll(static_cast<unsigned long long>(x)));
+#endif
+}
+
 struct StringInterner {
     std::unordered_set<std::string> pool;
-    std::mutex m;
+    std::unordered_map<std::string, TokenId> id_by_string;
+    std::vector<const std::string*> string_by_id;
+    mutable std::mutex m;
+
     const std::string* intern(const std::string &s){
         std::lock_guard<std::mutex> lk(m);
+
         auto [it, inserted] = pool.emplace(s);
+        if (inserted){
+            const TokenId id = static_cast<TokenId>(string_by_id.size());
+            id_by_string.emplace(*it, id);
+            string_by_id.push_back(&*it);
+        }
+
         return &*it;
     }
+
+    TokenId id_of(const std::string &s) const {
+        std::lock_guard<std::mutex> lk(m);
+
+        auto it = id_by_string.find(s);
+        return (it == id_by_string.end()) ? TOKEN_ID_INVALID : it->second;
+    }
+
+    TokenId id_of(const std::string *p) const {
+        return p ? id_of(*p) : TOKEN_ID_INVALID;
+    }
+
+    const std::string* ptr_from_id(TokenId id) const {
+        std::lock_guard<std::mutex> lk(m);
+
+        return (id < string_by_id.size()) ? string_by_id[(size_t)id] : nullptr;
+    }
+
+    std::size_t size() const {
+        std::lock_guard<std::mutex> lk(m);
+        return string_by_id.size();
+    }
 };
 
+static inline void bitset_set(std::uint64_t *bits, std::size_t words, TokenId id){
+    if (id == TOKEN_ID_INVALID) return;
+    const std::size_t idx = static_cast<std::size_t>(id >> 6);
+    if (idx >= words) return;
+    bits[idx] |= (1ULL << (id & 63u));
+}
+
+static inline std::size_t bitset_count(const std::uint64_t *bits, std::size_t words){
+    std::size_t total = 0;
+    for (std::size_t i = 0; i < words; ++i) total += popcount_u64(bits[i]);
+    return total;
+}
+
+static inline std::size_t bitset_intersection_count(const std::uint64_t *a, const std::uint64_t *b, std::size_t words){
+    std::size_t total = 0;
+    for (std::size_t i = 0; i < words; ++i) total += popcount_u64(a[i] & b[i]);
+    return total;
+}
+
 // ---------- Global parsed dictionary (populated once in main) ----------
 static void build_def_tokens_cache(){
     global_def_tokens_cache.clear();
@@ -332,7 +395,6 @@ static void build_def_tokens_cache(){
 }
 
 // --------------------------- Knowledge base (short methods) --------------
-using StrPtr = const std::string*;
 struct PtrHash { size_t operator()(StrPtr p) const noexcept { return std::hash<StrPtr>()(p); } };
 struct PtrEq   { bool operator()(StrPtr a, StrPtr b) const noexcept { return a == b; } };
 
@@ -453,28 +515,6 @@ intern_tokens(KnowledgeBase &kb, const std::vector<std::string> &tokens)
     return out;
 }
 
-static std::unordered_set<std::string>
-aggregate_sets(const std::vector<StrPtr> &tokens,
-               const std::unordered_map<StrPtr, std::vector<StrPtr>, PtrHash, PtrEq> &def_index)
-{
-    std::unordered_set<std::string> agg;
-
-    for (StrPtr t : tokens){
-        const std::string tk = normalize_dictionary_key(*t);
-        if (!tk.empty()) agg.insert(tk);
-
-        auto it = def_index.find(t);
-        if (it != def_index.end()){
-            for (StrPtr d : it->second){
-                const std::string dk = normalize_dictionary_key(*d);
-                if (!dk.empty()) agg.insert(dk);
-            }
-        }
-    }
-
-    return agg;
-}
-
 // --------------------------- Small JSON parse helpers ----------------------
 
 static inline bool json_valid_index(size_t i, size_t n){ return i < n; }
@@ -647,6 +687,7 @@ static std::string best_candidate_by_similarity(
     const std::vector<StrPtr> &prompt_ptrs,
     const std::vector<StrPtr> &resp_ptrs,
     const std::unordered_map<StrPtr, std::vector<StrPtr>, PtrHash, PtrEq> &def_index,
+    const StringInterner &interner,
     const std::unordered_map<std::string,int> &recent_counts,
     double repeat_penalty,
     const std::string &context_tok)
@@ -654,59 +695,108 @@ static std::string best_candidate_by_similarity(
     if (cands.empty()) return std::string();
     if (cands.size() == 1) return *cands[0];
 
-    auto agg = aggregate_sets(prompt_ptrs, def_index);
-    for (StrPtr r : resp_ptrs){
-        auto it = def_index.find(r);
+    const std::size_t words = std::max<std::size_t>(1, (interner.size() + 63u) / 64u);
+
+    std::vector<std::uint64_t> agg_words(words, 0ULL);
+
+    auto add_token_and_defs = [&](StrPtr t){
+        TokenId tid = interner.id_of(t);
+        if (tid != TOKEN_ID_INVALID) bitset_set(agg_words.data(), words, tid);
+
+        auto it = def_index.find(t);
         if (it != def_index.end()){
             for (StrPtr d : it->second){
-                const std::string dk = normalize_dictionary_key(*d);
-                if (!dk.empty()) agg.insert(dk);
+                TokenId did = interner.id_of(d);
+                if (did != TOKEN_ID_INVALID) bitset_set(agg_words.data(), words, did);
             }
         }
-    }
+    };
 
-    double best = -1e9;
-    std::string best_tok;
-    size_t M = cands.size();
-    std::vector<double> scores(M, 0.0);
+    for (StrPtr t : prompt_ptrs) add_token_and_defs(t);
+    for (StrPtr t : resp_ptrs) add_token_and_defs(t);
+
+    const std::size_t agg_count = bitset_count(agg_words.data(), words);
 
-    #pragma omp parallel for schedule(static)
-    for (ptrdiff_t i = 0; i < static_cast<ptrdiff_t>(M); ++i){
-        const StrPtr cand = cands[(size_t)i];
-        const std::string cand_key = normalize_dictionary_key(*cand);
+    const std::size_t M = cands.size();
+    std::vector<std::uint64_t> cand_words(M * words, 0ULL);
+    std::vector<std::size_t> cand_counts(M, 0);
 
-        size_t inter = (!cand_key.empty() && agg.count(cand_key)) ? 1 : 0;
-        size_t cand_size = 1;
+    for (std::size_t i = 0; i < M; ++i){
+        std::uint64_t *row = cand_words.data() + i * words;
+        const StrPtr cand = cands[i];
+
+        TokenId cid = interner.id_of(cand);
+        if (cid != TOKEN_ID_INVALID) bitset_set(row, words, cid);
 
         auto it = def_index.find(cand);
         if (it != def_index.end()){
-            cand_size += it->second.size();
             for (StrPtr d : it->second){
-                const std::string dk = normalize_dictionary_key(*d);
-                if (!dk.empty() && agg.count(dk)) ++inter;
-            }
-            if (std::find(it->second.begin(), it->second.end(), cand) != it->second.end()){
-                --cand_size;
+                TokenId did = interner.id_of(d);
+                if (did != TOKEN_ID_INVALID) bitset_set(row, words, did);
             }
         }
 
-        size_t uni = agg.size() + cand_size - inter;
-        double s = uni ? static_cast<double>(inter) / static_cast<double>(uni) : 0.0;
-        scores[(size_t)i] = s;
+        cand_counts[i] = bitset_count(row, words);
     }
 
-    for (size_t i = 0; i < M; ++i){
+    std::vector<double> scores(M, 0.0);
+
+#if defined(_OPENMP) && defined(CHATIPC_ENABLE_OMP_TARGET)
+    const bool use_target = (omp_get_num_devices() > 0) && (M >= 256);
+#else
+    const bool use_target = false;
+#endif
+
+    if (use_target){
+        std::uint64_t *agg_ptr = agg_words.data();
+        std::uint64_t *cand_ptr = cand_words.data();
+        std::size_t *count_ptr = cand_counts.data();
+        double *score_ptr = scores.data();
+        const std::size_t cand_words_total = cand_words.size();
+
+        #pragma omp target data map(to: agg_ptr[0:words], cand_ptr[0:cand_words_total], count_ptr[0:M]) map(from: score_ptr[0:M])
+        {
+            #pragma omp target teams distribute parallel for
+            for (ptrdiff_t i = 0; i < static_cast<ptrdiff_t>(M); ++i){
+                const std::uint64_t *row = cand_ptr + static_cast<std::size_t>(i) * words;
+
+                std::size_t inter = 0;
+                for (std::size_t w = 0; w < words; ++w){
+                    inter += popcount_u64(agg_ptr[w] & row[w]);
+                }
+
+                const std::size_t uni = agg_count + count_ptr[(size_t)i] - inter;
+                score_ptr[(size_t)i] = uni ? static_cast<double>(inter) / static_cast<double>(uni) : 0.0;
+            }
+        }
+    } else {
+        #ifdef _OPENMP
+        #pragma omp parallel for schedule(static)
+        #endif
+        for (ptrdiff_t i = 0; i < static_cast<ptrdiff_t>(M); ++i){
+            const std::uint64_t *row = cand_words.data() + static_cast<std::size_t>(i) * words;
+
+            const std::size_t inter = bitset_intersection_count(agg_words.data(), row, words);
+            const std::size_t uni = agg_count + cand_counts[(size_t)i] - inter;
+            scores[(size_t)i] = uni ? static_cast<double>(inter) / static_cast<double>(uni) : 0.0;
+        }
+    }
+
+    double best = -1e9;
+    std::string best_tok;
+
+    for (std::size_t i = 0; i < M; ++i){
         const std::string &tok = *cands[i];
         const std::string tok_key = normalize_dictionary_key(tok);
-        const std::string count_key = tok_key.empty() ? tok : tok;
+        const std::string count_key = tok_key.empty() ? tok : tok_key;
 
-        double s = scores[i];
         auto rc_it = recent_counts.find(count_key);
-        int cnt = (rc_it == recent_counts.end() ? 0 : rc_it->second);
+        const int cnt = (rc_it == recent_counts.end()) ? 0 : rc_it->second;
 
-        double adjusted = s
-                        + english_rule_bonus(context_tok, tok)
-                        - repeat_penalty * static_cast<double>(cnt);
+        const double adjusted =
+            scores[i] +
+            english_rule_bonus(context_tok, tok) -
+            repeat_penalty * static_cast<double>(cnt);
 
         if (adjusted > best || (adjusted == best && tok < best_tok)){
             best = adjusted;
@@ -787,7 +877,8 @@ static std::vector<std::string> construct_response(KnowledgeBase &kb,
         }
 
         std::string chosen = best_candidate_by_similarity(
-            candidates, prompt_ptrs, resp_ptrs, kb.def_index, recent_counts, repeat_penalty, context_tok
+            candidates, prompt_ptrs, resp_ptrs, kb.def_index, kb.interner,
+            recent_counts, repeat_penalty, context_tok
         );
 
         if (chosen.empty()) break;
@@ -1102,7 +1193,8 @@ int main(int argc, char **argv){
 
 
     if (!load_kb.empty()){
-        try { load_kb_binary(kb, load_kb, dict_depth); std::cerr << "Loaded KB: " << load_kb << "\n"; }
+        try { std::cerr << "Loading KB: " << load_kb << "\n";
+        load_kb_binary(kb, load_kb, dict_depth); std::cerr << "Loaded KB: " << load_kb << "\n"; }
         catch (const std::exception &e){ std::cerr << "Load KB error: " << e.what() << "\n"; }
     }
 
@@ -1125,15 +1217,11 @@ int main(int argc, char **argv){
             for (size_t i=1;i<combined.size();++i) kb.add_pair(combined[i-1], combined[i]);
         }
         if (!savefile.empty()){
-            try { save_kb_binary(kb, savefile); std::cerr << "Saved KB: " << savefile << "\n"; }
+            try { std::cerr << "Saving KB: " << savefile << "\n";
+            save_kb_binary(kb, savefile); std::cerr << "Saved KB: " << savefile << "\n"; }
             catch (const std::exception &e){ std::cerr << "Save KB error: " << e.what() << "\n"; }
         }
     }
 
-    if (!savefile.empty()){
-        try { save_kb_binary(kb, savefile); std::cerr << "Saved KB: " << savefile << "\n"; }
-        catch (const std::exception &e){ std::cerr << "Save KB error: " << e.what() << "\n"; }
-    }
-
     return 0;
 }