<commit_before>#ifndef V_SMC_CORE_SAMPLER_HPP
#define V_SMC_CORE_SAMPLER_HPP

#include <vSMC/internal/common.hpp>
#include <ostream>
#include <map>
#include <set>
#include <string>

namespace vSMC {

/// \brief SMC Sampler
///
/// \tparam T State state type. Requiment:
/// \li Consturctor: T (IntType N)
/// \li Method: copy (IntType from, IntType to)
template <typename T>
class Sampler
{
    public :

    /// The type of initialization functor
    typedef internal::function<unsigned (Particle<T> &, void *)>
        initialize_type;

    /// The type of move and mcmc functor
    typedef internal::function<unsigned (unsigned, Particle<T> &)>
        move_type;

    /// The type of ESS history vector
    typedef std::vector<double> ess_type;

    /// The type of resampling history vector
    typedef std::vector<bool> resampled_type;

    /// The type of accept count history vector
    typedef std::vector<std::vector<unsigned> > accept_type;

    /// The type of Monitor map
    typedef std::map<std::string, Monitor<T> > monitor_map_type;

    /// \brief Construct a sampler with given number of particles
    ///
    /// \param N The number of particles
    /// \param init The functor used to initialize the particles
    /// \param move The functor used to move the particles and weights
    /// \param scheme The resampling scheme. See ResampleScheme
    /// \param threshold The threshold for performing resampling
    /// \param seed The seed to the parallel RNG system
    explicit Sampler (
            typename Particle<T>::size_type N,
            const initialize_type &init = NULL,
            const move_type &move = NULL,
            ResampleScheme scheme = STRATIFIED,
            double threshold = 0.5,
            typename Particle<T>::seed_type seed = V_SMC_CRNG_SEED) :
        init_(init), move_(move),
        scheme_(scheme), threshold_(threshold),
        particle_(N, seed), iter_num_(0) {}

    /// \brief Size of the particle set
    ///
    /// \return The number of particles
    typename Particle<T>::size_type size () const
    {
        return particle_.size();
    }

    /// \brief The number of iterations recorded
    ///
    /// \return The number of iterations recorded so far (including the
    /// initialization)
    unsigned iter_size () const
    {
        return ess_.size();
    }

    /// \brief Get the current resampling scheme
    ///
    /// \return The current resampling scheme
    ResampleScheme resample_scheme () const
    {
        return scheme_;
    }

    /// \brief Set new resampling scheme
    ///
    /// \param scheme The new scheme for resampling
    void resample_scheme (ResampleScheme scheme)
    {
        scheme_ = scheme;
    }

    /// \brief Get the current threshold
    ///
    /// \return The current threshold for resmapling
    double resample_threshold () const
    {
        return threshold_;
    }

    /// \brief Set new resampling threshold
    ///
    /// \param threshold The new threshold for resampling
    void resample_threshold (double threshold)
    {
        threshold_ = threshold;
    }

    /// \brief ESS history
    ///
    /// \return A const reference to the history of ESS
    const ess_type &ess () const
    {
        return ess_;
    }

    /// \brief Resampling history
    ///
    /// \return A const reference to the history of resampling
    const resampled_type &resampled () const
    {
        return resampled_;
    }

    /// \brief Accept count history
    ///
    /// \return A const reference to the history of accept count
    const accept_type &accept () const
    {
        return accept_;
    }

    /// \brief Read and write access to the particle set
    ///
    /// \return A reference to the latest particle set
    Particle<T> &particle ()
    {
        return particle_;
    }

    /// \brief Read only access to the particle set
    ///
    /// \return A const reference to the latest particle set.
    const Particle<T> &particle () const
    {
        return particle_;
    }

    /// \brief Replace initialization functor
    ///
    /// \param new_init New Initialization functor
    void init (const initialize_type &new_init)
    {
        init_ = new_init;
    }

    /// \brief Replace iteration functor
    ///
    /// \param new_move New Move functor
    void move (const move_type &new_move)
    {
        move_ = new_move;
    }

    /// \brief Replace iteration functor
    ///
    /// \param new_mcmc New MCMC Move functor
    void mcmc (const move_type &new_mcmc)
    {
        mcmc_.push_back(new_mcmc);
    }

    /// \brief Clear all MCMC moves
    void clear_mcmc ()
    {
        mcmc_.clear();
    }

    /// \brief Initialize the particle set
    ///
    /// \param param Additional parameters passed to the initialization
    /// functor
    void initialize (void *param = NULL)
    {
        ess_.clear();
        resampled_.clear();
        accept_.clear();
        path_.clear();

        for (typename monitor_map_type::iterator
                m = monitor_.begin(); m != monitor_.end(); ++m)
            m->second.clear();

        iter_num_ = 0;
        if (bool(init_)) {
            accept_.push_back(
                    std::vector<unsigned>(1, init_(particle_, param)));
        }
#ifndef NDEBUG
        else {
            std::cerr << "vSMC Warning:" << std::endl;
            std::cerr << "\tSampler::initiliaize" << std::endl;
            std::cerr
                << "\tAttempt Initialization without a callable object"
                << std::endl;
        }
#endif
        post_move();
        particle_.reset_zconst();
    }

    /// \brief Perform iteration
    void iterate ()
    {
        ++iter_num_;
        std::vector<unsigned> acc;
        if (bool(move_)) {
            acc.push_back(move_(iter_num_, particle_));
        }
#ifndef V_SMC_NDEBUG
        else {
            std::cerr << "vSMC Warning:" << std::endl;
            std::cerr << "\tSampler::iterate" << std::endl;
            std::cerr
                << "\tAttempt Move without a callable object"
                << std::endl;
        }
#endif // V_SMC_NDEBUG
        for (typename std::vector<move_type>::iterator
                m = mcmc_.begin(); m != mcmc_.end(); ++m) {
            if (bool(*m)) {
                acc.push_back((*m)(iter_num_, particle_));
            }
#ifndef V_SMC_NDEBUG
            else {
                std::cerr << "vSMC Warning:" << std::endl;
                std::cerr << "\tSampler::iterate" << std::endl;
                std::cerr
                    << "\tAttempt MCMC without a callable object"
                    << std::endl;
            }
#endif // V_SMC_NDEBUG
        }
        accept_.push_back(acc);
        post_move();
    }

    /// \brief Perform iteration
    ///
    /// \param n The number of iterations to be performed
    void iterate (unsigned n)
    {
        for (unsigned i = 0; i != n; ++i)
            iterate();
    }

    /// \brief Add a monitor with an evaluation functor
    ///
    /// \param name The name of the monitor
    /// \param dim The dimension of the monitor, i.e., the number of variables
    /// \param eval The functor used to directly evaluate the results
    /// \param direct Whether or not eval return the integrands or the final
    void monitor (const std::string &name, unsigned dim,
            const typename Monitor<T>::eval_type &eval, bool direct = false)
    {
        monitor_.insert(std::make_pair(name, Monitor<T>(dim, eval, direct)));
        monitor_name_.insert(name);
    }

    /// \brief Read only access to a named monitor through iterator
    ///
    /// \param name The name of the monitor
    ///
    /// \return An const_iterator point to the monitor for the given name
    typename monitor_map_type::const_iterator monitor (
            const std::string &name) const
    {
        return monitor_.find(name);
    }

    /// \brief Read and write access to a named monitor through iterator
    ///
    /// \param name The name of the monitor
    ///
    /// \return An iterator point to the monitor for the given name
    typename monitor_map_type::iterator monitor (const std::string &name)
    {
        return monitor_.find(name);
    }

    /// \brief Read only access to all monitors
    ///
    /// \return A const reference to monitors
    const monitor_map_type &monitor () const
    {
        return monitor_;
    }

    /// \brief Read and write access to all monitors
    ///
    /// \return A reference to monitors
    monitor_map_type &monitor ()
    {
        return monitor_;
    }

    /// \brief Erase a named monitor
    ///
    /// \param name The name of the monitor
    void clear_monitor (const std::string &name)
    {
        monitor_.erase(name);
        monitor_name_.erase(name);
    }

    /// \brief Erase all monitors
    void clear_monitor ()
    {
        monitor_.clear();
        monitor_name_.clear();
    }

    /// \brief Read only access to the Path sampling monitor
    ///
    /// \return A const reference to the Path sampling monitor
    const Path<T> &path () const
    {
        return path_;
    }

    /// \brief Set the path sampling evaluation functor
    ///
    /// \param eval The functor used to compute the integrands or results
    /// \param direct Whether or not eval return the integrands or the final
    /// results
    void path_sampling (const typename Path<T>::eval_type &eval,
            bool direct = false)
    {
        path_.eval(eval, direct);
    }

    /// \brief Path sampling estimate of normalizing constant
    ///
    /// \return The log ratio of normalizing constants
    double path_sampling () const
    {
        return path_.zconst();
    }

    /// \brief SMC estimate of normalizing constant
    ///
    /// \return The log of SMC normalizng constant estimate
    double zconst () const
    {
        return particle_.zconst();
    }

