Home
last modified time | relevance | path

Searched refs:ShapedBuffer (Results 1 – 25 of 49) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dshaped_buffer.cc34 ShapedBuffer::ShapedBuffer(Shape on_device_shape, int device_ordinal) in ShapedBuffer() function in xla::ShapedBuffer
41 ShapedBuffer::ShapedBuffer(Shape on_host_shape, Shape on_device_shape, in ShapedBuffer() function in xla::ShapedBuffer
43 : ShapedBuffer(on_device_shape, device_ordinal) {} in ShapedBuffer()
45 ShapedBuffer::ShapedBuffer(ShapedBuffer&& s) in ShapedBuffer() function in xla::ShapedBuffer
56 ShapedBuffer& ShapedBuffer::operator=(ShapedBuffer&& s) { in operator =()
68 ShapedBuffer::~ShapedBuffer() {} in ~ShapedBuffer()
70 StatusOr<ShapedBuffer> ShapedBuffer::SubShapedBuffer( in SubShapedBuffer()
74 ShapedBuffer sub_shaped_buffer(*device_sub_shape, device_ordinal_); in SubShapedBuffer()
81 void ShapedBuffer::clear() { in clear()
88 string ShapedBuffer::ToString() const { in ToString()
[all …]
Dshaped_buffer.h38 class ShapedBuffer {
46 ShapedBuffer(Shape on_device_shape, int device_ordinal);
49 ShapedBuffer(Shape on_host_shape, Shape on_device_shape, int device_ordinal);
52 ShapedBuffer(ShapedBuffer&& s);
53 ShapedBuffer& operator=(ShapedBuffer&&);
54 ShapedBuffer(const ShapedBuffer&) = delete;
55 ShapedBuffer& operator=(const ShapedBuffer&) = delete;
58 ShapedBuffer(const ScopedShapedBuffer&) = delete;
59 ShapedBuffer& operator=(const ScopedShapedBuffer&) = delete;
61 virtual ~ShapedBuffer();
[all …]
Dtransfer_manager.h78 se::Stream* stream, const ShapedBuffer& device_buffer,
81 se::Stream* stream, const ShapedBuffer& device_buffer) { in TransferLiteralFromDevice()
85 se::Stream* stream, const ShapedBuffer& device_buffer,
89 const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
107 se::Stream* stream, const ShapedBuffer& device_buffer,
111 const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
131 const ShapedBuffer& device_buffer,
135 const ShapedBuffer& device_buffer) { in TransferLiteralToDevice()
157 const ShapedBuffer& device_buffer,
161 const ShapedBuffer& device_buffer) { in TransferLiteralToDeviceAsync()
[all …]
Dallocation_tracker.cc53 static ShapedBuffer ReleaseIfScopedShapedBuffer(ShapedBuffer b) { return b; } in ReleaseIfScopedShapedBuffer()
54 static ShapedBuffer ReleaseIfScopedShapedBuffer(ScopedShapedBuffer b) { in ReleaseIfScopedShapedBuffer()
61 static_assert(std::is_same<ShapedBufferTy, ShapedBuffer>::value || in RegisterInternal()
86 absl::make_unique<ShapedBuffer>( in RegisterInternal()
100 TF_ASSIGN_OR_RETURN(std::vector<const ShapedBuffer*> replicated_buffers, in Unregister()
132 TF_ASSIGN_OR_RETURN(std::vector<const ShapedBuffer*> replicated_buffers, in DeconstructTuple()
136 const ShapedBuffer* shaped_buffer = replicated_buffers[0]; in DeconstructTuple()
150 auto element_buffer = ShapedBuffer( in DeconstructTuple()
155 std::vector<ShapedBuffer> replicated_buffers; in DeconstructTuple()
166 StatusOr<std::vector<const ShapedBuffer*>> AllocationTracker::Resolve( in Resolve()
[all …]
Dallocation_tracker.h67 StatusOr<std::vector<const ShapedBuffer*>> Resolve(
73 StatusOr<const ShapedBuffer*> ResolveForReplica(const GlobalDataHandle& data,
89 StatusOr<std::vector<const ShapedBuffer*>> ResolveInternal(
148 absl::flat_hash_map<int64, std::vector<std::unique_ptr<ShapedBuffer>>>
Dexecutable.cc62 StatusOr<ShapedBuffer> ExecutionInput::ToShapedBuffer( in ToShapedBuffer()
65 ShapedBuffer shaped_buffer(input_shape, device_ordinal); in ToShapedBuffer()
83 absl::Span<const ShapedBuffer* const> arguments, in ExecuteOnStream()
94 const ShapedBuffer& shaped_buffer) { in MakeMaybeOwningDeviceMemoryTree()
105 absl::Span<const ShapedBuffer* const> arguments, in ExecuteAsyncOnStream()
109 for (const ShapedBuffer* arg : arguments) { in ExecuteAsyncOnStream()
132 absl::Span<const absl::Span<const ShapedBuffer* const>> arguments) { in ExecuteOnStreams()
164 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteOnStreamWrapper()
281 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteAsyncOnStreamWrapper()
Dservice.h199 absl::Span<const ShapedBuffer* const> arguments,
209 StatusOr<std::vector<std::vector<const ShapedBuffer*>>> GetArguments(
225 StatusOr<std::vector<std::vector<const ShapedBuffer*>>>
255 absl::Span<const std::vector<const ShapedBuffer*>> arguments,
264 absl::Span<const std::vector<std::vector<const ShapedBuffer*>>> arguments,
Dtransfer_manager.cc51 se::Stream* stream, const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
78 se::Stream* stream, const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
100 const ShapedBuffer& device_buffer, in TransferLiteralToDevice()
172 ShapedBuffer shaped_buffer(on_device_shape, in TransferArrayToDeviceAsync()
196 ShapedBuffer shaped_buffer(shape, stream->parent()->device_ordinal()); in TransferArrayFromDevice()
203 ShapedBuffer* device_buffer, in ReadDynamicShapes()
288 se::Stream* stream, const ShapedBuffer& device_buffer) { in WriteTupleIndexTables()
294 se::Stream* stream, const ShapedBuffer& device_buffer) { in WriteTupleIndexTablesAsync()
323 se::Stream* stream, const ShapedBuffer& device_buffer) { in WriteRootTupleIndexTable()
Dexecutable.h99 xla::StatusOr<xla::ShapedBuffer> ToShapedBuffer(
260 absl::Span<const ShapedBuffer* const> arguments,
283 absl::Span<const ShapedBuffer* const> arguments,
304 absl::Span<const absl::Span<const ShapedBuffer* const>> arguments);
321 absl::Span<const ShapedBuffer* const> arguments);
329 absl::Span<const ShapedBuffer* const> arguments);
Dservice.cc74 Status RecordArguments(const absl::Span<const ShapedBuffer* const> arguments, in RecordArguments()
78 for (const ShapedBuffer* argument : arguments) { in RecordArguments()
88 Status RecordResult(const ShapedBuffer& result, se::Stream* stream, in RecordResult()
231 StatusOr<std::vector<std::vector<const ShapedBuffer*>>>
236 std::vector<std::vector<const ShapedBuffer*>> replicated_arguments; in ResolveAndValidateArguments()
248 const ShapedBuffer* shaped_buffer = replicated_buffers[replica]; in ResolveAndValidateArguments()
274 absl::Span<const ShapedBuffer* const> arguments, in CreateModuleConfig()
346 absl::Span<const std::vector<std::vector<const ShapedBuffer*>>> arguments, in ExecuteParallelAndRegisterResult()
471 absl::Span<const std::vector<const ShapedBuffer*>> arguments, in ExecuteAndRegisterResult()
513 std::vector<absl::Span<const ShapedBuffer* const>> replicated_arguments; in ExecuteAndRegisterResult()
[all …]
Dgeneric_transfer_manager.h44 se::Stream* stream, const ShapedBuffer& device_buffer,
50 const ShapedBuffer& device_buffer,
Dhlo_runner.cc84 const ShapedBuffer& buffer) { in TransferLiteralFromDevice()
195 const std::vector<absl::Span<const ShapedBuffer* const>>&)> in ExecuteReplicatedImpl()
219 std::vector<const ShapedBuffer*> argument_buffer_ptrs(total_argument_count + in ExecuteReplicatedImpl()
221 std::vector<absl::Span<const ShapedBuffer* const>> argument_buffer_slices; in ExecuteReplicatedImpl()
324 const std::vector<absl::Span<const ShapedBuffer* const>>& in ExecuteReplicated()
378 const std::vector<absl::Span<const ShapedBuffer* const>>& in ExecuteReplicated()
Dhlo_runner.h64 StatusOr<Literal> TransferLiteralFromDevice(const ShapedBuffer& buffer);
156 const std::vector<absl::Span<const ShapedBuffer* const>>&)>
Dgeneric_transfer_manager.cc62 se::Stream* stream, const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
94 const ShapedBuffer& device_buffer, in TransferLiteralToDeviceAsync()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dlocal_client_test_base.h94 Literal ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer);
100 absl::Span<const ShapedBuffer* const> arguments);
103 absl::Span<const ShapedBuffer* const> arguments,
109 absl::Span<const ShapedBuffer* const> arguments);
112 absl::Span<const ShapedBuffer* const> arguments,
Dlocal_client_test_base.cc146 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral()
165 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteLocallyOrDie()
173 absl::Span<const ShapedBuffer* const> arguments, in ExecuteLocallyOrDie()
182 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteLocally()
189 absl::Span<const ShapedBuffer* const> arguments, in ExecuteLocally()
/external/tensorflow/tensorflow/compiler/xla/client/
Dlocal_client.cc173 const absl::Span<const ShapedBuffer* const> arguments, in Run()
177 for (const ShapedBuffer* const arg : arguments) { in Run()
201 const absl::Span<const ShapedBuffer* const> arguments, se::Stream* stream) { in DumpArguments()
205 for (const ShapedBuffer* arg : arguments) { in DumpArguments()
222 const ShapedBuffer& outputs, in DumpOutputsAndSaveSnapshot()
241 const absl::Span<const ShapedBuffer* const> arguments, in RunAsync()
245 for (const ShapedBuffer* const arg : arguments) { in RunAsync()
269 static ShapedBuffer MaybeOwningShapeTreeToShapedBuffer( in MaybeOwningShapeTreeToShapedBuffer()
271 ShapedBuffer result(tree.shape(), device_ordinal); in MaybeOwningShapeTreeToShapedBuffer()
295 std::vector<ShapedBuffer> shaped_buffers; in RunAsync()
[all …]
Dlocal_client.h51 const absl::Span<const ShapedBuffer* const> arguments,
62 const absl::Span<const ShapedBuffer* const> arguments,
90 StatusOr<Literal> LiteralFromShapedBuffer(const ShapedBuffer& shaped_buffer);
162 StatusOr<Literal> ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer);
166 StatusOr<const ShapedBuffer*> GlobalDataToShapedBuffer(
/external/tensorflow/tensorflow/stream_executor/tpu/
Dc_api_conversions.h71 void ToC(const xla::ShapedBuffer& buffer, XLA_ShapedBuffer* c_device_buffer);
72 xla::ShapedBuffer FromC(XLA_ShapedBuffer* c_buffer);
86 void ToC(const xla::ShapedBuffer& buffer, XLA_ShapedBuffer* c_device_buffer);
87 xla::ShapedBuffer FromC(XLA_ShapedBuffer* c_buffer);
Dtpu_transfer_manager.h46 const xla::ShapedBuffer& device_buffer,
50 stream_executor::Stream* stream, const xla::ShapedBuffer& device_buffer,
75 const xla::ShapedBuffer& device_buffer) const override;
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.h97 static Status CreateFromBuffer(const xla::ShapedBuffer& shaped_buffer,
105 static Status CreateFromBuffer(const xla::ShapedBuffer& shaped_buffer,
210 xla::StatusOr<xla::ShapedBuffer> ToShapedBuffer();
243 void InitializeFromShapedBuffer(const xla::ShapedBuffer& shaped_buffer,
Dxrt_state.cc221 const xla::ShapedBuffer& shaped_buffer, const xla::Shape& on_host_shape, in CreateFromBuffer()
235 const xla::ShapedBuffer& shaped_buffer, xla::Backend* backend, in CreateFromBuffer()
253 TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, ToShapedBuffer()); in StoreToLiteral()
272 TF_ASSIGN_OR_RETURN(xla::ShapedBuffer shaped_buffer, ToShapedBuffer()); in WriteLiteral()
503 xla::ShapedBuffer tuple_buffers = new_tuple_buffers->release(); in MakeTuple()
574 const xla::ShapedBuffer& shaped_buffer, in InitializeFromShapedBuffer()
586 xla::StatusOr<xla::ShapedBuffer> XRTTupleAllocation::ToShapedBuffer() { in ToShapedBuffer()
587 xla::ShapedBuffer shaped_buffer(on_host_shape(), on_device_shape(), in ToShapedBuffer()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.h61 const xla::ShapedBuffer& shaped_buffer() const { in shaped_buffer()
65 xla::ShapedBuffer& shaped_buffer() { in shaped_buffer()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtracked_device_buffer_test.cc67 ShapedBuffer shaped_a = a_buffer->AsShapedBuffer( in TEST()
69 ShapedBuffer shaped_b = b_buffer->AsShapedBuffer( in TEST()
71 ShapedBuffer shaped_c = c_buffer->AsShapedBuffer( in TEST()
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
Dexecutable_base.cc51 std::vector<ShapedBuffer> argument_buffers; in ExecuteAsyncOnStream()
59 argument_buffers.push_back(ShapedBuffer(buffers.shape(), in ExecuteAsyncOnStream()

12