src/jit/trace_cache.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "trace_cache.h" | ||
| 2 | #include "trace_dag.h" | ||
| 3 | |||
| 4 | #include <stdexcept> | ||
| 5 | |||
| 6 | namespace brotensor::jit { | ||
| 7 | |||
| 8 | 553 | void TraceHandle::execute() { | |
| 9 |
1/2✓ Branch 0 taken 553 times.
✗ Branch 1 not taken.
|
553 | if (!impl_) { |
| 10 | ✗ | throw std::runtime_error("brotensor::jit: cannot execute uninitialized TraceHandle"); | |
| 11 | } | ||
| 12 | 553 | impl_->execute(); | |
| 13 | 553 | } | |
| 14 | |||
| 15 | ✗ | uint64_t TraceHandle::signature_hash() const { | |
| 16 | ✗ | return impl_ ? impl_->signature_hash : 0; | |
| 17 | } | ||
| 18 | |||
| 19 | 1 | bool TraceHandle::is_cuda() const { | |
| 20 |
1/2✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
|
1 | return impl_ ? impl_->is_cuda : false; |
| 21 | } | ||
| 22 | |||
| 23 | 1 | size_t TraceHandle::node_count() const { | |
| 24 |
1/2✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
|
1 | return impl_ ? impl_->node_count : 0; |
| 25 | } | ||
| 26 | |||
| 27 | ✗ | double TraceHandle::last_duration_us() const { | |
| 28 | ✗ | return impl_ ? impl_->last_duration_us : 0.0; | |
| 29 | } | ||
| 30 | |||
| 31 | 2 | bool TraceHandle::is_cache_hit() const { | |
| 32 |
1/2✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
|
2 | return impl_ ? impl_->is_cache_hit : false; |
| 33 | } | ||
| 34 | |||
| 35 | 6 | size_t TraceHandle::launch_count() const { | |
| 36 |
1/2✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
|
6 | return impl_ ? impl_->launch_count : 0; |
| 37 | } | ||
| 38 | |||
| 39 | ✗ | double TraceHandle::compile_us() const { | |
| 40 | ✗ | return impl_ ? impl_->compile_us : 0.0; | |
| 41 | } | ||
| 42 | |||
| 43 | 5 | const char* TraceHandle::fusion_name() const { | |
| 44 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | return impl_ ? impl_->fusion_name : "none"; |
| 45 | } | ||
| 46 | |||
| 47 | 24 | TraceCache& TraceCache::instance() { | |
| 48 |
3/4✓ Branch 0 taken 1 time.
✓ Branch 1 taken 23 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
|
24 | static TraceCache cache; |
| 49 | 24 | return cache; | |
| 50 | } | ||
| 51 | |||
| 52 | 11 | std::shared_ptr<TraceHandleImpl> TraceCache::lookup(uint64_t hash) { | |
| 53 | 11 | const uint64_t k = normalize(hash); | |
| 54 | 11 | size_t i = slot_of(k); | |
| 55 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 11 times.
|
11 | for (size_t probe = 0; probe < kCapacity; ++probe, i = (i + 1) & (kCapacity - 1)) { |
| 56 | 11 | const uint64_t cur = table_[i].key.load(std::memory_order_acquire); | |
| 57 |
2/2✓ Branch 0 taken 1 time.
✓ Branch 1 taken 10 times.
|
11 | if (cur == 0) break; // empty slot ends the probe run |
| 58 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
|
1 | if (cur != k) continue; |
| 59 | // Claimed for this signature. A slot whose owner is not published yet | ||
| 60 | // belongs to a compile still in flight on another thread; report a | ||
| 61 | // miss and let this thread compile its own rather than spin. | ||
| 62 |
1/2✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
|
1 | if (table_[i].published.load(std::memory_order_acquire) == nullptr) break; |
| 63 | 1 | hit_count_.fetch_add(1, std::memory_order_relaxed); | |
| 64 | 1 | return table_[i].owner; | |
| 65 | } | ||
| 66 | 10 | miss_count_.fetch_add(1, std::memory_order_relaxed); | |
| 67 | 10 | return nullptr; | |
| 68 | 11 | } | |
| 69 | |||
| 70 | 10 | void TraceCache::insert(uint64_t hash, std::shared_ptr<TraceHandleImpl> handle) { | |
| 71 |
1/2✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
|
10 | if (!handle) return; |
| 72 | 10 | const uint64_t k = normalize(hash); | |
| 73 | 10 | size_t i = slot_of(k); | |
| 74 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
|
10 | for (size_t probe = 0; probe < kCapacity; ++probe, i = (i + 1) & (kCapacity - 1)) { |
| 75 | 10 | uint64_t expect = 0; | |
| 76 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
|
10 | if (table_[i].key.compare_exchange_strong(expect, k, std::memory_order_acq_rel, |
| 77 | std::memory_order_acquire)) { | ||
| 78 | 10 | table_[i].owner = std::move(handle); | |
| 79 | 10 | table_[i].published.store(table_[i].owner.get(), std::memory_order_release); | |
| 80 | 10 | return; | |
| 81 | } | ||
| 82 | ✗ | if (expect == k) { | |
| 83 | // Someone else owns this signature. If they have not published yet | ||
| 84 | // they are about to; either way their artifact is equivalent to | ||
| 85 | // ours, so drop ours. | ||
| 86 | ✗ | return; | |
| 87 | } | ||
| 88 | ✗ | } | |
| 89 | // Table full: the trace stays uncached and will be recompiled next time. | ||
| 90 | 10 | } | |
| 91 | |||
| 92 | 2 | size_t TraceCache::hit_count() const noexcept { | |
| 93 | 2 | return hit_count_.load(std::memory_order_relaxed); | |
| 94 | } | ||
| 95 | |||
| 96 | ✗ | size_t TraceCache::miss_count() const noexcept { | |
| 97 | ✗ | return miss_count_.load(std::memory_order_relaxed); | |
| 98 | } | ||
| 99 | |||
| 100 | 1 | void TraceCache::clear() { | |
| 101 |
2/2✓ Branch 0 taken 1024 times.
✓ Branch 1 taken 1 time.
|
1025 | for (auto& e : table_) { |
| 102 | 1024 | e.published.store(nullptr, std::memory_order_relaxed); | |
| 103 | 1024 | e.owner.reset(); | |
| 104 | 1024 | e.key.store(0, std::memory_order_relaxed); | |
| 105 | } | ||
| 106 | 1 | hit_count_.store(0, std::memory_order_relaxed); | |
| 107 | 1 | miss_count_.store(0, std::memory_order_relaxed); | |
| 108 | 1 | } | |
| 109 | |||
| 110 | 13 | void begin_trace() { | |
| 111 | 13 | TraceContext::current().begin(); | |
| 112 | 13 | } | |
| 113 | |||
| 114 | 35 | bool is_tracing() { | |
| 115 | 35 | return TraceContext::current().is_active(); | |
| 116 | } | ||
| 117 | |||
| 118 | 2 | void abort_trace() { | |
| 119 | 2 | TraceContext::current().discard(); | |
| 120 | 2 | } | |
| 121 | |||
| 122 | } // namespace brotensor::jit | ||
| 123 |