1 /* Copyright 2015 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 16 #ifndef TENSORFLOW_CORE_FRAMEWORK_FUNCTION_H_ 17 #define TENSORFLOW_CORE_FRAMEWORK_FUNCTION_H_ 18 19 #include <vector> 20 #include "tensorflow/core/framework/attr_value.pb.h" 21 #include "tensorflow/core/framework/attr_value_util.h" 22 #include "tensorflow/core/framework/function.pb.h" 23 #include "tensorflow/core/framework/node_def_util.h" 24 #include "tensorflow/core/framework/op.h" 25 #include "tensorflow/core/framework/op_kernel.h" 26 #include "tensorflow/core/framework/selective_registration.h" 27 #include "tensorflow/core/framework/types.h" 28 #include "tensorflow/core/lib/gtl/flatmap.h" 29 #include "tensorflow/core/lib/hash/hash.h" 30 #include "tensorflow/core/platform/env.h" 31 #include "tensorflow/core/platform/macros.h" 32 #include "tensorflow/core/platform/mutex.h" 33 #include "tensorflow/core/platform/protobuf.h" 34 #include "tensorflow/core/protobuf/config.pb.h" 35 36 namespace tensorflow { 37 38 class CancellationManager; 39 class CollectiveExecutor; 40 class DeviceSet; 41 class Graph; 42 class GraphDef; 43 class OpKernel; 44 class ProcessFunctionLibraryRuntime; 45 class ResourceMgr; 46 class Rendezvous; 47 class ScopedStepContainer; 48 class StepStatsCollectorInterface; 49 class Node; 50 51 // FunctionDefHelper::Create is a convenient helper to construct a 52 // FunctionDef proto. 53 // E.g., 54 // FunctionDef my_func = FunctionDefHelper::Create( 55 // "my_func_name", 56 // {"x:T", "y:T" /* one string per argument */}, 57 // {"z:T" /* one string per return value */}, 58 // {"T: {float, double}" /* one string per attribute */}, 59 // { 60 // {{"o"}, "Mul", {"x", "y"}, {{"T", "$T"}}} 61 // /* one entry per function node */ 62 // }, 63 // /* Mapping between function returns and function node outputs. */ 64 // {{"z", "o:z"}}); 65 // 66 // For the old Function::Node approach, use FunctionDefHelper::Define() 67 // E.g., 68 // FunctionDef my_func = FunctionDefHelper::Define( 69 // "my_func_name", 70 // {"x:T", "y:T" /* one string per argument */}, 71 // {"z:T" /* one string per return value */}, 72 // {"T: {float, double}" /* one string per attribute */}, 73 // { 74 // {{"z"}, "Mul", {"x", "y"}, {{"T", "$T"}}} 75 // /* one entry per function node */ 76 // }); 77 class FunctionDefHelper { 78 public: 79 // AttrValueWrapper has copy constructors for the type T so that 80 // it's easy to construct a simple AttrValue proto. 81 // 82 // If T is a string type (const char*, string, or StringPiece), and 83 // it starts with "$", we construct a AttrValue of "placeholder". 84 // 85 // E.g., 86 // std::<string, AttrValueWrapper> x = {"T", "$T"} 87 // is a named attr value placeholder. 88 struct AttrValueWrapper { 89 AttrValue proto; 90 91 AttrValueWrapper() {} 92 93 template <typename T> 94 AttrValueWrapper(T val) { // NOLINT(runtime/explicit) 95 SetAttrValue(val, &proto); 96 } 97 98 private: 99 void InitFromString(StringPiece val); 100 }; 101 102 // Constructs an AttrValue.func given the "name" and "attrs". 103 static AttrValueWrapper FunctionRef( 104 const string& name, 105 gtl::ArraySlice<std::pair<string, AttrValueWrapper>> attrs); 106 static AttrValueWrapper FunctionRef(const string& name) { 107 return FunctionRef(name, {}); 108 } 109 110 // Node is used to construct FunctionDef.Node using initialization 111 // lists. E.g., 112 // Node n = {{"z"}, "Mul", {"x", "y"}, {{"T", "$T"}}}; // z = x * y 113 struct Node { 114 // When constructing a NodeDef, the first entry in ret is used as 115 // the node name, the remaining values are ignored. 116 std::vector<string> ret; 117 string op; 118 std::vector<string> arg; 119 std::vector<std::pair<string, AttrValueWrapper>> attr; 120 std::vector<string> dep; 121 string device; 122 123 NodeDef ToNodeDef() const; 124 }; 125 126 // Creates a FunctionDef from the given parameters. Node inputs must use 127 // function encoding (node_name:output_name[:output_index]). 128 // - `ret_def` holds a mapping from the function output names from `out_def` 129 // to the node outputs from `node_def`. 130 // - `control_ret_def` holds a mapping from the function control 131 // output names to the nodes from `node_def`. 132 static FunctionDef Create( 133 const string& function_name, gtl::ArraySlice<string> in_def, 134 gtl::ArraySlice<string> out_def, gtl::ArraySlice<string> attr_def, 135 gtl::ArraySlice<Node> node_def, 136 gtl::ArraySlice<std::pair<string, string>> ret_def, 137 gtl::ArraySlice<std::pair<string, string>> control_ret_def); 138 139 // Creates a FunctionDef from the given parameters. Node inputs must use 140 // function encoding (node_name:output_name[:output_index]). 141 // - `ret_def` holds a mapping from the function output names from `out_def` 142 // to the node outputs from `node_def`. 143 static FunctionDef Create(const string& function_name, 144 gtl::ArraySlice<string> in_def, 145 gtl::ArraySlice<string> out_def, 146 gtl::ArraySlice<string> attr_def, 147 gtl::ArraySlice<Node> node_def, 148 gtl::ArraySlice<std::pair<string, string>> ret_def); 149 150 // TODO(josh11b): Get rid of these and transition to the one above. 151 static FunctionDef Define(const string& function_name, 152 gtl::ArraySlice<string> arg_def, 153 gtl::ArraySlice<string> ret_def, 154 gtl::ArraySlice<string> attr_def, 155 gtl::ArraySlice<Node> node_def); 156 157 // Defines an anonymous function. I.e., its name is not relevant. 158 static FunctionDef Define(gtl::ArraySlice<string> arg_def, 159 gtl::ArraySlice<string> ret_def, 160 gtl::ArraySlice<string> attr_def, 161 gtl::ArraySlice<Node> node_def); 162 163 // Helpers to construct a constant scalar. 164 template <typename T> 165 static Node Const(const string& name, const T& val) { 166 Node n = {{name}, "Const"}; 167 const DataType dtype = DataTypeToEnum<T>::value; 168 n.attr.push_back({"dtype", dtype}); 169 Tensor t(dtype, TensorShape({})); 170 t.scalar<T>()() = val; 171 n.attr.push_back({"value", t}); 172 return n; 173 } 174 175 template <typename T> 176 static Node Const(const string& name, gtl::ArraySlice<T> vals) { 177 Node n = {{name}, "Const"}; 178 const DataType dtype = DataTypeToEnum<T>::value; 179 n.attr.push_back({"dtype", dtype}); 180 int64 num = vals.size(); 181 Tensor t(dtype, TensorShape({num})); 182 for (size_t i = 0; i < vals.size(); ++i) { 183 t.flat<T>()(i) = vals[i]; 184 } 185 n.attr.push_back({"value", t}); 186 return n; 187 } 188 }; 189 190 template <> 191 inline FunctionDefHelper::AttrValueWrapper::AttrValueWrapper(const char* val) { 192 InitFromString(val); 193 } 194 195 template <> 196 inline FunctionDefHelper::AttrValueWrapper::AttrValueWrapper( 197 const string& val) { 198 InitFromString(val); 199 } 200 201 template <> 202 inline FunctionDefHelper::AttrValueWrapper::AttrValueWrapper(StringPiece val) { 203 InitFromString(val); 204 } 205 206 // Instantiate a function. 207 // 208 // "fdef" encodes a TF function with some attrs in fdef.signature.attr 209 // containing placeholders. InstantiateFunction binds these 210 // placeholders and produces an instantiated function encoded in 211 // "result.gdef". The value to substitute a placeholder is given by 212 // "attr_values", which is a map from a placeholder name to an attr 213 // value. 214 // 215 // InstantiateFunction calls "get_function" to find signatures of other 216 // functions and primitive ops. 217 218 // GetFunctionSignature(func name, opdef) returns OK if the func name is found 219 // and opdef is filled with a pointer to the corresponding signature 220 // (a OpDef proto). Otherwise, returns an error. 221 typedef std::function<Status(const string&, const OpDef**)> 222 GetFunctionSignature; 223 224 struct InstantiationResult { 225 DataTypeVector arg_types; 226 DataTypeVector ret_types; 227 std::vector<NodeDef> nodes; 228 }; 229 Status InstantiateFunction(const FunctionDef& fdef, AttrSlice attr_values, 230 GetFunctionSignature get_function, 231 InstantiationResult* result); 232 233 // Returns a debug string for a function definition. 234 // 235 // The returned text is multiple-line. It is intended to be 236 // human-readable rather than being friendly to parsers. It is _NOT_ 237 // intended to be the canonical string representation of "func_def". 238 // Particularly, it may not include all information presented in 239 // "func_def" (e.g., comments, description of the function arguments, 240 // etc.) 241 string DebugString(const FunctionDef& func_def); 242 string DebugString(const GraphDef& instantiated_func_def); 243 string DebugString(gtl::ArraySlice<NodeDef> instantiated_func_nodes); 244 245 // Returns a debug string for a top level graph (the main program and 246 // its supporting functions defined in its library). 247 string DebugStringWhole(const GraphDef& gdef); 248 249 // Returns true if f1 == f2. Compares all fields, including descriptions. Order 250 // of NodeDefs doesn't matter. 251 bool FunctionDefsEqual(const FunctionDef& f1, const FunctionDef& f2); 252 253 // Return a hash of `fdef` that is consistent with FunctionDefsEqual method. 254 // In other words, if two fdefs compare equal, their hash values will be the 255 // same. 256 uint64 FunctionDefHash(const FunctionDef& fdef); 257 258 class CallFrameInterface { 259 public: 260 virtual ~CallFrameInterface() {} 261 262 virtual size_t num_args() const = 0; 263 virtual size_t num_retvals() const = 0; 264 265 virtual Status GetArg(int index, Tensor* val) const = 0; 266 virtual Status SetRetval(int index, const Tensor& val) = 0; 267 }; 268 269 // Represents a function call frame. I.e., the data structure used to 270 // pass arguments to a function and retrieve its results. 271 // 272 // Runtime must arrange accesses to one FunctionCallFrame s.t. 273 // 1. SetArgs() happens before any GetArg(); 274 // 2. GetRetvals happens after all SetRetval(); 275 class FunctionCallFrame : public CallFrameInterface { 276 public: 277 FunctionCallFrame(DataTypeSlice arg_types, DataTypeSlice ret_types); 278 ~FunctionCallFrame(); 279 280 // Caller methods. 281 Status SetArgs(gtl::ArraySlice<Tensor> args); 282 Status GetRetvals(std::vector<Tensor>* rets) const; 283 284 // Moves the return values from the frame to rets. If allow_dead_tensors is 285 // false it will fail if any of the retvals do not have a value. 286 Status ConsumeRetvals(std::vector<Tensor>* rets, bool allow_dead_tensors); 287 288 size_t num_args() const override { return arg_types_.size(); } 289 size_t num_retvals() const override { return ret_types_.size(); } 290 291 // Callee methods. 292 Status GetArg(int index, Tensor* val) const override; 293 Status SetRetval(int index, const Tensor& val) override; 294 295 private: 296 DataTypeVector arg_types_; 297 DataTypeVector ret_types_; 298 gtl::InlinedVector<Tensor, 4> args_; 299 struct Retval { 300 bool has_val = false; 301 Tensor val; 302 }; 303 gtl::InlinedVector<Retval, 4> rets_; 304 305 TF_DISALLOW_COPY_AND_ASSIGN(FunctionCallFrame); 306 }; 307 308 // Helper to maintain a map between function names in a given 309 // FunctionDefLibrary and function definitions. 310 // 311 // This class is thread-safe. 312 class FunctionLibraryDefinition : public OpRegistryInterface { 313 public: 314 // Note: This constructor grabs `lib_def`'s lock in shared mode. 315 FunctionLibraryDefinition(const FunctionLibraryDefinition& lib_def); 316 FunctionLibraryDefinition(const OpRegistryInterface* default_registry, 317 const FunctionDefLibrary& lib_def); 318 ~FunctionLibraryDefinition() override; 319 320 FunctionLibraryDefinition& operator=(const FunctionLibraryDefinition&) = 321 delete; 322 323 // Returns True if the library contains `func`, False otherwise. 324 bool Contains(const string& func) const; 325 326 // Returns nullptr if "func" is not defined in "lib_def". Otherwise, 327 // returns its definition proto. 328 // 329 // NB: This function returns a borrowed pointer, which can be invalidated by a 330 // subsequent call to `ReplaceFunction()` with the given name. 331 const FunctionDef* Find(const string& func) const LOCKS_EXCLUDED(mu_); 332 333 // Adds function definition 'fdef' to this function library. 334 // Returns status 'ok' on success, or error otherwise. This is a no-op if 335 // 'fdef' already exists in this function library. 336 // If 'fdef' is successfully added to the library, it will be accessible 337 // from 'LookUp' and included in the proto returned by 'ToProto'. 338 // This operation is atomic. 339 Status AddFunctionDef(const FunctionDef& fdef) LOCKS_EXCLUDED(mu_); 340 341 // Adds gradient definition 'grad' to this function library. 342 // This is a no-op if 'grad' already exists in this function library. 343 // If 'grad' is successfully added, it will be accessible via 'FindGradient' 344 // and included in the proto returned by 'ToProto'. 345 // This operation is atomic. 346 Status AddGradientDef(const GradientDef& grad) LOCKS_EXCLUDED(mu_); 347 348 // Replaces the function corresponding to `func` with `fdef`. Returns 349 // a non-OK status if "func" was not found in the library, OK otherwise. 350 // Please be careful when replacing function: make sure all previous pointers 351 // returned by `Find()` are no longer in use. 352 Status ReplaceFunction(const string& func, const FunctionDef& fdef); 353 354 // Replaces the gradient corresponding to `grad.function_name()`. Returns 355 // a non-OK status if "grad.function_name()" was not found in the library, OK 356 // otherwise. 357 Status ReplaceGradient(const GradientDef& grad); 358 359 // Removes the function corresponding to 'func'. Returns a non-OK status if 360 // 'func' was not found in the library, OK otherwise. 361 // Please be careful when removing function: make sure there are no other 362 // nodes using the function, and all previous pointers returned by `Find()` 363 // are no longer in use. 364 Status RemoveFunction(const string& func); 365 366 // Adds the functions and gradients in 'other' to this function library. 367 // Duplicate functions and gradients are ignored. 368 // This operation is atomic. 369 Status AddLibrary(const FunctionLibraryDefinition& other) LOCKS_EXCLUDED(mu_); 370 371 // Adds the functions and gradients in 'lib_def' to this function library. 372 // Duplicate functions and gradients are ignored. 373 // This operation is atomic. 374 Status AddLibrary(const FunctionDefLibrary& lib_def) LOCKS_EXCLUDED(mu_); 375 376 // If the gradient function for 'func' is specified explicitly in 377 // the library, returns the gradient function name. Otherwise, 378 // returns an empty string. 379 string FindGradient(const string& func) const LOCKS_EXCLUDED(mu_); 380 381 // OpRegistryInterface method. Useful for constructing a Graph. 382 // 383 // If "op" is defined in the library, returns its signature. 384 // Otherwise, assume "op" is a primitive op and returns its op 385 // signature and shape inference function. 386 // 387 // NB: This function outputs a borrowed pointer, which can be invalidated by a 388 // subsequent call to `ReplaceFunction()` with the given name. 389 Status LookUp(const string& op_type_name, 390 const OpRegistrationData** op_reg_data) const override 391 LOCKS_EXCLUDED(mu_); 392 393 // Generates new function name with the specified prefix that is unique 394 // across this library. 395 string UniqueFunctionName(StringPiece prefix) const LOCKS_EXCLUDED(mu_); 396 397 // Ops created for function arguments bear the name given by `kArgOp`; those 398 // created for return values bear the name given by `kRetOp`. 399 static constexpr const char* const kArgOp = "_Arg"; 400 static constexpr const char* const kDeviceArgOp = "_DeviceArg"; 401 static constexpr const char* const kRetOp = "_Retval"; 402 static constexpr const char* const kDeviceRetOp = "_DeviceRetval"; 403 static constexpr const char* const kIntsOnDeviceAttr = 404 "experimental_ints_on_device"; 405 406 static constexpr const char* const kGradientOp = "SymbolicGradient"; 407 static constexpr const char* const kFuncAttr = "f"; 408 409 // Given a node def 'ndef', inspects attributes of the callee 410 // function to derive the attribute 'value' for 'attr'. Returns OK 411 // iff the attribute is given by the function's definition. 412 // TODO(irving): Remove; keep only the const Node& version. 413 template <typename T> 414 Status GetAttr(const NodeDef& ndef, const string& attr, T* value) const; 415 416 // Given a node, inspects attributes of the callee function to derive the 417 // attribute 'value' for 'attr'. Returns OK iff the attribute is given by the 418 // function's definition. 419 template <typename T> 420 Status GetAttr(const Node& node, const string& attr, T* value) const; 421 422 // Returns a proto representation of the state of this function library. 423 FunctionDefLibrary ToProto() const LOCKS_EXCLUDED(mu_); 424 425 size_t num_functions() const { 426 tf_shared_lock l(mu_); 427 return function_defs_.size(); 428 } 429 430 // Returns all the function names in the FunctionLibraryDefinition. 431 std::vector<string> ListFunctionNames() const LOCKS_EXCLUDED(mu_); 432 433 const OpRegistryInterface* default_registry() const { 434 return default_registry_; 435 } 436 437 // Returns a copy of `*this` with only the subset of functions that are 438 // reachable from the nodes of `graph` or `func`. 439 FunctionLibraryDefinition ReachableDefinitions(const GraphDef& graph) const; 440 FunctionLibraryDefinition ReachableDefinitions(const FunctionDef& func) const; 441 442 private: 443 // Shape inference for functions is handled separately by ShapeRefiner. 444 445 struct FunctionDefAndOpRegistration { 446 FunctionDefAndOpRegistration(const FunctionDef& fdef_in); 447 448 FunctionDef fdef; 449 OpRegistrationData op_registration_data; 450 }; 451 452 const FunctionDef* FindHelper(const string& func) const 453 SHARED_LOCKS_REQUIRED(mu_); 454 string FindGradientHelper(const string& func) const 455 SHARED_LOCKS_REQUIRED(mu_); 456 457 // Same as AddFunctionDef/AddGradientDef except these methods set 458 // `added` to true if the `fdef`/`grad` were actually added to this. 459 Status AddFunctionDefHelper(const FunctionDef& fdef, bool* added) 460 EXCLUSIVE_LOCKS_REQUIRED(mu_); 461 Status AddGradientDefHelper(const GradientDef& grad, bool* added) 462 EXCLUSIVE_LOCKS_REQUIRED(mu_); 463 464 mutable mutex mu_; 465 const OpRegistryInterface* const default_registry_; 466 gtl::FlatMap<string, std::unique_ptr<FunctionDefAndOpRegistration>> 467 function_defs_ GUARDED_BY(mu_); 468 gtl::FlatMap<string, string> func_grad_ GUARDED_BY(mu_); 469 470 // Helper function for GetAttr. Returns the FunctionDef* to get the 471 // attr from. 472 const FunctionDef* GetAttrImpl(const NodeDef& ndef) const LOCKS_EXCLUDED(mu_); 473 474 // Remove all functions in `funcs` and all gradients of functions in 475 // `funcs_with_grads` from this library. 476 void Remove(const std::vector<string>& funcs, 477 const std::vector<string>& funcs_with_grads) 478 EXCLUSIVE_LOCKS_REQUIRED(mu_); 479 480 // Remove `func` from the library. Returns non-OK Status unless `func` is in 481 // the library. This should only be called when there is a guarantee that the 482 // function being removed hasn't been retrieved with `Find`. 483 Status RemoveFunctionHelper(const string& func) EXCLUSIVE_LOCKS_REQUIRED(mu_); 484 485 // Remove gradient of function `func` from the library. Returns non-OK Status 486 // unless `func` has a gradient. 487 Status RemoveGradient(const string& func) EXCLUSIVE_LOCKS_REQUIRED(mu_); 488 }; 489 490 // Forward declare. Defined in common_runtime/function.h 491 struct FunctionBody; 492 493 // Forward declare. Defined in common_runtime/device.h 494 class Device; 495 // Forward declare. Defined in common_runtime/device_mgr.h 496 class DeviceMgr; 497 498 class FunctionLibraryRuntime { 499 public: 500 virtual ~FunctionLibraryRuntime() {} 501 502 // Instantiate a function with the given "attrs". 503 // 504 // Returns OK and fills in "handle" if the instantiation succeeds. 505 // Otherwise returns an error and "handle" is undefined. 506 struct InstantiateOptions { 507 // The canonical device name of the device on which the function 508 // should be instantiated. If empty, the function will be 509 // instantiated on the local device. 510 string target; 511 512 // Should the function be instantiated as a multi-device function? 513 bool is_multi_device_function = false; 514 515 // For multi-device functions, a vector of canonical device names for 516 // function's inputs. The device of resource inputs must be the device 517 // backing the resource, not the CPU device backing the resource handle. 518 // Must have the same length as number of inputs to the function. 519 std::vector<string> input_devices; 520 521 // For multi-device functions, a vector of canonical device names for 522 // function's outputs. The device of resource outputs should be the CPU 523 // device, not the device backing the resource. 524 // If specified, must have the same length as the number of function 525 // outputs. 526 // If not specified, output devices are picked automatically. If operations 527 // producing the output tensors have explicit device specification, they 528 // will be respected. These device specifications must identify a unique 529 // device, i.e. a general specification like "job:foo" matching multiple 530 // devices will result in an error. 531 std::vector<string> output_devices; 532 533 // This interface is EXPERIMENTAL and subject to change. 534 // 535 // If non-null, the runtime will use `overlay_lib` to resolve 536 // function(s) named in `function_name` and `attrs`. Otherwise, 537 // the runtime will use its internal library. 538 // NOTE(mrry): If provided, all functions defined in `overlay_lib` 539 // must be self-contained, and cannot refer to functions defined 540 // in other libraries. 541 // TODO(mrry): Provide a mechanism for sharing core functions 542 // between a set of libraries (e.g. by allowing a 543 // `FunctionLibraryDefinition` to store an `outer_scope` pointer 544 // and implementing name resolution across libraries). 545 const FunctionLibraryDefinition* overlay_lib = nullptr; 546 547 // This interface is EXPERIMENTAL and subject to change. 548 // 549 // If non-empty, the runtime will use `state_handle` to identify 550 // cached state related the instantiated function. Two functions 551 // of the same name and attrs, instantiated with the same 552 // `state_handle` will have the same handle and share the same 553 // state (in stateful kernels); and two functions with different 554 // values for `state_handle` will have independent state. 555 string state_handle; 556 557 // This interface is EXPERIMENTAL and subject to change. 558 // 559 // Instantiates the function using an executor of the given type. If empty, 560 // the default TensorFlow executor will be used. 561 string executor_type; 562 563 // If true, the runtime will attempt to create kernels for the function at 564 // instantiation time, rather than on the first run. This can be used to 565 // surface errors earlier. 566 bool create_kernels_eagerly = false; 567 568 // This interface is EXPERIMENTAL and subject to change. 569 // 570 // Instantiates the function with the provided config_proto. 571 ConfigProto config_proto; 572 573 // If provided, this optimization function will be invoked before 574 // the placer for multi-device functions. 575 std::function<Status(std::vector<string> /*ret_node_names*/, 576 std::vector<string> /*keep_node_names*/, 577 FunctionLibraryDefinition*, const DeviceSet&, 578 Device* /*cpu_device*/, std::unique_ptr<Graph>*)> 579 optimize_graph_fn; 580 581 // If set, partitioned functions will be added to `graph_collector`. 582 // `graph_collector` must be alive during the call to Instantiate. 583 GraphCollector* graph_collector = nullptr; 584 }; 585 typedef uint64 Handle; 586 virtual Status Instantiate(const string& function_name, AttrSlice attrs, 587 const InstantiateOptions& options, 588 Handle* handle) = 0; 589 Status Instantiate(const string& function_name, AttrSlice attrs, 590 Handle* handle) { 591 return Instantiate(function_name, attrs, {}, handle); 592 } 593 594 // Releases state associated with the handle. 595 virtual Status ReleaseHandle(Handle handle) = 0; 596 597 // Returns the function body for the instantiated function given its 598 // handle 'h'. Returns nullptr if "h" is not found. 599 // 600 // *this keeps the ownership of the returned object, which remains alive 601 // as long as *this. 602 virtual const FunctionBody* GetFunctionBody(Handle h) = 0; 603 604 // Asynchronously invokes the instantiated function identified by 605 // "handle". 606 // 607 // If function execution succeeds, "done" is called with OK and 608 // "*rets" is filled with the function's return values. Otheriwse, 609 // "done" is called with an error status. 610 // 611 // Does not take ownership of "rets". 612 // In the cross-process scenario, runner isn't used for making the Async 613 // RPC calls. 614 struct Options { 615 // The id of the step that is calling this function. 616 int64 step_id = 0; 617 Rendezvous* rendezvous = nullptr; 618 CancellationManager* cancellation_manager = nullptr; 619 CollectiveExecutor* collective_executor = nullptr; 620 ScopedStepContainer* step_container = nullptr; 621 StepStatsCollectorInterface* stats_collector = nullptr; 622 623 std::function<void(std::function<void()>)>* runner = nullptr; 624 625 // Parameters for remote function execution. 626 bool remote_execution = false; 627 string source_device = ""; // Fully specified device name. 628 629 // Allocator attributes specifying where the args are / rets should be put. 630 // These should either be {} or match the length of args / retvals. If {}, 631 // the default allocator attributes will be assumed for all args / retvals. 632 std::vector<AllocatorAttributes> args_alloc_attrs; 633 std::vector<AllocatorAttributes> rets_alloc_attrs; 634 635 // If true, we create a new IntraProcessRendezvous, else use the existing 636 // one. 637 bool create_rendezvous = false; 638 639 // If True, allow returning dead tensors. 640 bool allow_dead_tensors = false; 641 }; 642 typedef std::function<void(const Status&)> DoneCallback; 643 virtual void Run(const Options& opts, Handle handle, 644 gtl::ArraySlice<Tensor> args, std::vector<Tensor>* rets, 645 DoneCallback done) = 0; 646 virtual void Run(const Options& opts, Handle handle, 647 CallFrameInterface* call_frame, DoneCallback done) = 0; 648 649 // Creates a "kernel" for the given node def "ndef". 650 // 651 // If succeeds, returns OK and the caller takes the ownership of the 652 // returned "*kernel". Otherwise, returns an error. 653 virtual Status CreateKernel(const NodeDef& ndef, OpKernel** kernel) = 0; 654 655 // Returns true iff the function named `function_name` is stateful. 656 // NOTE(mrry): This method assumes that the runtime is associated with a 657 // default function library, and looks up `function_name` in that library. 658 // It does not support overlay libraries. 659 virtual bool IsStateful(const string& function_name) = 0; 660 661 // Returns the device on which the function executes. 662 virtual Device* device() = 0; 663 664 // Returns the default runner in which the ops should be launched. If the 665 // device on which the function executes has a private thread pool, return 666 // runner on the device local thread pool. 667 virtual std::function<void(std::function<void()>)>* runner() = 0; 668 669 // Get the DeviceMgr from which the device was obtained. 670 virtual const DeviceMgr* device_mgr() const = 0; 671 672 // Returns the function library definition that backs this runtime. 673 // NOTE(mrry): The returned library definition is the default function library 674 // for this runtime. The runtime may instantiate functions from separate 675 // overlay libraries, which are not returned by this function. 676 virtual const FunctionLibraryDefinition* GetFunctionLibraryDefinition() 677 const = 0; 678 679 // Returns the environment on which the function executes. 680 virtual Env* env() = 0; 681 682 // Returns a debug string showing the definition of the function of 683 // 'handle'. 684 virtual string DebugString(Handle handle) = 0; 685 686 // Returns the graph version number. 687 virtual int graph_def_version() = 0; 688 689 typedef uint64 LocalHandle; 690 691 virtual Status Clone(std::unique_ptr<FunctionLibraryDefinition>* out_lib_def, 692 std::unique_ptr<ProcessFunctionLibraryRuntime>* out_pflr, 693 FunctionLibraryRuntime** out_flr) = 0; 694 695 // Returns the name of the executor class (in the sense of 696 // `ExecutorFactory::GetFactory()`) that will be used based on the given 697 // dynamic `options` and static `attrs`. If none is specified, this method 698 // will return an empty string, which leaves the decision up to the runtime. 699 static string ExecutorType(const InstantiateOptions& options, 700 AttrSlice attrs); 701 }; 702 703 // Returns a canonicalized string for the instantiation of the 704 // function of the given "name", attributes "attrs", and "options". 705 // 706 // The returned string is guaranteed to be stable within one address 707 // space. But it may be change as the implementation 708 // evolves. Therefore, it should not be persisted or compared across 709 // address spaces. 710 string Canonicalize(const string& funcname, AttrSlice attrs, 711 const FunctionLibraryRuntime::InstantiateOptions& options); 712 inline string Canonicalize(const string& funcname, AttrSlice attrs) { 713 return Canonicalize(funcname, attrs, {}); 714 } 715 716 const FunctionLibraryRuntime::Handle kInvalidHandle = -1; 717 const FunctionLibraryRuntime::LocalHandle kInvalidLocalHandle = -1; 718 typedef std::function<Status(FunctionLibraryRuntime*, const NodeDef&, 719 std::unique_ptr<OpKernel>*)> 720 CustomKernelCreator; 721 722 // Used to instantiate and run functions in a distributed system. 723 class DistributedFunctionLibraryRuntime { 724 public: 725 virtual ~DistributedFunctionLibraryRuntime() {} 726 727 // The _target attr in attrs determines where the function is instantiated. 728 virtual Status Instantiate( 729 const string& function_name, const FunctionLibraryDefinition& lib_def, 730 AttrSlice attrs, 731 const FunctionLibraryRuntime::InstantiateOptions& options, 732 FunctionLibraryRuntime::LocalHandle* handle) = 0; 733 734 // opts.runner isn't used for execution. 735 virtual void Run(const FunctionLibraryRuntime::Options& opts, 736 FunctionLibraryRuntime::LocalHandle handle, 737 gtl::ArraySlice<Tensor> args, std::vector<Tensor>* rets, 738 FunctionLibraryRuntime::DoneCallback done) = 0; 739 }; 740 741 // Extracts the actual type from "attr_values" based on its definition 742 // "arg_def". 743 // 744 // If "arg_def" is a N*T type, *is_type_list is set to false, and 745 // *dtypes is set to be a vector of size N and each element is T. 746 // 747 // If "arg_def" is a list(type), *is_type_list is set to true, and 748 // *dtypes is set to be a vector of types specified in attrs for 749 // arg_def. 750 // 751 // Otherwise (arg_def is a simple type T), *is_type_list is set to 752 // false, and *dtypes is set to a single element vector, whose only 753 // element is T. 754 Status ArgNumType(AttrSlice attrs, const OpDef::ArgDef& arg_def, 755 bool* is_type_list, DataTypeVector* dtypes); 756 757 // To register a gradient function for a builtin op, one should use 758 // REGISTER_OP_GRADIENT(<op_name>, <c++ grad factory>); 759 // 760 // Typically, the c++ grad factory is a plan function that can be 761 // converted into ::tensorflow::gradient::Creator, which is 762 // std::function<Status(const AttrSlice&, FunctionDef*)>. 763 // 764 // A ::tensorflow::gradient::Creator should populate in FunctionDef* with a 765 // definition of a brain function which compute the gradient for the 766 // <op_name> when the <op_name> is instantiated with the given attrs. 767 // 768 // E.g., 769 // 770 // Status MatMulGrad(const AttrSlice& attrs, FunctionDef* g) { 771 // bool transpose_a; 772 // TF_RETURN_IF_ERROR(attrs.Get("transpose_a", &transpose_a)); 773 // bool transpose_b; 774 // TF_RETURN_IF_ERROR(attrs.Get("transpose_b", &transpose_b)); 775 // DataType dtype; 776 // TF_RETURN_IF_ERROR(attrs.Get("dtype", &dtype)); 777 // if (!transpose_a && !transpose_b) { 778 // *g = FunctionDefHelper::Define( 779 // "MatMulGrad", 780 // {"x:T ", "y:T", "dz:T"}, // Inputs to this function 781 // {"dx:T", "dy:T"}, // Outputs from this function 782 // {"T: {float, double}"}, // Attributes needed by this function 783 // { 784 // {{"x_t"}, "Transpose", {"x"}, {{"T", "$T"}}}, 785 // {{"y_t"}, "Transpose", {"y"}, {{"T", "$T"}}}, 786 // {{"dx"}, "MatMul", {"dz", "y_t"}, {{"T", "$T"}}}, 787 // {{"dy"}, "MatMul", {"x_", "dz"}, {{"T", "$T"}}}, 788 // }); 789 // } else { 790 // ... ... 791 // } 792 // return Status::OK(); 793 // } 794 // 795 // NOTE: $T is substituted with the type variable "T" when the 796 // gradient function MatMul is instantiated. 797 // 798 // TODO(zhifengc): Better documentation somewhere. 799 800 // Macros to define a gradient function factory for a primitive 801 // operation. 802 #define REGISTER_OP_GRADIENT(name, fn) \ 803 REGISTER_OP_GRADIENT_UNIQ_HELPER(__COUNTER__, name, fn) 804 805 #define REGISTER_OP_NO_GRADIENT(name) \ 806 REGISTER_OP_GRADIENT_UNIQ_HELPER(__COUNTER__, name, nullptr) 807 808 #define REGISTER_OP_GRADIENT_UNIQ_HELPER(ctr, name, fn) \ 809 REGISTER_OP_GRADIENT_UNIQ(ctr, name, fn) 810 811 #define REGISTER_OP_GRADIENT_UNIQ(ctr, name, fn) \ 812 static bool unused_grad_##ctr TF_ATTRIBUTE_UNUSED = \ 813 SHOULD_REGISTER_OP_GRADIENT && \ 814 ::tensorflow::gradient::RegisterOp(name, fn) 815 816 namespace gradient { 817 // Register a gradient creator for the "op". 818 typedef std::function<Status(const AttrSlice& attrs, FunctionDef*)> Creator; 819 bool RegisterOp(const string& op, Creator func); 820 821 // Returns OK the gradient creator for the "op" is found (may be 822 // nullptr if REGISTER_OP_NO_GRADIENT is used. 823 Status GetOpGradientCreator(const string& op, Creator* creator); 824 }; // namespace gradient 825 826 // Declare explicit instantiations of GetAttr 827 #define GET_ATTR(T) \ 828 extern template Status FunctionLibraryDefinition::GetAttr( \ 829 const Node&, const string&, T*) const; \ 830 extern template Status FunctionLibraryDefinition::GetAttr( \ 831 const NodeDef&, const string&, T*) const; 832 GET_ATTR(string) 833 GET_ATTR(bool) 834 #undef GET_ATTR 835 836 } // end namespace tensorflow 837 838 #endif // TENSORFLOW_CORE_FRAMEWORK_FUNCTION_H_ 839