src/execution/jit_cache.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "tenzr/execution/jit_cache.h"

#include <cstring>

#include "tenzr/runtime/hardware_manager.h"

namespace tenzr {

  CachedExecutable::~CachedExecutable() {
    if (executable) {
      auto* api = HardwareManager::Instance().Api();
      if (api && api->PJRT_LoadedExecutable_Destroy) {
        PJRT_LoadedExecutable_Destroy_Args args;
        std::memset(&args, 0, sizeof(args));
        args.struct_size = PJRT_LoadedExecutable_Destroy_Args_STRUCT_SIZE;
        args.executable = executable;
        (void)api->PJRT_LoadedExecutable_Destroy(&args);
      }
      executable = nullptr;
    }
  }

  JitCache& JitCache::Instance() {
    static JitCache cache;
    return cache;
  }

  JitCache::~JitCache() { Clear(); }

  std::shared_ptr<CachedExecutable> JitCache::Get(uint64_t hash) {
    std::lock_guard<std::mutex> lock(mutex_);
    auto it = cache_.find(hash);
    if (it != cache_.end()) {
      return it->second;
    }
    return nullptr;
  }

  void JitCache::Put(uint64_t hash, std::shared_ptr<CachedExecutable> exec) {
    std::lock_guard<std::mutex> lock(mutex_);
    cache_[hash] = std::move(exec);
  }

  bool JitCache::Contains(uint64_t hash) const {
    std::lock_guard<std::mutex> lock(mutex_);
    return cache_.find(hash) != cache_.end();
  }

  size_t JitCache::Size() const {
    std::lock_guard<std::mutex> lock(mutex_);
    return cache_.size();
  }

  void JitCache::Clear() {
    std::lock_guard<std::mutex> lock(mutex_);
    cache_.clear();
  }

}  // namespace tenzr