Skip to content

Commit

Permalink
Implement TT policies to control process mapping, priorities, and inl…
Browse files Browse the repository at this point in the history 
…ining

Policies replace keymap and priomaps and can be directly inlined if so
provided. Users of ttg::make_tt can use ttg::make_policy to define
policies with custom lambdas mapping keys to integers. Statically set
policies cannot be changed at runtime, all defaulted policies can be set.
Similarly, ttg::make_static_policy provides policies with statically
defaulted mapping functions, which cannot be changed at runtime.

This change does not implement task inlining but provides the policy to
control future implementations.

Signed-off-by: Joseph Schuchart <[email protected]>
  • Loading branch information
@devreal
devreal committed Apr 26, 2022
1 parent a05fc50 commit a311ac7
Show file tree
Hide file tree
Showing 14 changed files with 1,188 additions and 226 deletions.
17 changes: 11 additions & 6 deletions examples/madness/mrattg.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,8 @@ namespace detail {
using compress_out_type = std::tuple<Rout, Rout, cnodeOut<T, K, 1>>;
using compress_in_type = std::tuple<Rin, Rin>;
template <typename compfuncT>
using compwrap_type = ttg::CallableWrapTT<compfuncT, true, Key<1>, compress_out_type, Rin, Rin>;
using compwrap_type = ttg::CallableWrapTT<compfuncT, true, Key<1>, compress_out_type,
ttg::TTPolicyBase<Key<1>>, Rin, Rin>;
};

template <typename T, size_t K>
Expand All @@ -146,7 +147,8 @@ namespace detail {
using compress_out_type = std::tuple<Rout, Rout, Rout, Rout, cnodeOut<T, K, 2>>;
using compress_in_type = std::tuple<Rin, Rin, Rin, Rin>;
template <typename compfuncT>
using compwrap_type = ttg::CallableWrapTT<compfuncT, true, Key<2>, compress_out_type, Rin, Rin, Rin, Rin>;
using compwrap_type = ttg::CallableWrapTT<compfuncT, true, Key<2>, compress_out_type,
ttg::TTPolicyBase<Key<2>>, Rin, Rin, Rin, Rin>;
};

template <typename T, size_t K>
Expand All @@ -157,7 +159,8 @@ namespace detail {
using compress_in_type = std::tuple<Rin, Rin, Rin, Rin, Rin, Rin, Rin, Rin>;
template <typename compfuncT>
using compwrap_type =
ttg::CallableWrapTT<compfuncT, true, Key<3>, compress_out_type, Rin, Rin, Rin, Rin, Rin, Rin, Rin, Rin>;
ttg::CallableWrapTT<compfuncT, true, Key<3>, compress_out_type,
ttg::TTPolicyBase<Key<3>>, Rin, Rin, Rin, Rin, Rin, Rin, Rin, Rin>;
};
}; // namespace detail