    /// \brief Print the history of the sampler
    ///
    /// \param os The ostream to which the contents are printed
    /// \param print_header Print header if \b true
    template<typename CharT, typename Traits>
    void print (std::basic_ostream<CharT, Traits> &os = std::cout,
            bool print_header = true) const
    {
        const char sep = '\t';

        // Accept count
        unsigned accd = 0;
        for (accept_type::const_iterator
                a = accept_.begin(); a != accept_.end(); ++a) {
            if (a->size() > accd)
                accd = a->size();
        }
        Eigen::MatrixXd acc(iter_size(), accd);
        acc.setConstant(0);
        double accdnorm = static_cast<double>(size());
        for (unsigned r = 0; r != iter_size(); ++r)
            for (unsigned c = 0; c != accept_[r].size(); ++c)
                acc(r, c) = accept_[r][c] / accdnorm;

        // Path sampling
        Eigen::MatrixXd pa(iter_size(), 3);
        Eigen::MatrixXi pmask(iter_size(), 1);
        pmask.setConstant(0);
        for (unsigned d = 0; d != path_.iter_size(); ++d) {
            unsigned pr = path_.index()[d];
            pa(pr, 0) = path_.integrand()[d];
            pa(pr, 1) = path_.width()[d];
            pa(pr, 2) = path_.grid()[d];
            pmask(pr) = 1;
        }

        // Monitors
        unsigned mond = 0;
        for (typename monitor_map_type::const_iterator
                m = monitor_.begin(); m != monitor_.end(); ++m) {
            mond += m->second.dim();
        }
        Eigen::MatrixXd mon(iter_size(), mond);
        Eigen::MatrixXi mmask(iter_size(), monitor_.size());
        mmask.setConstant(0);
        unsigned mc = 0;
        unsigned mn = 0;
        for (typename monitor_map_type::const_iterator
                m = monitor_.begin(); m != monitor_.end(); ++m) {
            unsigned md = m->second.dim();
            for (unsigned d = 0; d != m->second.iter_size(); ++d) {
                unsigned mr = m->second.index()[d];
                for (unsigned c = 0; c != md; ++c)
                    mon(mr, c + mc) = m->second.record(c)[d];
                mmask(mr, mn) = 1;
            }
            mc += md;
            ++mn;
        }

        // Print header
        if (print_header) {
            os << "Iter" << sep << "ESS" << sep << "ResSam" << sep;
            if (accd == 1) {
                os << "Accept" << sep;
            } else {
                for (unsigned d = 0; d != accd; ++d)
                    os << "Accept." << d + 1 << sep;
            }
            os
                << "Path.Integrand" << sep
                << "Path.Width" << sep
                << "Path.Grid" << sep << "";
            for (typename monitor_map_type::const_iterator
                    m = monitor_.begin(); m != monitor_.end(); ++m) {
                if (m->second.dim() == 1) {
                    os << m->first << sep;
                } else {
                    for (unsigned d = 0; d != m->second.dim(); ++d)
                        os << m->first << '.' << d + 1 << sep;
                }
            }
            os << '\n';
        }

        // Print data
        for (unsigned iter = 0; iter != iter_size(); ++iter) {
            os << iter << sep;
            os << ess_[iter] / size() << sep;
            os << resampled_[iter] << sep;
            os << acc.row(iter) << sep;
            if (pmask(iter))
                os << pa.row(iter) << sep;
            else
                os << '.' << sep << '.' << sep << '.' << sep;
            unsigned mc = 0;
            unsigned mn = 0;
            for (typename monitor_map_type::const_iterator
                    m = monitor_.begin(); m != monitor_.end(); ++m) {
                unsigned md = m->second.dim();
                if (mmask(iter, mn)) {
                    os << mon.block(iter, mc, 1, md) << sep;
                } else {
                    for (unsigned m = 0; m != md; ++m)
                        os << '.' << sep;
                }
                mc += md;
                ++mn;
            }
            os << '\n';
        }
    }

    private :

    /// Initialization and movement
    initialize_type init_;
    move_type move_;
    std::vector<move_type> mcmc_;

    /// Resampling
    ResampleScheme scheme_;
    double threshold_;

    /// Particle sets
    Particle<T> particle_;
    unsigned iter_num_;
    ess_type ess_;
    resampled_type resampled_;
    accept_type accept_;

    /// Monte Carlo estimation by integration
    monitor_map_type monitor_;
    std::set<std::string> monitor_name_;

    /// Path sampling
    Path<T> path_;

    void post_move ()
    {
        bool do_resample = particle_.ess() < threshold_ * size();
        if (do_resample)
            particle_.resample(scheme_);

        if (!path_.empty())
            path_.eval(iter_num_, particle_);

        for (typename monitor_map_type::iterator
                m = monitor_.begin(); m != monitor_.end(); ++m) {
            if (!m->second.empty())
                m->second.eval(iter_num_, particle_);
        }

        ess_.push_back(particle_.ess());
        resampled_.push_back(do_resample);
        particle_.resampled(resampled_.back());
        particle_.accept(accept_.back().back());
    }
}; // class Sampler

} // namespace vSMC

namespace std {

/// \brief Print the sampler
///
/// \param os The ostream to which the contents are printed
/// \param sampler The sampler to be printed
///
/// \note This is the same as <tt>sampler.print(os)</tt>
template<typename T, typename CharT, typename Traits>
std::basic_ostream<CharT, Traits> &operator<< (
        std::basic_ostream<CharT, Traits> &os, const vSMC::Sampler<T> &sampler)
{
    sampler.print(os, true);
    return os;
}

} // namespace std

#endif // V_SMC_CORE_SAMPLER_HPP
<commit_msg>check edge cases of print<commit_after>#ifndef V_SMC_CORE_SAMPLER_HPP
#define V_SMC_CORE_SAMPLER_HPP

#include <vSMC/internal/common.hpp>
#include <ostream>
#include <map>
#include <set>
#include <string>

namespace vSMC {

/// \brief SMC Sampler
///
/// \tparam T State state type. Requiment:
/// \li Consturctor: T (IntType N)
/// \li Method: copy (IntType from, IntType to)
template <typename T>
class Sampler
{
    public :

    /// The type of initialization functor
    typedef internal::function<unsigned (Particle<T> &, void *)>
        initialize_type;

    /// The type of move and mcmc functor
    typedef internal::function<unsigned (unsigned, Particle<T> &)>
        move_type;

    /// The type of ESS history vector
    typedef std::vector<double> ess_type;

    /// The type of resampling history vector
    typedef std::vector<bool> resampled_type;

    /// The type of accept count history vector
    typedef std::vector<std::vector<unsigned> > accept_type;

    /// The type of Monitor map
    typedef std::map<std::string, Monitor<T> > monitor_map_type;

    /// \brief Construct a sampler with given number of particles
    ///
    /// \param N The number of particles
    /// \param init The functor used to initialize the particles
    /// \param move The functor used to move the particles and weights
    /// \param scheme The resampling scheme. See ResampleScheme
    /// \param threshold The threshold for performing resampling
    /// \param seed The seed to the parallel RNG system
    explicit Sampler (
            typename Particle<T>::size_type N,
            const initialize_type &init = NULL,
            const move_type &move = NULL,
            ResampleScheme scheme = STRATIFIED,
            double threshold = 0.5,
            typename Particle<T>::seed_type seed = V_SMC_CRNG_SEED) :
        init_(init), move_(move),
        scheme_(scheme), threshold_(threshold),
        particle_(N, seed), iter_num_(0) {}

    /// \brief Size of the particle set
    ///
    /// \return The number of particles
    typename Particle<T>::size_type size () const
    {
        return particle_.size();
    }

    /// \brief The number of iterations recorded
    ///
    /// \return The number of iterations recorded so far (including the
    /// initialization)
    unsigned iter_size () const
    {
        return ess_.size();
    }

    /// \brief Get the current resampling scheme
    ///
    /// \return The current resampling scheme
    ResampleScheme resample_scheme () const
    {
        return scheme_;
    }

    /// \brief Set new resampling scheme
    ///
    /// \param scheme The new scheme for resampling
    void resample_scheme (ResampleScheme scheme)
    {
        scheme_ = scheme;
    }

    /// \brief Get the current threshold
    ///
    /// \return The current threshold for resmapling
    double resample_threshold () const
    {
        return threshold_;
    }

    /// \brief Set new resampling threshold
    ///
    /// \param threshold The new threshold for resampling
    void resample_threshold (double threshold)
    {
        threshold_ = threshold;
    }

    /// \brief ESS history
    ///
    /// \return A const reference to the history of ESS
    const ess_type &ess () const
    {
        return ess_;
    }

    /// \brief Resampling history
    ///
    /// \return A const reference to the history of resampling
    const resampled_type &resampled () const
    {
        return resampled_;
    }

    /// \brief Accept count history
    ///
    /// \return A const reference to the history of accept count
    const accept_type &accept () const
    {
        return accept_;
    }

    /// \brief Read and write access to the particle set
    ///
    /// \return A reference to the latest particle set
    Particle<T> &particle ()
    {
        return particle_;
    }

    /// \brief Read only access to the particle set
    ///
    /// \return A const reference to the latest particle set.
    const Particle<T> &particle () const
    {
        return particle_;
    }

    /// \brief Replace initialization functor
    ///
    /// \param new_init New Initialization functor
    void init (const initialize_type &new_init)
    {
        init_ = new_init;
    }

    /// \brief Replace iteration functor
    ///
    /// \param new_move New Move functor
    void move (const move_type &new_move)
    {
        move_ = new_move;
    }

    /// \brief Replace iteration functor
    ///
    /// \param new_mcmc New MCMC Move functor
    void mcmc (const move_type &new_mcmc)
    {
        mcmc_.push_back(new_mcmc);
    }

    /// \brief Clear all MCMC moves
    void clear_mcmc ()
    {
        mcmc_.clear();
    }

    /// \brief Initialize the particle set
    ///
    /// \param param Additional parameters passed to the initialization
    /// functor
    void initialize (void *param = NULL)
    {
        ess_.clear();
        resampled_.clear();
        accept_.clear();
        path_.clear();

        for (typename monitor_map_type::iterator
                m = monitor_.begin(); m != monitor_.end(); ++m)
            m->second.clear();

        iter_num_ = 0;
        if (bool(init_)) {
            accept_.push_back(
                    std::vector<unsigned>(1, init_(particle_, param)));
        }
#ifndef NDEBUG
        else {
            std::cerr << "vSMC Warning:" << std::endl;
            std::cerr << "\tSampler::initiliaize" << std::endl;
            std::cerr
                << "\tAttempt Initialization without a callable object"
                << std::endl;
        }
#endif
        post_move();
        particle_.reset_zconst();
    }

    /// \brief Perform iteration
    void iterate ()
    {
        ++iter_num_;
        std::vector<unsigned> acc;
        if (bool(move_)) {
            acc.push_back(move_(iter_num_, particle_));
        }
#ifndef V_SMC_NDEBUG
        else {
            std::cerr << "vSMC Warning:" << std::endl;
            std::cerr << "\tSampler::iterate" << std::endl;
            std::cerr
                << "\tAttempt Move without a callable object"
                << std::endl;
        }
#endif // V_SMC_NDEBUG
        for (typename std::vector<move_type>::iterator
                m = mcmc_.begin(); m != mcmc_.end(); ++m) {
            if (bool(*m)) {
                acc.push_back((*m)(iter_num_, particle_));
            }
#ifndef V_SMC_NDEBUG
            else {
                std::cerr << "vSMC Warning:" << std::endl;
                std::cerr << "\tSampler::iterate" << std::endl;
                std::cerr
                    << "\tAttempt MCMC without a callable object"
                    << std::endl;
            }
#endif // V_SMC_NDEBUG
        }
        accept_.push_back(acc);
        post_move();
    }

    /// \brief Perform iteration
    ///
    /// \param n The number of iterations to be performed
    void iterate (unsigned n)
    {
        for (unsigned i = 0; i != n; ++i)
            iterate();
    }

    /// \brief Add a monitor with an evaluation functor
    ///
    /// \param name The name of the monitor
    /// \param dim The dimension of the monitor, i.e., the number of variables
    /// \param eval The functor used to directly evaluate the results
    /// \param direct Whether or not eval return the integrands or the final
    void monitor (const std::string &name, unsigned dim,
            const typename Monitor<T>::eval_type &eval, bool direct = false)
    {
        monitor_.insert(std::make_pair(name, Monitor<T>(dim, eval, direct)));
        monitor_name_.insert(name);
    }

