Home
last modified time | relevance | path

Searched refs:XlaCompiler (Results 1 – 25 of 45) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/tf2xla/
Dxla_compiler_test.cc71 XlaCompiler::Options DefaultOptions() { in DefaultOptions()
72 XlaCompiler::Options options; in DefaultOptions()
79 FunctionLibraryDefinition* LocalFlibDef(XlaCompiler* compiler) { in LocalFlibDef()
182 XlaCompiler compiler(DefaultOptions()); in TEST_F()
185 XlaCompiler::CompilationResult result; in TEST_F()
186 TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", in TEST_F()
205 std::vector<XlaCompiler::Argument> args(2); in TEST_F()
206 args[0].kind = XlaCompiler::Argument::kParameter; in TEST_F()
209 args[1].kind = XlaCompiler::Argument::kParameter; in TEST_F()
214 XlaCompiler compiler(DefaultOptions()); in TEST_F()
[all …]
Dxla_compiler.cc68 absl::Span<const XlaCompiler::Argument> args) { in CheckSignature()
165 const std::vector<XlaCompiler::Argument>& args, in BuildComputation()
171 const XlaCompiler::ShapeRepresentationFn& shape_representation_fn, in BuildComputation()
176 std::vector<XlaCompiler::OutputDescription>* outputs, in BuildComputation()
177 std::vector<XlaCompiler::ResourceUpdate>* resource_updates, in BuildComputation()
203 XlaCompiler::OutputDescription& output = (*outputs)[i]; in BuildComputation()
286 const XlaCompiler::Argument& arg = args[resource->arg_num()]; in BuildComputation()
300 XlaCompiler::ResourceUpdate& update = resource_updates->back(); in BuildComputation()
433 string XlaCompiler::Argument::HumanString() const { in HumanString()
474 std::vector<int64> XlaCompiler::Argument::DimensionSizes() const { in DimensionSizes()
[all …]
Dgraph_compiler.cc59 std::vector<XlaCompiler::Argument>* args) { in PrepareArguments()
69 XlaCompiler::Argument& arg = (*args)[i]; in PrepareArguments()
75 arg.kind = XlaCompiler::Argument::kConstant; in PrepareArguments()
88 arg.kind = XlaCompiler::Argument::kConstant; in PrepareArguments()
91 arg.kind = XlaCompiler::Argument::kParameter; in PrepareArguments()
96 XlaCompiler::PopulateArgumentFromResource(*resource, &arg); in PrepareArguments()
100 arg.kind = XlaCompiler::Argument::kTensorList; in PrepareArguments()
236 XlaCompiler* compiler = xla_op_context.compiler(); in CompileFunctionalNode()
250 std::vector<XlaCompiler::Argument> arguments; in CompileFunctionalNode()
262 XlaCompiler::CompileOptions compile_options; in CompileFunctionalNode()
[all …]
Dtf2xla.cc65 std::vector<XlaCompiler::Argument> xla_args; in ConvertGraphToXla()
70 XlaCompiler::Options compiler_options; in ConvertGraphToXla()
79 XlaCompiler compiler(compiler_options); in ConvertGraphToXla()
81 XlaCompiler::CompilationResult result; in ConvertGraphToXla()
83 XlaCompiler::CompileOptions options; in ConvertGraphToXla()
115 for (const XlaCompiler::ResourceUpdate& update : result.resource_updates) { in ConvertGraphToXla()
Dxla_context.h36 class XlaCompiler; variable
48 XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder,
54 XlaCompiler* compiler() const { return compiler_; } in compiler()
108 XlaCompiler* const compiler_;
Dxla_compiler.h100 class XlaCompiler {
205 explicit XlaCompiler(Options options);
207 ~XlaCompiler();
304 const std::vector<XlaCompiler::Argument>& args,
367 TF_DISALLOW_COPY_AND_ASSIGN(XlaCompiler);
/external/tensorflow/tensorflow/compiler/jit/
Dxla_compilation_cache.h72 Status Compile(const XlaCompiler::Options& options,
74 absl::Span<const XlaCompiler::Argument> args,
75 const XlaCompiler::CompileOptions& compile_options,
77 const XlaCompiler::CompilationResult** out_compilation_result,
85 const XlaCompiler::Options& options,
86 absl::Span<const XlaCompiler::Argument> args, OpKernelContext* ctx,
87 const XlaCompiler::CompileOptions& compile_options,
88 const XlaCompiler::CompilationResult** out_compilation_result,
123 absl::Span<const XlaCompiler::Argument> args);
128 const XlaCompiler::Options& options, const NameAttrList& function,
[all …]
Dxla_compilation_cache.cc133 absl::Span<const XlaCompiler::Argument> args) { in BuildSignature()
137 for (const XlaCompiler::Argument& arg : args) { in BuildSignature()
139 case XlaCompiler::Argument::kConstant: in BuildSignature()
140 case XlaCompiler::Argument::kConstantResource: in BuildSignature()
143 case XlaCompiler::Argument::kParameter: in BuildSignature()
144 case XlaCompiler::Argument::kResource: in BuildSignature()
158 const XlaCompiler::Options& options, in BuildExecutable()
159 const XlaCompiler::CompilationResult& result, in BuildExecutable()
186 const XlaCompiler::Options& options, const NameAttrList& function, in Compile()
187 absl::Span<const XlaCompiler::Argument> args, in Compile()
[all …]
Dget_compiler_ir.cc37 const XlaCompiler::Options& options, in GetLocalExecutable()
38 const XlaCompiler::CompileOptions& compile_options, in GetLocalExecutable()
40 absl::Span<XlaCompiler::Argument const> args, const XlaCompiler& compiler) { in GetLocalExecutable()
41 const XlaCompiler::CompilationResult* compilation_result = nullptr; in GetLocalExecutable()
105 XlaCompiler::Options options = in GetCompilerIr()
110 XlaCompiler::CompileOptions compile_options; in GetCompilerIr()
114 XlaCompiler compiler(options); in GetCompilerIr()
116 xla::StatusOr<std::vector<XlaCompiler::Argument>> args = in GetCompilerIr()
124 XlaCompiler::CompilationResult result; in GetCompilerIr()
Dxla_compilation_cache_test.cc30 std::vector<XlaCompiler::Argument> args(1); in TEST()
31 args[0].kind = XlaCompiler::Argument::kConstant; in TEST()
66 const XlaCompiler::CompilationResult* compilation_result; in TEST()
72 Status status = cache->Compile(XlaCompiler::Options{}, fn, {}, in TEST()
73 XlaCompiler::CompileOptions{}, in TEST()
89 std::vector<XlaCompiler::Argument> args(n_args); in BM_BuildSignature()
91 args[i].kind = (((i % 3) == 0) ? XlaCompiler::Argument::kConstant in BM_BuildSignature()
92 : XlaCompiler::Argument::kParameter); in BM_BuildSignature()
Dxla_compile_on_demand_op.h45 XlaCompiler::Argument CreateCompilerArgument(OpKernelContext* ctx, int64 i);
47 const XlaCompiler::CompilationResult** result,
53 const XlaCompiler::CompilationResult* result,
Dxla_compile_on_demand_op.cc46 const XlaCompiler::CompilationResult* result, in Run()
103 OpKernelContext* ctx, const XlaCompiler::CompilationResult** result, in Compile()
130 XlaCompiler::Options options = GenerateCompilerOptions( in Compile()
137 XlaCompiler::CompileOptions compile_options; in Compile()
144 xla::StatusOr<std::vector<XlaCompiler::Argument>> args; in Compile()
166 const XlaCompiler::CompilationResult* result; in Compute()
Dxla_launch_util.cc222 const XlaCompiler::CompilationResult* compilation_result, in PopulateInputs()
243 [&](const XlaCompiler::ResourceUpdate& update) { in PopulateInputs()
345 const XlaCompiler::CompilationResult* compilation_result, int output_num) { in SetOutputForConstant()
388 const XlaCompiler::ResourceUpdate& write) { in GetOrCreateResourceVar()
400 const XlaCompiler::CompilationResult& compilation_result, in GatherVariableInfo()
405 const XlaCompiler::ResourceUpdate& write = in GatherVariableInfo()
422 const XlaCompiler::CompilationResult* compilation_result, in PopulateOutputs()
538 const XlaCompiler::ResourceUpdate& write = in PopulateOutputs()
578 xla::StatusOr<std::vector<XlaCompiler::Argument>>
584 std::vector<XlaCompiler::Argument> out; in BuildXlaCompilerArguments()
[all …]
Dxla_device.h63 XlaCompiler::ShapeRepresentationFn shape_representation_fn,
72 const XlaCompiler::ShapeRepresentationFn& shape_representation_fn() const { in shape_representation_fn()
83 XlaCompiler::ShapeRepresentationFn shape_representation_fn_;
128 XlaCompiler::ShapeRepresentationFn shape_representation_fn;
245 const XlaCompiler::ShapeRepresentationFn shape_representation_fn_;
Dxla_launch_util.h87 const XlaCompiler::CompilationResult& compilation_result,
143 static xla::StatusOr<std::vector<XlaCompiler::Argument>>
158 const XlaCompiler::CompilationResult* compilation_result,
175 const XlaCompiler::CompilationResult* compilation_result,
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dif_op.cc54 absl::Span<XlaCompiler::Argument> arguments, in PopulateTensorArrayGradients()
55 XlaCompiler::CompilationResult* then_result, in PopulateTensorArrayGradients()
56 XlaCompiler::CompilationResult* else_result) { in PopulateTensorArrayGradients()
58 for (XlaCompiler::CompilationResult* result : {then_result, else_result}) { in PopulateTensorArrayGradients()
59 for (const XlaCompiler::ResourceUpdate& update : result->resource_updates) { in PopulateTensorArrayGradients()
63 XlaCompiler::Argument& arg = arguments[update.input_index]; in PopulateTensorArrayGradients()
88 XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& then_result, in ValidateShapes()
89 const XlaCompiler::CompilationResult& else_result) { in ValidateShapes()
180 std::vector<XlaCompiler::Argument> arguments(input_types_.size()); in Compile()
183 XlaCompiler::Argument& arg = arguments[i]; in Compile()
[all …]
Dxla_select_and_scatter_op.cc61 XlaCompiler::CompileOptions compile_options; in Compile()
67 XlaCompiler::Argument select_arg; in Compile()
68 select_arg.kind = XlaCompiler::Argument::kParameter; in Compile()
72 XlaCompiler::CompilationResult select; in Compile()
88 XlaCompiler::Argument scatter_arg; in Compile()
89 scatter_arg.kind = XlaCompiler::Argument::kParameter; in Compile()
93 XlaCompiler::CompilationResult scatter; in Compile()
Dcase_op.cc90 std::vector<XlaCompiler::Argument> arguments(input_types_.size()); in Compile()
93 XlaCompiler::Argument& arg = arguments[i]; in Compile()
99 XlaCompiler::PopulateArgumentFromResource(*resource, &arg); in Compile()
109 arg.kind = XlaCompiler::Argument::kParameter; in Compile()
137 if (arguments[arg_idx].kind != XlaCompiler::Argument::kParameter) { in Compile()
155 XlaCompiler::CompileOptions options; in Compile()
160 XlaCompiler* compiler = ctx->compiler(); in Compile()
162 std::vector<XlaCompiler::CompilationResult> branch_results(num_branches); in Compile()
170 for (XlaCompiler::CompilationResult& result : branch_results) { in Compile()
171 for (const XlaCompiler::ResourceUpdate& update : result.resource_updates) { in Compile()
[all …]
Dwhile_op.cc67 XlaOpKernelContext* ctx, std::vector<XlaCompiler::Argument>* args, in MakeXlaCompilerArgumentsFromInputs()
78 XlaCompiler::Argument& arg = (*args)[i]; in MakeXlaCompilerArgumentsFromInputs()
85 XlaCompiler::PopulateArgumentFromResource(*resource, &arg); in MakeXlaCompilerArgumentsFromInputs()
97 arg.kind = XlaCompiler::Argument::kParameter; in MakeXlaCompilerArgumentsFromInputs()
137 std::vector<XlaCompiler::Argument>* args, in ConvertLoopInvariantsToConst()
153 XlaCompiler::Argument& arg = (*args)[arg_idx]; in ConvertLoopInvariantsToConst()
154 return arg.kind != XlaCompiler::Argument::kResource && in ConvertLoopInvariantsToConst()
173 const XlaCompiler::CompilationResult& body, bool has_token_input_output) { in VerifyBodyInputAndOutputShapeMatch()
198 XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& cond) { in BuildWrappedCond()
209 XlaOpKernelContext* ctx, const XlaCompiler::CompilationResult& body, in BuildWrappedBody()
[all …]
Dgather_scatter_ops.cc72 XlaCompiler::Argument update_computation_arg; in Compile()
73 update_computation_arg.kind = XlaCompiler::Argument::kParameter; in Compile()
77 XlaCompiler::CompileOptions compile_options; in Compile()
81 XlaCompiler::CompilationResult update_computation; in Compile()
Dsort_ops.cc77 std::vector<XlaCompiler::Argument> comparator_args(2 * input_types_.size()); in Compile()
82 XlaCompiler::Argument comparator_arg; in Compile()
83 comparator_arg.kind = XlaCompiler::Argument::kParameter; in Compile()
91 XlaCompiler::CompilationResult comparator; in Compile()
92 XlaCompiler::CompileOptions compile_options; in Compile()
Dxla_reduce_op.cc73 XlaCompiler::Argument reducer_arg; in Compile()
74 reducer_arg.kind = XlaCompiler::Argument::kParameter; in Compile()
78 XlaCompiler::CompileOptions compile_options; in Compile()
82 XlaCompiler::CompilationResult reducer; in Compile()
87 std::vector<XlaCompiler::Argument>(n_ * 2, reducer_arg), &reducer)); in Compile()
Dreduce_window_op.cc77 XlaCompiler::Argument reducer_arg; in Compile()
78 reducer_arg.kind = XlaCompiler::Argument::kParameter; in Compile()
82 XlaCompiler::CompileOptions compile_options; in Compile()
86 XlaCompiler::CompilationResult reducer; in Compile()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_common.cc181 const GuaranteedConsts& guaranteed_constants, const XlaCompiler& compiler, in BuildComputationArgumentDescriptions()
182 std::vector<XlaCompiler::Argument>* args, in BuildComputationArgumentDescriptions()
190 args->push_back(XlaCompiler::Argument()); in BuildComputationArgumentDescriptions()
191 XlaCompiler::Argument& arg = args->back(); in BuildComputationArgumentDescriptions()
197 arg.kind = XlaCompiler::Argument::kParameter; in BuildComputationArgumentDescriptions()
200 arg.kind = XlaCompiler::Argument::kResource; in BuildComputationArgumentDescriptions()
206 arg.kind = XlaCompiler::Argument::kConstant; in BuildComputationArgumentDescriptions()
233 if (arg.kind == XlaCompiler::Argument::kInvalid) { in BuildComputationArgumentDescriptions()
236 if (arg.kind == XlaCompiler::Argument::kConstant) { in BuildComputationArgumentDescriptions()
257 const XlaCompiler::ShapeRepresentationFn shape_representation_fn, in GetShardingInfo()
[all …]
Dtpu_compile_op_common.h174 const XlaCompiler::ShapeRepresentationFn shape_representation_fn,
182 XlaCompiler::CompilationResult* compilation_result);
188 const XlaCompiler::ShapeRepresentationFn shape_representation_fn,
200 const GuaranteedConsts& guaranteed_constants, const XlaCompiler& compiler,
201 std::vector<XlaCompiler::Argument>* args,

12