Update pass manager to handle multi-target compilation (#1672)

partially solves #1656
This PR only handles compilation part of multitarget.

Jenkinsfile

@@ -117,6 +117,10 @@ rocmtest clang_debug: rocmnode('vega') { cmake_build ->
     stage('HipRTC GPU Debug') {
         cmake_build(flags: "-DCMAKE_BUILD_TYPE=release -DMIGRAPHX_USE_HIPRTC=On", gpu_debug: true, hiprtc_workarounds:  true)
     }
+}, all_targets_debug : rocmnode('vega') { cmake_build ->
+    stage('All targets Release') {
+        cmake_build(flags: "-DCMAKE_BUILD_TYPE=release -DMIGRAPHX_ENABLE_GPU=On -DMIGRAPHX_ENABLE_CPU=On -DMIGRAPHX_ENABLE_FPGA=On") 
+    }
 }, mlir_debug: rocmnode('vega') { cmake_build ->
     stage('MLIR Debug') {
         withEnv(['MIGRAPHX_ENABLE_MLIR=1']) {

src/CMakeLists.txt

@@ -195,6 +195,7 @@ register_migraphx_ops(
     roialign
     round
     rsqrt
+    run_on_target
     scalar
     scatter_add
     scatter_mul

src/include/migraphx/instruction.hpp

@@ -136,6 +136,9 @@ struct instruction
 
     operation normalized_operator() const;
 
+    std::size_t get_target_id() const;
+    void set_target_id(std::size_t tid);
+
     void debug_print() const;
 
     static void print(std::ostream& os,
@@ -173,6 +176,7 @@ struct instruction
     std::vector<module_ref> module_args;
     literal lit;
     bool normalized       = false;
+    std::size_t target_id = 0;
 };
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/include/migraphx/op/run_on_target.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#ifndef MIGRAPHX_GUARD_RTGLIB_RUN_ON_TARGET_HPP
+#define MIGRAPHX_GUARD_RTGLIB_RUN_ON_TARGET_HPP
+#include <unordered_map>
+#include <vector>
+#include <set>
+#include <algorithm>
+#include <migraphx/config.hpp>
+#include <migraphx/errors.hpp>
+#include <migraphx/shape.hpp>
+#include <migraphx/argument.hpp>
+#include <migraphx/module.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+namespace op {
+struct run_on_target
+{
+
+    std::size_t target_id = 0;
+    std::string name() const { return "run_on_target"; }
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.target_id, "target_id"));
+    }
+
+    migraphx::shape compute_shape(const std::vector<migraphx::shape>& inputs,
+                                  std::vector<migraphx::module_ref> mods) const
+    {
+        if(mods.size() != 1)
+        {
+            MIGRAPHX_THROW("RUN_ON_TARGET: must have exactly 1 module argument");
+        }
+        auto* mod_input = mods.front();
+        if(inputs.size() != mod_input->get_parameter_shapes().size())
+        {
+            MIGRAPHX_THROW("RUN_ON_TARGET: Mismatched number of input parameters");
+        }
+        auto mod_out_shapes = mod_input->get_output_shapes();
+        return mod_out_shapes;
+    }
+
+    migraphx::argument
+    compute(const migraphx::shape&,
+            const std::vector<migraphx::argument>& args,
+            const std::vector<migraphx::module_ref>& mods,
+            const std::function<std::vector<migraphx::argument>(
+                migraphx::module_ref&, const std::unordered_map<std::string, migraphx::argument>&)>&
+                run) const
+
+    {
+        std::unordered_map<std::string, migraphx::argument> params;
+        std::set<std::string> pnames;
+        const auto* smod = mods.front();
+        assert(mods.size() == 1);
+        auto names = smod->get_parameter_names();
+        pnames.insert(names.begin(), names.end());
+        assert(pnames.size() == args.size());
+        std::transform(pnames.begin(),
+                       pnames.end(),
+                       args.begin(),
+                       std::inserter(params, params.end()),
+                       [](auto&& name, auto&& arg) { return std::make_pair(name, arg); });
+        auto* mod    = mods.front();
+        auto results = run(mod, params);
+        return migraphx::argument{results};
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+#endif

src/include/migraphx/pass_manager.hpp

@@ -26,6 +26,7 @@
 
 #include <migraphx/config.hpp>
 #include <migraphx/pass.hpp>
+#include <migraphx/module_ref.hpp>
 #include <migraphx/tracer.hpp>
 #include <vector>
 
@@ -46,6 +47,10 @@ struct module_pass_manager
     virtual ~module_pass_manager() {}
 };
 
+void run_passes(program& prog,
+                module_ref root_mod,
+                const std::vector<pass>& passes,
+                tracer trace = tracer{});
 void run_passes(module& mod, const std::vector<pass>& passes, tracer trace = tracer{});
 void run_passes(program& prog, const std::vector<pass>& passes, tracer trace = tracer{});
 

src/include/migraphx/program.hpp

@@ -92,6 +92,9 @@ struct program
 
     void compile(const target& t, compile_options options = compile_options{});
 
+    void compile(const std::vector<target>& targets,
+                 std::vector<compile_options> compile_opts = {});
+
     bool is_compiled() const;
 
     void finalize();

src/instruction.cpp

@@ -406,6 +406,9 @@ void instruction::print(std::ostream& os,
     // skip return instruction shape
     if(ins->name() != "@return")
         os << " -> " << ins->get_shape();
+    // print tid
+
+    os << ", target_id=" << ins->target_id;
 }
 
 static void debug_name(std::ostream& os, const instruction& ins)
@@ -469,7 +472,8 @@ operation instruction::normalized_operator() const
     }
     return o;
 }
-
+std::size_t instruction::get_target_id() const { return target_id; }
+void instruction::set_target_id(std::size_t tid) { this->target_id = tid; }
 std::vector<shape> to_shapes(const std::vector<instruction_ref>& args)
 {
     std::vector<shape> shapes(args.size());

src/pass_manager.cpp

@@ -123,27 +123,18 @@ struct module_pm : module_pass_manager
 
 module& get_module(module_pass_manager& mpm) { return mpm.get_module(); }
 
-void run_passes(module& mod, const std::vector<pass>& passes, tracer trace)
-{
-    if(enabled(MIGRAPHX_TRACE_PASSES{}))
-        trace = tracer{std::cout};
-    for(const auto& p : passes)
-    {
-        module_pm{&mod, &trace}.run_pass(p);
-    }
-}
-
-void run_passes(program& prog, const std::vector<pass>& passes, tracer trace)
+void run_passes(program& prog, module_ref root_mod, const std::vector<pass>& passes, tracer trace)
 {
     if(enabled(MIGRAPHX_TRACE_PASSES{}))
         trace = tracer{std::cout};
     std::unordered_set<module_ref> visited;
     for(const auto& p : passes)
     {
-        auto mods = prog.get_modules();
         auto tree                        = prog.get_module_tree();
+        std::vector<module_ref> sub_mods = root_mod->get_sub_modules();
+        sub_mods.insert(sub_mods.begin(), root_mod);
         visited.clear();
-        for(const auto& mod : reverse(mods))
+        for(const auto& mod : reverse(sub_mods))
         {
             if(mod->bypass())
                 continue;
@@ -167,5 +158,20 @@ void run_passes(program& prog, const std::vector<pass>& passes, tracer trace)
     }
 }
 
+void run_passes(module& mod, const std::vector<pass>& passes, tracer trace)
+{
+    if(enabled(MIGRAPHX_TRACE_PASSES{}))
+        trace = tracer{std::cout};
+    for(const auto& p : passes)
+    {
+        module_pm{&mod, &trace}.run_pass(p);
+    }
+}
+
+void run_passes(program& prog, const std::vector<pass>& passes, tracer trace)
+{
+    run_passes(prog, prog.get_main_module(), passes, trace);
+}
+
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/program.cpp

@@ -21,6 +21,7 @@
  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  * THE SOFTWARE.
  */
+#include <migraphx/compile_options.hpp>
 #include <migraphx/program.hpp>
 #include <migraphx/stringutils.hpp>
 #include <migraphx/instruction.hpp>
@@ -42,6 +43,7 @@
 #include <sstream>
 #include <algorithm>
 #include <set>
+#include <unordered_map>
 #include <utility>
 
 #include <unordered_set>
@@ -53,12 +55,24 @@ inline namespace MIGRAPHX_INLINE_NS {
 
 using milliseconds = std::chrono::duration<double, std::milli>;
 
+struct mark_instruction_target
+{
+    std::size_t target_id = 0;
+    std::string name() const { return "mark_instruction_target"; }
+    void apply(module& m) const
+    {
+        for(auto& ins : m)
+            ins.set_target_id(target_id);
+    }
+};
+
 struct program_impl
 {
     // A map is used to keep references to modules of the program
     std::unordered_map<std::string, module> modules;
     context ctx;
     std::string target_name;
+    std::vector<context> contexts;
 };
 
 program::program() : impl(std::make_unique<program_impl>()) { this->create_module("main"); }
@@ -205,8 +219,94 @@ target_assignments program::get_target_assignments(const std::vector<target>& ta
 
 bool program::is_compiled() const { return not this->impl->target_name.empty(); }
 
+void program::compile(const std::vector<target>& targets, std::vector<compile_options> compile_opts)
+{
+    // Gather all the target roots
+    std::unordered_multimap<std::size_t, module_ref> roots;
+    auto mods = this->get_modules();
+    for(auto* mod : mods)
+    {
+        for(const auto& ins : *mod)
+        {
+            if(ins.name() != "run_on_target")
+                continue;
+            auto v                     = ins.get_operator().to_value();
+            module_ref root            = ins.module_inputs().front();
+            std::size_t root_target_id = v.at("target_id").to<std::size_t>();
+            assert(root_target_id < targets.size());
+            roots.insert({root_target_id, root});
+        }
+    }
+
+    auto trace = tracer{};
+    // TODO: Add tracer based on compile options
+    if(enabled(MIGRAPHX_TRACE_COMPILE{}))
+        trace = tracer{std::cout};
+
+    trace(*this);
+    trace();
+    // It is assumed that all instructions outside of any root module would run on "ref" target
+    // Ref target may or may not be passed as one of the target for the "compile()".
+    // If it is not passed, Create one and add context of it into the map.
+    auto target_idx = [&](const std::string& t_name) {
+        return static_cast<std::size_t>(
+            std::find_if(
+                targets.begin(), targets.end(), [&](const auto& t) { return t.name() == t_name; }) -
+            targets.begin());
+    };
+
+    std::size_t ref_target_id = target_idx("ref");
+    if(ref_target_id == targets.size())
+    {
+        this->impl->contexts.resize(targets.size() + 1);
+        this->impl->contexts[ref_target_id] = migraphx::make_target("ref").get_context();
+        // users could pass lessers compile_ops than targets, in that case use default compile_opts
+        compile_opts.resize(targets.size() + 1, migraphx::compile_options{});
+    }
+    else
+    {
+        this->impl->contexts.resize(targets.size());
+        compile_opts.resize(targets.size(), migraphx::compile_options{});
+    }
+    // mark all the instruction as ref target first, later change target_id based on root-target
+    run_passes(*this, {mark_instruction_target{ref_target_id}});
+
+    // Run passes on each root target
+    for(const auto i : range(targets.size()))
+    {
+        const auto& root_target              = targets.at(i);
+        auto root_target_id                  = i;
+        auto root_modules_range              = roots.equal_range(root_target_id);
+        this->impl->contexts[root_target_id] = root_target.get_context();
+        for(const auto& [id, current_mod] : range(root_modules_range))
+        {
+            auto passes = root_target.get_passes(this->impl->contexts[root_target_id],
+                                                 compile_opts[root_target_id]);
+            passes.push_back(mark_instruction_target{static_cast<size_t>(root_target_id)});
+            run_passes(*this, current_mod, passes, trace);
+
+            auto invalid = current_mod->validate();
+            if(invalid != current_mod->end())
+            {
+                MIGRAPHX_THROW("Invalid module " + current_mod->name() +
+                               " from compilation at instruction " +
+                               std::to_string(std::distance(current_mod->begin(), invalid)));
+            }
+            auto dangling = current_mod->find_dangling_reference();
+            if(dangling != current_mod->end())
+            {
+                auto index = std::distance(current_mod->begin(), dangling);
+                MIGRAPHX_THROW("Dangling reference in module " + current_mod->name() +
+                               " from instruction " + std::to_string(index));
+            }
+            current_mod->finalize(this->impl->contexts[root_target_id]);
+        }
+    }
+}
+
 void program::compile(const target& t, compile_options options)
 {
+    // todo: combine with multi-target compile method
     assert(not this->is_compiled());
     this->impl->target_name = t.name();
     this->impl->ctx         = t.get_context();
@@ -366,7 +466,6 @@ std::vector<argument> generic_eval(const module* mod,
                     assert(results.find(i) != results.end());
                     return results[i];
                 });
-
             const auto& mod_args = ins->module_inputs();
             auto module_eval     = [&](module_ref smod,
                                    const std::unordered_map<std::string, argument>& inputs) {

src/targets/fpga/subgraph.cpp

@@ -113,8 +113,7 @@ void subgraph::apply(module_pass_manager& mpm) const
     // TODO(varunsh): this code may be replaceable by code in the fuse_pointwise pass
 
     // assuming all FPGA instructions are in one contiguous range
-    pm->insert_instructions(pm->end(), first, last, {});
-
+    pm->insert_instructions(pm->end(), first, std::next(last), {});
     migraphx::instruction_ref placeholder_ins;
     for(auto it : iterator_for(mod))
     {

src/targets/gpu/lowering.cpp

@@ -83,6 +83,7 @@ struct miopen_apply
         auto& ctx      = get_context();
         int8_x4_format = get_int8_x4_format(ctx);
         compute_fp32   = get_compute_fp32_flag();
+        // TODO: Set Offload copy based on root modules' compile options
         offload_copy = (mod->name() == "main") ? pass->offload_copy : false;
 
         add_generic_op("contiguous");

test/CMakeLists.txt

@@ -190,6 +190,25 @@ if(MIGRAPHX_ENABLE_PYTHON)
     add_subdirectory(py)
 endif()
 
+# multitarget test
+if(MIGRAPHX_ENABLE_GPU AND MIGRAPHX_ENABLE_CPU AND MIGRAPHX_ENABLE_FPGA)
+    set(TEST_MULTI_TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR}/multi_target)
+    file(GLOB MULTI_TARGET_TESTS ${CONFIGURE_DEPENDS} ${TEST_MULTI_TARGET_DIR}/*.cpp)
+
+    foreach(MULTI_TARGET_TEST ${MULTI_TARGET_TESTS})
+        get_filename_component(BASE_NAME ${MULTI_TARGET_TEST} NAME_WE)
+        set(TEST_NAME test_${BASE_NAME})
+        add_executable(${TEST_NAME} ${MULTI_TARGET_TEST})
+        rocm_clang_tidy_check(${TEST_NAME})
+        target_link_libraries(${TEST_NAME} migraphx migraphx_onnx migraphx_tf migraphx_all_targets)
+        target_include_directories(${TEST_NAME} PUBLIC include)
+        add_test(NAME ${TEST_NAME} COMMAND $<TARGET_FILE:${TEST_NAME}> WORKING_DIRECTORY ${TEST_MULTI_TARGET_DIR})
+        add_dependencies(tests ${TEST_NAME})
+        add_dependencies(check ${TEST_NAME})
+    endforeach()
+endif()
+
+
 function(test_header NAME HEADER)
     file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/header-main-include-${NAME}.cpp
         "#include <${HEADER}>\nint main() {}\n"
@@ -228,3 +247,4 @@ if(MIGRAPHX_ENABLE_FPGA)
     test_headers(migraphx/fpga ${CMAKE_SOURCE_DIR}/src/targets/fpga/include/migraphx/fpga/*.hpp)
 endif()
 
+

test/argument_test.cpp

@@ -193,6 +193,15 @@ TEST_CASE(value_argument)
     EXPECT(a4 == a2);
 }
 
+TEST_CASE(value_empty_argument)
+{
+    migraphx::argument a5;
+    EXPECT(a5.empty());
+    auto v3 = migraphx::to_value(a5);
+    auto a6 = migraphx::from_value<migraphx::argument>(v3);
+    EXPECT(a6 == a5);
+}
+
 TEST_CASE(value_tuple)
 {
     auto a1 = make_tuple(3, 3.0, make_tuple(3, 4));

test/multi_target/multitarget_test.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <iostream>
+#include <set>
+#include <unordered_set>
+#include <vector>
+#include <random>
+#include <cmath>
+#include <migraphx/program.hpp>
+#include <migraphx/target.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/shape.hpp>
+#include <migraphx/verify.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/functional.hpp>
+#include <basic_ops.hpp>
+#include <migraphx/compile_options.hpp>
+#include <migraphx/register_target.hpp>
+#include "test.hpp"
+
+// check if it is custom_op or run_on_module operator
+bool has_target_attr(const migraphx::instruction& ins)
+{
+    return ins.get_operator().attributes().contains("target");
+}
+
+auto nonprefixed_ops()
+{
+    // ops without prefixes
+    static std::unordered_set<std::string> op_map = {
+        "select_module", "load", "if", "nonmaxsuppression", "multibroadcast"};
+    return op_map;
+}
+
+bool is_compiled_gpu_module(const migraphx::module& m)
+{
+    return std::all_of(m.begin(), m.end(), [](auto ins) {
+        auto ins_name = ins.name();
+        if(not migraphx::starts_with(ins_name, "@"))
+        {
+            if(not migraphx::starts_with(ins_name, "gpu::") and
+               not migraphx::starts_with(ins_name, "hip::") and
+               not migraphx::starts_with(ins_name, "check_context") and
+               not migraphx::contains(nonprefixed_ops(), ins_name) and not has_target_attr(ins))
+            {
+                return false;
+            }
+        }
+        return true;
+    });
+}
+
+bool is_compiled_fpga_module(const migraphx::module& m)
+{
+    return std::all_of(m.begin(), m.end(), [](auto ins) {
+        auto ins_name = ins.name();
+        if(not migraphx::starts_with(ins_name, "@"))
+        {
+            if(not migraphx::starts_with(ins_name, "fpga::") and
+               not migraphx::starts_with(ins_name, "check_context") and
+               not migraphx::contains(nonprefixed_ops(), ins_name) and not has_target_attr(ins))
+            {
+                return false;
+            }
+        }
+        return true;
+    });
+}
+
+bool is_compiled_cpu_module(const migraphx::module& m)
+{
+    return std::all_of(m.begin(), m.end(), [](auto ins) {
+        auto ins_name = ins.name();
+        if(not migraphx::starts_with(ins_name, "@"))
+        {
+            if(not migraphx::starts_with(ins_name, "cpu::") and
+               not migraphx::starts_with(ins_name, "dnnl::") and
+               not migraphx::starts_with(ins_name, "check_context") and not has_target_attr(ins) and
+               not migraphx::contains(nonprefixed_ops(), ins_name))
+            {
+                return false;
+            }
+        }
+        return true;
+    });
+}
+
+bool is_compiled_ref_module(const migraphx::module& m)
+{
+    return std::all_of(m.begin(), m.end(), [](auto ins) {
+        auto ins_name = ins.name();
+        if(not migraphx::starts_with(ins_name, "@"))
+        {
+            if((not migraphx::starts_with(ins_name, "ref::") and
+                not migraphx::starts_with(ins_name, "check_context") and
+                not has_target_attr(ins)) and
+               not migraphx::contains(nonprefixed_ops(), ins_name))
+            {
+                return false;
+            }
+        }
+        return true;
+    });
+}
+
+// NOLINT
+bool check_compiled_program(const migraphx::program& p,
+                            const std::vector<migraphx::target>& targets)
+{
+    auto mods           = p.get_modules();
+    bool check_compiled = true;
+    for(const auto* mod : mods)
+    {
+        for(const auto& ins : *mod)
+        {
+            if(ins.name() == "run_on_target")
+            {
+                auto* mod_input = ins.module_inputs().front();
+                std::size_t target_id =
+                    ins.get_operator().to_value()["target_id"].to<std::size_t>();
+                auto target_name = targets.at(target_id).name();
+                if(target_name == "gpu")
+                    check_compiled &= is_compiled_gpu_module(*mod_input);
+                else if(target_name == "cpu")
+                    check_compiled &= is_compiled_cpu_module(*mod_input);
+                else if(target_name == "fpga")
+                    check_compiled &= is_compiled_fpga_module(*mod_input);
+                else if(target_name == "ref")
+                    check_compiled &= is_compiled_ref_module(*mod_input);
+            }
+        }
+    }
+    return check_compiled;
+}
+
+TEST_CASE(multitarget_compile_cpu_gpu)
+{
+    migraphx::program p;
+    auto* mm      = p.get_main_module();
+    auto* cpu_mod = p.create_module("cpu_mod");
+    auto s        = migraphx::shape{migraphx::shape::float_type, {8}};
+    auto x_cpu    = cpu_mod->add_parameter("cpu_x", s);
+    auto y_cpu    = cpu_mod->add_parameter("cpu_y", s);
+    auto cpu_add  = cpu_mod->add_instruction(migraphx::make_op("add"), x_cpu, y_cpu);
+    cpu_mod->add_return({cpu_add});
+
+    auto* gpu_mod = p.create_module("gpu_mod");
+    auto x_gpu    = gpu_mod->add_parameter("gpu_x", s);
+    auto y_gpu    = gpu_mod->add_parameter("gpu_y", s);
+    auto gpu_add  = gpu_mod->add_instruction(migraphx::make_op("add"), x_gpu, y_gpu);
+    gpu_mod->add_return({gpu_add});
+
+    auto x_param = mm->add_parameter("x", s);
+    auto y_param = mm->add_parameter("y", s);
+    auto z_param = mm->add_parameter("z", s);
+    auto cpu_ins = mm->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 1}}), {x_param, y_param}, {cpu_mod});
+    auto gpu_ins = mm->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 0}}), {cpu_ins, z_param}, {gpu_mod});
+    mm->add_return({gpu_ins});
+    p.compile({migraphx::make_target("gpu"), migraphx::make_target("cpu")});
+    EXPECT(check_compiled_program(p, {migraphx::make_target("gpu"), migraphx::make_target("cpu")}));
+}
+
+TEST_CASE(single_target_compile)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+
+    migraphx::shape boxes_s{migraphx::shape::float_type, {1, 6, 4}};
+
+    migraphx::shape scores_s{migraphx::shape::float_type, {1, 1, 6}};
+    std::vector<float> scores_vec = {0.9, 0.75, 0.6, 0.95, 0.5, 0.3};
+
+    auto boxes_l         = mm->add_parameter("boxes", boxes_s);
+    auto scores_l        = mm->add_literal(migraphx::literal(scores_s, scores_vec));
+    auto max_out_l       = mm->add_literal(int64_t{4});
+    auto iou_threshold   = mm->add_literal(0.5f);
+    auto score_threshold = mm->add_literal(0.0f);
+
+    auto r = mm->add_instruction(migraphx::make_op("nonmaxsuppression", {{"center_point_box", 1}}),
+                                 boxes_l,
+                                 scores_l,
+                                 max_out_l,
+                                 iou_threshold,
+                                 score_threshold);
+    mm->add_return({r});
+    p.compile(migraphx::make_target("gpu"));
+    EXPECT(is_compiled_gpu_module(*p.get_main_module()));
+}
+
+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 = {0.384804, -1.77948, -0.453775, 0.477438, -1.06333, -1.12893};
+    auto l1                  = then_mod->add_literal(migraphx::literal(ds, data1));
+    auto a1                  = then_mod->add_instruction(migraphx::make_op("add"), x, l1);
+    then_mod->add_return({a1});
+
+    auto* else_mod           = p.create_module("else_cpu_mod");
+    std::vector<float> data2 = {-0.258047, 0.360394, 0.536804, -0.577762, 1.0217, 1.02442};
+    auto l2                  = else_mod->add_literal(migraphx::literal(ds, data2));
+    auto a2                  = else_mod->add_instruction(migraphx::make_op("mul"), 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}}), {}, {else_mod});
+    run_on_cpu_mod->add_return({run_cpu_ins});
+
+    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}}), {}, {then_mod});
+    run_on_gpu_mod->add_return({run_gpu_ins});
+
+    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
+    p.compile({migraphx::make_target("gpu"), migraphx::make_target("cpu")});
+    EXPECT(check_compiled_program(p, {migraphx::make_target("gpu"), migraphx::make_target("cpu")}));
+}
+
+TEST_CASE(multitarget_compile_nested_if_then_else)
+{
+    float seed = 0.0f;
+    std::mt19937 gen(seed);
+    std::uniform_real_distribution<> dis(0.0, 1.0);
+    auto get_random_values = [&](size_t elements) {
+        std::vector<float> rand_samples(elements);
+        std::generate(rand_samples.begin(), rand_samples.end(), [&]() { return dis(gen); });
+        return rand_samples;
+    };
+
+    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               = mm->add_parameter("cond", 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 = get_random_values(ds.elements());
+        auto l1                 = test_mod->add_literal(migraphx::literal(ds, data));
+        auto test_mod_param     = test_mod->add_parameter(mod_name, ds);
+        // instruction with local literal and main_mod param as inputs
+        auto ins1 = test_mod->add_instruction(migraphx::make_op("add"), x, l1);
+        // instructinon with local param and local ins as inputs
+        auto ins2 = test_mod->add_instruction(migraphx::make_op("mul"), ins1, test_mod_param);
+        // instruction with local ins and parent ins as inputs
+        auto ins3 = test_mod->add_instruction(migraphx::make_op("sub"), ins2, inputs.front());
+        test_mod->add_return({ins3});
+        auto* run_on_target_mod = prog.create_module("run_on_" + mod_name);
+        run_on_target_mod->add_instruction(
+            migraphx::make_op("run_on_target", {{"target_id", tid}}), {inputs.front()}, {test_mod});
+        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_param   = then_mod->add_parameter("then_mod_param", ds);
+    auto then_mod_ref_ins = then_mod->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 3}}), {then_mod_param, y}, {ref_mod});
+    auto then_mod_ref_ins_0 = then_mod->add_instruction(
+        migraphx::make_op("get_tuple_elem", {{"index", 0}}), then_mod_ref_ins);
+    then_mod->add_instruction(
+        migraphx::make_op("if"),
+        {cond},
+        {create_test_module(p, {z}, 1), create_test_module(p, {then_mod_ref_ins_0}, 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_param    = else_mod->add_parameter("else_mod_param", ds);
+    auto else_mod_fpga_ins = else_mod->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 2}}), {else_mod_param, y}, {fpga_mod});
+    auto else_mod_fpga_ins_0 = else_mod->add_instruction(
+        migraphx::make_op("get_tuple_elem", {{"index", 0}}), else_mod_fpga_ins);
+
+    else_mod->add_instruction(migraphx::make_op("if"),
+                              {cond},
+                              {create_test_module(p, {else_mod_fpga_ins_0}, 0),
+                               create_test_module(p, {else_mod_param}, 1)});
+
+    // Create nested and multi-target main module using "If"
+    auto main_if_ins =
+        mm->add_instruction(migraphx::make_op("if"), {cond, x}, {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
+    p.compile({migraphx::make_target("gpu"),
+               migraphx::make_target("cpu"),
+               migraphx::make_target("fpga"),
+               migraphx::make_target("ref")});
+    EXPECT(check_compiled_program(p,
+                                  {migraphx::make_target("gpu"),
+                                   migraphx::make_target("cpu"),
+                                   migraphx::make_target("fpga"),
+                                   migraphx::make_target("ref")}));
+}
+
+TEST_CASE(multitarget_select_module)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape lit_s{migraphx::shape{migraphx::shape::float_type, {1}}};
+    auto literal_ins = mm->add_literal(migraphx::literal{lit_s, {6}});
+
+    // create batch submodules
+    auto create_submodule = [&](std::size_t batch_size, const std::string& module_name) {
+        auto* submod = p.create_module(module_name);
+        migraphx::shape sm_shape{migraphx::shape::float_type, {batch_size, 4}};
+        auto sm_input = submod->add_parameter("data", sm_shape);
+        auto broadcast_lit =
+            submod->add_instruction(migraphx::make_op("multibroadcast"), literal_ins, sm_input);
+        auto add_ins0 = submod->add_instruction(migraphx::make_op("add"), sm_input, broadcast_lit);
+        auto add_ins1 = submod->add_instruction(migraphx::make_op("add"), add_ins0, broadcast_lit);
+        submod->add_return({add_ins0, add_ins1});
+        return submod;
+    };
+    auto* batch1 = create_submodule(1, "batch_1");
+    auto* batch2 = create_submodule(2, "batch_2");
+    auto* batch3 = create_submodule(3, "batch_3");
+    auto* batch4 = create_submodule(4, "batch_4");
+
+    migraphx::shape s{migraphx::shape::float_type, {{1, 4}, {4, 4}}};
+    auto input        = mm->add_parameter("data", s);
+    auto* run_cpu_mod = p.create_module("cpu_mod");
+    auto cpu_param    = run_cpu_mod->add_parameter(
+        "cpu_data", migraphx::shape{migraphx::shape::float_type, {1, 4}});
+    auto run_cpu_ins = run_cpu_mod->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 1}}), {cpu_param}, {batch1});
+    run_cpu_mod->add_return({run_cpu_ins});
+
+    auto* run_gpu_mod = p.create_module("gpu_mod");
+    auto gpu_param    = run_gpu_mod->add_parameter(
+        "gpu_data", migraphx::shape{migraphx::shape::float_type, {2, 4}});
+    auto run_gpu_ins = run_gpu_mod->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 0}}), {gpu_param}, {batch2});
+    run_gpu_mod->add_return({run_gpu_ins});
+
+    auto* run_fpga_mod = p.create_module("fpga_mod");
+    auto fpga_param    = run_fpga_mod->add_parameter(
+        "fpga_data", migraphx::shape{migraphx::shape::float_type, {3, 4}});
+    auto run_fpga_ins = run_fpga_mod->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 2}}), {fpga_param}, {batch3});
+    run_fpga_mod->add_return({run_fpga_ins});
+
+    auto* run_ref_mod = p.create_module("ref_mod");
+    auto ref_param    = run_fpga_mod->add_parameter(
+        "ref_data", migraphx::shape{migraphx::shape::float_type, {4, 4}});
+    auto run_ref_ins = run_ref_mod->add_instruction(
+        migraphx::make_op("run_on_target", {{"target_id", 3}}), {ref_param}, {batch4});
+    run_ref_mod->add_return({run_ref_ins});
+
+    std::vector<migraphx::shape> sub_shapes = {};
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+    sub_shapes.push_back(migraphx::shape{migraphx::shape::float_type, {{1, 4}, {4, 4}}});
+    migraphx::shape out_attr = migraphx::shape{sub_shapes};
+    auto sm_ins              = mm->add_instruction(
+        migraphx::make_op("select_module", {{"output_dyn_shapes", migraphx::to_value(out_attr)}}),
+        {input},
+        {run_cpu_mod, run_gpu_mod, run_fpga_mod, run_ref_mod});
+    auto ret0 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), sm_ins);
+    auto ret1 = mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 1}}), sm_ins);
+    mm->add_return({ret0, ret1});
+    // compile
+    p.compile({migraphx::make_target("gpu"),
+               migraphx::make_target("cpu"),
+               migraphx::make_target("fpga"),
+               migraphx::make_target("ref")});
+    EXPECT(check_compiled_program(p,
+                                  {migraphx::make_target("gpu"),
+                                   migraphx::make_target("cpu"),
+                                   migraphx::make_target("fpga"),
+                                   migraphx::make_target("ref")}));
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/run_on_target_test.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/program.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/literal.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/verify.hpp>
+#include <migraphx/shape.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/register_target.hpp>
+#include "test.hpp"
+
+TEST_CASE(run_on_target_shape_tests)
+{
+    {
+        test::throws([]() { migraphx::make_op("run_on_target"); });
+    }
+    {
+        migraphx::program p;
+        auto s        = migraphx::shape{migraphx::shape::float_type, {12, 12}};
+        auto* mm      = p.get_main_module();
+        auto x        = mm->add_parameter("x", s);
+        auto* run_mod = p.create_module("run_mod");
+        run_mod->add_return({x, x});
+        test::throws(
+            [&]() { mm->add_instruction(migraphx::make_op("run_on_target"), {x}, {run_mod}); });
+        test::throws([&]() {
+            mm->add_instruction(migraphx::make_op("run_on_target"), {x}, {run_mod, run_mod});
+        });
+    }
+}
+
+TEST_CASE(eval_run_on_target)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::float_type, {3}};
+    auto l         = mm->add_literal(migraphx::literal{s, {4.0, 16.0, 64.0}});
+    auto* ref_mod  = p.create_module("ref_mod");
+    auto ref_rsqrt = ref_mod->add_instruction(migraphx::make_op("rsqrt"), l);
+    ref_mod->add_return({ref_rsqrt});
+    auto run_on_ins =
+        mm->add_instruction(migraphx::make_op("run_on_target", {{"target_id", 0}}), {}, {ref_mod});
+    mm->add_instruction(migraphx::make_op("get_tuple_elem", {{"index", 0}}), run_on_ins);
+    p.compile({migraphx::make_target("ref")});
+    auto result = p.eval({}).back();
+    std::vector<float> results_vector(3);
+    result.visit([&](auto output) { results_vector.assign(output.begin(), output.end()); });
+    std::vector<float> gold = {0.5, 0.25, 0.125};
+    EXPECT(migraphx::verify_range(results_vector, gold));
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/targets.cpp

@@ -47,7 +47,7 @@ TEST_CASE(targets)
     auto ref_target = migraphx::make_target("ref");
 #endif
     auto ts = migraphx::get_targets();
-    EXPECT(ts.size() == 1);
+    EXPECT(ts.size() >= 1);
 }
 
 int main(int argc, const char* argv[]) { test::run(argc, argv); }