    /// \brief Read only access to a named monitor through iterator
    ///
    /// \param name The name of the monitor
    ///
    /// \return An const_iterator point to the monitor for the given name
    typename monitor_map_type::const_iterator monitor (
            const std::string &name) const
    {
        return monitor_.find(name);
    }

    /// \brief Read and write access to a named monitor through iterator
    ///
    /// \param name The name of the monitor
    ///
    /// \return An iterator point to the monitor for the given name
    typename monitor_map_type::iterator monitor (const std::string &name)
    {
        return monitor_.find(name);
    }

    /// \brief Read only access to all monitors
    ///
    /// \return A const reference to monitors
    const monitor_map_type &monitor () const
    {
        return monitor_;
    }

    /// \brief Read and write access to all monitors
    ///
    /// \return A reference to monitors
    monitor_map_type &monitor ()
    {
        return monitor_;
    }

    /// \brief Erase a named monitor
    ///
    /// \param name The name of the monitor
    void clear_monitor (const std::string &name)
    {
        monitor_.erase(name);
        monitor_name_.erase(name);
    }

    /// \brief Erase all monitors
    void clear_monitor ()
    {
        monitor_.clear();
        monitor_name_.clear();
    }

    /// \brief Read only access to the Path sampling monitor
    ///
    /// \return A const reference to the Path sampling monitor
    const Path<T> &path () const
    {
        return path_;
    }

    /// \brief Set the path sampling evaluation functor
    ///
    /// \param eval The functor used to compute the integrands or results
    /// \param direct Whether or not eval return the integrands or the final
    /// results
    void path_sampling (const typename Path<T>::eval_type &eval,
            bool direct = false)
    {
        path_.eval(eval, direct);
    }

    /// \brief Path sampling estimate of normalizing constant
    ///
    /// \return The log ratio of normalizing constants
    double path_sampling () const
    {
        return path_.zconst();
    }

    /// \brief SMC estimate of normalizing constant
    ///
    /// \return The log of SMC normalizng constant estimate
    double zconst () const
    {
        return particle_.zconst();
    }

    /// \brief Print the history of the sampler
    ///
    /// \param os The ostream to which the contents are printed
    /// \param print_header Print header if \b true
    template<typename CharT, typename Traits>
    void print (std::basic_ostream<CharT, Traits> &os = std::cout,
            bool print_header = true) const
    {
        const char sep = '\t';

        // Accept count
        Eigen::MatrixXd acc;
        unsigned accd = 0;
        for (accept_type::const_iterator
                a = accept_.begin(); a != accept_.end(); ++a) {
            if (a->size() > accd)
                accd = a->size();
        }
        bool print_accept = accd > 0 && iter_size() > 0;
        if (print_accept) {
            acc.resize(iter_size(), accd);
            acc.setConstant(0);
            double accdnorm = static_cast<double>(size());
            for (unsigned r = 0; r != iter_size(); ++r)
                for (unsigned c = 0; c != accept_[r].size(); ++c)
                    acc(r, c) = accept_[r][c] / accdnorm;
        }

        // Path sampling
        Eigen::MatrixXd pa;
        Eigen::MatrixXi pmask;
        bool print_path = path_.iter_size() > 0 && iter_size() > 0;
        if (print_path) {
            pa.resize(iter_size(), 3);
            pmask.resize(iter_size(), 1);
            pmask.setConstant(0);
            for (unsigned d = 0; d != path_.iter_size(); ++d) {
                unsigned pr = path_.index()[d];
                pa(pr, 0) = path_.integrand()[d];
                pa(pr, 1) = path_.width()[d];
                pa(pr, 2) = path_.grid()[d];
                pmask(pr) = 1;
            }
        }

        // Monitors
        Eigen::MatrixXd mon;
        Eigen::MatrixXi mmask;
        unsigned mond = 0;
        for (typename monitor_map_type::const_iterator
                m = monitor_.begin(); m != monitor_.end(); ++m) {
            mond += m->second.dim();
        }
        bool print_monitor = mond > 0 && iter_size() > 0;
        if (print_monitor) {
            mon.resize(iter_size(), mond);
            mmask.resize(iter_size(), monitor_.size());
            mmask.setConstant(0);
            unsigned mc = 0;
            unsigned mn = 0;
            for (typename monitor_map_type::const_iterator
                    m = monitor_.begin(); m != monitor_.end(); ++m) {
                unsigned md = m->second.dim();
                for (unsigned d = 0; d != m->second.iter_size(); ++d) {
                    unsigned mr = m->second.index()[d];
                    for (unsigned c = 0; c != md; ++c)
                        mon(mr, c + mc) = m->second.record(c)[d];
                    mmask(mr, mn) = 1;
                }
                mc += md;
                ++mn;
            }
        }

        // Print header
        if (print_header) {
            os << "Iter" << sep << "ESS" << sep << "ResSam" << sep;
            if (print_accept) {
                if (accd == 1) {
                    os << "Accept" << sep;
                } else {
                    for (unsigned d = 0; d != accd; ++d)
                        os << "Accept." << d + 1 << sep;
                }
            }
            if (print_path) {
                os
                    << "Path.Integrand" << sep
                    << "Path.Width" << sep
                    << "Path.Grid" << sep << "";
            }
            if (print_monitor) {
                for (typename monitor_map_type::const_iterator
                        m = monitor_.begin(); m != monitor_.end(); ++m) {
                    if (m->second.dim() == 1) {
                        os << m->first << sep;
                    } else {
                        for (unsigned d = 0; d != m->second.dim(); ++d)
                            os << m->first << '.' << d + 1 << sep;
                    }
                }
            }
            os << '\n';
        }

        // Print data
        for (unsigned iter = 0; iter != iter_size(); ++iter) {
            os << iter << sep;
            os << ess_[iter] / size() << sep;
            os << resampled_[iter] << sep;
            if (print_accept)
                os << acc.row(iter) << sep;
            if (print_path) {
                if (pmask(iter))
                    os << pa.row(iter) << sep;
                else
                    os << '.' << sep << '.' << sep << '.' << sep;
            }
            if (print_monitor) {
                unsigned mc = 0;
                unsigned mn = 0;
                for (typename monitor_map_type::const_iterator
                        m = monitor_.begin(); m != monitor_.end(); ++m) {
                    unsigned md = m->second.dim();
                    if (mmask(iter, mn)) {
                        os << mon.block(iter, mc, 1, md) << sep;
                    } else {
                        for (unsigned m = 0; m != md; ++m)
                            os << '.' << sep;
                    }
                    mc += md;
                    ++mn;
                }
            }
            os << '\n';
        }
    }

    private :

    /// Initialization and movement
    initialize_type init_;
    move_type move_;
    std::vector<move_type> mcmc_;

    /// Resampling
    ResampleScheme scheme_;
    double threshold_;

    /// Particle sets
    Particle<T> particle_;
    unsigned iter_num_;
    ess_type ess_;
    resampled_type resampled_;
    accept_type accept_;

    /// Monte Carlo estimation by integration
    monitor_map_type monitor_;
    std::set<std::string> monitor_name_;

    /// Path sampling
    Path<T> path_;

    void post_move ()
    {
        bool do_resample = particle_.ess() < threshold_ * size();
        if (do_resample)
            particle_.resample(scheme_);

        if (!path_.empty())
            path_.eval(iter_num_, particle_);

        for (typename monitor_map_type::iterator
                m = monitor_.begin(); m != monitor_.end(); ++m) {
            if (!m->second.empty())
                m->second.eval(iter_num_, particle_);
        }

        ess_.push_back(particle_.ess());
        resampled_.push_back(do_resample);
        particle_.resampled(resampled_.back());
        particle_.accept(accept_.back().back());
    }
}; // class Sampler

} // namespace vSMC

namespace std {

/// \brief Print the sampler
///
/// \param os The ostream to which the contents are printed
/// \param sampler The sampler to be printed
///
/// \note This is the same as <tt>sampler.print(os)</tt>
template<typename T, typename CharT, typename Traits>
std::basic_ostream<CharT, Traits> &operator<< (
        std::basic_ostream<CharT, Traits> &os, const vSMC::Sampler<T> &sampler)
{
    sampler.print(os, true);
    return os;
}

} // namespace std

#endif // V_SMC_CORE_SAMPLER_HPP
<|endoftext|>


