| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A 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 std::string ComputationLayout::ToString() const { in ToString()
74 ProgramShape ComputationLayout::ComputeProgramShape() const { in ComputeProgramShape()
84 bool ComputationLayout::operator==(const ComputationLayout& other) const { in operator ==()
89 bool ComputationLayout::operator!=(const ComputationLayout& other) const { in operator !=()
|
| H A D | layout_assignment_test.cc | 54 void AssignLayouts(HloModule* m, ComputationLayout* entry_computation_layout, in AssignLayouts()
91 TEST_F(LayoutAssignmentTest, ComputationLayout) { in TEST_F() argument
112 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
147 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
194 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
235 ComputationLayout computation_layout( in TEST_F()
273 ComputationLayout computation_layout( in TEST_F()
337 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
376 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
414 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
[all …]
|
| H A D | layout_assignment.h | 137 ComputationLayout* computation_layout, in ComputationLayoutConstraint()
145 ? ComputationLayout(computation->ComputeProgramShape(), in ComputationLayoutConstraint()
149 const ComputationLayout& computation_layout() const { in computation_layout()
152 void ResetComputationLayout(const ComputationLayout& layout, int64_t priority, in ResetComputationLayout()
186 ComputationLayout computation_layout_;
248 ComputationLayout* entry_computation_layout,
269 ComputationLayout* computation_layout,
287 const ComputationLayout& computation_layout() const { in computation_layout()
385 ComputationLayout& saved_entry_computation_layout() { in saved_entry_computation_layout()
549 ComputationLayout* computation_layout);
[all …]
|
| H A 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;
96 friend H AbslHashValue(H h, const ComputationLayout& computation_layout) { in AbslHashValue()
|
| H A 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()
|
| H A D | hlo_module_config.h | 81 explicit HloModuleConfig(ComputationLayout entry_computation_layout);
99 const ComputationLayout& entry_computation_layout() const { in entry_computation_layout()
106 ComputationLayout* mutable_entry_computation_layout() { in mutable_entry_computation_layout()
352 std::optional<ComputationLayout> entry_computation_layout_;
|
| H A D | layout_assignment.cc | 131 HloComputation* computation, ComputationLayout* computation_layout, in LayoutConstraints()
668 const ComputationLayout& called_computation_layout = in AddMandatoryConstraints()
692 ComputationLayout body_layout = body_constraint->computation_layout(); in AddMandatoryConstraints()
696 ComputationLayout condition_layout = in AddMandatoryConstraints()
757 const ComputationLayout& best_branch_computation_layout = in AddMandatoryConstraints()
765 ComputationLayout branch_computation_layout = in AddMandatoryConstraints()
830 const ComputationLayout& computation_layout) { in CheckCallLayout()
864 const ComputationLayout& condition_computation_layout, in CheckWhileLayout()
865 const ComputationLayout& body_computation_layout) { in CheckWhileLayout()
885 absl::Span<const ComputationLayout> branch_computation_layouts) { in CheckConditionalLayout()
[all …]
|
| H A D | hlo_module.h | 143 ComputationLayout compute_computation_layout() const { in compute_computation_layout()
144 return ComputationLayout(entry_computation()->ComputeProgramShape(), in compute_computation_layout()
148 ComputationLayout* mutable_entry_computation_layout() { in mutable_entry_computation_layout()
152 const ComputationLayout& entry_computation_layout() const { in entry_computation_layout()
|
| H A D | hlo_module_util.cc | 50 ComputationLayout* computation_layout = in CreateModuleConfig()
|
| H A D | hlo_parser.cc | 515 bool ParseComputationLayout(ComputationLayout* computation_layout);
781 ComputationLayout* computation_layout) { in ParseComputationLayout()
926 std::optional<ComputationLayout> entry_computation_layout; in ParseHloModule()
4359 ComputationLayout computation_layout(ShapeLayout(Shape{})); in ParseAttributeHelper()
4363 static_cast<optional<ComputationLayout>*>(attr_out_ptr) in ParseAttributeHelper()
|
| H A D | hlo_module.cc | 571 ComputationLayout* entry_layout = in CreateModuleConfigFromShape()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A D | gpu_layout_assignment_test.cc | 73 ComputationLayout computation_layout( in TEST_F()
109 ComputationLayout computation_layout( in TEST_F()
135 ComputationLayout computation_layout( in TEST_F()
162 ComputationLayout computation_layout( in TEST_F()
188 ComputationLayout computation_layout( in TEST_F()
215 ComputationLayout computation_layout( in TEST_F()
242 ComputationLayout computation_layout( in TEST_F()
277 ComputationLayout computation_layout( in TEST_F()
302 ComputationLayout computation_layout( in TEST_F()
|
| H A D | gpu_layout_assignment.h | 33 ComputationLayout* entry_computation_layout,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| H A D | cpu_layout_assignment_test.cc | 52 ComputationLayout* entry_computation_layout) { in AssignLayouts()
78 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
119 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
163 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
197 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
237 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
312 ComputationLayout computation_layout(computation->ComputeProgramShape()); in RunDotOutputFusion()
456 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F()
|
| H A D | cpu_layout_assignment.h | 32 ComputationLayout* entry_computation_layout,
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| H A D | xrt_execute_op.cc | 78 const xla::ComputationLayout& computation_layout) { in GetDynamicInputInfo()
92 const xla::ComputationLayout& computation_layout = in GetInputTuples()
303 const xla::ComputationLayout& computation_layout = in RunExecutable()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_compile_op_support.cc | 33 using ::xla::ComputationLayout;
67 ComputationLayout* computation_layout = in CreateModuleConfig()
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/ |
| H A D | tpu_execute.cc | 456 xla::ComputationLayout computation_layout( in TPUExecute()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
| H A D | local_client.cc | 122 const ComputationLayout& computation_layout = in RunHelper()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/tpu/ |
| H A D | c_api_conversions.cc | 445 xla::ComputationLayout layout(result_layout); in ConfigWithLayout()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| H A D | spmd_partitioner.cc | 4557 ComputationLayout(new_program_shape, /*ignore_layouts=*/false); in Run()
|