| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A D | layout_assignment_test.cc | 112 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
113 *computation_layout.mutable_parameter_layout(0) = shape_layout; in TEST_F()
114 *computation_layout.mutable_parameter_layout(1) = shape_layout; in TEST_F()
115 *computation_layout.mutable_result_layout() = shape_layout; in TEST_F()
116 AssignLayouts(m.get(), &computation_layout); in TEST_F()
147 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
148 *computation_layout.mutable_parameter_layout(0) = col_major; in TEST_F()
149 *computation_layout.mutable_parameter_layout(1) = row_major; in TEST_F()
150 *computation_layout.mutable_result_layout() = col_major; in TEST_F()
152 AssignLayouts(m.get(), &computation_layout); in TEST_F()
[all …]
|
| H A D | layout_assignment.h | 137 ComputationLayout* computation_layout, in ComputationLayoutConstraint() argument
140 layout_state_((computation_layout == nullptr) in ComputationLayoutConstraint()
144 (computation_layout == nullptr) in ComputationLayoutConstraint()
147 : *computation_layout) {} in ComputationLayoutConstraint()
149 const ComputationLayout& computation_layout() const { in computation_layout() function
269 ComputationLayout* computation_layout,
287 const ComputationLayout& computation_layout() const { in computation_layout() function
288 return computation_constraint_.computation_layout(); in computation_layout()
549 ComputationLayout* computation_layout);
|
| H A D | layout_assignment.cc | 131 HloComputation* computation, ComputationLayout* computation_layout, in LayoutConstraints() argument
134 computation_constraint_(computation, computation_layout, priority) {} in LayoutConstraints()
458 ? &computation_layout().result_layout() in ResultLayout()
609 constraints->computation_layout().parameter_layout( in AddMandatoryConstraints()
670 ->computation_layout(); in AddMandatoryConstraints()
692 ComputationLayout body_layout = body_constraint->computation_layout(); in AddMandatoryConstraints()
697 condition_constraint->computation_layout(); in AddMandatoryConstraints()
760 ->computation_layout(); in AddMandatoryConstraints()
767 ->computation_layout(); in AddMandatoryConstraints()
830 const ComputationLayout& computation_layout) { in CheckCallLayout() argument
[all …]
|
| H A D | computation_layout.h | 96 friend H AbslHashValue(H h, const ComputationLayout& computation_layout) { in AbslHashValue() argument
97 h = H::combine(std::move(h), computation_layout.result_layout_.shape()); in AbslHashValue()
98 for (const auto& parameter_layout : computation_layout.parameter_layouts_) { in AbslHashValue()
101 h = H::combine(std::move(h), computation_layout.parameter_layouts_.size()); in AbslHashValue()
|
| H A D | hlo_module_util.cc | 50 ComputationLayout* computation_layout = in CreateModuleConfig() local
70 computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( in CreateModuleConfig()
80 computation_layout->mutable_result_layout()->CopyLayoutFromShape( in CreateModuleConfig()
84 computation_layout->mutable_result_layout()->SetToDefaultLayout(); in CreateModuleConfig()
|
| H A D | BUILD | 1140 ":computation_layout",
1189 ":computation_layout",
1225 ":computation_layout",
1346 ":computation_layout",
4389 ":computation_layout",
4840 ":computation_layout",
5219 ":computation_layout",
5233 name = "computation_layout",
5234 srcs = ["computation_layout.cc"],
5235 hdrs = ["computation_layout.h"],
[all …]
|
| H A D | hlo_parser.cc | 515 bool ParseComputationLayout(ComputationLayout* computation_layout);
781 ComputationLayout* computation_layout) { in ParseComputationLayout() argument
794 computation_layout->add_parameter_layout(ShapeLayout(param)); in ParseComputationLayout()
814 *computation_layout->mutable_result_layout() = ShapeLayout(result); in ParseComputationLayout()
4359 ComputationLayout computation_layout(ShapeLayout(Shape{})); in ParseAttributeHelper() local
4360 if (!ParseComputationLayout(&computation_layout)) { in ParseAttributeHelper()
4364 ->emplace(computation_layout); in ParseAttributeHelper()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| H A D | cpu_layout_assignment_test.cc | 78 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
79 *computation_layout.mutable_parameter_layout(0) = in TEST_F()
81 *computation_layout.mutable_result_layout() = in TEST_F()
83 AssignLayouts(module.get(), &computation_layout); in TEST_F()
119 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
120 *computation_layout.mutable_parameter_layout(0) = in TEST_F()
122 *computation_layout.mutable_result_layout() = in TEST_F()
124 AssignLayouts(module.get(), &computation_layout); in TEST_F()
163 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
164 *computation_layout.mutable_parameter_layout(0) = in TEST_F()
[all …]
|
| H A D | BUILD | 469 "//tensorflow/compiler/xla/service:computation_layout",
1197 "//tensorflow/compiler/xla/service:computation_layout",
1228 "//tensorflow/compiler/xla/service:computation_layout",
1327 "//tensorflow/compiler/xla/service:computation_layout",
|
| /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() local
75 *computation_layout.mutable_parameter_layout(0) = in TEST_F()
77 *computation_layout.mutable_parameter_layout(1) = in TEST_F()
79 *computation_layout.mutable_result_layout() = in TEST_F()
83 &computation_layout, backend().default_stream_executor()); in TEST_F()
109 ComputationLayout computation_layout( in TEST_F() local
112 GpuLayoutAssignment layout_assignment(&computation_layout, in TEST_F()
135 ComputationLayout computation_layout( in TEST_F() local
138 GpuLayoutAssignment layout_assignment(&computation_layout, in TEST_F()
162 ComputationLayout computation_layout( in TEST_F() local
[all …]
|
| H A D | BUILD | 1230 "//tensorflow/compiler/xla/service:computation_layout",
2020 "//tensorflow/compiler/xla/service:computation_layout",
2041 "//tensorflow/compiler/xla/service: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() argument
80 input_is_dynamic.reserve(computation_layout.parameter_count()); in GetDynamicInputInfo()
81 for (int64_t i = 0; i < computation_layout.parameter_count(); ++i) { in GetDynamicInputInfo()
83 !computation_layout.parameter_shape(i).is_static()); in GetDynamicInputInfo()
92 const xla::ComputationLayout& computation_layout = in GetInputTuples() local
96 input_coords, working_set, backend, computation_layout.parameter_count(), in GetInputTuples()
97 [&](int64_t i) { return computation_layout.parameter_shape(i); }, in GetInputTuples()
303 const xla::ComputationLayout& computation_layout = in RunExecutable() local
305 std::vector<bool> input_is_dynamic = GetDynamicInputInfo(computation_layout); in RunExecutable()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
| H A D | local_client.cc | 122 const ComputationLayout& computation_layout = in RunHelper() local
127 if (argument_shapes_size != computation_layout.parameter_count()) { in RunHelper()
130 computation_layout.parameter_count(), argument_shapes.size()); in RunHelper()
134 if (!computation_layout.parameter_layout(i).MatchesLayoutInShape( in RunHelper()
144 computation_layout.parameter_layout(i).shape()), in RunHelper()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| H A D | tfrt_cpu_pjrt_client.cc | 1185 const auto& computation_layout = in TfrtCpuExecutable() local
1187 if (computation_layout.parameter_count() == 0) { in TfrtCpuExecutable()
1192 if (computation_layout.parameter_count() > 1 || in TfrtCpuExecutable()
1193 !computation_layout.parameter_shape(0).IsTuple()) { in TfrtCpuExecutable()
1194 input_buffer_sizes_in_bytes_.reserve(computation_layout.parameter_count()); in TfrtCpuExecutable()
1195 for (int i = 0; i < computation_layout.parameter_count(); ++i) { in TfrtCpuExecutable()
1197 ShapeUtil::ByteSizeOf(computation_layout.parameter_shape(i))); in TfrtCpuExecutable()
1201 computation_layout.parameter_shape(0).tuple_shapes_size()); in TfrtCpuExecutable()
1203 i < computation_layout.parameter_shape(0).tuple_shapes_size(); ++i) { in TfrtCpuExecutable()
1205 computation_layout.parameter_shape(0).tuple_shapes(i))); in TfrtCpuExecutable()
|
| H A D | pjrt_stream_executor_client.cc | 1778 const auto& computation_layout = in PjRtStreamExecutorExecutable() local
1781 parameter_shapes.reserve(computation_layout.parameter_count()); in PjRtStreamExecutorExecutable()
1782 for (int i = 0; i < computation_layout.parameter_count(); ++i) { in PjRtStreamExecutorExecutable()
1784 computation_layout.parameter_shape(i))); in PjRtStreamExecutorExecutable()
|
| H A D | BUILD | 318 "//tensorflow/compiler/xla/service:computation_layout",
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
| H A D | tpu_compile_op_support.cc | 67 ComputationLayout* computation_layout = in CreateModuleConfig() local
86 computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( in CreateModuleConfig()
94 computation_layout->mutable_result_layout()->CopyLayoutFromShape( in CreateModuleConfig()
98 computation_layout->mutable_result_layout()->SetToDefaultLayout(); in CreateModuleConfig()
|
| H A D | BUILD | 299 "//tensorflow/compiler/xla/service:computation_layout",
|
| /aosp_15_r20/external/tensorflow/tensorflow/core/tpu/ |
| H A D | tpu_execute.cc | 456 xla::ComputationLayout computation_layout( in TPUExecute() local
460 computation_layout.add_parameter_layout(xla::ShapeLayout(shape)); in TPUExecute()
465 xla::HloModuleConfig(std::move(computation_layout))); in TPUExecute()
|
| H A D | BUILD | 318 "//tensorflow/compiler/xla/service:computation_layout",
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/ |
| H A D | BUILD | 164 "//tensorflow/compiler/xla/service:computation_layout",
|