Caffe2 - C++ API
A deep learning, cross platform ML framework
conv_op_shared.cc
1 
17 #include "conv_op_shared.h"
18 #include "caffe2/core/context.h"
19 #include "caffe2/core/flags.h"
20 #include "caffe2/core/workspace.h"
21 
22 CAFFE2_DEFINE_bool(
23  caffe2_force_shared_col_buffer,
24  false,
25  "Always use the shared col buffer");
26 
27 namespace caffe2 {
28 
29 template <>
30 void createSharedBuffer<CPUContext>(Workspace* ws) {
31  auto* mutexPtr = ws->CreateBlob("__CAFFE2_SHARED_CONV_BUFFER_CPU_MUTEX__")
32  ->GetMutable<std::unique_ptr<std::mutex>>();
33  mutexPtr->reset(new std::mutex());
34  ws->CreateBlob("__CAFFE2_SHARED_CONV_BUFFER_CPU__");
35 }
36 
37 template <>
38 void runWithSharedBuffer(
39  Workspace* ws,
40  std::function<void(Tensor<CPUContext>* buffer)> f) {
41  auto* mutexBlob = ws->GetBlob("__CAFFE2_SHARED_CONV_BUFFER_CPU_MUTEX__");
42  CAFFE_ENFORCE(mutexBlob, "Must call createSharedBuffer() first");
43 
44  auto* mutexPtr = mutexBlob->GetMutable<std::unique_ptr<std::mutex>>();
45  std::lock_guard<std::mutex> g(**mutexPtr);
46  auto* buffer =
47  ws->GetBlob("__CAFFE2_SHARED_CONV_BUFFER_CPU__")->GetMutable<TensorCPU>();
48  f(buffer);
49 }
50 }
Copyright (c) 2016-present, Facebook, Inc.
Copyright (c) 2016-present, Facebook, Inc.
Copyright (c) 2016-present, Facebook, Inc.