<commit_before>/*=========================================================================
 *
 *  Copyright RTK Consortium
 *
 *  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
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  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 "rtkConjugateGradientConeBeamReconstructionFilter.h"

namespace rtk
{

template<>
ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
::ConjugateGradientConeBeamReconstructionFilter()
{
  this->SetNumberOfRequiredInputs(3);

  // Set the default values of member parameters
  m_NumberOfIterations=3;
  m_MeasureExecutionTimes=false;
//  m_IterationCosts=false;

  m_Gamma = 0;
  m_Regularized = false;
  m_CudaConjugateGradient = true;
  m_DisableDisplacedDetectorFilter = false;

  // Create the filters
#ifdef RTK_USE_CUDA
  m_DisplacedDetectorFilter = rtk::CudaDisplacedDetectorImageFilter::New();
  m_ConstantVolumeSource     = rtk::CudaConstantVolumeSource::New();
#else
  m_DisplacedDetectorFilter = DisplacedDetectorFilterType::New();
  m_ConstantVolumeSource     = ConstantImageSourceType::New();
#endif
  m_CGOperator = CGOperatorFilterType::New();

  m_MultiplyVolumeFilter = MultiplyFilterType::New();
  m_MatrixVectorMultiplyFilter = MatrixVectorMultiplyFilterType::New();
  m_MultiplyOutputFilter = MultiplyFilterType::New();

  // Set permanent parameters
  m_ConstantVolumeSource->SetConstant(0);
  m_DisplacedDetectorFilter->SetPadOnTruncatedSide(false);
}

//template<>
//void
//ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
//::SetSupportMask(const itk::Image<float, 3> *SupportMask)
//{
//  this->SetInput("SupportMask", const_cast<itk::Image<float, 3>*>(SupportMask));
//}

//template<>
//typename itk::Image<float, 3>::ConstPointer
//ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
//::GetSupportMask()
//{
//  return static_cast< const itk::Image<float, 3> * >
//          ( this->itk::ProcessObject::GetInput("SupportMask") );
//}

template<>
void
ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
::GenerateOutputInformation()
{
  // Choose between cuda or non-cuda conjugate gradient filter
  m_ConjugateGradientFilter = ConjugateGradientFilterType::New();
#ifdef RTK_USE_CUDA
  if (m_CudaConjugateGradient)
    m_ConjugateGradientFilter = rtk::CudaConjugateGradientImageFilter_3f::New();
#endif
  m_ConjugateGradientFilter->SetA(m_CGOperator.GetPointer());
//  m_ConjugateGradientFilter->SetIterationCosts(m_IterationCosts);

  // Set runtime connections
  m_ConstantVolumeSource->SetInformationFromImage(this->GetInput(0));
  m_CGOperator->SetInput(1, this->GetInput(1));
  m_CGOperator->SetSupportMask(this->GetSupportMask());
  m_ConjugateGradientFilter->SetX(this->GetInput(0));
  m_DisplacedDetectorFilter->SetDisable(m_DisableDisplacedDetectorFilter);
  m_DisplacedDetectorFilter->SetInput(this->GetInput(2));

  // Links with the m_BackProjectionFilter should be set here and not
  // in the constructor, as m_BackProjectionFilter is set at runtime
  m_BackProjectionFilterForB->SetInput(0, m_ConstantVolumeSource->GetOutput());
  m_ConjugateGradientFilter->SetB(m_BackProjectionFilterForB->GetOutput());

  // Set the matrix vector multiply filter's inputs for multiplication
  // by the inverse covariance matrix (for GLS minimization)
  m_MatrixVectorMultiplyFilter->SetInput1(this->GetInput(1));
  m_MatrixVectorMultiplyFilter->SetInput2(m_DisplacedDetectorFilter->GetOutput());
  m_CGOperator->SetInput(2, m_DisplacedDetectorFilter->GetOutput());
  m_BackProjectionFilterForB->SetInput(1, m_MatrixVectorMultiplyFilter->GetOutput());

  // If a support mask is used, it serves as preconditioning weights
  if (this->GetSupportMask().IsNotNull())
    {
    // Multiply the volume by support mask, and pass it to the conjugate gradient operator
    m_MultiplyVolumeFilter->SetInput1(m_BackProjectionFilterForB->GetOutput());
    m_MultiplyVolumeFilter->SetInput2(this->GetSupportMask());
    m_CGOperator->SetSupportMask(this->GetSupportMask());
    m_ConjugateGradientFilter->SetB(m_MultiplyVolumeFilter->GetOutput());

    // Multiply the output by the support mask
    m_MultiplyOutputFilter->SetInput1(m_ConjugateGradientFilter->GetOutput());
    m_MultiplyOutputFilter->SetInput2(this->GetSupportMask());
    }

  // For the same reason, set geometry now
  m_CGOperator->SetGeometry(this->m_Geometry);
  m_BackProjectionFilterForB->SetGeometry(this->m_Geometry.GetPointer());
  m_DisplacedDetectorFilter->SetGeometry(this->m_Geometry);

  // Set runtime parameters
  m_ConjugateGradientFilter->SetNumberOfIterations(this->m_NumberOfIterations);
  m_CGOperator->SetGamma(m_Gamma);

  // Set memory management parameters
  m_MatrixVectorMultiplyFilter->ReleaseDataFlagOn();
  m_BackProjectionFilterForB->ReleaseDataFlagOn();
  if (this->GetSupportMask().IsNotNull())
    {
    m_MultiplyVolumeFilter->ReleaseDataFlagOn();
    m_MultiplyOutputFilter->ReleaseDataFlagOn();
    }

  // Have the last filter calculate its output information
  m_ConjugateGradientFilter->UpdateOutputInformation();

  // Copy it as the output information of the composite filter
  this->GetOutput()->CopyInformation( m_ConjugateGradientFilter->GetOutput() );
}

template<>
void
ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
::GenerateData()
{
  itk::TimeProbe ConjugateGradientTimeProbe;
//  typename StatisticsImageFilterType::Pointer StatisticsImageFilterForC = StatisticsImageFilterType::New();
//  typename MultiplyFilterType::Pointer MultiplyFilterForC = MultiplyFilterType::New();

//  if (m_IterationCosts)
//    {
//    MultiplyFilterForC->SetInput(0,this->GetInput(1));
//    MultiplyFilterForC->SetInput(1,this->GetInput(2));
//    MultiplyFilterForC->Update();
//    MultiplyFilterForC->SetInput(1,MultiplyFilterForC->GetOutput());
//    MultiplyFilterForC->Update();
//    StatisticsImageFilterForC->SetInput(MultiplyFilterForC->GetOutput());
//    StatisticsImageFilterForC->Update();
//    m_ConjugateGradientFilter->SetC(0.5*StatisticsImageFilterForC->GetSum());
//    }

  if(m_MeasureExecutionTimes)
    {
    std::cout << "Starting ConjugateGradient" << std::endl;
    ConjugateGradientTimeProbe.Start();
    }

  m_ConjugateGradientFilter->Update();

  if (this->GetSupportMask())
    {
    m_MultiplyOutputFilter->Update();
    }

  if(m_MeasureExecutionTimes)
    {
    ConjugateGradientTimeProbe.Stop();
    std::cout << "ConjugateGradient took " << ConjugateGradientTimeProbe.GetTotal() << ' ' << ConjugateGradientTimeProbe.GetUnit() << std::endl;
    }

  if (this->GetSupportMask())
    {
    this->GraftOutput( m_MultiplyOutputFilter->GetOutput() );
    }
    else
    {
    this->GraftOutput( m_ConjugateGradientFilter->GetOutput() );
    }
}

} // end namespace rtk
<commit_msg>Pass tikhonov regularization parameter to rtkReconstructionConjugateGradientOperator<commit_after>/*=========================================================================
 *
 *  Copyright RTK Consortium
 *
 *  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
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  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 "rtkConjugateGradientConeBeamReconstructionFilter.h"

namespace rtk
{

template<>
ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
::ConjugateGradientConeBeamReconstructionFilter()
{
  this->SetNumberOfRequiredInputs(3);

  // Set the default values of member parameters
  m_NumberOfIterations=3;
  m_MeasureExecutionTimes=false;
//  m_IterationCosts=false;

  m_Gamma = 0;
  m_Regularized = false;
  m_CudaConjugateGradient = true;
  m_DisableDisplacedDetectorFilter = false;

  // Create the filters
#ifdef RTK_USE_CUDA
  m_DisplacedDetectorFilter = rtk::CudaDisplacedDetectorImageFilter::New();
  m_ConstantVolumeSource     = rtk::CudaConstantVolumeSource::New();
#else
  m_DisplacedDetectorFilter = DisplacedDetectorFilterType::New();
  m_ConstantVolumeSource     = ConstantImageSourceType::New();
#endif
  m_CGOperator = CGOperatorFilterType::New();

  m_MultiplyVolumeFilter = MultiplyFilterType::New();
  m_MatrixVectorMultiplyFilter = MatrixVectorMultiplyFilterType::New();
  m_MultiplyOutputFilter = MultiplyFilterType::New();

  // Set permanent parameters
  m_ConstantVolumeSource->SetConstant(0);
  m_DisplacedDetectorFilter->SetPadOnTruncatedSide(false);
}

//template<>
//void
//ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
//::SetSupportMask(const itk::Image<float, 3> *SupportMask)
//{
//  this->SetInput("SupportMask", const_cast<itk::Image<float, 3>*>(SupportMask));
//}

//template<>
//typename itk::Image<float, 3>::ConstPointer
//ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
//::GetSupportMask()
//{
//  return static_cast< const itk::Image<float, 3> * >
//          ( this->itk::ProcessObject::GetInput("SupportMask") );
//}

template<>
void
ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
::GenerateOutputInformation()
{
  // Choose between cuda or non-cuda conjugate gradient filter
  m_ConjugateGradientFilter = ConjugateGradientFilterType::New();
#ifdef RTK_USE_CUDA
  if (m_CudaConjugateGradient)
    m_ConjugateGradientFilter = rtk::CudaConjugateGradientImageFilter_3f::New();
#endif
  m_ConjugateGradientFilter->SetA(m_CGOperator.GetPointer());
//  m_ConjugateGradientFilter->SetIterationCosts(m_IterationCosts);

  // Set runtime connections
  m_ConstantVolumeSource->SetInformationFromImage(this->GetInput(0));
  m_CGOperator->SetInput(1, this->GetInput(1));
  m_CGOperator->SetSupportMask(this->GetSupportMask());
  m_ConjugateGradientFilter->SetX(this->GetInput(0));
  m_DisplacedDetectorFilter->SetDisable(m_DisableDisplacedDetectorFilter);
  m_DisplacedDetectorFilter->SetInput(this->GetInput(2));

  // Links with the m_BackProjectionFilter should be set here and not
  // in the constructor, as m_BackProjectionFilter is set at runtime
  m_BackProjectionFilterForB->SetInput(0, m_ConstantVolumeSource->GetOutput());
  m_ConjugateGradientFilter->SetB(m_BackProjectionFilterForB->GetOutput());

  // Set the matrix vector multiply filter's inputs for multiplication
  // by the inverse covariance matrix (for GLS minimization)
  m_MatrixVectorMultiplyFilter->SetInput1(this->GetInput(1));
  m_MatrixVectorMultiplyFilter->SetInput2(m_DisplacedDetectorFilter->GetOutput());
  m_CGOperator->SetInput(2, m_DisplacedDetectorFilter->GetOutput());
  m_BackProjectionFilterForB->SetInput(1, m_MatrixVectorMultiplyFilter->GetOutput());

  // If a support mask is used, it serves as preconditioning weights
  if (this->GetSupportMask().IsNotNull())
    {
    // Multiply the volume by support mask, and pass it to the conjugate gradient operator
    m_MultiplyVolumeFilter->SetInput1(m_BackProjectionFilterForB->GetOutput());
    m_MultiplyVolumeFilter->SetInput2(this->GetSupportMask());
    m_CGOperator->SetSupportMask(this->GetSupportMask());
    m_ConjugateGradientFilter->SetB(m_MultiplyVolumeFilter->GetOutput());

    // Multiply the output by the support mask
    m_MultiplyOutputFilter->SetInput1(m_ConjugateGradientFilter->GetOutput());
    m_MultiplyOutputFilter->SetInput2(this->GetSupportMask());
    }

  // For the same reason, set geometry now
  m_CGOperator->SetGeometry(this->m_Geometry);
  m_BackProjectionFilterForB->SetGeometry(this->m_Geometry.GetPointer());
  m_DisplacedDetectorFilter->SetGeometry(this->m_Geometry);

  // Set runtime parameters
  m_ConjugateGradientFilter->SetNumberOfIterations(this->m_NumberOfIterations);
  m_CGOperator->SetGamma(m_Gamma);
  m_CGOperator->SetTikhonov(m_Tikhonov);

  // Set memory management parameters
  m_MatrixVectorMultiplyFilter->ReleaseDataFlagOn();
  m_BackProjectionFilterForB->ReleaseDataFlagOn();
  if (this->GetSupportMask().IsNotNull())
    {
    m_MultiplyVolumeFilter->ReleaseDataFlagOn();
    m_MultiplyOutputFilter->ReleaseDataFlagOn();
    }

  // Have the last filter calculate its output information
  m_ConjugateGradientFilter->UpdateOutputInformation();

  // Copy it as the output information of the composite filter
  this->GetOutput()->CopyInformation( m_ConjugateGradientFilter->GetOutput() );
}

template<>
void
ConjugateGradientConeBeamReconstructionFilter< itk::VectorImage<float, 3>, itk::Image<float, 3> >
::GenerateData()
{
  itk::TimeProbe ConjugateGradientTimeProbe;
//  typename StatisticsImageFilterType::Pointer StatisticsImageFilterForC = StatisticsImageFilterType::New();
//  typename MultiplyFilterType::Pointer MultiplyFilterForC = MultiplyFilterType::New();

//  if (m_IterationCosts)
//    {
//    MultiplyFilterForC->SetInput(0,this->GetInput(1));
//    MultiplyFilterForC->SetInput(1,this->GetInput(2));
//    MultiplyFilterForC->Update();
//    MultiplyFilterForC->SetInput(1,MultiplyFilterForC->GetOutput());
//    MultiplyFilterForC->Update();
//    StatisticsImageFilterForC->SetInput(MultiplyFilterForC->GetOutput());
//    StatisticsImageFilterForC->Update();
//    m_ConjugateGradientFilter->SetC(0.5*StatisticsImageFilterForC->GetSum());
//    }

  if(m_MeasureExecutionTimes)
    {
    std::cout << "Starting ConjugateGradient" << std::endl;
    ConjugateGradientTimeProbe.Start();
    }

  m_ConjugateGradientFilter->Update();

  if (this->GetSupportMask())
    {
    m_MultiplyOutputFilter->Update();
    }

  if(m_MeasureExecutionTimes)
    {
    ConjugateGradientTimeProbe.Stop();
    std::cout << "ConjugateGradient took " << ConjugateGradientTimeProbe.GetTotal() << ' ' << ConjugateGradientTimeProbe.GetUnit() << std::endl;
    }

  if (this->GetSupportMask())
    {
    this->GraftOutput( m_MultiplyOutputFilter->GetOutput() );
    }
    else
    {
    this->GraftOutput( m_ConjugateGradientFilter->GetOutput() );
    }
}

} // end namespace rtk
<|endoftext|>


<commit_before>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

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

    http://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 <algorithm>
#include <memory>
#include <vector>

#include "absl/memory/memory.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_instance.pb.h"  // NOLINT
#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT
#include "third_party/tensorrt/NvInfer.h"

namespace tensorflow {
namespace tensorrt {
using ::nvinfer1::IRuntime;

class CreateTRTResourceHandle : public OpKernel {
 public:
  explicit CreateTRTResourceHandle(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("resource_name", &resource_name_));
  }

  void Compute(OpKernelContext* ctx) override {
    {
      mutex_lock l(mutex_);
      if (!initialized_) {
        AllocatorAttributes attr;
        attr.set_on_host(true);
        OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
                                               &handle_, attr));

        VLOG(1) << "Creating TRT engine cache resource handle for op "
                << resource_name_ << " on device " << ctx->device()->name();
        handle_.scalar<ResourceHandle>()() =
            MakeResourceHandle<TRTEngineCacheResource>(
                ctx, std::string(kTfTrtContainerName), resource_name_);
        initialized_ = true;
      }
    }
    ctx->set_output(0, handle_);
  }

 private:
  string resource_name_;
  Tensor handle_;
  mutex mutex_;
  bool initialized_ TF_GUARDED_BY(mutex_) = false;

  TF_DISALLOW_COPY_AND_ASSIGN(CreateTRTResourceHandle);
};

REGISTER_KERNEL_BUILDER(Name("CreateTRTResourceHandle")
                            .Device(DEVICE_GPU)
                            .HostMemory("resource_handle"),
                        CreateTRTResourceHandle);

class InitializeTRTResource : public OpKernel {
 public:
  explicit InitializeTRTResource(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(
        ctx, ctx->GetAttr("max_cached_engines_count", &max_cached_engines_));
  }

  void Compute(OpKernelContext* ctx) override {
    ResourceHandle handle = HandleFromInput(ctx, 0);
    core::RefCountPtr<TRTEngineCacheResource> resource;
    OP_REQUIRES_OK(
        ctx, LookupOrCreateResource<TRTEngineCacheResource>(
                 ctx, handle, &resource,
                 [this, ctx](TRTEngineCacheResource** resource) -> Status {
                   *resource = new TRTEngineCacheResource(
                       ctx, this->max_cached_engines_);
                   return Status::OK();
                 }));

    auto allocator = resource->allocator_.get();
    OP_REQUIRES(ctx, allocator != nullptr,
                errors::Internal("Not able to initialize TRT engine cache when "
                                 "GPU allocator is empty."));
    OP_REQUIRES(ctx, resource->cache_.size() == 0,
                errors::Internal("Expect engine cache to be empty, but got ",
                                 resource->cache_.size(), " entries."));

    // Get the file name.
    const string& filename = ctx->input(1).scalar<tstring>()();
    OP_REQUIRES(ctx, !filename.empty(),
                errors::InvalidArgument("filename cannot be empty."));

    // Parse the serialized engines and add them to the cache.
    std::unique_ptr<RandomAccessFile> file;
    OP_REQUIRES_OK(ctx, ctx->env()->NewRandomAccessFile(filename, &file));
    auto reader = absl::make_unique<io::RecordReader>(file.get());

    uint64 offset = 0;
    int num_loaded_engine = 0;
    do {
      tstring record;
      Status status = reader->ReadRecord(&offset, &record);
      if (errors::IsOutOfRange(status)) break;

      TRTEngineInstance engine_instance;
      engine_instance.ParseFromString(record);
      std::vector<TensorShape> engine_input_shapes;
      auto input_shapes = engine_instance.input_shapes();
      engine_input_shapes.reserve(input_shapes.size());
      for (const TensorShapeProto& shape : input_shapes) {
        engine_input_shapes.emplace_back(shape);
      }

      TrtUniquePtrType<IRuntime> infer(
          nvinfer1::createInferRuntime(TRTEngineCacheResource::GetLogger()));
      infer->setGpuAllocator(allocator);
      TrtUniquePtrType<nvinfer1::ICudaEngine> engine(
          infer->deserializeCudaEngine(
              engine_instance.serialized_engine().c_str(),
              engine_instance.serialized_engine().size(), nullptr));
      auto raw_engine = engine.get();
      std::vector<ExecutionContext> ctx_vec;
      if (num_loaded_engine == 0) {
        // Restore profiles if there are any. Currently only 1 engine is allowed
        // in dynamic mode therefore we call this only for the 0th engine.
        // it is a no-op in implicit batch mode.
        OP_REQUIRES_OK(ctx, resource->profiles_.RestoreProfiles(raw_engine));
        OP_REQUIRES_OK(ctx, resource->profiles_.CreateExecutionContexts(
                                raw_engine, &ctx_vec));
      } else {
        // Multiple engines are only available in static mode. For each engine
        // we have only a single execution context.
        ctx_vec.push_back(ExecutionContext::Create(raw_engine));
      }
      resource->cache_.emplace(engine_input_shapes,
                               absl::make_unique<EngineContext>(
                                   std::move(engine), std::move(ctx_vec)));
      ++num_loaded_engine;
    } while (1);
    VLOG(1) << "Loaded " << num_loaded_engine << " TRT engines for op "
            << handle.name() << " on device " << ctx->device()->name()
            << " from file " << filename;
  }

 private:
  // Maximum number of cached engines
  int max_cached_engines_;

  TF_DISALLOW_COPY_AND_ASSIGN(InitializeTRTResource);
};

REGISTER_KERNEL_BUILDER(Name("InitializeTRTResource")
                            .Device(DEVICE_GPU)
                            .HostMemory("resource_handle"),
                        InitializeTRTResource);

class SerializeTRTResource : public OpKernel {
 public:
  explicit SerializeTRTResource(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("delete_resource", &delete_resource_));
  }

  void Compute(OpKernelContext* ctx) override {
    const string& resource_name = ctx->input(0).scalar<tstring>()();
    const string& filename = ctx->input(1).scalar<tstring>()();
    OP_REQUIRES(ctx, !filename.empty(),
                errors::InvalidArgument("filename cannot be empty."));

    // Lookup engine cache resource.
    TRTEngineCacheResource* resource = nullptr;
    OP_REQUIRES_OK(
        ctx, ctx->resource_manager()->Lookup(std::string(kTfTrtContainerName),
                                             resource_name, &resource));
    core::ScopedUnref unref_me(resource);

    // Terminate the calibration if any.
    if (resource->calib_ctx_) resource->calib_ctx_->TerminateCalibration();

    // Serialize the engines and write them to file.
    std::unique_ptr<WritableFile> file;
    OP_REQUIRES_OK(ctx, ctx->env()->NewWritableFile(filename, &file));
    auto writer = absl::make_unique<io::RecordWriter>(file.get());

    int num_serialized_engines = 0;
    for (const auto& pair : resource->cache_) {
      // Ignore engines that failed to build.
      const std::unique_ptr<EngineContext>& engine = pair.second;
      if (!engine || !engine->cuda_engine) continue;

      TRTEngineInstance engine_instance;
      // Add input shapes.
      const std::vector<TensorShape>& engine_input_shapes = pair.first;
      for (const TensorShape& shape : engine_input_shapes) {
        shape.AsProto(engine_instance.add_input_shapes());
      }
      // Add the serialized engine.
      TrtUniquePtrType<nvinfer1::IHostMemory> engine_data(
          engine->cuda_engine->serialize());
      engine_instance.set_serialized_engine(engine_data->data(),
                                            engine_data->size());

      OP_REQUIRES_OK(ctx,
                     writer->WriteRecord(engine_instance.SerializeAsString()));
      ++num_serialized_engines;
    }
    VLOG(1) << "Serialized " << num_serialized_engines << " TRT engines for op "
            << resource_name << " on device " << ctx->device()->name()
            << " to file " << filename;

    if (delete_resource_) {
      VLOG(1) << "Destroying TRT engine cache resource for op " << resource_name
              << " on device " << ctx->device()->name();
      OP_REQUIRES_OK(ctx,
                     ctx->resource_manager()->Delete<TRTEngineCacheResource>(
                         std::string(kTfTrtContainerName), resource_name));
    }
  }

 private:
  bool delete_resource_ = false;

  TF_DISALLOW_COPY_AND_ASSIGN(SerializeTRTResource);
};

REGISTER_KERNEL_BUILDER(Name("SerializeTRTResource").Device(DEVICE_GPU),
                        SerializeTRTResource);

}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
<commit_msg>[tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc] Use `const auto&` instead of `auto`<commit_after>/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

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

    http://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 <algorithm>
#include <memory>
#include <vector>

#include "absl/memory/memory.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_instance.pb.h"  // NOLINT
#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"

#if GOOGLE_CUDA && GOOGLE_TENSORRT
#include "third_party/tensorrt/NvInfer.h"

namespace tensorflow {
namespace tensorrt {
using ::nvinfer1::IRuntime;

class CreateTRTResourceHandle : public OpKernel {
 public:
  explicit CreateTRTResourceHandle(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("resource_name", &resource_name_));
  }

  void Compute(OpKernelContext* ctx) override {
    {
      mutex_lock l(mutex_);
      if (!initialized_) {
        AllocatorAttributes attr;
        attr.set_on_host(true);
        OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
                                               &handle_, attr));

        VLOG(1) << "Creating TRT engine cache resource handle for op "
                << resource_name_ << " on device " << ctx->device()->name();
        handle_.scalar<ResourceHandle>()() =
            MakeResourceHandle<TRTEngineCacheResource>(
                ctx, std::string(kTfTrtContainerName), resource_name_);
        initialized_ = true;
      }
    }
    ctx->set_output(0, handle_);
  }

 private:
  string resource_name_;
  Tensor handle_;
  mutex mutex_;
  bool initialized_ TF_GUARDED_BY(mutex_) = false;

  TF_DISALLOW_COPY_AND_ASSIGN(CreateTRTResourceHandle);
};

REGISTER_KERNEL_BUILDER(Name("CreateTRTResourceHandle")
                            .Device(DEVICE_GPU)
                            .HostMemory("resource_handle"),
                        CreateTRTResourceHandle);

class InitializeTRTResource : public OpKernel {
 public:
  explicit InitializeTRTResource(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(
        ctx, ctx->GetAttr("max_cached_engines_count", &max_cached_engines_));
  }

  void Compute(OpKernelContext* ctx) override {
    ResourceHandle handle = HandleFromInput(ctx, 0);
    core::RefCountPtr<TRTEngineCacheResource> resource;
    OP_REQUIRES_OK(
        ctx, LookupOrCreateResource<TRTEngineCacheResource>(
                 ctx, handle, &resource,
                 [this, ctx](TRTEngineCacheResource** resource) -> Status {
                   *resource = new TRTEngineCacheResource(
                       ctx, this->max_cached_engines_);
                   return Status::OK();
                 }));

    auto allocator = resource->allocator_.get();
    OP_REQUIRES(ctx, allocator != nullptr,
                errors::Internal("Not able to initialize TRT engine cache when "
                                 "GPU allocator is empty."));
    OP_REQUIRES(ctx, resource->cache_.size() == 0,
                errors::Internal("Expect engine cache to be empty, but got ",
                                 resource->cache_.size(), " entries."));

    // Get the file name.
    const string& filename = ctx->input(1).scalar<tstring>()();
    OP_REQUIRES(ctx, !filename.empty(),
                errors::InvalidArgument("filename cannot be empty."));

    // Parse the serialized engines and add them to the cache.
    std::unique_ptr<RandomAccessFile> file;
    OP_REQUIRES_OK(ctx, ctx->env()->NewRandomAccessFile(filename, &file));
    auto reader = absl::make_unique<io::RecordReader>(file.get());

    uint64 offset = 0;
    int num_loaded_engine = 0;
    do {
      tstring record;
      Status status = reader->ReadRecord(&offset, &record);
      if (errors::IsOutOfRange(status)) break;

      TRTEngineInstance engine_instance;
      engine_instance.ParseFromString(record);
      std::vector<TensorShape> engine_input_shapes;
      const auto& input_shapes = engine_instance.input_shapes();
      engine_input_shapes.reserve(input_shapes.size());
      for (const TensorShapeProto& shape : input_shapes) {
        engine_input_shapes.emplace_back(shape);
      }

      TrtUniquePtrType<IRuntime> infer(
          nvinfer1::createInferRuntime(TRTEngineCacheResource::GetLogger()));
      infer->setGpuAllocator(allocator);
      TrtUniquePtrType<nvinfer1::ICudaEngine> engine(
          infer->deserializeCudaEngine(
              engine_instance.serialized_engine().c_str(),
              engine_instance.serialized_engine().size(), nullptr));
      auto raw_engine = engine.get();
      std::vector<ExecutionContext> ctx_vec;
      if (num_loaded_engine == 0) {
        // Restore profiles if there are any. Currently only 1 engine is allowed
        // in dynamic mode therefore we call this only for the 0th engine.
        // it is a no-op in implicit batch mode.
        OP_REQUIRES_OK(ctx, resource->profiles_.RestoreProfiles(raw_engine));
        OP_REQUIRES_OK(ctx, resource->profiles_.CreateExecutionContexts(
                                raw_engine, &ctx_vec));
      } else {
        // Multiple engines are only available in static mode. For each engine
        // we have only a single execution context.
        ctx_vec.push_back(ExecutionContext::Create(raw_engine));
      }
      resource->cache_.emplace(engine_input_shapes,
                               absl::make_unique<EngineContext>(
                                   std::move(engine), std::move(ctx_vec)));
      ++num_loaded_engine;
    } while (1);
    VLOG(1) << "Loaded " << num_loaded_engine << " TRT engines for op "
            << handle.name() << " on device " << ctx->device()->name()
            << " from file " << filename;
  }

 private:
  // Maximum number of cached engines
  int max_cached_engines_;

  TF_DISALLOW_COPY_AND_ASSIGN(InitializeTRTResource);
};

REGISTER_KERNEL_BUILDER(Name("InitializeTRTResource")
                            .Device(DEVICE_GPU)
                            .HostMemory("resource_handle"),
                        InitializeTRTResource);

class SerializeTRTResource : public OpKernel {
 public:
  explicit SerializeTRTResource(OpKernelConstruction* ctx) : OpKernel(ctx) {
    OP_REQUIRES_OK(ctx, ctx->GetAttr("delete_resource", &delete_resource_));
  }

  void Compute(OpKernelContext* ctx) override {
    const string& resource_name = ctx->input(0).scalar<tstring>()();
    const string& filename = ctx->input(1).scalar<tstring>()();
    OP_REQUIRES(ctx, !filename.empty(),
                errors::InvalidArgument("filename cannot be empty."));

    // Lookup engine cache resource.
    TRTEngineCacheResource* resource = nullptr;
    OP_REQUIRES_OK(
        ctx, ctx->resource_manager()->Lookup(std::string(kTfTrtContainerName),
                                             resource_name, &resource));
    core::ScopedUnref unref_me(resource);

    // Terminate the calibration if any.
    if (resource->calib_ctx_) resource->calib_ctx_->TerminateCalibration();

    // Serialize the engines and write them to file.
    std::unique_ptr<WritableFile> file;
    OP_REQUIRES_OK(ctx, ctx->env()->NewWritableFile(filename, &file));
    auto writer = absl::make_unique<io::RecordWriter>(file.get());

    int num_serialized_engines = 0;
    for (const auto& pair : resource->cache_) {
      // Ignore engines that failed to build.
      const std::unique_ptr<EngineContext>& engine = pair.second;
      if (!engine || !engine->cuda_engine) continue;

      TRTEngineInstance engine_instance;
      // Add input shapes.
      const std::vector<TensorShape>& engine_input_shapes = pair.first;
      for (const TensorShape& shape : engine_input_shapes) {
        shape.AsProto(engine_instance.add_input_shapes());
      }
      // Add the serialized engine.
      TrtUniquePtrType<nvinfer1::IHostMemory> engine_data(
          engine->cuda_engine->serialize());
      engine_instance.set_serialized_engine(engine_data->data(),
                                            engine_data->size());

      OP_REQUIRES_OK(ctx,
                     writer->WriteRecord(engine_instance.SerializeAsString()));
      ++num_serialized_engines;
    }
    VLOG(1) << "Serialized " << num_serialized_engines << " TRT engines for op "
            << resource_name << " on device " << ctx->device()->name()
            << " to file " << filename;

    if (delete_resource_) {
      VLOG(1) << "Destroying TRT engine cache resource for op " << resource_name
              << " on device " << ctx->device()->name();
      OP_REQUIRES_OK(ctx,
                     ctx->resource_manager()->Delete<TRTEngineCacheResource>(
                         std::string(kTfTrtContainerName), resource_name));
    }
  }

 private:
  bool delete_resource_ = false;

  TF_DISALLOW_COPY_AND_ASSIGN(SerializeTRTResource);
};

REGISTER_KERNEL_BUILDER(Name("SerializeTRTResource").Device(DEVICE_GPU),
                        SerializeTRTResource);

}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT
<|endoftext|>


<commit_before>#include "core/tcp_file_server.h"

#include "core/array.h"
#include "core/free_list.h"
#include "core/static_array.h"
#include "core/tcp_file_device.h"
#include "core/task.h"
#include "core/tcp_acceptor.h"
#include "core/tcp_stream.h"
#include "core/os_file.h"

namespace Lux
{
	namespace FS
	{
		class TCPFileServerTask : public MT::Task
		{
		public:
			TCPFileServerTask() {}
			~TCPFileServerTask() {}

			int task()
			{
				bool quit = false;

				m_acceptor.start("127.0.0.1", 10001);
				Net::TCPStream* stream = m_acceptor.accept();

				while(!quit)
				{
					int32_t op = 0;
					stream->read(op);
					switch(op)
					{
					case TCPCommand::OpenFile:
						{
							int32_t mode = 0;
							int32_t len = 0;
							stream->read(mode);
							stream->read(m_buffer.data(), m_buffer.size());

							int32_t ret = -2;
							int32_t id = m_ids.alloc();
							if(id > 0)
							{
								OsFile* file = LUX_NEW(OsFile)();
								m_files[id] = file;

								ret = file->open(m_buffer.data(), mode) ? id : -1;
							}
							stream->write(ret);
						}
						break;
					case TCPCommand::Close:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];
							m_ids.release(id);

							file->close();
							LUX_DELETE(file);
						}
						break;
					case TCPCommand::Read:
						{
							bool read_successful = true;
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t size = 0;
							stream->read(size);

							while(size > 0)
							{
								int32_t read = (int32_t)size > m_buffer.size() ? m_buffer.size() : (int32_t)size;
								read_successful &= file->read((void*)m_buffer.data(), read);
								stream->write((const void*)m_buffer.data(), read);
								size -= read;
							}

							stream->write(read_successful);
						}
						break;
					case TCPCommand::Write:
						{
							bool write_successful = true;
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t size = 0;
							stream->read(size);

							while(size > 0)
							{
								int32_t read = (int32_t)size > m_buffer.size() ? m_buffer.size() : (int32_t)size;
								write_successful &= stream->read((void*)m_buffer.data(), read);
								file->write(m_buffer.data(), read);
								size -= read;
							}

							stream->write(write_successful);
						}
						break;
					case TCPCommand::Size:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t size = (uint32_t)file->size();
							stream->write(size);
						}
						break;
					case TCPCommand::Seek:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t base = 0;
							int32_t offset = 0;
							stream->read(base);
							stream->read(offset);

							uint32_t pos = (uint32_t)file->seek((SeekMode)base, offset);
							stream->write(pos);
						}
						break;
					case TCPCommand::Pos:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t pos = (uint32_t)file->pos();
							stream->write(pos);
						}
						break;
					case TCPCommand::Disconnect:
						{
							quit = true;
							break;
						}
					default:
						ASSERT(0);
						break;
					}
				}

				LUX_DELETE(stream);
				return 0;
			}

			void stop() {} // TODO: implement stop

		private:
			Net::TCPAcceptor			m_acceptor;
			StaticArray<char, 0x50000>	m_buffer;
			StaticArray<OsFile*, 0x50000> m_files;
			FreeList<int32_t, 0x50000>	m_ids;
		};

		struct TCPFileServerImpl
		{
			TCPFileServerTask m_task;
		};

		TCPFileServer::TCPFileServer()
		{
			m_impl = NULL;
		}

		TCPFileServer::~TCPFileServer()
		{
			LUX_DELETE(m_impl);
		}

		void TCPFileServer::start()
		{
			m_impl = LUX_NEW(TCPFileServerImpl);
			m_impl->m_task.create("TCP File Server Task");
			m_impl->m_task.run();
		}

		void TCPFileServer::stop()
		{
			m_impl->m_task.stop();
			m_impl->m_task.destroy();
			LUX_DELETE(m_impl);
			m_impl = NULL;
		}
	} // ~namespace FS
} // ~namespace Lux<commit_msg>file server fixed<commit_after>#include "core/tcp_file_server.h"

#include "core/array.h"
#include "core/free_list.h"
#include "core/static_array.h"
#include "core/tcp_file_device.h"
#include "core/task.h"
#include "core/tcp_acceptor.h"
#include "core/tcp_stream.h"
#include "core/os_file.h"

namespace Lux
{
	namespace FS
	{
		class TCPFileServerTask : public MT::Task
		{
		public:
			TCPFileServerTask() {}
			~TCPFileServerTask() {}

			int task()
			{
				bool quit = false;

				m_acceptor.start("127.0.0.1", 10001);
				Net::TCPStream* stream = m_acceptor.accept();

				while(!quit)
				{
					int32_t op = 0;
					stream->read(op);
					switch(op)
					{
					case TCPCommand::OpenFile:
						{
							int32_t mode = 0;
							int32_t len = 0;
							stream->read(mode);
							stream->readString(m_buffer.data(), m_buffer.size());

							int32_t ret = -2;
							int32_t id = m_ids.alloc();
							if(id > 0)
							{
								OsFile* file = LUX_NEW(OsFile)();
								m_files[id] = file;

								ret = file->open(m_buffer.data(), mode) ? id : -1;
							}
							stream->write(ret);
						}
						break;
					case TCPCommand::Close:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];
							m_ids.release(id);

							file->close();
							LUX_DELETE(file);
						}
						break;
					case TCPCommand::Read:
						{
							bool read_successful = true;
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t size = 0;
							stream->read(size);

							while(size > 0)
							{
								int32_t read = (int32_t)size > m_buffer.size() ? m_buffer.size() : (int32_t)size;
								read_successful &= file->read((void*)m_buffer.data(), read);
								stream->write((const void*)m_buffer.data(), read);
								size -= read;
							}

							stream->write(read_successful);
						}
						break;
					case TCPCommand::Write:
						{
							bool write_successful = true;
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t size = 0;
							stream->read(size);

							while(size > 0)
							{
								int32_t read = (int32_t)size > m_buffer.size() ? m_buffer.size() : (int32_t)size;
								write_successful &= stream->read((void*)m_buffer.data(), read);
								file->write(m_buffer.data(), read);
								size -= read;
							}

							stream->write(write_successful);
						}
						break;
					case TCPCommand::Size:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t size = (uint32_t)file->size();
							stream->write(size);
						}
						break;
					case TCPCommand::Seek:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t base = 0;
							int32_t offset = 0;
							stream->read(base);
							stream->read(offset);

							uint32_t pos = (uint32_t)file->seek((SeekMode)base, offset);
							stream->write(pos);
						}
						break;
					case TCPCommand::Pos:
						{
							uint32_t id = -1;
							stream->read(id);
							OsFile* file = m_files[id];

							uint32_t pos = (uint32_t)file->pos();
							stream->write(pos);
						}
						break;
					case TCPCommand::Disconnect:
						{
							quit = true;
							break;
						}
					default:
						ASSERT(0);
						break;
					}
				}

				LUX_DELETE(stream);
				return 0;
			}

			void stop() {} // TODO: implement stop

		private:
			Net::TCPAcceptor			m_acceptor;
			StaticArray<char, 0x50000>	m_buffer;
			StaticArray<OsFile*, 0x50000> m_files;
			FreeList<int32_t, 0x50000>	m_ids;
		};

		struct TCPFileServerImpl
		{
			TCPFileServerTask m_task;
		};

		TCPFileServer::TCPFileServer()
		{
			m_impl = NULL;
		}

		TCPFileServer::~TCPFileServer()
		{
			LUX_DELETE(m_impl);
		}

		void TCPFileServer::start()
		{
			m_impl = LUX_NEW(TCPFileServerImpl);
			m_impl->m_task.create("TCP File Server Task");
			m_impl->m_task.run();
		}

		void TCPFileServer::stop()
		{
			m_impl->m_task.stop();
			m_impl->m_task.destroy();
			LUX_DELETE(m_impl);
			m_impl = NULL;
		}
	} // ~namespace FS
} // ~namespace Lux<|endoftext|>


<commit_before>#include <Ab/Config.hpp>
#include <Ab/LinearMemory.hpp>
#include <gtest/gtest.h>

namespace Ab::Test {

TEST(LinearMemoryTest, InitAndKill) { LinearMemory m; }

TEST(LinearMemoryTest, GrowThreeTimes) {
	LinearMemory m;
	for (std::size_t i = 0; i < 3; i++) {
		m.grow();
	}
}

TEST(LinearMemoryTest, AssignToMemory) {
	LinearMemoryConfig cfg;
	cfg.page_count_min = 1;
	cfg.page_count_max = 1;

	LinearMemory m(cfg);

	auto ptr = to_ptr<int>(m.address());
	EXPECT_NE(ptr, nullptr);
	EXPECT_EQ(*ptr, 0);
	*ptr = 123;
	EXPECT_EQ(*ptr, 123);
}

}  // namespace Ab::Test
<commit_msg>Use LinearMemoryConfig when creating LinearMemory in test<commit_after>#include <Ab/Config.hpp>
#include <Ab/LinearMemory.hpp>
#include <gtest/gtest.h>

namespace Ab::Test {

TEST(LinearMemoryTest, InitAndKill) { LinearMemory m; }

TEST(LinearMemoryTest, GrowThreeTimes) {
	LinearMemoryConfig cfg;
	cfg.page_count_min = 0;
	cfg.page_count_max = 3;

	LinearMemory m(cfg);
	for (std::size_t i = 0; i < 3; i++) {
		m.grow();
	}
}

TEST(LinearMemoryTest, AssignToMemory) {
	LinearMemoryConfig cfg;
	cfg.page_count_min = 1;
	cfg.page_count_max = 1;

	LinearMemory m(cfg);

	auto ptr = to_ptr<int>(m.address());
	EXPECT_NE(ptr, nullptr);
	EXPECT_EQ(*ptr, 0);
	*ptr = 123;
	EXPECT_EQ(*ptr, 123);
}

}  // namespace Ab::Test
<|endoftext|>


<commit_before>/*************************************************************************
 *
 *  $RCSfile: servreg.cxx,v $
 *
 *  $Revision: 1.2 $
 *
 *  last change: $Author: jl $ $Date: 2001-06-27 10:56:03 $
 *
 *  The Contents of this file are made available subject to the terms of
 *  either of the following licenses
 *
 *         - GNU Lesser General Public License Version 2.1
 *         - Sun Industry Standards Source License Version 1.1
 *
 *  Sun Microsystems Inc., October, 2000
 *
 *  GNU Lesser General Public License Version 2.1
 *  =============================================
 *  Copyright 2000 by Sun Microsystems, Inc.
 *  901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *  This library is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU Lesser General Public
 *  License version 2.1, as published by the Free Software Foundation.
 *
 *  This library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *  Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public
 *  License along with this library; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *  MA  02111-1307  USA
 *
 *
 *  Sun Industry Standards Source License Version 1.1
 *  =================================================
 *  The contents of this file are subject to the Sun Industry Standards
 *  Source License Version 1.1 (the "License"); You may not use this file
 *  except in compliance with the License. You may obtain a copy of the
 *  License at http://www.openoffice.org/license.html.
 *
 *  Software provided under this License is provided on an "AS IS" basis,
 *  WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING,
 *  WITHOUT LIMITATION, WARRANTIES THAT THE SOFTWARE IS FREE OF DEFECTS,
 *  MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE, OR NON-INFRINGING.
 *  See the License for the specific provisions governing your rights and
 *  obligations concerning the Software.
 *
 *  The Initial Developer of the Original Code is: Sun Microsystems, Inc.
 *
 *  Copyright: 2000 by Sun Microsystems, Inc.
 *
 *  All Rights Reserved.
 *
 *  Contributor(s): _______________________________________
 *
 *
 ************************************************************************/
