| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | computation_layout.cc | 27 ComputationLayout::ComputationLayout(const ProgramShape& program_shape, in ComputationLayout() function in xla::ComputationLayout
40 void ComputationLayout::SetToDefaultLayout() { in SetToDefaultLayout()
47 void ComputationLayout::SetToDefaultLayoutIfEmpty() { in SetToDefaultLayoutIfEmpty()
58 bool ComputationLayout::LayoutIsSet() const { in LayoutIsSet()
64 string ComputationLayout::ToString() const { in ToString()
73 ProgramShape ComputationLayout::ComputeProgramShape() const { in ComputeProgramShape()
83 bool ComputationLayout::operator==(const ComputationLayout& other) const { in operator ==()
88 bool ComputationLayout::operator!=(const ComputationLayout& other) const { in operator !=()
93 uint64 ComputationLayout::Hash() const { in Hash()
|
| D | computation_layout.h | 33 class ComputationLayout {
36 explicit ComputationLayout(ShapeLayout result_layout) in ComputationLayout() function
42 explicit ComputationLayout(const ProgramShape& program_shape,
92 bool operator==(const ComputationLayout& other) const;
93 bool operator!=(const ComputationLayout& other) const;
|
| D | hlo_module_config.cc | 34 ComputationLayout(program_shape, ignore_layouts)) {} in HloModuleConfig()
36 HloModuleConfig::HloModuleConfig(ComputationLayout entry_computation_layout) in HloModuleConfig()
41 entry_computation_layout_ = ComputationLayout(program_shape); in SetDefaultComputationLayout()
46 entry_computation_layout_ = ComputationLayout(program_shape, in SetComputationLayoutIfExists()
|
| D | layout_assignment_test.cc | 54 void AssignLayouts(HloModule* m, ComputationLayout* entry_computation_layout, in AssignLayouts()
90 TEST_F(LayoutAssignmentTest, ComputationLayout) { in TEST_F() argument
111 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
146 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
193 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
234 ComputationLayout computation_layout( in TEST_F()
272 ComputationLayout computation_layout( in TEST_F()
307 ComputationLayout computation_layout( in TEST_F()
371 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
412 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
[all …]
|
| D | layout_assignment.h | 304 ComputationLayout* entry_computation_layout,
388 Status AddMandatoryConstraints(const ComputationLayout* computation_layout,
407 Status RunOnComputation(ComputationLayout* computation_layout,
448 ComputationLayout* computation_layout);
451 ComputationLayout* entry_computation_layout_;
455 ComputationLayout saved_entry_computation_layout_;
513 std::map<HloComputation*, ComputationLayout> computation_layouts_;
516 std::map<HloComputation*, ComputationLayout> conditional_mismatch_;
|
| D | hlo_module_config.h | 77 explicit HloModuleConfig(ComputationLayout entry_computation_layout);
95 const ComputationLayout& entry_computation_layout() const { in entry_computation_layout()
102 ComputationLayout* mutable_entry_computation_layout() { in mutable_entry_computation_layout()
245 absl::optional<ComputationLayout> entry_computation_layout_;
|
| D | hlo_module.h | 139 ComputationLayout compute_computation_layout() const { in compute_computation_layout()
140 return ComputationLayout(entry_computation()->ComputeProgramShape(), in compute_computation_layout()
144 ComputationLayout* mutable_entry_computation_layout() { in mutable_entry_computation_layout()
148 const ComputationLayout& entry_computation_layout() const { in entry_computation_layout()
|
| D | layout_assignment.cc | 445 const ComputationLayout* computation_layout, in AddMandatoryConstraints()
554 const ComputationLayout& called_computation_layout = in AddMandatoryConstraints()
572 ComputationLayout& body_layout = FindOrDie(computation_layouts_, body); in AddMandatoryConstraints()
573 ComputationLayout& condition_layout = in AddMandatoryConstraints()
622 ComputationLayout& best_branch_computation_layout = in AddMandatoryConstraints()
629 ComputationLayout& branch_computation_layout = in AddMandatoryConstraints()
675 const ComputationLayout& computation_layout) { in CheckCallLayout()
709 const ComputationLayout& condition_computation_layout, in CheckWhileLayout()
710 const ComputationLayout& body_computation_layout) { in CheckWhileLayout()
725 absl::Span<const ComputationLayout> branch_computation_layouts) { in CheckConditionalLayout()
[all …]
|
| D | hlo_module_util.cc | 50 ComputationLayout* computation_layout = in CreateModuleConfig()
|
| D | hlo_module.cc | 470 ComputationLayout* entry_layout = in CreateModuleConfigFromShape()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gpu_layout_assignment_test.cc | 70 ComputationLayout computation_layout( in TEST_F()
157 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
226 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
307 ComputationLayout computation_layout( in TEST_F()
355 ComputationLayout computation_layout( in TEST_F()
393 ComputationLayout computation_layout( in TEST_F()
420 ComputationLayout computation_layout( in TEST_F()
|
| D | gpu_layout_assignment.h | 32 explicit GpuLayoutAssignment(ComputationLayout* entry_computation_layout, in GpuLayoutAssignment()
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | cpu_layout_assignment_test.cc | 52 ComputationLayout* entry_computation_layout) { in AssignLayouts()
79 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
120 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
164 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
198 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
238 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
313 ComputationLayout computation_layout(computation->ComputeProgramShape()); in RunDotOutputFusion()
458 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
|
| D | cpu_layout_assignment.h | 32 ComputationLayout* entry_computation_layout, in CpuLayoutAssignment()
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| D | xrt_execute_op.cc | 78 const xla::ComputationLayout& computation_layout) { in GetDynamicInputInfo()
92 const xla::ComputationLayout& computation_layout = in GetInputTuples()
302 const xla::ComputationLayout& computation_layout = in RunExecutable()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_compile_op_support.cc | 31 using ::xla::ComputationLayout;
65 ComputationLayout* computation_layout = in CreateModuleConfig()
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 452 xla::ComputationLayout computation_layout( in TPUExecute()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | local_client.cc | 121 const ComputationLayout& computation_layout = in RunHelper()
|
| /external/tensorflow/tensorflow/stream_executor/tpu/ |
| D | c_api_conversions.cc | 405 xla::ComputationLayout layout(result_layout); in ConfigWithLayout()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 3712 ComputationLayout(new_program_shape, /*ignore_layouts=*/false); in Run()
|