Merge pull request #246 from ROCmSoftwarePlatform/adjust_gpu_allocation

Adjust gpu allocation

src/targets/cpu/lowering.cpp

@@ -602,6 +602,7 @@ struct cpu_unary
                 std::transform(input.begin(), input.end(), output.begin(), op.fcn());
             });
         });
+
         return result;
     }
 };

src/targets/gpu/CMakeLists.txt

@@ -65,6 +65,7 @@ add_library(migraphx_gpu
     gather.cpp
     lrn.cpp
     schedule_model.cpp
+    adjust_allocation.cpp
 )
 set_target_properties(migraphx_gpu PROPERTIES EXPORT_NAME gpu)
 rocm_clang_tidy_check(migraphx_gpu)

src/targets/gpu/abs.cpp

@@ -8,7 +8,7 @@ namespace gpu {
 shape miopen_abs::compute_shape(const std::vector<shape>& inputs) const
 {
     check_shapes{inputs, *this}.has(2).not_broadcasted();
-    return inputs.at(1);
+    return inputs.at(0);
 }
 
 argument miopen_abs::compute(context& ctx,

src/targets/gpu/adjust_allocation.cpp

+#include <migraphx/gpu/adjust_allocation.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <algorithm>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+void adjust_allocation::apply(program& p) const
+{
+    for(auto ins : iterator_for(p))
+    {
+        // skip instruction with no input
+        if(ins->inputs().empty())
+            continue;
+
+        if(ins->name() == "load")
+            continue;
+
+        auto alias_ins = instruction::get_output_alias(ins, true);
+        if(alias_ins->name() == "hip::allocate")
+        {
+            // shape allocated is different from actual shape
+            // of the instruction, reallocate and replace the previous one
+            if(alias_ins->get_shape() != ins->get_shape())
+            {
+                auto alloc_ins = p.insert_instruction(ins, hip_allocate{ins->get_shape()});
+                p.replace_instruction(alias_ins, alloc_ins);
+            }
+        }
+    }
+}
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

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

+#ifndef MIGRAPHX_GUARD_RTGLIB_ADJUST_ALLOCATION_HPP
+#define MIGRAPHX_GUARD_RTGLIB_ADJUST_ALLOCATION_HPP
+
+#include <migraphx/program.hpp>
+#include <migraphx/config.hpp>
+#include <migraphx/gpu/context.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+namespace gpu {
+
+struct adjust_allocation
+{
+    std::string name() const { return "gpu::adjust_allocation"; }
+    void apply(program& p) const;
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

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

@@ -45,7 +45,7 @@ struct unary_device : oper<Derived>
     shape compute_shape(const std::vector<shape>& inputs) const
     {
         check_shapes{inputs, *this}.has(2);
-        return inputs.at(0);
+        return inputs.at(1);
     }
 
     argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
@@ -63,7 +63,7 @@ struct binary_device : oper<Derived>
     shape compute_shape(const std::vector<shape>& inputs) const
     {
         check_shapes{inputs, *this}.has(3);
-        return inputs.at(0);
+        return inputs.at(2);
     }
 
     argument compute(context& ctx, const shape&, const std::vector<argument>& args) const

src/targets/gpu/tanh.cpp

@@ -8,7 +8,7 @@ namespace gpu {
 shape miopen_tanh::compute_shape(const std::vector<shape>& inputs) const
 {
     check_shapes{inputs, *this}.has(2).not_broadcasted();
-    return inputs.at(1);
+    return inputs.at(0);
 }
 
 argument miopen_tanh::compute(context& ctx,

src/targets/gpu/target.cpp

@@ -20,6 +20,7 @@
 #include <migraphx/eliminate_identity.hpp>
 #include <migraphx/gpu/concat_gpu_opt.hpp>
 #include <migraphx/gpu/schedule_model.hpp>
+#include <migraphx/gpu/adjust_allocation.hpp>
 #include <migraphx/eliminate_pad.hpp>
 #include <migraphx/schedule.hpp>
 
@@ -57,6 +58,8 @@ std::vector<pass> target::get_passes(migraphx::context& gctx) const
         dead_code_elimination{},
         eliminate_contiguous{},
         dead_code_elimination{},
+        adjust_allocation{},
+        dead_code_elimination{},
         fuse_ops{&ctx},
         dead_code_elimination{},
         write_literals{&ctx},

test/gpu/adjust_allocation.cpp

+#include <migraphx/gpu/adjust_allocation.hpp>
+#include <migraphx/gpu/target.hpp>
+#include <migraphx/gpu/lowering.hpp>
+#include <migraphx/gpu/context.hpp>
+#include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/auto_contiguous.hpp>
+#include <migraphx/eliminate_contiguous.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/op/add.hpp>
+#include <migraphx/op/transpose.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/pass_manager.hpp>
+#include <migraphx/op/tanh.hpp>
+#include <basic_ops.hpp>
+#include <test.hpp>
+
+struct lowering_target
+{
+    std::string name() const { return "gpu::lowering"; }
+    std::vector<migraphx::pass> get_passes(migraphx::context& gctx) const
+    {
+        auto& ctx = migraphx::any_cast<migraphx::gpu::context>(gctx);
+        return {migraphx::auto_contiguous{},
+                migraphx::gpu::lowering{ctx},
+                migraphx::dead_code_elimination{},
+                migraphx::eliminate_contiguous{},
+                migraphx::dead_code_elimination{}};
+    }
+    migraphx::gpu::context get_context() const { return migraphx::gpu::context{}; }
+};
+
+TEST_CASE(tanh_shape)
+{
+    auto create_program = [] {
+        migraphx::program p;
+        migraphx::shape s{migraphx::shape::float_type, {2, 3}};
+        auto x   = p.add_parameter("x", s);
+        auto tx  = p.add_instruction(migraphx::op::transpose{{1, 0}}, x);
+        auto txh = p.add_instruction(migraphx::op::tanh{}, tx);
+        p.add_instruction(migraphx::op::add{}, txh, txh);
+
+        return p;
+    };
+
+    auto p1 = create_program();
+    auto p2 = create_program();
+    EXPECT(p1 == p2);
+
+    p1.compile(lowering_target{});
+    p2.compile(lowering_target());
+
+    EXPECT(p1 == p2);
+
+    for(auto ins : iterator_for(p1))
+    {
+        if(ins->name() == "hip::allocate")
+        {
+            migraphx::shape wrong_s{migraphx::shape::float_type, {3, 2}, {1, 3}};
+            ins->replace(wrong_s);
+        }
+    }
+    EXPECT(p1 != p2);
+
+    migraphx::run_passes(p2,
+                         {migraphx::gpu::adjust_allocation{}, migraphx::dead_code_elimination{}});
+    EXPECT(p1 == p2);
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/gpu/miopen.cpp

@@ -327,6 +327,19 @@ struct test_tanh : verify_program<test_tanh>
     }
 };
 
+struct test_trans_tanh : verify_program<test_trans_tanh>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto x  = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {4, 3, 3, 3}});
+        auto tx = p.add_instruction(migraphx::op::transpose{{0, 1, 3, 2}}, x);
+        auto tanhx = p.add_instruction(migraphx::op::tanh{}, tx);
+        p.add_instruction(migraphx::op::add{}, tanhx, tanhx);
+        return p;
+    }
+};
+
 struct test_asin : verify_program<test_asin>
 {
     migraphx::program create_program() const
@@ -674,6 +687,19 @@ struct test_abs : verify_program<test_abs>
     }
 };
 
+struct test_trans_abs : verify_program<test_trans_abs>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto x  = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {4, 3, 3, 3}});
+        auto tx = p.add_instruction(migraphx::op::transpose{{0, 1, 3, 2}}, x);
+        auto tanhx = p.add_instruction(migraphx::op::abs{}, tx);
+        p.add_instruction(migraphx::op::add{}, tanhx, tanhx);
+        return p;
+    }
+};
+
 struct test_leaky_relu : verify_program<test_leaky_relu>
 {
     migraphx::program create_program() const