From 7f9a3ce72fb1e087a7856360e10c7bdeb4e07258 Mon Sep 17 00:00:00 2001 From: Adrian Seyboldt Date: Thu, 16 Jun 2022 11:24:40 +0200 Subject: [PATCH] Add ffi interface for logp-gradient function --- httpstan/models.py | 45 ++++++++++++- httpstan/stan_services.cpp | 133 +++++++++++++++++++++++++++++++++++++ 2 files changed, 176 insertions(+), 2 deletions(-) diff --git a/httpstan/models.py b/httpstan/models.py index cafb6dbaa..6911c4159 100644 --- a/httpstan/models.py +++ b/httpstan/models.py @@ -93,7 +93,7 @@ def import_services_extension_module(model_name: str) -> ModuleType: return module -async def build_services_extension_module(program_code: str, extra_compile_args: Optional[List[str]] = None) -> str: +def _prepare_build_services_extension_module(program_code: str, extra_compile_args: Optional[List[str]] = None): """Compile a model-specific stan::services extension module. Since compiling an extension module takes a long time, compilation takes @@ -164,8 +164,49 @@ async def build_services_extension_module(program_code: str, extra_compile_args: extensions = [extension] build_lib = str(model_directory_path) - # Building the model takes a long time. Run in a different thread. + return extensions, build_lib + + +async def build_services_extension_module(program_code: str, extra_compile_args: Optional[List[str]] = None) -> str: + """Compile a model-specific stan::services extension module. + + Since compiling an extension module takes a long time, compilation takes + place in a different thread. + + Messages generated by the compiler—normally sent to stderr—are collected + and saved. These messages are returned by the function. + + Returns compiler messages. + + This is a coroutine function. + + IMPORTANT NOTE: This function builds the extension module in the cache + directory, making it available for later `import`ing. This "side-effect" is + why there are no functions called `load_services_extension_module` and + `dump_services_extension_module`. + + """ + extensions, build_lib = _prepare_build_services_extension_module(program_code, extra_compile_args) compiler_output = await asyncio.get_running_loop().run_in_executor( None, httpstan.build_ext.run_build_ext, extensions, build_lib ) return compiler_output + + +def build_services_extension_module_sync(program_code: str, extra_compile_args: Optional[List[str]] = None) -> str: + """Compile a model-specific stan::services extension module. + + Messages generated by the compiler—normally sent to stderr—are collected + and saved. These messages are returned by the function. + + Returns compiler messages. + + IMPORTANT NOTE: This function builds the extension module in the cache + directory, making it available for later `import`ing. This "side-effect" is + why there are no functions called `load_services_extension_module` and + `dump_services_extension_module`. + + """ + extensions, build_lib = _prepare_build_services_extension_module(program_code, extra_compile_args) + compiler_output = httpstan.build_ext.run_build_ext(extensions, build_lib) + return compiler_output diff --git a/httpstan/stan_services.cpp b/httpstan/stan_services.cpp index 8cbbf032c..2afd776fb 100644 --- a/httpstan/stan_services.cpp +++ b/httpstan/stan_services.cpp @@ -1,4 +1,6 @@ +#include #include +#include #include #include @@ -11,8 +13,11 @@ #include #include +#include + #include #include +#include #include "socket_logger.hpp" #include "socket_writer.hpp" @@ -163,6 +168,128 @@ double log_prob(py::dict data, const std::vector &unconstrained_paramete return lp; } + +struct StanLogpFunctionCtx { + py::dict data; + stan::io::var_context *var_context; + stan::model::model_base *model; +}; + +extern "C" { + int logp_gradient(size_t ndim, const double *unconstrained_parameters, double *gradient, double *logp, void *ctx) { + try { + const auto func = reinterpret_cast(ctx); + + size_t num_params = func->model->num_params_r(); + + // Unfortunately this copies the data. But I think stan only accepts data that is owned by a vector. + std::vector params_r = std::vector(unconstrained_parameters, unconstrained_parameters + num_params); + std::vector gradient_vector = std::vector(num_params); + std::vector params_i(func->model->num_params_i(), 0); + + int returncode = 0; + try { + // params_i, the third argument, is unused but the function requires it (see model_base.hpp). + *logp = stan::model::log_prob_grad(*func->model, params_r, params_i, gradient_vector, &std::cout); + std::copy(gradient_vector.begin(), gradient_vector.end(), gradient); + } catch (std::exception &ex) { + returncode = 1; + } + + if (!isfinite(*logp)) { + returncode = 2; + } + + auto has_nan = std::any_of( + gradient_vector.begin(), + gradient_vector.end(), + [](double const& val) { return !isfinite(val); } + ); + + if (has_nan) { + returncode = 3; + } + + return returncode; + } catch (std::exception &ex) { + return -1; + } + } +} + +std::uintptr_t new_logp_ctx(py::dict data) { + stan::io::array_var_context &var_context = new_array_var_context(data); + stan::model::model_base &model = new_model(var_context, (unsigned int)1, &std::cout); + + stan::math::ChainableStack::instance_ = new stan::math::AutodiffStackSingleton::AutodiffStackStorage(); + + auto ctx = new StanLogpFunctionCtx { + data, + &var_context, + &model, + }; + + return reinterpret_cast(ctx); +} + +void free_logp_ctx(std::uintptr_t ctx) { + auto func = reinterpret_cast(ctx); + delete func->model; + delete func->var_context; + delete func; +} + +std::uintptr_t logp_func(std::uintptr_t ctx) { + return reinterpret_cast(&logp_gradient); +} + +size_t num_unconstrained_parameters(std::uintptr_t ctx) { + auto func = reinterpret_cast(ctx); + return func->model->num_params_r(); +} + +py::array_t write_array_ctx(std::uintptr_t ctx, const py::array_t unconstrained_parameters, + bool include_tparams = true, bool include_gqs = true, int seed = 0) { + auto func = reinterpret_cast(ctx); + boost::ecuyer1988 base_rng(seed); + std::vector params_r_constrained_vec; + if (unconstrained_parameters.size() != func->model->num_params_r()) { + throw std::runtime_error( + "The number of parameters does not match the number of unconstrained parameters in the model."); + } + + if (unconstrained_parameters.ndim() != 1) { + throw std::runtime_error( + "Array of unconstrained parameters must be one dimensional" + ); + } + + // The params_r parameter is incorrectly declared as non-const in Stan C++. + // Unconstrained_parameters are cast from const to non-const below, as required by Stan (see model_base.hpp). + std::vector params_r = std::vector(unconstrained_parameters.data(), unconstrained_parameters.data() + unconstrained_parameters.size()); + // constrain parameters to their defined support + std::exception_ptr p; + std::vector params_i(func->model->num_params_i(), 0); + try { + // params_i, the third argument, is unused but the function requires it (see model_base.hpp). + func->model->write_array(base_rng, params_r, params_i, params_r_constrained_vec, include_tparams, include_gqs, &std::cout); + } catch (std::exception &ex) { + p = std::current_exception(); + } + + if (p) + std::rethrow_exception(p); + + auto params_r_constrained = py::array_t(params_r_constrained_vec.size()); + double *ptr = static_cast(params_r_constrained.request().ptr); + + for (size_t idx = 0; idx < params_r_constrained_vec.size(); idx++) { + ptr[idx] = params_r_constrained_vec[idx]; + } + + return params_r_constrained; +} + // See exported docstring std::vector log_prob_grad(py::dict data, const std::vector &unconstrained_parameters, bool adjust_transform) { @@ -371,4 +498,10 @@ PYBIND11_MODULE(stan_services, m) { m.def("fixed_param_wrapper", &fixed_param_wrapper, py::arg("socket_filename"), py::arg("data"), py::arg("init"), py::arg("random_seed"), py::arg("chain"), py::arg("init_radius"), py::arg("num_samples"), py::arg("num_thin"), py::arg("refresh"), "Call stan::services::sample::fixed_param"); + m.def("new_logp_ctx", &new_logp_ctx, py::arg("data"), "Create new logp function context"); + m.def("free_logp_ctx", &free_logp_ctx, py::arg("ctx"), "Destroy a logp function context"); + m.def("logp_func", &logp_func, py::arg("ctx"), "Return a C-function for computing logp values and gradients."); + m.def("num_unconstrained_parameters", &num_unconstrained_parameters, py::arg("ctx"), "Get the number of unconstrained parameters"); + m.def("write_array_ctx", &write_array_ctx, py::arg("ctx"), py::arg("unconstrained_parameters"), py::arg("include_tparams"), + py::arg("include_gqs"), py::arg("seed"), "Save all parameters at unconstrained parameter position."); }