onednn-src 0.1.13

Source of oneAPI Deep Neural Network Library (oneDNN)
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

/*******************************************************************************
* Copyright 2019 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#include "gpu/intel/include/dispatch.h"
#include "gpu/intel/include/io.h"
#include "gpu/intel/include/types.h"

#if IS_FWD == 1
KERNEL_ATTR
__kernel void ref_lrn_fwd(__global const DATA_T *src,
#if IS_TRAINING == 1
        __global DEF_ACC_DATA_T *ws,
#endif
        __global DATA_T *dst) {
    const uint mb = GWS_GET_MB();
    const uint ic = GWS_GET_IC();
    const uint id = GWS_GET_ID();
    const uint ih = GWS_GET_IH();
    const uint iw = GWS_GET_IW();

    const uint src_index = SRC_OFF(mb, ic, id, ih, iw);
    const uint dst_index = DST_OFF(mb, ic, id, ih, iw);

    DEF_ACC_DATA_T sum = 0.0f;

#if ACROSS_CHANNEL
    for (int j = 0; j < LOCAL_SIZE; j++) {
        const int z_idx = (j + ic - PADDING);
        bool zero = (z_idx < 0 || z_idx >= IC);
        DEF_ACC_DATA_T val = 0;
        if (!zero) load(&val, src + SRC_OFF(mb, z_idx, id, ih, iw));
        sum += val * val;
    }
#else

    const int d = (int)id - PADDING;
    const int h = (int)ih - PADDING;
    const int w = (int)iw - PADDING;

    const int d_start = max(d, 0);
    const int h_start = max(h, 0);
    const int w_start = max(w, 0);
    const int d_end = min(d + LOCAL_SIZE, ID);
    const int h_end = min(h + LOCAL_SIZE, IH);
    const int w_end = min(w + LOCAL_SIZE, IW);

    for (int k = d_start; k < d_end; ++k) {
        for (int j = h_start; j < h_end; ++j) {
            for (int i = w_start; i < w_end; ++i) {
                DEF_ACC_DATA_T val = load(val, src + SRC_OFF(mb, ic, k, j, i));
                sum += val * val;
            }
        }
    }

#endif

    const DEF_ACC_DATA_T num_elements_div = NUM_ELEMENTS_DIV;
    const DEF_ACC_DATA_T base = (DEF_ACC_DATA_T)LRN_K
            + (DEF_ACC_DATA_T)LRN_ALPHA * sum * num_elements_div;
    const DEF_ACC_DATA_T normalization_factor
            = native_powr(base, (DEF_ACC_DATA_T)(-LRN_BETA));

    const DEF_ACC_DATA_T val = load(val, src + src_index);
    const DEF_ACC_DATA_T normres = val * normalization_factor;
#if IS_TRAINING == 1
    ws[dst_index] = base;
#endif
    write(dst + dst_index, normres);
}
#endif

#if IS_BWD == 1
KERNEL_ATTR
__kernel void ref_lrn_bwd(__global const DATA_T *src,
        __global const DATA_T *diff_dst, __global DEF_ACC_DATA_T *ws,
        __global DATA_T *diff_src) {
    const uint mb = GWS_GET_MB();
    const uint ic = GWS_GET_IC();
    const uint id = GWS_GET_ID();
    const uint ih = GWS_GET_IH();
    const uint iw = GWS_GET_IW();

    const uint src_index = SRC_OFF(mb, ic, id, ih, iw);
    const uint dst_index = DST_OFF(mb, ic, id, ih, iw);
    const DEF_ACC_DATA_T num_elements_div = NUM_ELEMENTS_DIV;
    DEF_ACC_DATA_T B = 0;

#if ACROSS_CHANNEL
    for (int j = 0; j < LOCAL_SIZE; j++) {
        const int z_idx = (j + ic - PADDING);
        bool zero = (z_idx < 0 || z_idx >= IC);
        if (!zero) {
            DEF_ACC_DATA_T val
                    = load(val, src + SRC_OFF(mb, z_idx, id, ih, iw));
            DEF_ACC_DATA_T omega
                    = load(omega, ws + SRC_OFF(mb, z_idx, id, ih, iw));
            DEF_ACC_DATA_T tmp = (DEF_ACC_DATA_T)1.0f
                    / native_powr(omega, (DEF_ACC_DATA_T)LRN_BETA + 1);
            B += tmp * val * load(B, diff_dst + DST_OFF(mb, z_idx, id, ih, iw));
        }
    }
#else

    const int d = (int)id - PADDING;
    const int h = (int)ih - PADDING;
    const int w = (int)iw - PADDING;
    const int d_start = max(d, 0);
    const int h_start = max(h, 0);
    const int w_start = max(w, 0);
    const int d_end = min(d + LOCAL_SIZE, ID);
    const int h_end = min(h + LOCAL_SIZE, IH);
    const int w_end = min(w + LOCAL_SIZE, IW);
    for (int k = d_start; k < d_end; ++k) {
        for (int j = h_start; j < h_end; ++j) {
            for (int i = w_start; i < w_end; ++i) {
                int data_off = SRC_OFF(mb, ic, k, j, i);
                DEF_ACC_DATA_T val = load(val, src + data_off);
                DEF_ACC_DATA_T omega = ws[data_off];
                DEF_ACC_DATA_T tmp = (DEF_ACC_DATA_T)1.0f
                        / native_powr(omega, (DEF_ACC_DATA_T)(LRN_BETA + 1));
                B += tmp * val * load(B, diff_dst + data_off);
            }
        }
    }
#endif
    DEF_ACC_DATA_T A = native_powr(ws[src_index], (DEF_ACC_DATA_T)-LRN_BETA)
            * load(A, diff_dst + dst_index);

    A -= load(A, src + src_index) * 2 * (DEF_ACC_DATA_T)LRN_ALPHA
            * (DEF_ACC_DATA_T)LRN_BETA * num_elements_div * B;
    write(diff_src + src_index, A);
}
#endif