Caffe2 - C++ API
A deep learning, cross platform ML framework
snpe_op.cc
1 
17 #include "caffe2/core/context.h"
18 #include "caffe2/core/logging.h"
19 #include "caffe2/core/operator.h"
20 #include "caffe2/core/timer.h"
21 #include "snpe_ffi.h"
22 #include <dlfcn.h>
23 
24 namespace caffe2 {
25 
26 template <typename T>
27 using deleted_unique_ptr = std::unique_ptr<T, std::function<void(T*)>>;
28 
29 class SNPEOp final : public Operator<CPUContext> {
30  public:
31  SNPEOp(const OperatorDef& def, Workspace* ws) : Operator<CPUContext>(def, ws),
32  model_buffer_(OperatorBase::GetSingleArgument<string>("model_buffer", "")),
33  input_name_(OperatorBase::GetSingleArgument<string>("input_name", "data"))
34  {
35  CAFFE_ENFORCE(gSNPELocation() != "", "SNPE library \"", gSNPELocation(), "\" does not exist.");
36  std::ostringstream snpe_ffi;
37  snpe_ffi << gSNPELocation() << "/" << snpe_ffi_so;
38  handle_ = deleted_unique_ptr<void>(dlopen(snpe_ffi.str().c_str(), RTLD_LAZY), [](void* handle) {
39  if (handle) {
40  dlclose(handle);
41  }
42  });
43  if (!handle_.get()) {
44  std::cerr << dlerror() << std::endl;
45  }
46 
47  OPERATOR_NEEDS_FEATURE(handle_.get(), "Couldn't find ", snpe_ffi.str());
48 
49 #define X(n) \
50  dlerror(); \
51  auto* n##_f = (decltype(&n))dlsym(handle_.get(), #n); \
52  OPERATOR_NEEDS_FEATURE(n##_f, dlerror());
53 
54  {
55  X(snpe_has_gpu);
56  X(snpe_create);
57  X(snpe_destroy);
58  X(snpe_get_input_dims);
59  X(snpe_run);
60  X(snpe_copy_output_to);
61  }
62 
63  X(snpe_has_gpu);
64  OPERATOR_NEEDS_FEATURE(snpe_has_gpu_f(), "No GPU found, cannot use SNPE.");
65 
66  X(snpe_create)
67 #undef X
68 
69 // Redefine to use CAFFE_ENFORCE instead of OPERATOR_NEEDS_FEATURE.
70 
71 #define X(n) \
72  dlerror(); \
73  auto* n##_f = (decltype(&n))dlsym(handle_.get(), #n); \
74  CAFFE_ENFORCE(n##_f, dlerror());
75 
76  CAFFE_ENFORCE(def.input_size(), "No inputs.");
77  if (input_name_ == "") {
78  input_name_ = def.input().Get(0);
79  }
80  ctx_ = deleted_unique_ptr<void>(snpe_create_f(reinterpret_cast<const unsigned char *>(model_buffer_.data()),
81  model_buffer_.length(), input_name_.c_str()), [this](void* ctx) {
82  if (ctx) {
83  X(snpe_destroy);
84  snpe_destroy_f(ctx);
85  }
86  });
87  }
88 
89  bool RunOnDevice() override {
90  CAFFE_ENFORCE(gSNPELocation() != "", "SNPE library was never loaded.");
91 
92  X(snpe_get_input_dims);
93  size_t const* dims;
94  size_t dimSize;
95  snpe_get_input_dims_f(ctx_.get(), &dims, &dimSize);
96  if (Input(0).ndim() != dimSize) {
97  if (dimSize == 3 && dimSize == Input(0).ndim() - 1 && Input(0).dim32(0) == 1) {
98  const int C = Input(0).dim32(1);
99  const int H = Input(0).dim32(2);
100  const int W = Input(0).dim32(3);
101  if (dims[0] != C ||
102  dims[1] != H ||
103  dims[2] != W) {
104  CAFFE_THROW("Input size must match what SNPE expects, which in this case is: ",
105  dims[0], " ", dims[1], " ", dims[2]);
106  }
107  } else {
108  CAFFE_THROW("SNPE input dimensions are not compatible.");
109  }
110  } else {
111  for (auto i = 0; i < Input(0).ndim(); ++i) {
112  CAFFE_ENFORCE_EQ(dims[i], Input(0).dim32(i), "SNPE input dimension is not compatible.");
113  }
114  }
115 
116  X(snpe_run);
117  CAFFE_ENFORCE(ctx_.get(), "SNPE context doesn't exist.");
118  snpe_run_f(ctx_.get(), Input(0).data<float>(), Input(0).size(), &dims, &dimSize);
119 
120  std::vector<int64_t> outputDims(dimSize + 1);
121  outputDims[0] = 1;
122  for (auto i = 0; i < dimSize; ++i) {
123  outputDims[i+1] = dims[i];
124  };
125 
126  Output(0)->Resize(outputDims);
127  X(snpe_copy_output_to);
128  snpe_copy_output_to_f(ctx_.get(), Output(0)->mutable_data<float>());
129 
130  CAFFE_ENFORCE(Output(0)->data<float>(), "nullptr where output should be!\n");
131  return true;
132  }
133 
134  private:
135  string model_buffer_;
136  string input_name_;
137  deleted_unique_ptr<void> handle_;
138  // needs to be destroyed *before* handle_
139  deleted_unique_ptr<void> ctx_;
140 };
141 
142 REGISTER_CPU_OPERATOR(SNPE, SNPEOp);
143 }
144 
145 #undef X
const T * data() const
Returns a typed pointer of the underlying storage.
Definition: tensor.h:500
int dim32(const int i) const
Returns the i-th dimension of the tensor in int.
Definition: tensor.h:673
TIndex size() const
Returns the size (i.e.
Definition: tensor.h:609
T * mutable_data()
Returns a typed pointer of the underlying storage.
Definition: tensor.h:594
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
void Resize(Ts...dim_source)
Resizes a tensor.
Definition: tensor.h:304
Copyright (c) 2016-present, Facebook, Inc.
int ndim() const
Returns the number of dimensions of the data.
Definition: tensor.h:605