Skip to content

GitLab

Commit 0b2c59f1 authored by umang yadav's avatar umang yadav
Browse files

Use struct to keep state

parent 3633f1a8
Show whitespace changes
Inline Side-by-side
Showing with 365 additions and 244 deletions
src/CMakeLists.txt
View file @ 0b2c59f1
......@@ -71,7 +71,7 @@ add_library(migraphx
operation.cpp
optimize_module.cpp
pad_calc.cpp
partitioner.cpp
generate_root_modules.cpp
pass_manager.cpp
permutation.cpp
preallocate_param.cpp
......
src/partitioner.cpp → src/generate_root_modules.cpp
View file @ 0b2c59f1
......@@ -21,7 +21,6 @@
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/
#include "migraphx/target_assignments.hpp"
#include <cstddef>
#include <limits>
#include <iterator>
......@@ -31,7 +30,7 @@
#include <migraphx/env.hpp>
#include <migraphx/algorithm.hpp>
#include <migraphx/stringutils.hpp>
#include <migraphx/partitioner.hpp>
#include <migraphx/generate_root_modules.hpp>
#include <migraphx/pass_manager.hpp>
#include <migraphx/dead_code_elimination.hpp>
#include <migraphx/instruction.hpp>
......@@ -44,6 +43,7 @@ MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DEBUG_PARTITIONER)
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
// copied from fuse_pointwise.cpp
static literal get_scalar(instruction_ref ins)
{
if(ins->name() == "contiguous")
......@@ -68,8 +68,48 @@ static literal get_scalar(instruction_ref ins)
return r;
}
void update_tid_counter(std::size_t tid, std::unordered_map<std::size_t, std::size_t>& tid_counter)
/*
Given target assignments (tass) for the instructions, generate run_on_target modules subgraphs
automatically. Input graph should be uncompiled migraphx program. target assignments (tass) map
should have a map of instruction to target_id. Instructions that are not inside tass map are
considered to be targeted for the "Ref" by default. params, literals and other builtins shouldn't be
part of the tass, only compute and reshaper instructions should be part of tass. Copy, sync and
alloc instructions would be generated by compiler at later stage, so those shouldn't be considered.
(TODO): CustomOps may require special handling.
Step 1:
Identify subgraph boundaries
Ref is used for instructions that do not have assignments.
Boundaries can happen in following cases.
1. Ref --> Target X --> Ref
2. Ref --> Target X --> Target Y
3. Target X --> Target Y --> Target Z , in this case target X and target Z can be same
4. Target X --> "@return"
5. Target X --> Ref --> "@return"
Each of those identified regions could have futher nested sub modules which needs to be handled
separately.
Step 2:
Collect parameters and return instructions for the subgraphs identified in Step 1.
Step 3:
Create modules using information collected in step 2 and insert run_on_target instructions.
*/
struct AutoGenRootModules
{
AutoGenRootModules(migraphx::program& p, const target_assignments& target_assignments)
: tass(target_assignments)
{
auto* mm = p.get_main_module();
find_subgraphs(mm, p);
dead_code_elimination{}.apply(p);
}
void update_tid_counter(std::size_t tid)
{
assert(tid != std::numeric_limits<std::size_t>::max());
if(tid_counter.find(tid) != tid_counter.end())
{
......@@ -79,18 +119,87 @@ void update_tid_counter(std::size_t tid, std::unordered_map<std::size_t, std::si
{
tid_counter[tid] = 0;
}
}
}
void find_subgraphs(migraphx::module_ref mm, migraphx::program& p)
{
// sort the graph in reverse post order DFS order
mm->sort();
if(enabled(MIGRAPHX_DEBUG_PARTITIONER{}))
{
std::cout << "sorted module: \n";
mm->debug_print();
}
size_t current_tid = std::numeric_limits<std::size_t>::max();
for(auto ins : iterator_for(*mm))
{
if(enabled(MIGRAPHX_DEBUG_PARTITIONER{}))
{
std::cout << "looking at instruction: \n";
ins->debug_print();
std::cout << "\n";
}
if(ins->name() == "@return")
{
generate_run_on_target_modules(mm, p, ins, current_tid);
}
// skip all params, literal and builtins other than return, skip "run_on_target_mod"
// ins
else if(starts_with(ins->name(), "@") or skip_ins.count(ins) != 0)
{
continue;
}
else if(tass.find(ins) == tass.end())
{
generate_run_on_target_modules(mm, p, ins, current_tid);
}
else if(current_tid == std::numeric_limits<std::size_t>::max())
{
current_tid = tass.at(ins);
update_tid_counter(current_tid);
same_tid_ins_vec.push_back(ins);
same_tid_ins_set.insert(ins);
}
else if(tass.at(ins) == current_tid)
{
same_tid_ins_vec.push_back(ins);
same_tid_ins_set.insert(ins);
}
else if(tass.at(ins) != current_tid)
{
generate_run_on_target_modules(mm, p, ins, current_tid);
}
else
{
MIGRAPHX_THROW("Partition: this case shouldn't occur");
}
void generate_run_on_target_modules(migraphx::module_ref mm,
if(skip_ins.find(ins) == skip_ins.end() and not ins->module_inputs().empty())
{
std::vector<instruction_ref> same_tid_ins_vec_copy = {};
std::unordered_set<instruction_ref> same_tid_ins_set_copy = {};
std::swap(same_tid_ins_set_copy, same_tid_ins_set);
std::swap(same_tid_ins_vec_copy, same_tid_ins_vec);
for(auto sub_mod : ins->module_inputs())
{
find_subgraphs(sub_mod, p);
}
std::swap(same_tid_ins_set_copy, same_tid_ins_set);
std::swap(same_tid_ins_vec_copy, same_tid_ins_vec);
mm->replace_instruction(
ins, ins->get_operator(), ins->inputs(), ins->module_inputs());
}
}
assert(same_tid_ins_set.empty() and same_tid_ins_vec.empty());
}
void generate_run_on_target_modules(migraphx::module_ref mm,
migraphx::program& p,
migraphx::instruction_ref ins,
std::size_t& current_tid,
const target_assignments& tass,
std::unordered_set<instruction_ref>& skip_ins,
std::unordered_map<std::size_t, std::size_t>& tid_counter,
std::vector<instruction_ref>& same_tid_ins_vec,
std::unordered_set<instruction_ref>& same_tid_ins_set)
{
std::size_t& current_tid)
{
assert(same_tid_ins_vec.size() == same_tid_ins_set.size());
if(same_tid_ins_vec.empty())
{
......@@ -127,8 +236,7 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
{
(*tmp)->debug_print();
}
std::cout << "\n";
std::cout << "return ins: \n";
std::cout << "\n return ins: \n";
for(auto tmp : iterator_for(return_ins))
{
(*tmp)->debug_print();
......@@ -136,9 +244,10 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
std::cout << "\n";
}
std::cout << "1\n";
auto* tmod = p.create_module("target_mod_" + std::to_string(current_tid) + "_" +
std::to_string(tid_counter[current_tid]));
update_tid_counter(current_tid, tid_counter);
update_tid_counter(current_tid);
std::unordered_map<instruction_ref, instruction_ref> params_map;
std::size_t param_counter = 0;
std::vector<instruction_ref> rot_ins_params;
......@@ -156,7 +265,9 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
params_map[*pins] = tmod->add_literal(scalar);
}
}
std::cout << "2\n";
// TODO: what if params_map is empty ?
assert(not params_map.empty());
for(auto tins : iterator_for(same_tid_ins_vec))
{
auto inputs = (*tins)->inputs();
......@@ -165,23 +276,25 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
inputs.end(),
std::back_inserter(new_inputs),
[&](auto input_ins) { return params_map.at(input_ins); });
// [TODO]: what if it is has module args ?
auto tmod_tins =
tmod->add_instruction((*tins)->get_operator(), new_inputs, (*tins)->module_inputs());
auto tmod_tins = tmod->add_instruction(
(*tins)->get_operator(), new_inputs, (*tins)->module_inputs());
// add parameter to params map so that it can be looked up by other insturctions
params_map[*tins] = tmod_tins;
}
std::cout << "3\n";
std::vector<instruction_ref> rins;
std::unordered_map<instruction_ref, std::size_t> return_ins_idx_map;
std::cout << "4\n";
for(auto ritr : iterator_for(return_ins))
{
rins.push_back(params_map.at(*ritr));
return_ins_idx_map[*ritr] = std::distance(ritr, return_ins.begin());
}
tmod->add_return(rins);
std::cout << "5\n";
if(enabled(MIGRAPHX_DEBUG_PARTITIONER{}))
{
std::cout << "tmod: \n";
std::cout << "Created target module: " << tmod->name() << "\n";
tmod->debug_print();
}
// add run_on_target ins
......@@ -192,7 +305,9 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
for(auto mm_rins : iterator_for(return_ins))
{
auto tuple_elem_ins = mm->insert_instruction(
ins, make_op("get_tuple_elem", {{"index", return_ins_idx_map.at(*mm_rins)}}), rot_ins);
ins,
make_op("get_tuple_elem", {{"index", return_ins_idx_map.at(*mm_rins)}}),
rot_ins);
skip_ins.insert(tuple_elem_ins);
// replace returns from tmod in main module
mm->replace_instruction(*mm_rins, tuple_elem_ins);
......@@ -204,7 +319,7 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
if(tass.find(ins) != tass.end())
{
current_tid = tass.at(ins);
update_tid_counter(current_tid, tid_counter);
update_tid_counter(current_tid);
same_tid_ins_set.insert(ins);
same_tid_ins_vec.push_back(ins);
}
......@@ -214,146 +329,22 @@ void generate_run_on_target_modules(migraphx::module_ref mm,
}
if(enabled(MIGRAPHX_DEBUG_PARTITIONER{}))
{
std::cout << "module after creation of tmod and rot: \n";
std::cout << "Main module after creation of target module: " << tmod->name() << "\n";
mm->debug_print();
}
}
/*
Given target assignments (tass) for the instructions, generate run_on_target modules subgraphs
automatically. Input graph should be uncompiled migraphx program. target assignments (tass) map
should have a map of instruction to target_id. Instructions that are not inside tass map are
considered to be targeted for the "Ref" by default. params, literals and other builtins shouldn't be
part of the tass, only compute and reshaper instructions should be part of tass. Copy, sync and
alloc instructions would be generated by compiler at later stage, so those shouldn't be considered.
(TODO): CustomOps may require special handling.
Step 1:
Identify subgraph boundaries
Ref is used for instructions that do not have assignments.
Boundaries can happen in following cases.
1. Ref --> Target X --> Ref
2. Ref --> Target X --> Target Y
3. Target X --> Target Y --> Target Z , in this case target X and target Z can be same
4. Target X --> "@return"
5. Target X --> Ref --> "@return"
Each of those identified regions could have futher nested sub modules which needs to be handled
separately.
Step 2:
Collect parameters and return instructions for the subgraphs identified in Step 1.
Step 3:
Create modules using information collected in step 2 and insert run_on_target instructions.
*/
void partition(migraphx::module_ref mm,
migraphx::program& p,
const target_assignments& tass,
std::unordered_map<std::size_t, std::size_t>& tid_counter)
{
mm->sort();
if(enabled(MIGRAPHX_DEBUG_PARTITIONER{}))
{
std::cout << "sorted module: \n";
mm->debug_print();
}
private:
const target_assignments tass;
std::unordered_map<std::size_t, std::size_t> tid_counter;
std::unordered_set<instruction_ref> skip_ins;
std::vector<instruction_ref> same_tid_ins_vec;
std::unordered_set<instruction_ref> same_tid_ins_set;
// walk the graph in reverse-DFS order
size_t current_tid = std::numeric_limits<std::size_t>::max();
std::unordered_set<instruction_ref> skip_ins;
for(auto ins : iterator_for(*mm))
{
// gather instructions belonging to the same target_id
// for now, make sure that all the inputs to the insturctions are also from the same
// target_id, if not create another module
// skip all the builtins
if(enabled(MIGRAPHX_DEBUG_PARTITIONER{}))
{
std::cout << "currently processing: \n";
ins->debug_print();
std::cout << "\n";
}
if(skip_ins.count(ins) == 0)
{
if(not ins->module_inputs().empty())
{
for(auto sub_mod : ins->module_inputs())
{
partition(sub_mod, p, tass, tid_counter);
}
mm->replace_instruction(
ins, ins->get_operator(), ins->inputs(), ins->module_inputs());
}
}
if(ins->name() == "@return")
{
generate_run_on_target_modules(mm,
p,
ins,
current_tid,
tass,
skip_ins,
tid_counter,
same_tid_ins_vec,
same_tid_ins_set);
}
// skip all params, literal and builitins other than return, skip "run_on_target_mod" ins
else if(starts_with(ins->name(), "@") or skip_ins.count(ins) != 0)
{
continue;
}
else if(tass.find(ins) == tass.end())
{
generate_run_on_target_modules(mm,
p,
ins,
current_tid,
tass,
skip_ins,
tid_counter,
same_tid_ins_vec,
same_tid_ins_set);
}
else if(current_tid == std::numeric_limits<std::size_t>::max())
{
current_tid = tass.at(ins);
update_tid_counter(current_tid, tid_counter);
same_tid_ins_vec.push_back(ins);
same_tid_ins_set.insert(ins);
}
else if(tass.at(ins) == current_tid)
{
same_tid_ins_vec.push_back(ins);
same_tid_ins_set.insert(ins);
}
else if(tass.at(ins) != current_tid)
{
generate_run_on_target_modules(mm,
p,
ins,
current_tid,
tass,
skip_ins,
tid_counter,
same_tid_ins_vec,
same_tid_ins_set);
}
else
{
MIGRAPHX_THROW("Partition: this shouldn't occur");
}
}
}
};
void partition(migraphx::program& p, const target_assignments& tass)
void generate_root_modules(migraphx::program& p, const target_assignments& tass)
{
auto* mm = p.get_main_module();
// sort the graph in reverse post order DFS order
std::unordered_map<std::size_t, std::size_t> tid_counter;
partition(mm, p, tass, tid_counter);
dead_code_elimination{}.apply(p);
AutoGenRootModules(p, tass);
}
} // namespace MIGRAPHX_INLINE_NS
......
src/include/migraphx/partitioner.hpp → src/include/migraphx/generate_root_modules.hpp
View file @ 0b2c59f1
......@@ -30,9 +30,9 @@
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
/*
given target_assignments, paritions the migraphx program into separate root modules.
given target_assignments, generates root modules for each individual targets inside main module.
*/
void partition(migraphx::program& p, const migraphx::target_assignments& tass);
void generate_root_modules(migraphx::program& p, const migraphx::target_assignments& tass);
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx
......
test/multi_target/multitarget_test.cpp
View file @ 0b2c59f1
......@@ -41,7 +41,7 @@
#include <migraphx/compile_options.hpp>
#include <migraphx/register_target.hpp>
#include <migraphx/generate.hpp>
#include "migraphx/partitioner.hpp"
#include <migraphx/generate_root_modules.hpp>
#include "migraphx/target_assignments.hpp"
#include "test.hpp"
......@@ -174,7 +174,7 @@ TEST_CASE(multitarget_compile_cpu_gpu)
migraphx::target_assignments tass;
tass.insert(tass.begin(), std::make_pair(cpu_ins, 1));
tass.insert(tass.begin(), std::make_pair(gpu_ins, 0));
migraphx::partition(p, tass);
migraphx::generate_root_modules(p, tass);
migraphx::compile_options gpu_opts;
gpu_opts.offload_copy = true;
p.compile({migraphx::make_target("gpu"), migraphx::make_target("cpu")}, {gpu_opts});
......@@ -210,10 +210,10 @@ TEST_CASE(single_target_multi_compile)
iou_threshold,
score_threshold);
mm->add_return({r});
// do partition
// do target assignments
migraphx::target_assignments tass;
tass.insert(tass.begin(), std::make_pair(r, 0));
migraphx::partition(p, tass);
migraphx::generate_root_modules(p, tass);
// compile using multi-target compilation path
migraphx::compile_options gpu_opts;
gpu_opts.offload_copy = true;
......@@ -232,7 +232,7 @@ TEST_CASE(single_target_multi_compile)
EXPECT(output == gold);
}
TEST_CASE(multitarget_compile_if_then_else_partition)
TEST_CASE(multitarget_compile_if_then_else)
{
migraphx::program p;
auto* mm = p.get_main_module();
......@@ -264,8 +264,7 @@ TEST_CASE(multitarget_compile_if_then_else_partition)
tass.insert(tass.begin(), std::make_pair(a1, 0));
tass.insert(tass.begin(), std::make_pair(l2, 1));
tass.insert(tass.begin(), std::make_pair(a2, 1));
migraphx::partition(p, tass);
p.debug_print();
migraphx::generate_root_modules(p, tass);
// compile
migraphx::compile_options gpu_opts;
gpu_opts.offload_copy = true;
......@@ -288,230 +287,8 @@ TEST_CASE(multitarget_compile_if_then_else_partition)
}
}
TEST_CASE(multitarget_compile_if_then_else)
{
migraphx::program p;
auto* mm = p.get_main_module();
migraphx::shape cond_s{migraphx::shape::bool_type};
auto cond = mm->add_parameter("cond", cond_s);
migraphx::shape ds{migraphx::shape::float_type, {2, 3}};
auto x = mm->add_parameter("x", ds);
auto y = mm->add_parameter("y", ds);
auto* then_mod = p.create_module("if_gpu_mod");
std::vector<float> data1(ds.elements(), 1);
auto l1 = then_mod->add_literal(migraphx::literal(ds, data1));
auto gpu_x = then_mod->add_parameter("gpu_x", ds);
auto a1 = then_mod->add_instruction(migraphx::make_op("add"), gpu_x, l1);
then_mod->add_return({a1});
auto* else_mod = p.create_module("else_cpu_mod");
std::vector<float> data2(ds.elements(), 2);
auto l2 = else_mod->add_literal(migraphx::literal(ds, data2));
auto cpu_y = else_mod->add_parameter("cpu_y", ds);
auto a2 = else_mod->add_instruction(migraphx::make_op("mul"), cpu_y, l2);
else_mod->add_return({a2});
auto* run_on_cpu_mod = p.create_module("run_on_cpu");
auto run_cpu_ins = run_on_cpu_mod->add_instruction(
migraphx::make_op("run_on_target", {{"target_id", 1}}), {y}, {else_mod});
auto run_cpu_ins_0 = run_on_cpu_mod->add_instruction(
migraphx::make_op("get_tuple_elem", {{"index", 0}}), run_cpu_ins);
run_on_cpu_mod->add_return({run_cpu_ins_0});
auto* run_on_gpu_mod = p.create_module("run_on_gpu");
auto run_gpu_ins = run_on_gpu_mod->add_instruction(
migraphx::make_op("run_on_target", {{"target_id", 0}}), {x}, {then_mod});
auto run_gpu_ins_0 = run_on_gpu_mod->add_instruction(
migraphx::make_op("get_tuple_elem", {{"index", 0}}), run_gpu_ins);
run_on_gpu_mod->add_return({run_gpu_ins_0});
auto ret =
mm->add_instruction(migraphx::make_op("if"), {cond}, {run_on_gpu_mod, run_on_cpu_mod});
auto r = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), ret);
mm->add_return({r});
// compile
migraphx::compile_options gpu_opts;
gpu_opts.offload_copy = true;
p.compile({migraphx::make_target("gpu"), migraphx::make_target("cpu")}, {gpu_opts});
EXPECT(check_compiled_program(p, {migraphx::make_target("gpu"), migraphx::make_target("cpu")}));
migraphx::parameter_map params;
params["x"] = migraphx::fill_argument(ds, 2);
params["y"] = migraphx::fill_argument(ds, 3);
for(bool cond_val : {true, false})
{
params["cond"] = migraphx::argument(cond_s, &cond_val);
auto result = p.eval(params).back();
auto gold = migraphx::fill_argument(ds, (cond_val ? 3 : 6));
EXPECT(gold == result);
}
}
// TODO : FPGA compilation is broken right now, below test mentions fpga but doesn't compile for it
TEST_CASE(multitarget_compile_nested_if_then_else)
{
std::unordered_map<std::size_t, std::size_t> counter_map = {{0, 0}, {1, 0}};
migraphx::shape ds{migraphx::shape::float_type, {2, 3}};
migraphx::program p;
auto* mm = p.get_main_module();
migraphx::shape cond_s{migraphx::shape::bool_type};
auto cond_0 = mm->add_parameter("cond_0", cond_s);
auto cond_1 = mm->add_parameter("cond_1", cond_s);
auto x = mm->add_parameter("x", ds);
auto y = mm->add_parameter("y", ds);
auto z = mm->add_parameter("z", ds);
auto create_test_module = [&](migraphx::program& prog,
const std::vector<migraphx::instruction_ref>& inputs,
std::size_t tid) {
std::string mod_name =
"target_" + std::to_string(tid) + "_" + std::to_string(counter_map[tid]++);
auto* test_mod = prog.create_module(mod_name);
std::vector<float> data(ds.elements(), -1);
auto l1 = test_mod->add_literal(migraphx::literal(ds, data));
auto test_mod_param_0 = test_mod->add_parameter(mod_name + "_param_0", ds);
auto test_mod_param_1 = test_mod->add_parameter(mod_name + "_param_1", ds);
auto test_mod_param_2 = test_mod->add_parameter(mod_name + "_param_2", ds);
auto ins1 = test_mod->add_instruction(migraphx::make_op("add"), test_mod_param_0, l1);
auto ins2 = test_mod->add_instruction(migraphx::make_op("mul"), ins1, test_mod_param_1);
auto ins3 = test_mod->add_instruction(migraphx::make_op("sub"), ins2, test_mod_param_2);
test_mod->add_return({ins3});
auto* run_on_target_mod = prog.create_module("run_on_" + mod_name);
auto run_ins = run_on_target_mod->add_instruction(
migraphx::make_op("run_on_target", {{"target_id", tid}}), inputs, {test_mod});
auto run_ins_0 = run_on_target_mod->add_instruction(
migraphx::make_op("get_tuple_elem", {{"index", 0}}), run_ins);
run_on_target_mod->add_return({run_ins_0});
return run_on_target_mod;
};
// create nested module with multiple targets.
// then_mod has one instruction that runs a module on "ref" and another instruction that
// creates nested modules using "If" that runs on "cpu" and "gpu"
auto* ref_mod = p.create_module("ref_mod");
auto ref_x = ref_mod->add_parameter("ref_x", ds);
auto ref_y = ref_mod->add_parameter("ref_y", ds);
auto ref_add = ref_mod->add_instruction(migraphx::make_op("add"), ref_x, ref_y);
ref_mod->add_return({ref_add});
auto* then_mod = p.create_module("then_mod");
auto then_mod_cond = then_mod->add_parameter("then_mod_cond", cond_s);
auto then_mod_param_0 = then_mod->add_parameter("then_mod_param_0", ds);
auto then_mod_param_1 = then_mod->add_parameter("then_mod_param_1", ds);
auto then_mod_param_2 = then_mod->add_parameter("then_mod_param_2", ds);
auto then_mod_ref_ins =
then_mod->add_instruction(migraphx::make_op("run_on_target", {{"target_id", 3}}),
{then_mod_param_0, then_mod_param_1},
{ref_mod});
auto then_mod_ref_ins_0 = then_mod->add_instruction(
migraphx::make_op("get_tuple_elem", {{"index", 0}}), then_mod_ref_ins);
auto then_mod_if = then_mod->add_instruction(
migraphx::make_op("if"),
{then_mod_cond,
then_mod_param_0,
then_mod_param_1,
then_mod_param_2,
then_mod_ref_ins_0,
then_mod_param_1,
then_mod_param_2},
{create_test_module(p, {then_mod_param_0, then_mod_param_1, then_mod_param_2}, 1),
create_test_module(p, {then_mod_ref_ins_0, then_mod_param_1, then_mod_param_2}, 0)});
auto then_mod_if_0 =
then_mod->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), then_mod_if);
then_mod->add_return({then_mod_if_0});
// create nested else_mod with multiple targets.
// else_mod has one instruction that runs a module on "fpga" and another instruction that
// creates nested modules using "If" that runs on "cpu" and "gpu"
auto* fpga_mod = p.create_module("fpga_mod");
auto fpga_x = fpga_mod->add_parameter("fpga_x", ds);
auto fpga_y = fpga_mod->add_parameter("fpga_y", ds);
auto fpga_add = fpga_mod->add_instruction(migraphx::make_op("add"), fpga_x, fpga_y);
fpga_mod->add_return({fpga_add});
auto* else_mod = p.create_module("else_mod");
auto else_mod_cond = else_mod->add_parameter("else_mod_cond", cond_s);
auto else_mod_param_0 = else_mod->add_parameter("else_mod_param_0", ds);
auto else_mod_param_1 = else_mod->add_parameter("else_mod_param_1", ds);
auto else_mod_param_2 = else_mod->add_parameter("else_mod_param_2", ds);
auto else_mod_fpga_ins =
else_mod->add_instruction(migraphx::make_op("run_on_target", {{"target_id", 2}}),
{else_mod_param_0, else_mod_param_2},
{fpga_mod});
auto else_mod_fpga_ins_0 = else_mod->add_instruction(
migraphx::make_op("get_tuple_elem", {{"index", 0}}), else_mod_fpga_ins);
auto else_mod_if = else_mod->add_instruction(
migraphx::make_op("if"),
{else_mod_cond,
else_mod_fpga_ins_0,
else_mod_param_0,
else_mod_param_1,
else_mod_param_2,
else_mod_param_1,
else_mod_param_0},
{create_test_module(p, {else_mod_fpga_ins_0, else_mod_param_0, else_mod_param_1}, 0),
create_test_module(p, {else_mod_param_2, else_mod_param_1, else_mod_param_0}, 1)});
auto else_mod_if_0 =
else_mod->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), else_mod_if);
else_mod->add_return({else_mod_if_0});
// Create nested and multi-target main module using "If"
auto main_if_ins = mm->add_instruction(
migraphx::make_op("if"), {cond_0, cond_1, x, y, z, cond_1, x, y, z}, {then_mod, else_mod});
auto r = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), main_if_ins);
mm->add_return({r});
// compile
migraphx::compile_options gpu_opts;
gpu_opts.offload_copy = true;
p.compile({migraphx::make_target("gpu"),
migraphx::make_target("cpu"),
migraphx::make_target("ref"),
migraphx::make_target("ref")},
{gpu_opts});
EXPECT(check_compiled_program(p,
{migraphx::make_target("gpu"),
migraphx::make_target("cpu"),
migraphx::make_target("ref"),
migraphx::make_target("ref")}));
// do evaluation using different conditions
migraphx::parameter_map params;
float x_i = 2.0;
float y_i = 3.0;
float z_i = 4.0;
params["x"] = migraphx::fill_argument(ds, x_i);
params["y"] = migraphx::fill_argument(ds, y_i);
params["z"] = migraphx::fill_argument(ds, z_i);
// cover all paths with different combination of conditions
std::vector<std::pair<bool, bool>> test_conds = {
{true, true}, {true, false}, {false, true}, {false, false}};
for(auto [cond_val_0, cond_val_1] : test_conds)
{
params["cond_0"] = migraphx::argument(cond_s, &cond_val_0);
params["cond_1"] = migraphx::argument(cond_s, &cond_val_1);
auto result = p.eval(params).back();
// main has one instruction that is : if_then_else
// then mod is doing : {tmp = x+y; (cond) ? (((x-1)*y)-z) : (((tmp-1)*y)-z);}
// else mod is doing : {tmp = x+z; (cond) ? (((tmp-1)*x)-y) : (((z-1)*y)-x);}
float gold_i = -1.0;
if(cond_val_0)
{
float tmp_i = x_i + y_i;
gold_i = (cond_val_1) ? (((x_i - 1) * y_i) - z_i) : (((tmp_i - 1) * y_i) - z_i);
}
else
{
float tmp_i = x_i + z_i;
gold_i = (cond_val_1) ? (((tmp_i - 1) * x_i) - y_i) : (((z_i - 1) * y_i) - x_i);
}
auto gold = migraphx::fill_argument(ds, gold_i);
EXPECT(gold == result);
}
}
// TODO : FPGA compilation is broken right now, below test mentions fpga but doesn't compile for it
TEST_CASE(multitarget_compile_nested_if_then_else_partition)
{
std::unordered_map<std::size_t, std::size_t> counter_map = {{0, 0}, {1, 0}};
migraphx::shape ds{migraphx::shape::float_type, {2, 3}};
......@@ -525,7 +302,6 @@ TEST_CASE(multitarget_compile_nested_if_then_else_partition)
auto y = mm->add_parameter("y", ds);
auto z = mm->add_parameter("z", ds);
auto create_test_module = [&](migraphx::program& prog,
const std::vector<migraphx::instruction_ref>& inputs,
std::size_t tid) {
std::string mod_name =
"target_" + std::to_string(tid) + "_" + std::to_string(counter_map[tid]++);
......@@ -562,8 +338,8 @@ TEST_CASE(multitarget_compile_nested_if_then_else_partition)
then_mod_ref_ins,
then_mod_param_1,
then_mod_param_2},
{create_test_module(p, {then_mod_param_0, then_mod_param_1, then_mod_param_2}, 1),
create_test_module(p, {then_mod_ref_ins, then_mod_param_1, then_mod_param_2}, 0)});
{create_test_module(p, 1),
create_test_module(p, 0)});
auto then_mod_if_0 =
then_mod->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), then_mod_if);
then_mod->add_return({then_mod_if_0});
......@@ -588,8 +364,8 @@ TEST_CASE(multitarget_compile_nested_if_then_else_partition)
else_mod_param_2,
else_mod_param_1,
else_mod_param_0},
{create_test_module(p, {else_mod_fpga_ins, else_mod_param_0, else_mod_param_1}, 0),
create_test_module(p, {else_mod_param_2, else_mod_param_1, else_mod_param_0}, 1)});
{create_test_module(p, 0),
create_test_module(p, 1)});
auto else_mod_if_0 =
else_mod->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), else_mod_if);
else_mod->add_return({else_mod_if_0});
......@@ -603,18 +379,12 @@ TEST_CASE(multitarget_compile_nested_if_then_else_partition)
// compile
migraphx::compile_options gpu_opts;
gpu_opts.offload_copy = true;
std::cout << "before parition\n";
p.debug_print();
migraphx::partition(p, tass);
std::cout << "after partition\n";
p.debug_print();
migraphx::generate_root_modules(p, tass);
p.compile({migraphx::make_target("gpu"),
migraphx::make_target("cpu"),
migraphx::make_target("ref"),
migraphx::make_target("ref")},
{gpu_opts});
std::cout << "after compilation\n";
p.debug_print();
EXPECT(check_compiled_program(p,
{migraphx::make_target("gpu"),
migraphx::make_target("cpu"),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment