shear.hpp

Go to the documentation of this file.

#include <list>

#include "ade/ade.hpp"

#ifndef OPT_SHEAR_HPP
#define OPT_SHEAR_HPP

namespace opt
{

template <typename T>
using GetLeafValT = std::function<T(ade::iLeaf*)>;

using ParentMapT = std::unordered_map<
    ade::iTensor*,std::unordered_set<size_t>>;

using PruneFuncT = std::function<ade::TensptrT(ade::iFunctor*,\
    std::unordered_set<size_t>,ade::ArgsT)>;

template <typename T>
struct LeafFinder final : public ade::iTraveler
{
    LeafFinder (T target, GetLeafValT<T> get_leaf) :
        target_(target), get_leaf_(get_leaf) {}

    void visit (ade::iLeaf* leaf) override
    {
        if (target_ == get_leaf_(leaf))
        {
            founds_.emplace(leaf);
        }
    }

    void visit (ade::iFunctor* func) override
    {
        if (parents_.end() == parents_.find(func))
        {
            auto& children = func->get_children();
            size_t n = children.size();
            std::unordered_set<size_t> path;
            for (size_t i = 0; i < n; ++i)
            {
                ade::TensptrT tens = children[i].get_tensor();
                tens->accept(*this);
                if (parents_.end() != parents_.find(tens.get()) ||
                    founds_.end() != founds_.find(tens.get()))
                {
                    path.emplace(i);
                }
            }
            if (false == path.empty())
            {
                parents_[func] = path;
            }
        }
    }

    T target_;

    GetLeafValT<T> get_leaf_;

    std::unordered_set<ade::iTensor*> founds_;

    ParentMapT parents_;
};

template <typename T>
struct TargetPruner
{
    TargetPruner (T target, GetLeafValT<T> find_target, PruneFuncT pruner) :
        finder_(target, find_target), pruner_(pruner) {}

    ade::TensptrT prune (ade::TensptrT root)
    {
        // assert that context will be unaffected by prune,
        // since source will never be touched
        root->accept(finder_);
        auto& pathmap = finder_.parents_;
        if (pathmap.empty()) // not path to target or root is not a parent
        {
            return root;
        }
        ade::GraphStat stat;
        root->accept(stat);
        // grab the intersection of stat.funcs_ and pathmap
        std::list<ade::iFunctor*> parents;
        std::transform(pathmap.begin(), pathmap.end(),
            std::back_inserter(parents),
            [](std::pair<ade::iTensor*,std::unordered_set<size_t>> parent)
            {
                return static_cast<ade::iFunctor*>(parent.first);
            });
        parents.sort(
            [&](ade::iTensor* a, ade::iTensor* b)
            {
                return stat.graphsize_[a] < stat.graphsize_[b];
            });

        // each proceeding node in parents list is closer to target
        // start pruning according to each parent node in order
        std::unordered_set<ade::iTensor*> targets = finder_.founds_;

        std::unordered_map<ade::iTensor*,ade::TensptrT> mapping;
        std::unordered_map<ade::iTensor*,bool> targetmap;
        for (ade::iFunctor* func : parents)
        {
            ade::ArgsT children = func->get_children();
            // indices lead to target nodes
            std::unordered_set<size_t> indices = pathmap[func];
            for (auto it = indices.begin(), et = indices.end(); it != et;)
            {
                ade::MappedTensor& child = children[*it];
                ade::iTensor* tens = child.get_tensor().get();
                // child is not target, so erase ot from indices
                auto zit = targets.find(tens);
                if (targets.end() == zit)
                {
                    it = indices.erase(it);
                }
                else
                {
                    ++it;
                }
            }
            auto fwd = pruner_(func, indices, children);
            mapping.emplace(func, fwd);
            // func maps to target, so store in targets
            if (ade::iLeaf* fwdleaf = dynamic_cast<ade::iLeaf*>(fwd.get()))
            {
                if (finder_.get_leaf_(fwdleaf) == finder_.target_)
                {
                    targets.emplace(func);
                }
            }
        }
        auto it = mapping.find(root.get());
        if (mapping.end() == it)
        {
            logs::fatal(
                "GraphStat failed to identify children of root subgraph");
        }
        return it->second;
    }

private:
    LeafFinder<T> finder_;

    PruneFuncT pruner_;
};

}

#endif // OPT_SHEAR_HPP