#include <rtl/unload.h>
#include <osl/time.h>
#include "ole2uno.hxx"
#include "servprov.hxx"
#ifndef _RTL_USTRING_
#include <rtl/ustring>
#endif
#ifndef _CPPUHELPER_FACTORY_HXX_
#include <cppuhelper/factory.hxx>
#endif
using namespace rtl;
using namespace ole_adapter;
using namespace cppu;

namespace ole_adapter
{
rtl_StandardModuleCount globalModuleCount= MODULE_COUNT_INIT;

Reference<XInterface> SAL_CALL ConverterProvider_CreateInstance2(   const Reference<XMultiServiceFactory> & xSMgr)
                            throw(Exception)
{
    Reference<XInterface> xService = *new OleConverter_Impl2( xSMgr);
    return xService;
}

Reference<XInterface> SAL_CALL ConverterProvider_CreateInstanceVar1(    const Reference<XMultiServiceFactory> & xSMgr)
                            throw(Exception)
{
    Reference<XInterface> xService = *new OleConverter_Impl2( xSMgr, UNO_OBJECT_WRAPPER_REMOTE_OPT, IUNKNOWN_WRAPPER_IMPL);
    return xService;
}

Reference<XInterface> SAL_CALL OleClient_CreateInstance( const Reference<XMultiServiceFactory> & xSMgr)
                            throw(Exception)
{
    Reference<XInterface> xService = *new OleClient_Impl( xSMgr);
    return xService;
}

Reference<XInterface> SAL_CALL OleServer_CreateInstance( const Reference<XMultiServiceFactory> & xSMgr)
                            throw (Exception)
{
    Reference<XInterface > xService = *new OleServer_Impl(xSMgr);
    return xService;
}
} // end namespace

