shape.cpp 11.7 KB
Newer Older
Paul's avatar
Paul committed
1

Paul's avatar
Paul committed
2
3
#include <migraphx/shape.hpp>
#include <migraphx/stringutils.hpp>
4
#include <migraphx/serialize.hpp>
5
#include <migraphx/permutation.hpp>
Paul's avatar
Paul committed
6
7
8
#include <numeric>
#include <algorithm>
#include <functional>
9
#include <unordered_map>
Paul's avatar
Paul committed
10
#include <iostream>
Paul's avatar
Paul committed
11

Paul's avatar
Paul committed
12
namespace migraphx {
Paul's avatar
Paul committed
13
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
14

Paul's avatar
Paul committed
15
16
struct shape_impl
{
Paul's avatar
Paul committed
17
18
    static std::shared_ptr<shape_impl> default_shape()
    {
19
        static const std::shared_ptr<shape_impl> result = std::make_shared<shape_impl>();
Paul's avatar
Paul committed
20
21
22
        return result;
    }

Paul Fultz II's avatar
Paul Fultz II committed
23
    shape_impl() : m_type(shape::float_type) {}
Paul's avatar
Paul committed
24

Paul Fultz II's avatar
Paul Fultz II committed
25
26
27
28
    shape_impl(shape::type_t t) : m_type(t), m_lens({1}), m_strides({0}), m_standard(true)
    {
        assert(t != shape::tuple_type);
    }
Paul's avatar
Paul committed
29
30
31
    shape_impl(shape::type_t t, std::vector<std::size_t> l)
        : m_type(t), m_lens(std::move(l)), m_standard(true)
    {
Paul Fultz II's avatar
Paul Fultz II committed
32
        assert(t != shape::tuple_type);
Paul's avatar
Paul committed
33
34
35
36
37
38
        this->calculate_strides();
        assert(m_lens.size() == m_strides.size());
    }
    shape_impl(shape::type_t t, std::vector<std::size_t> l, std::vector<std::size_t> s)
        : m_type(t), m_lens(std::move(l)), m_strides(std::move(s))
    {
Paul Fultz II's avatar
Paul Fultz II committed
39
        assert(t != shape::tuple_type);
Paul's avatar
Paul committed
40
        assert(m_lens.size() == m_strides.size());
Khalique's avatar
Khalique committed
41
42
        // assert(std::any_of(m_strides.begin(), m_strides.end(), [](auto x) { return x > 0; }) and
        //        "At least one stride must be non-zero");
43
44
        m_standard = this->elements() == this->element_space() and
                     std::is_sorted(m_strides.rbegin(), m_strides.rend());
Paul's avatar
Paul committed
45
    }
Paul Fultz II's avatar
Paul Fultz II committed
46
47

    shape_impl(const std::vector<shape>& subs) : m_type(shape::tuple_type), m_shapes(subs) {}
Paul's avatar
Paul committed
48
    shape::type_t m_type;
Shucai Xiao's avatar
Shucai Xiao committed
49
50
51
52
    std::vector<std::size_t> m_lens                    = {};
    std::vector<std::size_t> m_strides                 = {};
    std::vector<shape> m_shapes                        = {};
    bool m_standard                                    = false;
53
    std::vector<shape::dynamic_dimension> dynamic_dims = {};
Paul's avatar
Paul committed
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

    void calculate_strides()
    {
        m_strides.clear();
        m_strides.resize(m_lens.size(), 0);
        if(m_strides.empty())
            return;
        m_strides.back() = 1;
        std::partial_sum(m_lens.rbegin(),
                         m_lens.rend() - 1,
                         m_strides.rbegin() + 1,
                         std::multiplies<std::size_t>());
    }

    std::size_t element_space() const
    {
        assert(m_lens.size() == m_strides.size());
        if(m_lens.empty())
            return 0;
        return std::inner_product(m_lens.begin(),
                                  m_lens.end(),
                                  m_strides.begin(),
                                  std::size_t{0},
                                  std::plus<std::size_t>{},
                                  [](std::size_t l, std::size_t s) { return (l - 1) * s; }) +
               1;
    }

    std::size_t elements() const
    {
        assert(m_lens.size() == m_strides.size());
        if(m_lens.empty())
            return 0;
        return std::accumulate(
            m_lens.begin(), m_lens.end(), std::size_t{1}, std::multiplies<std::size_t>());
    }
90
91

