/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()
|