tensor.hpp 7.85 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
#pragma once
Chao Liu's avatar
Chao Liu committed
2
3
4
#include <thread>
#include <vector>
#include <numeric>
Chao Liu's avatar
Chao Liu committed
5
#include <algorithm>
Chao Liu's avatar
Chao Liu committed
6
#include <utility>
Chao Liu's avatar
Chao Liu committed
7
8
#include <cassert>
#include <iostream>
Chao Liu's avatar
Chao Liu committed
9
10
#include "cuda_runtime.h"
#include "helper_cuda.h"
Chao Liu's avatar
Chao Liu committed
11

Chao Liu's avatar
Chao Liu committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
template <class Range>
std::ostream& LogRange(std::ostream& os, Range&& r, std::string delim)
{
    bool first = true;
    for(auto&& x : r)
    {
        if(first)
            first = false;
        else
            os << delim;
        os << x;
    }
    return os;
}

Chao Liu's avatar
Chao Liu committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
typedef enum
{
    Half  = 0,
    Float = 1,
} DataType_t;

template <class T>
struct DataType;

template <>
struct DataType<float> : std::integral_constant<DataType_t, DataType_t::Float>
{
};

Chao Liu's avatar
Chao Liu committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
template <class F, class T, std::size_t... Is>
auto call_f_unpack_args_impl(F f, T args, std::index_sequence<Is...>)
{
    return f(std::get<Is>(args)...);
}

template <class F, class T>
auto call_f_unpack_args(F f, T args)
{
    constexpr std::size_t N = std::tuple_size<T>::value;

    return call_f_unpack_args_impl(f, args, std::make_index_sequence<N>{});
}

template <class F, class T, std::size_t... Is>
auto construct_f_unpack_args_impl(T args, std::index_sequence<Is...>)
{
    return F(std::get<Is>(args)...);
}

template <class F, class T>
auto construct_f_unpack_args(F, T args)
{
    constexpr std::size_t N = std::tuple_size<T>::value;

    return construct_f_unpack_args_impl<F>(args, std::make_index_sequence<N>{});
}

Chao Liu's avatar
Chao Liu committed
69
70
71
struct TensorDescriptor
{
    TensorDescriptor() = delete;
Chao Liu's avatar
Chao Liu committed
72
73
    TensorDescriptor(std::initializer_list<std::size_t> lens);
    TensorDescriptor(std::initializer_list<std::size_t> lens,
Chao Liu's avatar
Chao Liu committed
74
                     std::initializer_list<std::size_t> strides);
Chao Liu's avatar
Chao Liu committed
75
    TensorDescriptor(std::vector<std::size_t> lens, std::vector<std::size_t> strides);
Chao Liu's avatar
Chao Liu committed
76
77
78
79

    void CalculateStrides();

    template <class Range>
Chao Liu's avatar
Chao Liu committed
80
    TensorDescriptor(const Range& lens) : mLens(lens.begin(), lens.end())
Chao Liu's avatar
Chao Liu committed
81
82
83
84
    {
        this->CalculateStrides();
    }

Chao Liu's avatar
Chao Liu committed
85
    template <class Range1, class Range2>
Chao Liu's avatar
Chao Liu committed
86
87
    TensorDescriptor(const Range1& lens, const Range2& strides)
        : mLens(lens.begin(), lens.end()), mStrides(strides.begin(), strides.end())
Chao Liu's avatar
Chao Liu committed
88
89
    {
    }
Chao Liu's avatar
Chao Liu committed
90
91
92
93
94

    std::size_t GetDimension() const;
    std::size_t GetElementSize() const;
    std::size_t GetElementSpace() const;

Chao Liu's avatar
Chao Liu committed
95
96
97
    const std::vector<std::size_t>& GetLengths() const;
    const std::vector<std::size_t>& GetStrides() const;

Chao Liu's avatar
Chao Liu committed
98
99
    template <class... Is>
    std::size_t Get1dIndex(Is... is) const
Chao Liu's avatar
Chao Liu committed
100
    {
Chao Liu's avatar
Chao Liu committed
101
102
103
        assert(sizeof...(Is) == this->GetDimension());
        std::initializer_list<std::size_t> iss{static_cast<std::size_t>(is)...};
        return std::inner_product(iss.begin(), iss.end(), mStrides.begin(), std::size_t{0});
Chao Liu's avatar
Chao Liu committed
104
105
106
107
108
109
110
    }

    private:
    std::vector<std::size_t> mLens;
    std::vector<std::size_t> mStrides;
};

