Add option to do offload copying automatically (#403)

* Add compiler options

* Add copy operators

* Formatting

* Use run_passes in tests

* Formatting

* Use run_pass in schedule test

* Formatting

* Add compile_options to get_passes in target

* Formatting

* Offload copy option

* Formatting

* Copy using pinned memory

* Formatting

* Improve performance of gpu copying

* Formatting

* Dont copy

* Formatting

* Always make an extra copy

* Formatting

* Remove unused write op

* Add missing include

* Remove copy_to_gpu function in python api

* Make offload copy disabled by default on C++

* Formatting

* Fix tidy issues

* Formatting

* Fix namespace

* Fix python tests

* Turn clang format off since its broken

* Fix compile error on gcc 5

* Remove commented code

src/targets/gpu/hip.cpp

@@ -3,6 +3,7 @@
 
 #include <migraphx/manage_ptr.hpp>
 #include <migraphx/gpu/context.hpp>
+#include <migraphx/gpu/device/contiguous.hpp>
 #include <miopen/miopen.h>
 
 #include <vector>
@@ -12,6 +13,7 @@ inline namespace MIGRAPHX_INLINE_NS {
 namespace gpu {
 
 using hip_ptr      = MIGRAPHX_MANAGE_PTR(void, hipFree);
+using hip_host_ptr = MIGRAPHX_MANAGE_PTR(void, hipHostUnregister);
 
 std::string hip_error(int error) { return hipGetErrorString(static_cast<hipError_t>(error)); }
 
@@ -25,6 +27,15 @@ std::size_t get_available_gpu_memory()
     return free;
 }
 
+void* get_device_ptr(void* hptr)
+{
+    void* result = nullptr;
+    auto status  = hipHostGetDevicePointer(&result, hptr, 0);
+    if(status != hipSuccess)
+        MIGRAPHX_THROW("Failed getting device pointer: " + hip_error(status));
+    return result;
+}
+
 hip_ptr allocate_gpu(std::size_t sz, bool host = false)
 {
     if(sz > get_available_gpu_memory())
@@ -41,6 +52,14 @@ hip_ptr allocate_gpu(std::size_t sz, bool host = false)
     return hip_ptr{result};
 }
 
+hip_host_ptr register_on_gpu(void* ptr, std::size_t sz)
+{
+    auto status = hipHostRegister(ptr, sz, hipHostRegisterMapped);
+    if(status != hipSuccess)
+        MIGRAPHX_THROW("Gpu register failed: " + hip_error(status));
+    return hip_host_ptr{ptr};
+}
+
 template <class T>
 std::vector<T> read_from_gpu(const void* x, std::size_t sz)
 {
@@ -76,10 +95,20 @@ argument allocate_gpu(const shape& s, bool host)
     return {s, [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
 }
 
+argument register_on_gpu(const argument& arg)
+{
+    auto arg_shared = arg.share();
+    auto p          = share(register_on_gpu(arg_shared.data(), arg_shared.get_shape().bytes()));
+    return {arg_shared.get_shape(),
+            [ p, a = std::move(arg_shared) ]() mutable {return get_device_ptr(p.get());
+}
+}; // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+
 argument to_gpu(const argument& arg, bool host)
 {
     auto p = share(write_to_gpu(arg.data(), arg.get_shape().bytes(), host));
-    return {arg.get_shape(), [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
+    return {arg.get_shape(), p};
 }
 
 argument from_gpu(const argument& arg)
@@ -88,7 +117,8 @@ argument from_gpu(const argument& arg)
     arg.visit([&](auto x) {
         using type = typename decltype(x)::value_type;
         auto v     = read_from_gpu<type>(arg.data(), x.get_shape().bytes() / sizeof(type));
-        result     = {x.get_shape(), [v]() mutable { return reinterpret_cast<char*>(v.data()); }};
+        // cppcheck-suppress returnDanglingLifetime
+        result = {x.get_shape(), [v]() mutable { return v.data(); }};
     });
     return result;
 }
@@ -102,29 +132,40 @@ void set_device(std::size_t id)
 
 void gpu_sync() { hipDeviceSynchronize(); }
 
-void copy_to_gpu(const argument& src, const argument& dst)
+void hip_async_copy(context& ctx, const argument& src, const argument& dst, hipMemcpyKind kind)
 {
     std::size_t src_size = src.get_shape().bytes();
     std::size_t dst_size = dst.get_shape().bytes();
     if(src_size > dst_size)
         MIGRAPHX_THROW("Not enough memory available in destination to do copy");
-    auto status = hipMemcpy(dst.data(), src.data(), src_size, hipMemcpyHostToDevice);
+    auto status = hipMemcpyAsync(dst.data(), src.data(), src_size, kind, ctx.get_stream().get());
     if(status != hipSuccess)
-        MIGRAPHX_THROW("Copy to gpu failed: " + hip_error(status));
+        MIGRAPHX_THROW("Gpu copy failed: " + hip_error(status));
 }
 
 void gpu_copy(context& ctx, const argument& src, const argument& dst)
 {
-    std::size_t src_size = src.get_shape().bytes();
-    std::size_t dst_size = dst.get_shape().bytes();
-    if(src_size > dst_size)
-        MIGRAPHX_THROW("Not enough memory available in destination to do copy");
-    auto status = hipMemcpyAsync(
-        dst.data(), src.data(), src_size, hipMemcpyDeviceToDevice, ctx.get_stream().get());
-    if(status != hipSuccess)
-        MIGRAPHX_THROW("Gpu copy failed: " + hip_error(status));
+    // Workaround: Use contiguous as hip's memcpy is broken
+    device::contiguous(ctx.get_stream().get(), dst, src);
+    // hip_async_copy(ctx, src, dst, hipMemcpyDeviceToDevice);
+}
+
+void copy_to_gpu(context& ctx, const argument& src, const argument& dst)
+{
+    gpu_copy(ctx, register_on_gpu(src), dst);
+}
+
+void copy_from_gpu(context& ctx, const argument& src, const argument& dst)
+{
+    gpu_copy(ctx, src, register_on_gpu(dst));
+}
+
+argument get_preallocation(context& ctx, const std::string& id)
+{
+    return ctx.get_current_device().preallocations.at(id);
 }
 
+// clang-format off
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/gpu/include/migraphx/gpu/context.hpp

@@ -6,6 +6,7 @@
 #include <migraphx/gpu/hip.hpp>
 #include <migraphx/env.hpp>
 #include <migraphx/config.hpp>
+#include <unordered_map>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -123,6 +124,9 @@ struct hip_device
     std::size_t device_id      = 0;
     std::size_t current_stream = 0;
     std::vector<stream> streams;
+
+    public:
+    std::unordered_map<std::string, argument> preallocations{};
 };
 
 struct context

src/targets/gpu/include/migraphx/gpu/device/arg_op.hpp

@@ -7,7 +7,6 @@
 #include <migraphx/gpu/device/launch.hpp>
 #include <migraphx/gpu/device/types.hpp>
 #include <migraphx/gpu/device/reduce.hpp>
-#include <migraphx/gpu/hip.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {

src/targets/gpu/include/migraphx/gpu/hip.hpp

@@ -14,6 +14,8 @@ struct context;
 
 argument allocate_gpu(const shape& s, bool host = false);
 
+argument register_on_gpu(const argument& arg);
+
 argument to_gpu(const argument& arg, bool host = false);
 
 argument from_gpu(const argument& arg);
@@ -22,9 +24,11 @@ void set_device(std::size_t id);
 
 void gpu_sync();
 
-void copy_to_gpu(const argument& src, const argument& dst);
-
 void gpu_copy(context& ctx, const argument& src, const argument& dst);
+void copy_to_gpu(context& ctx, const argument& src, const argument& dst);
+void copy_from_gpu(context& ctx, const argument& src, const argument& dst);
+
+argument get_preallocation(context& ctx, const std::string& id);
 
 struct hip_allocate
 {
@@ -77,40 +81,63 @@ struct hip_sync
     }
 };
 
-struct hip_write
+struct hip_copy_to_gpu
 {
-    std::string name() const { return "hip::write"; }
-    shape compute_shape(const std::vector<shape>& inputs) const
+    std::string name() const { return "hip::copy_to_gpu"; }
+    shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs}.has(1);
-        return inputs.front();
+        check_shapes{inputs}.has(1, 2);
+        return inputs.at(0);
     }
-    argument compute(context&, const shape&, const std::vector<argument>& args) const
+    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
     {
-        return to_gpu(args.front());
+        auto input = register_on_gpu(args[0]);
+        if(args.size() == 1)
+            return input;
+        argument result = args[1].share();
+        gpu_copy(ctx, input, result);
+        // Associate the input since it was registered with hip
+        return {result.get_shape(), [input, result]() mutable { return result.data(); }};
+    }
+    std::ptrdiff_t output_alias(const std::vector<shape>& args) const
+    {
+        if(args.size() == 1)
+            return -1;
+        return 1;
     }
-    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };
 
-struct hip_copy_to_gpu
+struct hip_copy_from_gpu
 {
-    std::string name() const { return "hip_copy_to_gpu"; }
+    std::string name() const { return "hip::copy_from_gpu"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs}.has(2);
-        return inputs.at(1);
+        check_shapes{inputs}.has(1, 2);
+        return inputs.at(0);
     }
-    argument compute(context&, const shape&, std::vector<argument> args) const
+    argument
+    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
+    {
+        if(args.size() == 1)
         {
-        copy_to_gpu(args[0], args[1]);
+            argument result = allocate_gpu(output_shape, true);
+            gpu_copy(ctx, args[0], result);
+            return result;
+        }
+        copy_from_gpu(ctx, args[0], args[1]);
         return args[1];
     }
-    std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 1; }
+    std::ptrdiff_t output_alias(const std::vector<shape>& args) const
+    {
+        if(args.size() == 1)
+            return -1;
+        return 1;
+    }
 };
 
 struct hip_copy
 {
-    std::string name() const { return "hip_copy"; }
+    std::string name() const { return "hip::copy"; }
     shape compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs}.has(2).standard();
@@ -124,6 +151,29 @@ struct hip_copy
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 1; }
 };
 
+struct hip_load_memory
+{
+    shape s;
+    std::string id{};
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.s, "shape"), f(self.id, "id"));
+    }
+
+    std::string name() const { return "hip::hip_load_memory"; }
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        check_shapes{inputs}.has(0);
+        return s;
+    }
+    argument compute(context& ctx, const shape&, const std::vector<argument>&) const
+    {
+        return get_preallocation(ctx, id);
+    }
+};
+
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/gpu/include/migraphx/gpu/lowering.hpp

@@ -11,7 +11,8 @@ namespace gpu {
 
 struct lowering
 {
-    context ctx;
+    context* ctx;
+    bool offload_copy;
     std::string name() const { return "gpu::lowering"; }
     void apply(program& p) const;
 };

src/targets/gpu/include/migraphx/gpu/preallocate_param.hpp

+#ifndef MIGRAPHX_GUARD_RTGLIB_GPU_PREALLOCATE_PARAM_HPP
+#define MIGRAPHX_GUARD_RTGLIB_GPU_PREALLOCATE_PARAM_HPP
+
+#include <string>
+#include <migraphx/instruction_ref.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+struct program;
+
+namespace gpu {
+
+struct preallocate_param
+{
+    std::string param{};
+    context* ctx = nullptr;
+    std::string name() const { return "preallocate_param"; }
+    void apply(program& p) const;
+};
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/targets/gpu/include/migraphx/gpu/target.hpp

@@ -2,6 +2,7 @@
 #define MIGRAPHX_GUARD_MIGRAPHLIB_MIOPEN_TARGET_HPP
 
 #include <migraphx/program.hpp>
+#include <migraphx/compile_options.hpp>
 #include <migraphx/config.hpp>
 
 namespace migraphx {
@@ -11,7 +12,7 @@ namespace gpu {
 struct target
 {
     std::string name() const;
-    std::vector<pass> get_passes(migraphx::context& gctx) const;
+    std::vector<pass> get_passes(migraphx::context& gctx, const compile_options& options) const;
     migraphx::context get_context() const;
 
     argument copy_to(const argument& arg) const;

src/targets/gpu/lowering.cpp

@@ -76,10 +76,17 @@ namespace gpu {
 struct miopen_apply
 {
     program* prog        = nullptr;
-    context ctx{};
+    const lowering* pass = nullptr;
     std::unordered_map<std::string, std::function<instruction_ref(instruction_ref)>> apply_map{};
     instruction_ref last{};
 
+    context& get_context()
+    {
+        assert(pass != nullptr);
+        assert(pass->ctx != nullptr);
+        return *pass->ctx;
+    }
+
     void check_shape(shape x, instruction_ref i)
     {
         assert(x == i->get_shape());
@@ -89,6 +96,8 @@ struct miopen_apply
 
     void init()
     {
+        assert(prog != nullptr);
+        assert(pass != nullptr);
         this->last = instruction::get_output_alias(std::prev(prog->end()));
 
         add_miopen_simple_op<miopen_abs>("abs", make_abs);
@@ -149,6 +158,23 @@ struct miopen_apply
         add_batch_norm_inference_op();
     }
 
+    void copy_params()
+    {
+        if(not pass->offload_copy)
+            return;
+        for(auto ins : iterator_for(*prog))
+        {
+            if(ins->name() != "@param")
+                continue;
+            auto pos = std::next(ins);
+            auto a   = insert_allocation(pos, ins->get_shape());
+            auto c   = prog->insert_instruction(pos, hip_copy_to_gpu{}, ins, a);
+            prog->replace_instruction(ins, c);
+        }
+        auto end = std::prev(prog->end());
+        prog->add_instruction(hip_copy_from_gpu{}, end);
+    }
+
     void apply()
     {
         init();
@@ -160,11 +186,12 @@ struct miopen_apply
                 check_shape(s, apply_map.at(it->name())(it));
             }
         }
+        copy_params();
     }
 
     instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
     {
-        if(ins == last and tag.empty())
+        if(not pass->offload_copy and ins == last and tag.empty())
         {
             return prog->add_parameter("output", s);
         }
@@ -181,7 +208,7 @@ struct miopen_apply
             auto&& op = any_cast<op::convolution>(ins->get_operator());
 
             auto conv = miopen_convolution{op, make_conv(op)};
-            auto ws   = conv.compile(ctx, ins->get_shape(), to_shapes(ins->inputs()));
+            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
 
             auto workspace = insert_allocation(ins, ws, "workspace");
             auto output    = insert_allocation(ins, ins->get_shape());
@@ -229,7 +256,7 @@ struct miopen_apply
         apply_map.emplace("quant_convolution", [=](instruction_ref ins) {
             auto&& op = any_cast<op::quant_convolution>(ins->get_operator());
             auto conv = miopen_quant_convolution{op, make_conv(op)};
-            auto ws   = conv.compile(ctx, ins->get_shape(), to_shapes(ins->inputs()));
+            auto ws   = conv.compile(get_context(), ins->get_shape(), to_shapes(ins->inputs()));
 
             auto args      = ins->inputs();
             auto workspace = insert_allocation(ins, ws, "workspace");
@@ -334,7 +361,7 @@ struct miopen_apply
     }
 };
 
-void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
+void lowering::apply(program& p) const { miopen_apply{&p, this}.apply(); }
 } // namespace gpu
 } // namespace MIGRAPHX_INLINE_NS
 } // namespace migraphx

src/targets/gpu/preallocate_param.cpp

+#include <migraphx/gpu/preallocate_param.hpp>
+#include <migraphx/gpu/hip.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/stringutils.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+void preallocate_param::apply(program& p) const
+{
+    for(auto ins : iterator_for(p))
+    {
+        if(ins->name() != "@param")
+            continue;
+        std::string id = any_cast<builtin::param>(ins->get_operator()).parameter;
+        if(id != param)
+            continue;
+        argument a                                   = allocate_gpu(ins->get_shape());
+        ctx->get_current_device().preallocations[id] = a;
+        auto r = p.insert_instruction(ins, hip_load_memory{a.get_shape(), id});
+        p.replace_instruction(ins, r);
+    }
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/schedule_model.cpp

@@ -96,6 +96,7 @@ void schedule_model::record(program& p, instruction_ref ins, std::size_t wait_id
 static std::unordered_map<std::string, std::size_t> create_weight_map()
 {
     return {{"hip::load_literal", 0},
+            {"hip::hip_load_memory", 0},
             {"hip::allocate", 0},
             {"gpu::convolution", 8},
             {"gpu::conv_bias_relu", 8},

src/targets/gpu/target.cpp

@@ -22,6 +22,7 @@
 #include <migraphx/gpu/concat_gpu_opt.hpp>
 #include <migraphx/gpu/schedule_model.hpp>
 #include <migraphx/gpu/adjust_allocation.hpp>
+#include <migraphx/gpu/preallocate_param.hpp>
 #include <migraphx/gpu/pack_int8_args.hpp>
 #include <migraphx/eliminate_pad.hpp>
 #include <migraphx/schedule.hpp>
@@ -32,7 +33,7 @@ namespace gpu {
 
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_SCHEDULE_PASS)
 
-std::vector<pass> target::get_passes(migraphx::context& gctx) const
+std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_options& options) const
 {
     auto& ctx = any_cast<context>(gctx);
     // clang-format off
@@ -58,7 +59,7 @@ std::vector<pass> target::get_passes(migraphx::context& gctx) const
         dead_code_elimination{},
         propagate_constant{},
         dead_code_elimination{},
-        lowering{ctx},
+        lowering{&ctx, options.offload_copy},
         eliminate_contiguous{},
         dead_code_elimination{},
         eliminate_concat{concat_gpu_optimization{}},
@@ -72,6 +73,7 @@ std::vector<pass> target::get_passes(migraphx::context& gctx) const
         write_literals{&ctx},
         schedule{gpu::schedule_model{ctx.get_current_device().nstreams()}, not enabled(MIGRAPHX_DISABLE_SCHEDULE_PASS{})},
         memory_coloring{"hip::allocate"},
+        preallocate_param{"scratch", &ctx},
         dead_code_elimination{},
         eliminate_workspace{},
         eliminate_allocation{"hip::allocate"},

src/targets/gpu/write_literals.cpp

@@ -10,32 +10,10 @@ namespace gpu {
 
 MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_COPY_LITERALS)
 
-struct hip_load_literal
-{
-    shape s;
-    std::size_t n = 0;
-
-    template <class Self, class F>
-    static auto reflect(Self& self, F f)
-    {
-        return pack(f(self.s, "shape"), f(self.n, "id"));
-    }
-
-    std::string name() const { return "hip::load_literal"; }
-    shape compute_shape(const std::vector<shape>& inputs) const
-    {
-        check_shapes{inputs}.has(0);
-        return s;
-    }
-    argument compute(context& ctx, const shape&, const std::vector<argument>&) const
-    {
-        return ctx.literals.at(n);
-    }
-};
-
 void write_literals::apply(program& p) const
 {
     assert(ctx != nullptr);
+    std::size_t n = 0;
     for(auto ins : iterator_for(p))
     {
         if(ins->name() == "@literal")
@@ -49,10 +27,11 @@ void write_literals::apply(program& p) const
             }
             else
             {
+                std::string id = "@literal:" + std::to_string(n);
                 argument a     = to_gpu(ins->get_literal().get_argument());
-                std::size_t n = ctx->literals.size();
-                ctx->literals.push_back(a);
-                p.replace_instruction(ins, hip_load_literal{a.get_shape(), n});
+                ctx->get_current_device().preallocations[id] = a;
+                p.replace_instruction(ins, hip_load_memory{a.get_shape(), id});
+                n++;
             }
         }
     }

test/auto_contiguous_test.cpp

@@ -2,18 +2,11 @@
 #include <migraphx/op/transpose.hpp>
 #include <migraphx/op/broadcast.hpp>
 #include <migraphx/instruction.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <basic_ops.hpp>
 #include <test.hpp>
 
-struct contiguous_target
-{
-    std::string name() const { return "contiguous"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::auto_contiguous{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+void run_pass(migraphx::program& p) { migraphx::run_passes(p, {migraphx::auto_contiguous{}}); }
 
 // TODO: Add this test case
 void literal_broadcast()
@@ -22,7 +15,7 @@ void literal_broadcast()
     p.add_literal(get_2_broadcasted());
     EXPECT(not p.get_shape().standard());
     EXPECT(p.get_shape().broadcasted());
-    p.compile(contiguous_target{});
+    run_pass(p);
     EXPECT(p.get_shape().standard());
     EXPECT(not p.get_shape().broadcasted());
 }
@@ -33,7 +26,7 @@ TEST_CASE(literal_transpose)
     p.add_literal(get_2x2_transposed());
     EXPECT(not p.get_shape().standard());
     EXPECT(p.get_shape().transposed());
-    p.compile(contiguous_target{});
+    run_pass(p);
     EXPECT(p.get_shape().standard());
     EXPECT(not p.get_shape().transposed());
 }
@@ -48,7 +41,7 @@ TEST_CASE(after_literal_transpose)
     p.add_instruction(pass_op{}, t);
     EXPECT(not p.get_shape().standard());
     EXPECT(p.get_shape().transposed());
-    p.compile(contiguous_target{});
+    run_pass(p);
     EXPECT(p.get_shape().standard());
     EXPECT(not p.get_shape().transposed());
 }
@@ -64,7 +57,7 @@ TEST_CASE(after_literal_broadcast)
     p.add_instruction(pass_op{}, b);
     EXPECT(not p.get_shape().standard());
     EXPECT(p.get_shape().broadcasted());
-    p.compile(contiguous_target{});
+    run_pass(p);
     EXPECT(p.get_shape().standard());
     EXPECT(not p.get_shape().broadcasted());
 }
@@ -79,7 +72,7 @@ TEST_CASE(after_param_transpose)
     p.add_instruction(pass_op{}, t);
     EXPECT(not p.get_shape().standard());
     EXPECT(p.get_shape().transposed());
-    p.compile(contiguous_target{});
+    run_pass(p);
     EXPECT(p.get_shape().standard());
     EXPECT(not p.get_shape().transposed());
 }
@@ -95,7 +88,7 @@ TEST_CASE(after_param_broadcast)
     p.add_instruction(pass_op{}, b);
     EXPECT(not p.get_shape().standard());
     EXPECT(p.get_shape().broadcasted());
-    p.compile(contiguous_target{});
+    run_pass(p);
     EXPECT(p.get_shape().standard());
     EXPECT(not p.get_shape().broadcasted());
 }

test/dead_code_elimination_test.cpp

 #include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <basic_ops.hpp>
 #include <migraphx/op/abnormal_ops.hpp>
 #include <migraphx/op/add.hpp>
 #include <migraphx/op/identity.hpp>
 #include <test.hpp>
 
-struct dce_target
+void run_pass(migraphx::program& p)
 {
-    std::string name() const { return "dce"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::dead_code_elimination{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+    migraphx::run_passes(p, {migraphx::dead_code_elimination{}});
+}
 
 TEST_CASE(simple_test)
 {
@@ -23,7 +19,7 @@ TEST_CASE(simple_test)
     auto two = p.add_literal(2);
     p.add_instruction(sum_op{}, one, two);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count);
     auto result = p.eval({});
     EXPECT(result == migraphx::literal{3});
@@ -39,7 +35,7 @@ TEST_CASE(simple_test_nop)
     p.add_instruction(nop{});
     p.add_instruction(sum_op{}, one, two);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count);
     auto result = p.eval({});
     EXPECT(result == migraphx::literal{3});
@@ -55,7 +51,7 @@ TEST_CASE(simple_test_nop2)
     p.add_instruction(nop{});
     p.add_instruction(sum_op{}, one, two);
     p.add_instruction(nop{});
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == 2);
     auto result = p.eval({});
     EXPECT(result == migraphx::literal{});
@@ -71,7 +67,7 @@ TEST_CASE(duplicate_test1)
     p.add_instruction(sum_op{}, one, two);
     p.add_instruction(sum_op{}, one, two);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == (count - 1));
     auto result = p.eval({});
     EXPECT(result == migraphx::literal{3});
@@ -88,7 +84,7 @@ TEST_CASE(duplicate_test2)
     p.add_instruction(minus_op{}, one, two);
     p.add_instruction(sum_op{}, one, two);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == (count - 2));
     auto result = p.eval({});
     EXPECT(result == migraphx::literal{3});
@@ -107,7 +103,7 @@ TEST_CASE(depth_test)
     p.add_instruction(minus_op{}, x1, x2);
     p.add_instruction(sum_op{}, one, two);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == (count - 4));
     auto result = p.eval({});
     EXPECT(result == migraphx::literal{3});
@@ -123,7 +119,7 @@ TEST_CASE(undefined_test)
     auto undef = p.add_instruction(migraphx::op::undefined{});
     p.add_instruction(sum_op{}, one, two);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count - 1);
     EXPECT(not p.has_instruction(undef));
     auto result = p.eval({});
@@ -140,7 +136,7 @@ TEST_CASE(duplicate_args1)
     p.add_instruction(migraphx::op::add{}, l3, l3);
     p.add_instruction(migraphx::op::identity{}, l0);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) != count);
     EXPECT(std::distance(p.begin(), p.end()) == 2);
     auto result = p.eval({});
@@ -157,7 +153,7 @@ TEST_CASE(duplicate_args2)
     p.add_instruction(migraphx::op::add{}, sum1, l3);
     p.add_instruction(migraphx::op::identity{}, l0);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) != count);
     EXPECT(std::distance(p.begin(), p.end()) == 2);
     auto result = p.eval({});
@@ -175,7 +171,7 @@ TEST_CASE(duplicate_args3)
     p.add_instruction(migraphx::op::add{}, sum2, l3);
     p.add_instruction(migraphx::op::identity{}, l0);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(dce_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) != count);
     EXPECT(std::distance(p.begin(), p.end()) == 2);
     auto result = p.eval({});

test/eliminate_allocation_test.cpp

 #include <migraphx/eliminate_allocation.hpp>
 #include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/check_shapes.hpp>
 #include <migraphx/argument.hpp>
 #include <basic_ops.hpp>
 #include <test.hpp>
 
-struct eliminate_allocation_target
+void run_pass(migraphx::program& p, std::size_t align = 32)
 {
-    std::size_t align = 32;
-    std::string name() const { return "eliminate_allocation"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::eliminate_allocation{"allocate", align},
-                migraphx::dead_code_elimination{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+    migraphx::run_passes(
+        p, {migraphx::eliminate_allocation{"allocate", align}, migraphx::dead_code_elimination{}});
+}
 
 struct allocate
 {
@@ -53,7 +48,7 @@ TEST_CASE(basic)
     auto a3 = p.add_instruction(allocate{migraphx::shape{migraphx::shape::float_type, {200}}});
     p.add_instruction(pass_op{}, a3, p2);
 
-    p.compile(eliminate_allocation_target{});
+    run_pass(p);
     EXPECT(p.get_shape() == migraphx::shape{migraphx::shape::float_type, {200}});
     EXPECT(p.get_parameter_shape("memory").bytes() == (8 * 4 + 40 * 4 + 200 * 4));
 }
@@ -70,7 +65,7 @@ TEST_CASE(aligned)
     auto a3 = p.add_instruction(allocate{migraphx::shape{migraphx::shape::float_type, {200}}});
     p.add_instruction(pass_op{}, a3, p2);
 
-    p.compile(eliminate_allocation_target{});
+    run_pass(p);
     EXPECT(p.get_shape() == migraphx::shape{migraphx::shape::float_type, {200}});
     EXPECT(p.get_parameter_shape("memory").bytes() == (32 + 32 + 200 * 4));
 }
@@ -87,7 +82,7 @@ TEST_CASE(unaligned)
     auto a3 = p.add_instruction(allocate{migraphx::shape{migraphx::shape::float_type, {200}}});
     p.add_instruction(pass_op{}, a3, p2);
 
-    p.compile(eliminate_allocation_target{1});
+    run_pass(p, 1);
     EXPECT(p.get_shape() == migraphx::shape{migraphx::shape::float_type, {200}});
     EXPECT(p.get_parameter_shape("memory").bytes() == (1 * 4 + 2 * 4 + 200 * 4));
 }
