experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add DCHECK to ScopedDeviceMemory constructor.

Browse Source 
This lets some error cases fail faster.

PiperOrigin-RevId: 333314395
Change-Id: I3f6b91b39b67d7afdfd37945092c553b04fca9d3
This commit is contained in:
Skye Wanderman-Milne 2020-09-23 10:02:27 -07:00 committed by TensorFlower Gardener
parent c27f8ed531
commit 29266d4c3d
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tensorflow/stream_executor/device_memory_allocator.h
View File

@ -59,7 +59,9 @@ class ScopedDeviceMemory {
// out of scope. // out of scope.
ScopedDeviceMemory(DeviceMemoryBase mem, int device_ordinal, ScopedDeviceMemory(DeviceMemoryBase mem, int device_ordinal,
DeviceMemoryAllocator *allocator) DeviceMemoryAllocator *allocator)
: wrapped_(mem), device_ordinal_(device_ordinal), allocator_(allocator) {} : wrapped_(mem), device_ordinal_(device_ordinal), allocator_(allocator) {
DCHECK_GE(device_ordinal_, 0);
}
// A helper constructor to generate a scoped device memory given an already // A helper constructor to generate a scoped device memory given an already
// allocated memory and a stream executor. // allocated memory and a stream executor.
@ -79,8 +81,9 @@ class ScopedDeviceMemory {
// //
// Postcondition: other == nullptr. // Postcondition: other == nullptr.
ScopedDeviceMemory(ScopedDeviceMemory &&other) ScopedDeviceMemory(ScopedDeviceMemory &&other)
: ScopedDeviceMemory(other.Release(), other.device_ordinal_, : wrapped_(other.Release()),
other.allocator_) {} device_ordinal_(other.device_ordinal_),
allocator_(other.allocator_) {}
// Releases the memory that was provided in the constructor, through the // Releases the memory that was provided in the constructor, through the
// "parent" StreamExecutor. // "parent" StreamExecutor.