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/tf2xla/xla_compiler.h" 24 #include "tensorflow/compiler/xla/client/local_client.h" 25 #include "tensorflow/core/framework/tensor.h" 26 #include "tensorflow/core/framework/types.h" 27 #include "tensorflow/core/lib/core/status.h" 28 29 namespace tensorflow { 30 31 // An OpKernel that compiles an op to an XLA computation and runs it. Unlike 32 // XlaLaunch this doesn't rely on any rewrites of the graphdef - it will run a 33 // vanilla TensorFlow op as long as the bridge supports it. 34 class XlaCompileOnDemandOp : public OpKernel { 35 public: XlaCompileOnDemandOp(OpKernelConstruction * ctx)36 explicit XlaCompileOnDemandOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} 37 void Compute(OpKernelContext* ctx) override; 38 39 private: 40 XlaCompiler::Argument CreateCompilerArgument(OpKernelContext* ctx, int64 i); 41 Status ShouldArgumentBeConstant(const OpKernel* op_kernel, int64 argument_idx, 42 bool* result); 43 Status MustArgumentBeConstant(const OpKernel* op_kernel, int64 argument_idx, 44 bool* result); 45 Status Compile(OpKernelContext* ctx, const XlaDevice::Metadata& metadata, 46 const XlaCompiler::CompilationResult** result, 47 xla::LocalExecutable** executable); 48 Status Run(OpKernelContext* ctx, const XlaDevice::Metadata& metadata, 49 const XlaCompiler::CompilationResult* result, 50 xla::LocalExecutable* executable); 51 }; 52 53 } // namespace tensorflow 54 55 #endif // TENSORFLOW_COMPILER_JIT_XLA_COMPILE_ON_DEMAND_OP_H_ 56