extern "C" void * SAL_CALL component_getFactory(
    const sal_Char * pImplName, void * pServiceManager, void * pRegistryKey )
{
    void * pRet = 0;

    OUString aImplName( OUString::createFromAscii( pImplName ) );
    Reference< XSingleServiceFactory > xFactory;
    Sequence<OUString> seqServiceNames;
    if (pServiceManager && aImplName.equals(  L"com.sun.star.comp.ole.OleConverter2"  ))
    {
        xFactory=  createSingleFactory( reinterpret_cast< XMultiServiceFactory*>(pServiceManager),
                                         OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleConverter2")),
                                         ConverterProvider_CreateInstance2, seqServiceNames,
                                         &globalModuleCount.modCnt );
    }
    else if (pServiceManager && aImplName.equals(  L"com.sun.star.comp.ole.OleConverterVar1"  ))
    {
        xFactory= createSingleFactory( reinterpret_cast<XMultiServiceFactory*>(pServiceManager),
                                       OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleConverterVar1")),
                                       ConverterProvider_CreateInstanceVar1, seqServiceNames,
                                       &globalModuleCount.modCnt );
    }
    else if(pServiceManager && aImplName.equals(L"com.sun.star.comp.ole.OleClient"))
    {
        xFactory= createSingleFactory( reinterpret_cast< XMultiServiceFactory*>(pServiceManager),
                                       OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleClient")),
                                       OleClient_CreateInstance, seqServiceNames,
                                       &globalModuleCount.modCnt);
    }
    else if(pServiceManager && aImplName.equals(L"com.sun.star.comp.ole.OleServer"))
    {
        xFactory= createOneInstanceFactory( reinterpret_cast< XMultiServiceFactory*>(pServiceManager),
                                            OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleServer")),
                                            OleServer_CreateInstance, seqServiceNames,
                                            &globalModuleCount.modCnt);
    }

    if (xFactory.is())
    {
        xFactory->acquire();
        pRet = xFactory.get();
    }

    return pRet;
}

