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

/**
 * \file dnn/src/x86/elemwise/avx_util/avx_util.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/arch.h"
#include "megdnn/basic_types.h"

#include <cstddef>

namespace megdnn {
namespace x86 {
namespace detail {
/*
 *  Case 1. Unary Optrs
 */

#define AVX_ELEMENT_OPTR_UNARY(optr_type)                                            \
    void avx_element_##optr_type(                                                    \
            const TensorND& src_tensor, const TensorND& dst_tensor);                 \
    void avx_element_##optr_type(size_t tsize, const float* src_ptr, float* dst_ptr) \
            MEGDNN_ATTRIBUTE_TARGET("avx");

AVX_ELEMENT_OPTR_UNARY(sigmoid)
AVX_ELEMENT_OPTR_UNARY(exp)
AVX_ELEMENT_OPTR_UNARY(tanh)
AVX_ELEMENT_OPTR_UNARY(fast_tanh)
AVX_ELEMENT_OPTR_UNARY(relu)
AVX_ELEMENT_OPTR_UNARY(set)

/*
 * Case 2. Binary Optrs
 * 2.1 src0 has the same size with src1.
 * 2.2 src1 is a scalar.
 * 2.3 shape of src1 is {1, C, 1, 1}.
 * 2.4 some other optrs only for dtype float32
 */

// Case 2.1 src0 has the same size with src1.

#define AVX_ELEMENT_OPTR_BINARY(optr_type)                                  \
    void avx_element_##optr_type(                                           \
            const TensorND& src0_tensor, const TensorND& src1_tensor,       \
            const TensorND& dst_tensor);                                    \
    void avx_element_##optr_type(                                           \
            size_t tsize, float* src0_ptr, float* src1_ptr, float* dst_ptr) \
            MEGDNN_ATTRIBUTE_TARGET("avx");

AVX_ELEMENT_OPTR_BINARY(add)
AVX_ELEMENT_OPTR_BINARY(bias_sigmoid)
AVX_ELEMENT_OPTR_BINARY(bias_relu)
AVX_ELEMENT_OPTR_BINARY(bias_tanh)

// Case 2.2 src1 is a scalar.

void avx_element_add_scalar(
        const size_t tsize, float* src_ptr, float* dst_ptr, const float bias)
        MEGDNN_ATTRIBUTE_TARGET("avx");

// Case 2.3 Shape of tensor src1 is 1C11.
#define AVX_ELEMENT_OPTR_BINARY_1C11(optr_type)                            \
    void avx_element_##optr_type(                                          \
            size_t batch_size, size_t channel_size, size_t channel_stride, \
            const TensorND& src0_tensor, const TensorND& src1_tensor,      \
            const TensorND& dst_tensor);                                   \
    void avx_element_##optr_type(                                          \
            size_t batch_size, size_t channel_size, size_t channel_stride, \
            float* src1_ptr, float* src2_ptr, float* dst_ptr)              \
            MEGDNN_ATTRIBUTE_TARGET("avx");

AVX_ELEMENT_OPTR_BINARY_1C11(add_1c11)
AVX_ELEMENT_OPTR_BINARY_1C11(bias_sigmoid_1c11)
AVX_ELEMENT_OPTR_BINARY_1C11(bias_relu_1c11)

// void avx_element_add_1c11(const TensorND &src1_tensor,
//                                const TensorND &src2_tensor,
//                                const TensorND &dst_tensor)
//    MEGDNN_ATTRIBUTE_TARGET("avx");
//
// void avx_element_add_1c11 (
//    size_t batch_size, size_t channel_size, size_t channel_stride,
//    float *src1_ptr, float *src2_ptr, float *dst_ptr)
//    MEGDNN_ATTRIBUTE_TARGET("avx");
//
// void avx_element_bias_relu_1c11(TensorND dst_tensor,
//        TensorND bias_tensor) MEGDNN_ATTRIBUTE_TARGET("avx");
//
// void avx_element_bias_sigmoid_1c11(TensorND dst_tensor,
//        TensorND bias_tensor) MEGDNN_ATTRIBUTE_TARGET("avx");

/*
 * Case 3. Ternary Optrs
 * 3.1 src0, src1 and src2 has the same size.
 * 3.2 src0, src1 has the same size, src2 is a scalar
 * 3.3 shape of src0 and src2 is (1,C,1,1).
 */

#define AVX_ELEMENT_OPTR_TERNARY(optr_type)                                  \
    void avx_element_##optr_type(                                            \
            const TensorND& src0_tensor, const TensorND& src1_tensor,        \
            const TensorND& src2_tensor, const TensorND& dst_tensor);        \
    void avx_element_##optr_type(                                            \
            size_t tsize, float* src0_ptr, float* src1_ptr, float* src2_ptr, \
            float* dst_ptr) MEGDNN_ATTRIBUTE_TARGET("avx");

// Case 3.1 src0, src1 and src2 has the same size.
AVX_ELEMENT_OPTR_TERNARY(fma3)

// Case 3.2 src1 is a scalar.
AVX_ELEMENT_OPTR_TERNARY(fma3_scalar)

// Case 3.3 shape of src0 and src2 is (1,C,1,1).
#define AVX_ELEMENT_OPTR_TERNARY_1C11(optr_type)                               \
    void avx_element_##optr_type(                                              \
            size_t batch_size, size_t channel_size, size_t channel_stride,     \
            const TensorND& src0_tensor, const TensorND& src1_tensor,          \
            const TensorND& src2_tensor, const TensorND& dst_tensor);          \
    void avx_element_##optr_type(                                              \
            size_t batch_size, size_t channel_size, size_t channel_stride,     \
            float* src0_ptr, float* src1_ptr, float* src2_ptr, float* dst_ptr) \
            MEGDNN_ATTRIBUTE_TARGET("avx");

AVX_ELEMENT_OPTR_TERNARY_1C11(fma3_1c11)

/*
 * Other Optrs
 */

// src1 is contiguous, with shape [batch_size, channel_size, channel_stride]
// src2 is contiguous, with shape [1, channel_size, 1]
void avx_element_add_in_a_channel(
        float* output_ptr, float* bias_ptr, size_t channel_size, size_t channel_stride)
        MEGDNN_ATTRIBUTE_TARGET("avx");

}  // namespace detail
}  // namespace x86
}  // namespace megdnn