| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | layout_assignment_test.cc | 111 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
112 *computation_layout.mutable_parameter_layout(0) = shape_layout; in TEST_F()
113 *computation_layout.mutable_parameter_layout(1) = shape_layout; in TEST_F()
114 *computation_layout.mutable_result_layout() = shape_layout; in TEST_F()
115 AssignLayouts(m.get(), &computation_layout); in TEST_F()
146 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
147 *computation_layout.mutable_parameter_layout(0) = col_major; in TEST_F()
148 *computation_layout.mutable_parameter_layout(1) = row_major; in TEST_F()
149 *computation_layout.mutable_result_layout() = col_major; in TEST_F()
151 AssignLayouts(m.get(), &computation_layout); in TEST_F()
[all …]
|
| 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()
|
| D | layout_assignment.cc | 511 const ComputationLayout* computation_layout, in AddMandatoryConstraints() argument
538 if (computation_layout != nullptr) { in AddMandatoryConstraints()
540 computation_layout->parameter_layout( in AddMandatoryConstraints()
743 } else if (computation_layout != nullptr) { in AddMandatoryConstraints()
744 const ShapeLayout& result_layout = computation_layout->result_layout(); in AddMandatoryConstraints()
762 const ComputationLayout& computation_layout) { in CheckCallLayout() argument
766 TF_RET_CHECK(computation_layout.parameter_layout(i).MatchesLayoutInShape( in CheckCallLayout()
769 TF_RET_CHECK(computation_layout.result_layout().MatchesLayoutInShape( in CheckCallLayout()
851 const ComputationLayout& computation_layout) { in CheckParameterLayout() argument
853 computation_layout.parameter_layout(parameter->parameter_number()); in CheckParameterLayout()
[all …]
|
| D | layout_assignment.h | 407 Status AddMandatoryConstraints(const ComputationLayout* computation_layout,
426 Status RunOnComputation(ComputationLayout* computation_layout,
467 ComputationLayout* computation_layout);
|
| D | BUILD | 1047 ":computation_layout",
1096 ":computation_layout",
1132 ":computation_layout",
1247 ":computation_layout",
4010 ":computation_layout",
4447 ":computation_layout",
4824 ":computation_layout",
4838 name = "computation_layout",
4839 srcs = ["computation_layout.cc"],
4840 hdrs = ["computation_layout.h"],
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| 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 …]
|
| D | BUILD | 368 "//tensorflow/compiler/xla/service:computation_layout",
982 "//tensorflow/compiler/xla/service:computation_layout",
1004 "//tensorflow/compiler/xla/service:computation_layout",
1102 "//tensorflow/compiler/xla/service:computation_layout",
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | gpu_layout_assignment_test.cc | 70 ComputationLayout computation_layout( in TEST_F() local
72 *computation_layout.mutable_parameter_layout(0) = in TEST_F()
74 *computation_layout.mutable_parameter_layout(1) = in TEST_F()
76 *computation_layout.mutable_result_layout() = in TEST_F()
80 &computation_layout, backend().default_stream_executor()); in TEST_F()
156 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
159 *computation_layout.mutable_parameter_layout(0) = in TEST_F()
164 *computation_layout.mutable_result_layout() = ShapeLayout(result_shape); in TEST_F()
168 &computation_layout, backend().default_stream_executor()); in TEST_F()
224 ComputationLayout computation_layout(computation->ComputeProgramShape()); in TEST_F() local
[all …]
|
| D | BUILD | 1172 "//tensorflow/compiler/xla/service:computation_layout",
1758 "//tensorflow/compiler/xla/service:computation_layout",
1778 "//tensorflow/compiler/xla/service:computation_layout",
|
| /external/tensorflow/tensorflow/compiler/xrt/kernels/ |
| 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()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | local_client.cc | 121 const ComputationLayout& computation_layout = in RunHelper() local
126 if (argument_shapes_size != computation_layout.parameter_count()) { in RunHelper()
129 computation_layout.parameter_count(), argument_shapes.size()); in RunHelper()
133 if (!computation_layout.parameter_layout(i).MatchesLayoutInShape( in RunHelper()
143 computation_layout.parameter_layout(i).shape()), in RunHelper()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| 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()
|
| D | BUILD | 205 "//tensorflow/compiler/xla/service:computation_layout",
|
| /external/tensorflow/tensorflow/core/tpu/ |
| D | tpu_execute.cc | 452 xla::ComputationLayout computation_layout( in TPUExecute() local
456 computation_layout.add_parameter_layout(xla::ShapeLayout(shape)); in TPUExecute()
461 xla::HloModuleConfig(std::move(computation_layout))); in TPUExecute()
|
| D | BUILD | 256 "//tensorflow/compiler/xla/service:computation_layout",
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | tfrt_cpu_pjrt_client.cc | 1218 const auto& computation_layout = in TfrtCpuExecutable() local
1220 if (computation_layout.parameter_count() == 0) { in TfrtCpuExecutable()
1225 if (computation_layout.parameter_count() > 1 || in TfrtCpuExecutable()
1226 !computation_layout.parameter_shape(0).IsTuple()) { in TfrtCpuExecutable()
1227 input_buffer_sizes_in_bytes_.reserve(computation_layout.parameter_count()); in TfrtCpuExecutable()
1228 for (int i = 0; i < computation_layout.parameter_count(); ++i) { in TfrtCpuExecutable()
1230 ShapeUtil::ByteSizeOf(computation_layout.parameter_shape(i))); in TfrtCpuExecutable()
1234 computation_layout.parameter_shape(0).tuple_shapes_size()); in TfrtCpuExecutable()
1236 i < computation_layout.parameter_shape(0).tuple_shapes_size(); ++i) { in TfrtCpuExecutable()
1238 computation_layout.parameter_shape(0).tuple_shapes(i))); in TfrtCpuExecutable()
|
| D | pjrt_stream_executor_client.cc | 1732 const auto& computation_layout = in PjRtStreamExecutorExecutable() local
1735 parameter_shapes.reserve(computation_layout.parameter_count()); in PjRtStreamExecutorExecutable()
1736 for (int i = 0; i < computation_layout.parameter_count(); ++i) { in PjRtStreamExecutorExecutable()
1738 computation_layout.parameter_shape(i))); in PjRtStreamExecutorExecutable()
|
| D | BUILD | 216 "//tensorflow/compiler/xla/service:computation_layout",
|
| /external/tensorflow/tensorflow/compiler/xla/tests/ |
| D | BUILD | 166 "//tensorflow/compiler/xla/service:computation_layout",
|