| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | maybe_owning_device_memory.h | 28 class MaybeOwningDeviceMemory {
30 MaybeOwningDeviceMemory() = default;
31 explicit MaybeOwningDeviceMemory(tensorflow::se::OwningDeviceMemory owned) in MaybeOwningDeviceMemory() function
33 explicit MaybeOwningDeviceMemory(tensorflow::se::DeviceMemoryBase unowned) in MaybeOwningDeviceMemory() function
35 MaybeOwningDeviceMemory(MaybeOwningDeviceMemory&&) = default;
36 ~MaybeOwningDeviceMemory() = default;
38 MaybeOwningDeviceMemory& operator=(tensorflow::se::DeviceMemoryBase unowned) {
43 MaybeOwningDeviceMemory& operator=(tensorflow::se::OwningDeviceMemory owned) {
48 MaybeOwningDeviceMemory& operator=(MaybeOwningDeviceMemory&&) = default;
|
| D | maybe_owning_device_memory.cc | 22 tensorflow::se::DeviceMemoryBase MaybeOwningDeviceMemory::AsDeviceMemoryBase() in AsDeviceMemoryBase()
31 bool MaybeOwningDeviceMemory::HasOwnership() const { in HasOwnership()
36 MaybeOwningDeviceMemory::Release() { in Release()
44 MaybeOwningDeviceMemory::AsOwningDeviceMemory() const { in AsOwningDeviceMemory()
|
| D | executable.h | 72 explicit ExecutionInput(ShapeTree<MaybeOwningDeviceMemory> buffers) in ExecutionInput()
77 ExecutionInput(ShapeTree<MaybeOwningDeviceMemory> buffers, in ExecutionInput()
102 void SetBuffer(const ShapeIndex& index, MaybeOwningDeviceMemory buffer) { in SetBuffer()
107 MaybeOwningDeviceMemory buffer);
119 const ShapeTree<MaybeOwningDeviceMemory>& Buffers() const { return buffers_; } in Buffers()
121 ShapeTree<MaybeOwningDeviceMemory>* MutableBuffers() { return &buffers_; } in MutableBuffers()
123 MaybeOwningDeviceMemory* MutableBuffer(const ShapeIndex& index) { in MutableBuffer()
127 const MaybeOwningDeviceMemory& Buffer(const ShapeIndex& index) const { in Buffer()
138 ShapeTree<MaybeOwningDeviceMemory> buffers_;
|
| D | executable.cc | 57 MaybeOwningDeviceMemory buffer) { in SetUnownedBuffer()
98 result.SetBuffer(index, MaybeOwningDeviceMemory(mem)); in MakeMaybeOwningDeviceMemoryTree()
|
| D | transfer_manager.h | 230 const ShapeTree<MaybeOwningDeviceMemory>& buffer_tree);
|
| D | transfer_manager.cc | 342 se::Stream* stream, const ShapeTree<MaybeOwningDeviceMemory>& buffer_tree) { in WriteRootTupleIndexTable()
|
| D | hlo_runner.cc | 130 ShapeTree<MaybeOwningDeviceMemory> buffer_tree( in ExecutionInputsFromScopedShapedBuffers()
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | cpu_executable.cc | 91 static StatusOr<MaybeOwningDeviceMemory> MemoryForAllocation( in MemoryForAllocation()
104 return MaybeOwningDeviceMemory{out}; in MemoryForAllocation()
107 return MaybeOwningDeviceMemory{se::DeviceMemoryBase{}}; in MemoryForAllocation()
110 return MaybeOwningDeviceMemory{se::DeviceMemoryBase{}}; in MemoryForAllocation()
124 return MaybeOwningDeviceMemory{std::move(out)}; in MemoryForAllocation()
127 StatusOr<std::vector<MaybeOwningDeviceMemory>> CpuExecutable::CreateBufferTable( in CreateBufferTable()
130 std::vector<MaybeOwningDeviceMemory> buffers( in CreateBufferTable()
150 absl::Span<MaybeOwningDeviceMemory const> buffers, in ExecuteComputeFunction()
227 absl::Span<MaybeOwningDeviceMemory> buffers, in CreateResultShapedBuffer()
264 MaybeOwningDeviceMemory* maybe_owning_memory = in CreateResultShapedBuffer()
[all …]
|
| D | cpu_executable.h | 104 StatusOr<std::vector<MaybeOwningDeviceMemory>> CreateBufferTable(
112 absl::Span<MaybeOwningDeviceMemory const> buffers,
121 absl::Span<MaybeOwningDeviceMemory> buffers,
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tracked_device_buffer.cc | 138 ShapeTree<MaybeOwningDeviceMemory>::iterator* iterator, in AddToInputAsImmutable()
139 const ShapeTree<MaybeOwningDeviceMemory>::iterator& end) const { in AddToInputAsImmutable()
143 (*iterator)->second = MaybeOwningDeviceMemory(buf); in AddToInputAsImmutable()
149 ShapeTree<MaybeOwningDeviceMemory>::iterator* iterator, in AddToInputAsDonated()
150 const ShapeTree<MaybeOwningDeviceMemory>::iterator& end, in AddToInputAsDonated()
156 (*iterator)->second = MaybeOwningDeviceMemory( in AddToInputAsDonated()
|
| D | tracked_device_buffer.h | 149 ShapeTree<MaybeOwningDeviceMemory>::iterator* iterator,
150 const ShapeTree<MaybeOwningDeviceMemory>::iterator& end) const;
161 ShapeTree<MaybeOwningDeviceMemory>::iterator* iterator,
162 const ShapeTree<MaybeOwningDeviceMemory>::iterator& end,
|
| D | pjrt_stream_executor_client.cc | 564 ShapeTree<MaybeOwningDeviceMemory>::iterator* iterator, in AddToInput()
565 const ShapeTree<MaybeOwningDeviceMemory>::iterator& end, in AddToInput()
1517 ShapeTree<MaybeOwningDeviceMemory>::iterator input_iterator = in MakeTupleHelper()
1519 ShapeTree<MaybeOwningDeviceMemory>::iterator iterator_end = in MakeTupleHelper()
1524 MaybeOwningDeviceMemory(std::move(root_table_memory))); in MakeTupleHelper()
1659 ShapeTree<MaybeOwningDeviceMemory>::iterator input_iterator = in MakeExecutionInputsAndWaitForEvents()
1661 ShapeTree<MaybeOwningDeviceMemory>::iterator iterator_end = in MakeExecutionInputsAndWaitForEvents()
|
| D | pjrt_stream_executor_client.h | 412 void AddToInput(ShapeTree<MaybeOwningDeviceMemory>::iterator* iterator,
413 const ShapeTree<MaybeOwningDeviceMemory>::iterator& end,
|
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | c_api_conversions.h | 90 xla::MaybeOwningDeviceMemory FromC(
102 SE_MaybeOwningDeviceMemory ToC(xla::MaybeOwningDeviceMemory& mem, bool aliased);
|
| D | tpu_executable_interface.cc | 70 const MaybeOwningDeviceMemory& original_input = in AllocateOutputMemoryWithInputReuse()
117 MaybeOwningDeviceMemory* device_memory = in AllocateOutputMemoryWithInputReuse()
|
| D | c_api_conversions.cc | 43 SE_MaybeOwningDeviceMemory ToC(xla::MaybeOwningDeviceMemory& mem, in ToC()
65 xla::MaybeOwningDeviceMemory FromC( in FromC()
69 return xla::MaybeOwningDeviceMemory( in FromC()
73 return xla::MaybeOwningDeviceMemory(ApiConverter::FromC(se_mem->memory)); in FromC()
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | buffer_donation_test.cc | 109 ShapeTree<MaybeOwningDeviceMemory> owned_buffers( in RunAndCheck()
112 [&](const ShapeIndex& index, MaybeOwningDeviceMemory* device_memory) { in RunAndCheck()
|
| /external/tensorflow/tensorflow/compiler/xla/service/interpreter/ |
| D | executable_base.cc | 58 const ShapeTree<MaybeOwningDeviceMemory>& buffers = argument.Buffers(); in ExecuteAsyncOnStream()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 225 xla::MaybeOwningDeviceMemory* mutable_input_mem = in UpdateDynamicInputs()
278 xla::MaybeOwningDeviceMemory(std::move(new_input)); in UpdateDynamicInputs()
|
| /external/tensorflow/tensorflow/compiler/xla/g3doc/ |
| D | aliasing.md | 56 which in turn contain a tree of `MaybeOwningDeviceMemory`. If memory is
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/tests/ |
| D | mlir_gpu_test_base.cc | 96 execution_inputs.back().SetBuffer({}, MaybeOwningDeviceMemory(arg)); in RunMlirModule()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gpu_executable.cc | 471 MaybeOwningDeviceMemory* maybe_owning_memory = in ExecuteAsyncOnStreamImpl()
472 [&]() -> xla::MaybeOwningDeviceMemory* { in ExecuteAsyncOnStreamImpl()
|
| /external/tensorflow/tensorflow/compiler/xrt/ |
| D | xrt_state.cc | 666 xla::MaybeOwningDeviceMemory(index_buffer.second->allocation())); in ToExecutionInput()
671 xla::MaybeOwningDeviceMemory(se::OwningDeviceMemory( in ToExecutionInput()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 177 [](xla::MaybeOwningDeviceMemory* buffer) { in ToShapedBuffer()
185 xla::ShapeTree<xla::MaybeOwningDeviceMemory> buffers;
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | local_client.cc | 270 const ShapeTree<MaybeOwningDeviceMemory>& tree, int device_ordinal) { in MaybeOwningShapeTreeToShapedBuffer()
|