Chao Liu's avatar
Chao Liu committed
111
struct DeviceMem
Chao Liu's avatar
Chao Liu committed
112
{
Chao Liu's avatar
Chao Liu committed
113
114
    DeviceMem() = delete;
    DeviceMem(std::size_t mem_size) : mMemSize(mem_size)
Chao Liu's avatar
Chao Liu committed
115
    {
Chao Liu's avatar
Chao Liu committed
116
        cudaMalloc(static_cast<void**>(&mpDeviceBuf), mMemSize);
Chao Liu's avatar
Chao Liu committed
117
118
    }

Chao Liu's avatar
Chao Liu committed
119
120
121
    void* GetDeviceBuffer() { return mpDeviceBuf; }

    int ToDevice(const void* p)
Chao Liu's avatar
Chao Liu committed
122
    {
Chao Liu's avatar
Chao Liu committed
123
124
        return static_cast<int>(
            cudaMemcpy(mpDeviceBuf, const_cast<void*>(p), mMemSize, cudaMemcpyHostToDevice));
Chao Liu's avatar
Chao Liu committed
125
126
    }

Chao Liu's avatar
Chao Liu committed
127
    int FromDevice(void* p)
Chao Liu's avatar
Chao Liu committed
128
    {
Chao Liu's avatar
Chao Liu committed
129
        return static_cast<int>(cudaMemcpy(p, mpDeviceBuf, mMemSize, cudaMemcpyDeviceToHost));
Chao Liu's avatar
Chao Liu committed
130
    }
Chao Liu's avatar
Chao Liu committed
131

Chao Liu's avatar
Chao Liu committed
132
    ~DeviceMem() { cudaFree(mpDeviceBuf); }
Chao Liu's avatar
Chao Liu committed
133

Chao Liu's avatar
Chao Liu committed
134
135
    void* mpDeviceBuf;
    std::size_t mMemSize;
Chao Liu's avatar
Chao Liu committed
136
137
};

Chao Liu's avatar
Chao Liu committed
138
struct joinable_thread : std::thread
Chao Liu's avatar
Chao Liu committed
139
{
Chao Liu's avatar
Chao Liu committed
140
141
142
143
    template <class... Xs>
    joinable_thread(Xs&&... xs) : std::thread(std::forward<Xs>(xs)...)
    {
    }
Chao Liu's avatar
Chao Liu committed
144

Chao Liu's avatar
Chao Liu committed
145
146
    joinable_thread(joinable_thread&&) = default;
    joinable_thread& operator=(joinable_thread&&) = default;
Chao Liu's avatar
Chao Liu committed
147

Chao Liu's avatar
Chao Liu committed
148
149
150
151
152
153
    ~joinable_thread()
    {
        if(this->joinable())
            this->join();
    }
};
Chao Liu's avatar
Chao Liu committed
154
155
156
157
158

template <class F, class... Xs>
struct ParallelTensorFunctor
{
    F mF;
Chao Liu's avatar
Chao Liu committed
159
    static constexpr std::size_t NDIM = sizeof...(Xs);
Chao Liu's avatar
Chao Liu committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
    std::array<std::size_t, NDIM> mLens;
    std::array<std::size_t, NDIM> mStrides;
    std::size_t mN1d;

    ParallelTensorFunctor(F f, Xs... xs) : mF(f), mLens({static_cast<std::size_t>(xs)...})
    {
        mStrides.back() = 1;
        std::partial_sum(mLens.rbegin(),
                         mLens.rend() - 1,
                         mStrides.rbegin() + 1,
                         std::multiplies<std::size_t>());
        mN1d = mStrides[0] * mLens[0];
    }

Chao Liu's avatar
Chao Liu committed
174
175
176
177
178
179
180
181
182
183
184
185
186
    std::array<std::size_t, NDIM> GetNdIndices(std::size_t i) const
    {
        std::array<std::size_t, NDIM> indices;

        for(int idim = 0; idim < NDIM; ++idim)
        {
            indices[idim] = i / mStrides[idim];
            i -= indices[idim] * mStrides[idim];
        }

        return indices;
    }

Chao Liu's avatar
Chao Liu committed
187
    void operator()(std::size_t num_thread) const
Chao Liu's avatar
Chao Liu committed
188
189
190
191
192
193
194
195
    {
        std::size_t work_per_thread = (mN1d + num_thread - 1) / num_thread;

        std::vector<joinable_thread> threads(num_thread);

        for(std::size_t it = 0; it < num_thread; ++it)
        {
            std::size_t iw_begin = it * work_per_thread;
Chao Liu's avatar
Chao Liu committed
196
            std::size_t iw_end   = std::min((it + 1) * work_per_thread, mN1d);
Chao Liu's avatar
Chao Liu committed
197
198
199

            auto f = [=] {
                for(std::size_t iw = iw_begin; iw < iw_end; ++iw)
Chao Liu's avatar
Chao Liu committed
200
201
202
                {
                    call_f_unpack_args(mF, GetNdIndices(iw));
                }
Chao Liu's avatar
Chao Liu committed
203
204
205
206
207
208
            };
            threads[it] = joinable_thread(f);
        }
    }
};

