Unify kernel and firmware JIT build deduplication into JitBuildCache (#37452)

ruizhangTT · cursoragent · web-flow · commit bfd3b818fef7 · 2026-02-13T18:48:06.000Z
### Ticket #31189 ### Problem description Multi-thread concurrency handling for JIT build for the same target is handled by HashLookup for kernels, and mutex+set for firmware at call sites. This logic can be absorbed into jit build and the call sites can be cleaned up. ### What's changed Replace two separate concurrency/deduplication mechanisms (HashLookup for kernels, mutex+set for firmware) with a single JitBuildCache class in the jit_build module. Callers now use jit_build_once(hash, fn) which guarantees exactly-once execution per hash, with concurrent callers blocking until the build completes. Fixed a subtle bug in the old HashLookup::clear() where a thread waiting in wait_for_bin_generated() could sleep forever if clear() was called concurrently. ### Checklist - [ ] [![All post-commit tests](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml/badge.svg?branch=ruizhang/jit_build_cache)](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml?query=branch:ruizhang/jit_build_cache) - [ ] [![Blackhole Post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml/badge.svg?branch=ruizhang/jit_build_cache)](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml?query=branch:ruizhang/jit_build_cache) - [ ] [![cpp-unit-tests](https://github.com/tenstorrent/tt-metal/actions/workflows/tt-metal-l2-nightly.yaml/badge.svg?branch=ruizhang/jit_build_cache)](https://github.com/tenstorrent/tt-metal/actions/workflows/tt-metal-l2-nightly.yaml?query=branch:ruizhang/jit_build_cache) (failures in tools test is due to [recent change](#37147 (comment))) - [ ] New/Existing tests provide coverage for changes --------- Co-authored-by: Cursor <cursoragent@cursor.com>
diff --git a/tests/tt_metal/tt_metal/api/test_kernel_compile_cache.cpp b/tests/tt_metal/tt_metal/api/test_kernel_compile_cache.cpp
@@ -13,7 +13,6 @@
 
 #include <tt-metalium/core_coord.hpp>
 #include <tt-metalium/data_types.hpp>
-#include "detail/kernel_cache.hpp"
 #include <tt-metalium/device.hpp>
 #include "device_fixture.hpp"
 #include <tt-metalium/hal.hpp>
@@ -35,7 +34,7 @@ TEST_F(MeshDeviceFixture, TensixTestEquivalentDataMovementKernelsWithDifferentPr
 
     for (const auto& mesh_device : this->devices_) {
         auto* device = mesh_device->get_devices()[0];
-        detail::ClearKernelCache();
+        jit_build_cache_clear();
 
         DataMovementConfig config_riscv_0 = {.processor = DataMovementProcessor::RISCV_0};
         DataMovementConfig config_riscv_1 = {.processor = DataMovementProcessor::RISCV_1};
diff --git a/tests/tt_metal/tt_metal/test_compile_program.cpp b/tests/tt_metal/tt_metal/test_compile_program.cpp
@@ -17,7 +17,6 @@
 
 #include "hal_types.hpp"
 #include "jit_build/build.hpp"
-#include "tt_metal/detail/kernel_cache.hpp"
 #include "tt_metal/jit_build/build_env_manager.hpp"
 #include <umd/device/types/arch.hpp>
 
@@ -38,7 +37,7 @@ struct KernelCacheStatus {
 
 void ClearKernelCache(const std::string& kernel_root_path) {
     std::filesystem::remove_all(kernel_root_path);
-    detail::HashLookup::inst().clear();
+    jit_build_cache_clear();
 }
 
 // This assumes binaries are written to specific location: kernel_compile_outpath / kernel_name / hash
diff --git a/tests/tt_metal/tt_metal/test_compile_sets_kernel_binaries.cpp b/tests/tt_metal/tt_metal/test_compile_sets_kernel_binaries.cpp
@@ -25,7 +25,6 @@
 #include "impl/program/program_impl.hpp"
 #include "impl/kernels/kernel.hpp"
 #include "tt_memory.h"
-#include "tt_metal/detail/kernel_cache.hpp"
 #include "tt_metal/jit_build/build_env_manager.hpp"
 #include <umd/device/types/arch.hpp>
 
@@ -184,7 +183,7 @@ TEST_F(CompileSetsKernelBinariesFixture, CompileSetsKernelBinaries) {
                     kernel_name);
             }
         }
-        detail::ClearKernelCache();
+        jit_build_cache_clear();
         std::vector<Program> new_programs;
         for (int i = 0; i < num_devices_; i++) {
             auto& device = devices_[i];
diff --git a/tt_metal/detail/kernel_cache.hpp b/tt_metal/detail/kernel_cache.hpp
diff --git a/tt_metal/impl/context/metal_context.cpp b/tt_metal/impl/context/metal_context.cpp
@@ -32,6 +32,7 @@
 #include "dispatch/topology.hpp"
 #include "dispatch/dispatch_core_common.hpp"
 #include "profiler/profiler_state_manager.hpp"
+#include "jit_build/build.hpp"
 #include "jit_build/build_env_manager.hpp"
 #include "llrt/get_platform_architecture.hpp"
 #include "llrt/llrt.hpp"
@@ -273,19 +274,13 @@ void MetalContext::initialize(
                     BuildEnvManager::get_instance().add_build_env(device_id, num_hw_cqs_);
                     // fw_build_key is a combination of build_key and fw_compile_hash
                     // If fw_compile_hash changes, the fw_build_key will change and FW will be rebuilt
-                    // if it's not already in firmware_built_keys_
                     // Combine build_key and fw_compile_hash using XOR to create unique firmware build key
                     // Uses full 64-bit fw_compile_hash for proper change detection
                     uint64_t fw_build_key =
                         BuildEnvManager::get_instance().get_device_build_env(device_id).build_key() ^ fw_compile_hash;
 
-                    {
-                        std::lock_guard<std::mutex> lock(firmware_built_keys_mutex_);
-                        if (!firmware_built_keys_.contains(fw_build_key)) {
-                            BuildEnvManager::get_instance().build_firmware(device_id);
-                            firmware_built_keys_.insert(fw_build_key);
-                        }
-                    }
+                    jit_build_once(
+                        fw_build_key, [device_id] { BuildEnvManager::get_instance().build_firmware(device_id); });
 
                     // Clear the entire launch message ring buffer on ethernet cores before application firmware is
                     // activated. This is required since ethernet cores context switch between application and routing
diff --git a/tt_metal/impl/context/metal_context.hpp b/tt_metal/impl/context/metal_context.hpp
@@ -201,10 +201,6 @@ class MetalContext {
     size_t worker_l1_unreserved_start_ = 0;
     size_t fw_compile_hash_ = 0;  // To check if FW recompilation is needed
 
-    // Used to track which FW has been built already
-    std::unordered_set<uint64_t> firmware_built_keys_;
-    std::mutex firmware_built_keys_mutex_;
-
     // Mutex to protect control_plane_ for thread-safe access
     std::mutex control_plane_mutex_;
 
diff --git a/tt_metal/impl/program/program.cpp b/tt_metal/impl/program/program.cpp
@@ -65,7 +65,6 @@
 #include "sub_device_types.hpp"
 #include "tile.hpp"
 #include "tt_memory.h"
-#include "tt_metal/detail/kernel_cache.hpp"
 #include "tt_metal/impl/debug/inspector/inspector.hpp"
 #include "tt_metal/impl/dispatch/data_collection.hpp"
 #include "tt_metal/impl/dispatch/device_command.hpp"
@@ -190,7 +189,7 @@ size_t KernelCompileHash(const std::shared_ptr<Kernel>& kernel, JitBuildOptions&
 
 namespace experimental {
 
-void ClearKernelCache() { detail::HashLookup::inst().clear(); }
+void ClearKernelCache() { jit_build_cache_clear(); }
 
 }  // namespace experimental
 
@@ -1487,17 +1486,15 @@ void detail::ProgramImpl::compile(IDevice* device, bool force_slow_dispatch) {
                     bool is_mock = tt::tt_metal::MetalContext::instance().get_cluster().get_target_device_type() ==
                                    tt::TargetDevice::Mock;
 
-                    if (detail::HashLookup::inst().add(kernel_hash)) {
+                    jit_build_once(kernel_hash, [&] {
                         if (!is_mock) {
                             GenerateBinaries(device, build_options, kernel);
                         } else {
                             // Create empty stub binaries for mock devices
                             std::vector<const ll_api::memory*> empty_binaries(kernel->expected_num_binaries(), nullptr);
                             kernel->set_binaries(build_env.build_key(), std::move(empty_binaries));
                         }
-                        detail::HashLookup::inst().add_generated_bin(kernel_hash);
-                    }
-                    detail::HashLookup::inst().wait_for_bin_generated(kernel_hash);
+                    });
 
                     Inspector::program_kernel_compile_finished(this, device, kernel, build_options);
                 },
diff --git a/tt_metal/jit_build/CMakeLists.txt b/tt_metal/jit_build/CMakeLists.txt
@@ -4,6 +4,7 @@ set(JIT_BUILD_SRCS
     ${CMAKE_CURRENT_SOURCE_DIR}/data_format.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/depend.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/genfiles.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/jit_build_cache.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/kernel_args.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/jit_build_options.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/jit_build_utils.cpp
diff --git a/tt_metal/jit_build/build.cpp b/tt_metal/jit_build/build.cpp
@@ -4,6 +4,8 @@
 
 #include "build.hpp"
 
+#include "jit_build_cache.hpp"
+
 #include <algorithm>
 #include <array>
 #include <atomic>
@@ -748,4 +750,10 @@ void sync_build_steps(std::vector<std::shared_future<void>>& events) {
     }
 }
 
+void jit_build_once(size_t hash, const std::function<void()>& build_fn) {
+    JitBuildCache::inst().build_once(hash, build_fn);
+}
+
+void jit_build_cache_clear() { JitBuildCache::inst().clear(); }
+
 }  // namespace tt::tt_metal
diff --git a/tt_metal/jit_build/build.hpp b/tt_metal/jit_build/build.hpp
@@ -168,4 +168,13 @@ void jit_link_additional_processor(
 void launch_build_step(const std::function<void()>& build_func, std::vector<std::shared_future<void>>& events);
 void sync_build_steps(std::vector<std::shared_future<void>>& events);
 
+// Execute build_fn exactly once for a given hash.
+// Concurrent callers with the same hash block until the build completes.
+// Returns immediately if hash was already built.
+// If build_fn throws, subsequent callers will retry.
+void jit_build_once(size_t hash, const std::function<void()>& build_fn);
+
+// Clear the JIT build cache so that subsequent jit_build_once() calls re-execute.
+void jit_build_cache_clear();
+
 }  // namespace tt::tt_metal
diff --git a/tt_metal/jit_build/jit_build_cache.cpp b/tt_metal/jit_build/jit_build_cache.cpp
@@ -0,0 +1,59 @@
+// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "jit_build_cache.hpp"
+
+namespace tt::tt_metal {
+
+void JitBuildCache::build_once(size_t hash, const std::function<void()>& build_fn) {
+    std::unique_lock<std::mutex> lock(mutex_);
+
+    while (true) {
+        auto it = entries_.find(hash);
+        if (it != entries_.end()) {
+            if (it->second == State::Built) {
+                // Already built, nothing to do.
+                return;
+            }
+            // Another thread is building this hash. Wait and re-check.
+            cv_.wait(lock);
+            continue;
+        }
+
+        // Hash not present -- we are the builder.
+        entries_.emplace(hash, State::Building);
+        break;
+    }
+
+    lock.unlock();
+
+    try {
+        build_fn();
+    } catch (...) {
+        // Build failed -- remove the entry so subsequent callers can retry.
+        std::lock_guard<std::mutex> guard(mutex_);
+        entries_.erase(hash);
+        cv_.notify_all();
+        throw;
+    }
+
+    {
+        std::lock_guard<std::mutex> guard(mutex_);
+        entries_[hash] = State::Built;
+    }
+    cv_.notify_all();
+}
+
+void JitBuildCache::clear() {
+    {
+        std::unique_lock<std::mutex> lock(mutex_);
+        // Only erase Built entries. In-flight builds (Building) are preserved so
+        // that waiters continue to wait for the current builder rather than
+        // starting a duplicate build for the same hash.
+        std::erase_if(entries_, [](const auto& p) { return p.second == State::Built; });
+    }
+    cv_.notify_all();
+}
+
+}  // namespace tt::tt_metal
diff --git a/tt_metal/jit_build/jit_build_cache.hpp b/tt_metal/jit_build/jit_build_cache.hpp
@@ -0,0 +1,49 @@
+// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <condition_variable>
+#include <functional>
+#include <mutex>
+#include <stdint.h>
+#include <unordered_map>
+
+namespace tt::tt_metal {
+
+// Thread-safe build-once cache for JIT compilation.
+//
+// Ensures that for a given hash (representing a build target), the build function
+// is executed exactly once. Concurrent callers with the same hash block until the
+// build completes. Callers arriving after the build is done return immediately.
+//
+// Used to deduplicate both kernel and firmware JIT builds across threads.
+class JitBuildCache {
+public:
+    static JitBuildCache& inst() {
+        static JitBuildCache instance;
+        return instance;
+    }
+
+    // Execute build_fn exactly once for a given hash.
+    // Concurrent callers with the same hash block until the build completes.
+    // Returns immediately if hash was already built.
+    // If build_fn throws, the entry is removed so subsequent callers can retry.
+    void build_once(size_t hash, const std::function<void()>& build_fn);
+
+    // Clear completed entries. After clear(), the next build_once() for any hash
+    // will re-execute the build function.
+    void clear();
+
+private:
+    JitBuildCache() = default;
+
+    enum class State { Building, Built };
+
+    std::mutex mutex_;
+    std::condition_variable cv_;
+    std::unordered_map<size_t, State> entries_;
+};
+
+}  // namespace tt::tt_metal
