tensor_view.hpp 4.15 KB
Newer Older
Paul's avatar
Paul committed
1
2
#ifndef MIGRAPH_GUARD_TENSOR_VIEW_HPP
#define MIGRAPH_GUARD_TENSOR_VIEW_HPP
Paul's avatar
Paul committed
3

Paul's avatar
Paul committed
4
5
6
#include <migraph/shape.hpp>
#include <migraph/float_equal.hpp>
#include <migraph/requires.hpp>
Paul's avatar
Paul committed
7
8

#include <iostream>
Paul's avatar
Paul committed
9
#include <utility>
Paul's avatar
Paul committed
10

Paul's avatar
Paul committed
11
namespace migraph {
Paul's avatar
Paul committed
12

Paul's avatar
Paul committed
13
template <class T>
Paul's avatar
Paul committed
14
15
struct tensor_view
{
16
    using value_type = T;
Paul's avatar
Paul committed
17
    tensor_view() : m_data(nullptr) {}
Paul's avatar
Paul committed
18
    tensor_view(shape s, T* d) : m_data(d), m_shape(std::move(s)) {}
Paul's avatar
Paul committed
19

Paul's avatar
Paul committed
20
    const shape& get_shape() const { return this->m_shape; }
Paul's avatar
Paul committed
21

Paul's avatar
Paul committed
22
    bool empty() const { return m_data == nullptr || m_shape.lens().empty(); }
Paul's avatar
Paul committed
23

Paul's avatar
Paul committed
24
    std::size_t size() const { return m_shape.elements(); }
Paul's avatar
Paul committed
25

Paul's avatar
Paul committed
26
    T* data() { return this->m_data; }
Paul's avatar
Paul committed
27

Paul's avatar
Paul committed
28
    const T* data() const { return this->m_data; }
Paul's avatar
Paul committed
29

Paul's avatar
Paul committed
30
    template <class... Ts, MIGRAPH_REQUIRES(std::is_integral<Ts>{}...)>
Paul's avatar
Paul committed
31
32
    const T& operator()(Ts... xs) const
    {
Paul's avatar
Paul committed
33
        assert(std::vector<std::size_t>{static_cast<std::size_t>(xs)...} < m_shape.lens());
Scott Thornton's avatar
Scott Thornton committed
34
        assert(m_shape.index({static_cast<std::size_t>(xs)...}) < m_shape.bytes() / sizeof(T));
Paul's avatar
Paul committed
35
        return m_data[m_shape.index({static_cast<std::size_t>(xs)...})];
Paul's avatar
Paul committed
36
37
    }

Paul's avatar
Paul committed
38
    template <class... Ts, MIGRAPH_REQUIRES(std::is_integral<Ts>{}...)>
Paul's avatar
Paul committed
39
40
    T& operator()(Ts... xs)
    {
Paul's avatar
Paul committed
41
        assert(std::vector<std::size_t>{static_cast<std::size_t>(xs)...} < m_shape.lens());
Scott Thornton's avatar
Scott Thornton committed
42
        assert(m_shape.index({static_cast<std::size_t>(xs)...}) < m_shape.bytes() / sizeof(T));
Paul's avatar
Paul committed
43
        return m_data[m_shape.index({static_cast<std::size_t>(xs)...})];
Paul's avatar
Paul committed
44
45
    }

Paul's avatar
Paul committed
46
    template <class Iterator, MIGRAPH_REQUIRES(not std::is_integral<Iterator>{})>
Paul's avatar
Paul committed
47
48
49
50
51
    const T& operator()(Iterator start, Iterator last) const
    {
        return m_data[m_shape.index(start, last)];
    }

Paul's avatar
Paul committed
52
    template <class Iterator, MIGRAPH_REQUIRES(not std::is_integral<Iterator>{})>
Paul's avatar
Paul committed
53
54
55
56
57
    T& operator()(Iterator start, Iterator last)
    {
        return m_data[m_shape.index(start, last)];
    }

Paul's avatar
Paul committed
58
59
60
    T& operator[](std::size_t i)
    {
        assert(!this->empty() && i < this->size());
Paul's avatar
Paul committed
61
        return m_data[m_shape.index(i)];
Paul's avatar
Paul committed
62
63
64
65
66
    }

    const T& operator[](std::size_t i) const
    {
        assert(!this->empty() && i < this->size());
Paul's avatar
Paul committed
67
        return m_data[m_shape.index(i)];
Paul's avatar
Paul committed
68
69
70
71
72
    }

    T& front()
    {
        assert(!this->empty());
Paul's avatar
Paul committed
73
        return m_data[0];
Paul's avatar
Paul committed
74
75
76
77
78
    }

    const T& front() const
    {
        assert(!this->empty());
Paul's avatar
Paul committed
79
        return m_data[0];
Paul's avatar
Paul committed
80
81
82
83
84
    }

    T& back()
    {
        assert(!this->empty());
Paul's avatar
Paul committed
85
        return m_data[m_shape.index(this->size() - 1)];
Paul's avatar
Paul committed
86
87
88
89
90
    }

    const T& back() const
    {
        assert(!this->empty());
Paul's avatar
Paul committed
91
        return m_data[m_shape.index(this->size() - 1)];
Paul's avatar
Paul committed
92
93
    }

Paul's avatar
Paul committed
94
    // TODO: Add iterators so it can handle nonstandard tensors
Paul's avatar
Paul committed
95
96
    T* begin()
    {
Paul's avatar
Paul committed
97
        assert(this->m_shape.standard() or this->empty());
Paul's avatar
Paul committed
98
        return m_data;
Paul's avatar
Paul committed
99
100
101
102
    }

    T* end()
    {
Paul's avatar
Paul committed
103
        assert(this->m_shape.standard() or this->empty());
Paul's avatar
Paul committed
104
        if(this->empty())
Paul's avatar
Paul committed
105
            return m_data;
Paul's avatar
Paul committed
106
        else
Paul's avatar
Paul committed
107
            return m_data + this->size();
Paul's avatar
Paul committed
108
109
110
111
    }

    const T* begin() const
    {
Paul's avatar
Paul committed
112
        assert(this->m_shape.standard() or this->empty());
Paul's avatar
Paul committed
113
        return m_data;
Paul's avatar
Paul committed
114
115
116
117
    }

    const T* end() const
    {
Paul's avatar
Paul committed
118
        assert(this->m_shape.standard() or this->empty());
Paul's avatar
Paul committed
119
        if(this->empty())
Paul's avatar
Paul committed
120
            return m_data;
Paul's avatar
Paul committed
121
        else
Paul's avatar
Paul committed
122
            return m_data + this->size();
Paul's avatar
Paul committed
123
124
    }

Paul's avatar
Paul committed
125
126
127
128
129
    friend std::ostream& operator<<(std::ostream& os, const tensor_view<T>& x)
    {
        if(!x.empty())
        {
            os << x.front();
Paul's avatar
Paul committed
130
            for(std::size_t i = 1; i < x.m_shape.elements(); i++)
Paul's avatar
Paul committed
131
            {
Paul's avatar
Paul committed
132
                os << ", " << x.m_data[x.m_shape.index(i)];
Paul's avatar
Paul committed
133
134
135
136
137
            }
        }
        return os;
    }

Paul's avatar
Paul committed
138
    private:
Paul's avatar
Paul committed
139
140
    T* m_data;
    shape m_shape;
Paul's avatar
Paul committed
141
142
};

Paul's avatar
Paul committed
143
template <class T, class U>
144
145
146
147
148
149
150
151
152
153
154
155
156
157
bool operator==(const tensor_view<T>& x, const tensor_view<U>& y)
{
    if(x.get_shape() == y.get_shape())
    {
        for(std::size_t i = 0; i < x.get_shape().elements(); i++)
        {
            if(!float_equal(x[i], y[i]))
                return false;
        }
        return true;
    }
    return false;
}

Paul's avatar
Paul committed
158
159
160
161
162
template <class T, class U>
bool operator!=(const tensor_view<T>& x, const tensor_view<U>& y)
{
    return !(x == y);
}
163

Paul's avatar
Paul committed
164
template <class T>
Paul's avatar
Paul committed
165
166
167
168
169
tensor_view<T> make_view(shape s, T* data)
{
    return {s, data};
}

Paul's avatar
Paul committed
170
} // namespace migraph
Paul's avatar
Paul committed
171
172

#endif