Home
last modified time | relevance | path

Searched refs:PjRtBuffer (Results 1 – 15 of 15) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dpjrt_client.h101 class PjRtBuffer; variable
110 std::unique_ptr<PjRtBuffer> buffer;
193 virtual StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
225 virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
233 virtual StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
244 virtual StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
274 class PjRtBuffer {
276 virtual ~PjRtBuffer() = default;
368 virtual StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
460 virtual StatusOr<std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>>
[all …]
Dpjrt_stream_executor_client.h181 StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
183 StatusOr<std::unique_ptr<PjRtBuffer>> CreateUninitializedBuffer(
187 StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostBuffer(
193 StatusOr<std::unique_ptr<PjRtBuffer>> BufferFromHostLiteral(
200 StatusOr<std::unique_ptr<PjRtBuffer>> CreateViewOfDeviceBuffer(
237 std::vector<std::unique_ptr<PjRtBuffer>>&& buffers, in EnqueueCrossHostReceive()
244 PjRtBuffer* buffer, absl::string_view serialized_descriptor) const { in CopyToRemoteDevice()
283 class PjRtStreamExecutorBuffer : public PjRtBuffer {
487 using PjRtBuffer::ToLiteral;
518 StatusOr<std::unique_ptr<PjRtBuffer>> CopyToDevice(
[all …]
Dpjrt_stream_executor_client.cc585 class ScopedHoldAsExternalReference : public PjRtBuffer::ExternalReference {
603 StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
613 : public PjRtBuffer::ExternalReference {
627 StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>>
638 std::unique_ptr<PjRtBuffer::ExternalReference> ref; in ReleaseDeviceMemoryOwnership()
646 StatusOr<std::unique_ptr<PjRtBuffer>>
714 return std::unique_ptr<PjRtBuffer>( in BufferFromHostBuffer()
824 return std::unique_ptr<PjRtBuffer>(std::move(py_buffer)); in BufferFromHostBuffer()
827 StatusOr<std::unique_ptr<PjRtBuffer>>
833 StatusOr<std::unique_ptr<PjRtBuffer>>
[all …]
Dpjrt_client.cc20 PjRtBuffer::ExternalReference::~ExternalReference() = default;
/external/tensorflow/tensorflow/compiler/xla/python/
Djax_jit.h126 std::vector<xla::PjRtBuffer*> arg_buffers;
131 std::vector<std::unique_ptr<xla::PjRtBuffer>> keep_alive;
142 explicit DevicePutResult(xla::PjRtBuffer* b, bool weak_type) in DevicePutResult()
144 DevicePutResult(std::unique_ptr<xla::PjRtBuffer> new_buffer, bool weak_type) in DevicePutResult()
149 xla::PjRtBuffer* buffer;
151 std::unique_ptr<xla::PjRtBuffer> owned_buffer;
Dpy_executable.cc74 const std::vector<PjRtBuffer*>& args) { in PjRtExecute()
75 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in PjRtExecute()
92 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in Execute()
95 std::vector<PjRtBuffer*> arg_buffers(args.size()); in Execute()
114 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in ExecuteOnLocalDevices()
117 std::vector<std::vector<PjRtBuffer*>> arg_buffers(args.size()); in ExecuteOnLocalDevices()
142 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in ExecuteShardedOnLocalDevices()
159 std::vector<std::vector<PjRtBuffer*>> arg_buffers(num_computations); in ExecuteShardedOnLocalDevices()
Dpy_buffer.cc31 std::unique_ptr<PjRtBuffer> buffer, in PyBuffer()
75 std::unique_ptr<PjRtBuffer> out; in CopyToDevice()
121 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference_hold; in AsNumPyArray()
161 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference_hold, in UnsafeBufferPointer()
191 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference_hold, in CudaArrayInterface()
210 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference_hold) in ExtraBufferInfo()
219 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference_hold;
250 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference_hold, in PjRtBufferGetBuffer()
Dpy_buffer.h55 PyBuffer(std::shared_ptr<PyClient> client, std::unique_ptr<PjRtBuffer> buffer,
60 PjRtBuffer* buffer() const { return buffer_.get(); } in buffer()
121 std::unique_ptr<PjRtBuffer> buffer_;
Dpy_client.cc102 StatusOr<std::unique_ptr<PjRtBuffer>> PyClient::PjRtBufferFromPyval( in PjRtBufferFromPyval()
136 std::unique_ptr<PjRtBuffer> buffer; in PjRtBufferFromPyval()
150 std::unique_ptr<PjRtBuffer> buffer, in BufferFromPyval()
Dpy_executable.h64 const std::vector<PjRtBuffer*>& args);
Djax_jit.cc358 std::unique_ptr<xla::PjRtBuffer> ConvertToScalarBuffer( in ConvertToScalarBuffer()
655 std::unique_ptr<xla::PjRtBuffer> copied_buffer = in HandleDeviceArray()
666 std::unique_ptr<xla::PjRtBuffer> buffer = in HandleBufferFromPyval()
690 std::unique_ptr<xla::PjRtBuffer> buffer = in HandleUint64()
701 std::unique_ptr<xla::PjRtBuffer> buffer = in HandleUint64()
730 std::unique_ptr<xla::PjRtBuffer> buffer = in HandleNdarray()
945 std::vector<xla::PjRtBuffer*>& arg_buffers = arguments.arg_buffers; in ConvertArgsToBuffers()
1012 xla::PjRtBuffer* buffer = on_device.buffer; in ConvertArgsToBuffers()
Ddlpack.cc47 std::unique_ptr<PjRtBuffer::ExternalReference> external_reference;
264 StatusOr<std::unique_ptr<PjRtBuffer::ExternalReference>> buffer_or = in BufferToDLPackManagedTensor()
Dpy_client.h128 StatusOr<std::unique_ptr<PjRtBuffer>> PjRtBufferFromPyval(
Doutfeed_receiver.cc412 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers, in SendShutdownOutfeedHeader()
Doutfeed_receiver_test.cc50 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers, in CompileAndExecute()