Add test for eliminate allocation

src/CMakeLists.txt

@@ -2,6 +2,7 @@
 add_library(migraph 
     auto_contiguous.cpp
     dead_code_elimination.cpp
+    eliminate_allocation.cpp
     eliminate_contiguous.cpp
     fwd_conv_batchnorm_rewrite.cpp
     env.cpp

src/targets/gpu/eliminate_allocation.cpp → src/eliminate_allocation.cpp

-#include <migraph/gpu/eliminate_allocation.hpp>
-#include <migraph/gpu/hip.hpp>
+#include <migraph/eliminate_allocation.hpp>
 #include <migraph/program.hpp>
 #include <migraph/instruction.hpp>
 #include <migraph/operators.hpp>
 #include <migraph/iterator_for.hpp>
 #include <migraph/ranges.hpp>
-#include <migraph/stringutils.hpp>
 
 namespace migraph {
-namespace gpu {
 
 void eliminate_allocation::apply(program& p) const
 {
+    assert(alignment > 0);
     std::size_t n = 0;
     std::vector<std::pair<instruction_ref, std::size_t>> allocs;
     for(auto ins : iterator_for(p))
     {
-        if(ins->op.name() != "hip::allocate")
+        if(ins->op.name() != allocation_op)
             continue;
         allocs.emplace_back(ins, n);
         std::size_t size = ins->get_shape().bytes();
-        n += size + (size % 32);
+        std::size_t padding = (alignment - (size % alignment)) % alignment;
+        n += size + padding;
     }
     auto mem = p.add_parameter("memory", shape{shape::int8_type, {n}});
     for(auto&& pp : allocs)
@@ -28,8 +27,7 @@ void eliminate_allocation::apply(program& p) const
         auto ins    = pp.first;
         auto s      = ins->get_shape();
         auto offset = pp.second;
-        p.replace_instruction(ins, hip_load{s, offset}, mem);
+        p.replace_instruction(ins, load{s, offset}, mem);
     }
 }
-} // namespace gpu
 } // namespace migraph

src/targets/gpu/include/migraph/gpu/eliminate_allocation.hpp → src/include/migraph/eliminate_allocation.hpp

@@ -7,14 +7,13 @@
 namespace migraph {
 struct program;
 
-namespace gpu {
-
 struct eliminate_allocation
 {
+    std::string allocation_op{};
+    std::size_t alignment = 32;
     std::string name() const { return "eliminate_allocation"; }
     void apply(program& p) const;
 };
-} // namespace gpu
 } // namespace migraph
 
 #endif

src/include/migraph/operators.hpp

@@ -534,6 +534,22 @@ struct div : binary
     std::string name() const { return "div"; }
 };
 
