tf_cuda_library(
name = "utils",
srcs = ["utils.cc"],
- hdrs = [
- "utils.h",
- ],
+ hdrs = ["utils.h"],
visibility = ["//visibility:public"],
deps = [
"//third_party/eigen3",
"//tensorflow/core:framework",
+ "//tensorflow/core:gpu_id",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
] + select({
"//tensorflow/core:core_cpu_lib",
"//tensorflow/core:direct_session",
"//tensorflow/core:framework",
+ "//tensorflow/core:gpu_id",
"//tensorflow/core:lib",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/kernels:ops_util",
#include "tensorflow/cc/training/queue_runner.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/grappler/clusters/utils.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/kernels/ops_util.h"
std::vector<DeviceAttributes> devices;
TF_RETURN_IF_ERROR(session_->ListDevices(&devices));
- int gpu_id = 0;
for (const auto& dev : devices) {
DeviceProperties attr;
if (dev.device_type() == "CPU") {
attr = GetLocalCPUInfo();
} else if (dev.device_type() == "GPU") {
- attr = GetLocalGPUInfo(gpu_id++);
+ DeviceNameUtils::ParsedName parsed;
+ if (!DeviceNameUtils::ParseFullName(dev.name(), &parsed)) {
+ return errors::InvalidArgument(
+ strings::StrCat("Not able to parse GPU device name: ", dev.name()));
+ }
+ attr = GetLocalGPUInfo(TfGpuId(parsed.id));
} else {
attr.set_type(dev.device_type());
}
#include "include/libxsmm.h"
#endif
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/cpu_info.h"
return device;
}
-DeviceProperties GetLocalGPUInfo(int gpu_id) {
+DeviceProperties GetLocalGPUInfo(TfGpuId tf_gpu_id) {
DeviceProperties device;
device.set_type("GPU");
#if GOOGLE_CUDA
cudaDeviceProp properties;
- cudaError_t error = cudaGetDeviceProperties(&properties, gpu_id);
+ CudaGpuId cuda_gpu_id = GpuIdManager::TfToCudaGpuId(tf_gpu_id);
+ cudaError_t error = cudaGetDeviceProperties(&properties, cuda_gpu_id.value());
if (error == cudaSuccess) {
device.set_vendor("NVidia");
device.set_model(properties.name);
// double data rate (DDR).
device.set_bandwidth(properties.memoryBusWidth / 8 *
properties.memoryClockRate * 2);
+ } else {
+ LOG(ERROR) << "Failed to get device properties, error code: " << error;
+ device.set_type("UNKNOWN");
+ return device;
}
(*device.mutable_environment())["architecture"] =
return GetLocalCPUInfo();
} else if (device.type == "GPU") {
if (device.has_id) {
- return GetLocalGPUInfo(device.id);
+ return GetLocalGPUInfo(TfGpuId(device.id));
} else {
- return GetLocalGPUInfo(0);
+ return GetLocalGPUInfo(TfGpuId(0));
}
}
DeviceProperties result;
#ifndef TENSORFLOW_GRAPPLER_CLUSTERS_UTILS_H_
#define TENSORFLOW_GRAPPLER_CLUSTERS_UTILS_H_
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/util/device_name_utils.h"
// Returns the DeviceProperties for the specified GPU attached to the server on
// which grappler is running.
-DeviceProperties GetLocalGPUInfo(int gpu_id);
+DeviceProperties GetLocalGPUInfo(TfGpuId tf_gpu_id);
// Returns the DeviceProperties of the specified device
DeviceProperties GetDeviceInfo(const DeviceNameUtils::ParsedName& device);
"//third_party/eigen3",
"//tensorflow/core:framework",
"//tensorflow/core:graph",
+ "//tensorflow/core:gpu_id",
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
#include "cuda/include/cudnn.h"
#endif
+#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/op.h"
DeviceNameUtils::ParsedName parsed;
if (DeviceNameUtils::ParseFullName(device_str, &parsed)) {
if (parsed.type == "GPU") {
- return GetLocalGPUInfo(parsed.id);
+ return GetLocalGPUInfo(TfGpuId(parsed.id));
} else if (parsed.type == "CPU") {
return GetLocalCPUInfo();
}