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 // ConvOpAttrs contains all of the metadata necessary to specify a TF or XLA
39 // convolution.
40 struct ConvOpAttrs {
41   // Constructs a ConvOpAttrs, reading most of the attributes from `ctx`.
42   static xla::StatusOr<ConvOpAttrs> Create(int num_spatial_dims, bool depthwise,
43                                            OpKernelConstruction* ctx);
44 
45   bool depthwise;
46   int num_spatial_dims;
47   std::vector<int32> dilations;
48   std::vector<int32> strides;
49   Padding padding;
50   std::vector<int64> explicit_paddings;
51   TensorFormat data_format;
52 };
53 
54 // Creates a new XLA forward or backward convolution with the given inputs and
55 // attributes.
56 xla::StatusOr<xla::XlaOp> MakeXlaForwardConvOp(StringPiece type_string,
57                                                xla::XlaOp conv_input,
58                                                xla::XlaOp filter,
59                                                const ConvOpAttrs& attrs);
60 xla::StatusOr<xla::XlaOp> MakeXlaBackpropInputConvOp(
61     StringPiece type_string, const xla::Shape& input_shape, xla::XlaOp filter,
62     xla::XlaOp out_backprop, const ConvOpAttrs& attrs);
63 xla::StatusOr<xla::XlaOp> MakeXlaBackpropFilterConvOp(
64     StringPiece type_string, xla::XlaOp activations,
65     const xla::Shape& filter_shape, xla::XlaOp gradients,
66     const ConvOpAttrs& attrs);
67 
68 }  // namespace tensorflow
69 
70 #endif  // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CONV_OP_HELPERS_H_
71