Chao Liu's avatar
Chao Liu committed
209
210
template <class F, class... Xs>
auto make_ParallelTensorFunctor(F f, Xs... xs)
Chao Liu's avatar
Chao Liu committed
211
{
Chao Liu's avatar
Chao Liu committed
212
    return ParallelTensorFunctor<F, Xs...>(f, xs...);
Chao Liu's avatar
Chao Liu committed
213
214
}

Chao Liu's avatar
Chao Liu committed
215
216
template <class T>
struct Tensor
Chao Liu's avatar
Chao Liu committed
217
{
Chao Liu's avatar
Chao Liu committed
218
    template <class X>
Chao Liu's avatar
Chao Liu committed
219
    Tensor(std::initializer_list<X> lens) : mDesc(lens), mData(mDesc.GetElementSpace())
Chao Liu's avatar
Chao Liu committed
220
221
    {
    }
Chao Liu's avatar
Chao Liu committed
222

Chao Liu's avatar
Chao Liu committed
223
    template <class X>
Chao Liu's avatar
Chao Liu committed
224
    Tensor(std::vector<X> lens) : mDesc(lens), mData(mDesc.GetElementSpace())
Chao Liu's avatar
Chao Liu committed
225
226
    {
    }
Chao Liu's avatar
Chao Liu committed
227

Chao Liu's avatar
Chao Liu committed
228
229
    template <class X, class Y>
    Tensor(std::vector<X> lens, std::vector<Y> strides)
Chao Liu's avatar
Chao Liu committed
230
        : mDesc(lens, strides), mData(mDesc.GetElementSpace())
Chao Liu's avatar
Chao Liu committed
231
232
    {
    }
Chao Liu's avatar
Chao Liu committed
233

Chao Liu's avatar
Chao Liu committed
234
235
    Tensor(const TensorDescriptor& desc) : mDesc(desc), mData(mDesc.GetElementSpace()) {}

Chao Liu's avatar
Chao Liu committed
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
    template <class G>
    void GenerateTensorValue(G g, std::size_t num_thread = 1)
    {
        switch(mDesc.GetDimension())
        {
        case 1:
        {
            auto f = [&](auto i) { (*this)(i) = g(i); };
            make_ParallelTensorFunctor(f, mDesc.GetLengths()[0])(num_thread);
            break;
        }
        case 2:
        {
            auto f = [&](auto i0, auto i1) { (*this)(i0, i1) = g(i0, i1); };
            make_ParallelTensorFunctor(f, mDesc.GetLengths()[0], mDesc.GetLengths()[1])(num_thread);
            break;
        }
        case 3:
        {
            auto f = [&](auto i0, auto i1, auto i2) { (*this)(i0, i1, i2) = g(i0, i1, i2); };
            make_ParallelTensorFunctor(
                f, mDesc.GetLengths()[0], mDesc.GetLengths()[1], mDesc.GetLengths()[2])(num_thread);
            break;
        }
        case 4:
        {
            auto f = [&](auto i0, auto i1, auto i2, auto i3) {
                (*this)(i0, i1, i2, i3) = g(i0, i1, i2, i3);
            };
            make_ParallelTensorFunctor(f,
                                       mDesc.GetLengths()[0],
                                       mDesc.GetLengths()[1],
                                       mDesc.GetLengths()[2],
                                       mDesc.GetLengths()[3])(num_thread);
            break;
        }
        default: throw std::runtime_error("unspported dimension");
        }
    }

    template <class... Is>
    T& operator()(Is... is)
    {
        return mData[mDesc.Get1dIndex(is...)];
    }

    template <class... Is>
    const T& operator()(Is... is) const
    {
        return mData[mDesc.Get1dIndex(is...)];
    }

    typename std::vector<T>::iterator begin() { return mData.begin(); }

    typename std::vector<T>::iterator end() { return mData.end(); }

    typename std::vector<T>::const_iterator begin() const { return mData.begin(); }

    typename std::vector<T>::const_iterator end() const { return mData.end(); }

    TensorDescriptor mDesc;
    std::vector<T> mData;
};