wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/c/experimental/grappler/BUILD
+77
View File
@@ -0,0 +1,77 @@
# Description:
# Graph C API.
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "grappler_hdrs",
hdrs = ["grappler.h"],
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/c:c_api",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_status_headers",
],
)
cc_library(
name = "grappler",
srcs = ["grappler.cc"],
hdrs = [
"grappler.h",
"grappler_internal.h",
],
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/c:c_api_internal",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_buffer",
"//tensorflow/c:tf_buffer_internal",
"//tensorflow/c:tf_status",
"//tensorflow/c:tf_status_helper",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/costs:graph_properties",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
"//tensorflow/core/platform:logging",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/status",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "grappler_test",
srcs = ["grappler_test.cc"],
deps = [
":grappler",
"//tensorflow/c:tf_buffer_internal",
"//tensorflow/c:tf_status_headers",
"//tensorflow/core:framework",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/clusters:single_machine",
"//tensorflow/core/grappler/costs:op_performance_data_cc",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/log:check",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tensorflow/c/experimental/grappler/grappler.cc
+396
View File
@@ -0,0 +1,396 @@
/* Copyright 2021 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.
==============================================================================*/
// This file extends/implements core graph optimizer base classes in terms of
// the C API defined in grappler.h. A class "CSomething" represents a
// "Something" that can be manipulated via calls in the C interface and a C
// struct called "TP_Something".
#include "tensorflow/c/experimental/grappler/grappler.h"
#include <algorithm>
#include <cstddef>
#include <cstring>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/grappler/grappler_internal.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_buffer_internal.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_helper.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace {
#define VALIDATE_STRUCT_SIZE(STRUCT_NAME, STRUCT_OBJ, SIZE_VALUE_NAME) \
do { \
if (STRUCT_OBJ.struct_size == 0) { \
return absl::Status(absl::StatusCode::kFailedPrecondition, \
"struct_size field in " #STRUCT_NAME \
" must be set to " #SIZE_VALUE_NAME "."); \
} \
} while (0)
#define VALIDATE_MEMBER(STRUCT_NAME, STRUCT_OBJ, NAME) \
do { \
if (STRUCT_OBJ.NAME == 0) { \
return absl::Status(absl::StatusCode::kFailedPrecondition, \
"'" #NAME "' field in " #STRUCT_NAME \
" must be set."); \
} \
} while (0)
absl::Status ValidateTPOptimizerRegistrationParams(
const TP_OptimizerRegistrationParams& params) {
VALIDATE_STRUCT_SIZE(TP_OptimizerRegistrationParams, params,
TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE);
VALIDATE_MEMBER(TP_OptimizerRegistrationParams, params, device_type);
return absl::OkStatus();
}
absl::Status ValidateTPOptimizer(const TP_Optimizer& optimizer) {
VALIDATE_STRUCT_SIZE(TP_Optimizer, optimizer, TP_OPTIMIZER_STRUCT_SIZE);
VALIDATE_MEMBER(TP_Optimizer, optimizer, optimize_func);
return absl::OkStatus();
}
absl::Status ValidateTPOptimizerConfigs(const TP_OptimizerConfigs& configs) {
VALIDATE_STRUCT_SIZE(TP_OptimizerConfigs, configs,
TP_OPTIMIZER_CONFIGS_STRUCT_SIZE);
return absl::OkStatus();
}
#undef VALIDATE_MEMBER
#undef VALIDATE_STRUCT_SIZE
} // namespace
namespace tensorflow {
namespace grappler {
absl::Status CGraphOptimizer::Optimize(Cluster* cluster,
const GrapplerItem& item,
GraphDef* optimized_graph_def) {
OwnedTFStatus c_status(TF_NewStatus());
OwnedTFBuffer graph_buf(TF_NewBuffer());
OwnedTFBuffer optimized_graph_buf(TF_NewBuffer());
TF_RETURN_IF_ERROR(MessageToBuffer(item.graph, graph_buf.get()));
optimizer_.optimize_func(c_optimizer_, graph_buf.get(),
reinterpret_cast<const TF_GrapplerItem*>(&item),
optimized_graph_buf.get(), c_status.get());
TF_RETURN_IF_ERROR(tsl::StatusFromTF_Status(c_status.get()));
TF_RETURN_IF_ERROR(
BufferToMessage(optimized_graph_buf.get(), optimized_graph_def));
return absl::OkStatus();
}
#define CONFIG_TOGGLE(optimizer) \
if (tp_configs.optimizer == TF_TriState_Off) \
configs.toggle_config[#optimizer] = RewriterConfig::OFF; \
else \
configs.toggle_config[#optimizer] = RewriterConfig::ON;
void CGraphOptimizerRegister(
const PluginGraphOptimizerRegistry::Creator& creator,
const TP_OptimizerConfigs tp_configs, const char* device_type) {
ConfigList configs;
// disable_model_pruning is turned off by default.
if (tp_configs.disable_model_pruning == TF_TriState_On)
configs.disable_model_pruning = true;
else
configs.disable_model_pruning = false;
// The other configs are turned on by default.
CONFIG_TOGGLE(implementation_selector);
CONFIG_TOGGLE(function_optimization);
CONFIG_TOGGLE(common_subgraph_elimination);
CONFIG_TOGGLE(arithmetic_optimization);
CONFIG_TOGGLE(debug_stripper);
CONFIG_TOGGLE(constant_folding);
CONFIG_TOGGLE(shape_optimization);
CONFIG_TOGGLE(auto_mixed_precision);
CONFIG_TOGGLE(auto_mixed_precision_onednn_bfloat16);
CONFIG_TOGGLE(auto_mixed_precision_mkl);
CONFIG_TOGGLE(pin_to_host_optimization);
CONFIG_TOGGLE(layout_optimizer);
CONFIG_TOGGLE(remapping);
CONFIG_TOGGLE(loop_optimization);
CONFIG_TOGGLE(dependency_optimization);
CONFIG_TOGGLE(auto_parallel);
CONFIG_TOGGLE(memory_optimization);
CONFIG_TOGGLE(scoped_allocator_optimization);
PluginGraphOptimizerRegistry::RegisterPluginOptimizerOrDie(
creator, device_type, configs);
}
#undef CONFIG_TOGGLE
absl::Status InitGraphPlugin(void* dso_handle) {
tsl::Env* env = tsl::Env::Default();
// Step 1: Load symbol for `TF_InitPlugin`
void* dso_symbol;
TF_RETURN_IF_ERROR(
env->GetSymbolFromLibrary(dso_handle, "TF_InitGraph", &dso_symbol));
// Step 2: Call `TF_InitPlugin`
auto init_fn = reinterpret_cast<TFInitGraphPluginFn>(dso_symbol);
return InitGraphPlugin(init_fn);
}
absl::Status InitGraphPlugin(TFInitGraphPluginFn init_fn) {
TP_OptimizerRegistrationParams params{
TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE};
TP_Optimizer optimizer{TP_OPTIMIZER_STRUCT_SIZE};
TP_OptimizerConfigs optimizer_configs{TP_OPTIMIZER_CONFIGS_STRUCT_SIZE};
params.major_version = GO_MAJOR;
params.minor_version = GO_MINOR;
params.patch_version = GO_PATCH;
params.optimizer = &optimizer;
params.optimizer_configs = &optimizer_configs;
OwnedTFStatus c_status(TF_NewStatus());
init_fn(&params, c_status.get());
TF_RETURN_IF_ERROR(tsl::StatusFromTF_Status(c_status.get()));
TF_RETURN_IF_ERROR(ValidateTPOptimizerRegistrationParams(params));
TF_RETURN_IF_ERROR(ValidateTPOptimizer(optimizer));
TF_RETURN_IF_ERROR(ValidateTPOptimizerConfigs(optimizer_configs));
CGraphOptimizerRegister(
[=]() { return new CGraphOptimizer(optimizer, params.device_type); },
optimizer_configs, params.device_type);
return absl::OkStatus();
}
} // namespace grappler
} // namespace tensorflow
void TF_GetNodesToPreserveListSize(const TF_GrapplerItem* item, int* num_values,
size_t* storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::unordered_set<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)
->NodesToPreserve();
*num_values = nodes.size();
*storage_size = 0;
for (const std::string& str : nodes) {
*storage_size += str.size();
}
}
void TF_GetNodesToPreserveList(const TF_GrapplerItem* item, char** values,
size_t* lengths, int num_values, void* storage,
size_t storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::unordered_set<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)
->NodesToPreserve();
char* p = static_cast<char*>(storage);
int index = 0;
for (const std::string& s : nodes) {
if (index >= num_values) break;
values[index] = p;
lengths[index] = s.size();
if ((p + s.size()) > (static_cast<char*>(storage) + storage_size)) {
tsl::Set_TF_Status_from_Status(
status,
absl::InvalidArgumentError(
"Not enough storage to hold the requested list of nodes"));
return;
}
memcpy(values[index], s.data(), s.size());
p += s.size();
index++;
}
}
void TF_GetFetchNodesListSize(const TF_GrapplerItem* item, int* num_values,
size_t* storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)->fetch;
*num_values = nodes.size();
*storage_size = 0;
for (const std::string& str : nodes) {
*storage_size += str.size();
}
}
void TF_GetFetchNodesList(const TF_GrapplerItem* item, char** values,
size_t* lengths, int num_values, void* storage,
size_t storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)->fetch;
const int len = std::min(num_values, static_cast<int>(nodes.size()));
char* p = static_cast<char*>(storage);
for (int index = 0; index < len; ++index) {
const std::string& s = nodes[index];
values[index] = p;
lengths[index] = s.size();
if ((p + s.size()) > (static_cast<char*>(storage) + storage_size)) {
tsl::Set_TF_Status_from_Status(
status,
absl::InvalidArgumentError(
"Not enough storage to hold the requested list of nodes"));
return;
}
memcpy(values[index], s.data(), s.size());
p += s.size();
}
}
TF_GraphProperties* TF_NewGraphProperties(const TF_GrapplerItem* item) {
return reinterpret_cast<TF_GraphProperties*>(
new tensorflow::grappler::GraphProperties(
*reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)));
}
void TF_DeleteGraphProperties(TF_GraphProperties* graph_properties) {
if (graph_properties == nullptr) return;
delete reinterpret_cast<tensorflow::grappler::GraphProperties*>(
graph_properties);
}
void TF_InferStatically(TF_GraphProperties* graph_properties,
TF_Bool assume_valid_feeds,
TF_Bool aggressive_shape_inference,
TF_Bool include_input_tensor_values,
TF_Bool include_output_tensor_values,
TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
absl::Status s =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->InferStatically(assume_valid_feeds, aggressive_shape_inference,
include_input_tensor_values,
include_output_tensor_values);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
}
}
void TF_GetInputPropertiesListSize(TF_GraphProperties* graph_properties,
const char* name, int* num_values,
TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
*num_values =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetInputProperties(name)
.size();
}
void TF_GetOutputPropertiesListSize(TF_GraphProperties* graph_properties,
const char* name, int* num_values,
TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
*num_values =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetOutputProperties(name)
.size();
}
void TF_GetInputPropertiesList(TF_GraphProperties* graph_properties,
const char* name, TF_Buffer** properties,
int num_values, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<tensorflow::OpInfo::TensorProperties>& tensor_properties =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetInputProperties(name);
const int len =
std::min(num_values, static_cast<int>(tensor_properties.size()));
for (int i = 0; i < len; ++i) {
absl::Status s =
tensorflow::MessageToBuffer(tensor_properties[i], properties[i]);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
}
}
void TF_GetOutputPropertiesList(TF_GraphProperties* graph_properties,
const char* name, TF_Buffer** properties,
int num_values, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<tensorflow::OpInfo::TensorProperties>& tensor_properties =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetOutputProperties(name);
const int len =
std::min(num_values, static_cast<int>(tensor_properties.size()));
for (int i = 0; i < len; ++i) {
absl::Status s =
tensorflow::MessageToBuffer(tensor_properties[i], properties[i]);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
}
}
TF_FunctionLibraryDefinition* TF_NewFunctionLibraryDefinition(
const TF_Buffer* graph_buf, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
tensorflow::GraphDef graph_def;
absl::Status s = tensorflow::BufferToMessage(graph_buf, &graph_def);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return nullptr;
}
return reinterpret_cast<TF_FunctionLibraryDefinition*>(
new tensorflow::FunctionLibraryDefinition(
tensorflow::OpRegistry::Global(), graph_def.library()));
}
void TF_DeleteFunctionLibraryDefinition(TF_FunctionLibraryDefinition* fn_lib) {
if (fn_lib == nullptr) return;
delete reinterpret_cast<tensorflow::FunctionLibraryDefinition*>(fn_lib);
}
void TF_LookUpOpDef(TF_FunctionLibraryDefinition* fn_lib, const char* name,
TF_Buffer* buf, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const tensorflow::OpDef* op_def_ptr = nullptr;
absl::Status s =
reinterpret_cast<tensorflow::FunctionLibraryDefinition*>(fn_lib)
->LookUpOpDef(name, &op_def_ptr);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
s = tensorflow::MessageToBuffer(*op_def_ptr, buf);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
}
tensorflow/c/experimental/grappler/grappler.h
+294
View File
@@ -0,0 +1,294 @@
/* Copyright 2021 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_
#include <stddef.h>
#include <stdint.h>
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_status.h"
// --------------------------------------------------------------------------
// C API for Graph. The API is under active development and eventually
// should allow registering a plugin graph optimizer with TensorFlow.
//
// Conventions:
// * Struct prefix indicates whether struct fields should be filled by the
// plugin or core implementation:
// * Struct that should be filled by the plugin: `TP_OptimizerConfigs`,
// `TP_Optimizer`, `TP_OptimizerRegistrationParams`
// * Struct that should be filled by the proper: `TF_GrapplerItem`,
// `TF_GraphProperties`, `TF_FunctionLibraryDefinition`
// * We use `struct_size` for version checking. It should be set both by
// core and the plugin.
// * For example, `TF_InitGraph` function receives
// `TP_OptimizerRegistrationParams*` as input with `struct_size`
// populated by core. The plugin is responsible for setting
// `struct_size` as well, along with all other fields.
// * Refer to "TensorFlow Versioning Strategy" section at
// https://github.com/tensorflow/community/pull/257/files.
// * Note that the API is still under active development and doesn't have
// versioning guarantees yet.
// * `void* ext` is a free-form field that can be populated by
// a plugin in `TP_*` structs or potential future extension points .
//
// Example usage:
//
// /* Sample TensorFlow code below, exact implementation might differ. */
// // Version checking uses `struct_size`. It should be set both by core
// // and the plugin.
// TP_OptimizerRegistrationParams params{
// TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE};
// TP_Optimizer optimizer{TP_OPTIMIZER_STRUCT_SIZE};
// TP_OptimizerConfigs configs{TP_OPTIMIZER_CONFIGS_STRUCT_SIZE};
// params.optimizer = &optimizer;
// params.configs = &configs;
//
// /* Plugin code below */
// void TF_InitGraph(TP_OptimizerRegistrationParams* params,
// TF_Status* status) {
// params->struct_size = TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE;
// params->device_type = "MY_DEVICE";
//
// // Disable certain optimizer.
// params->optimizer_configs->struct_size =
// TP_OPTIMIZER_CONFIGS_STRUCT_SIZE; params->optimizer_configs->remapping =
// TF_TriState_Off;
//
// // Set functions to create a new optimizer.
// params->optimizer->struct_size = TP_OPTIMIZER_STRUCT_SIZE;
// params->optimizer->create_func = (My_optimizer::create_func);
// }
#define GO_MAJOR 0
#define GO_MINOR 0
#define GO_PATCH 1
#ifdef __cplusplus
extern "C" {
#endif
// TF_TriState is the C API typedef for tri-state.
typedef enum TF_TriState {
TF_TriState_Default = 0,
TF_TriState_Off,
TF_TriState_On,
} TF_TriState;
// TF_GrapplerItem represents a combination of a graph, one of more fetch nodes,
// and potentially a set of nodes to feed.
typedef struct TF_GrapplerItem TF_GrapplerItem;
// Flags indicating whether existing optimizers should be turned off.
// It's optional for plugin to set functions to return true/false. If not
// set, proper uses configuration set by user.
typedef struct TP_OptimizerConfigs {
size_t struct_size;
void* ext; // reserved for future use
TF_TriState disable_model_pruning;
TF_TriState implementation_selector;
TF_TriState function_optimization;
TF_TriState common_subgraph_elimination;
TF_TriState arithmetic_optimization;
TF_TriState debug_stripper;
TF_TriState constant_folding;
TF_TriState shape_optimization;
TF_TriState auto_mixed_precision;
TF_TriState auto_mixed_precision_onednn_bfloat16;
TF_TriState auto_mixed_precision_mkl;
TF_TriState pin_to_host_optimization;
TF_TriState layout_optimizer;
TF_TriState remapping;
TF_TriState loop_optimization;
TF_TriState dependency_optimization;
TF_TriState auto_parallel;
TF_TriState memory_optimization;
TF_TriState scoped_allocator_optimization;
} TP_OptimizerConfigs;
#define TP_OPTIMIZER_CONFIGS_STRUCT_SIZE \
TF_OFFSET_OF_END(TP_OptimizerConfigs, scoped_allocator_optimization)
// Struct for Optimizer. Plugin authors must provide an optimize function.
// Creation and deletion functions are optional.
typedef struct TP_Optimizer {
size_t struct_size;
void* ext; // reserved for future use
// [Optional]
// Create function for optimizer.
void* (*create_func)();
// Optimizer function for optimizer. The first param is an optimizer created
// by create_func. The second param is input graph. The third param is
// GrapplerItem. The fourth param is output graph.
void (*optimize_func)(void*, const TF_Buffer*, const TF_GrapplerItem*,
TF_Buffer*, TF_Status*);
// [Optional]
// Destroy function for optimizer. If Create function is provided, destroy
// function is must.
void (*destroy_func)(void*);
} TP_Optimizer;
#define TP_OPTIMIZER_STRUCT_SIZE TF_OFFSET_OF_END(TP_Optimizer, destroy_func)
typedef struct TP_OptimizerRegistrationParams {
size_t struct_size;
void* ext; // reserved for future use
// Graph C API version.
int32_t major_version;
int32_t minor_version;
int32_t patch_version;
// Backend device type supported by the optimizer.
const char* device_type;
TP_OptimizerConfigs* optimizer_configs; // output, set by plugin
TP_Optimizer* optimizer; // output, set by plugin
} TP_OptimizerRegistrationParams;
#define TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE \
TF_OFFSET_OF_END(TP_OptimizerRegistrationParams, optimizer)
// TF_InitGraph is used to do graph optimizer registration.
// Plugin should implement TF_InitGraph to register graph optimizers.
void TF_InitGraph(TP_OptimizerRegistrationParams* params, TF_Status* status);
// Get a set of node names that must be preserved. They can not be transformed
// or removed during the graph transformation. This includes feed and fetch
// nodes, keep_ops, init_ops. Fills in `num_values` and `storage_size`, they
// will be used in `TF_GetNodesToPreserveList`.
TF_CAPI_EXPORT extern void TF_GetNodesToPreserveListSize(
const TF_GrapplerItem* item, int* num_values, size_t* storage_size,
TF_Status* status);
// Get a set of node names that must be preserved. They can not be transformed
// or removed during the graph transformation. This includes feed and fetch
// nodes, keep_ops, init_ops. Fills in `values` and `lengths`, each of which
// must point to an array of length at least `num_values`.
//
// The elements of values will point to addresses in `storage` which must be at
// least `storage_size` bytes in length. `num_values` and `storage` can be
// obtained from TF_GetNodesToPreserveSize
//
// Fails if storage_size is too small to hold the requested number of strings.
TF_CAPI_EXPORT extern void TF_GetNodesToPreserveList(
const TF_GrapplerItem* item, char** values, size_t* lengths, int num_values,
void* storage, size_t storage_size, TF_Status* status);
// Get a set of node names for fetch nodes. Fills in `values` and `lengths`,
// they will be used in `TF_GetFetchNodesList`
TF_CAPI_EXPORT extern void TF_GetFetchNodesListSize(const TF_GrapplerItem* item,
int* num_values,
size_t* storage_size,
TF_Status* status);
// Get a set of node names for fetch nodes. Fills in `values` and `lengths`,
// each of which must point to an array of length at least `num_values`.
//
// The elements of values will point to addresses in `storage` which must be at
// least `storage_size` bytes in length. `num_values` and `storage` can be
// obtained from TF_GetFetchNodesSize
//
// Fails if storage_size is too small to hold the requested number of strings.
TF_CAPI_EXPORT extern void TF_GetFetchNodesList(const TF_GrapplerItem* item,
char** values, size_t* lengths,
int num_values, void* storage,
size_t storage_size,
TF_Status* status);
// Infer OpInfo::TensorProperties for graph nodes inputs/outputs.
//
// Typical use case, is to infer tensor properties from a graph, before doing
// optimization pass. Nodes modified during optimization pass have to be
// invalidated, to prevent further incorrect optimizations based on wrong shape
// and data type properties.
typedef struct TF_GraphProperties TF_GraphProperties;
// Create GraphProperties. The item must outlive the properties.
TF_CAPI_EXPORT extern TF_GraphProperties* TF_NewGraphProperties(
const TF_GrapplerItem* item);
// Delete GraphProperties.
TF_CAPI_EXPORT extern void TF_DeleteGraphProperties(
TF_GraphProperties* graph_properties);
// Infer tensor shapes through abstract interpretation.
// If assume_valid_feeds is true, it can help infer shapes in the fanout of fed
// nodes. This may cause incorrectness in graph analyses, but is useful for
// simulation or scheduling.
// If aggressive_shape_inference is true, nodes are executed on the host to
// identify output values when possible and does other aggressive strategies.
// This may cause incorrectness in graph analyses, but is useful for simulation
// or scheduling.
// If include_input_tensor_values is true, the values of constant
// tensors will included in the input properties.
// If include_output_tensor_values is true, the values of constant tensors will
// be included in the output properties.
TF_CAPI_EXPORT extern void TF_InferStatically(
TF_GraphProperties* graph_properties, TF_Bool assume_valid_feeds,
TF_Bool aggressive_shape_inference, TF_Bool include_input_tensor_values,
TF_Bool include_output_tensor_values, TF_Status* s);
// Get the size of input OpInfo::TensorProperties given node name.
TF_CAPI_EXPORT extern void TF_GetInputPropertiesListSize(
TF_GraphProperties* graph_properties, const char* name, int* num_values,
TF_Status* status);
// Get the size of output OpInfo::TensorProperties given node name.
TF_CAPI_EXPORT extern void TF_GetOutputPropertiesListSize(
TF_GraphProperties* graph_properties, const char* name, int* num_values,
TF_Status* status);
// Get a list of input OpInfo::TensorProperties given node name.
// Return the serialized list `properties`.
TF_CAPI_EXPORT extern void TF_GetInputPropertiesList(
TF_GraphProperties* graph_properties, const char* name,
TF_Buffer** properties, int num_values, TF_Status* status);
// Get a list of output OpInfo::TensorProperties given node name.
// Return the serialized list `properties`.
TF_CAPI_EXPORT extern void TF_GetOutputPropertiesList(
TF_GraphProperties* graph_properties, const char* name,
TF_Buffer** properties, int num_values, TF_Status* status);
// Helper to maintain a map between function names in a given
// FunctionDefLibrary and function definitions.
// Typical use case, is to look up an OpDef by type name.
typedef struct TF_FunctionLibraryDefinition TF_FunctionLibraryDefinition;
// Create NewFunctionLibraryDefinition.
TF_CAPI_EXPORT extern TF_FunctionLibraryDefinition*
TF_NewFunctionLibraryDefinition(const TF_Buffer* graph_buf, TF_Status* status);
// Delete NewFunctionLibraryDefinition.
TF_CAPI_EXPORT extern void TF_DeleteFunctionLibraryDefinition(
TF_FunctionLibraryDefinition* fn_lib);
// Shorthand for calling LookUp to get the OpDef from FunctionLibraryDefinition
// given op name. The returned OpDef is represented by TF_Buffer.
TF_CAPI_EXPORT extern void TF_LookUpOpDef(TF_FunctionLibraryDefinition* fn_lib,
const char* name, TF_Buffer* buf,
TF_Status* s);
#ifdef __cplusplus
} // extern "C"
#endif
#endif // TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_
tensorflow/c/experimental/grappler/grappler_internal.h
+104
View File
@@ -0,0 +1,104 @@
/* Copyright 2021 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.
==============================================================================*/
// Classes and utilities that work with Graph C API for internal use.
// This includes functions used for optimizer registration and interfaces needed
// for testing.
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_INTERNAL_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_INTERNAL_H_
#include <functional>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/experimental/grappler/grappler.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace grappler {
// Plugin initialization function that a device plugin
// must define.
typedef void (*TFInitGraphPluginFn)(TP_OptimizerRegistrationParams* const,
TF_Status* const);
// Registers Graph optimizers.
Status InitGraphPlugin(void* dso_handle);
// Allow registering a graph optimizer using a function (used for
// testing).
Status InitGraphPlugin(TFInitGraphPluginFn init_fn);
struct GrapplerItem;
class Cluster;
struct TFStatusDeleter {
void operator()(TF_Status* s) const { TF_DeleteStatus(s); }
};
using OwnedTFStatus = std::unique_ptr<TF_Status, TFStatusDeleter>;
struct TFBufferDeleter {
void operator()(TF_Buffer* buf) const { TF_DeleteBuffer(buf); }
};
using OwnedTFBuffer = std::unique_ptr<TF_Buffer, TFBufferDeleter>;
class CGraphOptimizer : public CustomGraphOptimizer {
public:
explicit CGraphOptimizer(TP_Optimizer optimizer, const char* device_type)
: optimizer_(optimizer), device_type_(device_type) {
if (optimizer.create_func != nullptr) {
c_optimizer_ = (*optimizer_.create_func)();
} else {
c_optimizer_ = nullptr;
}
}
std::string name() const override { return "PluggableGraphOptimizer"; }
bool UsesFunctionLibrary() const override { return false; }
Status Init(
const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override {
return OkStatus();
}
Status Optimize(Cluster* cluster, const GrapplerItem& item,
GraphDef* optimized_graph_def) override;
~CGraphOptimizer() override {
if (optimizer_.destroy_func != nullptr) {
(*optimizer_.destroy_func)(c_optimizer_);
}
}
private:
TP_Optimizer optimizer_;
std::string device_type_;
void* c_optimizer_;
};
// Registration function to register a CGraphOptimizer along with plugin configs
// and device type.
void CGraphOptimizerRegister(
const PluginGraphOptimizerRegistry::Creator& creator,
const TP_OptimizerConfigs tp_configs, const char* device_type);
} // namespace grappler
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_INTERNAL_H_
tensorflow/c/experimental/grappler/grappler_test.cc
+329
View File
@@ -0,0 +1,329 @@
/* Copyright 2021 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 "tensorflow/c/experimental/grappler/grappler.h"
#include <cstddef>
#include <memory>
#include <set>
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/log/check.h"
#include "tensorflow/c/experimental/grappler/grappler_internal.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_buffer_internal.h"
#include "tensorflow/c/tf_status.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/grappler/clusters/single_machine.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
void optimize_func(void* optimizer, const TF_Buffer* graph_buf,
const TF_GrapplerItem* item, TF_Buffer* optimized_graph_buf,
TF_Status* tf_status) {}
void PopulateDefaultParam(TP_OptimizerRegistrationParams* params) {
params->struct_size = TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE;
params->optimizer_configs->struct_size = TP_OPTIMIZER_CONFIGS_STRUCT_SIZE;
params->optimizer->struct_size = TP_OPTIMIZER_STRUCT_SIZE;
params->optimizer->create_func = nullptr;
params->optimizer->optimize_func = optimize_func;
params->optimizer->destroy_func = nullptr;
}
TEST(Grappler, SuccessfulRegistration) {
auto plugin_init = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "Success";
params->optimizer_configs->remapping = TF_TriState_Off;
};
TF_ASSERT_OK(InitGraphPlugin(plugin_init));
ASSERT_EQ(PluginGraphOptimizerRegistry::CreateOptimizers(
std::set<std::string>{"Success"})
.size(),
1);
ConfigList config = PluginGraphOptimizerRegistry::GetPluginConfigs(
true, std::set<std::string>{"Success"});
ASSERT_EQ(config.toggle_config["remapping"], RewriterConfig::OFF);
}
TEST(Grappler, MultiplePluginRegistration) {
auto plugin_init_0 = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "Device0";
};
auto plugin_init_1 = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "Device1";
};
TF_ASSERT_OK(InitGraphPlugin(plugin_init_0));
TF_ASSERT_OK(InitGraphPlugin(plugin_init_1));
ASSERT_EQ(PluginGraphOptimizerRegistry::CreateOptimizers(
std::set<std::string>{"Device0", "Device1"})
.size(),
2);
}
TEST(Grappler, DeviceTypeNotSet) {
auto plugin_init = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = nullptr;
};
absl::Status status = InitGraphPlugin(plugin_init);
ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
ASSERT_EQ(
status.message(),
"'device_type' field in TP_OptimizerRegistrationParams must be set.");
}
TEST(Grappler, OptimizeFuncNotSet) {
auto plugin_init = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "FuncNotSet";
params->optimizer->optimize_func = nullptr;
};
absl::Status status = InitGraphPlugin(plugin_init);
ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
ASSERT_EQ(status.message(),
"'optimize_func' field in TP_Optimizer must be set.");
}
TEST(TF_GrapplerItem, NodesToPreserve) {
GrapplerItem item;
item.fetch = std::vector<std::string>{"Conv", "BiasAdd"};
std::unordered_set<std::string> nodes_preserved = item.NodesToPreserve();
TF_GrapplerItem* c_item = reinterpret_cast<TF_GrapplerItem*>(&item);
int list_total_size = 0;
for (const std::string& s : nodes_preserved) {
list_total_size += s.size();
}
size_t storage_size = 0;
int num_values = 0;
TF_Status* status = TF_NewStatus();
TF_GetNodesToPreserveListSize(c_item, &num_values, &storage_size, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(nodes_preserved.size(), num_values);
EXPECT_EQ(list_total_size, storage_size);
std::unique_ptr<char*[]> values(new char*[nodes_preserved.size()]);
std::unique_ptr<size_t[]> lens(new size_t[nodes_preserved.size()]);
std::unique_ptr<char[]> storage(new char[storage_size]);
TF_GetNodesToPreserveList(c_item, values.get(), lens.get(),
nodes_preserved.size(), storage.get(), storage_size,
status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (size_t i = 0; i < nodes_preserved.size(); ++i) {
EXPECT_EQ(
nodes_preserved.find(std::string(static_cast<const char*>(values[i]),
lens[i])) != nodes_preserved.end(),
true);
}
TF_DeleteStatus(status);
}
TEST(TF_GrapplerItem, FetchNodes) {
GrapplerItem item;
item.fetch = std::vector<std::string>{"Conv", "BiasAdd"};
TF_GrapplerItem* c_item = reinterpret_cast<TF_GrapplerItem*>(&item);
int list_total_size = 0;
for (const std::string& s : item.fetch) {
list_total_size += s.size();
}
size_t storage_size = 0;
int num_values = 0;
TF_Status* status = TF_NewStatus();
TF_GetFetchNodesListSize(c_item, &num_values, &storage_size, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(item.fetch.size(), num_values);
EXPECT_EQ(list_total_size, storage_size);
std::unique_ptr<char*[]> values(new char*[item.fetch.size()]);
std::unique_ptr<size_t[]> lens(new size_t[item.fetch.size()]);
std::unique_ptr<char[]> storage(new char[storage_size]);
TF_GetFetchNodesList(c_item, values.get(), lens.get(), item.fetch.size(),
storage.get(), storage_size, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (size_t i = 0; i < item.fetch.size(); ++i) {
EXPECT_EQ(item.fetch[i].size(), lens[i]) << i;
EXPECT_EQ(item.fetch[i],
std::string(static_cast<const char*>(values[i]), lens[i]))
<< i;
}
TF_DeleteStatus(status);
}
TEST(TF_GraphProperties, InputProperties) {
std::unique_ptr<SingleMachine> cluster(new SingleMachine(5 * 60, 3, 0));
TF_ASSERT_OK(cluster->Provision());
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_Status* status = TF_NewStatus();
TF_GraphProperties* graph_properties =
TF_NewGraphProperties(reinterpret_cast<TF_GrapplerItem*>(&item));
TF_InferStatically(graph_properties, true, false, false, false, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (const NodeDef& node : item.graph.node()) {
if (node.op() == "AddN") {
int num_values = 0;
TF_GetInputPropertiesListSize(graph_properties, node.name().c_str(),
&num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(num_values, 1);
std::vector<TF_Buffer*> in_props_buf(num_values, TF_NewBuffer());
TF_GetInputPropertiesList(graph_properties, node.name().c_str(),
in_props_buf.data(), num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
tensorflow::OpInfo::TensorProperties in_props;
absl::Status s = tensorflow::BufferToMessage(in_props_buf[0], &in_props);
TF_ASSERT_OK(s);
EXPECT_EQ(DT_FLOAT, in_props.dtype());
EXPECT_FALSE(in_props.shape().unknown_rank());
EXPECT_EQ(2, in_props.shape().dim_size());
EXPECT_EQ(10, in_props.shape().dim(0).size());
EXPECT_EQ(1, in_props.shape().dim(1).size());
for (int i = 0; i < in_props_buf.size(); i++)
TF_DeleteBuffer(in_props_buf[i]);
}
}
TF_DeleteGraphProperties(graph_properties);
TF_DeleteStatus(status);
TF_ASSERT_OK(cluster->Shutdown());
}
TEST(TF_GraphProperties, OutputProperties) {
std::unique_ptr<SingleMachine> cluster(new SingleMachine(5 * 60, 3, 0));
TF_ASSERT_OK(cluster->Provision());
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_Status* status = TF_NewStatus();
TF_GraphProperties* graph_properties =
TF_NewGraphProperties(reinterpret_cast<TF_GrapplerItem*>(&item));
TF_InferStatically(graph_properties, true, false, false, false, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (const NodeDef& node : item.graph.node()) {
if (node.op() == "AddN") {
int num_values = 0;
TF_GetOutputPropertiesListSize(graph_properties, node.name().c_str(),
&num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(num_values, 1);
std::vector<TF_Buffer*> out_props_buf(num_values, TF_NewBuffer());
TF_GetOutputPropertiesList(graph_properties, node.name().c_str(),
out_props_buf.data(), num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
tensorflow::OpInfo::TensorProperties out_props;
absl::Status s =
tensorflow::BufferToMessage(out_props_buf[0], &out_props);
TF_ASSERT_OK(s);
EXPECT_EQ(DT_FLOAT, out_props.dtype());
EXPECT_FALSE(out_props.shape().unknown_rank());
EXPECT_EQ(2, out_props.shape().dim_size());
EXPECT_EQ(10, out_props.shape().dim(0).size());
EXPECT_EQ(1, out_props.shape().dim(1).size());
for (int i = 0; i < out_props_buf.size(); i++)
TF_DeleteBuffer(out_props_buf[i]);
}
}
TF_DeleteStatus(status);
TF_DeleteGraphProperties(graph_properties);
TF_ASSERT_OK(cluster->Shutdown());
}
TEST(TF_FunctionLibraryDefinition, LookUpOpDef) {
TF_Buffer* g_buf = TF_NewBuffer();
TF_Buffer* op_buf = TF_NewBuffer();
TF_Status* status = TF_NewStatus();
GraphDef g_def;
absl::Status s = MessageToBuffer(g_def, g_buf);
TF_ASSERT_OK(s);
TF_FunctionLibraryDefinition* func =
TF_NewFunctionLibraryDefinition(g_buf, status);
TF_LookUpOpDef(func, "Add", op_buf, status);
std::string actual_string(reinterpret_cast<const char*>(op_buf->data),
op_buf->length);
ASSERT_EQ(TF_OK, TF_GetCode(status));
const OpDef* expected_op_def;
TF_ASSERT_OK(OpRegistry::Global()->LookUpOpDef("Add", &expected_op_def));
std::string expected_serialized;
expected_op_def->SerializeToString(&expected_serialized);
EXPECT_EQ(expected_serialized, actual_string);
TF_DeleteBuffer(g_buf);
TF_DeleteBuffer(op_buf);
TF_DeleteStatus(status);
TF_DeleteFunctionLibraryDefinition(func);
}
} // namespace
} // namespace grappler
} // namespace tensorflow