megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

/**
 * \file dnn/src/naive/topk/opr_impl.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "./opr_impl.h"
#include "src/common/utils.h"
#include "src/naive/handle.h"

#include <algorithm>
#include <cmath>
#include <limits>

using namespace megdnn;
using namespace naive;

namespace {
constexpr ptrdiff_t rowoff(ptrdiff_t i, ptrdiff_t lda) {
    return i * lda;
}
}  // namespace

template <typename ctype>
void TopKImpl::dispatch_with_ctype(
        int k, size_t m, size_t n, ptrdiff_t lda, const ctype* data, ctype* values,
        int* indices, void* workspace) {
    using CmpGt = std::greater<std::pair<ctype, uint32_t>>;
    thin_function<void()> compute;
    switch (param().mode) {
        case Param::Mode::KTH_ONLY:
            compute = [k0 = k, m, n, lda, data, values,
                       wk = static_cast<ctype*>(workspace)]() {
                int k = k0 < 0 ? n + k0 : k0 - 1;
                for (size_t i = 0; i < m; ++i) {
                    memcpy(wk, data + rowoff(i, lda), sizeof(ctype) * n);
                    std::nth_element(wk, wk + k, wk + n);
                    values[i] = wk[k];
                }
            };
            break;
        case Param::Mode::VALUE_IDX_NOSORT:
            megdnn_assert(n <= std::numeric_limits<uint32_t>::max());
            compute = [k, m, n, lda, data, values, indices,
                       wk = static_cast<std::pair<ctype, uint32_t>*>(workspace)]() {
                int ow = std::abs(k);
                for (size_t i = 0; i < m; ++i) {
                    for (size_t j = 0; j < n; ++j) {
                        wk[j].first = data[rowoff(i, lda) + j];
                        wk[j].second = j;
                    }
                    if (k < 0) {
                        std::nth_element(wk, wk - k - 1, wk + n, CmpGt{});
                    } else {
                        std::nth_element(wk, wk + k - 1, wk + n);
                    }
                    for (int j = 0; j < ow; ++j) {
                        values[i * ow + j] = wk[j].first;
                        indices[i * ow + j] = wk[j].second;
                    }
                }
            };
            break;
        case Param::Mode::VALUE_IDX_SORTED:
            megdnn_assert(n <= std::numeric_limits<uint32_t>::max());
            compute = [k, m, n, lda, data, values, indices,
                       wk = static_cast<std::pair<ctype, uint32_t>*>(workspace)]() {
                for (size_t i = 0; i < m; ++i) {
                    for (size_t j = 0; j < n; ++j) {
                        wk[j].first = data[rowoff(i, lda) + j];
                        wk[j].second = j;
                    }
                    if (k < 0) {
                        std::partial_sort(wk, wk - k, wk + n, CmpGt{});
                    } else {
                        std::partial_sort(wk, wk + k, wk + n);
                    }
                    for (int j = 0, jt = std::abs(k); j < jt; ++j) {
                        values[i * jt + j] = wk[j].first;
                        indices[i * jt + j] = wk[j].second;
                    }
                }
            };
            break;
        default:
            megdnn_throw("invalid TopK mode");
    }

    static_cast<HandleImpl*>(handle())->dispatch_kern(std::move(compute));
}

void TopKImpl::do_exec(
        int k, _megdnn_tensor_in data, _megdnn_tensor_out values, int32_t* indices,
        _megdnn_workspace workspace) {
    size_t m = data.layout[0], n = data.layout[1];
    ptrdiff_t lda = data.layout.stride[0];
    switch (data.layout.dtype.enumv()) {
#define cb(t)                                                                \
    case DTypeTrait<t>::enumv:                                               \
        do {                                                                 \
            using ct = DTypeTrait<t>::ctype;                                 \
            dispatch_with_ctype<ct>(                                         \
                    k, m, n, lda, data.ptr<ct>(), values.ptr<ct>(), indices, \
                    workspace.raw_ptr);                                      \
            return;                                                          \
        } while (0);
        MEGDNN_FOREACH_COMPUTING_DTYPE(cb);
        default:
            megdnn_throw("unsupported dtype in naive TopKImpl");
    }
}

size_t TopKImpl::get_workspace_in_bytes(
        int k, const TensorLayout& data, const TensorLayout& values,
        const TensorLayout& indices) {
    MEGDNN_MARK_USED_VAR(k);
    MEGDNN_MARK_USED_VAR(values);
    MEGDNN_MARK_USED_VAR(indices);
    return std::max(sizeof(uint32_t), data.dtype.size()) * 2 * data[1];
}