    std::shared_ptr<shape_impl> copy() const { return std::make_shared<shape_impl>(*this); }
Paul's avatar
Paul committed
92
};
Paul's avatar
Paul committed
93

94
95
96
97
const std::vector<shape::type_t>& shape::types()
{
    static const std::vector<shape::type_t> result = {
#define MIGRAPHX_GENERATE_TYPE_VECTOR(x, t) x,
Paul Fultz II's avatar
Paul Fultz II committed
98
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_GENERATE_TYPE_VECTOR) tuple_type};
99
100
101
    return result;
}

102
103
104
105
std::string shape::name(shape::type_t t)
{
    switch(t)
    {
Paul Fultz II's avatar
Paul Fultz II committed
106
    case tuple_type: return "tuple_type";
107
108
109
110
111
112
113
114
115
116
117
#define MIGRAPHX_SHAPE_GENERATE_TYPE_NAME_CASE(x, t) \
    case x: return #x;
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_TYPE_NAME_CASE)
#undef MIGRAPHX_SHAPE_GENERATE_TYPE_NAME_CASE
    }
    MIGRAPHX_THROW("Invalid type");
}
std::string shape::cpp_type(shape::type_t t)
{
    switch(t)
    {
Paul Fultz II's avatar
Paul Fultz II committed
118
    case tuple_type: MIGRAPHX_THROW("No C++ type for tuple");
119
120
121
122
123
124
125
126
#define MIGRAPHX_SHAPE_GENERATE_CPP_TYPE_CASE(x, t) \
    case x: return #t;
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_CPP_TYPE_CASE)
#undef MIGRAPHX_SHAPE_GENERATE_CPP_TYPE_CASE
    }
    MIGRAPHX_THROW("Invalid type");
}

Paul's avatar
Paul committed
127
128
129
shape::shape() : impl(shape_impl::default_shape()) {}

shape::shape(type_t t) : impl(std::make_shared<shape_impl>(t)) {}
Paul's avatar
Paul committed
130
shape::shape(type_t t, std::vector<std::size_t> l)
Paul's avatar
Paul committed
131
    : impl(std::make_shared<shape_impl>(t, std::move(l)))
Paul's avatar
Paul committed
132
133
134
{
}
shape::shape(type_t t, std::vector<std::size_t> l, std::vector<std::size_t> s)
Paul's avatar
Paul committed
135
    : impl(std::make_shared<shape_impl>(t, std::move(l), std::move(s)))
Paul's avatar
Paul committed
136
137
138
{
}

Paul Fultz II's avatar
Paul Fultz II committed
139
140
shape::shape(const std::vector<shape>& subs) : impl(std::make_shared<shape_impl>(subs)) {}

141
142
shape::shape(std::shared_ptr<shape_impl> pimpl) : impl(std::move(pimpl)) {}

143
144
145
146
147
148
149
150
151
152
shape shape::from_permutation(type_t t,
                              const std::vector<std::size_t>& l,
                              const std::vector<int64_t>& perm)
{
    auto new_lens = reorder_dims(l, perm);
    shape result  = reorder_shape({t, new_lens}, invert_permutation(perm));
    assert(result.lens() == l);
    return result;
}

