1 /* Copyright 2022 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 #ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_GPU_TF_KERNEL_CUSTOM_CALL_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_KERNELS_GPU_TF_KERNEL_CUSTOM_CALL_H_ 18 19 #include <functional> 20 21 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" 22 #include "tensorflow/core/platform/status.h" 23 namespace tensorflow { 24 25 // Using std::map as the maps are presumed to be tiny, and we want a 26 // deterministic iteration order. 27 // 28 // Dimension -> bound. 29 using DimensionBoundsMap = std::map<int, int>; 30 31 // Output -> dimension -> bound. 32 using OutputDimensionBoundsMap = std::map<int, DimensionBoundsMap>; 33 34 // Generic kernel for registering TF2XLA kernels which call back into the TF 35 // runtime to run a given kernel defined by the wrapped node. 36 // 37 // Cf. example usages in light_outside_compilation_kernels_for_test.cc. 38 // 39 // Currently does not support dynamic shape or resource variables. Currently 40 // works only on GPU. 41 class LightOutsideCompilationOp : public XlaOpKernel { 42 public: 43 explicit LightOutsideCompilationOp(OpKernelConstruction* context); 44 void Compile(XlaOpKernelContext* ctx) override; 45 46 // Override to provide statically known bounds on output in case of dynamic 47 // shapes. DynamicOutputDimensions(const NodeDef & ndef,XlaOpKernelContext * ctx)48 virtual StatusOr<OutputDimensionBoundsMap> DynamicOutputDimensions( 49 const NodeDef& ndef, XlaOpKernelContext* ctx) const { 50 return OutputDimensionBoundsMap{}; 51 } 52 53 private: 54 Status CompileToCustomCallCallingTfKernel(int graph_def_version, 55 const NodeDef& node_def, 56 XlaOpKernelContext* ctx) const; 57 58 NodeDef def_; 59 int graph_def_version_; 60 }; 61 62 } // namespace tensorflow 63 64 #endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_GPU_TF_KERNEL_CUSTOM_CALL_H_ 65