dc86bb5296
If pybind11 is installed on the system its headers are already captured by @local_config_python//:python_headers, so the system lib only needs to depend on that. When installed correctly, includes should be #include "pybind11/...", the bundled pybind11 is based off the source repo which does not match the install paths. Use bazels strip_include_prefix to align the bundled headers correctly. Signed-off-by: Jason Zaman <jason@perfinion.com>
73 lines
3.0 KiB
C++
73 lines
3.0 KiB
C++
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#include <string>
|
|
|
|
#include "pybind11/cast.h"
|
|
#include "pybind11/pybind11.h"
|
|
#include "pybind11/pytypes.h"
|
|
#include "pybind11/stl.h"
|
|
#include "tensorflow/compiler/aot/compile.h"
|
|
#include "tensorflow/compiler/aot/flags.h"
|
|
#include "tensorflow/python/lib/core/pybind11_lib.h"
|
|
#include "tensorflow/python/lib/core/pybind11_status.h"
|
|
|
|
namespace py = pybind11;
|
|
|
|
PYBIND11_MODULE(_pywrap_tfcompile, m) {
|
|
m.doc() = R"pbdoc(
|
|
_pywrap_tfcompile
|
|
-----
|
|
)pbdoc";
|
|
|
|
m.def(
|
|
"Compile",
|
|
[](std::string graph, std::string config, std::string target_triple,
|
|
std::string target_cpu, std::string target_features,
|
|
std::string entry_point, std::string cpp_class,
|
|
std::string out_function_object, std::string out_metadata_object,
|
|
std::string out_header, std::string out_session_module,
|
|
std::string mlir_components, bool gen_name_to_index,
|
|
bool gen_program_shape) {
|
|
tensorflow::tfcompile::MainFlags flags;
|
|
flags.graph = std::move(graph);
|
|
flags.config = std::move(config);
|
|
flags.target_triple = std::move(target_triple);
|
|
flags.target_cpu = std::move(target_cpu);
|
|
flags.target_features = std::move(target_features);
|
|
flags.entry_point = std::move(entry_point);
|
|
flags.cpp_class = std::move(cpp_class);
|
|
flags.out_function_object = std::move(out_function_object);
|
|
flags.out_metadata_object = std::move(out_metadata_object);
|
|
flags.out_header = std::move(out_header);
|
|
flags.out_session_module = std::move(out_session_module);
|
|
flags.mlir_components = std::move(mlir_components);
|
|
|
|
// C++ codegen options
|
|
flags.gen_name_to_index = gen_name_to_index;
|
|
flags.gen_program_shape = gen_program_shape;
|
|
|
|
tensorflow::MaybeRaiseFromStatus(tensorflow::tfcompile::Main(flags));
|
|
},
|
|
py::arg("graph") = "", py::arg("config") = "",
|
|
py::arg("target_triple") = "x86_64-pc-linux", py::arg("target_cpu") = "",
|
|
py::arg("target_features") = "", py::arg("entry_point") = "entry",
|
|
py::arg("cpp_class") = "", py::arg("out_function_object") = "out_model.o",
|
|
py::arg("out_metadata_object") = "out_helper.o",
|
|
py::arg("out_header") = "out.h", py::arg("out_session_module") = "",
|
|
py::arg("mlir_components") = "", py::arg("gen_name_to_index") = false,
|
|
py::arg("gen_program_shape") = false);
|
|
}
|