shape.hpp 13.8 KB
Newer Older
1
2
3
/*
 * The MIT License (MIT)
 *
4
 * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
Paul's avatar
Paul committed
24
25
#ifndef MIGRAPHX_GUARD_MIGRAPHLIB_SHAPE_HPP
#define MIGRAPHX_GUARD_MIGRAPHLIB_SHAPE_HPP
Paul's avatar
Paul committed
26
27
28

#include <vector>
#include <cassert>
Paul's avatar
Paul committed
29
#include <ostream>
Paul's avatar
Paul committed
30
#include <numeric>
Paul's avatar
Paul committed
31
#include <memory>
32
#include <set>
Paul's avatar
Paul committed
33

34
#include <migraphx/functional.hpp>
Paul's avatar
Paul committed
35
36
#include <migraphx/errors.hpp>
#include <migraphx/half.hpp>
37
#include <migraphx/serialize.hpp>
Paul's avatar
Paul committed
38
#include <migraphx/config.hpp>
Paul's avatar
Paul committed
39

Paul's avatar
Paul committed
40
namespace migraphx {
Paul's avatar
Paul committed
41
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
42

43
struct value;
Paul's avatar
Paul committed
44
45
struct shape_impl;

46
struct MIGRAPHX_EXPORT shape
Paul's avatar
Paul committed
47
{
Paul's avatar
Paul committed
48
49

// Add new types here
Paul's avatar
Paul committed
50
// clang-format off
Paul's avatar
Paul committed
51
#define MIGRAPHX_SHAPE_VISIT_TYPES(m) \
52
    m(bool_type, bool) \
Paul's avatar
Paul committed
53
    m(half_type, half) \
Paul's avatar
Paul committed
54
    m(float_type, float) \
Paul's avatar
Paul committed
55
56
57
58
59
60
61
62
    m(double_type, double) \
    m(uint8_type, uint8_t) \
    m(int8_type, int8_t) \
    m(uint16_type, uint16_t) \
    m(int16_type, int16_t) \
    m(int32_type, int32_t) \
    m(int64_type, int64_t) \
    m(uint32_type, uint32_t) \
Paul's avatar
Paul committed
63
    m(uint64_type, uint64_t)
bpickrel's avatar
bpickrel committed
64
    // clang-format on
Paul's avatar
Paul committed
65

Paul's avatar
Paul committed
66
#define MIGRAPHX_SHAPE_GENERATE_ENUM_TYPES(x, t) x,
Paul's avatar
Paul committed
67
68
    enum type_t
    {
Paul Fultz II's avatar
Paul Fultz II committed
69
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_ENUM_TYPES) tuple_type
Paul's avatar
Paul committed
70
    };
Paul's avatar
Paul committed
71
#undef MIGRAPHX_SHAPE_GENERATE_ENUM_TYPES
Paul's avatar
Paul committed
72

Paul's avatar
Paul committed
73
    template <class T, class = void>
Paul's avatar
Paul committed
74
    struct get_type;
Paul's avatar
Paul committed
75
#define MIGRAPHX_SHAPE_GENERATE_GET_TYPE(x, t)                \
Paul's avatar
Paul committed
76
    template <class T>                                        \
Paul's avatar
Paul committed
77
    struct get_type<t, T> : std::integral_constant<type_t, x> \
Paul's avatar
Paul committed
78
79
    {                                                         \
    };
Paul's avatar
Paul committed
80
81
    MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_GET_TYPE)
#undef MIGRAPHX_SHAPE_GENERATE_GET_TYPE
Paul's avatar
Paul committed
82

wsttiger's avatar
wsttiger committed
83
84
85
86
87
    template <class T>
    struct get_type<const T> : get_type<T>
    {
    };

88
    struct MIGRAPHX_EXPORT dynamic_dimension
Charlie Lin's avatar
Charlie Lin committed
89
90
91
    {
        std::size_t min = 0;
        std::size_t max = 0;
92
        std::set<std::size_t> optimals{};
Charlie Lin's avatar
Charlie Lin committed
93
94

        template <class Self, class F>
95
96
        static auto reflect(Self& self, F f)
        {
97
            return pack(f(self.min, "min"), f(self.max, "max"), f(self.optimals, "optimals"));
98
        }
Charlie Lin's avatar
Charlie Lin committed
99
100
101
102

        bool is_fixed() const;
        bool has_optimal() const;

103
104
105
106
107
108
        MIGRAPHX_EXPORT friend bool operator==(const dynamic_dimension& x,
                                               const dynamic_dimension& y);
        MIGRAPHX_EXPORT friend bool operator!=(const dynamic_dimension& x,
                                               const dynamic_dimension& y);
        MIGRAPHX_EXPORT friend std::ostream& operator<<(std::ostream& os,
                                                        const dynamic_dimension& x);
109
110

        // compare to fixed std::size_t dimension
111
112
113
114
        MIGRAPHX_EXPORT friend bool operator==(const dynamic_dimension& x, const std::size_t& y);
        MIGRAPHX_EXPORT friend bool operator==(const std::size_t& x, const dynamic_dimension& y);
        MIGRAPHX_EXPORT friend bool operator!=(const dynamic_dimension& x, const std::size_t& y);
        MIGRAPHX_EXPORT friend bool operator!=(const std::size_t& x, const dynamic_dimension& y);
Charlie Lin's avatar
Charlie Lin committed
115
116
117
118

        // add and subtract fixed std::size_t dimension
        dynamic_dimension& operator+=(const std::size_t& x);
        dynamic_dimension& operator-=(const std::size_t& x);
119
120
121
122
123
124
        MIGRAPHX_EXPORT friend dynamic_dimension operator+(const dynamic_dimension& x,
                                                           const std::size_t& y);
        MIGRAPHX_EXPORT friend dynamic_dimension operator+(const std::size_t& x,
                                                           const dynamic_dimension& y);
        MIGRAPHX_EXPORT friend dynamic_dimension operator-(const dynamic_dimension& x,
                                                           const std::size_t& y);
Charlie Lin's avatar
Charlie Lin committed
125
126
    };

127
128
    static const std::vector<type_t>& types();

129
130
131
    static std::string name(type_t t);
    static std::string cpp_type(type_t t);

Paul's avatar
Paul committed
132
133
134
135
    shape();
    shape(type_t t);
    shape(type_t t, std::vector<std::size_t> l);
    shape(type_t t, std::vector<std::size_t> l, std::vector<std::size_t> s);
Paul's avatar
Paul committed
136

Charlie Lin's avatar
Charlie Lin committed
137
138
139
140
141
142
    // Force all calls of the format `shape( type_t, { size_t compatibles } )` to map to
    // shape(type_t, std::vector<std::size_t> l)
    shape(type_t t, std::initializer_list<std::size_t> d);

    shape(type_t t, std::vector<dynamic_dimension> dims);

143
144
    // Construct a dynamic shape from vectors of mins, maxes, and optimals.
    // optimals_list is a vector of optimals that corresponds to each min and max.
145
146
147
    shape(type_t t,
          std::vector<std::size_t> mins,
          std::vector<std::size_t> maxes,
148
          std::vector<std::set<std::size_t>> optimals_list);
149

Paul's avatar
Paul committed
150
151
152
153
154
155
    template <class Range>
    shape(type_t t, const Range& l) : shape(t, std::vector<std::size_t>(l.begin(), l.end()))
    {
    }

    template <class Range1, class Range2>
Paul's avatar
Paul committed
156
    shape(type_t t, const Range1& l, const Range2& s)
Paul's avatar
Paul committed
157
158
159
160
161
        : shape(t,
                std::vector<std::size_t>(l.begin(), l.end()),
                std::vector<std::size_t>(s.begin(), s.end()))
    {
    }
Paul's avatar
Paul committed
162

Paul Fultz II's avatar
Paul Fultz II committed
163
164
    shape(const std::vector<shape>& subs);

165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
    /**
     * Creates an output shape with dimensions equal to the input lengths and strides determined
     * by the permutation argument such that find_permutation() of the output shape returns the
     * inputted permuation.
     *
     * 2D example:
     *   parameters:
     *     l = [2, 3], perm = [1, 0]
     *   therefore:
     *     "original" shape = {lens = [3, 2], strides = [2, 1]}
     *     output_shape = {lens = [2, 3], strides = [1, 2]
     *
     * 3D example:
     *   parameters:
     *     l = [2, 3, 4], perm = [1, 2, 0]
     *   therefore:
     *     "original" shape = {lens = [3, 4, 2], strides = [8, 2, 1]}
     *     output_shape = {lens = [2, 3, 4], strides = [1, 8, 2]}
     */
184
185
    static shape
    from_permutation(type_t t, const std::vector<std::size_t>& l, const std::vector<int64_t>& perm);
186

Paul's avatar
Paul committed
187
    type_t type() const;
Paul's avatar
Paul committed
188
189
    const std::vector<std::size_t>& lens() const;
    const std::vector<std::size_t>& strides() const;
Charlie Lin's avatar
Charlie Lin committed
190

191
    /*!
192
     * The number of dimensions in the shape, either static or dynamic.
193
194
195
196
     * Same as the number of indices required to get a data value.
     */
    std::size_t ndim() const;

Charlie Lin's avatar
Charlie Lin committed
197
198
199
    /*!
     * Return the number of elements in the tensor.
     */
Paul's avatar
Paul committed
200
    std::size_t elements() const;
Charlie Lin's avatar
Charlie Lin committed
201
202
203
204
205

    /*!
     * Return the number of total bytes used for storage of the tensor data; includes subshapes.
     * For dynamic shape, returns the maximum number of bytes presuming a packed shape.
     */
Paul's avatar
Paul committed
206
    std::size_t bytes() const;
Charlie Lin's avatar
Charlie Lin committed
207
208
209
210
211

    /*!
     * Return the size of the type of the main shape.
     * Returns 0 if there are subshapes.
     */
Scott Thornton's avatar
Scott Thornton committed
212
    std::size_t type_size() const;
Paul's avatar
Paul committed
213

Charlie Lin's avatar
Charlie Lin committed
214
215
216
217
    const std::vector<dynamic_dimension>& dyn_dims() const;

    /*!
     * Minimum lengths for dynamic shape.
218
     * lens() for static shape.
Charlie Lin's avatar
Charlie Lin committed
219
220
221
222
223
     */
    std::vector<std::size_t> min_lens() const;

    /*!
     * Maximum lengths for dynamic shape.
224
     * lens() for static shape.
Charlie Lin's avatar
Charlie Lin committed
225
226
227
228
229
     */
    std::vector<std::size_t> max_lens() const;

    /*!
     * Optimum lengths for dynamic shape.
230
     * Empty for static shape.
Charlie Lin's avatar
Charlie Lin committed
231
     */
232
    std::vector<std::set<std::size_t>> opt_lens() const;
Charlie Lin's avatar
Charlie Lin committed
233

Paul's avatar
Paul committed
234
    /// Map multiple indices to space index
Paul's avatar
Paul committed
235
    std::size_t index(std::initializer_list<std::size_t> l) const;
Paul's avatar
Paul committed
236
    /// Map multiple indices to space index
Paul's avatar
Paul committed
237
    std::size_t index(const std::vector<std::size_t>& l) const;
Paul's avatar
Paul committed
238

Paul's avatar
Paul committed
239
    /// Map multiple indices from a range of iterator to a space index
Paul's avatar
Paul committed
240
    template <class Iterator>
Paul's avatar
Paul committed
241
242
    std::size_t index(Iterator start, Iterator last) const
    {
243
244
245
246
        if(this->dynamic())
        {
            MIGRAPHX_THROW("SHAPE: index() called on dynamic shape");
        }
Paul's avatar
Paul committed
247
248
        assert(std::distance(start, last) <= this->lens().size());
        assert(this->lens().size() == this->strides().size());
Paul Fultz II's avatar
Paul Fultz II committed
249
        return std::inner_product(start, last, this->strides().begin(), std::size_t{0}); // NOLINT
Paul's avatar
Paul committed
250
    }
Paul's avatar
Paul committed
251

Paul's avatar
Paul committed
252
    /// Map element index to space index
Paul's avatar
Paul committed
253
254
    std::size_t index(std::size_t i) const;

