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 16 #ifndef TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_PARAMETERS_H_ 17 #define TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_PARAMETERS_H_ 18 19 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM 20 #include "absl/types/optional.h" 21 #include "tensorflow/core/platform/stream_executor.h" 22 #include "tensorflow/core/util/autotune_maps/conv_parameters.pb.h" 23 24 namespace tensorflow { 25 // Uniquely identifies a convolution operation that runs on a particular device 26 // model. 27 // 28 // This can serve as a hashtable key, where the value might be the autotuned 29 // algorithm we choose for the conv. 30 // 31 // All of the data in this class other than the device_id is stored in the 32 // ConvParametersProto, so it can be easily serialized (for the purposes of 33 // ahead-of-time autotuning). 34 // 35 // When using the cudnn frontend API, two autotuning results for two different 36 // GPUs of the same model are not interchangeable, because an autotuning result 37 // includes a cudnn execution plan, which is tied to the GPU. As a result, we 38 // need to create separate ConvParameters objects for them. 39 class ConvParameters { 40 public: 41 struct FusionInfo { 42 bool has_side_input = false; 43 stream_executor::dnn::ActivationMode activation_mode; 44 bool is_contrib; 45 }; 46 47 // We have three kinds of convolutions today. Vanilla unfused convolutions, 48 // fused convolutions, and fused convolutions as implemented in the `contrib` 49 // directory. The two fused convolutions ultimately correspond to the same 50 // cudnn calls, but have slightly different semantics (e.g. they interpret 51 // padding differently). 52 ConvParameters( 53 int64_t batch, int64_t in_depths, absl::Span<const int64_t> in, 54 int data_format, int64_t out_depths, absl::Span<const int64_t> filter, 55 absl::Span<const int64_t> dilation, absl::Span<const int64_t> stride, 56 absl::Span<const int64_t> padding, DataType dtype, int device_id, 57 int group_count, 58 absl::optional<FusionInfo> fusion_info = absl::optional<FusionInfo>()); 59 60 ConvParameters(int device_id, const ConvParametersProto& proto); 61 62 bool operator==(const ConvParameters& other) const; 63 64 bool operator!=(const ConvParameters& other) const { 65 return !(*this == other); 66 } hash()67 uint64 hash() const { return hash_code_; } 68 69 string ToString() const; 70 proto()71 const ConvParametersProto& proto() const { return proto_; } 72 73 private: 74 int device_id_; 75 ConvParametersProto proto_; 76 uint64 hash_code_; 77 }; 78 } // namespace tensorflow 79 #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM 80 81 #endif // TENSORFLOW_CORE_UTIL_AUTOTUNE_MAPS_CONV_PARAMETERS_H_ 82