| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | tpu_executable_interface.cc | 55 const Shape& host_shape, const HloInputOutputAliasConfig& alias_config, in AllocateOutputMemoryWithInputReuse() argument
62 << " host_shape = " << ShapeUtil::HumanStringWithLayout(host_shape); in AllocateOutputMemoryWithInputReuse()
63 Shape device_shape = HostShapeToDeviceShape(host_shape); in AllocateOutputMemoryWithInputReuse()
85 << " host_shape = " << ShapeUtil::HumanStringWithLayout(host_shape); in AllocateOutputMemoryWithInputReuse()
86 if (!Shape::Equal().MinorToMajorOnlyInLayout()(host_shape, device_shape)) { in AllocateOutputMemoryWithInputReuse()
88 << ShapeUtil::HumanStringWithLayout(host_shape) << " -> " in AllocateOutputMemoryWithInputReuse()
|
| D | tpu_executable_interface.h | 67 const Shape& host_shape, const HloInputOutputAliasConfig& alias_config,
82 virtual Shape HostShapeToDeviceShape(const Shape& host_shape) = 0;
|
| D | tpu_transfer_manager.h | 42 const xla::Shape& host_shape) const override;
71 const xla::Shape& host_shape) const override;
|
| D | tpu_transfer_manager.cc | 53 const xla::Shape& host_shape) const { in HostShapeToDeviceShape()
57 ApiConverter::ToC(host_shape, &c_host_shape); in HostShapeToDeviceShape()
225 const xla::Shape& host_shape) const { in ChooseCompactLayoutForShape()
227 ApiConverter::ToC(host_shape, &c_host_shape); in ChooseCompactLayoutForShape()
|
| D | tpu_executable.cc | 106 Shape TpuExecutable::HostShapeToDeviceShape(const Shape& host_shape) { in HostShapeToDeviceShape() argument
109 ApiConverter::ToC(host_shape, &c_host_shape); in HostShapeToDeviceShape()
|
| D | tpu_executable.h | 59 Shape HostShapeToDeviceShape(const Shape& host_shape) override;
|
| D | tpu_executor_c_api.h | 180 XLA_Shape* host_shape,
192 XLA_TransferManager* manager, XLA_Shape* host_shape, XLA_Shape* output,
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | executable.h | 67 ExecutionInput(xla::Shape shape, xla::Shape host_shape) in ExecutionInput() argument
78 xla::Shape host_shape) in ExecutionInput() argument
93 const Shape& host_shape() const { in host_shape() function
132 void SetHostShape(xla::Shape host_shape) { in SetHostShape() argument
133 if (shape() != host_shape) { in SetHostShape()
134 host_shape_ = absl::make_unique<Shape>(std::move(host_shape)); in SetHostShape()
|
| D | transfer_manager.h | 53 virtual Shape HostShapeToDeviceShape(const Shape& host_shape) const { in HostShapeToDeviceShape() argument
58 return ShapeUtil::DeviceShapeToHostShape(host_shape); in HostShapeToDeviceShape()
244 const Shape& host_shape) const;
|
| D | transfer_manager.cc | 421 const Shape& host_shape) const { in ChooseCompactLayoutForShape()
422 return LayoutUtil::GetWithDefaultLayout(host_shape); in ChooseCompactLayoutForShape()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 170 xla::ShapedBuffer ToShapedBuffer(xla::Shape host_shape, in ToShapedBuffer()
174 xla::ShapedBuffer shaped_buffer(std::move(host_shape), buffers.shape(), in ToShapedBuffer()
321 const xla::Shape& host_shape = in BuildComputationInputs() local
325 host_shape, allocator, device_ordinal)); in BuildComputationInputs()
663 xla::Shape host_shape(executable.input_shapes(0)); in DoWork() local
673 BuildComputationInputs(context, host_shape, variable_update_map, backend, in DoWork()
690 auto shaped_buffer = input_buffers->ToShapedBuffer(std::move(host_shape), in DoWork()
|
| /external/tensorflow/tensorflow/compiler/xrt/ |
| D | xrt_state.cc | 415 se::DeviceMemoryAllocator* allocator, xla::Shape* host_shape, in ExpandTreeOfTuples() argument
419 *host_shape = elements.shape(); in ExpandTreeOfTuples()
438 *xla::ShapeUtil::GetMutableSubshape(host_shape, index) = in ExpandTreeOfTuples()
462 xla::Shape host_shape; in MakeTuple() local
465 &host_shape, &device_shape)); in MakeTuple()
471 device_ordinal, allocator, host_shape, device_shape); in MakeTuple()
|
| D | xrt_state.h | 269 se::DeviceMemoryAllocator* allocator, xla::Shape* host_shape,
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 106 xla::Shape HostShapeToDeviceShape(const xla::Shape& host_shape) { in HostShapeToDeviceShape() argument
109 ApiConverter::ToC(host_shape, &c_host_shape); in HostShapeToDeviceShape()
|
| D | tpu_ops_c_api.h | 178 XLA_Shape* host_shape, XLA_Shape* device_shape);
|
| D | tpu_on_demand_compiler.cc | 198 Shape HostShapeToDeviceShape(const Shape& host_shape) override { in HostShapeToDeviceShape() argument
|