extern "C" sal_Bool SAL_CALL component_writeInfo(   void * pServiceManager, void * pRegistryKey )
{
    if (pRegistryKey)
    {
        try
        {
            Reference<XRegistryKey> xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleConverter2/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleBridgeSupplier2");

            xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleConverterVar1/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleBridgeSupplierVar1");

            xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleClient/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleObjectFactory");

            xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleServer/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleApplicationRegistration");

            return sal_True;
        }
        catch (InvalidRegistryException &)
        {
            OSL_ENSURE( sal_False, "### InvalidRegistryException!" );
        }
    }
    return sal_False;
}

extern "C" void SAL_CALL component_getImplementationEnvironment(
    const sal_Char ** ppEnvTypeName, uno_Environment ** ppEnv )
{
    *ppEnvTypeName = CPPU_CURRENT_LANGUAGE_BINDING_NAME;
}

extern "C"  sal_Bool component_canUnload( TimeValue* libUnused)
{
    return globalModuleCount.canUnload( &globalModuleCount, libUnused);
}
<commit_msg>INTEGRATION: CWS jl5vba (1.2.244); FILE MERGED 2004/01/19 11:06:44 jl 1.2.244.1: #112439# join cws ab02vba<commit_after>/*************************************************************************
 *
 *  $RCSfile: servreg.cxx,v $
 *
 *  $Revision: 1.3 $
 *
 *  last change: $Author: obo $ $Date: 2004-03-17 13:08:41 $
 *
 *  The Contents of this file are made available subject to the terms of
 *  either of the following licenses
 *
 *         - GNU Lesser General Public License Version 2.1
 *         - Sun Industry Standards Source License Version 1.1
 *
 *  Sun Microsystems Inc., October, 2000
 *
 *  GNU Lesser General Public License Version 2.1
 *  =============================================
 *  Copyright 2000 by Sun Microsystems, Inc.
 *  901 San Antonio Road, Palo Alto, CA 94303, USA
 *
 *  This library is free software; you can redistribute it and/or
 *  modify it under the terms of the GNU Lesser General Public
 *  License version 2.1, as published by the Free Software Foundation.
 *
 *  This library is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 *  Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public
 *  License along with this library; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston,
 *  MA  02111-1307  USA
 *
 *
 *  Sun Industry Standards Source License Version 1.1
 *  =================================================
 *  The contents of this file are subject to the Sun Industry Standards
 *  Source License Version 1.1 (the "License"); You may not use this file
 *  except in compliance with the License. You may obtain a copy of the
 *  License at http://www.openoffice.org/license.html.
 *
 *  Software provided under this License is provided on an "AS IS" basis,
 *  WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING,
 *  WITHOUT LIMITATION, WARRANTIES THAT THE SOFTWARE IS FREE OF DEFECTS,
 *  MERCHANTABLE, FIT FOR A PARTICULAR PURPOSE, OR NON-INFRINGING.
 *  See the License for the specific provisions governing your rights and
 *  obligations concerning the Software.
 *
 *  The Initial Developer of the Original Code is: Sun Microsystems, Inc.
 *
 *  Copyright: 2000 by Sun Microsystems, Inc.
 *
 *  All Rights Reserved.
 *
 *  Contributor(s): _______________________________________
 *
 *
 ************************************************************************/
