Home
last modified time | relevance | path

Searched refs:program_shape (Results 1 – 25 of 76) sorted by relevance

1234

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/
H A Dshape_test.cc146 ProgramShape program_shape; in TEST_F() local
147 *program_shape.add_parameters() = ShapeUtil::MakeShape(F32, {1, 2, 3}); in TEST_F()
148 *program_shape.add_parameters() = ShapeUtil::MakeTokenShape(); in TEST_F()
149 *program_shape.add_parameters() = ShapeUtil::MakeShape(S64, {}); in TEST_F()
150 *program_shape.add_parameters() = ShapeUtil::MakeTupleShape( in TEST_F()
155 *program_shape.mutable_result() = ShapeUtil::MakeShape(F32, {7}); in TEST_F()
157 program_shape.add_parameter_names("foo"); in TEST_F()
158 program_shape.add_parameter_names("bar"); in TEST_F()
159 program_shape.add_parameter_names("baz"); in TEST_F()
160 program_shape.add_parameter_names("qux qux"); in TEST_F()
[all …]
H A Dlayout_util.cc168 /* static */ void LayoutUtil::SetToDefaultLayout(ProgramShape* program_shape) { in SetToDefaultLayout() argument
169 for (auto& parameter_shape : *program_shape->mutable_parameters()) { in SetToDefaultLayout()
172 LayoutUtil::SetToDefaultLayout(program_shape->mutable_result()); in SetToDefaultLayout()
278 /* static */ void LayoutUtil::ClearLayout(ProgramShape* program_shape) { in ClearLayout() argument
279 for (auto& parameter_shape : *program_shape->mutable_parameters()) { in ClearLayout()
282 LayoutUtil::ClearLayout(program_shape->mutable_result()); in ClearLayout()
371 /* static */ bool LayoutUtil::HasLayout(const ProgramShape& program_shape) { in HasLayout() argument
372 for (auto& parameter_shape : program_shape.parameters()) { in HasLayout()
377 return LayoutUtil::HasLayout(program_shape.result()); in HasLayout()
H A Dlayout_util.h74 static void SetToDefaultLayout(ProgramShape* program_shape);
92 static void ClearLayout(ProgramShape* program_shape);
152 static bool HasLayout(const ProgramShape& program_shape);
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dcomputation_layout.cc27 ComputationLayout::ComputationLayout(const ProgramShape& program_shape, in ComputationLayout() argument
29 : result_layout_(program_shape.result()) { in ComputationLayout()
30 for (auto& shape : program_shape.parameters()) { in ComputationLayout()
75 ProgramShape program_shape; in ComputeProgramShape() local
77 *program_shape.add_parameters() = parameter_layouts_[i].shape(); in ComputeProgramShape()
78 *program_shape.add_parameter_names() = absl::StrCat("p", i); in ComputeProgramShape()
80 *program_shape.mutable_result() = result_layout_.shape(); in ComputeProgramShape()
81 return program_shape; in ComputeProgramShape()
H A Dhlo_module_config.cc31 HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape, in HloModuleConfig() argument
34 ComputationLayout(program_shape, ignore_layouts)) {} in HloModuleConfig()
40 const ProgramShape& program_shape) { in SetDefaultComputationLayout() argument
41 entry_computation_layout_ = ComputationLayout(program_shape); in SetDefaultComputationLayout()
45 const ProgramShape& program_shape) { in SetComputationLayoutIfExists() argument
46 entry_computation_layout_ = ComputationLayout(program_shape, in SetComputationLayoutIfExists()
H A Dlocal_service.cc105 ProgramShape program_shape(proto.host_program_shape()); in GetHloModuleConfig() local
108 if (argument_layouts.size() != program_shape.parameters_size()) { in GetHloModuleConfig()
111 program_shape.parameters_size(), argument_layouts.size()); in GetHloModuleConfig()
118 if (!ShapeUtil::Compatible(argument_shape, program_shape.parameters(i))) { in GetHloModuleConfig()
135 ShapeUtil::HumanString(program_shape.parameters(i)), in GetHloModuleConfig()
141 program_shape.result())); in GetHloModuleConfig()
145 CreateExecutionOptions(build_options, &program_shape); in GetHloModuleConfig()
147 return CreateModuleConfig(program_shape, argument_layouts, in GetHloModuleConfig()
H A Dhlo_module_util.cc45 const ProgramShape& program_shape, in CreateModuleConfig() argument
49 auto config = std::make_unique<HloModuleConfig>(program_shape); in CreateModuleConfig()
53 if (program_shape.parameters_size() != argument_shapes_size) { in CreateModuleConfig()
55 program_shape.parameters_size(), in CreateModuleConfig()
62 program_shape.parameters(i))) { in CreateModuleConfig()
66 i, ShapeUtil::HumanString(program_shape.parameters(i)), in CreateModuleConfig()
78 ValidateResultShape(shape_with_output_layout, program_shape.result())); in CreateModuleConfig()
H A Dhlo_module_config.h78 explicit HloModuleConfig(const ProgramShape& program_shape,
91 void SetDefaultComputationLayout(const ProgramShape& program_shape);
95 void SetComputationLayoutIfExists(const ProgramShape& program_shape);
H A Dhlo_proto_util.cc66 const auto& program_shape = hlo_proto.hlo_module().host_program_shape(); in EntryComputationParameterShapes() local
67 for (const ShapeProto& shape : program_shape.parameters()) { in EntryComputationParameterShapes()
H A Dservice.cc261 const ProgramShape& program_shape, in CreateModuleConfig() argument
272 return xla::CreateModuleConfig(program_shape, argument_shapes, in CreateModuleConfig()
278 const ProgramShape& program_shape, in CreateModuleConfig() argument
286 return CreateModuleConfig(program_shape, argument_shapes, &execution_options, in CreateModuleConfig()
1131 ProgramShape program_shape(arg->computation().host_program_shape()); in ComputeConstantGraph() local
1132 TF_DCHECK_OK(ShapeUtil::ValidateShape(program_shape.result())); in ComputeConstantGraph()
1137 *output_layout, program_shape.result())); in ComputeConstantGraph()
1140 HloModuleConfig config(program_shape); in ComputeConstantGraph()
H A Dhlo_module.cc525 const ProgramShape& program_shape, const DebugOptions& debug_options, in CreateModuleConfigFromShape() argument
527 HloModuleConfig module_config(ProgramShape{program_shape}); in CreateModuleConfigFromShape()
576 program_shape.parameters(i))); in CreateModuleConfigFromShape()
579 program_shape.result())); in CreateModuleConfigFromShape()
591 ProgramShape program_shape(module.host_program_shape()); in CreateModuleConfigFromProto() local
593 CreateModuleConfigFromShape(program_shape, debug_options, in CreateModuleConfigFromProto()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dutils.cc77 const XlaComputation& computation, const ProgramShape& program_shape) { in GetShardedProgramShapes() argument
79 arg_shapes.resize(program_shape.parameters_size()); in GetShardedProgramShapes()
87 if (instr.parameter_number() >= program_shape.parameters_size()) { in GetShardedProgramShapes()
90 instr.parameter_number(), program_shape.parameters_size()); in GetShardedProgramShapes()
155 TF_ASSIGN_OR_RETURN(ProgramShape program_shape, in DetermineArgumentLayoutsFromCompileOptions()
158 argument_layouts.emplace(program_shape.parameters()); in DetermineArgumentLayoutsFromCompileOptions()
162 } else if (argument_layouts->size() != program_shape.parameters_size()) { in DetermineArgumentLayoutsFromCompileOptions()
166 argument_layouts->size(), program_shape.parameters_size()); in DetermineArgumentLayoutsFromCompileOptions()
190 GetShardedProgramShapes(computation, program_shape)); in DetermineArgumentLayoutsFromCompileOptions()
203 result_layout = program_shape.result(); in DetermineArgumentLayoutsFromCompileOptions()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dxla_jit_compiled_cpu_function_test.cc220 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST() local
221 ASSERT_EQ(program_shape.parameters_size(), 2); in TEST()
222 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(0), s32)); in TEST()
223 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(1), s32)); in TEST()
225 const xla::Shape& result = program_shape.result(); in TEST()
271 const xla::ProgramShape program_shape(*function.ProgramShape()); in TEST() local
272 ASSERT_EQ(program_shape.parameters_size(), 2); in TEST()
273 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(0), s32)); in TEST()
274 EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(1), s32_1)); in TEST()
276 const xla::Shape& result = program_shape.result(); in TEST()
H A Dxla_jit_compiled_cpu_function.cc97 TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::ProgramShape> program_shape, in Compile()
99 if (program_shape->result().element_type() != xla::TUPLE) { in Compile()
108 program_shape->clear_parameter_names(); in Compile()
112 arg_shapes.reserve(program_shape->parameters_size()); in Compile()
113 for (int i = 0; i < program_shape->parameters_size(); ++i) { in Compile()
114 arg_shapes.push_back(&program_shape->parameters(i)); in Compile()
146 std::make_unique<xla::ProgramShapeProto>(program_shape->ToProto()); in Compile()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/kernels/
H A Dxrt_compile_ops.cc154 config.program_shape().parameters_size()); in Compile()
156 config.program_shape().parameters_size()); in Compile()
157 for (int i = 0; i < config.program_shape().parameters_size(); ++i) { in Compile()
158 argument_layouts[i] = xla::Shape(config.program_shape().parameters(i)); in Compile()
164 build_options.set_result_layout(xla::Shape(config.program_shape().result())); in Compile()
230 xla::ProgramShapeProto program_shape = executable->executable() in Compute() local
237 program_shape_output.vec<tstring>()(0) = program_shape.SerializeAsString(); in Compute()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tools/
H A Ddumped_computation_to_operation_list.cc79 std::unique_ptr<ProgramShape> program_shape = in RealMain() local
83 layouts.reserve(program_shape->parameters_size()); in RealMain()
84 for (int i = 0; i < program_shape->parameters_size(); ++i) { in RealMain()
85 layouts.push_back(&program_shape->parameters(i)); in RealMain()
89 build_options.set_result_layout(program_shape->result()); in RealMain()
H A Ddumped_computation_to_text.cc55 std::unique_ptr<ProgramShape> program_shape = in RealMain() local
59 layouts.reserve(program_shape->parameters_size()); in RealMain()
60 for (int i = 0; i < program_shape->parameters_size(); ++i) { in RealMain()
61 layouts.push_back(&program_shape->parameters(i)); in RealMain()
66 build_options.set_result_layout(program_shape->result()); in RealMain()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_compile_op_support.cc59 const ProgramShape& program_shape, absl::Span<const Shape> argument_shapes, in CreateModuleConfig() argument
66 auto config = absl::make_unique<HloModuleConfig>(program_shape); in CreateModuleConfig()
69 if (program_shape.parameters_size() != argument_shapes.size()) { in CreateModuleConfig()
71 program_shape.parameters_size(), in CreateModuleConfig()
78 program_shape.parameters(i))) { in CreateModuleConfig()
82 i, ShapeUtil::HumanString(program_shape.parameters(i)), in CreateModuleConfig()
92 ValidateResultShape(result_layout.value(), program_shape.result())); in CreateModuleConfig()
137 const xla::ProgramShape& program_shape, in CreateModuleConfig() argument
142 return CreateModuleConfig(program_shape, argument_shapes, result_layout, in CreateModuleConfig()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/tests/
H A Dcpu_topk_test.cc37 TF_ASSERT_OK_AND_ASSIGN(ProgramShape program_shape, in TEST_F()
39 HloModuleConfig config(program_shape); in TEST_F()
64 TF_ASSERT_OK_AND_ASSIGN(ProgramShape program_shape, in TEST_F()
66 HloModuleConfig config(program_shape); in TEST_F()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xrt/tests/
H A Draw_api_test.cc1109 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
1117 xla::ProgramShapeProto program_shape; in TEST() local
1118 EXPECT_TRUE(ParseFromTString(outputs[1].vec<tstring>()(0), &program_shape)); in TEST()
1119 EXPECT_EQ(program_shape.parameters_size(), 2); in TEST()
1166 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
1223 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
1278 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
1348 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
1399 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
1443 TF_EXPECT_OK(session.Run({read_back, c_handle.program_shape}, &outputs)); in TEST()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/aot/
H A Dcodegen.cc417 const xla::ProgramShapeProto& ps = compile_result.program_shape; in GenerateHeader()
748 std::unique_ptr<xla::ProgramShapeProto> program_shape; in GenerateMetadata() local
751 program_shape = in GenerateMetadata()
752 absl::make_unique<xla::ProgramShapeProto>(compile_result.program_shape); in GenerateMetadata()
757 program_shape->clear_parameter_names(); in GenerateMetadata()
765 "::xla::ProgramShapeProto", program_shape.get()}; in GenerateMetadata()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/
H A Dhlo_test_base.cc64 ProgramShape program_shape; in GetProgramShapeWithLayout() local
67 *program_shape.add_parameters() = param->shape(); in GetProgramShapeWithLayout()
68 *program_shape.add_parameter_names() = param->name(); in GetProgramShapeWithLayout()
70 *program_shape.mutable_result() = entry->root_instruction()->shape(); in GetProgramShapeWithLayout()
71 return program_shape; in GetProgramShapeWithLayout()
295 const auto& program_shape = GetProgramShapeWithLayout(test_module); in MakeReferenceModule() local
299 if (!ProgramShapesEqual(program_shape, in MakeReferenceModule()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/
H A Dpod_tpu_driver.cc197 xla::Status program_shape(xla::ProgramShapeProto* program_shape) override;
800 xla::ProgramShapeProto* program_shape) in GetCompiledProgramShape() argument
810 return underlying_cph_[op_id][0]->program_shape(program_shape); in GetCompiledProgramShape()
966 xla::Status PodCompiledProgramHandle::program_shape( in program_shape() function in tpu_driver::__anonf243d9070111::PodCompiledProgramHandle
967 xla::ProgramShapeProto* program_shape) { in program_shape() argument
968 return driver_->GetCompiledProgramShape(operation_id(), program_shape); in program_shape()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/
H A Dcompile_only_client.cc28 const ProgramShape& program_shape, in CreateModuleConfig() argument
31 return compiler_service_->CreateModuleConfig(program_shape, argument_shapes, in CreateModuleConfig()
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/
H A Dtpu_embedding_optimization_parameters_utils.cc168 const xla::ProgramShapeProto& program_shape = in GetBaseAuxiliaryParameterCount() local
171 const int num_inputs = program_shape.parameters_size(); in GetBaseAuxiliaryParameterCount()
172 const int num_outputs = program_shape.result().tuple_shapes_size(); in GetBaseAuxiliaryParameterCount()

1234