156 lines
6.5 KiB
C++
156 lines
6.5 KiB
C++
/* @file pyjit.cpp */
|
|
|
|
#include "pyjit.hpp"
|
|
#include "xo/pyexpression/pyexpression.hpp"
|
|
#include "xo/jit/MachPipeline.hpp"
|
|
#include "xo/pyutil/pycaller.hpp"
|
|
#include "xo/pyutil/pyutil.hpp"
|
|
#include <pybind11/stl.h>
|
|
|
|
namespace xo {
|
|
struct XferDbl2DblFn : public ref::Refcount {
|
|
using fptr_type = double (*) (double);
|
|
|
|
explicit XferDbl2DblFn(fptr_type fptr) : fptr_{fptr} {}
|
|
|
|
double operator() (double x) { return (*fptr_)(x); }
|
|
|
|
fptr_type fptr_;
|
|
}; /*XferDbl2DblFn*/
|
|
|
|
struct XferDblDbl2DblFn : public ref::Refcount {
|
|
using fptr_type = double (*) (double, double);
|
|
|
|
explicit XferDblDbl2DblFn(fptr_type fptr) : fptr_{fptr} {}
|
|
|
|
double operator() (double x, double y) { return (*fptr_)(x, y); }
|
|
|
|
fptr_type fptr_;
|
|
}; /*XferDblDbl2DblFn*/
|
|
|
|
namespace jit {
|
|
using xo::ast::Expression;
|
|
using xo::pyutil::pycaller_base;
|
|
using xo::pyutil::pycaller;
|
|
using xo::ref::rp;
|
|
//using xo::ref::Refcount;
|
|
using xo::ref::unowned_ptr;
|
|
namespace py = pybind11;
|
|
|
|
PYBIND11_MODULE(XO_PYJIT_MODULE_NAME(), m) {
|
|
// e.g. for xo::ast::Expression
|
|
XO_PYEXPRESSION_IMPORT_MODULE(); // py::module_::import("pyexpression");
|
|
|
|
m.doc() = "pybind11 plugin for xo-jit";
|
|
|
|
pycaller<double, double>::declare_once(m);
|
|
pycaller<double, double, double>::declare_once(m);
|
|
|
|
py::class_<MachPipeline, rp<MachPipeline>>(m, "MachPipeline")
|
|
.def_static("make", &MachPipeline::make,
|
|
py::doc("Create machine pipeline for in-process code generation"
|
|
" and execution. Not threadsafe.\n"
|
|
"Does not share resources with any other instance"))
|
|
|
|
.def_property_readonly("target_triple", &MachPipeline::target_triple,
|
|
py::doc("string describing target host for code generation"))
|
|
.def("get_function_name_v", &MachPipeline::get_function_name_v,
|
|
py::doc("get vector of function names defined in jit module"))
|
|
.def("dump_execution_session", &MachPipeline::dump_execution_session,
|
|
py::doc("write to console with state of all jit-owned dynamic libraries"))
|
|
.def("codegen",
|
|
[](MachPipeline & jit, const rp<Expression> & expr) {
|
|
return jit.codegen(expr.borrow());
|
|
},
|
|
py::arg("x"),
|
|
py::doc("generate llvm (IR) code for Expression x"),
|
|
/* we're assuming llvm-generated code lives for as long as the Jit
|
|
* instance that created it.
|
|
*
|
|
* RC 14jun2024 - I think this is true, modulo use of llvm resource trackers.
|
|
*/
|
|
py::return_value_policy::reference_internal)
|
|
.def("machgen_current_module", &MachPipeline::machgen_current_module,
|
|
py::doc("Make current module available for execution via the jit.\n"
|
|
"Adds all functions generated since last call to this method."))
|
|
.def("dump_current_module", &MachPipeline::dump_current_module,
|
|
py::doc("Dump contents of current module to console"))
|
|
|
|
/* double -> double */
|
|
.def("lookup_dbl2dbl_fn",
|
|
[](MachPipeline & jit, const std::string & symbol) {
|
|
auto llvm_addr = jit.lookup_symbol(symbol);
|
|
|
|
auto fn_addr = llvm_addr.toPtr<double (*) (double)>();
|
|
|
|
return new XferDbl2DblFn(fn_addr);
|
|
})
|
|
|
|
/* (double x double) -> double */
|
|
.def("lookup_dbldbl2dbl_fn",
|
|
[](MachPipeline & jit, const std::string & symbol) {
|
|
auto llvm_addr = jit.lookup_symbol(symbol);
|
|
|
|
auto fn_addr = llvm_addr.toPtr<double (*) (double, double)>();
|
|
|
|
return new XferDblDbl2DblFn(fn_addr);
|
|
})
|
|
|
|
.def("lookup_fn",
|
|
[](MachPipeline & jit, const std::string & prototype, const std::string & symbol) -> pycaller_base* {
|
|
auto llvm_addr = jit.lookup_symbol(symbol);
|
|
|
|
/* note: llvm_addr.toPtr<..> always succeeds,
|
|
* event if pointer refers to an object of incompatible type
|
|
*
|
|
* note: return value policy is for python to own the wrapper
|
|
*/
|
|
|
|
if((prototype == "double(double,double)") || (prototype == "double(*)(double,double)")) {
|
|
auto fn_addr = llvm_addr.toPtr<double(*)(double,double)>();
|
|
|
|
return new pycaller<double, double, double>(fn_addr);
|
|
//return new XferDblDbl2DblFn(fn_addr);
|
|
} else if ((prototype == "double(double)") || (prototype == "double(*)(double)")) {
|
|
auto fn_addr = llvm_addr.toPtr<double(*)(double)>();
|
|
|
|
return new pycaller<double, double>(fn_addr);
|
|
} else {
|
|
throw std::runtime_error(tostr("MachPipeline.lookup_fn: unknown function prototype",
|
|
xtag("p", prototype)));
|
|
}})
|
|
;
|
|
|
|
|
|
py::class_<XferDbl2DblFn, rp<XferDbl2DblFn>>(m, "XferDbl2DblFn")
|
|
.def("__call__",
|
|
[](XferDbl2DblFn & self, double x) { return self(x); }
|
|
)
|
|
;
|
|
py::class_<XferDblDbl2DblFn, rp<XferDblDbl2DblFn>>(m, "XferDblDbl2DblFn")
|
|
.def("__call__",
|
|
[](XferDblDbl2DblFn & self, double x, double y) { return self(x, y); }
|
|
)
|
|
;
|
|
|
|
py::class_<llvm::Value,
|
|
unowned_ptr<llvm::Value>>(m, "llvm_Value")
|
|
.def("print",
|
|
[](llvm::Value & x) {
|
|
std::string buf;
|
|
llvm::raw_string_ostream ss(buf);
|
|
x.print(ss);
|
|
return buf;
|
|
})
|
|
// .def("__repr__",
|
|
// &Jit::display_string)
|
|
;
|
|
|
|
}
|
|
|
|
|
|
} /*namespace jit*/
|
|
} /*namespace xo*/
|
|
|
|
|
|
/* end pyjit.cpp */
|