1 #ifndef CAFFE2_UTILS_MKL_CONTEXT_H_ 2 #define CAFFE2_UTILS_MKL_CONTEXT_H_ 8 #include "caffe2/core/context.h" 23 explicit MKLContext(
const DeviceOption& option)
25 option.has_random_seed() ? option.random_seed()
27 CAFFE_ENFORCE_EQ(option.device_type(), MKLDNN);
32 inline void SwitchToDevice(
int = 0) {}
34 inline void WaitEvent(
const Event& ev) {
35 ev.Wait(MKLDNN,
this);
38 inline void Record(
Event* ev,
const char* err_msg =
nullptr)
const {
39 CAFFE_ENFORCE(ev,
"Event must not be null.");
40 ev->Record(MKLDNN,
this, err_msg);
43 inline void FinishDeviceComputation() {}
45 inline std::mt19937& RandGenerator() {
46 if (!random_generator_.get()) {
47 random_generator_.reset(
new std::mt19937(random_seed_));
49 return *random_generator_.get();
52 inline static std::pair<void*, MemoryDeleter> New(
size_t nbytes) {
53 return GetCPUAllocator()->New(nbytes);
57 template <
class SrcContext,
class DstContext>
58 inline void CopyBytes(
size_t nbytes,
const void* src,
void* dst);
60 template <
typename T,
class SrcContext,
class DstContext>
61 inline void Copy(
size_t n,
const T* src, T* dst) {
62 if (std::is_fundamental<T>::value) {
63 CopyBytes<SrcContext, DstContext>(
65 static_cast<const void*>(src),
66 static_cast<void*
>(dst));
68 for (
int i = 0; i < n; ++i) {
74 template <
class SrcContext,
class DstContext>
76 CopyItems(
const TypeMeta& meta,
size_t n,
const void* src,
void* dst) {
78 meta.
copy()(src, dst, n);
80 CopyBytes<SrcContext, DstContext>(n * meta.
itemsize(), src, dst);
85 static bool HasAsyncPartDefault() {
89 static bool SupportsAsyncScheduling() {
93 static bool IsStreamFree(
const DeviceOption& ,
int ) {
99 int random_seed_{1701};
100 std::unique_ptr<std::mt19937> random_generator_;
104 inline void MKLContext::CopyBytes<MKLContext, MKLContext>(
108 memcpy(dst, src, nbytes);
112 inline void MKLContext::CopyBytes<CPUContext, MKLContext>(
116 memcpy(dst, src, nbytes);
120 inline void MKLContext::CopyBytes<MKLContext, CPUContext>(
124 memcpy(dst, src, nbytes);
128 #endif // CAFFE2_UTILS_MKL_CONTEXT_H_ The MKL Context, which is largely the same as the CPUContext.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
uint32_t RandomNumberSeed()
A function to generate a random number seed that is unique in a best-effort basis, using an ever-incrementing seed and the current time.