Add onnxifi interface for set/get options (#52388)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/52388
Pull Request resolved: https://github.com/pytorch/glow/pull/5364
This allows us to change global variables through onnxifi calls. And add python bindings along with it. Note that we supply a dummy backend_id as it's not needed by glow due to setting being global.
#codemod
Test Plan:
```
buck test mode/dev //glow/fb/test:test_onnxifi_optionnnpi
```
Reviewed By: jfix71, khabinov
Differential Revision: D26481652
fbshipit-source-id: 19b8201c77f653cf7d93ad68760aa7fb5ec45ff4
diff --git a/caffe2/opt/onnxifi_transformer.cc b/caffe2/opt/onnxifi_transformer.cc
index 79634bb..525f564 100644
--- a/caffe2/opt/onnxifi_transformer.cc
+++ b/caffe2/opt/onnxifi_transformer.cc
@@ -483,6 +483,71 @@
new_net.Swap(net);
}
+OnnxifiOptionHelper::OnnxifiOptionHelper() {
+ lib_ = onnx::initOnnxifiLibrary();
+ CAFFE_ENFORCE(lib_, "Cannot initialize ONNXIFI library");
+}
+
+bool OnnxifiOptionHelper::setOnnxifiOption(
+ const std::string& option,
+ const std::string& value) {
+#ifdef ONNXIFI_ENABLE_EXT
+ onnxStatus (*onnxSetOptionFunctionPointer)(
+ const char* optionName, const char* optionValue) = nullptr;
+ union {
+ onnxExtensionFunctionPointer p;
+ decltype(onnxSetOptionFunctionPointer) set;
+ } u{};
+ onnxBackendID backend_id = nullptr;
+ if (lib_->onnxGetExtensionFunctionAddress(
+ backend_id, "onnxSetOptionFunction", &u.p) !=
+ ONNXIFI_STATUS_SUCCESS) {
+ LOG(ERROR) << "Cannot find onnxSetOptionFunction";
+ return false;
+ } else {
+ onnxSetOptionFunctionPointer = u.set;
+ }
+ if (onnxSetOptionFunctionPointer != nullptr &&
+ (*onnxSetOptionFunctionPointer)(option.c_str(), value.c_str()) ==
+ ONNXIFI_STATUS_SUCCESS) {
+ return true;
+ }
+#endif
+ return false;
+}
+
+std::string OnnxifiOptionHelper::getOnnxifiOption(const std::string& option) {
+#ifdef ONNXIFI_ENABLE_EXT
+ onnxStatus (*onnxGetOptionFunctionPointer)(
+ const char* optionName, char* optionValue, size_t* optionValueLength) =
+ nullptr;
+ union {
+ onnxExtensionFunctionPointer p;
+ decltype(onnxGetOptionFunctionPointer) get;
+ } u{};
+ onnxBackendID backend_id = nullptr;
+ if (lib_->onnxGetExtensionFunctionAddress(
+ backend_id, "onnxGetOptionFunction", &u.p) !=
+ ONNXIFI_STATUS_SUCCESS) {
+ LOG(ERROR) << "Cannot find onnxGetOptionFunction";
+ return "";
+ } else {
+ onnxGetOptionFunctionPointer = u.get;
+ }
+
+ constexpr size_t ll = 1024;
+ char buf[ll];
+ size_t len = ll;
+ if (onnxGetOptionFunctionPointer != nullptr &&
+ (*onnxGetOptionFunctionPointer)(option.c_str(), buf, &len) ==
+ ONNXIFI_STATUS_SUCCESS) {
+ return std::string(buf, len);
+ }
+#endif
+
+ return "";
+}
+
OnnxifiTransformer::OnnxifiTransformer(const OnnxifiTransformerOptions& opts)
: BackendTransformerBase(), opts_(opts) {
lib_ = onnx::initOnnxifiLibrary();
@@ -537,7 +602,7 @@
const std::vector<std::string>& external_inputs,
const std::vector<std::string>& external_outputs,
const ShapeInfoMap& shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs) {
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs) {
OperatorDef op;
op.set_type("Onnxifi");
auto* onnx_model_arg = op.add_arg();
@@ -577,7 +642,10 @@
if (!initialization_list.count(input)) {
const auto it = shape_hints_max_bs.find(input);
CAFFE_ENFORCE(
- it != shape_hints_max_bs.end(), "Input shape for ", input, " not found");
+ it != shape_hints_max_bs.end(),
+ "Input shape for ",
+ input,
+ " not found");
const auto& info = it->second;
if (info.getDimType(0) == TensorBoundShape_DimType_BATCH &&
getBlob1stDimSize(info) == max_batch_size) {
@@ -610,7 +678,8 @@
// Add output size hints per batch size
if (canPassOutputShapeHintsPerBs(op, shape_hints_per_bs)) {
- VLOG(2) << "Passing in output shape hints for batch sizes in [1, " << opts_.bound_shape_spec.max_batch_size << ")";
+ VLOG(2) << "Passing in output shape hints for batch sizes in [1, "
+ << opts_.bound_shape_spec.max_batch_size << ")";
AddArgument("use_passed_output_shapes", 1, &op);
for (int bs = 1; bs < opts_.bound_shape_spec.max_batch_size; ++bs) {
@@ -625,9 +694,11 @@
const auto& output_name = op.output(output_idx);
const auto& shape_hint = shape_hints.find(output_name)->second;
if (!shape_hint.is_quantized) {
- output_shape_arg->mutable_tensors()->Add()->CopyFrom(wrapShapeInfoIntoTensorProto(output_name, shape_hint));
+ output_shape_arg->mutable_tensors()->Add()->CopyFrom(
+ wrapShapeInfoIntoTensorProto(output_name, shape_hint));
} else {
- output_shape_arg->mutable_qtensors()->Add()->CopyFrom(wrapShapeInfoIntoQTensorProto(output_name, shape_hint));
+ output_shape_arg->mutable_qtensors()->Add()->CopyFrom(
+ wrapShapeInfoIntoQTensorProto(output_name, shape_hint));
}
}
}
@@ -663,7 +734,7 @@
const caffe2::NetDef& net,
const std::unordered_set<std::string>& weights_in_ws,
const ShapeInfoMap& shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs) {
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs) {
int onnxifi_op_id = onnxifi_op_id_;
if (opts_.debug) {
WriteProtoToTextFile(
@@ -802,7 +873,7 @@
Workspace* ws,
onnx::OnnxExporter* exporter,
ShapeInfoMap* shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs) {
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs) {
if (opts_.min_ops > net.op_size()) {
return net;
}
@@ -1222,22 +1293,24 @@
const std::unordered_set<std::string>& weights,
const std::unordered_set<int>& blocklisted_ops,
const ShapeInfoMap& shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs) {
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs) {
onnxBackendID backend_id = backend_ids_[idx_];
- auto c2_supports = [this,
- &shape_hints_max_bs,
- &blocklisted_ops,
- backend_id,
- &weights](const caffe2::OperatorDef& op) {
- return supportOpC2(op, shape_hints_max_bs, weights, blocklisted_ops, backend_id);
- };
-
- auto c2_converter =
- [this, &weights, &shape_hints_max_bs, &shape_hints_per_bs](const caffe2::NetDef& net) {
- return SubnetToOnnxifiOpViaC2(net, weights, shape_hints_max_bs, shape_hints_per_bs);
+ auto c2_supports =
+ [this, &shape_hints_max_bs, &blocklisted_ops, backend_id, &weights](
+ const caffe2::OperatorDef& op) {
+ return supportOpC2(
+ op, shape_hints_max_bs, weights, blocklisted_ops, backend_id);
};
+ auto c2_converter = [this,
+ &weights,
+ &shape_hints_max_bs,
+ &shape_hints_per_bs](const caffe2::NetDef& net) {
+ return SubnetToOnnxifiOpViaC2(
+ net, weights, shape_hints_max_bs, shape_hints_per_bs);
+ };
+
return opt::OptimizeForBackend(
*pred_net, c2_supports, c2_converter, opts_.debug);
}
@@ -1248,7 +1321,7 @@
const std::unordered_set<std::string>& weights,
const std::unordered_set<int>& blocklisted_ops,
ShapeInfoMap* shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs) {
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs) {
onnxBackendID backend_id = backend_ids_[idx_];
// function to tell whether the ONNXIFI backend supports a given C2 op or not
@@ -1262,9 +1335,15 @@
// the same exporter throughout the process to avoid duplicated dummy name
// generation
onnx::OnnxExporter exporter2(nullptr);
- auto onnx_converter = [this, ws, &weights, shape_hints_max_bs, &exporter2, &shape_hints_per_bs](
+ auto onnx_converter = [this,
+ ws,
+ &weights,
+ shape_hints_max_bs,
+ &exporter2,
+ &shape_hints_per_bs](
const caffe2::NetDef& net) mutable {
- return SubnetToOnnxifiOpViaOnnx(net, weights, ws, &exporter2, shape_hints_max_bs, shape_hints_per_bs);
+ return SubnetToOnnxifiOpViaOnnx(
+ net, weights, ws, &exporter2, shape_hints_max_bs, shape_hints_per_bs);
};
return opt::OptimizeForBackend(
@@ -1334,7 +1413,10 @@
}
if (opts_.merge_fp32_inputs_into_fp16) {
mergeFp32InputsAndConvertToFp16(
- opts_.bound_shape_spec.max_batch_size, weights, pred_net, &shape_hints_max_bs);
+ opts_.bound_shape_spec.max_batch_size,
+ weights,
+ pred_net,
+ &shape_hints_max_bs);
}
if (opts_.debug) {
@@ -1355,13 +1437,23 @@
// Apply some filtering rules
std::unordered_set<int> new_blocklisted_ops(
blocklisted_ops.begin(), blocklisted_ops.end());
- applyFilteringRules(*pred_net, shape_hints_max_bs, weights, &new_blocklisted_ops);
+ applyFilteringRules(
+ *pred_net, shape_hints_max_bs, weights, &new_blocklisted_ops);
// Transform the net
- NetDef net_opt = opts_.use_onnx
- ? TransformViaOnnx(
- ws, pred_net, weights, new_blocklisted_ops, &shape_hints_max_bs, opts_.shape_hints_per_bs)
- : TransformViaC2(pred_net, weights, new_blocklisted_ops, shape_hints_max_bs, opts_.shape_hints_per_bs);
+ NetDef net_opt = opts_.use_onnx ? TransformViaOnnx(
+ ws,
+ pred_net,
+ weights,
+ new_blocklisted_ops,
+ &shape_hints_max_bs,
+ opts_.shape_hints_per_bs)
+ : TransformViaC2(
+ pred_net,
+ weights,
+ new_blocklisted_ops,
+ shape_hints_max_bs,
+ opts_.shape_hints_per_bs);
// Need to figure out a proper place to handle device option
net_opt.mutable_device_option()->CopyFrom(pred_net->device_option());
diff --git a/caffe2/opt/onnxifi_transformer.h b/caffe2/opt/onnxifi_transformer.h
index d1af173..56bd1cd 100644
--- a/caffe2/opt/onnxifi_transformer.h
+++ b/caffe2/opt/onnxifi_transformer.h
@@ -49,6 +49,21 @@
std::unordered_map<int, ShapeInfoMap> shape_hints_per_bs;
};
+class TORCH_API OnnxifiOptionHelper final {
+ public:
+ OnnxifiOptionHelper();
+
+ // Set Onnxifi option
+ bool setOnnxifiOption(const std::string& option, const std::string& value);
+
+ // Get Onnxifi option
+ std::string getOnnxifiOption(const std::string& option);
+
+ private:
+ // Pointer to loaded onnxifi library
+ onnxifi_library* lib_{nullptr};
+};
+
class TORCH_API OnnxifiTransformer final : public BackendTransformerBase {
public:
explicit OnnxifiTransformer(const OnnxifiTransformerOptions& opts);
@@ -84,14 +99,14 @@
Workspace* ws,
onnx::OnnxExporter* exporter,
ShapeInfoMap* shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs);
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs);
// Convert a cutoff subgraph net to an Onnxifi op
caffe2::NetDef SubnetToOnnxifiOpViaC2(
const caffe2::NetDef& net,
const std::unordered_set<std::string>& weights_in_ws,
const ShapeInfoMap& shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs);
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs);
// Check that output shape hints are present to ensure we can pass them to
// OnnxifiOp
@@ -106,7 +121,7 @@
const std::vector<std::string>& external_inputs,
const std::vector<std::string>& external_outputs,
const ShapeInfoMap& shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs);
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs);
// Transform by passing C2 proto to backend
NetDef TransformViaC2(
@@ -114,7 +129,7 @@
const std::unordered_set<std::string>& weights,
const std::unordered_set<int>& blocklisted_ops,
const ShapeInfoMap& shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs);
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs);
// Transform by passing ONNX proto to backend
NetDef TransformViaOnnx(
@@ -123,7 +138,7 @@
const std::unordered_set<std::string>& weights,
const std::unordered_set<int>& blocklisted_ops,
ShapeInfoMap* shape_hints_max_bs,
- const std::unordered_map<int, ShapeInfoMap> &shape_hints_per_bs);
+ const std::unordered_map<int, ShapeInfoMap>& shape_hints_per_bs);
// Query whether an operator is supported by passing ONNX protobuf
bool supportOpOnnx(
diff --git a/caffe2/python/onnx/onnxifi.py b/caffe2/python/onnx/onnxifi.py
index 3e67c49..973d063 100644
--- a/caffe2/python/onnx/onnxifi.py
+++ b/caffe2/python/onnx/onnxifi.py
@@ -5,15 +5,23 @@
ONNXIFI a Caffe2 net
"""
-
-
-
-
-
from caffe2.proto import caffe2_pb2
import caffe2.python._import_c_extension as C
+def onnxifi_set_option(option_name, option_value):
+ """
+ Set onnxifi option
+ """
+ return C.onnxifi_set_option(option_name, str(option_value))
+
+
+def onnxifi_get_option(option_name):
+ """
+ Get onnxifi option
+ """
+ return C.onnxifi_get_option(option_name)
+
def onnxifi_caffe2_net(
pred_net,
input_shapes,
diff --git a/caffe2/python/pybind_state.cc b/caffe2/python/pybind_state.cc
index 65a246e..87f95d9 100644
--- a/caffe2/python/pybind_state.cc
+++ b/caffe2/python/pybind_state.cc
@@ -23,13 +23,13 @@
#include "caffe2/onnx/offline_tensor.h"
#include "caffe2/onnx/onnx_exporter.h"
#include "caffe2/opt/converter.h"
+#include "caffe2/opt/fakefp16_transform.h"
#include "caffe2/opt/fusion.h"
#include "caffe2/opt/mobile.h"
#include "caffe2/opt/onnxifi_transformer.h"
#include "caffe2/opt/optimize_ideep.h"
#include "caffe2/opt/passes.h"
#include "caffe2/opt/shape_info.h"
-#include "caffe2/opt/fakefp16_transform.h"
#include "caffe2/predictor/emulator/data_filler.h"
#include "caffe2/predictor/predictor.h"
#include "caffe2/python/pybind_state_registry.h"
@@ -85,14 +85,13 @@
return gWorkspace;
}
-Workspace* GetWorkspaceByName(const std::string &name) {
+Workspace* GetWorkspaceByName(const std::string& name) {
if (gWorkspaces.count(name)) {
return gWorkspaces[name].get();
}
return nullptr;
}
-
class StringFetcher : public BlobFetcherBase {
public:
py::object Fetch(const Blob& blob) override {
@@ -376,10 +375,9 @@
py::gil_scoped_release g;
CAFFE_ENFORCE(net->Run());
})
- .def("cancel",
- [](NetBase* net) {
- py::gil_scoped_release g;
- net->Cancel();
+ .def("cancel", [](NetBase* net) {
+ py::gil_scoped_release g;
+ net->Cancel();
});
py::class_<ObserverBase<NetBase>>(m, "Observer")
@@ -1763,6 +1761,17 @@
return true;
});
m.def(
+ "onnxifi_set_option",
+ [](const std::string& optionName,
+ const std::string& optionValue) -> bool {
+ OnnxifiOptionHelper ts;
+ return ts.setOnnxifiOption(optionName, optionValue);
+ });
+ m.def("onnxifi_get_option", [](const std::string& optionName) -> std::string {
+ OnnxifiOptionHelper ts;
+ return ts.getOnnxifiOption(optionName);
+ });
+ m.def(
"onnxifi",
[](const py::bytes& pred_net_str,
const py::bytes& shapes_str,
@@ -1840,8 +1849,7 @@
m.def("fakeFp16FuseOps", [](const py::bytes& net_str) {
caffe2::NetDef netDef;
CAFFE_ENFORCE(
- ParseProtoFromLargeString(
- net_str.cast<std::string>(), &netDef),
+ ParseProtoFromLargeString(net_str.cast<std::string>(), &netDef),
"broken pred_net protobuf");
opt::fakeFp16FuseOps(&netDef);
std::string out_net;
diff --git a/third_party/foxi b/third_party/foxi
index 6a4e19a..bd6feb6 160000
--- a/third_party/foxi
+++ b/third_party/foxi
@@ -1 +1 @@
-Subproject commit 6a4e19a2aaf7ae4b9fa9597526e65b395d5e79ad
+Subproject commit bd6feb6d0d3fc903df42b4feb82a602a5fcb1fd5