Expand Down Expand Up @@ -237,7 +240,7 @@ auto make_compress(rnodeEdge<T, K, NDIM>& in, cnodeEdge<T, K, NDIM>& out, const
using sendfuncT = decltype(&send_leaves_up<T, K, NDIM>);
using sendwrapT =
ttg::CallableWrapTT<sendfuncT, true, Key<NDIM>, typename ::detail::tree_types<T, K, NDIM>::compress_out_type,
FunctionReconstructedNode<T, K, NDIM>>;
ttg::TTPolicyBase<Key<NDIM>>, FunctionReconstructedNode<T, K, NDIM>>;
using compfuncT = decltype(&do_compress<T, K, NDIM>);
using compwrapT = typename ::detail::tree_types<T, K, NDIM>::template compwrap_type<compfuncT>;

Expand All @@ -250,8 +253,10 @@ auto make_compress(rnodeEdge<T, K, NDIM>& in, cnodeEdge<T, K, NDIM>& out, const
outnames.push_back("output");

auto s =
std::unique_ptr<sendwrapT>(new sendwrapT(&send_leaves_up<T, K, NDIM>, "send_leaves_up", {"input"}, outnames));
auto c = std::unique_ptr<compwrapT>(new compwrapT(&do_compress<T, K, NDIM>, "do_compress", innames, outnames));
std::unique_ptr<sendwrapT>(new sendwrapT(&send_leaves_up<T, K, NDIM>, ttg::make_policy<Key<NDIM>>(),
"send_leaves_up", {"input"}, outnames));
auto c = std::unique_ptr<compwrapT>(new compwrapT(&do_compress<T, K, NDIM>, ttg::make_policy<Key<NDIM>>(),
"do_compress", innames, outnames));

in.set_out(s->template in<0>()); // Connect input to s
out.set_in(s->template out<num_children>()); // Connect s result to output
Expand Down
53 changes: 30 additions & 23 deletions examples/madness/mrattg_streaming.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ class LevelPmapX {
LevelPmapX(size_t nproc) : nproc(nproc) {}

/// Find the owner of a given key
HashValue operator()(const Key<3>& key) const {
int operator()(const Key<3>& key) const {
Level n = key.level();
if (n == 0) return 0;
madness::hashT hash;
Expand Down Expand Up @@ -112,7 +112,7 @@ class PartitionPmap {
}

/// Find the owner of a given key
HashValue operator()(const Key<NDIM>& key) const {
int operator()(const Key<NDIM>& key) const {
HashValue hash;
if (key.level() <= target_level) {
hash = key.hash();
Expand Down Expand Up @@ -313,12 +313,14 @@ auto make_start(const ctlEdge<NDIM>& ctl) {
/// Constructs an operator that adaptively projects the provided function into the basis

/// Returns an std::unique_ptr to the object
template <typename functorT, typename T, size_t K, Dimension NDIM>
template <typename functorT, typename T, size_t K, Dimension NDIM,
typename PmapT = ttg::detail::default_keymap<Key<NDIM>>>
auto make_project(functorT& f,
const T thresh, /// should be scalar value not complex
ctlEdge<NDIM>& ctl,
rnodeEdge<T,K,NDIM>& result,
const std::string& name = "project") {
const std::string& name = "project",
PmapT&& pmap = PmapT()) {

auto F = [f, thresh](const Key<NDIM>& key, std::tuple<ctlOut<NDIM>, rnodeOut<T,K,NDIM>>& out) {
FunctionReconstructedNodeWrap<T,K,NDIM> node(key); // Our eventual result
Expand Down Expand Up @@ -348,7 +350,8 @@ auto make_project(functorT& f,
ttg::send<1>(key, std::move(node), out); // always produce a result
};
ctlEdge<NDIM> refine("refine");
return ttg::make_tt(F, edges(fuse(refine, ctl)), ttg::edges(refine, result), name, {"control"}, {"refine", "result"});
return ttg::make_tt(F, edges(fuse(refine, ctl)), ttg::edges(refine, result),
ttg::make_static_policy(pmap), name, {"control"}, {"refine", "result"});
}

namespace detail {
Expand Down Expand Up @@ -471,13 +474,20 @@ std::string int2bitstring(size_t i, size_t width) {
}

/// Make a composite operator that implements compression for a single function
template <typename T, size_t K, Dimension NDIM>
auto make_compress(rnodeEdge<T,K,NDIM>& in, cnodeEdge<T,K,NDIM>& out, const std::string& name = "compress") {
template <typename T, size_t K, Dimension NDIM, typename PmapT = ttg::detail::default_keymap<Key<NDIM>>>
auto make_compress(rnodeEdge<T,K,NDIM>& in, cnodeEdge<T,K,NDIM>& out,
PmapT&& pmap, const std::string& name = "compress") {
rnodeEdge<T,K,NDIM> children1("children1"), children2("children2");

return std::make_tuple(ttg::make_tt(&send_leaves_up<T,K,NDIM>, edges(in), edges(children1, out), "send_leaves_up", {"input"}, {"children1", "output"}),
ttg::make_tt(&reduce_leaves<T,K,NDIM>, edges(children1), edges(children2), "reduce_leaves", {"children1"}, {"children2"}),
ttg::make_tt(&do_compress<T,K,NDIM>, edges(children2), edges(children1,out), "do_compress", {"children2"}, {"recur","output"}));
return std::make_tuple(ttg::make_tt(&send_leaves_up<T,K,NDIM>, edges(in), edges(children1, out),
ttg::make_static_policy(pmap),
"send_leaves_up", {"input"}, {"children1", "output"}),
ttg::make_tt(&reduce_leaves<T,K,NDIM>, edges(children1), edges(children2),
ttg::make_static_policy(pmap),
"reduce_leaves", {"children1"}, {"children2"}),
ttg::make_tt(&do_compress<T,K,NDIM>, edges(children2), edges(children1,out),
ttg::make_static_policy(pmap),
"do_compress", {"children2"}, {"recur","output"}));
}

template <typename T, size_t K, Dimension NDIM>
Expand Down Expand Up @@ -519,11 +529,13 @@ void do_reconstruct(const Key<NDIM>& key,
ttg::broadcast<1>(bcast_keys[1], std::move(r), out);
}

template <typename T, size_t K, Dimension NDIM>
auto make_reconstruct(const cnodeEdge<T,K,NDIM>& in, rnodeEdge<T,K,NDIM>& out, const std::string& name = "reconstruct") {
template <typename T, size_t K, Dimension NDIM, typename PmapT = ttg::detail::default_keymap<Key<NDIM>>>
auto make_reconstruct(const cnodeEdge<T,K,NDIM>& in, rnodeEdge<T,K,NDIM>& out,
PmapT&& pmap = PmapT(), const std::string& name = "reconstruct") {
ttg::Edge<Key<NDIM>,FixedTensor<T,K,NDIM>> S("S"); // passes scaling functions down

auto s = ttg::make_tt_tpl(&do_reconstruct<T,K,NDIM>, ttg::edges(in, S), ttg::edges(S, out), name, {"input", "s"}, {"s", "output"});
auto s = ttg::make_tt_tpl(&do_reconstruct<T,K,NDIM>, ttg::edges(in, S), ttg::edges(S, out),
ttg::make_static_policy(pmap), name, {"input", "s"}, {"s", "output"});

if (ttg::default_execution_context().rank() == 0) {
s->template in<1>()->send(Key<NDIM>{0,{0}}, FixedTensor<T,K,NDIM>()); // Prime the flow of scaling functions
Expand Down Expand Up @@ -764,16 +776,12 @@ void test2(size_t nfunc, T thresh = 1e-6) {
rnodeEdge<T,K,NDIM> a("a"), c("c");
cnodeEdge<T,K,NDIM> b("b");

auto p1 = make_project(ff, T(thresh), ctl, a, "project A");
p1->set_keymap(pmap);
auto p1 = make_project(ff, T(thresh), ctl, a, "project A", pmap);

auto compress = make_compress<T,K,NDIM>(a, b);
std::get<0>(compress)->set_keymap(pmap);
std::get<1>(compress)->set_keymap(pmap);
std::get<2>(compress)->set_keymap(pmap);
auto compress = make_compress<T,K,NDIM>(a, b, pmap);

auto &reduce_leaves_op = std::get<1>(compress);
reduce_leaves_op->template set_input_reducer<0>([](FunctionReconstructedNodeWrap<T,K,NDIM> &node,
reduce_leaves_op->template set_input_reducer<0>([](FunctionReconstructedNodeWrap<T,K,NDIM> &node,
const FunctionReconstructedNodeWrap<T,K,NDIM> &another)
{
//Update self values into the array.
Expand All @@ -786,8 +794,7 @@ void test2(size_t nfunc, T thresh = 1e-6) {
});
reduce_leaves_op->template set_static_argstream_size<0>(1 << NDIM);

auto recon = make_reconstruct<T,K,NDIM>(b,c);
recon->set_keymap(pmap);
auto recon = make_reconstruct<T,K,NDIM>(b,c,pmap);

//auto printer = make_printer(a,"projected ", true);
// auto printer2 = make_printer(b,"compressed ", false);
Expand All @@ -814,7 +821,7 @@ void test2(size_t nfunc, T thresh = 1e-6) {
//std::cout << "==== begin dot ====\n";
//std::cout << Dot()(start.get()) << std::endl;
//std::cout << "==== end dot ====\n";
beg = std::chrono::high_resolution_clock::now();
beg = std::chrono::high_resolution_clock::now();
// This kicks off the entire computation
start->invoke(Key<NDIM>(0, {0}));
}
Expand Down
9 changes: 4 additions & 5 deletions examples/spmm/spmm.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,7 @@ class Read_SpMatrix : public TT<Key<2>, std::tuple<Out<Key<2>, Blk>>, Read_SpMat
auto rank = ttg::default_execution_context().rank();
for (int k = 0; k < matrix_.outerSize(); ++k) {
for (typename SpMatrix<Blk>::InnerIterator it(matrix_, k); it; ++it) {
if (rank == this->get_keymap()(Key<2>({it.row(), it.col()})))
if (rank == this->procmap(Key<2>({it.row(), it.col()})))
::send<0>(Key<2>({it.row(), it.col()}), it.value(), out);
}
}
Expand Down Expand Up @@ -340,9 +340,8 @@ class SpMM {
if (k >= b_rowidx_to_colidx_.size()) return;
auto world = default_execution_context();
std::vector<bool> procmap(world.size());
auto keymap = baseT::get_keymap();
for (auto &j : b_rowidx_to_colidx_[k]) {
const long p = keymap(Key<2>({i, j}));
const long p = this->procmap(Key<2>({i, j}));
if (!procmap[p]) {
ttg::trace("Broadcasting A[", i, "][", k, "] to proc ", p);
ikp_keys.emplace_back(Key<3>({i, k, p}));
Expand Down Expand Up @@ -413,9 +412,9 @@ class SpMM {
auto world = default_execution_context();
std::vector<bool> procmap(world.size());
for (auto &i : a_colidx_to_rowidx_[k]) {
long p = baseT::get_keymap()(Key<2>({i, j}));
long p = baseT::procmap(Key<2>({i, j}));
if (!procmap[p]) {
ttg::trace("Broadcasting B[", k, "][", j, "] to proc ", p);
ttg::trace("Broadcasting A[", k, "][", j, "] to proc ", p);
kjp_keys.emplace_back(Key<3>({k, j, p}));
procmap[p] = true;
}
Expand Down
113 changes: 113 additions & 0 deletions tests/unit/tt.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,35 @@ namespace tt_i_iv {
};
} // namespace tt_i_iv

// {task_id,data} = {int, int, void}
namespace tt_i_i_p {

struct Policy : public ttg::TTPolicyBase<int> {

Policy() : TTPolicyBase() {}

int procmap(const int&) const { return 0; }

};

class tt : public ttg::TT<int, std::tuple<>, tt, ttg::typelist<const int>, Policy> {
using baseT = typename tt::ttT;

public:
tt(const typename baseT::input_edges_type &inedges, const typename baseT::output_edges_type &outedges,
const std::string &name)
: baseT(inedges, outedges, name, {"int"}, {}, Policy()) {}

static constexpr const bool have_cuda_op = false;

void op(const int &key, const baseT::input_refs_tuple_type &data, baseT::output_terminals_type &outs) {}

~tt() {}
};
} // namespace tt_i_iv



TEST_CASE("TemplateTask", "[core]") {
SECTION("constructors") {
{ // void task id, void data
Expand Down Expand Up @@ -249,5 +278,89 @@ TEST_CASE("TemplateTask", "[core]") {
},
ttg::edges(in), ttg::edges()));
}
{ // nonvoid task id, nonvoid data, w/ policies
ttg::Edge<int, int> in;
CHECK_NOTHROW(std::make_unique<tt_i_i_p::tt>(ttg::edges(in), ttg::edges(), ""));
CHECK_NOTHROW(
ttg::make_tt([](const int &key, const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges(),
tt_i_i_p::Policy()));

auto tt = ttg::make_tt([](const int &key, const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges(),
tt_i_i_p::Policy());
auto policy = tt->get_policy();
auto procmap = tt->get_procmap();
}
}

SECTION("policies") {
{ // default policy
ttg::Edge<int, int> in;
auto tt = ttg::make_tt([](const int &key, const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges());
auto policy = tt->get_policy();
CHECK(tt->procmap(0) == 0);
CHECK(tt->get_procmap()(0) == 0);
CHECK(policy.procmap(0) == 0);
}
{ // custom procmap
ttg::Edge<int, int> in;
auto tt = ttg::make_tt([](const int &key, const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges(),
ttg::make_policy([](const int& key){ return key*1000; }));
auto policy = tt->get_policy();
CHECK(tt->procmap(1) == 1000);
CHECK(tt->get_procmap()(2) == 2000);
CHECK(policy.procmap(3) == 3000);
tt->set_priomap([](const int&){ return 0; });
}
{ // custom all maps
ttg::Edge<int, int> in;
auto tt = ttg::make_tt([](const int &key, const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges(),
ttg::make_policy([](const int& key){ return key*1000; },
[](const int& key){ return key*1000; },
[](const int& key){ return key*1000; }));
auto policy = tt->get_policy();
CHECK(policy.procmap(1) == 1000);
CHECK(tt->procmap(2) == 2000);
CHECK(tt->get_procmap()(3) == 3000);

CHECK(tt->priomap(1) == 1000);
CHECK(tt->get_priomap()(2) == 2000);

CHECK(tt->inlinemap(1) == 1000);
CHECK(tt->get_inlinemap()(2) == 2000);
}
{ // custom all maps static
ttg::Edge<int, int> in;
auto tt = ttg::make_tt([](const int &key, const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges(),
ttg::make_static_policy([](const int& key){ return key*1000; },
[](const int& key){ return key*1000; },
[](const int& key){ return key*1000; }));
auto policy = tt->get_policy();
CHECK(policy.procmap(1) == 1000);
CHECK(tt->procmap(2) == 2000);
CHECK(tt->get_procmap()(3) == 3000);

CHECK(tt->priomap(1) == 1000);
CHECK(tt->get_priomap()(2) == 2000);

CHECK(tt->inlinemap(1) == 1000);
CHECK(tt->get_inlinemap()(2) == 2000);
}
{ // custom all maps static, void key
ttg::Edge<void, int> in;
auto tt = ttg::make_tt([](const int &datum, std::tuple<> &outs) {}, ttg::edges(in), ttg::edges(),
ttg::make_static_policy([](){ return 1000; },
[](){ return 2000; },
[](){ return 3000; }));
auto policy = tt->get_policy();
CHECK(policy.procmap() == 1000);
CHECK(tt->procmap() == 1000);
CHECK(tt->get_procmap()() == 1000);

CHECK(tt->priomap() == 2000);
CHECK(tt->get_priomap()() == 2000);

CHECK(tt->inlinemap() == 3000);
CHECK(tt->get_inlinemap()() == 3000);
}
}
}
1 change: 1 addition & 0 deletions ttg/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ set(ttg-impl-headers
${CMAKE_CURRENT_SOURCE_DIR}/ttg/func.h
${CMAKE_CURRENT_SOURCE_DIR}/ttg/fwd.h
${CMAKE_CURRENT_SOURCE_DIR}/ttg/impl_selector.h
${CMAKE_CURRENT_SOURCE_DIR}/ttg/policies.h
${CMAKE_CURRENT_SOURCE_DIR}/ttg/tt.h
${CMAKE_CURRENT_SOURCE_DIR}/ttg/reduce.h
${CMAKE_CURRENT_SOURCE_DIR}/ttg/run.h
Expand Down
14 changes: 14 additions & 0 deletions ttg/ttg/base/keymap.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,20 @@ namespace ttg {
operator()() const { return 0; }
};


/// the default inline map implementation
template <typename keyT>
struct default_inlinemap_impl {
default_inlinemap_impl() = default;

template <typename Key = keyT>
std::enable_if_t<!meta::is_void_v<Key>,int>
operator()(const Key &key) const { return 0; }
template <typename Key = keyT>
std::enable_if_t<meta::is_void_v<Key>,int>
operator()() const { return 0; }
};

} // namespace detail

} // namespace ttg
Expand Down
5 changes: 4 additions & 1 deletion ttg/ttg/madness/fwd.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,15 @@

#include "ttg/fwd.h"
#include "ttg/util/typelist.h"
#include "ttg/policies.h"

#include <future>

namespace ttg_madness {

template <typename keyT, typename output_terminalsT, typename derivedT, typename input_valueTs = ttg::typelist<>>
template <typename keyT, typename output_terminalsT, typename derivedT,
typename input_valueTs = ttg::typelist<>,
typename policiesT = ttg::TTPolicyBase<keyT>>
class TT;

/// \internal the OG name
Expand Down
Loading

0 comments on commit a311ac7

Please sign in to comment.