@@ -104,13 +99,9 @@ TEST_CASE(float_aligned)
     auto a3 = p.add_instruction(allocate{migraphx::shape{migraphx::shape::float_type, {200}}});
     p.add_instruction(pass_op{}, a3, p2);
 
-    p.compile(eliminate_allocation_target{4});
+    run_pass(p, 4);
     EXPECT(p.get_shape() == migraphx::shape{migraphx::shape::float_type, {200}});
     EXPECT(p.get_parameter_shape("memory").bytes() == (1 * 4 + 2 * 4 + 200 * 4));
 }
 
-int main(int argc, const char* argv[])
-{
-    setenv("MIGRAPHX_DISABLE_MEMORY_COLORING", "1", 1);
-    test::run(argc, argv);
-}
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/eliminate_common_subexpression_test.cpp

 #include <migraphx/eliminate_common_subexpression.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/dead_code_elimination.hpp>
 #include <migraphx/op/add.hpp>
 #include <basic_ops.hpp>
 #include <test.hpp>
 
-struct cse_target
+void run_pass(migraphx::program& p)
 {
-    std::string name() const { return "dce"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::eliminate_common_subexpression{}, migraphx::dead_code_elimination{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+    migraphx::run_passes(
+        p, {migraphx::eliminate_common_subexpression{}, migraphx::dead_code_elimination{}});
+}
 
 TEST_CASE(cse_test1)
 {
@@ -25,7 +22,7 @@ TEST_CASE(cse_test1)
         auto sum3 = p1.add_instruction(migraphx::op::add{}, sum1, sum2);
         p1.add_instruction(pass_op{}, sum3);
     }
-    p1.compile(cse_target{});
+    run_pass(p1);
 
     migraphx::program p2;
     {
@@ -49,7 +46,7 @@ TEST_CASE(cse_test2)
         auto sum3 = p1.add_instruction(migraphx::op::add{}, sum1, sum2);
         p1.add_instruction(pass_op{}, sum3);
     }
-    p1.compile(cse_target{});
+    run_pass(p1);
 
     migraphx::program p2;
     {
@@ -74,7 +71,7 @@ TEST_CASE(cse_test3)
         auto sum3 = p1.add_instruction(migraphx::op::add{}, sum1, sum2);
         p1.add_instruction(pass_op{}, sum3);
     }
-    p1.compile(cse_target{});
+    run_pass(p1);
 
     migraphx::program p2;
     {
@@ -99,7 +96,7 @@ TEST_CASE(cse_test4)
         auto sum5 = p1.add_instruction(migraphx::op::add{}, sum4, sum3);
         p1.add_instruction(pass_op{}, sum5);
     }
-    p1.compile(cse_target{});
+    run_pass(p1);
 
     migraphx::program p2;
     {

test/eliminate_concat_test.cpp

 #include <migraphx/eliminate_concat.hpp>
 #include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/op/concat.hpp>
 #include <migraphx/op/load.hpp>
 #include <migraphx/op/identity.hpp>
@@ -43,17 +44,12 @@ struct concat_test_optimization
     }
 };
 
-struct eliminate_concat_target
+void run_pass(migraphx::program& p)
 {
-    std::size_t align = 32;
-    std::string name() const { return "eliminate_target"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::eliminate_concat{concat_test_optimization{}},
-                migraphx::dead_code_elimination{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+    migraphx::run_passes(p,
+                         {migraphx::eliminate_concat{concat_test_optimization{}},
+                          migraphx::dead_code_elimination{}});
+}
 
 struct allocate
 {
@@ -131,7 +127,7 @@ TEST_CASE(simple)
 
     auto p1 = create_test_program();
     auto p2 = create_control_program();
-    p1.compile(eliminate_concat_target{});
+    run_pass(p1);
 
     EXPECT(p1 == p2);
 }
@@ -176,7 +172,7 @@ TEST_CASE(nested)
 
     auto p1 = create_test_program();
     auto p2 = create_control_program();
-    p1.compile(eliminate_concat_target{});
+    run_pass(p1);
 
     EXPECT(p1 == p2);
 }
@@ -219,7 +215,7 @@ TEST_CASE(basic)
 
     auto p1 = create_test_program();
     auto p2 = create_control_program();
-    p1.compile(eliminate_concat_target{});
+    run_pass(p1);
 
     EXPECT(p1 == p2);
 }
@@ -263,7 +259,7 @@ TEST_CASE(wont_work)
 
     auto p1 = create_test_program();
     auto p2 = create_control_program();
-    p1.compile(eliminate_concat_target{});
+    run_pass(p1);
 
     EXPECT(p1 == p2);
 }

