Home
last modified time | relevance | path

Searched refs:ScopedShapedBuffer (Results 1 – 25 of 43) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dshaped_buffer.h33 class ScopedShapedBuffer; variable
55 ShapedBuffer(const ScopedShapedBuffer&) = delete;
56 ShapedBuffer& operator=(const ScopedShapedBuffer&) = delete;
133 class ScopedShapedBuffer : public ShapedBuffer {
136 explicit ScopedShapedBuffer(const Shape& on_host_shape,
143 explicit ScopedShapedBuffer(ShapedBuffer shaped_buffer,
147 ScopedShapedBuffer(ScopedShapedBuffer&& s);
148 ScopedShapedBuffer& operator=(ScopedShapedBuffer&&);
149 ScopedShapedBuffer(const ScopedShapedBuffer&) = delete;
150 ScopedShapedBuffer& operator=(const ScopedShapedBuffer&) = delete;
[all …]
Dshaped_buffer.cc106 ScopedShapedBuffer::ScopedShapedBuffer(const Shape& on_host_shape, in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
114 ScopedShapedBuffer::ScopedShapedBuffer(ShapedBuffer shaped_buffer, in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
118 ScopedShapedBuffer::ScopedShapedBuffer(ScopedShapedBuffer&& s) in ScopedShapedBuffer() function in xla::ScopedShapedBuffer
124 ScopedShapedBuffer& ScopedShapedBuffer::operator=(ScopedShapedBuffer&& s) { in operator =()
134 ScopedShapedBuffer::~ScopedShapedBuffer() { Deallocate(); } in ~ScopedShapedBuffer()
136 ShapedBuffer ScopedShapedBuffer::release() { in release()
142 void ScopedShapedBuffer::Deallocate() { in Deallocate()
160 ScopedShapedBuffer ScopedShapedBuffer::TakeSubTree(ShapeIndexView index) { in TakeSubTree()
166 ScopedShapedBuffer output(sub_on_host_shape, sub_on_device_shape, in TakeSubTree()
Dhlo_runner.cc97 StatusOr<ScopedShapedBuffer> HloRunner::TransferLiteralToDevice( in TransferLiteralToDevice()
99 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer buffer, in TransferLiteralToDevice()
110 StatusOr<std::vector<ScopedShapedBuffer>> HloRunner::TransferLiteralsToDevice( in TransferLiteralsToDevice()
112 std::vector<ScopedShapedBuffer> buffers; in TransferLiteralsToDevice()
115 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer buffer, in TransferLiteralsToDevice()
122 StatusOr<std::vector<ScopedShapedBuffer>> HloRunner::TransferLiteralsToDevice( in TransferLiteralsToDevice()
144 TF_ASSIGN_OR_RETURN(std::vector<ScopedShapedBuffer> argument_buffers, in Execute()
146 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, in Execute()
176 TF_ASSIGN_OR_RETURN(std::vector<ScopedShapedBuffer> argument_buffers, in Execute()
178 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, in Execute()
[all …]
Dhlo_runner.h105 StatusOr<ScopedShapedBuffer> TransferLiteralToDevice(const Literal& literal);
106 StatusOr<std::vector<ScopedShapedBuffer>> TransferLiteralsToDevice(
108 StatusOr<std::vector<ScopedShapedBuffer>> TransferLiteralsToDevice(
137 StatusOr<ScopedShapedBuffer> ExecuteWithDeviceBuffers(
142 StatusOr<ScopedShapedBuffer> ExecuteWithDeviceBuffers(
144 const absl::Span<const ScopedShapedBuffer> arguments,
150 StatusOr<ScopedShapedBuffer> ExecuteWithDeviceBuffers(
155 StatusOr<ScopedShapedBuffer> ExecuteWithDeviceBuffers(
157 const absl::Span<const ScopedShapedBuffer> arguments,
Dexecutable.h49 ExecutionOutput(ScopedShapedBuffer result, in ExecutionOutput()
52 ScopedShapedBuffer result;
82 virtual StatusOr<ScopedShapedBuffer> ExecuteOnStream(
89 virtual StatusOr<ScopedShapedBuffer> ExecuteAsyncOnStream(
121 virtual StatusOr<std::vector<ScopedShapedBuffer>> ExecuteOnStreams(
137 StatusOr<ScopedShapedBuffer> ExecuteOnStreamWrapper(
Dshaped_buffer_test.cc40 auto scoped_buffer = absl::make_unique<xla::ScopedShapedBuffer>( in TEST()
97 ScopedShapedBuffer sb1(s, s, &allocator, /*device_ordinal=*/0); in TEST()
102 ScopedShapedBuffer sb2(s, s, &allocator, /*device_ordinal=*/1); in TEST()
119 ScopedShapedBuffer sb(s, s, &allocator, /*device_ordinal=*/0); in TEST()
131 ScopedShapedBuffer output = sb.TakeSubTree(subtree_index); in TEST()
161 xla::ScopedShapedBuffer shaped_buffer(shape, shape, /*allocator=*/&allocator, in BM_TakeSubTree()
Dexecutable.cc32 StatusOr<std::vector<ScopedShapedBuffer>> Executable::ExecuteOnStreams( in ExecuteOnStreams()
37 std::vector<ScopedShapedBuffer> return_values; in ExecuteOnStreams()
63 StatusOr<ScopedShapedBuffer> Executable::ExecuteOnStreamWrapper( in ExecuteOnStreamWrapper()
83 StatusOr<ScopedShapedBuffer> return_value = in ExecuteOnStreamWrapper()
Dallocation_tracker.cc34 ScopedShapedBuffer shaped_buffer, const string& tag) { in Register()
37 std::vector<ScopedShapedBuffer> replicated_buffers; in Register()
43 std::vector<ScopedShapedBuffer> replicated_buffers, const string& tag) { in RegisterReplicatedBuffers()
53 static ShapedBuffer ReleaseIfScopedShapedBuffer(ScopedShapedBuffer b) { in ReleaseIfScopedShapedBuffer()
61 std::is_same<ShapedBufferTy, ScopedShapedBuffer>::value, in RegisterInternal()
Dallocation_tracker.h49 StatusOr<GlobalDataHandle> Register(ScopedShapedBuffer shaped_buffer,
55 std::vector<ScopedShapedBuffer> replicated_buffers, const string& tag);
/external/tensorflow/tensorflow/compiler/xla/tests/
Dlocal_client_test_base.h85 ScopedShapedBuffer LiteralToShapedBuffer(const Literal& literal);
93 StatusOr<ScopedShapedBuffer> ExecuteLocally(
96 StatusOr<ScopedShapedBuffer> ExecuteLocally(
102 ScopedShapedBuffer ExecuteLocallyOrDie(
105 ScopedShapedBuffer ExecuteLocallyOrDie(
Dlocal_client_execute_test.cc59 ScopedShapedBuffer result = in XLA_TEST_F()
72 ScopedShapedBuffer result = in XLA_TEST_F()
85 ScopedShapedBuffer result = in XLA_TEST_F()
99 ScopedShapedBuffer result = in XLA_TEST_F()
114 ScopedShapedBuffer result = ExecuteLocallyOrDie( in XLA_TEST_F()
142 ScopedShapedBuffer result_colmaj = in XLA_TEST_F()
149 ScopedShapedBuffer result_param_swap = in XLA_TEST_F()
169 ScopedShapedBuffer result_colmaj = ExecuteLocallyOrDie( in XLA_TEST_F()
181 ScopedShapedBuffer result_rowmaj = ExecuteLocallyOrDie( in XLA_TEST_F()
205 ScopedShapedBuffer result = in XLA_TEST_F()
[all …]
Dlocal_client_test_base.cc132 ScopedShapedBuffer LocalClientTestBase::LiteralToShapedBuffer( in LiteralToShapedBuffer()
157 ScopedShapedBuffer LocalClientTestBase::ExecuteLocallyOrDie( in ExecuteLocallyOrDie()
165 ScopedShapedBuffer LocalClientTestBase::ExecuteLocallyOrDie( in ExecuteLocallyOrDie()
174 StatusOr<ScopedShapedBuffer> LocalClientTestBase::ExecuteLocally( in ExecuteLocally()
181 StatusOr<ScopedShapedBuffer> LocalClientTestBase::ExecuteLocally( in ExecuteLocally()
Dmultiple_devices_on_host_test.cc39 std::vector<std::pair<int, StatusOr<ScopedShapedBuffer>>>* results) { in CompileAndExecute()
48 StatusOr<ScopedShapedBuffer> result = executable->Run({}, execute_options); in CompileAndExecute()
72 std::vector<std::pair<int, StatusOr<ScopedShapedBuffer>>> results; in TestWithDeviceCount()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_executable.h58 StatusOr<ScopedShapedBuffer> ExecuteOnStream(
63 StatusOr<ScopedShapedBuffer> ExecuteAsyncOnStream(
94 StatusOr<ScopedShapedBuffer> ExecuteAsyncOnStreamImpl(
127 StatusOr<ScopedShapedBuffer> CreateResultShapedBuffer(
Dcpu_executable.cc208 StatusOr<ScopedShapedBuffer> CpuExecutable::CreateResultShapedBuffer( in CreateResultShapedBuffer()
212 ScopedShapedBuffer result_buffer( in CreateResultShapedBuffer()
267 StatusOr<ScopedShapedBuffer> CpuExecutable::ExecuteOnStream( in ExecuteOnStream()
278 StatusOr<ScopedShapedBuffer> CpuExecutable::ExecuteAsyncOnStream( in ExecuteAsyncOnStream()
289 StatusOr<ScopedShapedBuffer> CpuExecutable::ExecuteAsyncOnStreamImpl( in ExecuteAsyncOnStreamImpl()
309 ScopedShapedBuffer result, in ExecuteAsyncOnStreamImpl()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.h72 void set_shaped_buffer(xla::ScopedShapedBuffer shaped_buffer) { in set_shaped_buffer()
74 absl::make_unique<xla::ScopedShapedBuffer>(std::move(shaped_buffer)); in set_shaped_buffer()
112 std::unique_ptr<xla::ScopedShapedBuffer> shaped_buffer_;
/external/tensorflow/tensorflow/compiler/xla/client/
Dlocal_client.cc143 StatusOr<ScopedShapedBuffer> LocalExecutable::Run( in Run()
178 StatusOr<ScopedShapedBuffer> LocalExecutable::ExecuteAndDump( in ExecuteAndDump()
185 ScopedShapedBuffer result, in ExecuteAndDump()
262 StatusOr<ScopedShapedBuffer> LocalClient::LiteralToShapedBuffer( in LiteralToShapedBuffer()
318 ::xla::ScopedShapedBuffer shaped_buffer, in TransferToLocalServer()
325 std::vector<::xla::ScopedShapedBuffer> replicated_buffer; in TransferToLocalServer()
Dlocal_client.h42 StatusOr<ScopedShapedBuffer> Run(
74 StatusOr<ScopedShapedBuffer> ExecuteAndDump(
128 StatusOr<ScopedShapedBuffer> LiteralToShapedBuffer(
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
Dexecutable.cc48 StatusOr<ScopedShapedBuffer> InterpreterExecutable::ExecuteOnStream( in ExecuteOnStream()
106 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, in ExecuteOnStream()
124 StatusOr<ScopedShapedBuffer> InterpreterExecutable::ExecuteAsyncOnStream( in ExecuteAsyncOnStream()
Dexecutable.h49 StatusOr<ScopedShapedBuffer> ExecuteOnStream(
55 StatusOr<ScopedShapedBuffer> ExecuteAsyncOnStream(
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dbuffer_comparator.h54 ScopedShapedBuffer ref_buffer) in F16BufferComparator()
65 ScopedShapedBuffer ref_buffer_;
Dbuffer_comparator.cc106 auto shaped_buffer, ([&]() -> StatusOr<ScopedShapedBuffer> { in Create()
115 ScopedShapedBuffer ret(shape, shape, allocator, device_ordinal); in Create()
134 auto result_buffer, ([&]() -> StatusOr<ScopedShapedBuffer> { in CompareEqualImpl()
Dgpu_executable.h79 StatusOr<ScopedShapedBuffer> ExecuteOnStream(
84 StatusOr<ScopedShapedBuffer> ExecuteAsyncOnStream(
Dgpu_executable.cc236 StatusOr<ScopedShapedBuffer> GpuExecutable::ExecuteOnStream( in ExecuteOnStream()
290 ScopedShapedBuffer shaped_buffer(root->shape(), root->shape(), in ExecuteOnStream()
349 StatusOr<ScopedShapedBuffer> GpuExecutable::ExecuteAsyncOnStream( in ExecuteAsyncOnStream()
/external/tensorflow/tensorflow/compiler/xla/tools/
Dreplay_computation.cc205 std::vector<ScopedShapedBuffer> scoped_shaped_buffer_arguments; in ReplayComputation()
224 ScopedShapedBuffer data, in ReplayComputation()
277 absl::optional<ScopedShapedBuffer> final_result; in ReplayComputation()
290 TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, in ReplayComputation()

12