255
256
    /// Map element index to multi-dimensional index
    std::vector<std::size_t> multi(std::size_t idx) const;
257

258
259
260
261
262
263
    /// Map element index to multi-dimensional index and put them them into location provided by
    /// pointers
    void multi_copy(std::size_t idx, std::size_t* start, const std::size_t* end) const;

    /// Returns true if the shape is packed (number of elements and buffer size the same) with
    /// no padding
Paul's avatar
Paul committed
264
    bool packed() const;
Charlie Lin's avatar
Charlie Lin committed
265

266
    /// Returns true if the shape has been transposed. That is the strides are not in descending
Paul's avatar
Paul committed
267
    /// order
Paul's avatar
Paul committed
268
    bool transposed() const;
Charlie Lin's avatar
Charlie Lin committed
269

Paul's avatar
Paul committed
270
    /// Returns true if the shape is broadcasting a dimension. That is, one of the strides are zero
Paul's avatar
Paul committed
271
    bool broadcasted() const;
Charlie Lin's avatar
Charlie Lin committed
272

Paul's avatar
Paul committed
273
274
    /// Returns true if the shape is in its standard format. That is, the shape is both packed and
    /// not transposed.
Paul's avatar
Paul committed
275
    bool standard() const;
Charlie Lin's avatar
Charlie Lin committed
276

Khalique's avatar
Khalique committed
277
278
    /// Returns true if all strides are equal to 0 (scalar tensor)
    bool scalar() const;
Paul's avatar
Paul committed
279

Charlie Lin's avatar
Charlie Lin committed
280
281
282
    /// Return true if the shape is dynamic
    bool dynamic() const;

Charlie Lin's avatar
Charlie Lin committed
283
284
285
    /// Return true if this shape or any of the sub_shapes are dynamic
    bool any_of_dynamic() const;

286
287
    shape normalize_standard() const;

288
289
290
    shape with_lens(type_t t, const std::vector<std::size_t>& l) const;
    shape with_lens(const std::vector<std::size_t>& l) const;

291
292
    shape with_type(type_t t) const;

293
    // convert the shape to an equivalent dynamic shape with empty optimals
294
295
    shape to_dynamic() const;

296
297
298
    // convert the shape to a static one setting any non-fixed dynamic_dimensions to x
    shape to_static(std::size_t x) const;

299
300
301
    MIGRAPHX_EXPORT friend bool operator==(const shape& x, const shape& y);
    MIGRAPHX_EXPORT friend bool operator!=(const shape& x, const shape& y);
    MIGRAPHX_EXPORT friend std::ostream& operator<<(std::ostream& os, const shape& x);
Paul's avatar
Paul committed
302

Paul's avatar
Paul committed
303
    template <class T>
Paul's avatar
Paul committed
304
305
    struct as
    {
306
        using type = std::conditional_t<std::is_same<T, bool>{}, int8_t, T>;
Paul's avatar
Paul committed
307

308
309
310
311
        type max() const { return std::numeric_limits<type>::max(); }

        type min() const { return std::numeric_limits<type>::lowest(); }

312
313
        type nan() const { return std::numeric_limits<type>::quiet_NaN(); }

Paul's avatar
Paul committed
314
        template <class U>
315
        type operator()(U u) const
Paul's avatar
Paul committed
316
        {
317
            return type(u);
Paul's avatar
Paul committed
318
319
        }

Paul's avatar
Paul committed
320
        template <class U>
321
        type* operator()(U* u) const
Paul's avatar
Paul committed
322
        {
323
            return static_cast<type*>(u);
Paul's avatar
Paul committed
324
325
        }

Paul's avatar
Paul committed
326
        template <class U>
327
        const type* operator()(const U* u) const
Paul's avatar
Paul committed
328
        {
329
            return static_cast<type*>(u);
Paul's avatar
Paul committed
330
331
        }

332
        type operator()() const { return {}; }
Paul's avatar
Paul committed
333

334
        std::size_t size(std::size_t n = 1) const { return sizeof(type) * n; }
Paul's avatar
Paul committed
335

Paul Fultz II's avatar
Paul Fultz II committed
336
337
338
339
        auto is_integral() const { return std::is_integral<type>{}; }
        auto is_signed() const { return std::is_signed<type>{}; }
        auto is_unsigned() const { return std::is_unsigned<type>{}; }

Paul's avatar
Paul committed
340
        template <class U>
341
        type* from(U* buffer, std::size_t n = 0) const
Paul's avatar
Paul committed
342
        {
343
            return reinterpret_cast<type*>(buffer) + n;
Paul's avatar
Paul committed
344
        }
Paul's avatar
Paul committed
345

Paul's avatar
Paul committed
346
        template <class U>
347
        const type* from(const U* buffer, std::size_t n = 0) const
Paul's avatar
Paul committed
348
        {
349
            return reinterpret_cast<const type*>(buffer) + n;
Paul's avatar
Paul committed
350
        }
Paul's avatar
Paul committed
351

352
        type_t type_enum() const { return get_type<type>{}; }
Paul's avatar
Paul committed
353
354
    };

