blob: d11b58813dbb8ae6d03f6ae49f81e3cbabb0af79 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
|
#include "tensorflow/stream_executor/temporary_device_memory.h"
#include "tensorflow/stream_executor/stream.h"
namespace perftools {
namespace gputools {
TemporaryDeviceMemoryBase::~TemporaryDeviceMemoryBase() {
parent_->temporary_memory_manager()->MarkFinalized(device_memory_,
allocation_generation_,
/*must_exist=*/false);
}
DeviceMemoryBase* TemporaryDeviceMemoryBase::mutable_device_memory() {
DCHECK(!IsFinalized())
<< "should not access device memory after finalization";
return &device_memory_;
}
const DeviceMemoryBase& TemporaryDeviceMemoryBase::device_memory() const {
DCHECK(!IsFinalized())
<< "should not access device memory after finalization";
return device_memory_;
}
void TemporaryDeviceMemoryBase::Finalize() {
DCHECK(!IsFinalized()) << "should not finalize more than once";
parent_->temporary_memory_manager()->MarkFinalized(device_memory_,
allocation_generation_,
/*must_exist=*/true);
}
bool TemporaryDeviceMemoryBase::IsFinalized() const {
return parent_->temporary_memory_manager()->IsFinalized(
device_memory_, allocation_generation_);
}
bool TemporaryDeviceMemoryBase::IsAllocated() const {
return parent_->temporary_memory_manager()->HasAllocated(
device_memory_, allocation_generation_);
}
TemporaryDeviceMemoryBase::TemporaryDeviceMemoryBase(
Stream* parent, DeviceMemoryBase device_memory,
uint64 allocation_generation)
: device_memory_(device_memory),
allocation_generation_(allocation_generation),
parent_(parent) {
DCHECK(IsAllocated());
}
} // namespace gputools
} // namespace perftools
|