wehub-resource-sync
927 lines
22 KiB
Python
927 lines
22 KiB
Python
# 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",
|
|
)
|