test/eliminate_contiguous_test.cpp

 #include <migraphx/eliminate_contiguous.hpp>
 #include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/op/identity.hpp>
 #include <migraphx/op/dot.hpp>
 #include <migraphx/op/sin.hpp>
@@ -9,15 +10,10 @@
 #include <basic_ops.hpp>
 #include <test.hpp>
 
-struct eliminate_contiguous_target
+void run_pass(migraphx::program& p)
 {
-    std::string name() const { return "eliminate_contiguous"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::eliminate_contiguous{}, migraphx::dead_code_elimination{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+    migraphx::run_passes(p, {migraphx::eliminate_contiguous{}, migraphx::dead_code_elimination{}});
+}
 
 TEST_CASE(standard_op)
 {
@@ -27,7 +23,7 @@ TEST_CASE(standard_op)
     auto c = p.add_instruction(migraphx::op::contiguous{}, t);
     p.add_instruction(pass_standard_op{}, c);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count);
 }
 
@@ -38,7 +34,7 @@ TEST_CASE(standard_op_const)
     auto t = p.add_instruction(migraphx::op::transpose{{1, 0}}, l);
     auto c = p.add_instruction(migraphx::op::contiguous{}, t);
     p.add_instruction(pass_standard_op{}, c);
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == 2);
 }
 
