chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,926 @@
|
||||
# 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",
|
||||
)
|
||||
Reference in New Issue
Block a user