-
Notifications
You must be signed in to change notification settings - Fork 555
Expand file tree
/
Copy pathkernel_cache.cpp
More file actions
146 lines (129 loc) · 4.82 KB
/
Copy pathkernel_cache.cpp
File metadata and controls
146 lines (129 loc) · 4.82 KB
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
/*******************************************************
* Copyright (c) 2020, ArrayFire
* All rights reserved.
*
* This file is distributed under 3-clause BSD license.
* The complete license agreement can be obtained at:
* http://arrayfire.com/licenses/BSD-3-Clause
********************************************************/
#if !defined(AF_CPU) && !defined(AF_ONEAPI)
#include <common/compile_module.hpp>
#include <common/deterministicHash.hpp>
#include <common/kernel_cache.hpp>
#include <device_manager.hpp>
#include <platform.hpp>
#include <nonstd/span.hpp>
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <vector>
using detail::Kernel;
using detail::Module;
using nonstd::span;
using std::array;
using std::back_inserter;
using std::shared_lock;
using std::shared_timed_mutex;
using std::string;
using std::to_string;
using std::transform;
using std::unique_lock;
using std::unordered_map;
using std::vector;
namespace arrayfire {
namespace common {
using ModuleMap = unordered_map<size_t, Module>;
shared_timed_mutex& getCacheMutex(const int device) {
static shared_timed_mutex mutexes[detail::DeviceManager::MAX_DEVICES];
return mutexes[device];
}
ModuleMap& getCache(const int device) {
static ModuleMap* caches =
new ModuleMap[detail::DeviceManager::MAX_DEVICES];
return caches[device];
}
Module findModule(const int device, const size_t& key) {
shared_lock<shared_timed_mutex> readLock(getCacheMutex(device));
auto& cache = getCache(device);
auto iter = cache.find(key);
if (iter != cache.end()) { return iter->second; }
return Module{};
}
Kernel getKernel(const string& kernelName, span<const common::Source> sources,
span<const TemplateArg> targs, span<const string> options,
const bool sourceIsJIT) {
string tInstance = kernelName;
#if defined(AF_CUDA)
auto targsIt = targs.begin();
auto targsEnd = targs.end();
if (targsIt != targsEnd) {
tInstance += '<' + targsIt->_tparam;
while (++targsIt != targsEnd) { tInstance += ',' + targsIt->_tparam; }
tInstance += '>';
}
#else
UNUSED(targs);
#endif
// The JIT kernel uses the hashing of the kernelName (tInstance) only to
// speed up to search for its cached kernel. All the other kernels have the
// full source code linked in, and will hash the full code + options
// instead.
size_t moduleKeyCache = 0;
if (sourceIsJIT) {
moduleKeyCache = deterministicHash(tInstance);
} else {
moduleKeyCache = (sources.size() == 1 && sources[0].hash)
? sources[0].hash
: deterministicHash(sources);
moduleKeyCache = deterministicHash(options, moduleKeyCache);
#if defined(AF_CUDA)
moduleKeyCache = deterministicHash(tInstance, moduleKeyCache);
#endif
}
const int device = detail::getActiveDeviceId();
Module currModule = findModule(device, moduleKeyCache);
if (!currModule) {
// When saving on disk, the moduleKeyDisk has to correspond with the
// full code + optinos (in all circumstances). A recalculation for JIT
// is necessary, while for the others we can reuse the moduleKeyCache.
size_t moduleKeyDisk = 0;
if (sourceIsJIT) {
moduleKeyDisk = (sources.size() == 1 && sources[0].hash)
? sources[0].hash
: deterministicHash(sources);
moduleKeyDisk = deterministicHash(options, moduleKeyDisk);
#if defined(AF_CUDA)
moduleKeyDisk = deterministicHash(tInstance, moduleKeyDisk);
#endif
} else {
moduleKeyDisk = moduleKeyCache;
}
currModule =
loadModuleFromDisk(device, to_string(moduleKeyDisk), sourceIsJIT);
if (!currModule) {
vector<string> sources_str;
for (const auto& s : sources) {
sources_str.push_back({s.ptr, s.length});
}
currModule = compileModule(to_string(moduleKeyDisk), sources_str,
options, array{tInstance}, sourceIsJIT);
}
unique_lock<shared_timed_mutex> writeLock(getCacheMutex(device));
auto& cache = getCache(device);
auto iter = cache.find(moduleKeyCache);
if (iter == cache.end()) {
// If not found, this thread is the first one to compile
// this kernel. Keep the generated module.
Module mod = currModule;
getCache(device).emplace(moduleKeyCache, mod);
} else {
currModule.unload(); // dump the current threads extra
// compilation
currModule = iter->second;
}
}
return getKernel(currModule, tInstance, sourceIsJIT);
}
} // namespace common
} // namespace arrayfire
#endif