hip.cpp 3.36 KB
Newer Older
Paul's avatar
Paul committed
1

Paul's avatar
Paul committed
2
#include <migraphx/gpu/hip.hpp>
Paul's avatar
Paul committed
3

Paul's avatar
Paul committed
4
#include <migraphx/manage_ptr.hpp>
Paul's avatar
Paul committed
5
6
7
8
#include <miopen/miopen.h>

#include <vector>

Paul's avatar
Paul committed
9
namespace migraphx {
Paul's avatar
Paul committed
10
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
11
namespace gpu {
Paul's avatar
Paul committed
12

Paul's avatar
Paul committed
13
using hip_ptr = MIGRAPHX_MANAGE_PTR(void, hipFree);
Paul's avatar
Paul committed
14

Paul's avatar
Paul committed
15
std::string hip_error(int error) { return hipGetErrorString(static_cast<hipError_t>(error)); }
Paul's avatar
Paul committed
16

Paul's avatar
Paul committed
17
std::size_t get_available_gpu_memory()
Paul's avatar
Paul committed
18
{
Paul's avatar
Paul committed
19
20
    size_t free;
    size_t total;
Paul's avatar
Paul committed
21
    auto status = hipMemGetInfo(&free, &total);
Paul's avatar
Paul committed
22
    if(status != hipSuccess)
Paul's avatar
Paul committed
23
        MIGRAPHX_THROW("Failed getting available memory: " + hip_error(status));
Paul's avatar
Paul committed
24
25
26
    return free;
}

Paul's avatar
Paul committed
27
hip_ptr allocate_gpu(std::size_t sz, bool host = false)
Paul's avatar
Paul committed
28
29
{
    if(sz > get_available_gpu_memory())
Paul's avatar
Paul committed
30
        MIGRAPHX_THROW("Memory not available to allocate buffer: " + std::to_string(sz));
Paul's avatar
Paul committed
31
32
    void* result;
    auto status = host ? hipHostMalloc(&result, sz) : hipMalloc(&result, sz);
Paul's avatar
Paul committed
33
34
35
    if(status != hipSuccess)
    {
        if(host)
Paul's avatar
Paul committed
36
            MIGRAPHX_THROW("Gpu allocation failed: " + hip_error(status));
Paul's avatar
Paul committed
37
38
39
        else
            allocate_gpu(sz, true);
    }
Paul's avatar
Paul committed
40
41
42
    return hip_ptr{result};
}

Paul's avatar
Paul committed
43
44
45
template <class T>
std::vector<T> read_from_gpu(const void* x, std::size_t sz)
{
46
    gpu_sync();
Paul's avatar
Paul committed
47
    std::vector<T> result(sz);
Paul's avatar
Paul committed
48
49
    auto status = hipMemcpy(result.data(), x, sz * sizeof(T), hipMemcpyDeviceToHost);
    if(status != hipSuccess)
Paul's avatar
Paul committed
50
        MIGRAPHX_THROW("Copy from gpu failed: " + hip_error(status)); // NOLINT
Paul's avatar
Paul committed
51
52
53
    return result;
}

Paul's avatar
Paul committed
54
hip_ptr write_to_gpu(const void* x, std::size_t sz, bool host = false)
Paul's avatar
Paul committed
55
{
56
    gpu_sync();
Paul's avatar
Paul committed
57
    auto result = allocate_gpu(sz, host);
Paul's avatar
Paul committed
58
59
    auto status = hipMemcpy(result.get(), x, sz, hipMemcpyHostToDevice);
    if(status != hipSuccess)
Paul's avatar
Paul committed
60
        MIGRAPHX_THROW("Copy to gpu failed: " + hip_error(status));
Paul's avatar
Paul committed
61
62
63
    return result;
}

Paul's avatar
Paul committed
64
65
66
67
68
69
70
71
template <class T>
hip_ptr write_to_gpu(const T& x)
{
    using type = typename T::value_type;
    auto size  = x.size() * sizeof(type);
    return write_to_gpu(x.data(), size);
}

Paul's avatar
Paul committed
72
argument allocate_gpu(const shape& s, bool host)
Paul's avatar
Paul committed
73
{
Paul's avatar
Paul committed
74
    auto p = share(allocate_gpu(s.bytes() + 1, host));
Paul's avatar
Paul committed
75
76
77
    return {s, [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
}

Paul's avatar
Paul committed
78
argument to_gpu(const argument& arg, bool host)
Paul's avatar
Paul committed
79
{
Paul's avatar
Paul committed
80
    auto p = share(write_to_gpu(arg.data(), arg.get_shape().bytes(), host));
Paul's avatar
Paul committed
81
82
83
    return {arg.get_shape(), [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
}

Paul's avatar
Paul committed
84
argument from_gpu(const argument& arg)
Paul's avatar
Paul committed
85
{
Paul's avatar
Paul committed
86
    argument result;
Paul's avatar
Paul committed
87
88
89
90
91
92
93
94
    arg.visit([&](auto x) {
        using type = typename decltype(x)::value_type;
        auto v     = read_from_gpu<type>(arg.data(), x.get_shape().bytes() / sizeof(type));
        result     = {x.get_shape(), [v]() mutable { return reinterpret_cast<char*>(v.data()); }};
    });
    return result;
}

Paul's avatar
Paul committed
95
96
97
98
void set_device(std::size_t id)
{
    auto status = hipSetDevice(id);
    if(status != hipSuccess)
Paul's avatar
Paul committed
99
        MIGRAPHX_THROW("Error setting device");
Paul's avatar
Paul committed
100
101
}

Paul's avatar
Paul committed
102
103
void gpu_sync() { hipDeviceSynchronize(); }

Paul's avatar
Paul committed
104
void copy_to_gpu(const argument& src, const argument& dst)
mei-ye's avatar
mei-ye committed
105
{
106
107
108
    std::size_t src_size = src.get_shape().bytes();
    std::size_t dst_size = dst.get_shape().bytes();
    if(src_size > dst_size)
Paul's avatar
Paul committed
109
        MIGRAPHX_THROW("Not enough memory available in destination to do copy");
110
111
    auto status = hipMemcpy(dst.data(), src.data(), src_size, hipMemcpyHostToDevice);
    if(status != hipSuccess)
Paul's avatar
Paul committed
112
        MIGRAPHX_THROW("Copy to gpu failed: " + hip_error(status));
mei-ye's avatar
mei-ye committed
113
}
114

Paul's avatar
Paul committed
115
} // namespace gpu
Paul's avatar
Paul committed
116
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
117
} // namespace migraphx