1cf07600653c01675fe339d604f42000074d9a976Benoit Steiner/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 3cf07600653c01675fe339d604f42000074d9a976Benoit SteinerLicensed under the Apache License, Version 2.0 (the "License"); 4cf07600653c01675fe339d604f42000074d9a976Benoit Steineryou may not use this file except in compliance with the License. 5cf07600653c01675fe339d604f42000074d9a976Benoit SteinerYou may obtain a copy of the License at 6cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 7cf07600653c01675fe339d604f42000074d9a976Benoit Steiner http://www.apache.org/licenses/LICENSE-2.0 8cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 9cf07600653c01675fe339d604f42000074d9a976Benoit SteinerUnless required by applicable law or agreed to in writing, software 10cf07600653c01675fe339d604f42000074d9a976Benoit Steinerdistributed under the License is distributed on an "AS IS" BASIS, 11cf07600653c01675fe339d604f42000074d9a976Benoit SteinerWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12cf07600653c01675fe339d604f42000074d9a976Benoit SteinerSee the License for the specific language governing permissions and 13cf07600653c01675fe339d604f42000074d9a976Benoit Steinerlimitations under the License. 14cf07600653c01675fe339d604f42000074d9a976Benoit Steiner==============================================================================*/ 15cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 16cf07600653c01675fe339d604f42000074d9a976Benoit Steiner%include "tensorflow/python/platform/base.i" 174e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner%include <std_shared_ptr.i> 184e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner%include "item.i" 194e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner 204e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner// Wrap the cluster into an object that swig can manipulate. This ensures it will call the object 214e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner// destructor upon garbage collection instead of leaking memory. 224e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstruct GCluster { 234e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner std::shared_ptr<tensorflow::grappler::Cluster> cluster_; 244e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner}; 25cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 26f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang%{ 27f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang#include "tensorflow/core/protobuf/device_properties.pb.h" 28f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang 29f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhangtemplate <> 30f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhangbool _PyObjAs(PyObject *input, tensorflow::NamedDevice *out) { 31f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang char* c_string; 32f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang Py_ssize_t py_size; 33f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang if (PyBytes_AsStringAndSize(input, &c_string, &py_size) == -1) { 34f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang // Python has raised an error (likely TypeError or UnicodeEncodeError). 35f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang return false; 36f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang } 37f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang 38f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::NamedDevice named_device; 39f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang if (!named_device.ParseFromString(string(c_string, py_size))) { 40f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang PyErr_SetString( 41f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang PyExc_TypeError, 42f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang "The NamedDevice could not be parsed as a valid protocol buffer"); 43f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang return false; 44f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang } 45f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang if (out) *out = named_device; 46f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang return true; 47f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang} 48f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang%} 49f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang 50f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang%typemap(in) const std::vector<tensorflow::NamedDevice>& (std::vector<tensorflow::NamedDevice> temp) { 51f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang if (!tf_vector_input_helper($input, &temp, &_PyObjAs<tensorflow::NamedDevice>)) { 52f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang SWIG_fail; 53f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang } 54f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang $1 = &temp; 55f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang} 56f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang 575582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner%typemap(in) const tensorflow::NamedDevice& (tensorflow::NamedDevice temp) { 585582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner char* c_string; 595582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner Py_ssize_t py_size; 605582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner if (PyBytes_AsStringAndSize($input, &c_string, &py_size) == -1) { 615582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner // Python has raised an error (likely TypeError or UnicodeEncodeError). 625582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner SWIG_fail; 635582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner } 645582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner 655582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner if (!temp.ParseFromString(string(c_string, py_size))) { 665582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner PyErr_SetString( 675582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner PyExc_TypeError, 685582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner "The NamedDevice could not be parsed as a valid protocol buffer"); 695582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner SWIG_fail; 705582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner } 715582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner $1 = &temp; 725582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner} 735582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner 74cf07600653c01675fe339d604f42000074d9a976Benoit Steiner%typemap(in) const tensorflow::RunMetadata& (tensorflow::RunMetadata temp) { 75cf07600653c01675fe339d604f42000074d9a976Benoit Steiner char* c_string; 76cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_ssize_t py_size; 77cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (PyBytes_AsStringAndSize($input, &c_string, &py_size) == -1) { 78cf07600653c01675fe339d604f42000074d9a976Benoit Steiner // Python has raised an error (likely TypeError or UnicodeEncodeError). 79cf07600653c01675fe339d604f42000074d9a976Benoit Steiner SWIG_fail; 80cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 81cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 82cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (!temp.ParseFromString(string(c_string, py_size))) { 83cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyErr_SetString( 84cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyExc_TypeError, 85f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang "The RunMetadata could not be parsed as a valid protocol buffer"); 86cf07600653c01675fe339d604f42000074d9a976Benoit Steiner SWIG_fail; 87cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 88cf07600653c01675fe339d604f42000074d9a976Benoit Steiner $1 = &temp; 89cf07600653c01675fe339d604f42000074d9a976Benoit Steiner} 90cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 91cf07600653c01675fe339d604f42000074d9a976Benoit Steiner%typemap(in) const string& (string temp) { 92cf07600653c01675fe339d604f42000074d9a976Benoit Steiner char *buf; 93cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_ssize_t len; 94cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (PyBytes_AsStringAndSize($input, &buf, &len) == -1) return NULL; 95cf07600653c01675fe339d604f42000074d9a976Benoit Steiner temp.assign(buf, len); 96cf07600653c01675fe339d604f42000074d9a976Benoit Steiner $1 = &temp; 97cf07600653c01675fe339d604f42000074d9a976Benoit Steiner} 98cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 99cf07600653c01675fe339d604f42000074d9a976Benoit Steiner%{ 1004e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner#include <memory> 10163a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner#include <vector> 102cf07600653c01675fe339d604f42000074d9a976Benoit Steiner#include "tensorflow/core/grappler/devices.h" 103dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner#include "tensorflow/core/grappler/utils.h" 104cf07600653c01675fe339d604f42000074d9a976Benoit Steiner#include "tensorflow/core/grappler/clusters/single_machine.h" 105f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang#include "tensorflow/core/grappler/clusters/virtual_cluster.h" 1068fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner#include "tensorflow/core/grappler/costs/graph_memory.h" 107cf07600653c01675fe339d604f42000074d9a976Benoit Steiner#include "tensorflow/core/grappler/costs/op_performance_data.pb.h" 108cf07600653c01675fe339d604f42000074d9a976Benoit Steiner#include "tensorflow/core/grappler/costs/measuring_cost_estimator.h" 109cf07600653c01675fe339d604f42000074d9a976Benoit Steiner#include "tensorflow/core/grappler/costs/utils.h" 11063a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner#include "tensorflow/core/protobuf/device_properties.pb.h" 111dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner#include "tensorflow/core/framework/kernel_def.pb.h" 112dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner#include "tensorflow/core/framework/memory_types.h" 113cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 1144e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner// Provide the implementation of the GCluster struct here. 1154e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstruct GCluster { 1164e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GCluster() {} 1174e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GCluster(tensorflow::grappler::Cluster* cluster) : cluster_(cluster) {} 1184e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner 1194e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::grappler::Cluster* operator->() const { 1204e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner return cluster_.get(); 1214e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner } 1224e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::grappler::Cluster* get() const { 1234e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner return cluster_.get(); 1244e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner } 1254e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner bool is_none() const { 1264e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner return cluster_.get() == nullptr; 1274e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner } 1284e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner 1294e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner std::shared_ptr<tensorflow::grappler::Cluster> cluster_; 1304e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner}; 1314e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner 1324e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner 1334e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic GCluster TF_NewCluster(bool allow_soft_placement, 1344e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner bool disable_detailed_stats, TF_Status* out_status) { 1354e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner int num_cpu_cores = tensorflow::grappler::GetNumAvailableLogicalCPUCores(); 1364e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner int num_gpus = tensorflow::grappler::GetNumAvailableGPUs(); 137cf07600653c01675fe339d604f42000074d9a976Benoit Steiner int timeout_s = 60 * 10; 1384e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::grappler::Cluster* cluster_ = 139f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang new tensorflow::grappler::SingleMachine( 140f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang timeout_s, num_cpu_cores, num_gpus); 1414e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner cluster_->DisableDetailedStats(disable_detailed_stats); 1424e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner cluster_->AllowSoftPlacement(allow_soft_placement); 1437ebc013313876748af71fde5af03859b9728181aYao Zhang cluster_->SetNumWarmupSteps(10); 1444e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::Status status = cluster_->Provision(); 145cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::Set_TF_Status_from_Status(out_status, status); 1464e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner return GCluster(cluster_); 147cf07600653c01675fe339d604f42000074d9a976Benoit Steiner} 148cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 1494e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic GCluster TF_NewVirtualCluster( 150f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang const std::vector<tensorflow::NamedDevice>& named_devices, 151f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang TF_Status* out_status) { 152f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang std::unordered_map<string, tensorflow::DeviceProperties> devices; 153f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang for (const auto& named_device : named_devices) { 154f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang devices[named_device.name()]= named_device.properties(); 155f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang } 1564e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::grappler::Cluster*cluster_ = 157f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang new tensorflow::grappler::VirtualCluster(devices); 158f7837905c0a24b18841198863a64e8a878c53b44Benoit Steiner PyGILState_STATE gstate = PyGILState_Ensure(); 1594e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::Status status = cluster_->Provision(); 160f7837905c0a24b18841198863a64e8a878c53b44Benoit Steiner PyGILState_Release(gstate); 161f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::Set_TF_Status_from_Status(out_status, status); 1624e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner return GCluster(cluster_); 163f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang} 164f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang 1654e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic void TF_ShutdownCluster(GCluster cluster) { 166f7837905c0a24b18841198863a64e8a878c53b44Benoit Steiner PyGILState_STATE gstate = PyGILState_Ensure(); 167cf07600653c01675fe339d604f42000074d9a976Benoit Steiner cluster->Shutdown(); 168f7837905c0a24b18841198863a64e8a878c53b44Benoit Steiner PyGILState_Release(gstate); 169cf07600653c01675fe339d604f42000074d9a976Benoit Steiner} 170cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 171f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhangtensorflow::Status _GetOpPerformanceDataAndRunTime( 172f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang const tensorflow::grappler::GrapplerItem& item, 173f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::grappler::CostEstimator* cost_measure, 174f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::OpPerformanceList* op_performance_data, 175f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::grappler::Costs* costs) { 176cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::Status status = cost_measure->Initialize(item); 177cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (!status.ok()) return status; 178cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 179cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::CostGraphDef cost_graph; 180cf07600653c01675fe339d604f42000074d9a976Benoit Steiner TF_RETURN_IF_ERROR( 181cf07600653c01675fe339d604f42000074d9a976Benoit Steiner cost_measure->PredictCosts(item.graph, &cost_graph, costs)); 182cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 183cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (op_performance_data) { 184cf07600653c01675fe339d604f42000074d9a976Benoit Steiner *op_performance_data = tensorflow::grappler::CostGraphToOpPerformanceData( 185cf07600653c01675fe339d604f42000074d9a976Benoit Steiner cost_graph, item.graph); 186cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 187cf07600653c01675fe339d604f42000074d9a976Benoit Steiner return tensorflow::Status::OK(); 188cf07600653c01675fe339d604f42000074d9a976Benoit Steiner} 189cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 1904e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic PyObject* TF_ListDevices(GCluster cluster) { 19163a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner const std::unordered_map<string, tensorflow::DeviceProperties>& devices = cluster->GetDevices(); 192806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner PyGILState_STATE gstate = PyGILState_Ensure(); 19363a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner PyObject* result = PyList_New(devices.size()); 19463a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner int i = 0; 19563a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner for (auto& dev : devices) { 19663a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner tensorflow::NamedDevice d; 19763a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner d.set_name(dev.first); 19863a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner *d.mutable_properties() = dev.second; 19963a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner string dev_str = d.SerializeAsString(); 20063a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner PyObject* dev_obj = PyBytes_FromStringAndSize(dev_str.data(), 20163a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner dev_str.size()); 20263a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner PyList_SetItem(result, i, dev_obj); 20363a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner ++i; 20463a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner } 205806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner PyGILState_Release(gstate); 20663a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner return result; 20763a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner} 20863a69b8c4a30f4458b2696993463d1c78b5c922eBenoit Steiner 20981fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steinerstatic std::vector<string> TF_ListAvailableOps() { 21081fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner tensorflow::OpRegistry* registry = tensorflow::OpRegistry::Global(); 21181fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner std::vector<tensorflow::OpDef> ops; 21281fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner registry->GetRegisteredOps(&ops); 21381fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner std::vector<string> op_names; 21481fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner for (const tensorflow::OpDef& op : ops) { 21581fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner op_names.push_back(op.name()); 21681fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner } 21781fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner std::sort(op_names.begin(), op_names.end()); 21881fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner return op_names; 21981fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner} 22081fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steiner 221dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steinerstatic PyObject* TF_GetSupportedDevices(GCluster cluster, GItem item) { 222dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner if (cluster.is_none() || item.is_none()) { 223dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner Py_RETURN_NONE; 224dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 225dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const std::unordered_map<string, tensorflow::DeviceProperties>& devices = cluster->GetDevices(); 226dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::unordered_map<string, std::vector<string>> device_types; 227dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (const auto& dev : devices) { 228dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner device_types[dev.second.type()].push_back(dev.first); 229dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 230dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 231dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::unordered_map<string, std::set<string>> supported_device_types; 232dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::unordered_map<string, std::set<string>> device_restrictions; 233dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 234dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (const auto& node : item->graph.node()) { 235dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (const auto& dev : device_types) { 236dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const string& type = dev.first; 237dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner if (cluster->type() != "single_machine") { 238dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner // The actual kernel may not be linked in this binary. 239dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner supported_device_types[node.name()].insert(type); 240dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } else { 241dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner // Check the kernel capabilities 242dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const tensorflow::DeviceType dev_type(type); 243dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner tensorflow::Status s = tensorflow::FindKernelDef(dev_type, node, nullptr, nullptr); 244dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner if (s.ok()) { 245dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner supported_device_types[node.name()].insert(type); 246dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 247dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner // Check which inputs are restricted to reside on the host. 248dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner // TODO: extends this to support outputs as well 249dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner tensorflow::MemoryTypeVector inp_mtypes; 250dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner tensorflow::MemoryTypeVector out_mtypes; 251dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner s = tensorflow::MemoryTypesForNode(tensorflow::OpRegistry::Global(), dev_type, node, 252dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner &inp_mtypes, &out_mtypes); 253dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner if (s.ok()) { 254dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (int i = 0; i < inp_mtypes.size(); ++i) { 255dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner if (inp_mtypes[i] == tensorflow::HOST_MEMORY) { 256dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner device_restrictions[tensorflow::grappler::NodeName(node.input(i))].insert("CPU"); 257dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner break; 258dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 259dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 260dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 261dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 262dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 263dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 264dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 265dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 266dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner PyGILState_STATE gstate = PyGILState_Ensure(); 267dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner PyObject* result = PyDict_New(); 268dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 269dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (const auto& supported_dev : supported_device_types) { 270dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const string& node = supported_dev.first; 271dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::set<string> feasible; 272dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const auto it = device_restrictions.find(node); 273dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner if (it != device_restrictions.end()) { 274dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const std::set<string>& candidates = supported_dev.second; 275dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner const std::set<string>& valid = it->second; 276dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::set_intersection(candidates.begin(), candidates.end(), valid.begin(), valid.end(), 277dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::inserter(feasible, feasible.begin())); 278dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } else { 279dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner feasible = supported_dev.second; 280dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 281dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 282dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner std::vector<string> device_names; 283dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (const string& type : feasible) { 284dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner auto it = device_types.find(type); 285dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner CHECK(it != device_types.end()); 286dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (const string& name : it->second) { 287dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner device_names.push_back(name); 288dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 289dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 290dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 291dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner PyObject* dev = PyList_New(device_names.size()); 292dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner for (int i = 0; i < device_names.size(); ++i) { 293dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner PyList_SetItem(dev, i, PyString_FromString(device_names[i].c_str())); 294dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 295dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner CHECK_EQ(0, PyDict_SetItem(result, PyString_FromString(node.c_str()), dev)); 296dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner } 297dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner PyGILState_Release(gstate); 298dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner return result; 299dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner} 300dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 301dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steiner 3025582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steinerstatic double TF_EstimatePerformance(const tensorflow::NamedDevice& device) { 3035582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner tensorflow::grappler::OpLevelCostEstimator estimator; 3045582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner tensorflow::grappler::OpLevelCostEstimator::DeviceInfo info = 3055582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner estimator.GetDeviceInfo(device.properties()); 3065582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner return info.gigaops; 3075582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner} 3085582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steiner 309cf07600653c01675fe339d604f42000074d9a976Benoit Steinerstatic PyObject* TF_MeasureCosts( 3104e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GItem item, 3114e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GCluster cluster, 312cf07600653c01675fe339d604f42000074d9a976Benoit Steiner bool generate_timeline, TF_Status* out_status) { 313cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::OpPerformanceList op_performance_data; 314cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::StepStats step_stats; 315cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 3164e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner tensorflow::grappler::MeasuringCostEstimator cost_measure(cluster.get(), 10, 0); 317cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 318cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::grappler::Costs costs; 319f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::Status status = _GetOpPerformanceDataAndRunTime( 320f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang *item, &cost_measure, &op_performance_data, &costs); 321cf07600653c01675fe339d604f42000074d9a976Benoit Steiner double run_time = FLT_MAX; 322cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (status.ok()) { 323cf07600653c01675fe339d604f42000074d9a976Benoit Steiner run_time = static_cast<double>(costs.execution_time.count()) / 1e9; 324cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 325cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (generate_timeline) { 326cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::RunMetadata metadata; 327f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang tensorflow::Status s = cluster->Run( 328f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang item->graph, item->feed, item->fetch, &metadata); 329cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (s.ok()) { 330cf07600653c01675fe339d604f42000074d9a976Benoit Steiner step_stats = metadata.step_stats(); 331cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } else { 332cf07600653c01675fe339d604f42000074d9a976Benoit Steiner status = s; 333cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 334cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 335cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 336cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::Set_TF_Status_from_Status(out_status, status); 337cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (!status.ok()) { 338cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_RETURN_NONE; 339cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 340806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner PyGILState_STATE gstate = PyGILState_Ensure(); 341f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang PyObject* op_perf_objs = PyList_New( 342f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang op_performance_data.op_performance_size()); 343cf07600653c01675fe339d604f42000074d9a976Benoit Steiner for (int i = 0; i < op_performance_data.op_performance_size(); i++) { 344f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang string op_perf_str = 345f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang op_performance_data.op_performance(i).SerializeAsString(); 346cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyObject* op_perf_obj = PyBytes_FromStringAndSize(op_perf_str.data(), 347cf07600653c01675fe339d604f42000074d9a976Benoit Steiner op_perf_str.size()); 348cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyList_SetItem(op_perf_objs, i, op_perf_obj); 349cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 350cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 351cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyObject* run_time_obj = PyFloat_FromDouble(run_time); 352cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 353cf07600653c01675fe339d604f42000074d9a976Benoit Steiner string step_stats_str = step_stats.SerializeAsString(); 354cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyObject* metadata_obj = PyBytes_FromStringAndSize(step_stats_str.data(), 355cf07600653c01675fe339d604f42000074d9a976Benoit Steiner step_stats_str.size()); 356cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 357cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyObject* ret = PyTuple_New(3); 358cf07600653c01675fe339d604f42000074d9a976Benoit Steiner if (PyTuple_SetItem(ret, 0, op_perf_objs) != 0 || 359cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyTuple_SetItem(ret, 1, run_time_obj) != 0 || 360cf07600653c01675fe339d604f42000074d9a976Benoit Steiner PyTuple_SetItem(ret, 2, metadata_obj) != 0) { 361cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_DECREF(ret); 362cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_XDECREF(op_perf_objs); 363cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_XDECREF(run_time_obj); 364cf07600653c01675fe339d604f42000074d9a976Benoit Steiner Py_XDECREF(metadata_obj); 365cf07600653c01675fe339d604f42000074d9a976Benoit Steiner status = tensorflow::Status(tensorflow::error::Code::INTERNAL, 366cf07600653c01675fe339d604f42000074d9a976Benoit Steiner "Error setting return tuples."); 367cf07600653c01675fe339d604f42000074d9a976Benoit Steiner tensorflow::Set_TF_Status_from_Status(out_status, status); 368806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner Py_INCREF(Py_None); 369806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner ret = Py_None; 370cf07600653c01675fe339d604f42000074d9a976Benoit Steiner } 371806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner PyGILState_Release(gstate); 372cf07600653c01675fe339d604f42000074d9a976Benoit Steiner return ret; 373cf07600653c01675fe339d604f42000074d9a976Benoit Steiner} 374cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 3758fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner 3768fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steinerstatic PyObject* TF_DeterminePeakMemoryUsage( 3774e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GItem item, 3784e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GCluster cluster, 3798fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner TF_Status* out_status) { 3804e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner if (item.is_none() || cluster.is_none()) { 3818fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner tensorflow::Status status(tensorflow::error::Code::INTERNAL, 3828fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner "You need both a cluster and an item to determine peak memory usage"); 3838fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner tensorflow::Set_TF_Status_from_Status(out_status, status); 3848fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner Py_RETURN_NONE; 3858fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner } 3868fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner tensorflow::grappler::GraphMemory memory(*item); 3878fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner 3888fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner tensorflow::Status status; 3898fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner if (cluster->DetailedStatsEnabled()) { 3904e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner status = memory.InferDynamically(cluster.get()); 3918fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner } else { 3928fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner status = memory.InferStatically(cluster->GetDevices()); 3938fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner } 3948fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner if (!status.ok()) { 3958fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner tensorflow::Set_TF_Status_from_Status(out_status, status); 3968fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner Py_RETURN_NONE; 3978fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner } 3988fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner 399806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner PyGILState_STATE gstate = PyGILState_Ensure(); 4008fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyObject* result = PyDict_New(); 4018fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner for (const auto& device : cluster->GetDevices()) { 4028fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner const tensorflow::grappler::GraphMemory::MemoryUsage& usage = 4038fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner memory.GetPeakMemoryUsage(device.first); 4048fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyObject* per_device = PyList_New(usage.live_tensors.size()); 4058fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner for (int i = 0; i < usage.live_tensors.size(); ++i) { 4068fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner const auto& live_tensor = usage.live_tensors[i]; 4078fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyObject* live = PyTuple_New(5); 4088fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(live, 0, PyString_FromString(live_tensor.node.c_str())); 4098fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(live, 1, PyInt_FromLong(live_tensor.output_id)); 4108fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(live, 2, PyLong_FromLong(live_tensor.memory_used)); 4118fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(live, 3, PyLong_FromLong(live_tensor.allocation_time.count())); 4128fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(live, 4, PyLong_FromLong(live_tensor.deallocation_time.count())); 4138fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyList_SetItem(per_device, i, live); 4148fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner 4158fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner } 4168fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyObject* ret = PyTuple_New(2); 4178fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(ret, 0, PyLong_FromLong(usage.used_memory)); 4188fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyTuple_SetItem(ret, 1, per_device); 4198fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner PyDict_SetItem(result, PyString_FromString(device.first.c_str()), ret); 4208fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner } 421806754888188e40430bc96ad33c5f51282c2d338Benoit Steiner PyGILState_Release(gstate); 4228fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner return result; 4238fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner} 4248fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner 425cf07600653c01675fe339d604f42000074d9a976Benoit Steiner%} 426cf07600653c01675fe339d604f42000074d9a976Benoit Steiner 427cf07600653c01675fe339d604f42000074d9a976Benoit Steiner// Wrap these functions. 4284e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic GCluster TF_NewCluster( 429cf07600653c01675fe339d604f42000074d9a976Benoit Steiner bool allow_soft_placement, bool disable_detailed_stats, TF_Status* out_status); 4304e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic GCluster TF_NewVirtualCluster( 431f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang const std::vector<tensorflow::NamedDevice>& named_devices, 432f0d1abbf2389aa2a29fe6fd090ba68ab6b8fd76fYao Zhang TF_Status* out_status); 4334e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic void TF_ShutdownCluster(GCluster cluster); 4344e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steinerstatic PyObject* TF_ListDevices(GCluster cluster); 43581fe1b8700d839c00914fb3ade0cb626b97f6a08Benoit Steinerstatic std::vector<string> TF_ListAvailableOps(); 436dddbe5f43a7a0688089cf4fc9472ca8893460b3dBenoit Steinerstatic PyObject* TF_GetSupportedDevices(GCluster cluster, GItem item); 4375582df1a70c00e34c1254f798585cfd3ca05d90cBenoit Steinerstatic float TF_EstimatePerformance(const tensorflow::NamedDevice& device); 438cf07600653c01675fe339d604f42000074d9a976Benoit Steinerstatic PyObject* TF_MeasureCosts( 4394e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GItem item, GCluster cluster, 440cf07600653c01675fe339d604f42000074d9a976Benoit Steiner bool generate_timeline, TF_Status* out_status); 4418fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steinerstatic PyObject* TF_DeterminePeakMemoryUsage( 4424e5534d3d35a72b87902212e2847ca2871cc7b75Benoit Steiner GItem item, GCluster cluster, 4438fe6ea5f32c4ad5a5feb7f54d0fba8ddab4927caBenoit Steiner TF_Status* out_status); 444