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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

/**
 * \file dnn/src/common/argmxx_helper.h
 * 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.
 */

#pragma once
#include "megdnn/dtype.h"

#if MEGDNN_CC_HOST
#include "megdnn/basic_types.h"
#endif

namespace megdnn {
namespace argmxx {

template <typename stype_, bool is_max>
struct ArgmxxOp {
    struct wtype {
        stype_ key;
        dt_int32 val;
        MEGDNN_HOST MEGDNN_DEVICE wtype() {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(stype_ key, dt_int32 val)
                : key(key), val(val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(wtype& rhs) : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(volatile wtype& rhs)
                : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(const wtype& rhs)
                : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(const volatile wtype& rhs)
                : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE volatile wtype& operator=(const wtype& rhs) volatile {
            this->key = rhs.key;
            this->val = rhs.val;
            return *this;
        }
    };
    MEGDNN_HOST MEGDNN_DEVICE
    ArgmxxOp(stype_* src, dt_int32* dst, uint32_t A, uint32_t B, uint32_t C)
            : src(src),
              dst(dst),
              A(A),
              B(B),
              C(C),
              INIT(wtype(
                      is_max ? DTypeTrait<stype_>::min() : DTypeTrait<stype_>::max(),
                      0)) {}
    MEGDNN_HOST MEGDNN_DEVICE wtype read(uint32_t idx) {
        wtype res;
        res.key = src[idx];
        res.val = idx / C % B;
        return res;
    }
    MEGDNN_HOST MEGDNN_DEVICE void write(uint32_t idx, wtype val) {
        dst[idx] = val.val;
    }
    static MEGDNN_HOST MEGDNN_DEVICE wtype apply(wtype lhs, wtype rhs) {
        if (is_max) {
            if (lhs.key > rhs.key)
                return lhs;
            else
                return rhs;
        } else {
            if (lhs.key < rhs.key)
                return lhs;
            else
                return rhs;
        }
    }
    stype_* src;
    dt_int32* dst;
    uint32_t A, B, C;
    const wtype INIT;
};

template <bool is_max>
struct ArgmxxOp<dt_float32, is_max> {
    using stype_ = dt_float32;
    struct wtype {
        stype_ key;
        dt_int32 val;
        MEGDNN_HOST MEGDNN_DEVICE wtype() {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(stype_ key, dt_int32 val)
                : key(key), val(val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(wtype& rhs) : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(volatile wtype& rhs)
                : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(const wtype& rhs)
                : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE wtype(const volatile wtype& rhs)
                : key(rhs.key), val(rhs.val) {}
        MEGDNN_HOST MEGDNN_DEVICE volatile wtype& operator=(const wtype& rhs) volatile {
            this->key = rhs.key;
            this->val = rhs.val;
            return *this;
        }
    };
    MEGDNN_HOST MEGDNN_DEVICE
    ArgmxxOp(stype_* src, dt_int32* dst, uint32_t A, uint32_t B, uint32_t C)
            : src(src),
              dst(dst),
              A(A),
              B(B),
              C(C),
              INIT(wtype(
                      is_max ? DTypeTrait<stype_>::min() : DTypeTrait<stype_>::max(),
                      0)) {}
    MEGDNN_HOST MEGDNN_DEVICE wtype read(uint32_t idx) {
        wtype res;
        res.key = src[idx];
        res.val = idx / C % B;
        return res;
    }
    MEGDNN_HOST MEGDNN_DEVICE void write(uint32_t idx, wtype val) {
        dst[idx] = val.val;
    }
    static MEGDNN_HOST MEGDNN_DEVICE wtype apply(wtype lhs, wtype rhs) {
#if defined(__CUDA_ARCH__)
        if (isnan(lhs.key))
#else
        if (std::isnan(lhs.key))
#endif
            return lhs;
        if (is_max) {
            if (lhs.key > rhs.key)
                return lhs;
            else
                return rhs;
        } else {
            if (lhs.key < rhs.key)
                return lhs;
            else
                return rhs;
        }
    }
    stype_* src;
    dt_int32* dst;
    uint32_t A, B, C;
    const wtype INIT;
};

}  // namespace argmxx
}  // namespace megdnn
// vim: syntax=cpp.doxygen