Skip to content

Commit 0325a38

Browse files
annarevtensorflower-gardener
annarev
authored and
tensorflower-gardener
committed
Move gpu/gpu_id_utils.h to device/device_id_utils.h so that it can be reused for pluggable device.
PiperOrigin-RevId: 342988876 Change-Id: I9fabd153914e3251432cb07fa174f3d35abfac09
1 parent 570cb3d commit 0325a38

File tree

12 files changed

+191
-168
lines changed

12 files changed

+191
-168
lines changed

‎tensorflow/core/common_runtime/device/BUILD‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ package(
2424
"//tensorflow:internal",
2525
"//tensorflow_models:__subpackages__",
2626
],
27+
# features = ["-parse_headers"],
2728
licenses = ["notice"], # Apache 2.0
2829
)
2930

@@ -59,6 +60,7 @@ filegroup(
5960
srcs = [
6061
"device_id.h",
6162
"device_id_manager.h",
63+
"device_id_utils.h",
6264
],
6365
)
6466

‎tensorflow/core/common_runtime/device/device_id_utils.h‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
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_COMMON_RUNTIME_DEVICE_DEVICE_ID_UTILS_H_
17+
#define TENSORFLOW_CORE_COMMON_RUNTIME_DEVICE_DEVICE_ID_UTILS_H_
18+
19+
#include "tensorflow/core/common_runtime/device/device_id.h"
20+
#include "tensorflow/core/common_runtime/device/device_id_manager.h"
21+
#include "tensorflow/core/lib/gtl/int_type.h"
22+
#include "tensorflow/core/platform/types.h"
23+
#include "tensorflow/stream_executor/platform.h"
24+
#include "tensorflow/stream_executor/stream_executor.h"
25+
26+
namespace tensorflow {
27+
28+
// Utility methods for translation between TensorFlow device ids and platform
29+
// device ids.
30+
class DeviceIdUtil {
31+
public:
32+
// Convenient methods for getting the associated executor given a TfDeviceId
33+
// or PlatformDeviceId.
34+
static se::port::StatusOr<se::StreamExecutor*> ExecutorForPlatformDeviceId(
35+
se::Platform* device_manager, PlatformDeviceId platform_device_id) {
36+
return device_manager->ExecutorForDevice(platform_device_id.value());
37+
}
38+
static se::port::StatusOr<se::StreamExecutor*> ExecutorForTfDeviceId(
39+
const DeviceType& type, se::Platform* device_manager,
40+
TfDeviceId tf_device_id) {
41+
PlatformDeviceId platform_device_id;
42+
TF_RETURN_IF_ERROR(DeviceIdManager::TfToPlatformDeviceId(
43+
type, tf_device_id, &platform_device_id));
44+
return ExecutorForPlatformDeviceId(device_manager, platform_device_id);
45+
}
46+
47+
// Verify that the platform_device_id associated with a TfDeviceId is
48+
// legitimate.
49+
static void CheckValidTfDeviceId(const DeviceType& type,
50+
se::Platform* device_manager,
51+
TfDeviceId tf_device_id) {
52+
PlatformDeviceId platform_device_id;
53+
TF_CHECK_OK(DeviceIdManager::TfToPlatformDeviceId(type, tf_device_id,
54+
&platform_device_id));
55+
const int visible_device_count = device_manager->VisibleDeviceCount();
56+
CHECK_LT(platform_device_id.value(), visible_device_count)
57+
<< "platform_device_id is outside discovered device range."
58+
<< " TF " << type << " id: " << tf_device_id << ", platform " << type
59+
<< " id: " << platform_device_id
60+
<< ", visible device count: " << visible_device_count;
61+
}
62+
};
63+
64+
} // namespace tensorflow
65+
66+
#endif // TENSORFLOW_CORE_COMMON_RUNTIME_DEVICE_DEVICE_ID_UTILS_H_

