wehub-resource-sync
71 lines
2.3 KiB
C++
71 lines
2.3 KiB
C++
/* Copyright 2022 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.
|
|
==============================================================================*/
|
|
|
|
// Usage:
|
|
// offset_counter [ --out_path FILENAME ] [ SRC_FILE1 SRC_FILE2 ... ]
|
|
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "absl/strings/string_view.h"
|
|
#include "xla/tsl/platform/errors.h"
|
|
#include "xla/tsl/platform/types.h"
|
|
#include "xla/tsl/util/command_line_flags.h"
|
|
#include "tensorflow/python/framework/offset_counter_helper.h"
|
|
#include "tensorflow/python/framework/op_reg_offset.pb.h"
|
|
#include "tsl/platform/init_main.h"
|
|
#include "tsl/platform/strcat.h"
|
|
|
|
inline constexpr absl::string_view kUsage =
|
|
"offset_counter reads C++ source codes, scans for the location of where "
|
|
"the REGISTER_OP gets called, and outputs a OpRegOffsets proto to stdout "
|
|
"or a file.";
|
|
|
|
int main(int argc, char* argv[]) {
|
|
std::string out_path;
|
|
std::vector<tsl::Flag> flag_list = {
|
|
tsl::Flag("out_path", &out_path, "Output file path."),
|
|
};
|
|
const std::string usage_string =
|
|
absl::StrCat(kUsage, "\n\n", tsl::Flags::Usage(argv[0], flag_list));
|
|
const bool parse_result = tsl::Flags::Parse(&argc, argv, flag_list);
|
|
tsl::port::InitMain(usage_string.c_str(), &argc, &argv);
|
|
if (!parse_result) {
|
|
LOG(ERROR) << usage_string;
|
|
return -1;
|
|
}
|
|
|
|
tensorflow::OpRegOffsets op_reg_offsets;
|
|
for (size_t i = 1; i < argc; ++i) {
|
|
TF_CHECK_OK(tensorflow::FindOpRegistationFromFile(argv[i], op_reg_offsets));
|
|
}
|
|
|
|
if (out_path.empty()) {
|
|
std::cout << op_reg_offsets.SerializeAsString();
|
|
} else {
|
|
std::ofstream f(out_path);
|
|
if (f.bad()) {
|
|
TF_CHECK_OK(tsl::errors::IOError(
|
|
absl::StrCat("Cannot open file: ", out_path), errno));
|
|
}
|
|
f << op_reg_offsets.SerializeAsString();
|
|
f.close();
|
|
}
|
|
|
|
return 0;
|
|
}
|