@@ -50,7 +46,7 @@ TEST_CASE(non_standard_op)
     auto c = p.add_instruction(migraphx::op::contiguous{}, t);
     p.add_instruction(pass_op{}, c);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count);
 }
 
@@ -61,7 +57,7 @@ TEST_CASE(non_standard_op_const)
     auto t = p.add_instruction(migraphx::op::transpose{{1, 0}}, l);
     auto c = p.add_instruction(migraphx::op::contiguous{}, t);
     p.add_instruction(pass_op{}, c);
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == 2);
 }
 
@@ -74,7 +70,7 @@ TEST_CASE(transpose_gemm)
     auto ic = p.add_instruction(migraphx::op::identity{}, c);
     p.add_instruction(migraphx::op::dot{}, ic, l);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == (count - 1));
 }
 
@@ -87,7 +83,7 @@ TEST_CASE(transpose_standard_op)
     auto sn = p.add_instruction(migraphx::op::sin{}, c);
     p.add_instruction(pass_standard_op{}, sn);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count);
 }
 
@@ -99,7 +95,7 @@ TEST_CASE(transpose_standard_op_const)
     auto c  = p.add_instruction(migraphx::op::contiguous{}, t);
     auto sn = p.add_instruction(migraphx::op::sin{}, c);
     p.add_instruction(pass_standard_op{}, sn);
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == 3);
 }
 
