session.hpp

Go to the documentation of this file.

#include <list>
#include <unordered_set>

#include "teq/traveler.hpp"

#include "opt/optimize.hpp"

#include "eteq/constant.hpp"
#include "eteq/functor.hpp"

#ifndef ETEQ_SESSION_HPP
#define ETEQ_SESSION_HPP

namespace eteq
{

struct iSession
{
    virtual ~iSession (void) = default;

    virtual void track (teq::TensptrsT roots) = 0;

    virtual void update (teq::TensSetT ignored = {}) = 0;

    virtual void update_target (teq::TensSetT target, teq::TensSetT ignored = {}) = 0;
};

struct Session final : public iSession
{
    void track (teq::TensptrsT roots) override
    {
        ops_.clear();
        tracked_.insert(roots.begin(), roots.end());

        teq::GraphStat stat;
        for (teq::TensptrT& root : roots)
        {
            root->accept(stat);
        }
        auto& statmap = stat.graphsize_;

        for (auto& statpair : statmap)
        {
            if (0 < statpair.second.upper_)
            {
                // ensure we only track operable functors
                auto op = dynamic_cast<teq::iOperableFunc*>(statpair.first);
                if (nullptr == op)
                {
                    logs::fatalf("cannot track non-operable functor %s",
                        statpair.first->to_string().c_str());
                }
                ops_.push_back(op);
            }
        }
        std::sort(ops_.begin(), ops_.end(),
            [&statmap](teq::iOperableFunc* a, teq::iOperableFunc* b)
            { return statmap[a].upper_ < statmap[b].upper_; });
    }

    void update (teq::TensSetT ignored = {}) override
    {
        std::list<teq::iOperableFunc*> reqs;
        teq::TensSetT acceptable;
        for (auto& root : tracked_)
        {
            acceptable.emplace(root.get());
        }
        // ignored tensors will never populate reqs
        for (auto rit = ops_.rbegin(), ret = ops_.rend();
            rit != ret; ++rit)
        {
            auto& op = *rit;
            if (estd::has(acceptable, op) &&
                false == estd::has(ignored, op))
            {
                reqs.push_front(op);
                auto& children = op->get_children();
                for (auto& child : children)
                {
                    acceptable.emplace(child.get_tensor().get());
                }
            }
        }

        for (auto& op : reqs)
        {
            op->update();
        }
    }

    void update_target (teq::TensSetT target, teq::TensSetT ignored = {}) override
    {
        std::list<teq::iOperableFunc*> reqs;
        teq::TensSetT acceptable;
        for (auto& root : target)
        {
            acceptable.emplace(root);
        }
        // ignored tensors will never populate reqs
        for (auto rit = ops_.rbegin(), ret = ops_.rend();
            rit != ret; ++rit)
        {
            auto& op = *rit;
            if (estd::has(acceptable, op) &&
                false == estd::has(ignored, op))
            {
                reqs.push_front(op);
                auto& children = op->get_children();
                for (auto& child : children)
                {
                    acceptable.emplace(child.get_tensor().get());
                }
            }
        }

        for (auto& op : reqs)
        {
            op->update();
        }
    }

    void optimize (const opt::OptCtx& rules)
    {
        teq::TensptrsT tracked(tracked_.begin(), tracked_.end());
        opt::optimize(tracked, rules);
        track(tracked);
    }

    teq::TensptrSetT tracked_;

    std::vector<teq::iOperableFunc*> ops_;
};

}

#endif // ETEQ_SESSION_HPP