src/jit/trace_cache.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <brotensor/jit/trace.h> | ||
| 4 | |||
| 5 | #include <atomic> | ||
| 6 | #include <chrono> | ||
| 7 | #include <cstddef> | ||
| 8 | #include <cstdint> | ||
| 9 | #include <functional> | ||
| 10 | #include <memory> | ||
| 11 | #include <vector> | ||
| 12 | |||
| 13 | namespace brotensor::jit { | ||
| 14 | |||
| 15 | 11 | class TraceHandleImpl { | |
| 16 | public: | ||
| 17 | 11 | uint64_t signature_hash = 0; | |
| 18 | 11 | bool is_cuda = false; | |
| 19 | 11 | size_t node_count = 0; | |
| 20 | 11 | double last_duration_us = 0.0; | |
| 21 | 11 | bool is_cache_hit = false; | |
| 22 | |||
| 23 | // Kernel launches one execute() issues, and what the compiler called the | ||
| 24 | // fusion it picked. Both are read by the benchmarks and the tests. | ||
| 25 | 11 | size_t launch_count = 0; | |
| 26 | 11 | double compile_us = 0.0; | |
| 27 | 11 | const char* fusion_name = "unfused"; | |
| 28 | |||
| 29 | // Callable kernel dispatcher with bound buffer pointers | ||
| 30 | std::function<void()> execute_fn; | ||
| 31 | |||
| 32 | // Factory/builder to re-bind pointers when reusing compiled artifact | ||
| 33 | std::function<std::shared_ptr<TraceHandleImpl>(const std::vector<void*>& inputs, | ||
| 34 | const std::vector<void*>& outputs)> rebind_fn; | ||
| 35 | |||
| 36 | 564 | void execute() { | |
| 37 |
1/2✓ Branch 0 taken 564 times.
✗ Branch 1 not taken.
|
564 | if (!execute_fn) { |
| 38 | ✗ | throw std::runtime_error("brotensor::jit: TraceHandle has no executable function bound"); | |
| 39 | } | ||
| 40 | 564 | auto t0 = std::chrono::steady_clock::now(); | |
| 41 | 564 | execute_fn(); | |
| 42 | 564 | auto t1 = std::chrono::steady_clock::now(); | |
| 43 | 564 | last_duration_us = std::chrono::duration<double, std::micro>(t1 - t0).count(); | |
| 44 | 564 | } | |
| 45 | }; | ||
| 46 | |||
| 47 | // ─── TraceCache ───────────────────────────────────────────────────────────── | ||
| 48 | // | ||
| 49 | // Compiled-trace lookup by DAG signature, with no lock on any path. | ||
| 50 | // | ||
| 51 | // The table is a fixed-capacity open-addressed array of write-once entries. | ||
| 52 | // An insert claims a slot by compare-exchanging its key from 0 to the | ||
| 53 | // signature; the winner then writes the owning shared_ptr and publishes it by | ||
| 54 | // storing the raw pointer with release ordering. A reader that sees a | ||
| 55 | // published pointer (acquire) is guaranteed to see the fully written | ||
| 56 | // shared_ptr beside it, and nothing ever rewrites an entry, so the copy it | ||
| 57 | // takes cannot race with a writer. | ||
| 58 | // | ||
| 59 | // Two threads that race on the same signature both compile; the loser's | ||
| 60 | // artifact is simply dropped and it uses the winner's. That costs a duplicate | ||
| 61 | // PTX compile in a case that does not arise in practice (a trace is compiled | ||
| 62 | // once, on the first step) and buys a table with no blocking and no spin. | ||
| 63 | // | ||
| 64 | // Capacity is fixed because a trace population is small and static: a model's | ||
| 65 | // worth of fused patterns is tens of entries. A full table degrades to "always | ||
| 66 | // recompile", never to a wrong answer. | ||
| 67 | 1 | class TraceCache { | |
| 68 | public: | ||
| 69 | static constexpr size_t kCapacity = 1024; // power of two | ||
| 70 | |||
| 71 | static TraceCache& instance(); | ||
| 72 | |||
| 73 | std::shared_ptr<TraceHandleImpl> lookup(uint64_t hash); | ||
| 74 | void insert(uint64_t hash, std::shared_ptr<TraceHandleImpl> handle); | ||
| 75 | |||
| 76 | size_t hit_count() const noexcept; | ||
| 77 | size_t miss_count() const noexcept; | ||
| 78 | |||
| 79 | // Drops every entry. Not safe against concurrent lookup/insert — it exists | ||
| 80 | // for tests that want a cold cache, which run single-threaded. | ||
| 81 | void clear(); | ||
| 82 | |||
| 83 | private: | ||
| 84 | TraceCache() = default; | ||
| 85 | |||
| 86 | struct Entry { | ||
| 87 | std::atomic<uint64_t> key{0}; | ||
| 88 | std::atomic<TraceHandleImpl*> published{nullptr}; | ||
| 89 | std::shared_ptr<TraceHandleImpl> owner; | ||
| 90 | }; | ||
| 91 | |||
| 92 | // A signature of 0 means "empty", so fold it onto a different value. | ||
| 93 |
1/2✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
|
21 | static uint64_t normalize(uint64_t h) { return h ? h : 0x9E3779B97F4A7C15ULL; } |
| 94 | 21 | static size_t slot_of(uint64_t h) { return static_cast<size_t>(h) & (kCapacity - 1); } | |
| 95 | |||
| 96 | // A plain array, not a vector: the entries hold atomics, so the storage | ||
| 97 | // must never be relocated. | ||
| 98 | Entry table_[kCapacity]; | ||
| 99 | mutable std::atomic<size_t> hit_count_{0}; | ||
| 100 | mutable std::atomic<size_t> miss_count_{0}; | ||
| 101 | }; | ||
| 102 | |||
| 103 | } // namespace brotensor::jit | ||
| 104 |