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

/*******************************************************************************
* Copyright 2026 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 <mutex>
#include <unordered_map>

#include "common/c_types_map.hpp"
#include "common/utils.hpp"

#include "gpu/intel/compute/kernel.hpp"
#include "gpu/intel/compute/kernel_ctx.hpp"
#include "gpu/intel/engine.hpp"
#include "gpu/intel/stream.hpp"

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

static status_t get_cached_kernel(
        intel::engine_t *engine, compute::kernel_t &kernel) {
    static std::unordered_map<engine_id_t, compute::kernel_t> cache;
    static std::mutex mutex;

    std::lock_guard<std::mutex> lock(mutex);
    auto it = cache.find(engine->engine_id());
    if (it != cache.end()) {
        kernel = it->second;
        return status::success;
    }

    compute::kernel_ctx_t ctx;
    std::vector<compute::kernel_t> kernels;
    CHECK(engine->create_kernels(&kernels, {"fill_random"}, ctx));
    kernel = cache.emplace(engine->engine_id(), kernels[0]).first->second;
    return status::success;
}

status_t fill_random(impl::stream_t *stream, size_t size,
        impl::memory_t *memory, int buffer_index, uint32_t seed) {
    if (size == 0) return status::success;

    auto *intel_stream = utils::downcast<intel::stream_t *>(stream);
    auto *intel_engine = utils::downcast<intel::engine_t *>(stream->engine());
    compute::kernel_t kernel;
    CHECK(get_cached_kernel(intel_engine, kernel));

    // Each subgroup (16 work-items) processes 256 bytes (16 * 4 * sizeof(uint)).
    static constexpr size_t subgroup_size = 16;
    static constexpr size_t bytes_per_subgroup
            = subgroup_size * 4 * sizeof(uint32_t);
    size_t num_subgroups = utils::div_up(size, bytes_per_subgroup);
    compute::nd_range_t nd_range({num_subgroups * subgroup_size, 1, 1});
    compute::kernel_arg_list_t arg_list;
    arg_list.set(0, *memory->memory_storage(buffer_index));
    arg_list.set(1, seed);
    arg_list.set(2, static_cast<uint64_t>(size));

    CHECK(kernel.parallel_for(*stream, nd_range, arg_list,
            intel_stream->ctx().get_deps(), intel_stream->ctx().get_deps()));
    return status::success;
}

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

extern "C" dnnl::impl::status_t DNNL_API dnnl_impl_gpu_fill_random(
        dnnl::impl::stream_t *stream, size_t size, dnnl::impl::memory_t *memory,
        int buffer_index, uint32_t seed) {
    return dnnl::impl::gpu::intel::fill_random(
            stream, size, memory, buffer_index, seed);
}