+struct load
+{
+    shape s;
+    std::size_t offset = 0;
+    std::string name() const { return "load"; }
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        check_shapes{inputs}.has(1);
+        return s;
+    }
+    argument compute(context&, const shape&, const std::vector<argument>& args) const
+    {
+        return {s, args[0].data() + offset};
+    }
+};
+
 struct outline
 {
     shape s;

src/targets/gpu/CMakeLists.txt

@@ -21,7 +21,6 @@ target_include_directories(migraph_device PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRE
 target_include_directories(migraph_device PRIVATE $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/device/include>)
 
 add_library(migraph_gpu
-    eliminate_allocation.cpp
     eliminate_workspace.cpp
     fuse_ops.cpp
     hip.cpp

src/targets/gpu/target.cpp

@@ -3,7 +3,7 @@
 #include <migraph/gpu/write_literals.hpp>
 #include <migraph/gpu/context.hpp>
 #include <migraph/gpu/eliminate_workspace.hpp>
-#include <migraph/gpu/eliminate_allocation.hpp>
+#include <migraph/eliminate_allocation.hpp>
 #include <migraph/gpu/fuse_ops.hpp>
 #include <migraph/check_context.hpp>
 #include <migraph/auto_contiguous.hpp>
@@ -34,7 +34,7 @@ std::vector<pass> target::get_passes(migraph::context& gctx) const
         eliminate_contiguous{},
         dead_code_elimination{},
         write_literals{&ctx},
-        eliminate_allocation{},
+        eliminate_allocation{""},
         check_context<context>{},
         dead_code_elimination{}
     };

test/eliminate_allocation_test.cpp

+#include <migraph/eliminate_allocation.hpp>
+#include <migraph/dead_code_elimination.hpp>
+#include <migraph/operators.hpp>
+#include <basic_ops.hpp>
+#include <test.hpp>
+
+struct eliminate_allocation_target
+{
+    std::size_t align = 32;
+    std::string name() const { return "eliminate_allocation"; }
+    std::vector<migraph::pass> get_passes(migraph::context&) const
+    {
+        return {migraph::eliminate_allocation{"allocate", align}, migraph::dead_code_elimination{}};
+    }
+    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}.has(0);
+        return s;
+    }
+    migraph::argument compute(migraph::context&, const migraph::shape& output_shape, const std::vector<migraph::argument>&) const
+    {
+        return {output_shape};
+    }
+};
+
+void basic()
+{
+    migraph::program p;
+    auto a1 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {8}}});
+    auto p1 = p.add_instruction(pass_op{}, a1);
+
+    auto a2 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {40}}});
+    auto p2 = p.add_instruction(pass_op{}, a2, p1);
+
+    auto a3 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {200}}});
+    p.add_instruction(pass_op{}, a3, p2);
+
+    p.compile(eliminate_allocation_target{});
+    EXPECT(p.get_shape() == migraph::shape{migraph::shape::float_type, {200}});
+    EXPECT(p.get_parameter_shape("memory").bytes() == (8*4 + 40*4 + 200*4));
+}
+
+void aligned()
+{
+    migraph::program p;
+    auto a1 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {1}}});
+    auto p1 = p.add_instruction(pass_op{}, a1);
+
+    auto a2 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {2}}});
+    auto p2 = p.add_instruction(pass_op{}, a2, p1);
+
+    auto a3 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {200}}});
+    p.add_instruction(pass_op{}, a3, p2);
+
+    p.compile(eliminate_allocation_target{});
+    EXPECT(p.get_shape() == migraph::shape{migraph::shape::float_type, {200}});
+    EXPECT(p.get_parameter_shape("memory").bytes() == (32 + 32 + 200*4));
+}
+
+void unaligned()
+{
+    migraph::program p;
+    auto a1 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {1}}});
+    auto p1 = p.add_instruction(pass_op{}, a1);
+
+    auto a2 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {2}}});
+    auto p2 = p.add_instruction(pass_op{}, a2, p1);
+
+    auto a3 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {200}}});
+    p.add_instruction(pass_op{}, a3, p2);
+
+    p.compile(eliminate_allocation_target{1});
+    EXPECT(p.get_shape() == migraph::shape{migraph::shape::float_type, {200}});
+    EXPECT(p.get_parameter_shape("memory").bytes() == (1*4 + 2*4 + 200*4));
+}
+
+void float_aligned()
+{
+    migraph::program p;
+    auto a1 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {1}}});
+    auto p1 = p.add_instruction(pass_op{}, a1);
+
+    auto a2 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {2}}});
+    auto p2 = p.add_instruction(pass_op{}, a2, p1);
+
+    auto a3 = p.add_instruction(allocate{migraph::shape{migraph::shape::float_type, {200}}});
+    p.add_instruction(pass_op{}, a3, p2);
+
+    p.compile(eliminate_allocation_target{4});
+    EXPECT(p.get_shape() == migraph::shape{migraph::shape::float_type, {200}});
+    EXPECT(p.get_parameter_shape("memory").bytes() == (1*4 + 2*4 + 200*4));
+}
+
+int main()
+{
+    basic();
+    aligned();
+    unaligned();
+    float_aligned();
+}