1 /* Copyright 2021 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 #ifndef TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_ 16 #define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_ 17 #if GOOGLE_CUDA && GOOGLE_TENSORRT 18 19 #include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h" 20 #include "tensorflow/core/lib/core/status.h" 21 #include "tensorflow/core/util/strided_slice_op.h" 22 23 namespace tensorflow { 24 namespace tensorrt { 25 namespace convert { 26 using SliceDims = absl::InlinedVector<int64, 4>; 27 28 // Creates a strided slice operation using the given information. This function 29 // expects that the begin, stride, and end vectors have already been validated. 30 // This function converts the [begin:stride:end] specification to the TensorRT 31 // [begin:stride:size] ISliceLayer specification. The following algorithm is 32 // used to perform this conversion: 1) The given (input_dims, 33 // [begin:stride:end]) specification is dividied into 34 // "static dimensions" and "dynamic dimensions". "Dynamic dimensions" 35 // includes all dimensions of the slice where input_dims[i] == -1. 36 // 2a) If there are no dynamic dimensions, then the "begin", "stride", and 37 // "size" variables are passed to the ISLiceLayer creation as build-time 38 // constants in the form of nvinfer1::Dims objects. 39 // 2b) If there are any dynamic dimensions, then the "begin", "stride", and 40 // "size" variables are treated as runtime dynamic shape Tensors in the 41 // TensorRT graph. In this case, we must calculate "size" at runtime for all 42 // dynamic dimensions, while static dimensions use the constant values. 43 // 44 // Note that when any dynamic indices are present (2b), the "strided_slice_spec" 45 // must be specified. This structure can be obtained through the 46 // "tensorflow::ValidateStridedSliceOp" function, or it can be constructed 47 // directly. When the ValidateStridedSliceOp helper function is used, it will 48 // also return the "begin", "stride", and "end" vectors. When all dimensions are 49 // static (2a), the "strided_slice_spec" variable is not required. 50 // 51 // If the "final_shape" variable is specified, then a reshape operation will be 52 // added to the graph to achieve this shape. The shape must be fully specified. 53 // 54 // "op_instance" is only required if the caller needs to pass this variable 55 // through to the Converter functions optionally accept it (SetLayerName, 56 // PrepareTensorForShape). 57 Status ConvertStridedSliceHelper( 58 OpConverterParams* params, const TRT_TensorOrWeights& input, 59 const PartialTensorShape& input_dims, const SliceDims& begin, 60 const SliceDims& stride, const SliceDims& end, 61 std::optional<nvinfer1::Dims> final_shape = std::nullopt, 62 std::optional<int> op_instance = std::nullopt, 63 std::optional<StridedSliceShapeSpec> strided_slice_spec = std::nullopt); 64 65 } // namespace convert 66 } // namespace tensorrt 67 } // namespace tensorflow 68 69 #endif // GOOGLE_CUDA && GOOGLE_TENSORRT 70 #endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_OPS_SLICE_OPS_H_ 71