tensor_view.hpp 4.06 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
9

#include <iostream>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Paul's avatar
Paul committed
142
template <class T, class U>
143
144
145
146
147
148
149
150
151
152
153
154
155
156
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
157
158
159
160
161
template <class T, class U>
bool operator!=(const tensor_view<T>& x, const tensor_view<U>& y)
{
    return !(x == y);
}
162

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

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

#endif