Paul's avatar
Paul committed
153
154
155
shape::type_t shape::type() const { return impl->m_type; }
const std::vector<std::size_t>& shape::lens() const { return impl->m_lens; }
const std::vector<std::size_t>& shape::strides() const { return impl->m_strides; }
Paul's avatar
Paul committed
156
std::size_t shape::elements() const { return impl->elements(); }
Paul's avatar
Paul committed
157
158
std::size_t shape::bytes() const
{
Paul Fultz II's avatar
Paul Fultz II committed
159
160
161
162
163
164
165
166
167
168
169
170
171
    if(this->sub_shapes().empty())
    {
        std::size_t n = 0;
        this->visit_type([&](auto as) { n = as.size(); });
        return n * this->element_space();
    }
    else
    {
        return std::accumulate(this->sub_shapes().begin(),
                               this->sub_shapes().end(),
                               std::size_t{0},
                               [&](auto x, auto y) { return x + y.bytes(); });
    }
Paul's avatar
Paul committed
172
}
Scott Thornton's avatar
Scott Thornton committed
173
174
175
std::size_t shape::type_size() const
{
    std::size_t n = 0;
Paul Fultz II's avatar
Paul Fultz II committed
176
177
    if(this->sub_shapes().empty())
        this->visit_type([&](auto as) { n = as.size(); });
Scott Thornton's avatar
Scott Thornton committed
178
179
    return n;
}
Paul's avatar
Paul committed
180
181
182
183
184
185
186
187
188
189
190
191
std::size_t shape::index(std::initializer_list<std::size_t> l) const
{
    assert(l.size() <= this->lens().size());
    assert(this->lens().size() == this->strides().size());
    return std::inner_product(l.begin(), l.end(), this->strides().begin(), std::size_t{0});
}
std::size_t shape::index(const std::vector<std::size_t>& l) const
{
    assert(l.size() <= this->lens().size());
    assert(this->lens().size() == this->strides().size());
    return std::inner_product(l.begin(), l.end(), this->strides().begin(), std::size_t{0});
}
Paul's avatar
Paul committed
192
193
194
std::size_t shape::index(std::size_t i) const
{
    assert(this->lens().size() == this->strides().size());
Paul's avatar
Paul committed
195
    if(this->standard())
Paul's avatar
Paul committed
196
197
        return i;
    else
Paul's avatar
Paul committed
198
    {
Paul's avatar
Paul committed
199
        std::size_t s      = 1;
Paul's avatar
Paul committed
200
        std::size_t result = 0;
Paul's avatar
Paul committed
201
        for(std::size_t j = 0; j < this->lens().size(); j++)
Paul's avatar
Paul committed
202
        {
Paul's avatar
Paul committed
203
            const std::size_t k      = this->lens().size() - j - 1;
Paul's avatar
Paul committed
204
            const std::size_t stride = this->strides()[k];
Paul's avatar
Paul committed
205
206
            const std::size_t len    = this->lens()[k];
            const std::size_t idx    = (i % (s * len)) / s;
Paul's avatar
Paul committed
207
208
209
210
211
            result += stride * idx;
            s *= len;
        }
        return result;
    }
Paul's avatar
Paul committed
212
}
213
214
215
216
217
218

std::vector<std::size_t> shape::multi(std::size_t i) const
{
    assert(this->standard());

    std::vector<std::size_t> indices(lens().size());
219
220
221
222
223
224
225
226
227
228
    multi_copy(i, indices.data(), indices.data() + lens().size());

    return indices;
}

void shape::multi_copy(std::size_t i, std::size_t* start, const std::size_t* end) const
{
    assert(this->standard());
    (void)end;
    assert(lens().size() <= (end - start));
229
    std::transform(strides().begin(),
Shucai Xiao's avatar
Shucai Xiao committed
230
231
                   strides().end(),
                   lens().begin(),
232
                   start,
Shucai Xiao's avatar
Shucai Xiao committed
233
234
235
236
                   [&](std::size_t stride, std::size_t len) {
                       assert(len > 0 and stride > 0);
                       return (i / stride) % len;
                   });
237
238
}

Paul Fultz II's avatar
Paul Fultz II committed
239
240
241
242
bool shape::packed() const
{
    return this->sub_shapes().empty() and this->elements() == this->element_space();
}
Paul's avatar
Paul committed
243

Paul's avatar
Paul committed
244
245
bool shape::transposed() const
{
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
    if(this->broadcasted())
    {
        // TODO: Use a filter_iterator instead
        std::vector<std::size_t> s;
        s.reserve(this->strides().size());
        std::copy_if(this->strides().begin(),
                     this->strides().end(),
                     std::back_inserter(s),
                     [](std::size_t x) { return x != 0; });
        return not std::is_sorted(s.rbegin(), s.rend());
    }
    else
    {
        return not std::is_sorted(this->strides().rbegin(), this->strides().rend());
    }
Paul's avatar
Paul committed
261
}
Paul's avatar
Paul committed
262
263
264
265

bool shape::broadcasted() const
{
    assert(this->lens().size() == this->strides().size());
Paul's avatar
Paul committed
266
267
268
269
    return std::accumulate(this->strides().begin(),
                           this->strides().end(),
                           std::size_t{1},
                           std::multiplies<std::size_t>()) == 0;
Paul's avatar
Paul committed
270
271
}

