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_XLA_SERVICE_GPU_CUBLAS_GEMM_PAD_FOR_TENSOR_CORES_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUBLAS_GEMM_PAD_FOR_TENSOR_CORES_H_ 18 19 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" 20 21 namespace xla { 22 namespace gpu { 23 24 // Adds padding to dot operations to make them run faster on GPUs with 25 // tensor cores (https://devblogs.nvidia.com/programming-tensor-cores-cuda-9/). 26 // 27 // f16 dots are padded to have input/output shapes with dimensions that 28 // are multiples of 8, so that we can use tensor cores. 29 // 30 // Don't run this pass on GPUs without tensor cores -- it will make them slower! 31 // 32 // This pass depends on xla::DotDecomposer pass, 33 // so it should go strictly later. 34 class CublasGemmPadForTensorCores : public HloModulePass { 35 public: name()36 absl::string_view name() const override { 37 return "cublas-gemm-pad-for-speed"; 38 } 39 40 StatusOr<bool> Run(HloModule* module) override; 41 }; 42 43 } // namespace gpu 44 } // namespace xla 45 46 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_CUBLAS_GEMM_PAD_FOR_TENSOR_CORES_H_ 47