@@ -112,7 +108,7 @@ TEST_CASE(no_packed_unary_op)
     auto sn = p.add_instruction(migraphx::op::sin{}, c);
     p.add_instruction(pass_standard_op{}, sn);
     auto count = std::distance(p.begin(), p.end());
-    p.compile(eliminate_contiguous_target{});
+    run_pass(p);
     EXPECT(std::distance(p.begin(), p.end()) == count - 1);
 }
 

test/eliminate_identity_test.cpp

 #include <migraphx/dead_code_elimination.hpp>
 #include <migraphx/eliminate_identity.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/instruction.hpp>
 #include <basic_ops.hpp>
 #include <migraphx/op/identity.hpp>
 #include <test.hpp>
 
-struct eliminate_identity_target
-{
-    std::string name() const { return "eliminate_identity"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::eliminate_identity{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+void run_pass(migraphx::program& p) { migraphx::run_passes(p, {migraphx::eliminate_identity{}}); }
 
 TEST_CASE(simple_test)
 {
@@ -24,7 +17,7 @@ TEST_CASE(simple_test)
     auto two          = p.add_literal(2);
     auto two_identity = p.add_instruction(migraphx::op::identity{}, two);
     p.add_instruction(sum_op{}, one_identity, two_identity);
-    p.compile(eliminate_identity_target{});
+    run_pass(p);
     EXPECT(std::none_of(p.begin(), p.end(), [](const migraphx::instruction& ins) {
         return ins.name() == "identity";
     }));
@@ -40,7 +33,7 @@ TEST_CASE(simple_test_end)
     auto two = p.add_literal(2);
     auto ans = p.add_instruction(sum_op{}, one, two);
     p.add_instruction(migraphx::op::identity{}, ans);
-    p.compile(eliminate_identity_target{});
+    run_pass(p);
     EXPECT(std::none_of(p.begin(), p.end(), [](const migraphx::instruction& ins) {
         return ins.name() == "identity";
     }));
@@ -58,7 +51,7 @@ TEST_CASE(simple_test_end_dependency)
     auto ans   = p.add_instruction(sum_op{}, one, two);
     p.add_instruction(sum_op{}, ans, three);
     p.add_instruction(migraphx::op::identity{}, ans);
-    p.compile(eliminate_identity_target{});
+    run_pass(p);
     EXPECT(std::any_of(p.begin(), p.end(), [](const migraphx::instruction& ins) {
         return ins.name() == "identity";
     }));

test/eliminate_pad_test.cpp

 #include <migraphx/dead_code_elimination.hpp>
 #include <migraphx/eliminate_pad.hpp>
+#include <migraphx/pass_manager.hpp>
 #include <migraphx/instruction.hpp>
 #include <basic_ops.hpp>
 #include <migraphx/operators.hpp>
 #include <test.hpp>
 
-struct eliminate_pad_target
+void run_pass(migraphx::program& p)
 {
-    std::string name() const { return "eliminate_pad"; }
-    std::vector<migraphx::pass> get_passes(migraphx::context&) const
-    {
-        return {migraphx::eliminate_pad{}, migraphx::dead_code_elimination{}};
-    }
-    migraphx::context get_context() const { return {}; }
-};
+    migraphx::run_passes(p, {migraphx::eliminate_pad{}, migraphx::dead_code_elimination{}});
+}
 
 migraphx::instruction_ref
 create_im2col(migraphx::instruction_ref& l_img, size_t channels, migraphx::program& p)
@@ -59,7 +55,7 @@ TEST_CASE(rewrite_test)
     auto l2 = p.add_instruction(migraphx::op::pooling{}, padded_img);
     p.add_instruction(migraphx::op::identity{}, l0, l1, l2);
 
-    p.compile(eliminate_pad_target{});
+    run_pass(p);
     EXPECT(std::none_of(
         p.begin(), p.end(), [](const migraphx::instruction& ins) { return ins.name() == "pad"; }));
 }
@@ -78,7 +74,7 @@ TEST_CASE(rewrite_test_asymmetric)
 
     create_im2col(padded_img, channels, p);
 
-    p.compile(eliminate_pad_target{});
+    run_pass(p);
     EXPECT(std::any_of(
         p.begin(), p.end(), [](const migraphx::instruction& ins) { return ins.name() == "pad"; }));
 }