include/tenzr/txpr/jit.h¶

// 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.

#ifndef TENZR_TXPR_JIT_H_
#define TENZR_TXPR_JIT_H_

#include <atomic>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <tuple>
#include <type_traits>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/common/trace.h"
#include "tenzr/execution/execution_engine.h"
#include "tenzr/execution/jit_cache.h"
#include "tenzr/frontend/array.h"
#include "tenzr/runtime/hardware_manager.h"
#include "tenzr/txpr/emitter.h"
#include "tenzr/txpr/tracer.h"
#include "tenzr/txpr/txpr.h"

namespace tenzr::txpr {

  namespace detail {

    inline uint64_t NextJitToken() noexcept {
      static std::atomic<uint64_t> counter{1};
      return counter.fetch_add(1, std::memory_order_relaxed);
    }

    inline void HashCombine(uint64_t& seed, uint64_t val) noexcept { seed ^= val + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); }

    inline uint64_t HashSignature(uint64_t fn_id, const std::vector<Array>& args) noexcept {
      uint64_t seed = fn_id;
      for (const auto& a : args) {
        HashCombine(seed, static_cast<uint64_t>(a.Dtype()));
        HashCombine(seed, a.shape().Rank());
        for (size_t d = 0; d < a.shape().Rank(); ++d) {
          HashCombine(seed, static_cast<uint64_t>(a.shape().Dim(d)));
        }
      }
      return seed;
    }

  }  // namespace detail

  ExecutionResult ExecuteJittedExecutable(const std::shared_ptr<CachedExecutable>& cached, const std::vector<Array>& inputs);

  std::shared_ptr<CachedExecutable> CompileTxprModule(uint64_t hash, const Txpr& txpr);

  inline std::shared_ptr<CachedExecutable> CompileTapeModule(uint64_t hash, const Txpr& txpr) { return CompileTxprModule(hash, txpr); }

  template <typename Fn> class JittedFunction {
   public:
    explicit JittedFunction(Fn fn) : fn_(std::move(fn)), fn_token_(detail::NextJitToken()) {}

    template <typename... Args> Array operator()(const Args&... args) {
      if (TracerContext::IsActive() || (args.IsSymbolic() || ...)) {
        return fn_(args...);
      }

      TENZR_TRACE_SCOPE_CAT("jit", "JittedFunction::dispatch");

      std::vector<Array> in_args = {args...};
      uint64_t sig_hash = detail::HashSignature(fn_token_, in_args);
      auto cached = JitCache::Instance().Get(JitCache::DomainKey(JitCache::kJitKeyDomain, sig_hash));

      if (!cached) {
        TENZR_TRACE_SCOPE_CAT("tracing", "JittedFunction::trace_on_miss");
        Txpr txpr;
        std::vector<Array> sym_args;
        sym_args.reserve(in_args.size());
        for (const auto& arr : in_args) {
          VarId vid = txpr.AddVar(arr.shape(), arr.Dtype());
          txpr.MarkInput(vid);
          sym_args.push_back(Array::Symbolic(vid, arr.shape(), arr.Dtype()));
        }

        auto sym_tuple = [&sym_args]<size_t... Is>(std::index_sequence<Is...>) {
          return std::forward_as_tuple(sym_args[Is]...);
        }(std::make_index_sequence<sizeof...(Args)>{});

        Array sym_out;
        {
          TracerScope scope(txpr);
          sym_out = std::apply(fn_, sym_tuple);
        }
        TENZR_CHECK(sym_out.IsSymbolic(), "JIT function output must be symbolic during tracing");
        txpr.MarkOutput(sym_out.SymbolicVarId());

        cached = CompileTxprModule(sig_hash, txpr);
      }

      ExecutionResult exec_res = ExecuteJittedExecutable(cached, in_args);
      return exec_res.AsArray();
    }

   private:
    Fn fn_;
    uintptr_t fn_token_{0};
  };

  template <typename Fn>
    requires(!std::is_same_v<std::decay_t<Fn>, Array>)
  auto Jit(Fn&& fn) {
    return JittedFunction<std::decay_t<Fn>>(std::forward<Fn>(fn));
  }

}  // namespace tenzr::txpr

namespace tenzr {
  using txpr::CompileTapeModule;
  using txpr::CompileTxprModule;
  using txpr::ExecuteJittedExecutable;
  using txpr::Jit;
  using txpr::JittedFunction;
}  // namespace tenzr

#endif  // TENZR_TXPR_JIT_H_