"model/sentencepiece.go" did not exist on "4b037a97dcc0387d2b94d293d4bd20e7e262d609"
common_dims.cpp 2.27 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
Format  
Paul committed
34
35
common_dims common_dims::compute(const std::vector<std::size_t>& dims1,
                                 const std::vector<std::size_t>& dims2)
Paul's avatar
Paul committed
36
37
38
{
    assert(elements(dims1) == elements(dims2));
    common_dims cd;
Paul's avatar
Format  
Paul committed
39
40
    auto it1         = dims1.begin();
    auto it2         = dims2.begin();
Paul's avatar
Paul committed
41
42
43
44
45
46
    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
47
        if(d1 == d2)
Paul's avatar
Paul committed
48
49
50
51
52
53
54
        {
            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
55
        else if(d1 < d2)
Paul's avatar
Paul committed
56
57
        {
            auto dim_end = compute_end_dim(it1, dims1.begin(), d2);
Paul's avatar
Format  
Paul committed
58
59
60
            auto dims    = range(it1, dim_end);
            auto n       = elements(dims);
            if(n != d2)
Paul's avatar
Paul committed
61
62
            {
                // If not divisible then we can't compute a common dims
Paul's avatar
Format  
Paul committed
63
                if((d2 % n) != 0)
Paul's avatar
Paul committed
64
65
66
67
68
69
70
71
72
                    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
73
            if(rem1 != 1)
Paul's avatar
Paul committed
74
75
76
77
78
79
80
81
82
83
                cd.dims.push_back(rem1);
            it1 += distance(dims);
            it2++;
        }
    }
    return cd;
}

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx