add a pass to adjust gpu memory allocation

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/adjust_allocation.cpp

+#include <migraphx/gpu/adjust_allocation.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/program.hpp>
+#include <migraphx/iterator_for.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace gpu {
+
+void adjust_allocation::apply(program& p) const
+{
+    for(auto ins : iterator_for(p))
+    {
+        auto alias_ins = instruction::get_output_alias(ins, true);
+        if(alias_ins->name() == "hip::allocate")
+        {
+            std::cout << "====================" << std::endl;
+            std::cout << "ins_name = " << ins->name() << ", shape = " << ins->get_shape() << std::endl;
+            std::cout << "alias_ins_name = " << alias_ins->name() << ", shape = " << alias_ins->get_shape() << std::endl;
+            // 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
+{
+    context* ctx = nullptr;
+    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/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},