| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | shaped_buffer.cc | 116 ScopedShapedBuffer::ScopedShapedBuffer(Shape on_device_shape, in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
122 ScopedShapedBuffer::ScopedShapedBuffer(Shape on_host_shape, in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
126 : ScopedShapedBuffer(std::move(on_device_shape), allocator, in ScopedShapedBuffer()
129 ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
133 ScopedShapedBuffer::ScopedShapedBuffer(ScopedShapedBuffer&& s) in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
139 ScopedShapedBuffer& ScopedShapedBuffer::operator=(ScopedShapedBuffer&& s) { in operator =()
149 ScopedShapedBuffer::~ScopedShapedBuffer() { Deallocate(); } in ~ScopedShapedBuffer()
151 ShapedBuffer ScopedShapedBuffer::release() { in release()
157 void ScopedShapedBuffer::Deallocate() { in Deallocate()
175 ScopedShapedBuffer ScopedShapedBuffer::TakeSubTree(ShapeIndexView index) { in TakeSubTree()
[all …]
|
| D | shaped_buffer.h | 34 class ScopedShapedBuffer; variable
58 ShapedBuffer(const ScopedShapedBuffer&) = delete;
59 ShapedBuffer& operator=(const ScopedShapedBuffer&) = delete;
149 class ScopedShapedBuffer : public ShapedBuffer {
152 explicit ScopedShapedBuffer(Shape on_device_shape,
156 explicit ScopedShapedBuffer(Shape on_host_shape, Shape on_device_shape,
162 explicit ScopedShapedBuffer(ShapedBuffer shaped_buffer,
166 ScopedShapedBuffer(ScopedShapedBuffer&& s);
167 ScopedShapedBuffer& operator=(ScopedShapedBuffer&&);
168 ScopedShapedBuffer(const ScopedShapedBuffer&) = delete;
[all …]
|
| D | hlo_runner.cc | 48 StatusOr<ScopedShapedBuffer> HloRunner::TransferLiteralToDevice( in TransferLiteralToDevice()
50 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer buffer, in TransferLiteralToDevice()
61 StatusOr<std::vector<ScopedShapedBuffer>> HloRunner::TransferLiteralsToDevice( in TransferLiteralsToDevice()
63 std::vector<ScopedShapedBuffer> buffers; in TransferLiteralsToDevice()
66 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer buffer, in TransferLiteralsToDevice()
73 StatusOr<std::vector<ScopedShapedBuffer>> HloRunner::TransferLiteralsToDevice( in TransferLiteralsToDevice()
95 TF_ASSIGN_OR_RETURN(std::vector<ScopedShapedBuffer> argument_buffers, in Execute()
109 TF_ASSIGN_OR_RETURN(std::vector<ScopedShapedBuffer> argument_buffers, in ExecuteWithExecutable()
122 absl::Span<ScopedShapedBuffer const> inputs, in ExecutionInputsFromScopedShapedBuffers()
129 const ScopedShapedBuffer& input_buffer = inputs[param_num]; in ExecutionInputsFromScopedShapedBuffers()
[all …]
|
| D | executable.h | 150 explicit ExecutionOutput(ScopedShapedBuffer result) in ExecutionOutput()
152 ExecutionOutput(ScopedShapedBuffer result, in ExecutionOutput()
190 const ScopedShapedBuffer& Result() const { return result_; } in Result()
192 ScopedShapedBuffer* MutableResult() { return &result_; } in MutableResult()
194 ScopedShapedBuffer ConsumeResult() { in ConsumeResult()
214 ScopedShapedBuffer result_;
258 StatusOr<ScopedShapedBuffer> ExecuteOnStream(
281 virtual StatusOr<ScopedShapedBuffer> ExecuteAsyncOnStream(
302 virtual StatusOr<std::vector<ScopedShapedBuffer>> ExecuteOnStreams(
319 StatusOr<ScopedShapedBuffer> ExecuteOnStreamWrapper(
[all …]
|
| D | hlo_runner.h | 59 StatusOr<ScopedShapedBuffer> TransferLiteralToDevice(const Literal& literal);
60 StatusOr<std::vector<ScopedShapedBuffer>> TransferLiteralsToDevice(
62 StatusOr<std::vector<ScopedShapedBuffer>> TransferLiteralsToDevice(
90 absl::Span<ScopedShapedBuffer const> arguments,
94 Executable* executable, absl::Span<ScopedShapedBuffer const> arguments,
154 std::function<StatusOr<std::vector<ScopedShapedBuffer>>(
|
| D | shaped_buffer_test.cc | 38 auto scoped_buffer = absl::make_unique<xla::ScopedShapedBuffer>( in TEST()
100 ScopedShapedBuffer sb1(s, &allocator, /*device_ordinal=*/0); in TEST()
105 ScopedShapedBuffer sb2(s, &allocator, /*device_ordinal=*/1); in TEST()
122 ScopedShapedBuffer sb(s, &allocator, /*device_ordinal=*/0); in TEST()
134 ScopedShapedBuffer output = sb.TakeSubTree(subtree_index); in TEST()
159 ScopedShapedBuffer sb(tuple_shape, &allocator, /*device_ordinal=*/0); in TEST()
186 xla::ScopedShapedBuffer shaped_buffer(shape, /*allocator=*/&allocator, in BM_TakeSubTree()
|
| D | executable.cc | 81 StatusOr<ScopedShapedBuffer> Executable::ExecuteOnStream( in ExecuteOnStream()
85 StatusOr<ScopedShapedBuffer> result = in ExecuteOnStream()
103 StatusOr<ScopedShapedBuffer> Executable::ExecuteAsyncOnStream( in ExecuteAsyncOnStream()
130 StatusOr<std::vector<ScopedShapedBuffer>> Executable::ExecuteOnStreams( in ExecuteOnStreams()
135 std::vector<ScopedShapedBuffer> return_values; in ExecuteOnStreams()
162 StatusOr<ScopedShapedBuffer> Executable::ExecuteOnStreamWrapper( in ExecuteOnStreamWrapper()
165 StatusOr<ScopedShapedBuffer> result = in ExecuteOnStreamWrapper()
279 StatusOr<ScopedShapedBuffer> Executable::ExecuteAsyncOnStreamWrapper( in ExecuteAsyncOnStreamWrapper()
283 StatusOr<ScopedShapedBuffer> return_value = in ExecuteAsyncOnStreamWrapper()
|
| D | allocation_tracker.cc | 35 ScopedShapedBuffer shaped_buffer, const string& tag) { in Register()
38 std::vector<ScopedShapedBuffer> replicated_buffers; in Register()
44 std::vector<ScopedShapedBuffer> replicated_buffers, const string& tag) { in RegisterReplicatedBuffers()
54 static ShapedBuffer ReleaseIfScopedShapedBuffer(ScopedShapedBuffer b) { in ReleaseIfScopedShapedBuffer()
62 std::is_same<ShapedBufferTy, ScopedShapedBuffer>::value, in RegisterInternal()
|
| D | allocation_tracker.h | 49 StatusOr<GlobalDataHandle> Register(ScopedShapedBuffer shaped_buffer,
55 std::vector<ScopedShapedBuffer> replicated_buffers, const string& tag);
|
| D | local_service.h | 70 std::vector<ScopedShapedBuffer> replicated_buffers, const string& tag);
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | local_client_test_base.h | 90 ScopedShapedBuffer LiteralToShapedBuffer(const Literal& literal);
98 StatusOr<ScopedShapedBuffer> ExecuteLocally(
101 StatusOr<ScopedShapedBuffer> ExecuteLocally(
107 ScopedShapedBuffer ExecuteLocallyOrDie(
110 ScopedShapedBuffer ExecuteLocallyOrDie(
|
| D | local_client_execute_test.cc | 60 ScopedShapedBuffer result = in XLA_TEST_F()
73 ScopedShapedBuffer result = in XLA_TEST_F()
86 ScopedShapedBuffer result = in XLA_TEST_F()
100 ScopedShapedBuffer result = in XLA_TEST_F()
115 ScopedShapedBuffer result = ExecuteLocallyOrDie( in XLA_TEST_F()
143 ScopedShapedBuffer result_colmaj = in XLA_TEST_F()
150 ScopedShapedBuffer result_param_swap = in XLA_TEST_F()
170 ScopedShapedBuffer result_colmaj = ExecuteLocallyOrDie( in XLA_TEST_F()
183 ScopedShapedBuffer result_rowmaj = ExecuteLocallyOrDie( in XLA_TEST_F()
208 ScopedShapedBuffer result = in XLA_TEST_F()
[all …]
|
| D | local_client_test_base.cc | 138 ScopedShapedBuffer LocalClientTestBase::LiteralToShapedBuffer( in LiteralToShapedBuffer()
163 ScopedShapedBuffer LocalClientTestBase::ExecuteLocallyOrDie( in ExecuteLocallyOrDie()
171 ScopedShapedBuffer LocalClientTestBase::ExecuteLocallyOrDie( in ExecuteLocallyOrDie()
180 StatusOr<ScopedShapedBuffer> LocalClientTestBase::ExecuteLocally( in ExecuteLocally()
187 StatusOr<ScopedShapedBuffer> LocalClientTestBase::ExecuteLocally( in ExecuteLocally()
|
| D | multiple_devices_on_host_test.cc | 39 std::vector<std::pair<int, StatusOr<ScopedShapedBuffer>>>* results) { in CompileAndExecute()
48 StatusOr<ScopedShapedBuffer> result = in CompileAndExecute()
74 std::vector<std::pair<int, StatusOr<ScopedShapedBuffer>>> results; in TestWithDeviceCount()
|
| D | local_client_allocation_test.cc | 56 absl::optional<ScopedShapedBuffer> result = in XLA_TEST_F()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | local_client.h | 50 StatusOr<ScopedShapedBuffer> Run(
61 StatusOr<ScopedShapedBuffer> RunAsync(
152 StatusOr<ScopedShapedBuffer> LiteralToShapedBuffer(
|
| D | local_client.cc | 172 StatusOr<ScopedShapedBuffer> LocalExecutable::Run( in Run()
180 return AsyncCallAndBlockHostUntilDone<xla::ScopedShapedBuffer>( in Run()
240 StatusOr<ScopedShapedBuffer> LocalExecutable::RunAsync( in RunAsync()
257 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer outputs, in RunAsync()
399 StatusOr<ScopedShapedBuffer> LocalClient::LiteralToShapedBuffer( in LiteralToShapedBuffer()
452 TF_ASSIGN_OR_RETURN(::xla::ScopedShapedBuffer shaped_buffer, in TransferToLocalServer()
459 std::vector<::xla::ScopedShapedBuffer> replicated_buffer; in TransferToLocalServer()
|
| /external/tensorflow/tensorflow/compiler/jit/ |
| D | xla_tensor.h | 70 void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) { in set_shaped_buffer()
99 absl::optional<xla::ScopedShapedBuffer> shaped_buffer_;
|
| D | xla_launch_util.cc | 44 using xla::ScopedShapedBuffer;
331 xla::ScopedShapedBuffer* output, int output_num, in PopulateXlaTensor()
423 ScopedShapedBuffer output, int missing_ctx_input_prefix, in PopulateOutputs()
449 output = ScopedShapedBuffer(std::move(buffer), output.memory_allocator()); in PopulateOutputs()
|
| D | xla_tensor.cc | 50 xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape, in AllocateShapedBuffer()
|
| D | xla_launch_util.h | 176 xla::ScopedShapedBuffer output, int missing_ctx_input_prefix,
|
| /external/tensorflow/tensorflow/compiler/xrt/ |
| D | xrt_state.cc | 74 const xla::Shape& shape, std::unique_ptr<xla::ScopedShapedBuffer>* buffer) { in AllocateScopedShapedBuffer()
95 *buffer = absl::make_unique<xla::ScopedShapedBuffer>( in AllocateScopedShapedBuffer()
174 std::unique_ptr<xla::ScopedShapedBuffer> scoped_buffer; in CreateAndTransfer()
201 std::unique_ptr<xla::ScopedShapedBuffer> scoped_buffer; in CreateUninitialized()
306 std::unique_ptr<xla::ScopedShapedBuffer> scoped_buffer; in SwapIn()
480 auto new_tuple_buffers = absl::make_unique<xla::ScopedShapedBuffer>( in MakeTuple()
|
| /external/tensorflow/tensorflow/compiler/xla/tools/ |
| D | replay_computation.cc | 230 std::vector<ScopedShapedBuffer> scoped_shaped_buffer_arguments; in ReplayComputation()
256 ScopedShapedBuffer data, in ReplayComputation()
309 absl::optional<ScopedShapedBuffer> final_result; in ReplayComputation()
331 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, in ReplayComputation()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 323 TF_ASSIGN_OR_RETURN(xla::ScopedShapedBuffer buffers, in BuildComputationInputs()
362 OutputBuffers(xla::ScopedShapedBuffer b, se::DeviceMemoryAllocator* allocator) in OutputBuffers()
395 OpKernelContext* context, xla::ScopedShapedBuffer scoped_buffers, in AllocateOutputTensors()
467 xla::ScopedShapedBuffer shaped_buffer(device_shape, allocator, in AllocateOutputTensors()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tracked_device_buffer.h | 136 ScopedShapedBuffer* shaped_buffer,
|