‎tensorflow/core/common_runtime/gpu/BUILD‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,6 @@ filegroup(
116116
"gpu_host_allocator.h",
117117
"gpu_id.h",
118118
"gpu_id_manager.h",
119-
"gpu_id_utils.h",
120119
"gpu_init.h",
121120
"gpu_managed_allocator.h",
122121
"gpu_mem_allocator.h",
@@ -446,5 +445,6 @@ tf_cuda_cc_test(
446445
"//tensorflow/core/common_runtime:core_cpu_internal",
447446
"//tensorflow/core/common_runtime:direct_session_internal",
448447
"//tensorflow/core/kernels:ops_util",
448+
"//tensorflow/stream_executor:platform",
449449
],
450450
)

‎tensorflow/core/common_runtime/gpu/gpu_bfc_allocator_test.cc‎

Lines changed: 49 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@ limitations under the License.
2121
#include <algorithm>
2222
#include <vector>
2323

24+
#include "tensorflow/core/common_runtime/device/device_id_utils.h"
2425
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
25-
#include "tensorflow/core/common_runtime/gpu/gpu_id_utils.h"
2626
#include "tensorflow/core/common_runtime/gpu/gpu_init.h"
2727
#include "tensorflow/core/framework/typed_allocator.h"
2828
#include "tensorflow/core/lib/core/threadpool.h"
@@ -52,11 +52,18 @@ static void CheckStats(Allocator* a, int64 num_allocs, int64 bytes_in_use,
5252
EXPECT_EQ(stats->largest_alloc_size, largest_alloc_size);
5353
}
5454

55+
se::StreamExecutor* ExecutorForPlatformGpuId(
56+
PlatformDeviceId platform_device_id) {
57+
return DeviceIdUtil::ExecutorForPlatformDeviceId(GPUMachineManager(),
58+
platform_device_id)
59+
.ValueOrDie();
60+
}
61+
5562
TEST(GPUBFCAllocatorTest, NoDups) {
5663
PlatformGpuId platform_gpu_id(0);
5764
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
58-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
59-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
65+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
66+
false /*use_unified_memory*/, {}, {});
6067
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
6168
CheckStats(&a, 0, 0, 0, 0);
6269

