# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # ruff: noqa: F821 """External function interface to CuDNN v7 library.""" # pylint: disable-msg=C0103 import ctypes import numpy as np import tvm_ffi import tvm from tvm import te # algos can be read from cudnn.h _FWD_ALGOS = [ "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM", "CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM", "CUDNN_CONVOLUTION_FWD_ALGO_GEMM", "CUDNN_CONVOLUTION_FWD_ALGO_DIRECT", "CUDNN_CONVOLUTION_FWD_ALGO_FFT", "CUDNN_CONVOLUTION_FWD_ALGO_FFT_TILING", "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD", "CUDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED", "CUDNN_CONVOLUTION_FWD_ALGO_COUNT", ] def exists(): """ Checks whether the local machine can use CuDNN. Returns ------- exists: bool True if CuDNN support is enabled and a CuDNN-capable GPU exists. Otherwise, False. """ func = tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True) if func is None: return False return bool(func()) def algo_to_index(algo_type, algo_name): """Return a index represents the algorithm, which can be used in calling CuDNN function Parameters ---------- algo_type : str One of ``"fwd"``, ``"bwd_filter"``, or ``"bwd_data"``. algo_name : str Algorithm name as defined in cuDNN. For example: * fwd: ``CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM``, etc. * bwd_filter: ``CUDNN_CONVOLUTION_BWD_FILTER_ALGO_0``, etc. * bwd_data: ``CUDNN_CONVOLUTION_BWD_DATA_ALGO_0``, etc. Returns ------- algo : int Algorithm index """ idx = -1 if algo_type == "fwd": idx = _FWD_ALGOS.index(algo_name) elif algo_type == "bwd_filter": idx = _BWD_FILTER_ALGOS.index(algo_name) elif algo_type == "bwd_data": idx = _BWD_DATA_ALGOS.index(algo_name) assert idx >= 0 return idx def _get_np_int32_array_handle(arr): """Return a void_p handle for a numpy array Parameters ---------- arr: numpy.Tensor source numpy array Returns ------- ptr: ctypes.c_void_p pointer to the data """ assert arr.dtype == np.int32 ptr = arr.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)) return ctypes.cast(ptr, ctypes.c_void_p) def _prepare_global_func_params(dims, pad, stride, dilation, x_shape=None, w_shape=None): full_dims = dims + 2 if x_shape: assert isinstance(x_shape, list) assert len(x_shape) == full_dims if w_shape: assert isinstance(w_shape, list) assert len(w_shape) == full_dims pad = ( np.full(dims, pad, dtype=np.int32) if isinstance(pad, int) else np.array(pad, dtype=np.int32) ) stride = ( np.full(dims, stride, dtype=np.int32) if isinstance(stride, int) else np.array(stride, dtype=np.int32) ) dilation = ( np.full(dims, dilation, dtype=np.int32) if isinstance(dilation, int) else np.array(dilation, dtype=np.int32) ) xshape = np.array(x_shape, dtype=np.int32) if x_shape else None wshape = np.array(w_shape, dtype=np.int32) if x_shape else None return pad, stride, dilation, xshape, wshape def conv_output_shape( tensor_format, pad, stride, dilation, x_shape, w_shape, data_dtype, conv_dtype, groups=1 ): """Get output shape of 2D or 3D convolution Paramters --------- tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC 2: CUDNN_TENSOR_NCHW_VECT_C pad: int or list padding stride: int or list stride dilation: int or list dilation x_shape: list input shape w_shape: list weight shape data_dtype: str data type conv_dtype: str convolution type groups: int number of groups Returns ------- oshape: list output shape """ assert len(x_shape) == len(w_shape) assert len(x_shape) in (4, 5) if tensor_format == 0: n_output = x_shape[0] c_output = w_shape[0] x_chan = x_shape[1] w_chan_input = w_shape[1] x_shape = x_shape[2:] w_shape = w_shape[2:] elif tensor_format == 1: n_output = x_shape[0] c_output = w_shape[0] x_chan = x_shape[-1] w_chan_input = w_shape[-1] assert len(x_shape) == 4, "CuDNN layout NHWC is only well-defined for 4d tensors" x_shape = x_shape[1:-1] w_shape = w_shape[1:-1] elif tensor_format == 2: n_output = x_shape[0] c_output = w_shape[0] x_chan = x_shape[1] w_chan_input = w_shape[1] w_lanes = tvm.runtime.DataType(conv_dtype).lanes assert w_lanes == 1 x_shape = x_shape[2:] w_shape = w_shape[2:] else: raise ValueError(f"Unknown CuDNN tensor format: '{tensor_format}'") x_lanes = tvm.runtime.DataType(data_dtype).lanes assert x_chan * x_lanes == w_chan_input * groups, ( f"Mismatched dimensions, data has {x_chan // groups} channels/group " f"(dimension {x_chan} with {x_lanes} lanes/value, {groups} groups), " f"but weights require {w_chan_input} input channels/group" ) output_dims = [] for x_shape_i, w_shape_i, pad_i, stride_i, dilation_i in zip( x_shape, w_shape, pad, stride, dilation ): output_dim = 1 + (x_shape_i + 2 * pad_i - (((w_shape_i - 1) * dilation_i) + 1)) // stride_i output_dims.append(output_dim) if tensor_format in [0, 2]: output = [n_output, c_output, *output_dims] elif tensor_format == 1: output = [n_output, *output_dims, c_output] else: raise ValueError(f"Unknown CuDNN tensor format: '{tensor_format}'") return output def conv_dgrad_shape( tensor_format, pad, stride, dilation, dy_shape, w_shape, output_padding=(0, 0), groups=1 ): """Get output shape of conv2d gradient with respect to data Paramters --------- tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC pad: int or list padding stride: int or list stride dilation: int or list dilation dy_shape: list output gradient shape w_shape: list weight shape data_dtype: str data type conv_dtype: str convolution type groups: int number of groups Returns ------- oshape: list output shape """ assert len(dy_shape) == len(w_shape) assert len(dy_shape) == 4 if tensor_format == 0: N = dy_shape[0] C = w_shape[1] * groups dy_shape = dy_shape[2:] w_shape = w_shape[2:] elif tensor_format == 1: N = dy_shape[0] C = w_shape[-1] * groups dy_shape = dy_shape[1:-1] w_shape = w_shape[1:-1] else: raise ValueError(f"Unsupported CuDNN tensor format: '{tensor_format}'") input_dims = [] for dy_shape_i, w_shape_i, pad_i, stride_i, dilation_i, out_pad in zip( dy_shape, w_shape, pad, stride, dilation, output_padding ): input_dim = ( (dy_shape_i - 1) * stride_i - 2 * pad_i + (((w_shape_i - 1) * dilation_i) + 1) + out_pad ) input_dims.append(input_dim) if tensor_format == 0: output = [N, C, *input_dims] else: output = [N, *input_dims, C] return output def _conv_find_algo( func_name, tensor_format, pad, stride, dilation, x_shape, w_shape, y_shape, data_dtype, conv_dtype, groups=1, verbose=False, ): """ Common function to choose the best cudnn convolution algorithm for the given input and the convolution type. """ dims = len(x_shape) assert dims in (4, 5) pad, stride, dilation, xshape, wshape = _prepare_global_func_params( dims - 2, pad, stride, dilation, x_shape, w_shape ) yshape = np.array(y_shape, dtype=np.int32) func = tvm_ffi.get_global_func(func_name) return func( tensor_format, dims - 2, _get_np_int32_array_handle(pad), _get_np_int32_array_handle(stride), _get_np_int32_array_handle(dilation), _get_np_int32_array_handle(xshape), _get_np_int32_array_handle(wshape), _get_np_int32_array_handle(yshape), data_dtype, conv_dtype, groups, verbose, ) def conv_forward_find_algo( tensor_format, pad, stride, dilation, x_shape, w_shape, y_shape, data_dtype, conv_dtype, groups=1, verbose=True, ): """Choose the best forward algorithm for the given input. Paramters --------- tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC 2: CUDNN_TENSOR_NCHW_VECT_C pad: int or list padding stride: int or list stride dilation: int or list dilation x_shape: list input shape w_shape: list weight shape y_shape: list output shape data_dtype: str data type conv_dtype: str convolution type groups: int number of groups Returns ------- algo: int algo chosen by CUDNN """ return _conv_find_algo( "tvm.contrib.cudnn.conv.forward_find_algo", tensor_format, pad, stride, dilation, x_shape, w_shape, y_shape, data_dtype, conv_dtype, groups, verbose, ) def conv_backward_data_find_algo( tensor_format, pad, stride, dilation, dy_shape, w_shape, dx_shape, data_dtype, conv_dtype, groups=1, verbose=True, ): """Choose the best backward data algorithm for the given input. Paramters --------- tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC 2: CUDNN_TENSOR_NCHW_VECT_C pad: int or list padding stride: int or list stride dilation: int or list dilation dy_shape: list output gradient shape w_shape: list weight shape dx_shape: list dgrad shape data_dtype: str data type conv_dtype: str convolution type groups: int number of groups verbose: bool whether to show the selection trials Returns ------- algo: int algo chosen by CUDNN """ return _conv_find_algo( "tvm.contrib.cudnn.conv.backward_data_find_algo", tensor_format, pad, stride, dilation, dy_shape, w_shape, dx_shape, data_dtype, conv_dtype, groups, verbose, ) def conv_backward_filter_find_algo( tensor_format, pad, stride, dilation, dy_shape, x_shape, dw_shape, data_dtype, conv_dtype, groups=1, verbose=True, ): """Choose the best backward filter algorithm for the given input. Paramters --------- tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC 2: CUDNN_TENSOR_NCHW_VECT_C pad: int or list padding stride: int or list stride dilation: int or list dilation dy_shape: list output gradient shape x_shape: list weight shape dw_shape: list wgrad shape data_dtype: str data type conv_dtype: str convolution type groups: int number of groups verbose: bool whether to show the selection trials Returns ------- algo: int algo chosen by CUDNN """ return _conv_find_algo( "tvm.contrib.cudnn.conv.backward_filter_find_algo", tensor_format, pad, stride, dilation, dy_shape, x_shape, dw_shape, data_dtype, conv_dtype, groups, verbose, ) def conv_forward( x, w, pad, stride, dilation, conv_mode, tensor_format, algo, conv_dtype, groups=1, verbose=True ): """Create an extern op that compute 2D or 3D convolution with CuDNN Parameters ---------- x: Tensor input feature map w: Tensor convolution weight pad: int or list padding stride: int or list stride dilation: int or list dilation conv_mode: int 0: CUDNN_CONVOLUTION 1: CUDNN_CROSS_CORRELATION tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC 2: CUDNN_TENSOR_NCHW_VECT_C algo: int Forward algorithm, get index from ```algo_to_index``` function if algo == -1, the best algo will be chosen by CUDNN conv_dtype: str convolution type groups: int the number of groups verbose: bool whether to show the selection trials Returns ------- y: Tensor The result tensor """ dims = len(x.shape) assert dims in (4, 5) conv_dtype = x.dtype if conv_dtype is None else conv_dtype pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation) x_shape = list(x.shape) if isinstance(x.shape[0], tvm.tirx.expr.IntImm): oshape = conv_output_shape( tensor_format, pad, stride, dilation, x_shape, list(w.shape), x.dtype, conv_dtype, groups, ) if algo == -1: # For now if we try to call `cudnnFindConvolutionForwardAlgorithm` when # using INT8 data type, CuDNN will crash down. # On the other hand, CuDNN only support IMPLICIT_PRECOMP_GEMM at NHWC format if tensor_format == 1 and conv_dtype == "int32": algo = 1 else: algo = conv_forward_find_algo( tensor_format, pad, stride, dilation, list(x.shape), list(w.shape), oshape, x.dtype, conv_dtype, groups, verbose, ) else: # The dynamic batch size case, pretend this is a single batch x_shape[0] = 1 oshape = conv_output_shape( tensor_format, pad, stride, dilation, x_shape, list(w.shape), x.dtype, conv_dtype, groups, ) oshape[0] = x.shape[0] # This picks CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM # It seems this is the fastest among algorithms that are always applicable algo = 1 if dims == 4: return te.extern( oshape, [x, w], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cudnn.conv2d.forward", conv_mode, tensor_format, algo, pad[0], pad[1], stride[0], stride[1], dilation[0], dilation[1], ins[0], ins[1], outs[0], conv_dtype, groups, ), name="y", ) return te.extern( oshape, [x, w], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cudnn.conv3d.forward", conv_mode, tensor_format, algo, pad[0], pad[1], pad[2], stride[0], stride[1], stride[2], dilation[0], dilation[1], dilation[2], ins[0], ins[1], outs[0], conv_dtype, groups, ), name="y", ) def conv_backward_data( dy, w, pad, stride, dilation, conv_mode, tensor_format, conv_dtype, groups=1, output_padding=(0, 0), ): """Create a CuDNN extern op that computes the gradient of 2D convolution with respect to data. Parameters ---------- dy: Tensor output gradient w: Tensor convolution weight pad: int or list padding stride: int or list stride dilation: int or list dilation conv_mode: int 0: CUDNN_CONVOLUTION 1: CUDNN_CROSS_CORRELATION tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC conv_dtype: str convolution type groups: int the number of groups Returns ------- dx: Tensor dgrad tensor """ dims = len(dy.shape) assert dims == 4 conv_dtype = dy.dtype if conv_dtype is None else conv_dtype pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation) assert isinstance(dy.shape[0], tvm.tirx.expr.IntImm), ( "Dynamic batch is not supported for cudnn conv2d backwad data yet." ) dx_shape = conv_dgrad_shape( tensor_format, pad, stride, dilation, dy.shape, w.shape, output_padding, groups ) if exists(): # When cudnn exists, find the backward data algo algo = conv_backward_data_find_algo( tensor_format, pad, stride, dilation, list(dy.shape), list(w.shape), dx_shape, dy.dtype, conv_dtype, groups, True, ) else: algo = 1 return te.extern( dx_shape, [dy, w], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cudnn.conv2d.backward_data", conv_mode, tensor_format, algo, pad[0], pad[1], stride[0], stride[1], dilation[0], dilation[1], ins[0], ins[1], outs[0], conv_dtype, groups, ), name="dx", ) def conv_backward_filter( dy, x, kernel_size, pad, stride, dilation, conv_mode, tensor_format, conv_dtype, groups=1 ): """Create a CuDNN extern op that computes the gradient of 2D convolution with respect to weight. Parameters ---------- dy: Tensor output gradient x: Tensor input tensor kernel_size: a pair of int The spatial size of the corresponding forward convolution kernel pad: int or list padding stride: int or list stride dilation: int or list dilation conv_mode: int 0: CUDNN_CONVOLUTION 1: CUDNN_CROSS_CORRELATION tensor_format: int 0: CUDNN_TENSOR_NCHW 1: CUDNN_TENSOR_NHWC conv_dtype: str convolution type groups: int the number of groups Returns ------- dw: Tensor wgrad tensor """ dims = len(x.shape) assert dims == 4 conv_dtype = x.dtype if conv_dtype is None else conv_dtype pad, stride, dilation, _, _ = _prepare_global_func_params(dims - 2, pad, stride, dilation) filter_h, filter_w = kernel_size x_shape = list(x.shape) assert isinstance(x.shape[0], tvm.tirx.expr.IntImm), ( "Dynamic batch is not supported for cudnn conv2d backwad filter yet." ) ic_ind = 1 if tensor_format == 0 else 3 if groups > 1: assert x_shape[ic_ind] == dy.shape[ic_ind] and x_shape[ic_ind] == groups, ( "Only depthwise wgrad supported for groups > 1." ) ic = 1 else: ic = x_shape[ic_ind] if tensor_format == 0: dw_shape = [dy.shape[1], ic, filter_h, filter_w] else: dw_shape = [dy.shape[3], filter_h, filter_w, ic] algo = conv_backward_filter_find_algo( tensor_format, pad, stride, dilation, list(dy.shape), list(x.shape), dw_shape, x.dtype, conv_dtype, groups, True, ) return te.extern( dw_shape, [dy, x], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cudnn.conv2d.backward_filter", conv_mode, tensor_format, algo, pad[0], pad[1], stride[0], stride[1], dilation[0], dilation[1], ins[0], ins[1], outs[0], conv_dtype, groups, ), name="dw", ) def softmax(x, axis=-1): """Compute softmax using CuDNN Parameters ---------- x : tvm.te.Tensor The input tensor axis : int The axis to compute the softmax Returns ------- ret : tvm.te.Tensor The result tensor """ return te.extern( x.shape, [x], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cudnn.softmax.forward", ins[0], outs[0], axis ), name="y", ) def log_softmax(x, axis=-1): """Compute log_softmax using CuDNN Parameters ---------- x : tvm.te.Tensor The input tensor axis : int The axis to compute log softmax over Returns ------- ret : tvm.te.Tensor The result tensor """ return te.extern( x.shape, [x], lambda ins, outs: tvm.tirx.call_packed( "tvm.contrib.cudnn.log_softmax.forward", ins[0], outs[0], axis ), name="y", )