Paul Fultz II's avatar
Paul Fultz II committed
355
356
    template <class Visitor, class TupleVisitor>
    static void visit(type_t t, Visitor v, TupleVisitor tv)
Paul's avatar
Paul committed
357
    {
358
        switch(t)
Paul's avatar
Paul committed
359
        {
bpickrel's avatar
bpickrel committed
360
        case tuple_type: {
Paul Fultz II's avatar
Paul Fultz II committed
361
362
363
            tv();
            return;
        }
Paul's avatar
Paul committed
364
#define MIGRAPHX_SHAPE_GENERATE_VISITOR_CASE(x, t) \
Paul's avatar
Paul committed
365
    case x: v(as<t>()); return;
Paul's avatar
Paul committed
366
367
            MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_VISITOR_CASE)
#undef MIGRAPHX_SHAPE_GENERATE_VISITOR_CASE
Paul's avatar
Paul committed
368
        }
Paul's avatar
Paul committed
369
        MIGRAPHX_THROW("Unknown type");
Paul's avatar
Paul committed
370
    }
Paul's avatar
Paul committed
371

372
    template <class Visitor>
Paul Fultz II's avatar
Paul Fultz II committed
373
    static void visit(type_t t, Visitor v)
374
    {
Paul Fultz II's avatar
Paul Fultz II committed
375
376
377
378
379
380
381
        return visit(t, v, [] { MIGRAPHX_THROW("Tuple cannot be visited."); });
    }

    template <class... Visitors>
    void visit_type(Visitors... vs) const
    {
        visit(this->type(), vs...);
382
383
    }

Paul's avatar
Paul committed
384
385
386
    template <class Visitor>
    static void visit_types(Visitor v)
    {
Paul's avatar
Paul committed
387
#define MIGRAPHX_SHAPE_GENERATE_VISITOR_ALL(x, t) v(as<t>());
Paul's avatar
Paul committed
388
389
390
391
        MIGRAPHX_SHAPE_VISIT_TYPES(MIGRAPHX_SHAPE_GENERATE_VISITOR_ALL)
#undef MIGRAPHX_SHAPE_GENERATE_VISITOR_ALL
    }

392
    std::string type_string() const;
393
    static type_t parse_type(const std::string& s);
394

Paul Fultz II's avatar
Paul Fultz II committed
395
396
    const std::vector<shape>& sub_shapes() const;

Charlie Lin's avatar
Charlie Lin committed
397
398
399
400
401
    /*!
     * Returns the number of elements in the data buffer.
     * For a dynamic shape, returns the maximum number of elements of the data buffer and assumes it
     * is packed.
     */
402
403
    std::size_t element_space() const;

Paul's avatar
Paul committed
404
    private:
405
    shape(std::shared_ptr<shape_impl> pimpl);
Paul's avatar
Paul committed
406
    std::shared_ptr<const shape_impl> impl;
Paul's avatar
Paul committed
407
408
};

409
410
MIGRAPHX_EXPORT void migraphx_to_value(value& v, const shape& s);
MIGRAPHX_EXPORT void migraphx_from_value(const value& v, shape& s);
411

Paul's avatar
Paul committed
412
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
413
} // namespace migraphx
Paul's avatar
Paul committed
414
415

#endif