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 #ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_CONV_OP_HELPERS_H_ 17 #define TENSORFLOW_COMPILER_TF2XLA_KERNELS_CONV_OP_HELPERS_H_ 18 19 #include <vector> 20 21 #include "tensorflow/compiler/xla/client/xla_builder.h" 22 #include "tensorflow/compiler/xla/statusor.h" 23 #include "tensorflow/core/framework/op_kernel.h" 24 #include "tensorflow/core/framework/types.h" 25 #include "tensorflow/core/util/padding.h" 26 #include "tensorflow/core/util/tensor_format.h" 27 28 // This header exposes utilities for translating TensorFlow convolution ops into 29 // XLA ops. 30 // 31 // conv_ops.cc contains lowerings for many of these TF convolution ops (e.g. 32 // Conv2D, Conv3DBackpropFilterV2), but you might want to use the utilities in 33 // this header to implement a new and exciting convolution op, for example a 34 // fused TensorFlow op that contains a convolution and other things. 35 36 namespace tensorflow { 37 38 // We don't support integers for convolutions, so we list the supported types 39 // here. 40 std::vector<DataType> GetXlaConvTypes(); 41 42 // ConvOpAttrs contains all of the metadata necessary to specify a TF or XLA 43 // convolution. 44 struct ConvOpAttrs { 45 // Constructs a ConvOpAttrs, reading most of the attributes from `ctx`. 46 static xla::StatusOr<ConvOpAttrs> Create(int num_spatial_dims, bool depthwise, 47 OpKernelConstruction* ctx); 48 49 bool depthwise; 50 int num_spatial_dims; 51 std::vector<int32> dilations; 52 std::vector<int32> strides; 53 Padding padding; 54 std::vector<int64> explicit_paddings; 55 TensorFormat data_format; 56 }; 57 58 // Creates a new XLA forward or backward convolution with the given inputs and 59 // attributes. 60 xla::StatusOr<xla::XlaOp> MakeXlaForwardConvOp( 61 StringPiece type_string, xla::XlaOp conv_input, xla::XlaOp filter, 62 const ConvOpAttrs& attrs, 63 const xla::PrecisionConfig* precision_config = nullptr); 64 xla::StatusOr<xla::XlaOp> MakeXlaBackpropInputConvOp( 65 StringPiece type_string, const xla::Shape& input_shape, xla::XlaOp filter, 66 xla::XlaOp out_backprop, const ConvOpAttrs& attrs, 67 const xla::PrecisionConfig* precision_config = nullptr, 68 xla::XlaOp* input_sizes = nullptr); 69 xla::StatusOr<xla::XlaOp> MakeXlaBackpropFilterConvOp( 70 StringPiece type_string, xla::XlaOp activations, 71 const xla::Shape& filter_shape, xla::XlaOp gradients, 72 const ConvOpAttrs& attrs, 73 const xla::PrecisionConfig* precision_config = nullptr); 74 75 } // namespace tensorflow 76 77 #endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CONV_OP_HELPERS_H_ 78