#include <rtl/unload.h>
#include <osl/time.h>
#include "ole2uno.hxx"
#include "servprov.hxx"
#ifndef _RTL_USTRING_
#include <rtl/ustring>
#endif
#ifndef _CPPUHELPER_FACTORY_HXX_
#include <cppuhelper/factory.hxx>
#endif
using namespace rtl;
using namespace ole_adapter;
using namespace cppu;

namespace ole_adapter
{
rtl_StandardModuleCount globalModuleCount= MODULE_COUNT_INIT;

Reference<XInterface> SAL_CALL ConverterProvider_CreateInstance2(   const Reference<XMultiServiceFactory> & xSMgr)
                            throw(Exception)
{
    Reference<XInterface> xService = *new OleConverter_Impl2( xSMgr);
    return xService;
}

Reference<XInterface> SAL_CALL ConverterProvider_CreateInstanceVar1(    const Reference<XMultiServiceFactory> & xSMgr)
                            throw(Exception)
{
    Reference<XInterface> xService = *new OleConverter_Impl2( xSMgr, UNO_OBJECT_WRAPPER_REMOTE_OPT, IUNKNOWN_WRAPPER_IMPL);
    return xService;
}

Reference<XInterface> SAL_CALL OleClient_CreateInstance( const Reference<XMultiServiceFactory> & xSMgr)
                            throw(Exception)
{
    Reference<XInterface> xService = *new OleClient_Impl( xSMgr);
    return xService;
}

Reference<XInterface> SAL_CALL OleServer_CreateInstance( const Reference<XMultiServiceFactory> & xSMgr)
                            throw (Exception)
{
    Reference<XInterface > xService = *new OleServer_Impl(xSMgr);
    return xService;
}
} // end namespace

extern "C" void * SAL_CALL component_getFactory(
    const sal_Char * pImplName, void * pServiceManager, void * pRegistryKey )
{
    void * pRet = 0;

    OUString aImplName( OUString::createFromAscii( pImplName ) );
    Reference< XSingleServiceFactory > xFactory;
    Sequence<OUString> seqServiceNames;
    if (pServiceManager && aImplName.equals(  L"com.sun.star.comp.ole.OleConverter2"  ))
    {
        xFactory=  createSingleFactory( reinterpret_cast< XMultiServiceFactory*>(pServiceManager),
                                         OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleConverter2")),
                                         ConverterProvider_CreateInstance2, seqServiceNames,
                                         &globalModuleCount.modCnt );
    }
    else if (pServiceManager && aImplName.equals(  L"com.sun.star.comp.ole.OleConverterVar1"  ))
    {
        xFactory= createSingleFactory( reinterpret_cast<XMultiServiceFactory*>(pServiceManager),
                                       OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleConverterVar1")),
                                       ConverterProvider_CreateInstanceVar1, seqServiceNames,
                                       &globalModuleCount.modCnt );
    }
    else if(pServiceManager && aImplName.equals(L"com.sun.star.comp.ole.OleClient"))
    {
        xFactory= createSingleFactory( reinterpret_cast< XMultiServiceFactory*>(pServiceManager),
                                       OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleClient")),
                                       OleClient_CreateInstance, seqServiceNames,
                                       &globalModuleCount.modCnt);
    }
    else if(pServiceManager && aImplName.equals(L"com.sun.star.comp.ole.OleServer"))
    {
        xFactory= createOneInstanceFactory( reinterpret_cast< XMultiServiceFactory*>(pServiceManager),
                                            OUString( RTL_CONSTASCII_USTRINGPARAM("com.sun.star.comp.ole.OleServer")),
                                            OleServer_CreateInstance, seqServiceNames,
                                            &globalModuleCount.modCnt);
    }

    if (xFactory.is())
    {
        xFactory->acquire();
        pRet = xFactory.get();
    }

    return pRet;
}

extern "C" sal_Bool SAL_CALL component_writeInfo(   void * pServiceManager, void * pRegistryKey )
{
    if (pRegistryKey)
    {
        try
        {
            //deprecated
            Reference<XRegistryKey> xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleConverter2/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleBridgeSupplier2");

            xNewKey->createKey(L"com.sun.star.bridge.oleautomation.BridgeSupplier");

            //deprecated
            xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleConverterVar1/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleBridgeSupplierVar1");

            //deprecated
            xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleClient/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleObjectFactory");

            xNewKey->createKey(L"com.sun.star.bridge.oleautomation.Factory");
            //deprecated
            xNewKey =
                reinterpret_cast<XRegistryKey*>( pRegistryKey)->createKey(L"/com.sun.star.comp.ole.OleServer/UNO/SERVICES");
            xNewKey->createKey(L"com.sun.star.bridge.OleApplicationRegistration");

            xNewKey->createKey(L"com.sun.star.bridge.oleautomation.ApplicationRegistration");

            return sal_True;
        }
        catch (InvalidRegistryException &)
        {
            OSL_ENSURE( sal_False, "### InvalidRegistryException!" );
        }
    }
    return sal_False;
}

extern "C" void SAL_CALL component_getImplementationEnvironment(
    const sal_Char ** ppEnvTypeName, uno_Environment ** ppEnv )
{
    *ppEnvTypeName = CPPU_CURRENT_LANGUAGE_BINDING_NAME;
}

extern "C"  sal_Bool component_canUnload( TimeValue* libUnused)
{
    return globalModuleCount.canUnload( &globalModuleCount, libUnused);
}
<|endoftext|>