Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dshaped_buffer.cc34 ShapedBuffer::ShapedBuffer(const Shape& on_host_shape, in ShapedBuffer() function in xla::ShapedBuffer
43 ShapedBuffer::ShapedBuffer(ShapedBuffer&& s) in ShapedBuffer() function in xla::ShapedBuffer
55 ShapedBuffer& ShapedBuffer::operator=(ShapedBuffer&& s) { in operator =()
68 ShapedBuffer::~ShapedBuffer() {} in ~ShapedBuffer()
70 void ShapedBuffer::clear() { in clear()
77 string ShapedBuffer::ToString() const { in ToString()
101 std::ostream& operator<<(std::ostream& out, const ShapedBuffer& buffer) { in operator <<()
110 : ShapedBuffer(on_host_shape, on_device_shape, allocator->platform(), in ScopedShapedBuffer()
114 ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, in ScopedShapedBuffer()
116 : ShapedBuffer(std::move(shaped_buffer)), allocator_(allocator) {} in ScopedShapedBuffer()
[all …]
Dshaped_buffer.h37 class ShapedBuffer {
45 ShapedBuffer(const Shape& on_host_shape, const Shape& on_device_shape,
49 ShapedBuffer(ShapedBuffer&& s);
50 ShapedBuffer& operator=(ShapedBuffer&&);
51 ShapedBuffer(const ShapedBuffer&) = delete;
52 ShapedBuffer& operator=(const ShapedBuffer&) = delete;
55 ShapedBuffer(const ScopedShapedBuffer&) = delete;
56 ShapedBuffer& operator=(const ScopedShapedBuffer&) = delete;
58 virtual ~ShapedBuffer();
122 std::ostream& operator<<(std::ostream& out, const ShapedBuffer& buffer);
[all …]
Dtransfer_manager.h72 se::Stream* stream, const ShapedBuffer& device_buffer,
75 se::Stream* stream, const ShapedBuffer& device_buffer) { in TransferLiteralFromDevice()
79 se::Stream* stream, const ShapedBuffer& device_buffer,
83 const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
101 se::Stream* stream, const ShapedBuffer& device_buffer,
105 const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
125 const ShapedBuffer& device_buffer,
129 const ShapedBuffer& device_buffer) { in TransferLiteralToDevice()
151 const ShapedBuffer& device_buffer,
155 const ShapedBuffer& device_buffer) { in TransferLiteralToDeviceAsync()
[all …]
Dallocation_tracker.cc52 static ShapedBuffer ReleaseIfScopedShapedBuffer(ShapedBuffer b) { return b; } in ReleaseIfScopedShapedBuffer()
53 static ShapedBuffer ReleaseIfScopedShapedBuffer(ScopedShapedBuffer b) { in ReleaseIfScopedShapedBuffer()
60 static_assert(std::is_same<ShapedBufferTy, ShapedBuffer>::value || in RegisterInternal()
94 absl::make_unique<ShapedBuffer>( in RegisterInternal()
108 TF_ASSIGN_OR_RETURN(std::vector<const ShapedBuffer*> replicated_buffers, in Unregister()
140 TF_ASSIGN_OR_RETURN(std::vector<const ShapedBuffer*> replicated_buffers, in DeconstructTuple()
144 const ShapedBuffer* shaped_buffer = replicated_buffers[0]; in DeconstructTuple()
161 auto element_buffer = ShapedBuffer( in DeconstructTuple()
167 std::vector<ShapedBuffer> replicated_buffers; in DeconstructTuple()
178 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(
147 absl::flat_hash_map<int64, std::vector<std::unique_ptr<ShapedBuffer>>>
Dtransfer_manager.cc48 se::Stream* stream, const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
75 se::Stream* stream, const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
97 const ShapedBuffer& device_buffer, in TransferLiteralToDevice()
169 ShapedBuffer shaped_buffer(/*on_host_shape=*/literal.shape(), on_device_shape, in TransferArrayToDeviceAsync()
193 ShapedBuffer shaped_buffer(/*on_host_shape=*/shape, shape, in TransferArrayFromDevice()
234 se::Stream* stream, const ShapedBuffer& device_buffer) { in WriteTupleIndexTables()
240 se::Stream* stream, const ShapedBuffer& device_buffer) { in WriteTupleIndexTablesAsync()
268 se::Stream* stream, const ShapedBuffer& device_buffer) { in WriteRootTupleIndexTable()
Dexecutable.h84 absl::Span<const ShapedBuffer* const> arguments,
91 absl::Span<const ShapedBuffer* const> arguments) = 0;
123 absl::Span<const absl::Span<const ShapedBuffer* const>> arguments);
139 absl::Span<const ShapedBuffer* const> arguments);
Dgeneric_transfer_manager.h44 se::Stream* stream, const ShapedBuffer& device_buffer,
50 const ShapedBuffer& device_buffer,
70 const ShapedBuffer& device_buffer,
Dservice.h191 absl::Span<const ShapedBuffer* const> arguments,
200 StatusOr<std::vector<std::vector<const ShapedBuffer*>>> GetArguments(
216 StatusOr<std::vector<std::vector<const ShapedBuffer*>>>
253 absl::Span<const std::vector<const ShapedBuffer*>> arguments,
262 absl::Span<const std::vector<std::vector<const ShapedBuffer*>>> arguments,
Dservice.cc73 Status RecordArguments(const absl::Span<const ShapedBuffer* const> arguments, in RecordArguments()
77 for (const ShapedBuffer* argument : arguments) { in RecordArguments()
87 Status RecordResult(const ShapedBuffer& result, se::Stream* stream, in RecordResult()
229 StatusOr<std::vector<std::vector<const ShapedBuffer*>>>
234 std::vector<std::vector<const ShapedBuffer*>> replicated_arguments; in ResolveAndValidateArguments()
246 const ShapedBuffer* shaped_buffer = replicated_buffers[replica]; in ResolveAndValidateArguments()
337 absl::Span<const ShapedBuffer* const> arguments, in CreateModuleConfig()
396 absl::Span<const std::vector<std::vector<const ShapedBuffer*>>> arguments, in ExecuteParallelAndRegisterResult()
519 absl::Span<const std::vector<const ShapedBuffer*>> arguments, in ExecuteAndRegisterResult()
561 std::vector<absl::Span<const ShapedBuffer* const>> replicated_arguments; in ExecuteAndRegisterResult()
[all …]
Dhlo_runner.h110 StatusOr<Literal> TransferLiteralFromDevice(const ShapedBuffer& buffer);
139 const absl::Span<const ShapedBuffer* const> arguments,
152 const absl::Span<const ShapedBuffer* const> arguments,
Dgeneric_transfer_manager.cc59 se::Stream* stream, const ShapedBuffer& device_buffer, in TransferLiteralFromDevice()
72 se::StreamExecutor* executor, const ShapedBuffer& device_buffer, in TransferLiteralFromDeviceInternal()
101 const ShapedBuffer& device_buffer, in TransferLiteralToDeviceAsync()
Dhlo_runner.cc133 const ShapedBuffer& buffer) { in TransferLiteralFromDevice()
203 const absl::Span<const ShapedBuffer* const> arguments, bool run_hlo_passes, in ExecuteWithDeviceBuffers()
226 std::vector<const ShapedBuffer*> argument_pointers; in ExecuteWithDeviceBuffers()
240 const absl::Span<const ShapedBuffer* const> arguments, in ExecuteWithDeviceBuffers()
261 std::vector<const ShapedBuffer*> argument_pointers; in ExecuteWithDeviceBuffers()
289 std::vector<const ShapedBuffer*> argument_buffer_ptrs( in ExecuteReplicated()
291 std::vector<absl::Span<const ShapedBuffer* const>> argument_buffer_slices; in ExecuteReplicated()
/external/tensorflow/tensorflow/compiler/xla/client/
Dlocal_client.h43 const absl::Span<const ShapedBuffer* const> arguments,
66 const absl::Span<const ShapedBuffer* const> arguments,
76 const absl::Span<const ShapedBuffer* const> arguments);
80 Status RecordArguments(const absl::Span<const ShapedBuffer* const> arguments,
84 Status RecordResult(const ShapedBuffer* result, HloSnapshot* hlo_snapshot);
87 StatusOr<Literal> LiteralFromShapedBuffer(const ShapedBuffer& shaped_buffer);
138 StatusOr<Literal> ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer);
142 StatusOr<const ShapedBuffer*> GlobalDataToShapedBuffer(
Dlocal_client.cc55 const absl::Span<const ShapedBuffer* const> arguments, in ValidateExecutionOptions()
144 const absl::Span<const ShapedBuffer* const> arguments, in Run()
180 const absl::Span<const ShapedBuffer* const> arguments) { in ExecuteAndDump()
194 const absl::Span<const ShapedBuffer* const> arguments, in RecordArguments()
197 for (const ShapedBuffer* argument : arguments) { in RecordArguments()
204 Status LocalExecutable::RecordResult(const ShapedBuffer* result, in RecordResult()
213 const ShapedBuffer& shaped_buffer) { in LiteralFromShapedBuffer()
279 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral()
286 StatusOr<const ShapedBuffer*> LocalClient::GlobalDataToShapedBuffer( in GlobalDataToShapedBuffer()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dlocal_client_test_base.h89 Literal ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer);
95 absl::Span<const ShapedBuffer* const> arguments);
98 absl::Span<const ShapedBuffer* const> arguments,
104 absl::Span<const ShapedBuffer* const> arguments);
107 absl::Span<const ShapedBuffer* const> arguments,
Dlocal_client_test_base.cc140 const ShapedBuffer& shaped_buffer) { in ShapedBufferToLiteral()
159 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteLocallyOrDie()
167 absl::Span<const ShapedBuffer* const> arguments, in ExecuteLocallyOrDie()
176 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteLocally()
183 absl::Span<const ShapedBuffer* const> arguments, in ExecuteLocally()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_executable.h60 absl::Span<const ShapedBuffer* const> arguments,
65 absl::Span<const ShapedBuffer* const> arguments) override;
96 absl::Span<const ShapedBuffer* const> arguments,
116 absl::Span<const ShapedBuffer* const> arguments);
Dcpu_executable.cc79 absl::Span<const ShapedBuffer* const> arguments) { in CreateBufferTable()
269 absl::Span<const ShapedBuffer* const> arguments, in ExecuteOnStream()
280 absl::Span<const ShapedBuffer* const> arguments) { in ExecuteAsyncOnStream()
291 absl::Span<const ShapedBuffer* const> arguments, in ExecuteAsyncOnStreamImpl()
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.h96 static Status CreateFromBuffer(const xla::ShapedBuffer& shaped_buffer,
177 xla::ShapedBuffer ToShapedBuffer();
202 void InitializeFromShapedBuffer(const xla::ShapedBuffer& shaped_buffer,
Dxrt_state.cc210 const xla::ShapedBuffer& shaped_buffer, xla::Backend* backend, in CreateFromBuffer()
229 xla::ShapedBuffer shaped_buffer = ToShapedBuffer(); in ToLiteral()
429 xla::ShapedBuffer tuple_buffers = new_tuple_buffers->release(); in MakeTuple()
500 const xla::ShapedBuffer& shaped_buffer, in InitializeFromShapedBuffer()
509 xla::ShapedBuffer XRTTupleAllocation::ToShapedBuffer() { in ToShapedBuffer()
510 xla::ShapedBuffer shaped_buffer(on_host_shape(), on_device_shape(), in ToShapedBuffer()
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
Dexecutable.h51 absl::Span<const ShapedBuffer* const> arguments,
57 absl::Span<const ShapedBuffer* const> arguments) override;
/external/tensorflow/tensorflow/compiler/jit/
Dxla_launch_util.h185 const std::vector<xla::ShapedBuffer*>& arguments() const { return arg_ptrs_; } in arguments()
192 std::vector<std::unique_ptr<xla::ShapedBuffer>> arg_buffers_;
193 std::vector<xla::ShapedBuffer*> arg_ptrs_;
Dxla_tensor.h63 const xla::ShapedBuffer& shaped_buffer() const { in shaped_buffer()
67 xla::ShapedBuffer& shaped_buffer() { in shaped_buffer()
/external/tensorflow/tensorflow/compiler/xrt/kernels/
Dxrt_execute_op.cc71 std::vector<xla::ShapedBuffer>* input_allocations, in GetComputationInputs()
72 std::vector<xla::ShapedBuffer*>* input_pointers) { in GetComputationInputs()
184 std::vector<xla::ShapedBuffer> input_allocations; in DoWork()
185 std::vector<xla::ShapedBuffer*> input_pointers; in DoWork()

12