diff options
Diffstat (limited to 'tensorflow/compiler/xla/client/local_client.h')
| -rw-r--r-- | tensorflow/compiler/xla/client/local_client.h | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index e9eeaa0aa2..9f985ed527 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -37,6 +37,14 @@ namespace xla { // LocalClient::Compile. class ExecutableBuildOptions { public: + // If set, this is the platform to build the computation for. This must match + // the underlying platform of the service. A value of nullptr indicates the + // option has not been set. + // + // TODO(b/28616830): Support multiple platforms. + ExecutableBuildOptions& set_platform(perftools::gputools::Platform* platform); + perftools::gputools::Platform* platform() const; + // If set, this is the device to build the computation for. Valid // device_ordinal values are: 0 to # of devices - 1. These values are // identical to the device ordinal values used by StreamExecutor. The built @@ -53,10 +61,18 @@ class ExecutableBuildOptions { ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout); const Shape* result_layout() const; + // If set, the executable will be built to output a hybrid + // ShapedBuffer with top-level tuple pointers in host memory and + // result buffers in device memory. + ExecutableBuildOptions& set_has_hybrid_result(bool has_hybrid_result); + bool has_hybrid_result() const; + private: + perftools::gputools::Platform* platform_ = nullptr; int device_ordinal_ = -1; Shape result_layout_; bool result_layout_set_ = false; + bool has_hybrid_result_ = true; }; class LocalExecutable { |