@@ -88,8 +95,8 @@ TEST(GPUBFCAllocatorTest, NoDups) {
8895
TEST(GPUBFCAllocatorTest, AllocationsAndDeallocations) {
8996
PlatformGpuId platform_gpu_id(0);
9097
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
91-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
92-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
98+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
99+
false /*use_unified_memory*/, {}, {});
93100
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
94101
// Allocate 256 raw pointers of sizes between 100 bytes and about
95102
// a meg
@@ -150,8 +157,8 @@ TEST(GPUBFCAllocatorTest, AllocationsAndDeallocations) {
150157
TEST(GPUBFCAllocatorTest, ExerciseCoalescing) {
151158
PlatformGpuId platform_gpu_id(0);
152159
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
153-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
154-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
160+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
161+
false /*use_unified_memory*/, {}, {});
155162
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
156163
CheckStats(&a, 0, 0, 0, 0);
157164

@@ -189,8 +196,8 @@ TEST(GPUBFCAllocatorTest, ExerciseCoalescing) {
189196
TEST(GPUBFCAllocatorTest, AllocateZeroBufSize) {
190197
PlatformGpuId platform_gpu_id(0);
191198
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
192-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
193-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
199+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
200+
false /*use_unified_memory*/, {}, {});
194201
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
195202
float* ptr = TypedAllocator::Allocate<float>(&a, 0, {});
196203
EXPECT_EQ(nullptr, ptr);
@@ -199,17 +206,17 @@ TEST(GPUBFCAllocatorTest, AllocateZeroBufSize) {
199206
TEST(GPUBFCAllocatorTest, TracksSizes) {
200207
PlatformGpuId platform_gpu_id(0);
201208
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
202-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
203-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
209+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
210+
false /*use_unified_memory*/, {}, {});
204211
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
205212
EXPECT_EQ(true, a.TracksAllocationSizes());
206213
}
207214

208215
TEST(GPUBFCAllocatorTest, AllocatedVsRequested) {
209216
PlatformGpuId platform_gpu_id(0);
210217
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
211-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
212-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
218+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
219+
false /*use_unified_memory*/, {}, {});
213220
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
214221
float* t1 = TypedAllocator::Allocate<float>(&a, 1, {});
215222
EXPECT_EQ(4, a.RequestedSize(t1));
@@ -220,8 +227,8 @@ TEST(GPUBFCAllocatorTest, AllocatedVsRequested) {
220227
TEST(GPUBFCAllocatorTest, TestCustomMemoryLimit) {
221228
PlatformGpuId platform_gpu_id(0);
222229
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
223-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
224-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
230+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
231+
false /*use_unified_memory*/, {}, {});
225232
// Configure a 1MiB byte limit
226233
GPUBFCAllocator a(sub_allocator, 1 << 20, "GPU_0_bfc");
227234

@@ -240,8 +247,8 @@ TEST(GPUBFCAllocatorTest, AllocationsAndDeallocationsWithGrowth) {
240247
// Max of 2GiB, but starts out small.
241248
PlatformGpuId platform_gpu_id(0);
242249
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
243-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
244-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
250+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
251+
false /*use_unified_memory*/, {}, {});
245252
GPUBFCAllocator a(sub_allocator, 1LL << 31, "GPU_0_bfc");
246253

247254
// Allocate 10 raw pointers of sizes between 100 bytes and about
@@ -306,12 +313,12 @@ TEST(GPUBFCAllocatorTest, AllocationsAndDeallocationsWithGrowth) {
306313
TEST(GPUBFCAllocatorTest, DISABLED_AllocatorReceivesZeroMemory) {
307314
PlatformGpuId platform_gpu_id(0);
308315
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
309-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
310-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
316+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
317+
false /*use_unified_memory*/, {}, {});
311318
GPUBFCAllocator a(sub_allocator, 1UL << 60, "GPU_0_bfc");
312-
sub_allocator = new GPUMemAllocator(
313-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
314-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
319+
sub_allocator = new GPUMemAllocator(ExecutorForPlatformGpuId(platform_gpu_id),
320+
platform_gpu_id,
321+
false /*use_unified_memory*/, {}, {});
315322
GPUBFCAllocator b(sub_allocator, 1UL << 60, "GPU_0_bfc");
316323
void* amem = a.AllocateRaw(1, 1);
317324
void* bmem = b.AllocateRaw(1, 1 << 30);
@@ -322,8 +329,8 @@ TEST(GPUBFCAllocatorTest, DISABLED_AllocatorReceivesZeroMemory) {
322329
static void BM_Allocation(int iters) {
323330
PlatformGpuId platform_gpu_id(0);
324331
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
325-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
326-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
332+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
333+
false /*use_unified_memory*/, {}, {});
327334
GPUBFCAllocator a(sub_allocator, 1uLL << 33, "GPU_0_bfc");
328335
// Exercise a few different allocation sizes
329336
std::vector<size_t> sizes = {256, 4096, 16384, 524288,
@@ -342,8 +349,8 @@ BENCHMARK(BM_Allocation);
342349
static void BM_AllocationThreaded(int iters, int num_threads) {
343350
PlatformGpuId platform_gpu_id(0);
344351
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
345-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
346-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
352+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
353+
false /*use_unified_memory*/, {}, {});
347354
GPUBFCAllocator a(sub_allocator, 1uLL << 33, "GPU_0_bfc");
348355
thread::ThreadPool pool(Env::Default(), "test", num_threads);
349356
std::atomic_int_fast32_t count(iters);
@@ -382,8 +389,8 @@ BENCHMARK(BM_AllocationThreaded)->Arg(1)->Arg(4)->Arg(16);
382389
static void BM_AllocationDelayed(int iters, int delay) {
383390
PlatformGpuId platform_gpu_id(0);
384391
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
385-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
386-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
392+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
393+
false /*use_unified_memory*/, {}, {});
387394
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
388395
// Exercise a few different allocation sizes
389396
std::vector<int> sizes = {256, 4096, 16384, 4096, 512, 1024, 1024};
@@ -426,8 +433,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
426433
void TestBinDebugInfo() {
427434
PlatformGpuId platform_gpu_id(0);
428435
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
429-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
430-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
436+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
437+
false /*use_unified_memory*/, {}, {});
431438
GPUBFCAllocator a(sub_allocator, 1 << 30, "GPU_0_bfc");
432439

433440
std::vector<void*> initial_ptrs;
@@ -526,8 +533,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
526533
void TestLog2FloorNonZeroSlow() {
527534
PlatformGpuId platform_gpu_id(0);
528535
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
529-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
530-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
536+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
537+
false /*use_unified_memory*/, {}, {});
531538
GPUBFCAllocator a(sub_allocator, 1 /* total_memory */, "GPU_0_bfc");
532539
EXPECT_EQ(-1, a.Log2FloorNonZeroSlow(0));
533540
EXPECT_EQ(0, a.Log2FloorNonZeroSlow(1));
@@ -545,8 +552,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
545552
unsetenv("TF_FORCE_GPU_ALLOW_GROWTH");
546553
options.set_allow_growth(true);
547554
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
548-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
549-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
555+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
556+
false /*use_unified_memory*/, {}, {});
550557
GPUBFCAllocator unset_flag_allocator(sub_allocator, 1LL << 31, options,
551558
"GPU_0_bfc");
552559
EXPECT_EQ(GPUBFCAllocator::RoundedBytes(size_t{1048576}),
@@ -556,8 +563,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
556563
setenv("TF_FORCE_GPU_ALLOW_GROWTH", "unparseable", 1);
557564
options.set_allow_growth(true);
558565
sub_allocator = new GPUMemAllocator(
559-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
560-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
566+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
567+
false /*use_unified_memory*/, {}, {});
561568
GPUBFCAllocator unparsable_flag_allocator(sub_allocator, 1LL << 31, options,
562569
"GPU_1_bfc");
563570
EXPECT_EQ(GPUBFCAllocator::RoundedBytes(size_t{1048576}),
@@ -568,8 +575,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
568575
setenv("TF_FORCE_GPU_ALLOW_GROWTH", "true", 1);
569576
options.set_allow_growth(false);
570577
sub_allocator = new GPUMemAllocator(
571-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
572-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
578+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
579+
false /*use_unified_memory*/, {}, {});
573580
GPUBFCAllocator force_allow_growth_allocator(sub_allocator, 1LL << 31,
574581
options, "GPU_2_bfc");
575582
EXPECT_EQ(GPUBFCAllocator::RoundedBytes(size_t{1048576}),
@@ -580,8 +587,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
580587
setenv("TF_FORCE_GPU_ALLOW_GROWTH", "false", 1);
581588
options.set_allow_growth(true);
582589
sub_allocator = new GPUMemAllocator(
583-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
584-
platform_gpu_id, false /*use_unified_memory*/, {}, {});
590+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
591+
false /*use_unified_memory*/, {}, {});
585592
GPUBFCAllocator force_no_allow_growth_allocator(sub_allocator, 1LL << 31,
586593
options, "GPU_3_bfc");
587594
EXPECT_EQ(GPUBFCAllocator::RoundedBytes(1LL << 31),
@@ -595,8 +602,8 @@ class GPUBFCAllocatorPrivateMethodsTest : public ::testing::Test {
595602
// Max of 2GiB, but starts out small.
596603
PlatformGpuId platform_gpu_id(0);
597604
GPUMemAllocator* sub_allocator = new GPUMemAllocator(
598-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie(),
599-
platform_gpu_id, /*use_unified_memory=*/false, {}, {});
605+
ExecutorForPlatformGpuId(platform_gpu_id), platform_gpu_id,
606+
/*use_unified_memory=*/false, {}, {});
600607
GPUBFCAllocator a(sub_allocator, 1LL << 31, options, "GPU_0_bfc");
601608

602609
// Allocate 128 raw pointers of 4 megs.

‎tensorflow/core/common_runtime/gpu/gpu_cudamalloc_allocator.cc‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,9 @@ limitations under the License.
1818
#include "tensorflow/stream_executor/cuda/cuda_activation.h"
1919
#endif // GOOGLE_CUDA
2020

21+
#include "tensorflow/core/common_runtime/device/device_id_utils.h"
2122
#include "tensorflow/core/common_runtime/gpu/gpu_cudamalloc_allocator.h"
22-
2323
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
24-
#include "tensorflow/core/common_runtime/gpu/gpu_id_utils.h"
2524
#include "tensorflow/core/common_runtime/gpu/gpu_init.h"
2625
#include "tensorflow/core/platform/stream_executor.h"
2726

@@ -30,8 +29,9 @@ namespace tensorflow {
3029
GPUcudaMallocAllocator::GPUcudaMallocAllocator(Allocator* allocator,
3130
PlatformGpuId platform_gpu_id)
3231
: base_allocator_(allocator) {
33-
stream_exec_ =
34-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie();
32+
stream_exec_ = DeviceIdUtil::ExecutorForPlatformDeviceId(GPUMachineManager(),
33+
platform_gpu_id)
34+
.ValueOrDie();
3535
}
3636

3737
GPUcudaMallocAllocator::~GPUcudaMallocAllocator() { delete base_allocator_; }

‎tensorflow/core/common_runtime/gpu/gpu_debug_allocator.cc‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,8 @@ limitations under the License.
1818
#include <cstddef>
1919
#include <vector>
2020

21+
#include "tensorflow/core/common_runtime/device/device_id_utils.h"
2122
#include "tensorflow/core/common_runtime/gpu/gpu_id.h"
22-
#include "tensorflow/core/common_runtime/gpu/gpu_id_utils.h"
2323
#include "tensorflow/core/common_runtime/gpu/gpu_init.h"
2424
#include "tensorflow/core/platform/stream_executor.h"
2525

@@ -78,8 +78,9 @@ void InitMask(se::StreamExecutor* exec, void* ptr, int64* mask) {
7878
GPUDebugAllocator::GPUDebugAllocator(Allocator* allocator,
7979
PlatformGpuId platform_gpu_id)
8080
: base_allocator_(allocator) {
81-
stream_exec_ =
82-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie();
81+
stream_exec_ = DeviceIdUtil::ExecutorForPlatformDeviceId(GPUMachineManager(),
82+
platform_gpu_id)
83+
.ValueOrDie();
8384
}
8485

8586
GPUDebugAllocator::~GPUDebugAllocator() { delete base_allocator_; }
@@ -156,8 +157,9 @@ bool GPUDebugAllocator::CheckFooter(void* ptr) {
156157
GPUNanResetAllocator::GPUNanResetAllocator(Allocator* allocator,
157158
PlatformGpuId platform_gpu_id)
158159
: base_allocator_(allocator) {
159-
stream_exec_ =
160-
GpuIdUtil::ExecutorForPlatformGpuId(platform_gpu_id).ValueOrDie();
160+
stream_exec_ = DeviceIdUtil::ExecutorForPlatformDeviceId(GPUMachineManager(),
161+
platform_gpu_id)
162+
.ValueOrDie();
161163
}
162164

163165
GPUNanResetAllocator::~GPUNanResetAllocator() { delete base_allocator_; }

0 commit comments

Comments
 (0)