sync_device.cpp 866 Bytes
Newer Older
1
2
3
4
5
6
7
8
9
10
#include <migraphx/gpu/sync_device.hpp>
#include <migraphx/gpu/hip.hpp>
#include <migraphx/program.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/iterator_for.hpp>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {

11
void sync_device::apply(module& p) const
12
13
14
15
16
17
18
19
20
{
    auto last = std::prev(p.end());
    if(last->name() == "@return")
    {
        auto inputs = last->inputs();
        if(std::any_of(inputs.begin(), inputs.end(), [](auto i) {
               return (i->name() == "hip::copy_from_gpu");
           }))
        {
Shucai Xiao's avatar
Shucai Xiao committed
21
22
23
24
25
            auto sync_in = p.insert_instruction(last, hip_sync_device{}, inputs);
            if(not inputs.empty())
            {
                p.replace_instruction(inputs.front(), sync_in);
            }
26
27
28
29
30
31
32
        }
    }
}

} // namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx