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

/*******************************************************************************
* Copyright 2022 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/sum/many_inputs.hpp"

namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace sum {

status_t many_inputs_t::execute(const exec_ctx_t &ctx) const {
    auto &output = CTX_OUT_STORAGE(DNNL_ARG_DST);
    const int num_arrs = pd()->n_inputs()
            - 1; // input0 is copied over to output. Accumulate the rest.
    const memory_desc_wrapper o_d(pd()->dst_md());
    const size_t nelems = o_d.nelems(true);
    compute::kernel_arg_list_t arg_list;

    int num_batches = utils::div_up(num_arrs, max_num_arrs);

    auto &input0 = CTX_IN_STORAGE(DNNL_ARG_MULTIPLE_SRC);
    const bool is_inplace = (output.data_handle() == input0.data_handle());
    if (!is_inplace) {
        auto *intel_stream = utils::downcast<intel::stream_t *>(ctx.stream());
        CHECK(intel_stream->copy(input0, output, o_d.size(),
                intel_stream->ctx().get_deps(),
                intel_stream->ctx().get_deps()));
    }
    status_t status;
    for (int batch_iter = 0; batch_iter < num_batches; batch_iter++) {
        int kernel_num_arrs = max_num_arrs;
        if ((batch_iter == num_batches - 1) && (num_arrs % max_num_arrs)) {
            kernel_num_arrs = num_arrs % max_num_arrs;
        }
        for (int a = 0; a < max_num_arrs; ++a) {
            if (a < kernel_num_arrs) {
                auto &input = CTX_IN_STORAGE(DNNL_ARG_MULTIPLE_SRC + 1 + a
                        + batch_iter * max_num_arrs);
                arg_list.set(a, input);
            } else
                arg_list.set(a, memory_storage_t::empty_storage());
        }
        arg_list.set(many_inputs_t::max_num_arrs, output);
        arg_list.set(
                many_inputs_t::max_num_arrs + 1, CTX_GPU_RES_STORAGE(SCALES_));

        const size_t total_width = nelems * kernel_num_arrs;
        const size_t lws = utils::rnd_dn(256, kernel_num_arrs);

        compute::nd_range_t nd_range({utils::rnd_up(total_width, lws)}, {lws});
        if (batch_iter == num_batches - 1) {
            status = parallel_for(ctx, nd_range, kernel_, arg_list);
        } else {
            status = parallel_for(ctx, nd_range, batched_kernel_, arg_list);
        }
        if (status != dnnl_success) return status;
    }
    return status;
}

} // namespace sum
} // namespace intel
} // namespace gpu
} // namespace impl
} // namespace dnnl