Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/AMDMIGraphX into nhwc_workaround

d7b1895a · Khalique Ahmed · 1add453a · a83371ca · d7b1895a · d7b1895a
Commit d7b1895a authored Jul 07, 2023 by Khalique Ahmed
20 changed files
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -27,23 +27,20 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_BINARY_DIR}")
    message(FATAL_ERROR "The binary and source directroy cannot be the same")
 endif()
+get_property(_GENERATOR_IS_MULTI_CONFIG GLOBAL PROPERTY GENERATOR_IS_MULTI_CONFIG)
 # This has to be initialized before the project() command appears
 # Set the default of CMAKE_BUILD_TYPE to be release, unless user specifies with -D.  MSVC_IDE does not use CMAKE_BUILD_TYPE
-if( NOT MSVC_IDE AND NOT CMAKE_BUILD_TYPE )
+if(_GENERATOR_IS_MULTI_CONFIG)
-    set( CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build, options are: None Debug Release RelWithDebInfo MinSizeRel." )
+    if (NOT CMAKE_CONFIGURATION_TYPES)
-endif()
+        set(CMAKE_CONFIGURATION_TYPES "Debug;Release;RelWithDebInfo;MinSizeRel" CACHE STRING
+            "Available build types (configurations) on multi-config generators")
-# Setup valid strings for build type
+    endif()
-if (NOT CMAKE_CONFIGURATION_TYPES)
-    set(CMAKE_CONFIGURATION_TYPES "Debug;Release;RelWithDebInfo;MinSizeRel" CACHE STRING "Configs")
-endif()
-set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS ${CMAKE_CONFIGURATION_TYPES})
-# Default installation path
-if(WIN32)
-    set(CMAKE_INSTALL_PREFIX "/opt/rocm/x86_64-w64-mingw32" CACHE PATH "")
 else()
-    set(CMAKE_INSTALL_PREFIX "/opt/rocm" CACHE PATH "")
+    if(NOT CMAKE_BUILD_TYPE)
+        set(CMAKE_BUILD_TYPE Release CACHE STRING
+            "Choose the type of build, options are: None Debug Release RelWithDebInfo MinSizeRel.")
+    endif()
 endif()
 set(CMAKE_BUILD_RPATH "${CMAKE_BINARY_DIR}/lib")
@@ -130,6 +127,7 @@ rocm_enable_clang_tidy(
        -bugprone-implicit-widening-of-multiplication-result
        -bugprone-macro-parentheses
        -bugprone-signed-char-misuse
+        -bugprone-unchecked-optional-access
        # Disable the aliased reserved identifiers
        -cert-dcl37-c
        -cert-dcl51-cpp

--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -89,6 +89,8 @@ def rocmnodename(name) {
        node_name = "${rocmtest_name} && vega";
    } else if(name == "navi21") {
        node_name = "${rocmtest_name} && navi21";
+    } else if(name == "mi100+") {
+        node_name = "${rocmtest_name} && (gfx908 || gfx90a)";
    } else if(name == "anygpu") {
        node_name = "${rocmtest_name} && (gfx908 || gfx90a || vega)";
    } else if(name == "nogpu") {
@@ -134,6 +136,12 @@ rocmtest clang_debug: rocmnode('vega') { cmake_build ->
            cmake_build(flags: "-DCMAKE_BUILD_TYPE=debug -DMIGRAPHX_ENABLE_PYTHON=Off -DMIGRAPHX_ENABLE_MLIR=On -DCMAKE_CXX_FLAGS_DEBUG='${debug_flags}' -DCMAKE_C_FLAGS_DEBUG='${debug_flags}'")
        }
    }
+}, ck_release: rocmnode('mi100+') { cmake_build ->
+    stage('CK Release') {
+        withEnv(['MIGRAPHX_ENABLE_CK=1', 'MIGRAPHX_TUNE_CK=1']) {
+            cmake_build(flags: "-DCMAKE_BUILD_TYPE=release")
+        }
+    }
 }, clang_asan: rocmnode('nogpu') { cmake_build ->
    stage('Clang ASAN') {
        def sanitizers = "undefined,address"

--- a/requirements.txt
+++ b/requirements.txt
@@ -28,4 +28,4 @@ ROCmSoftwarePlatform/half@rocm-5.4.2
 pybind/pybind11@d159a563383d10c821ba7b2a71905d1207db6de4 --build
 msgpack/msgpack-c@cpp-3.3.0 -DMSGPACK_BUILD_TESTS=Off
 sqlite3@3.17 -DCMAKE_POSITION_INDEPENDENT_CODE=On
-ROCmSoftwarePlatform/composable_kernel@ac580f77a84c705c678816ef7195adfcc02bdda5 -DCK_BUILD_JIT_LIB=On -DCMAKE_POSITION_INDEPENDENT_CODE=On
+ROCmSoftwarePlatform/composable_kernel@5172ec5280f14974beee2acf1af1db3b2670244c -DCK_BUILD_JIT_LIB=On -DCMAKE_POSITION_INDEPENDENT_CODE=On
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -95,6 +95,7 @@ add_library(migraphx
    simplify_algebra.cpp
    simplify_reshapes.cpp
    split_single_dyn_dim.cpp
+    target.cpp
    tmp_dir.cpp
    value.cpp
    verify_args.cpp

--- a/src/common.cpp
+++ b/src/common.cpp
@@ -157,7 +157,7 @@ insert_common_args(module& m, instruction_ref ins, std::vector<instruction_ref>
                ins, make_op("multibroadcast", {{"out_dyn_dims", to_value(c_dyn_dims)}}), inputs);
        }
        std::transform(inputs.begin() + 1, inputs.end(), inputs.begin() + 1, [&](auto input) {
-            // uses previous multibroadcast to avoid recalculating the common shape from the
+            // uses previous input to avoid recalculating the common shape from the
            // full set of input shapes at runtime
            if(input->get_shape().dyn_dims() != c_dyn_dims)
            {

--- a/src/dead_code_elimination.cpp
+++ b/src/dead_code_elimination.cpp
@@ -49,8 +49,10 @@ void dead_code_elimination::apply(module& m) const
        if(i == last)
            break;
        // Skip instruction with empty shape as output unless its [dynamic, builtin, undefined,
-        // identity, allocate]
+        // identity, allocate or tuple_type]
-        if((not i->get_shape().dynamic() and i->get_shape().elements() == 0) and
+        if((not i->get_shape().dynamic() and
+            (i->get_shape().elements() == 0 and
+             i->get_shape().type() != migraphx::shape::tuple_type)) and
           not(i->name().front() == '@') and not contains({"identity", "allocate"}, i->name()) and
           not i->is_undefined())
            continue;

--- a/src/include/migraphx/instruction.hpp
+++ b/src/include/migraphx/instruction.hpp
@@ -137,6 +137,7 @@ 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;

--- a/src/include/migraphx/module.hpp
+++ b/src/include/migraphx/module.hpp
@@ -189,7 +189,7 @@ struct module
    instruction_ref validate() const;
    instruction_ref find_dangling_reference() const;
-    void finalize(context& ctx);
+    void finalize(std::vector<context>& contexts);
    void debug_print() const;
    void debug_print(instruction_ref ins) const;

--- a/src/include/migraphx/op/clip.hpp
+++ b/src/include/migraphx/op/clip.hpp
@@ -25,12 +25,13 @@
 #define MIGRAPHX_GUARD_OPERATORS_CLIP_HPP
 #include <array>
+#include <cmath>
 #include <migraphx/check_shapes.hpp>
 #include <migraphx/argument.hpp>
 #include <migraphx/par_for.hpp>
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
-#include <cmath>
+#include <migraphx/dyn_output.hpp>
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -48,15 +49,15 @@ struct clip
    shape compute_shape(std::vector<shape> inputs) const
    {
-        check_shapes{inputs, *this}.has(3).same_type().same_dims();
+        check_shapes{inputs, *this, true}.has(3).same_type().same_dims();
        return inputs.front();
    }
-    argument compute(const shape& output_shape, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
    {
-        argument result{output_shape};
+        argument result{dyn_out.computed_shape};
        visit_all(result, args[0], args[1], args[2])([&](auto output, auto x, auto min, auto max) {
-            par_for(output_shape.elements(),
+            par_for(dyn_out.computed_shape.elements(),
                    [&](auto i) { output[i] = std::min(std::max(min[i], x[i]), max[i]); });
        });

--- a/src/include/migraphx/operation.hpp
+++ b/src/include/migraphx/operation.hpp
@@ -261,11 +261,13 @@ auto compute_op(rank<1>,
 template <class T, class F>
 argument compute_op(rank<0>,
                    const T& x,
-                    const shape&,
+                    const shape& output,
-                    const std::vector<argument>&,
+                    const std::vector<argument>& inputs,
-                    const std::vector<module_ref>&,
+                    const std::vector<module_ref>& module_args,
                    F)
 {
+    if(module_args.empty())
+        return compute_op(x, output, inputs);
    std::string name = x.name();
    MIGRAPHX_THROW("Not computable: " + name);
 }

--- a/src/include/migraphx/program.hpp
+++ b/src/include/migraphx/program.hpp
@@ -79,6 +79,9 @@ struct program
    std::vector<argument> eval(parameter_map params,
                               execution_environment exec_env = execution_environment{}) const;
+    void finish() const;
    std::size_t size() const;
    std::vector<shape> get_output_shapes() const;

--- a/src/include/migraphx/raw_data.hpp
+++ b/src/include/migraphx/raw_data.hpp
@@ -187,6 +187,7 @@ struct raw_data : raw_data_base
    std::string to_string() const
    {
        std::stringstream ss;
+        ss.precision(std::numeric_limits<double>::max_digits10);
        ss << static_cast<const Derived&>(*this);
        return ss.str();
    }

--- a/src/include/migraphx/replace_allocate.hpp
+++ b/src/include/migraphx/replace_allocate.hpp
@@ -30,7 +30,7 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
-struct module;
+struct module_pass_manager;
 /**
 *  Replace `allocate` instructions with target allocations or output parameters.
@@ -40,7 +40,7 @@ struct replace_allocate
    allocation_model model;
    bool offload_copy = false;
    std::string name() const { return "replace_allocate"; }
-    void apply(module& m) const;
+    void apply(module_pass_manager& mpm) const;
 };
 } // namespace MIGRAPHX_INLINE_NS

--- a/src/include/migraphx/target.hpp
+++ b/src/include/migraphx/target.hpp
@@ -45,6 +45,8 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
+struct value;
 #ifdef DOXYGEN
 /// An interface for a compilation target
@@ -467,6 +469,9 @@ inline const ValueType& any_cast(const target& x)
 #endif
+void migraphx_to_value(value& v, const target& t);
+void migraphx_from_value(const value& v, target& t);
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

--- a/src/include/migraphx/value.hpp
+++ b/src/include/migraphx/value.hpp
@@ -32,6 +32,7 @@
 #include <algorithm>
 #include <cassert>
 #include <memory>
+#include <cstdint>
 #include <sstream>
 #include <type_traits>
 #include <tuple>

--- a/src/instruction.cpp
+++ b/src/instruction.cpp
@@ -473,7 +473,9 @@ 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());

--- a/src/module.cpp
+++ b/src/module.cpp
@@ -326,6 +326,8 @@ instruction_ref module::replace_instruction(instruction_ref ins, instruction_ref
    if(ins == std::prev(this->end()))
    {
+        // "rep" instruction could be used earlier in the program and moving it at the end
+        // may cause invalid program, therefore make an identity operation in this case.
        return replace_instruction(ins, make_op("identity"), rep);
    }
@@ -650,8 +652,9 @@ instruction_ref module::find_dangling_reference() const
    return end();
 }
-void module::finalize(context& ctx)
+void module::finalize(std::vector<context>& contexts)
 {
+    assert(not contexts.empty());
    const bool trace = enabled(MIGRAPHX_TRACE_FINALIZE{});
    for(auto ins : iterator_for(*this))
    {
@@ -660,10 +663,10 @@ void module::finalize(context& ctx)
            std::cout << "Finalize: ";
            this->debug_print(ins);
        }
-        ins->finalize(ctx);
+        ins->finalize(contexts[ins->get_target_id()]);
        for(const auto& smod : ins->module_inputs())
        {
-            smod->finalize(ctx);
+            smod->finalize(contexts);
        }
    }

--- a/src/onnx/onnx_parser.cpp
+++ b/src/onnx/onnx_parser.cpp
@@ -38,6 +38,9 @@
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_ONNX_PARSER)
 static shape shape_from_dyn_dims(shape::type_t shape_type,
                                 const std::vector<shape::dynamic_dimension>& dyn_dims)
@@ -53,8 +56,6 @@ static shape shape_from_dyn_dims(shape::type_t shape_type,
    return {shape_type, dyn_dims};
 }
-namespace onnx {
 static onnx_parser::attribute_map get_attributes(const onnx::NodeProto& node)
 {
    std::unordered_map<std::string, onnx::AttributeProto> result;
@@ -297,16 +298,48 @@ int64_t onnx_parser::get_opset_version(const onnx::ModelProto& model)
    return version;
 }
-std::vector<instruction_ref>
+void print_added_instructions(module* mod,
-onnx_parser::parse_graph(module* mod, const onnx::GraphProto& graph, bool inlining)
+                              const std::vector<instruction_ref>& args,
+                              const std::vector<instruction_ref>& result)
+{
+    // Print instructions added by the parser not in args
+    std::vector<instruction_ref> added_instructions;
+    fix([&](auto self, auto r) {
+        for(auto ins : r)
+        {
+            if(contains(args, ins))
+                continue;
+            if(contains(added_instructions, ins))
+                continue;
+            self(ins->inputs());
+            added_instructions.push_back(ins);
+        }
+    })(result);
+    mod->debug_print(added_instructions);
+}
+std::unordered_map<std::string, instruction_ref>
+parse_intializer(const onnx_parser& parser, module* mod, const onnx::GraphProto& graph)
 {
    std::unordered_map<std::string, instruction_ref> mod_insts;
    for(auto&& f : graph.initializer())
    {
+        if(enabled(MIGRAPHX_TRACE_ONNX_PARSER{}))
+            std::cout << "initializer: " << f.name() << std::endl;
        // backup instructions in parent mod
-        mod_insts[f.name()] = mod->add_literal(parse_tensor(f));
+        mod_insts[f.name()] = mod->add_literal(parser.parse_tensor(f));
+        if(enabled(MIGRAPHX_TRACE_ONNX_PARSER{}))
+            mod->debug_print(mod_insts[f.name()]);
    }
+    return mod_insts;
+}
+std::unordered_map<std::string, instruction_ref>
+parse_inputs(const onnx_parser& parser,
+             module* mod,
+             const onnx::GraphProto& graph,
+             std::unordered_map<std::string, instruction_ref> mod_insts)
+{
    for(auto&& input : graph.input())
    {
        const std::string& name = input.name();
@@ -317,7 +350,7 @@ onnx_parser::parse_graph(module* mod, const onnx::GraphProto& graph, bool inlini
            // scenario that a nested subgraph contains a parameter with the
            // name existed in its parent graph.
            // In the current implementation, MIGraphX throws an exception for that.
-            if(contains(instructions, name))
+            if(contains(parser.instructions, name))
            {
                MIGRAPHX_THROW("module \"" + mod->name() + "\" has parameter name \"" + name +
                               "\" existing in parent graph!");
@@ -325,28 +358,41 @@ onnx_parser::parse_graph(module* mod, const onnx::GraphProto& graph, bool inlini
            shape s;
            std::vector<std::size_t> dims;
-            if(map_input_dims.count(name) > 0)
+            if(parser.map_input_dims.count(name) > 0)
            {
-                dims = map_input_dims.at(name);
+                dims = parser.map_input_dims.at(name);
-                s    = parse_type(input.type(), dims);
+                s    = parser.parse_type(input.type(), dims);
            }
-            else if(map_dyn_input_dims.count(name) > 0)
+            else if(parser.map_dyn_input_dims.count(name) > 0)
            {
                shape::type_t shape_type = get_type(input.type().tensor_type().elem_type());
-                s = shape_from_dyn_dims(shape_type, map_dyn_input_dims.at(name));
+                s = shape_from_dyn_dims(shape_type, parser.map_dyn_input_dims.at(name));
            }
            else
            {
-                s = parse_type(input.type(), dims);
+                s = parser.parse_type(input.type(), dims);
            }
            mod_insts[name] = mod->add_parameter(name, s);
        }
    }
+    return mod_insts;
+}
+std::vector<instruction_ref>
+onnx_parser::parse_graph(module* mod, const onnx::GraphProto& graph, bool inlining)
+{
+    std::unordered_map<std::string, instruction_ref> mod_insts =
+        parse_intializer(*this, mod, graph);
+    mod_insts = parse_inputs(*this, mod, graph, mod_insts);
    std::copy(mod_insts.begin(), mod_insts.end(), std::inserter(instructions, instructions.end()));
    for(auto&& node : graph.node())
    {
+        if(enabled(MIGRAPHX_TRACE_ONNX_PARSER{}))
+            std::cout << "operator: " << node.op_type() << std::endl;
        std::vector<instruction_ref> args;
        for(auto&& input : node.input())
        {
@@ -384,6 +430,11 @@ onnx_parser::parse_graph(module* mod, const onnx::GraphProto& graph, bool inlini
                       result.begin(),
                       std::inserter(instructions, instructions.end()),
                       [](auto&& x, auto&& y) { return std::make_pair(x, y); });
+        if(enabled(MIGRAPHX_TRACE_ONNX_PARSER{}))
+        {
+            print_added_instructions(mod, args, result);
+        }
    }
    // Find instructions corresponding to the output

--- a/src/pass_manager.cpp
+++ b/src/pass_manager.cpp
@@ -68,12 +68,18 @@ void run_pass(program& prog, const pass& p, tracer trace)
 struct module_pm : module_pass_manager
 {
    module* mod           = nullptr;
+    module* root_mod      = nullptr;
    tracer* t             = nullptr;
    module* common_parent = nullptr;
    program* prog         = nullptr;
    module_pm(module* pmod = nullptr, tracer* pt = nullptr) : mod(pmod), t(pt) {}
+    module_pm(module* pmod = nullptr, module* rmod = nullptr, tracer* pt = nullptr)
+        : mod(pmod), root_mod(rmod), t(pt)
+    {
+    }
    template <class... Ts>
    void trace(Ts&&... xs) const
    {
@@ -97,6 +103,8 @@ struct module_pm : module_pass_manager
    virtual module* get_root_module() override
    {
+        if(root_mod != nullptr)
+            return root_mod;
        assert(prog);
        return prog->get_main_module();
    }
@@ -140,7 +148,7 @@ void run_passes(program& prog, module_ref root_mod, const std::vector<pass>& pas
                continue;
            if(not visited.insert(mod).second)
                continue;
-            module_pm mpm{mod, &trace};
+            module_pm mpm{mod, root_mod, &trace};
            mpm.prog      = &prog;
            auto parents  = range(tree.equal_range(mod));
            auto nparents = distance(parents);
@@ -164,7 +172,7 @@ void run_passes(module& mod, const std::vector<pass>& passes, tracer trace)
        trace = tracer{std::cout};
    for(const auto& p : passes)
    {
-        module_pm{&mod, &trace}.run_pass(p);
+        module_pm{&mod, &mod, &trace}.run_pass(p);
    }
 }

--- a/src/program.cpp
+++ b/src/program.cpp
@@ -70,9 +70,8 @@ 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;
+    std::vector<target> targets;
 };
 program::program() : impl(std::make_unique<program_impl>()) { this->create_module("main"); }
@@ -96,14 +95,8 @@ void program::assign(const program& p)
    {
        impl = std::make_unique<program_impl>();
    }
-    else if(not impl->modules.empty())
-    {
-        impl->modules.clear();
-    }
-    impl->ctx         = p.impl->ctx;
+    *impl = *p.impl;
-    impl->target_name = p.impl->target_name;
-    impl->modules     = p.impl->modules;
    // build a map from old ins to new ins
    // Build a map from old module to new module
@@ -166,7 +159,11 @@ std::vector<shape> program::get_output_shapes() const
    return mm->get_output_shapes();
 }
-context& program::get_context() const { return impl->ctx; }
+context& program::get_context() const
+{
+    assert(impl->contexts.size() == 1);
+    return impl->contexts.front();
+}
 instruction_ref program::validate() const
 {
@@ -217,7 +214,7 @@ target_assignments program::get_target_assignments(const std::vector<target>& ta
    return p;
 }
-bool program::is_compiled() const { return not this->impl->target_name.empty(); }
+bool program::is_compiled() const { return not this->impl->contexts.empty(); }
 void program::compile(const std::vector<target>& targets, std::vector<compile_options> compile_opts)
 {
@@ -299,24 +296,24 @@ void program::compile(const std::vector<target>& targets, std::vector<compile_op
                MIGRAPHX_THROW("Dangling reference in module " + current_mod->name() +
                               " from instruction " + std::to_string(index));
            }
-            current_mod->finalize(this->impl->contexts[root_target_id]);
        }
    }
+    this->finalize();
 }
 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->targets  = {t};
-    this->impl->ctx         = t.get_context();
+    this->impl->contexts = {t.get_context()};
    if(enabled(MIGRAPHX_TRACE_COMPILE{}))
        options.trace = tracer{std::cout};
    options.trace(*this);
    options.trace();
-    auto&& passes = t.get_passes(this->impl->ctx, options);
+    auto&& passes = t.get_passes(this->impl->contexts.front(), options);
    run_passes(*this, passes, options.trace);
    auto mods = this->get_modules();
    // Validate and finalize
@@ -335,14 +332,14 @@ void program::compile(const target& t, compile_options options)
            MIGRAPHX_THROW("Dangling reference in module " + mod->name() + " from instruction " +
                           std::to_string(index));
        }
-        mod->finalize(this->impl->ctx);
+        mod->finalize(this->impl->contexts);
    }
 }
 void program::finalize()
 {
    auto* mm = this->get_main_module();
-    mm->finalize(this->impl->ctx);
+    mm->finalize(this->impl->contexts);
 }
 template <class T>
@@ -366,7 +363,7 @@ void print_statistics(std::ostream& os, const argument& a)
            os << "Min value: " << *std::min_element(t.begin(), t.end()) << ", ";
            os << "Max value: " << *std::max_element(t.begin(), t.end()) << ", ";
            double num_elements = t.size();
-            auto mean           = std::reduce(t.begin(), t.end(), 0.0) / num_elements;
+            auto mean           = std::accumulate(t.begin(), t.end(), 0.0) / num_elements;
            auto stddev         = std::sqrt(
                std::accumulate(t.begin(),
                                t.end(),
@@ -429,16 +426,15 @@ void preview_argument(std::ostream& os, const argument& a)
 template <class F>
 std::vector<argument> generic_eval(const module* mod,
-                                   context& ctx,
+                                   std::vector<context>& ctx,
                                   std::unordered_map<std::string, argument> params,
                                   std::unordered_map<instruction_ref, argument> results,
-                                   F make_trace)
+                                   F trace)
 {
    assert(mod->validate() == mod->end());
    results.reserve(mod->size() * 2);
    std::vector<argument> values;
    values.reserve(16);
-    auto trace = make_trace(mod);
    for(auto ins : iterator_for(*mod))
    {
        assert(results.find(ins) == results.end());
@@ -494,13 +490,18 @@ std::vector<argument> generic_eval(const module* mod,
            const auto& mod_args = ins->module_inputs();
            auto module_eval     = [&](module_ref smod,
                                   const std::unordered_map<std::string, argument>& inputs) {
-                auto ssctx = ctx;
+                return generic_eval(smod, ctx, inputs, results, trace);
-                return generic_eval(smod, ssctx, inputs, results, make_trace);
            };
-            results.emplace(ins, trace(ins, [&] {
+            results.emplace(
-                                return ins->normalized_operator().compute(
+                ins, trace(ins, [&] {
-                                    ctx, ins->get_shape(), values, mod_args, module_eval);
+                    auto op = ins->normalized_operator();
+                    if(op.is_context_free())
+                        return op.compute(ins->get_shape(), values, mod_args, module_eval);
+                    if(ins->get_target_id() >= ctx.size())
+                        MIGRAPHX_THROW("No context available for " + op.name());
+                    return op.compute(
+                        ctx[ins->get_target_id()], ins->get_shape(), values, mod_args, module_eval);
                }));
        }
        assert(results.find(ins) != results.end());
@@ -514,44 +515,25 @@ std::vector<argument> generic_eval(const module* mod,
 template <class F>
 std::vector<argument> generic_eval(const program& p,
-                                   context& ctx,
+                                   std::vector<context>& ctx,
                                   std::unordered_map<std::string, argument> params,
-                                   F make_trace)
+                                   F trace)
 {
    const module* mm = p.get_main_module();
-    return generic_eval(mm, ctx, params, {}, make_trace);
+    return generic_eval(mm, ctx, params, {}, trace);
 }
 std::vector<argument> program::eval(parameter_map params, execution_environment exec_env) const
 {
-    auto& ctx = this->impl->ctx;
+    auto& contexts = this->impl->contexts;
-#ifndef NDEBUG
-    auto with_check_context = [&](auto f) {
-        return [=, &ctx](auto&&) {
-            auto sctx          = std::make_shared<context>(ctx);
-            auto check_context = [=, &ctx](auto g) {
-                assert(is_shared(ctx, *sctx));
-                auto x = g();
-                *sctx  = ctx;
-                return x;
-            };
-            return [=](auto&&... xs) { return f(xs..., check_context); };
-        };
-    };
-#else
-    auto with_check_context = [](auto f) {
-        return [=](auto&&) {
-            return [=](auto&&... xs) { return f(xs..., [](auto g) { return g(); }); };
-        };
-    };
-#endif
    auto trace_level = value_of(MIGRAPHX_TRACE_EVAL{});
    std::vector<argument> ret;
    if(exec_env.async)
    {
-        ctx.wait_for(exec_env.queue);
+        assert(contexts.size() == 1);
+        contexts.front().wait_for(exec_env.queue);
    }
    if(trace_level > 0)
@@ -563,16 +545,12 @@ std::vector<argument> program::eval(parameter_map params, execution_environment
            instruction::print(ss, x, ins_names);
            ins_out[x] = ss.str();
        });
+        ret = generic_eval(*this, contexts, std::move(params), [&](instruction_ref ins, auto f) {
-        ret = generic_eval(
+            auto& ctx = contexts[ins->get_target_id()];
-            *this,
-            ctx,
-            std::move(params),
-            with_check_context([&](auto& ins, auto f, auto&& check_context) {
            ctx.finish();
            std::cout << "Run instruction: " << ins_out.at(ins) << std::endl;
            timer t{};
-                auto result = check_context(f);
+            auto result = f();
            double t1   = t.record<milliseconds>();
            ctx.finish();
            double t2 = t.record<milliseconds>();
@@ -583,7 +561,7 @@ std::vector<argument> program::eval(parameter_map params, execution_environment
                migraphx::argument buffer;
                try
                {
-                        target tgt = make_target(this->impl->target_name);
+                    const target& tgt = this->impl->targets.at(ins->get_target_id());
                    buffer            = tgt.copy_from(result);
                }
                catch(const migraphx::exception&)
@@ -593,8 +571,7 @@ std::vector<argument> program::eval(parameter_map params, execution_environment
                }
                catch(...)
                {
-                        MIGRAPHX_THROW(
+                    MIGRAPHX_THROW("MIGraphX program execution with MIGRAPHX_TRACE_EVAL failed.\n");
-                            "MIGraphX program execution with MIGRAPHX_TRACE_EVAL failed.\n");
                }
                if(trace_level == 2)
                {
@@ -611,35 +588,36 @@ std::vector<argument> program::eval(parameter_map params, execution_environment
                }
            }
            return result;
-            }));
+        });
    }
    else
    {
-        ret = generic_eval(*this,
+        ret = generic_eval(*this, contexts, std::move(params), [&](auto&&, auto f) { return f(); });
-                           ctx,
-                           std::move(params),
-                           with_check_context([&](auto&, auto f, auto&& check_context) {
-                               return check_context(f);
-                           }));
    }
    if(exec_env.async)
    {
-        ctx.finish_on(exec_env.queue);
+        assert(contexts.size() == 1);
+        contexts.front().finish_on(exec_env.queue);
    }
    return ret;
 }
-const int program_file_version = 5;
+void program::finish() const
+{
+    for(const auto& ctx : this->impl->contexts)
+        ctx.finish();
+}
+const int program_file_version = 6;
 value program::to_value() const
 {
    value result;
    result["version"]  = program_file_version;
-    result["target"]  = this->impl->target_name;
+    result["targets"]  = migraphx::to_value(this->impl->targets);
-    if(not this->impl->target_name.empty())
+    result["contexts"] = migraphx::to_value(this->impl->contexts);
-        result["context"] = this->impl->ctx.to_value();
    value module_vals = value::object{};
    std::unordered_map<instruction_ref, std::string> names;
@@ -768,12 +746,12 @@ void program::from_value(const value& v)
        MIGRAPHX_THROW("Warning: Program version mismatch");
    }
-    this->impl->target_name = v.at("target").to<std::string>();
+    migraphx::from_value(v.at("targets"), this->impl->targets);
-    if(not this->impl->target_name.empty())
+    for(auto i : range(this->impl->targets.size()))
    {
-        target t        = make_target(this->impl->target_name);
+        this->impl->contexts.push_back(this->impl->targets[i].get_context());
-        this->impl->ctx = t.get_context();
+        this->impl->contexts.back().from_value(v.at("contexts")[i]);
-        this->impl->ctx.from_value(v.at("context"));
    }
    auto module_vals = v.at("modules");
@@ -794,6 +772,8 @@ void program::from_value(const value& v)
    auto* mm = get_main_module();
    mod_from_val(mm, module_vals, map_insts, map_mods);
+    // Finalize a compiled model
+    if(not this->impl->contexts.empty())
        this->finalize();
 }
@@ -814,19 +794,19 @@ std::string perf_group(const operation& op)
 void program::mark(const parameter_map& params, marker&& m)
 {
-    auto& ctx = this->impl->ctx;
+    auto& ctx = this->impl->contexts;
    // Run once by itself
    eval(params);
-    ctx.finish();
+    this->finish();
    // Start marking
    m.mark_start(*this);
-    generic_eval(*this, ctx, params, always([&](auto ins, auto f) {
+    generic_eval(*this, ctx, params, [&](auto ins, auto f) {
        argument result;
        m.mark_start(ins);
        result = f();
        m.mark_stop(ins);
        return result;
-    }));
+    });
    m.mark_stop(*this);
 }
@@ -835,10 +815,10 @@ void program::perf_report(std::ostream& os,
                          parameter_map params,
                          std::size_t batch) const
 {
-    auto& ctx = this->impl->ctx;
+    auto& ctx = this->impl->contexts;
    // Run once by itself
    eval(params);
-    ctx.finish();
+    this->finish();
    // Run and time entire program
    std::vector<double> total_vec;
    total_vec.reserve(n);
@@ -846,28 +826,28 @@ void program::perf_report(std::ostream& os,
    {
        total_vec.push_back(time<milliseconds>([&] {
            eval(params);
-            ctx.finish();
+            this->finish();
        }));
    }
    std::sort(total_vec.begin(), total_vec.end());
    std::unordered_map<instruction_ref, std::vector<double>> ins_vec;
    // Fill the map
-    generic_eval(*this, ctx, params, always([&](auto ins, auto) {
+    generic_eval(*this, ctx, params, [&](auto ins, auto) {
        ins_vec[ins].reserve(n);
        return argument{ins->get_shape(), nullptr};
-    }));
+    });
    // Run and time each instruction
    for(std::size_t i = 0; i < n; i++)
    {
-        generic_eval(*this, ctx, params, always([&](auto ins, auto f) {
+        generic_eval(*this, ctx, params, [&](auto ins, auto f) {
            argument result;
            ins_vec[ins].push_back(time<milliseconds>([&] {
                result = f();
-                ctx.finish();
+                this->impl->contexts[ins->get_target_id()].finish();
            }));
            return result;
-        }));
+        });
    }
    for(auto&& p : ins_vec)
        std::sort(p.second.begin(), p.second.end());
@@ -1035,10 +1015,10 @@ void program::print_cpp(std::ostream& os) const
 void program::dry_run(std::unordered_map<std::string, argument> params) const
 {
-    auto& ctx = this->impl->ctx;
+    auto& ctx = this->impl->contexts;
-    generic_eval(*this, ctx, std::move(params), always([](auto ins, auto&&...) {
+    generic_eval(*this, ctx, std::move(params), [](auto ins, auto&&...) {
        return argument{ins->get_shape(), nullptr};
-    }));
+    });
 }
 void program::annotate(std::ostream& os, const std::function<void(instruction_ref)>& a) const