tensor_view.hpp 3.68 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
        return m_data[m_shape.index({static_cast<std::size_t>(xs)...})];
Paul's avatar
Paul committed
33
34
    }

Paul's avatar
Paul committed
35
    template <class... Ts, MIGRAPH_REQUIRES(std::is_integral<Ts>{}...)>
Paul's avatar
Paul committed
36
37
    T& operator()(Ts... xs)
    {
Paul's avatar
Paul committed
38
        return m_data[m_shape.index({static_cast<std::size_t>(xs)...})];
Paul's avatar
Paul committed
39
40
    }

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

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

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

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

    T& front()
    {
        assert(!this->empty());
Paul's avatar
Paul committed
68
        return m_data[0];
Paul's avatar
Paul committed
69
70
71
72
73
    }

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

    T& back()
    {
        assert(!this->empty());
Paul's avatar
Paul committed
80
        return m_data[m_shape.index(this->size() - 1)];
Paul's avatar
Paul committed
81
82
83
84
85
    }

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

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

    T* end()
    {
Paul's avatar
Paul committed
98
        assert(this->m_shape.packed());
Paul's avatar
Paul committed
99
        if(this->empty())
Paul's avatar
Paul committed
100
            return m_data;
Paul's avatar
Paul committed
101
        else
Paul's avatar
Paul committed
102
            return m_data + this->size();
Paul's avatar
Paul committed
103
104
105
106
    }

    const T* begin() const
    {
Paul's avatar
Paul committed
107
108
        assert(this->m_shape.packed());
        return m_data;
Paul's avatar
Paul committed
109
110
111
112
    }

    const T* end() const
    {
Paul's avatar
Paul committed
113
        assert(this->m_shape.packed());
Paul's avatar
Paul committed
114
        if(this->empty())
Paul's avatar
Paul committed
115
            return m_data;
Paul's avatar
Paul committed
116
        else
Paul's avatar
Paul committed
117
            return m_data + this->size();
Paul's avatar
Paul committed
118
119
    }

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

Paul's avatar
Paul committed
133
    private:
Paul's avatar
Paul committed
134
135
    T* m_data;
    shape m_shape;
Paul's avatar
Paul committed
136
137
};

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

Paul's avatar
Paul committed
159
template <class T>
Paul's avatar
Paul committed
160
161
162
163
164
tensor_view<T> make_view(shape s, T* data)
{
    return {s, data};
}

Paul's avatar
Paul committed
165
} // namespace migraph
Paul's avatar
Paul committed
166
167

#endif