[jit] Refactor

src/jit/asmjit/asmjit.cc

@@ -124,7 +124,7 @@ namespace {
 
 using JitFunction = void (*)(CPUState &);
 
-class AsmJit : public JitEngine {
+class AsmJit : public Translator {
 public:
   AsmJit() = default;
 
@@ -297,5 +297,5 @@ JitFunction AsmJit::translate(const BBInfo &info) {
 }
 } // namespace
 
-std::unique_ptr<ExecEngine> makeAsmJit() { return std::make_unique<AsmJit>(); }
+std::unique_ptr<Translator> makeAsmJit() { return std::make_unique<AsmJit>(); }
 } // namespace prot::engine

src/jit/asmjit/include/prot/jit/asmjit.hh

@@ -3,10 +3,10 @@
 
 #include <memory>
 
-#include "prot/exec_engine.hh"
+#include "prot/jit/base.hh"
 
 namespace prot::engine {
-std::unique_ptr<ExecEngine> makeAsmJit();
+std::unique_ptr<Translator> makeAsmJit();
 }
 
 #endif // PROT_JIT_ASMJIT_HH_INCLUDED

src/jit/base/base.cc

@@ -19,10 +19,12 @@ void JitEngine::step(CPUState &cpu) {
 
     // colllect bb
     const auto pc = cpu.getPC();
-    auto found = m_tbCache.lookup(pc);
-    if (found != nullptr) [[likely]] {
-      found(cpu);
-      continue;
+    if (m_translator) {
+      if (const auto found = m_tbCache.lookup(pc); found != nullptr)
+          [[likely]] {
+        found(cpu);
+        continue;
+      }
     }
 
     auto [bbIt, wasNew] = m_cacheBB.try_emplace(pc);
@@ -46,13 +48,16 @@ void JitEngine::step(CPUState &cpu) {
         curAddr += isa::kWordSize;
       }
     }
-    if (bbIt->second.num_exec >= m_config.execThreshold) [[likely]] {
-      auto code = translate(bbIt->second);
-      m_tbCache.insert(pc, code);
-      if (code != nullptr) [[likely]] {
-        code(cpu);
-        continue;
+    if (m_translator && bbIt->second.num_exec >= m_config.execThreshold)
+        [[likely]] {
+      auto code = m_translator->translate(bbIt->second);
+      if (code == nullptr) [[unlikely]] {
+        throw std::runtime_error{
+            fmt::format("Failed to translate BB on pc: {:#x}", pc)};
       }
+
+      code(cpu);
+      m_tbCache.insert(pc, code);
     }
 
     interpret(cpu, bbIt->second);

src/jit/base/include/prot/jit/base.hh

@@ -9,17 +9,33 @@
 namespace prot::engine {
 using JitFunction = void (*)(CPUState &);
 
-class JitEngine : public Interpreter {
+// simple bb counting
+struct BBInfo final {
+  std::vector<isa::Instruction> insns;
+  std::size_t num_exec{};
+};
+
+struct Translator {
+  Translator() = default;
+  Translator(const Translator &) = delete;
+  Translator &operator=(const Translator &) = delete;
+
+  [[nodiscard]] virtual JitFunction translate(const BBInfo &info) = 0;
+  virtual ~Translator() = default;
+};
+
+class JitEngine final : public Interpreter {
 public:
   struct Config final {
     std::size_t execThreshold{};
     bool singleStep{false};
     bool enableDump{false};
   };
 
-  void step(CPUState &cpu) override;
+  JitEngine(const Config &config, std::unique_ptr<Translator> translator)
+      : m_config{config}, m_translator{std::move(translator)} {}
 
-  void setConfig(const Config &config) { m_config = config; }
+  void step(CPUState &cpu) override;
 
 protected:
   struct TbCache {
@@ -52,11 +68,6 @@ protected:
     std::array<Entry, kSize> m_cache;
   };
 
-  // simple bb counting
-  struct BBInfo final {
-    std::vector<isa::Instruction> insns;
-    std::size_t num_exec{};
-  };
   [[nodiscard]] const BBInfo *getBBInfo(isa::Addr pc) const;
 
 private:
@@ -66,19 +77,12 @@ private:
   }
 
 private:
-  [[nodiscard]] virtual JitFunction translate(const BBInfo &info) = 0;
-
   Config m_config{};
   TbCache m_tbCache;
+  std::unique_ptr<Translator> m_translator;
   std::unordered_map<isa::Addr, BBInfo> m_cacheBB;
 };
 
-class CachedInterpreter final : public JitEngine {
-  JitFunction translate(const BBInfo & /* unused */) override {
-    return nullptr;
-  }
-};
-
 // Helper class to store JITed code
 // Especially helpful for libraries w/out propper mem pool support
 class CodeHolder final {

src/jit/factory/factory.cc

@@ -13,12 +13,11 @@
 
 namespace prot::engine {
 const std::unordered_map<std::string_view,
-                         std::function<std::unique_ptr<ExecEngine>()>>
+                         std::function<std::unique_ptr<Translator>()>>
     JitFactory::kFactories = {
         {"xbyak", []() { return makeXbyak(); }},
         {"asmjit", []() { return makeAsmJit(); }},
-        {"cached-interp",
-         []() { return std::make_unique<CachedInterpreter>(); }},
+        {"cached-interp", []() { return std::unique_ptr<Translator>(); }},
         {"llvm", []() { return makeLLVMBasedJIT(); }},
         {"lightning", []() { return makeLightning(); }},
         {"mir", []() { return makeMirJit(); }},
@@ -30,20 +29,10 @@ std::vector<std::string_view> JitFactory::backends() {
   return res;
 }
 
-std::unique_ptr<ExecEngine>
-JitFactory::createEngine(const std::string &backend,
-                         const JitEngine::Config &config) {
-  auto it = kFactories.find(backend);
-  if (it != kFactories.end()) {
-    auto engine = it->second();
-    auto *jitBase = dynamic_cast<JitEngine *>(engine.get());
-    if (jitBase == nullptr) {
-      throw std::invalid_argument{"Not a JIT engine"};
-    }
-
-    jitBase->setConfig(config);
-
-    return engine;
+std::unique_ptr<Translator>
+JitFactory::createTranslator(const std::string &backend) {
+  if (const auto it = kFactories.find(backend); it != kFactories.end()) {
+    return it->second();
   }
 
   throw std::invalid_argument("Undefined JIT backend: " + backend);

src/jit/factory/include/prot/jit/factory.hh

@@ -15,13 +15,13 @@ namespace prot::engine {
 class JitFactory {
 public:
   [[nodiscard]] static std::vector<std::string_view> backends();
-  static std::unique_ptr<ExecEngine>
-  createEngine(const std::string &backend, const JitEngine::Config &config);
+  static std::unique_ptr<Translator>
+  createTranslator(const std::string &backend);
   static bool exist(const std::string &backend);
 
 private:
   static const std::unordered_map<std::string_view,
-                                  std::function<std::unique_ptr<ExecEngine>()>>
+                                  std::function<std::unique_ptr<Translator>()>>
       kFactories;
 };
 

src/jit/lightning/include/prot/jit/lightning.hh

@@ -3,10 +3,10 @@
 
 #include <memory>
 
-#include "prot/exec_engine.hh"
+#include "prot/jit/base.hh"
 
 namespace prot::engine {
-std::unique_ptr<ExecEngine> makeLightning();
+std::unique_ptr<Translator> makeLightning();
 }
 
 #endif // INCLUDE_PROT_JIT_LIGHTNING_HH_INCLUDED

src/jit/lightning/lightning.cc

@@ -15,7 +15,7 @@ extern "C" {
 namespace prot::engine {
 
 namespace {
-struct Lightning : public JitEngine {
+struct Lightning : public Translator {
   Lightning() { init_jit("JIT Research"); }
 
   [[nodiscard]] JitFunction translate(const BBInfo &info) override;
@@ -333,7 +333,7 @@ JitFunction Lightning::translate(const BBInfo &info) {
 }
 } // namespace
 
-std::unique_ptr<ExecEngine> makeLightning() {
+std::unique_ptr<Translator> makeLightning() {
   return std::make_unique<Lightning>();
 }
 } // namespace prot::engine

src/jit/llvm/jit/include/prot/jit/llvmbasedjit.hh

@@ -4,10 +4,10 @@
 #include <llvm/Support/Error.h>
 #include <memory>
 
-#include "prot/exec_engine.hh"
+#include "prot/jit/base.hh"
 
 namespace prot::engine {
-std::unique_ptr<ExecEngine> makeLLVMBasedJIT();
+std::unique_ptr<Translator> makeLLVMBasedJIT();
 } // end namespace prot::engine
 
 #endif // PROT_JIT_LLVMBASEDJIT_HH_INCLUDED

src/jit/llvm/jit/llvmbasedjit.cc

@@ -37,12 +37,10 @@
 
 namespace prot::engine {
 namespace {
-class LLVMBasedJIT : public JitEngine {
+class LLVMBasedJIT : public Translator {
   std::unique_ptr<llvm::orc::LLJIT> m_jit;
   std::size_t m_moduleId{};
 
-  using TBFunc = void (*)(CPUState &);
-
 public:
   LLVMBasedJIT(std::unique_ptr<llvm::orc::LLJIT> JIT);
 
@@ -97,7 +95,7 @@ LLVMBasedJIT::LLVMBasedJIT(std::unique_ptr<llvm::orc::LLJIT> JIT)
 
 } // namespace
 
-std::unique_ptr<ExecEngine> makeLLVMBasedJIT() {
+std::unique_ptr<Translator> makeLLVMBasedJIT() {
   LLVMInitializeNativeTarget();
   LLVMInitializeNativeAsmPrinter();
   LLVMInitializeNativeAsmParser();

src/jit/mir/include/prot/jit/mir.hh

@@ -3,10 +3,10 @@
 
 #include <memory>
 
-#include "prot/exec_engine.hh"
+#include "prot/jit/base.hh"
 
 namespace prot::engine {
-std::unique_ptr<ExecEngine> makeMirJit();
+std::unique_ptr<Translator> makeMirJit();
 }
 
 #endif // PROT_JIT_MIR_HH_INCLUDED

src/jit/mir/mir.cc

@@ -142,7 +142,7 @@ template <typename T> T loadHelper(CPUState &state, isa::Addr addr) {
 
 void syscallHelper(CPUState &state) { state.emulateSysCall(); }
 
-class MIRJit : public JitEngine {
+class MIRJit : public Translator {
 public:
   MIRJit() : ctx(MIR_init()) {
     MIR_gen_init(ctx);
@@ -395,5 +395,5 @@ JitFunction MIRJit::translate(const BBInfo &info) {
 
 } // namespace
 
-std::unique_ptr<ExecEngine> makeMirJit() { return std::make_unique<MIRJit>(); }
+std::unique_ptr<Translator> makeMirJit() { return std::make_unique<MIRJit>(); }
 } // namespace prot::engine

src/jit/tpde/include/prot/jit/tpde.hh

@@ -3,10 +3,10 @@
 
 #include <memory>
 
-#include "prot/exec_engine.hh"
+#include "prot/jit/base.hh"
 
 namespace prot::engine {
-std::unique_ptr<ExecEngine> makeTPDE();
+std::unique_ptr<Translator> makeTPDE();
 }
 
 #endif // PROT_JIT_TPDE_HH_INCLUDED

src/jit/tpde/tpde.cc

@@ -19,7 +19,7 @@ namespace prot::engine {
 
 namespace {
 
-class TPDEJit final : public JitEngine {
+class TPDEJit final : public Translator {
 public:
   TPDEJit()
       : m_jit([] {
@@ -56,5 +56,5 @@ class TPDEJit final : public JitEngine {
 };
 } // namespace
 
-std::unique_ptr<ExecEngine> makeTPDE() { return std::make_unique<TPDEJit>(); }
+std::unique_ptr<Translator> makeTPDE() { return std::make_unique<TPDEJit>(); }
 } // namespace prot::engine

src/jit/xbyak/include/prot/jit/xbyak.hh

@@ -3,10 +3,10 @@
 
 #include <memory>
 
-#include "prot/exec_engine.hh"
+#include "prot/jit/base.hh"
 
 namespace prot::engine {
-std::unique_ptr<ExecEngine> makeXbyak();
+std::unique_ptr<Translator> makeXbyak();
 }
 
 #endif // PROT_JIT_XBYAK_HH_INCLUDED

src/jit/xbyak/xbyak.cc

@@ -14,7 +14,7 @@
 
 namespace prot::engine {
 namespace {
-class XByakJit : public JitEngine, private Xbyak::CodeGenerator {
+class XByakJit : public Translator, private Xbyak::CodeGenerator {
 public:
   XByakJit()
       : Xbyak::CodeGenerator{Xbyak::DEFAULT_MAX_CODE_SIZE, Xbyak::AutoGrow} {}
@@ -289,5 +289,5 @@ JitFunction XByakJit::translate(const BBInfo &info) {
 } // namespace
 } // namespace
 
-std::unique_ptr<ExecEngine> makeXbyak() { return std::make_unique<XByakJit>(); }
+std::unique_ptr<Translator> makeXbyak() { return std::make_unique<XByakJit>(); }
 } // namespace prot::engine

tools/sim/sim_app.cpp

@@ -56,7 +56,8 @@ int main(int argc, const char *argv[]) try {
 
     auto engine = [&]() -> std::unique_ptr<prot::ExecEngine> {
       if (jitEnabled) {
-        return prot::engine::JitFactory::createEngine(jitBackend, jitConfig);
+        return std::make_unique<prot::engine::JitEngine>(
+            jitConfig, prot::engine::JitFactory::createTranslator(jitBackend));
       }
       return std::make_unique<prot::engine::Interpreter>();
     }();