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

/**
 * \file dnn/src/common/cond_take/predicate.cuh
 * 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/arch.h"
#include "src/common/opr_param_defs_enumv.cuh"

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

#ifndef __device__
#define __device__
#define __host__
#define def_device
#endif

#include <cmath>

namespace megdnn {
namespace cond_take {
typedef param_enumv::CondTake::Mode PEnum;

struct KParam {
    float val, eps;
#if MEGDNN_CC_HOST
    KParam(const param::CondTake& p) : val(p.val), eps(p.eps) {}
#endif
};

template <uint32_t mode, typename ctype>
struct Pred;

#define do_inst_eq_f(_ct)                                    \
    template <>                                              \
    struct Pred<PEnum::EQ, _ct> {                            \
        typedef _ct ctype;                                   \
        ctype val, eps;                                      \
        Pred(const KParam& p) : val(p.val), eps(p.eps) {}    \
        __device__ __host__ bool operator()(ctype x) const { \
            return fabsf(val - x) < eps;                     \
        }                                                    \
    };

#define do_inst_eq_i(_ct)                                                       \
    template <>                                                                 \
    struct Pred<PEnum::EQ, _ct> {                                               \
        typedef _ct ctype;                                                      \
        ctype val;                                                              \
        Pred(const KParam& p) : val(p.val) {}                                   \
        __device__ __host__ bool operator()(ctype x) const { return val == x; } \
    };

#define inst_eq_f(_dt) do_inst_eq_f(DTypeTrait<_dt>::ctype)
#define inst_eq_i(_dt) do_inst_eq_i(DTypeTrait<_dt>::ctype)
MEGDNN_FOREACH_COMPUTING_DTYPE_FLOAT(inst_eq_f)
MEGDNN_FOREACH_COMPUTING_DTYPE_INT(inst_eq_i)
inst_eq_i(::megdnn::dtype::Bool)
#undef inst_eq_f
#undef inst_eq_i

        template <typename ctype_>
        struct Pred<PEnum::NEQ, ctype_> {
    typedef ctype_ ctype;
    Pred<PEnum::EQ, ctype> eq;

    Pred(const KParam& p) : eq(p) {}

    __device__ __host__ bool operator()(ctype x) const { return !this->eq(x); }
};

#define DEF_OP(_name, _op)                                                       \
    template <typename ctype_>                                                   \
    struct Pred<PEnum::_name, ctype_> {                                          \
        typedef ctype_ ctype;                                                    \
        ctype val;                                                               \
        Pred(const KParam& p) : val(p.val) {}                                    \
        __device__ __host__ bool operator()(ctype x) const { return x _op val; } \
    }

DEF_OP(LT, <);
DEF_OP(LEQ, <=);
DEF_OP(GT, >);
DEF_OP(GEQ, >=);

#undef DEF_OP

#define MEGDNN_FOREACH_COND_TAKE_MODE(cb) cb(EQ) cb(NEQ) cb(LT) cb(LEQ) cb(GT) cb(GEQ)

}  // namespace cond_take
}  // namespace megdnn

#ifdef def_device
#undef __device__
#undef __host__
#endif

// vim: ft=cpp syntax=cpp.doxygen