"vscode:/vscode.git/clone" did not exist on "e1039a1cffdb0eadacaa4eb414865c65c6c4c680"
common_dims.cpp 2.84 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#include <migraphx/common_dims.hpp>
#include <migraphx/ranges.hpp>
#include <algorithm>
#include <cassert>
#include <numeric>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {

template <class Iterator>
static auto compute_end_dim(Iterator start, Iterator last, std::size_t dim)
{
    std::size_t x = 1;
    auto it       = std::find_if(start, last, [&](auto i) {
        x *= i;
        return x >= dim;
    });
    if(x != dim)
        return start;
    return it;
}

Paul's avatar
Format  
Paul committed
23
template <class Iterator>
Paul's avatar
Paul committed
24
25
26
27
static auto elements(Iterator start, Iterator last)
{
    return std::accumulate(start, last, std::size_t{1}, std::multiplies<>{});
}
Paul's avatar
Format  
Paul committed
28
template <class Range>
Paul's avatar
Paul committed
29
30
31
32
33
static auto elements(const Range& r)
{
    return elements(r.begin(), r.end());
}

Paul's avatar
Paul committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
struct common_dim_state
{
    common_dim_state(const std::vector<std::size_t>& pdims)
    : dims(&pdims), it(dims->begin())
    {}
    const std::vector<std::size_t>* dims;
    std::vector<std::size_t>::const_iterator it;
    std::size_t rem = 1;
    std::size_t get() const
    {
        return *it;
    }
    bool is_end() const
    {
        return it == dims->end();
    }
    void next(std::size_t i = 1)
    {
        it += i;
    }
    auto dims_for(std::size_t d) const
    {
        auto dim_end = compute_end_dim(it, dims->end(), d);
        return range(it, dim_end);
    }
}

Paul's avatar
Format  
Paul committed
61
62
common_dims common_dims::compute(const std::vector<std::size_t>& dims1,
                                 const std::vector<std::size_t>& dims2)
Paul's avatar
Paul committed
63
64
65
{
    assert(elements(dims1) == elements(dims2));
    common_dims cd;
Paul's avatar
Format  
Paul committed
66
67
    auto it1         = dims1.begin();
    auto it2         = dims2.begin();
Paul's avatar
Paul committed
68
69
70
71
72
73
    std::size_t rem1 = 1;
    std::size_t rem2 = 1;
    while(it1 != dims1.end() and it2 != dims2.end())
    {
        auto d1 = *it1;
        auto d2 = *it2;
Paul's avatar
Format  
Paul committed
74
        if(d1 == d2)
Paul's avatar
Paul committed
75
76
77
78
79
80
81
        {
            cd.axes_map1.push_back({cd.dims.size()});
            cd.axes_map2.push_back({cd.dims.size()});
            cd.dims.push_back(d1);
            it1++;
            it2++;
        }
Paul's avatar
Format  
Paul committed
82
        else if(d1 < d2)
Paul's avatar
Paul committed
83
84
        {
            auto dim_end = compute_end_dim(it1, dims1.begin(), d2);
Paul's avatar
Format  
Paul committed
85
86
87
            auto dims    = range(it1, dim_end);
            auto n       = elements(dims);
            if(n != d2)
Paul's avatar
Paul committed
88
89
            {
                // If not divisible then we can't compute a common dims
Paul's avatar
Format  
Paul committed
90
                if((d2 % n) != 0)
Paul's avatar
Paul committed
91
92
93
94
95
96
97
98
99
                    return {};
                rem1 = d2 / n;
            }
            std::vector<std::size_t> axes(distance(dims));
            std::iota(axes.begin(), axes.end(), cd.dims.size());
            cd.axes_map1.push_back(axes);
            cd.axes_map2.push_back(axes);

            cd.dims.insert(cd.dims.end(), dims.begin(), dims.end());
Paul's avatar
Format  
Paul committed
100
            if(rem1 != 1)
Paul's avatar
Paul committed
101
102
103
104
105
106
107
108
109
110
                cd.dims.push_back(rem1);
            it1 += distance(dims);
            it2++;
        }
    }
    return cd;
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx