Home
last modified time | relevance | path

Searched refs:ComputationLayout (Results 1 – 21 of 21) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dcomputation_layout.cc27 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 Dlayout_assignment_test.cc54 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 Dlayout_assignment.h137 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 Dcomputation_layout.h33 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 Dhlo_module_config.cc34 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 Dhlo_module_config.h81 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 Dlayout_assignment.cc131 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 Dhlo_module.h143 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 Dhlo_module_util.cc50 ComputationLayout* computation_layout = in CreateModuleConfig()
H A Dhlo_parser.cc515 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 Dhlo_module.cc571 ComputationLayout* entry_layout = in CreateModuleConfigFromShape()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dgpu_layout_assignment_test.cc73 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 Dgpu_layout_assignment.h33 ComputationLayout* entry_computation_layout,
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/
H A Dcpu_layout_assignment_test.cc52 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 Dcpu_layout_assignment.h32 ComputationLayout* entry_computation_layout,
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/
H A Dxrt_execute_op.cc78 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 Dtpu_compile_op_support.cc33 using ::xla::ComputationLayout;
67 ComputationLayout* computation_layout = in CreateModuleConfig()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/
H A Dtpu_execute.cc456 xla::ComputationLayout computation_layout( in TPUExecute()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dlocal_client.cc122 const ComputationLayout& computation_layout = in RunHelper()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/tpu/
H A Dc_api_conversions.cc445 xla::ComputationLayout layout(result_layout); in ConfigWithLayout()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dspmd_partitioner.cc4557 ComputationLayout(new_program_shape, /*ignore_layouts=*/false); in Run()