Merge branch 'mem_color_separate_literal'

src/targets/gpu/target.cpp

 #include <migraph/gpu/target.hpp>
 #include <migraph/gpu/lowering.hpp>
+#include <migraph/memory_coloring.hpp>
 #include <migraph/gpu/write_literals.hpp>
 #include <migraph/gpu/context.hpp>
 #include <migraph/gpu/eliminate_workspace.hpp>
@@ -28,6 +29,7 @@ std::vector<pass> target::get_passes(migraph::context& gctx) const
         simplify_reshapes{},
         dead_code_elimination{},
         lowering{ctx},
+        memory_coloring{"hip::allocate"},
         fuse_ops{},
         dead_code_elimination{},
         eliminate_contiguous{},
@@ -45,10 +47,8 @@ std::string target::name() const { return "miopen"; }
 
 migraph::context target::get_context() const
 {
-    return context{share(make_obj<miopen_handle>(&miopenCreate)),
-                   share(create_rocblas_handle_ptr())};
+    return context{
+        share(make_obj<miopen_handle>(&miopenCreate)), share(create_rocblas_handle_ptr()), {}};
 }
-
 } // namespace gpu
-
 } // namespace migraph

src/targets/gpu/write_literals.cpp

@@ -37,7 +37,5 @@ void write_literals::apply(program& p) const
         }
     }
 }
-
 } // namespace gpu
-
 } // namespace migraph

test/eliminate_allocation_test.cpp

@@ -102,6 +102,7 @@ void float_aligned()
 
 int main()
 {
+    setenv("MIGRAPH_DISABLE_MEMORY_COLORING", "1", 1);
     basic();
     aligned();
     unaligned();

test/memory_coloring_test.cpp

+#include <migraph/memory_coloring.hpp>
+#include <migraph/operators.hpp>
+#include <basic_ops.hpp>
+#include <test.hpp>
+
+struct memory_coloring_target
+{
+    std::string name() const { return "memory_coloring"; }
+    std::vector<migraph::pass> get_passes(migraph::context&) const
+    {
+        return {migraph::memory_coloring{"allocate"}};
+    }
+    migraph::context get_context() const { return {}; }
+};
+
+struct allocate
+{
+    migraph::shape s{};
+    std::string name() const { return "allocate"; }
+    migraph::shape compute_shape(const std::vector<migraph::shape>& inputs) const
+    {
+        migraph::check_shapes{inputs, *this}.has(1);
+        return inputs.front();
+    }
+    migraph::argument compute(migraph::context&,
+                              const migraph::shape& output_shape,
+                              const std::vector<migraph::argument>&) const
+    {
+        return {output_shape};
+    }
+};
+
+// A custom test operator that takes a single argument and an allocation
+// This operator's output is an operand alias of argument 1
+struct pass_memory
+{
+    std::string name() const { return "memory_coloring::pass_memory"; }
+    migraph::shape compute_shape(const std::vector<migraph::shape>& inputs) const
+    {
+        migraph::check_shapes{inputs, *this}.has(2);
+        return inputs.at(1);
+    }
+    migraph::argument compute(migraph::context&,
+                              const migraph::shape&,
+                              const std::vector<migraph::argument>& args) const
+    {
+        return args[1];
+    }
+};
+
+// The previous existing test
+void test1()
+{
+    migraph::program p;
+    auto a0 = p.add_outline(migraph::shape{migraph::shape::float_type, {8}});
+    auto a1 = p.add_instruction(allocate{}, a0);
+    auto p1 = p.add_instruction(pass_op{}, a1);
+    auto a2 = p.add_outline(migraph::shape{migraph::shape::float_type, {40}});
+    auto p2 = p.add_instruction(allocate{}, a2);
+    p.add_instruction(pass_op{}, p2, p1);
+    p.compile(memory_coloring_target{});
+    EXPECT(p.get_parameter_shape("scratch").bytes() == 192);
+}
+
+// This test uses the pass_memory operator
+void test2()
+{
+    migraph::program p;
+    auto input = p.add_parameter("input", migraph::shape{migraph::shape::float_type, {16}});
+
+    auto a0 = p.add_outline(migraph::shape{migraph::shape::float_type, {128}});
+    auto a1 = p.add_instruction(allocate{}, a0);
+    auto p1 = p.add_instruction(pass_memory{}, input, a1);
+    auto a2 = p.add_outline(migraph::shape{migraph::shape::float_type, {40}});
+    auto p2 = p.add_instruction(allocate{}, a2);
+    p.add_instruction(pass_memory{}, p1, p2);
+    p.compile(memory_coloring_target{});
+    EXPECT(p.get_parameter_shape("scratch").bytes() == 672);
+}
+
+// This test uses the pass_memory operator with two memory allocation passed together.
+// This is similar to allocations done for workspaces, that is one allocation is aliased and the
+// other is just used
+void test3()
+{
+    migraph::program p;
+    auto a0 = p.add_outline(migraph::shape{migraph::shape::float_type, {8}});
+    auto a1 = p.add_instruction(allocate{}, a0);
+    auto a2 = p.add_outline(migraph::shape{migraph::shape::float_type, {128}});
+    auto p2 = p.add_instruction(allocate{}, a2);
+    auto p1 = p.add_instruction(pass_memory{}, a1, p2);
+    auto a3 = p.add_outline(migraph::shape{migraph::shape::float_type, {40}});
+    auto p3 = p.add_instruction(allocate{}, a3);
+    p.add_instruction(pass_memory{}, p1, p3);
+    p.compile(memory_coloring_target{});
+    EXPECT(p.get_parameter_shape("scratch").bytes() == 704);
+}
+
+// Like the previous test, but this tests a zero workspace memory allocation
+void test4()
+{
+    migraph::program p;
+    auto a0 = p.add_outline(migraph::shape{migraph::shape::float_type, {0}});
+    auto a1 = p.add_instruction(allocate{}, a0);
+    auto a2 = p.add_outline(migraph::shape{migraph::shape::float_type, {128}});
+    auto p2 = p.add_instruction(allocate{}, a2);
+    auto p1 = p.add_instruction(pass_memory{}, a1, p2);
+    auto a3 = p.add_outline(migraph::shape{migraph::shape::float_type, {40}});
+    auto p3 = p.add_instruction(allocate{}, a3);
+    p.add_instruction(pass_memory{}, p1, p3);
+    p.compile(memory_coloring_target{});
+    EXPECT(p.get_parameter_shape("scratch").bytes() == 672);
+}
+
+int main()
+{
+    test1();
+    test2();
+    test3();
+    test4();
+}