1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // The XlaCompileOnDemandOp is an OpKernel that, when its Compute method is 17 // called, will generate an xla::Computation and run it asynchronously. 18 19 #ifndef TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_ 20 #define TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_ 21 22 #include "tensorflow/compiler/jit/xla_device.h" 23 #include "tensorflow/compiler/jit/xla_launch_util.h" 24 #include "tensorflow/compiler/jit/xla_platform_info.h" 25 #include "tensorflow/compiler/tf2xla/xla_compiler.h" 26 #include "tensorflow/compiler/xla/client/local_client.h" 27 #include "tensorflow/core/framework/function.h" 28 #include "tensorflow/core/framework/tensor.h" 29 #include "tensorflow/core/framework/types.h" 30 #include "tensorflow/core/lib/core/status.h" 31 32 namespace tensorflow { 33 34 // An OpKernel that compiles an op to an XLA computation and runs it. Unlike 35 // XlaLaunch this doesn't rely on any rewrites of the graphdef - it will run a 36 // vanilla TensorFlow op as long as the bridge supports it. 37 class XlaCompileOnDemandOp : public OpKernel { 38 public: XlaCompileOnDemandOp(OpKernelConstruction * ctx)39 explicit XlaCompileOnDemandOp(OpKernelConstruction* ctx) 40 : OpKernel(ctx), 41 platform_info_(XlaPlatformInfoFromDevice(ctx->device())) {} 42 void Compute(OpKernelContext* ctx) override; 43 44 private: 45 XlaCompiler::Argument CreateCompilerArgument(OpKernelContext* ctx, int64 i); 46 Status Compile(OpKernelContext* ctx, 47 const XlaCompiler::CompilationResult** result, 48 XlaCompilationCache** cache, 49 ResourceVarsSnapshot* variable_args, 50 xla::LocalExecutable** executable); 51 52 Status Run(OpKernelContext* ctx, XlaCompilationCache* cache, 53 const XlaCompiler::CompilationResult* result, 54 xla::LocalExecutable* executable, 55 const ResourceVarsSnapshot& variable_args); 56 57 const XlaPlatformInfo platform_info_; 58 }; 59 60 } // namespace tensorflow 61 62 #endif // TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_ 63