Khalique's avatar
Khalique committed
272
273
274
275
bool shape::scalar() const
{
    assert(this->lens().size() == this->strides().size());
    // if any stride > 0, then accumulate will return false
Paul Fultz II's avatar
Paul Fultz II committed
276
277
    return this->sub_shapes().empty() and
           std::accumulate(this->strides().begin(), this->strides().end(), std::size_t(0)) == 0;
Khalique's avatar
Khalique committed
278
279
}

Paul's avatar
Paul committed
280
bool shape::standard() const { return impl->m_standard; }
Paul's avatar
Paul committed
281

282
283
284
285
286
287
288
289
shape shape::normalize_standard() const
{
    if(this->standard())
        return {this->type(), this->lens()};
    else
        return *this;
}

290
291
292
293
294
295
296
297
298
299
300
301
shape shape::with_lens(type_t t, const std::vector<std::size_t>& l) const
{
    assert(l.size() == this->lens().size());
    auto perm = find_permutation(*this);
    return shape::from_permutation(t, l, perm);
}

shape shape::with_lens(const std::vector<std::size_t>& l) const
{
    return this->with_lens(this->type(), l);
}

302
303
304
305
306
307
308
shape shape::with_type(type_t t) const
{
    auto c    = impl->copy();
    c->m_type = t;
    return {c};
}

Paul's avatar
Paul committed
309
std::size_t shape::element_space() const { return impl->element_space(); }
Paul's avatar
Paul committed
310

311
std::string shape::type_string() const { return name(this->type()); }
Paul's avatar
Paul committed
312

Paul's avatar
Paul committed
313
314
bool operator==(const shape& x, const shape& y)
{
Paul Fultz II's avatar
Paul Fultz II committed
315
316
    return x.impl == y.impl or (x.type() == y.type() and x.lens() == y.lens() and
                                x.strides() == y.strides() and x.sub_shapes() == y.sub_shapes());
Paul's avatar
Paul committed
317
}
Paul's avatar
Paul committed
318
bool operator!=(const shape& x, const shape& y) { return !(x == y); }
Paul's avatar
Paul committed
319

Paul's avatar
Paul committed
320
321
std::ostream& operator<<(std::ostream& os, const shape& x)
{
Paul Fultz II's avatar
Paul Fultz II committed
322
323
324
325
326
327
328
329
330
331
    if(x.sub_shapes().empty())
    {
        os << x.type_string() << ", ";
        os << "{" << to_string_range(x.lens()) << "}, ";
        os << "{" << to_string_range(x.strides()) << "}";
    }
    else
    {
        os << "[" << to_string_range(x.sub_shapes()) << "]";
    }
Paul's avatar
Paul committed
332
333
334
    return os;
}

335
336
shape::type_t shape::parse_type(const std::string& s)
{
337
    static const std::unordered_map<std::string, shape::type_t> m = {
338
#define MIGRAPHX_SHAPE_GENERATE_TYPE_STRING_MAP(x, t) {#x, x}, {#t, x},
Paul Fultz II's avatar
Paul Fultz II committed
339
340
341
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_TYPE_STRING_MAP){"tuple_type",
                                                                            tuple_type},
        {"tuple", tuple_type}};
342
343
344
    return m.at(s);
}

Paul Fultz II's avatar
Paul Fultz II committed
345
346
const std::vector<shape>& shape::sub_shapes() const { return impl->m_shapes; }

347
348
349
void migraphx_to_value(value& v, const shape& s)
{
    value result;
Paul Fultz II's avatar
Paul Fultz II committed
350
351
352
353
354
    result["type"]       = migraphx::to_value(s.type_string());
    result["lens"]       = migraphx::to_value(s.lens());
    result["strides"]    = migraphx::to_value(s.strides());
    result["sub_shapes"] = migraphx::to_value(s.sub_shapes());
    v                    = result;
355
356
357
}
void migraphx_from_value(const value& v, shape& s)
{
Paul Fultz II's avatar
Paul Fultz II committed
358
359
360
361
362
363
364
365
366
367
368
    auto t = v.at("type").get_string();
    if(t == "tuple_type")
    {
        s = shape{migraphx::from_value<std::vector<migraphx::shape>>(v.at("sub_shapes"))};
    }
    else
    {
        s = shape{shape::parse_type(t),
                  v.at("lens").to_vector<std::size_t>(),
                  v.at("strides").to_vector<std::size_t>()};
    }
369
370
}

Paul's avatar
Paul committed
371
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
372
} // namespace migraphx