1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_ 16 #define TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_ 17 18 #include <stddef.h> 19 #include <stdint.h> 20 21 #include "tensorflow/c/c_api.h" 22 #include "tensorflow/c/c_api_macros.h" 23 #include "tensorflow/c/tf_status.h" 24 25 // -------------------------------------------------------------------------- 26 // C API for Graph. The API is under active development and eventually 27 // should allow registering a plugin graph optimizer with TensorFlow. 28 // 29 // Conventions: 30 // * Struct prefix indicates whether struct fields should be filled by the 31 // plugin or core implementation: 32 // * Struct that should be filled by the plugin: `TP_OptimizerConfigs`, 33 // `TP_Optimizer`, `TP_OptimizerRegistrationParams` 34 // * Struct that should be filled by the proper: `TF_GrapplerItem`, 35 // `TF_GraphProperties`, `TF_FunctionLibraryDefinition` 36 // * We use `struct_size` for version checking. It should be set both by 37 // core and the plugin. 38 // * For example, `TF_InitGraph` function receives 39 // `TP_OptimizerRegistrationParams*` as input with `struct_size` 40 // populated by core. The plugin is responsible for setting 41 // `struct_size` as well, along with all other fields. 42 // * Refer to "TensorFlow Versioning Strategy" section at 43 // https://github.com/tensorflow/community/pull/257/files. 44 // * Note that the API is still under active development and doesn't have 45 // versioning guarantees yet. 46 // * `void* ext` is a free-form field that can be populated by 47 // a plugin in `TP_*` structs or potential future extension points . 48 // 49 // Example usage: 50 // 51 // /* Sample TensorFlow code below, exact implementation might differ. */ 52 // // Version checking uses `struct_size`. It should be set both by core 53 // // and the plugin. 54 // TP_OptimizerRegistrationParams params{ 55 // TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE}; 56 // TP_Optimizer optimizer{TP_OPTIMIZER_STRUCT_SIZE}; 57 // TP_OptimizerConfigs configs{TP_OPTIMIZER_CONFIGS_STRUCT_SIZE}; 58 // params.optimizer = &optimizer; 59 // params.configs = &configs; 60 // 61 // /* Plugin code below */ 62 // void TF_InitGraph(TP_OptimizerRegistrationParams* params, 63 // TF_Status* status) { 64 // params->struct_size = TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE; 65 // params->device_type = "MY_DEVICE"; 66 // 67 // // Disable certain optimizer. 68 // params->optimizer_configs->struct_size = 69 // TP_OPTIMIZER_CONFIGS_STRUCT_SIZE; params->optimizer_configs->remapping = 70 // TF_TriState_Off; 71 // 72 // // Set functions to create a new optimizer. 73 // params->optimizer->struct_size = TP_OPTIMIZER_STRUCT_SIZE; 74 // params->optimizer->create_func = (My_optimizer::create_func); 75 // } 76 77 #define GO_MAJOR 0 78 #define GO_MINOR 0 79 #define GO_PATCH 1 80 81 #ifdef __cplusplus 82 extern "C" { 83 #endif 84 85 // TF_TriState is the C API typedef for tri-state. 86 typedef enum TF_TriState { 87 TF_TriState_Default = 0, 88 TF_TriState_Off, 89 TF_TriState_On, 90 } TF_TriState; 91 92 // TF_GrapplerItem represents a combination of a graph, one of more fetch nodes, 93 // and potentially a set of nodes to feed. 94 typedef struct TF_GrapplerItem TF_GrapplerItem; 95 96 // Flags indicating whether existing optimizers should be turned off. 97 // It's optional for plugin to set functions to return true/false. If not 98 // set, proper uses configuration set by user. 99 typedef struct TP_OptimizerConfigs { 100 size_t struct_size; 101 void* ext; // reserved for future use 102 TF_TriState disable_model_pruning; 103 TF_TriState implementation_selector; 104 TF_TriState function_optimization; 105 TF_TriState common_subgraph_elimination; 106 TF_TriState arithmetic_optimization; 107 TF_TriState debug_stripper; 108 TF_TriState constant_folding; 109 TF_TriState shape_optimization; 110 TF_TriState auto_mixed_precision; 111 TF_TriState auto_mixed_precision_onednn_bfloat16; 112 TF_TriState auto_mixed_precision_mkl; 113 TF_TriState pin_to_host_optimization; 114 TF_TriState layout_optimizer; 115 TF_TriState remapping; 116 TF_TriState loop_optimization; 117 TF_TriState dependency_optimization; 118 TF_TriState auto_parallel; 119 TF_TriState memory_optimization; 120 TF_TriState scoped_allocator_optimization; 121 } TP_OptimizerConfigs; 122 123 #define TP_OPTIMIZER_CONFIGS_STRUCT_SIZE \ 124 TF_OFFSET_OF_END(TP_OptimizerConfigs, scoped_allocator_optimization) 125 126 // Struct for Optimizer. Plugin authors must provide an optimize function. 127 // Creation and deletion functions are optional. 128 typedef struct TP_Optimizer { 129 size_t struct_size; 130 void* ext; // reserved for future use 131 132 // [Optional] 133 // Create function for optimizer. 134 void* (*create_func)(); 135 136 // Optimizer function for optimizer. The first param is an optimizer created 137 // by create_func. The second param is input graph. The third param is 138 // GrapplerItem. The fourth param is output graph. 139 void (*optimize_func)(void*, const TF_Buffer*, const TF_GrapplerItem*, 140 TF_Buffer*, TF_Status*); 141 142 // [Optional] 143 // Destroy function for optimizer. If Create function is provided, destroy 144 // function is must. 145 void (*destroy_func)(void*); 146 } TP_Optimizer; 147 148 #define TP_OPTIMIZER_STRUCT_SIZE TF_OFFSET_OF_END(TP_Optimizer, destroy_func) 149 150 typedef struct TP_OptimizerRegistrationParams { 151 size_t struct_size; 152 void* ext; // reserved for future use 153 154 // Graph C API version. 155 int32_t major_version; 156 int32_t minor_version; 157 int32_t patch_version; 158 159 // Backend device type supported by the optimizer. 160 const char* device_type; 161 TP_OptimizerConfigs* optimizer_configs; // output, set by plugin 162 TP_Optimizer* optimizer; // output, set by plugin 163 } TP_OptimizerRegistrationParams; 164 165 #define TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE \ 166 TF_OFFSET_OF_END(TP_OptimizerRegistrationParams, optimizer) 167 168 // TF_InitGraph is used to do graph optimizer registration. 169 // Plugin should implement TF_InitGraph to register graph optimizers. 170 void TF_InitGraph(TP_OptimizerRegistrationParams* params, TF_Status* status); 171 172 // Get a set of node names that must be preserved. They can not be transformed 173 // or removed during the graph transformation. This includes feed and fetch 174 // nodes, keep_ops, init_ops. Fills in `num_values` and `storage_size`, they 175 // will be used in `TF_GetNodesToPreserveList`. 176 TF_CAPI_EXPORT extern void TF_GetNodesToPreserveListSize( 177 const TF_GrapplerItem* item, int* num_values, size_t* storage_size, 178 TF_Status* status); 179 180 // Get a set of node names that must be preserved. They can not be transformed 181 // or removed during the graph transformation. This includes feed and fetch 182 // nodes, keep_ops, init_ops. Fills in `values` and `lengths`, each of which 183 // must point to an array of length at least `num_values`. 184 // 185 // The elements of values will point to addresses in `storage` which must be at 186 // least `storage_size` bytes in length. `num_values` and `storage` can be 187 // obtained from TF_GetNodesToPreserveSize 188 // 189 // Fails if storage_size is too small to hold the requested number of strings. 190 TF_CAPI_EXPORT extern void TF_GetNodesToPreserveList( 191 const TF_GrapplerItem* item, char** values, size_t* lengths, int num_values, 192 void* storage, size_t storage_size, TF_Status* status); 193 194 // Get a set of node names for fetch nodes. Fills in `values` and `lengths`, 195 // they will be used in `TF_GetFetchNodesList` 196 TF_CAPI_EXPORT extern void TF_GetFetchNodesListSize(const TF_GrapplerItem* item, 197 int* num_values, 198 size_t* storage_size, 199 TF_Status* status); 200 201 // Get a set of node names for fetch nodes. Fills in `values` and `lengths`, 202 // each of which must point to an array of length at least `num_values`. 203 // 204 // The elements of values will point to addresses in `storage` which must be at 205 // least `storage_size` bytes in length. `num_values` and `storage` can be 206 // obtained from TF_GetFetchNodesSize 207 // 208 // Fails if storage_size is too small to hold the requested number of strings. 209 TF_CAPI_EXPORT extern void TF_GetFetchNodesList(const TF_GrapplerItem* item, 210 char** values, size_t* lengths, 211 int num_values, void* storage, 212 size_t storage_size, 213 TF_Status* status); 214 215 // Infer OpInfo::TensorProperties for graph nodes inputs/outputs. 216 // 217 // Typical use case, is to infer tensor properties from a graph, before doing 218 // optimization pass. Nodes modified during optimization pass have to be 219 // invalidated, to prevent further incorrect optimizations based on wrong shape 220 // and data type properties. 221 typedef struct TF_GraphProperties TF_GraphProperties; 222 223 // Create GraphProperties. The item must outlive the properties. 224 TF_CAPI_EXPORT extern TF_GraphProperties* TF_NewGraphProperties( 225 const TF_GrapplerItem* item); 226 227 // Delete GraphProperties. 228 TF_CAPI_EXPORT extern void TF_DeleteGraphProperties( 229 TF_GraphProperties* graph_properties); 230 231 // Infer tensor shapes through abstract interpretation. 232 // If assume_valid_feeds is true, it can help infer shapes in the fanout of fed 233 // nodes. This may cause incorrectness in graph analyses, but is useful for 234 // simulation or scheduling. 235 // If aggressive_shape_inference is true, nodes are executed on the host to 236 // identify output values when possible and does other aggressive strategies. 237 // This may cause incorrectness in graph analyses, but is useful for simulation 238 // or scheduling. 239 // If include_input_tensor_values is true, the values of constant 240 // tensors will included in the input properties. 241 // If include_output_tensor_values is true, the values of constant tensors will 242 // be included in the output properties. 243 TF_CAPI_EXPORT extern void TF_InferStatically( 244 TF_GraphProperties* graph_properties, TF_Bool assume_valid_feeds, 245 TF_Bool aggressive_shape_inference, TF_Bool include_input_tensor_values, 246 TF_Bool include_output_tensor_values, TF_Status* s); 247 248 // Get the size of input OpInfo::TensorProperties given node name. 249 TF_CAPI_EXPORT extern void TF_GetInputPropertiesListSize( 250 TF_GraphProperties* graph_properties, const char* name, int* num_values, 251 TF_Status* status); 252 253 // Get the size of output OpInfo::TensorProperties given node name. 254 TF_CAPI_EXPORT extern void TF_GetOutputPropertiesListSize( 255 TF_GraphProperties* graph_properties, const char* name, int* num_values, 256 TF_Status* status); 257 258 // Get a list of input OpInfo::TensorProperties given node name. 259 // Return the serialized list `properties`. 260 TF_CAPI_EXPORT extern void TF_GetInputPropertiesList( 261 TF_GraphProperties* graph_properties, const char* name, 262 TF_Buffer** properties, int num_values, TF_Status* status); 263 264 // Get a list of output OpInfo::TensorProperties given node name. 265 // Return the serialized list `properties`. 266 TF_CAPI_EXPORT extern void TF_GetOutputPropertiesList( 267 TF_GraphProperties* graph_properties, const char* name, 268 TF_Buffer** properties, int num_values, TF_Status* status); 269 270 // Helper to maintain a map between function names in a given 271 // FunctionDefLibrary and function definitions. 272 // Typical use case, is to look up an OpDef by type name. 273 typedef struct TF_FunctionLibraryDefinition TF_FunctionLibraryDefinition; 274 275 // Create NewFunctionLibraryDefinition. 276 TF_CAPI_EXPORT extern TF_FunctionLibraryDefinition* 277 TF_NewFunctionLibraryDefinition(const TF_Buffer* graph_buf, TF_Status* status); 278 279 // Delete NewFunctionLibraryDefinition. 280 TF_CAPI_EXPORT extern void TF_DeleteFunctionLibraryDefinition( 281 TF_FunctionLibraryDefinition* fn_lib); 282 283 // Shorthand for calling LookUp to get the OpDef from FunctionLibraryDefinition 284 // given op name. The returned OpDef is represented by TF_Buffer. 285 TF_CAPI_EXPORT extern void TF_LookUpOpDef(TF_FunctionLibraryDefinition* fn_lib, 286 const char* name, TF_Buffer* buf, 287 TF_Status* s); 288 289 #ifdef __cplusplus 290 } // extern "C" 291 #endif 292 293 #endif // TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_ 294