diff --git a/cmake/config.cmake b/cmake/config.cmake index 0960a5ccdfa2..6b6e2623d681 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -494,5 +494,6 @@ SET(CMAKE_VS_PLATFORM_TOOLSET_HOST_ARCHITECTURE "x64") # Enable using flash-attn as a BYOC backend # Need to have USE_MACA=ON -set(USE_FLASHATTN ON) +set(USE_FLASHATTN OFF) set(USE_MCDNN ON) +set(USE_MCBLAS ON) \ No newline at end of file diff --git a/cmake/modules/LibInfo.cmake b/cmake/modules/LibInfo.cmake index 7cbd080f4172..662c498d7b11 100644 --- a/cmake/modules/LibInfo.cmake +++ b/cmake/modules/LibInfo.cmake @@ -119,6 +119,8 @@ function(add_lib_info src_file) TVM_INFO_USE_HIPBLAS="${USE_HIPBLAS}" TVM_INFO_USE_ROCM="${USE_ROCM}" TVM_INFO_USE_MACA="${USE_MACA}" + TVM_INFO_USE_MCBLAS="${USE_MCBLAS}" + TVM_INFO_USE_MCDNN="${USE_MCDNN}" TVM_INFO_USE_RCCL="${USE_RCCL}" TVM_INFO_USE_RPC="${USE_RPC}" TVM_INFO_USE_RTTI="${USE_RTTI}" diff --git a/cmake/modules/MACA.cmake b/cmake/modules/MACA.cmake index 909303cca84e..9b1884adc5e7 100644 --- a/cmake/modules/MACA.cmake +++ b/cmake/modules/MACA.cmake @@ -47,6 +47,29 @@ if(USE_MACA) list(APPEND RUNTIME_SRCS ${CONTRIB_FLASHATTN_SRCS}) list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_FLASHATTN_LIBRARY}) endif(USE_FLASHATTN) + + if(USE_MCBLAS) + message(STATUS "Build with mcBLAS support") + tvm_file_glob(GLOB MCBLAS_CONTRIB_SRC src/relay/backend/contrib/mcblas/*.cc src/relax/backend/contrib/mcblas/*.cc) + list(APPEND COMPILER_SRCS ${MCBLAS_CONTRIB_SRC}) + tvm_file_glob(GLOB CONTRIB_MCBLAS_SRCS src/runtime/contrib/mcblas/*.cc) + list(APPEND RUNTIME_SRCS ${CONTRIB_MCBLAS_SRCS}) + list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_MCBLAS_LIBRARY}) + if(NOT MACA_MCBLASLT_LIBRARY STREQUAL "MACA_MCBLASLT_LIBRARY-NOTFOUND") + list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_MCBLASLT_LIBRARY}) + endif() + endif(USE_MCBLAS) + + if(USE_MCDNN) + message(STATUS "Build with mcdnn support") + include_directories(SYSTEM ${MACA_INCLUDE_DIRS}/mcdnn) + tvm_file_glob(GLOB MCDNN_RELAY_CONTRIB_SRC src/relay/backend/contrib/mcdnn/*.cc src/relax/backend/contrib/mcdnn/*.cc) + list(APPEND COMPILER_SRCS ${MCDNN_RELAY_CONTRIB_SRC}) + tvm_file_glob(GLOB CONTRIB_MCDNN_SRCS src/runtime/contrib/mcdnn/*.cc) + list(APPEND RUNTIME_SRCS ${CONTRIB_MCDNN_SRCS}) + list(APPEND TVM_RUNTIME_LINKER_LIBS ${MACA_MCDNN_LIBRARY}) + endif(USE_MCDNN) + else(USE_MACA) list(APPEND COMPILER_SRCS src/target/opt/build_maca_off.cc) endif(USE_MACA) diff --git a/cmake/utils/FindMACA.cmake b/cmake/utils/FindMACA.cmake index 53a557de28c8..62166faba974 100644 --- a/cmake/utils/FindMACA.cmake +++ b/cmake/utils/FindMACA.cmake @@ -48,7 +48,9 @@ macro(find_maca use_maca) find_library(MACA_MACAMCC_LIBRARY mcruntime ${__maca_sdk}/lib) find_library(MACA_HCA_LIBRARY mxc-runtime64 ${__maca_sdk}/lib) find_library(MACA_FLASHATTN_LIBRARY mcFlashAttn ${__maca_sdk}/lib) - + find_library(MACA_MCBLAS_LIBRARY mcblas ${__maca_sdk}/lib) + find_library(MACA_MCBLASLT_LIBRARY mcblasLt ${__maca_sdk}/lib) + find_library(MACA_MCDNN_LIBRARY mcdnn ${__maca_sdk}/lib) if(MACA_MACAMCC_LIBRARY) set(MACA_FOUND TRUE) endif() @@ -57,6 +59,8 @@ macro(find_maca use_maca) message(STATUS "Found MACA_INCLUDE_DIRS=" ${MACA_INCLUDE_DIRS}) message(STATUS "Found MACA_MACAMCC_LIBRARY=" ${MACA_MACAMCC_LIBRARY}) message(STATUS "Found MACA_FLASHATTN_LIBRARY=" ${MACA_FLASHATTN_LIBRARY}) + message(STATUS "Found MACA_MCBLAS_LIBRARY=" ${MACA_MCBLAS_LIBRARY}) + message(STATUS "Found MACA_MCDNN_LIBRARY=" ${MACA_MCDNN_LIBRARY}) endif(MACA_FOUND) endmacro(find_maca) diff --git a/include/tvm/topi/contrib/mcblas.h b/include/tvm/topi/contrib/mcblas.h new file mode 100644 index 000000000000..d6858c30806d --- /dev/null +++ b/include/tvm/topi/contrib/mcblas.h @@ -0,0 +1,88 @@ +/* + * 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. + */ + +/*! + * \brief External function interface to mcBLAS libraries + * \file mcblas.h + */ +#ifndef TVM_TOPI_CONTRIB_MCBLAS_H_ +#define TVM_TOPI_CONTRIB_MCBLAS_H_ + +#include +#include + +namespace tvm { +namespace topi { +namespace contrib { + +using namespace tvm::te; +using namespace topi::detail; +/*! + * \brief Create an op that multiplies lhs and rhs with mcBLAS + * + * \param lhs The left matrix operand + * \param rhs The right matrix operand + * \param transa Whether to transpose lhs + * \param transb Whether to transpose rhs + * + * \return The output tensor + */ +inline Tensor mcblas_matmul(const Tensor& lhs, const Tensor& rhs, bool transa, bool transb) { + auto n = transa ? lhs->shape[1] : lhs->shape[0]; + auto m = transb ? rhs->shape[0] : rhs->shape[1]; + + return make_extern( + {{n, m}}, {lhs->dtype}, {lhs, rhs}, + [&](Array ins, Array outs) { + return call_packed({StringImm("tvm.contrib.mcblas.matmul"), pack_buffer(ins[0]), + pack_buffer(ins[1]), pack_buffer(outs[0]), transa, transb}); + }, + "C", "", {})[0]; +} + +/*! + * \brief Create an op that multiplies batch matrices + * lhs and rhs with mcBLAS + * + * \param lhs The left matrix operand + * \param rhs The right matrix operand + * \param transa Whether to transpose lhs + * \param transb Whether to transpose rhs + * + * \return The output tensor + */ +inline Tensor mcblas_batch_matmul(const Tensor& lhs, const Tensor& rhs, bool transa, bool transb) { + auto b = lhs->shape[0]; + auto n = transa ? lhs->shape[2] : lhs->shape[1]; + auto m = transb ? rhs->shape[1] : rhs->shape[2]; + + return make_extern( + {{b, n, m}}, {lhs->dtype}, {lhs, rhs}, + [&](Array ins, Array outs) { + return call_packed({StringImm("tvm.contrib.mcblas.batch_matmul"), pack_buffer(ins[0]), + pack_buffer(ins[1]), pack_buffer(outs[0]), transa, transb}); + }, + "C", "", {})[0]; +} + +} // namespace contrib +} // namespace topi +} // namespace tvm + +#endif // TVM_TOPI_CONTRIB_MCBLAS_H_ diff --git a/include/tvm/topi/maca/dense.h b/include/tvm/topi/maca/dense.h new file mode 100644 index 000000000000..88fe258fc7b4 --- /dev/null +++ b/include/tvm/topi/maca/dense.h @@ -0,0 +1,99 @@ +/* + * 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. + */ + +/*! + * \file maca/dense.h + * \brief MACA schedule for dense operation + */ +#ifndef TVM_TOPI_MACA_DENSE_H_ +#define TVM_TOPI_MACA_DENSE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace topi { + +using namespace tvm::te; + +namespace maca { +/*! + * \brief Implementation of dense for MACA backend + * + * \param target The target device + * \param data Tensor with shape [batch, in_dim] + * \param weight Tensor with shape [out_dim, in_dim] + * \param bias Tensor with shape [out_dim]. Optional; to omit bias, pass Tensor() + * \param out_dtype Output data type. Used for mixed precision. + * + * \return Tensor with shape [batch, out_dim] + */ +inline tvm::te::Tensor dense_maca(const Target& target, const tvm::te::Tensor& data, + const tvm::te::Tensor& weight, const tvm::te::Tensor& bias, + const DataType& out_dtype) { + ICHECK_EQ(data->shape.size(), 2) << "dense requires 2-D data"; + ICHECK_EQ(weight->shape.size(), 2) << "dense requires 2-D weight"; + if (bias.defined()) { + ICHECK_EQ(bias->shape.size(), 1) << "dense requires 1-D bias"; + } + + auto batch = data->shape[0]; + auto in_dim = data->shape[1]; + auto out_dim = weight->shape[0]; + + if (target->GetLibs().count("mcblas")) { + ICHECK_EQ(data->dtype, out_dtype) << "Mixed precision not supported."; + auto mm = topi::contrib::mcblas_matmul(data, weight, false, true); + if (bias.defined()) { + mm = tvm::te::compute( + {batch, out_dim}, [&](Var i, Var j) { return mm(i, j) + bias(j); }, "tensor", kBroadcast); + } + + return mm; + } else { + return topi::nn::dense(data, weight, bias, out_dtype); + } +} + +/*! + * \brief Create a MACA schedule for dense + * + * \param target The target to generate a schedule for. + * \param outs The output tensors. + * + * \return A schedule for the given ops. + */ +inline Schedule schedule_dense(const Target& target, const Array& outs) { + if (target->kind->name == "maca" && target->GetLibs().count("mcblas")) { + return topi::generic::schedule_extern(target, outs); + } + return topi::cuda::schedule_dense(target, outs); +} + +} // namespace maca +} // namespace topi +} // namespace tvm +#endif // TVM_TOPI_MACA_DENSE_H_ diff --git a/python/tvm/contrib/mcblas.py b/python/tvm/contrib/mcblas.py new file mode 100644 index 000000000000..d54e9a81d745 --- /dev/null +++ b/python/tvm/contrib/mcblas.py @@ -0,0 +1,86 @@ +# 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. +"""External function interface to mcBLAS libraries.""" +import tvm +from tvm import te + + +def matmul(lhs, rhs, transa=False, transb=False, dtype=None): + """Create an extern op that compute matrix mult of A and rhs with mcBLAS + + Parameters + ---------- + lhs : Tensor + The left matrix operand + rhs : Tensor + The right matrix operand + transa : bool + Whether transpose lhs + transb : bool + Whether transpose rhs + + Returns + ------- + C : Tensor + The result tensor. + """ + n = lhs.shape[1] if transa else lhs.shape[0] + m = rhs.shape[0] if transb else rhs.shape[1] + dtype = dtype if dtype is not None else lhs.dtype + return te.extern( + (n, m), + [lhs, rhs], + lambda ins, outs: tvm.tir.call_packed( + "tvm.contrib.mcblas.matmul", ins[0], ins[1], outs[0], transa, transb + ), + dtype=dtype, + name="matmul_mcblas", + ) + + +def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None): + """Create an extern op that compute batch matrix mult of A and rhs with mcBLAS + + Parameters + ---------- + lhs : Tensor + The left matrix operand + rhs : Tensor + The right matrix operand + transa : bool + Whether transpose lhs + transb : bool + Whether transpose rhs + + Returns + ------- + C : Tensor + The result tensor. + """ + b = lhs.shape[0] + n = lhs.shape[2] if transa else lhs.shape[1] + m = rhs.shape[1] if transb else rhs.shape[2] + dtype = dtype if dtype is not None else lhs.dtype + return te.extern( + (b, n, m), + [lhs, rhs], + lambda ins, outs: tvm.tir.call_packed( + "tvm.contrib.mcblas.batch_matmul", ins[0], ins[1], outs[0], transa, transb + ), + dtype=dtype, + name="batch_matmul_mcblas", + ) diff --git a/python/tvm/contrib/mcblaslt.py b/python/tvm/contrib/mcblaslt.py new file mode 100644 index 000000000000..24091320c537 --- /dev/null +++ b/python/tvm/contrib/mcblaslt.py @@ -0,0 +1,54 @@ +# 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. +"""External function interface to mcBLASlt libraries.""" +import tvm +from tvm import te + + +def matmul(lhs, rhs, transa=False, transb=False, n=0, m=0, dtype=None): + """Create an extern op that compute matrix mult of A and rhs with mcBLAS + + Parameters + ---------- + lhs : Tensor + The left matrix operand + rhs : Tensor + The right matrix operand + transa : bool + Whether transpose lhs + transb : bool + Whether transpose rhs + + Returns + ------- + C : Tensor + The result tensor. + """ + if n == 0: + n = lhs.shape[1] if transa else lhs.shape[0] + if m == 0: + m = rhs.shape[0] if transb else rhs.shape[1] + dtype = dtype if dtype is not None else lhs.dtype + return te.extern( + (n, m), + [lhs, rhs], + lambda ins, outs: tvm.tir.call_packed( + "tvm.contrib.mcblaslt.matmul", ins[0], ins[1], outs[0], transa, transb + ), + dtype=dtype, + name="C", + ) diff --git a/python/tvm/contrib/mcdnn.py b/python/tvm/contrib/mcdnn.py new file mode 100644 index 000000000000..36c6db10b78d --- /dev/null +++ b/python/tvm/contrib/mcdnn.py @@ -0,0 +1,953 @@ +# 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. +"""External function interface to McDNN v7 library.""" +# pylint: disable-msg=C0103 +import ctypes +import numpy as np +import tvm + +import tvm._ffi +from tvm import te + +# algos can be read from mcdnn.h +_FWD_ALGOS = [ + "MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_DIRECT", + "MCDNN_CONVOLUTION_FWD_ALGO_FFT", + "MCDNN_CONVOLUTION_FWD_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_FWD_ALGO_WINOGRAD", + "MCDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_FWD_ALGO_COUNT", +] + + +def exists(): + """ + Checks whether the local machine can use McDNN. + + Returns + ------- + exists: bool + + True if McDNN support is enabled and a McDNN-capable GPU + exists. Otherwise, False. + """ + func = tvm.get_global_func("tvm.contrib.mcdnn.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 McDNN function + + Parameters + ---------- + algo_type : str + ["fwd", "bwd_filter", "bwd_data] + + algo_name : str + algorithm name in mcdnn definition + fwd = [ + "MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_DIRECT", + "MCDNN_CONVOLUTION_FWD_ALGO_FFT", + "MCDNN_CONVOLUTION_FWD_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_FWD_ALGO_WINOGRAD", + "MCDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_FWD_ALGO_COUNT", + ] + bwd_filter = [ + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_0", + # non-deterministic + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_1", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_3", + # non-deterministic, algo0 with workspaceS + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD", + # not implemented + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT", + ] + bwd_data = [ + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_0", + # non-deterministic + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_1", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_FFT", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT", + ] + + 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.NDArray + 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: MCDNN_TENSOR_NCHW + 1: MCDNN_TENSOR_NHWC + 2: MCDNN_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, "McDNN 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 McDNN tensor format: '{tensor_format}'") + + x_lanes = tvm.runtime.DataType(data_dtype).lanes + assert x_chan * x_lanes == w_chan_input * groups, ( + "Mismatched dimensions, data has {} channels/group " + "(dimension {} with {} lanes/value, {} groups), " + "but weights require {} input channels/group" + ).format(x_chan // groups, x_chan, x_lanes, groups, w_chan_input) + + 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 McDNN 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: MCDNN_TENSOR_NCHW + 1: MCDNN_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 McDNN 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 mcdnn 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: MCDNN_TENSOR_NCHW + 1: MCDNN_TENSOR_NHWC + 2: MCDNN_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 MCDNN + """ + return _conv_find_algo( + "tvm.contrib.mcdnn.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: MCDNN_TENSOR_NCHW + 1: MCDNN_TENSOR_NHWC + 2: MCDNN_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 MCDNN + """ + return _conv_find_algo( + "tvm.contrib.mcdnn.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: MCDNN_TENSOR_NCHW + 1: MCDNN_TENSOR_NHWC + 2: MCDNN_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 MCDNN + """ + return _conv_find_algo( + "tvm.contrib.mcdnn.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 McDNN + + 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: MCDNN_CONVOLUTION + 1: MCDNN_CROSS_CORRELATION + tensor_format: int + 0: MCDNN_TENSOR_NCHW + 1: MCDNN_TENSOR_NHWC + 2: MCDNN_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 MCDNN + 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.tir.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 `mcdnnFindConvolutionForwardAlgorithm` when + # using INT8 data type, McDNN will crash down. + # On the other hand, McDNN 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 MCDNN_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.tir.call_packed( + "tvm.contrib.mcdnn.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.tir.call_packed( + "tvm.contrib.mcdnn.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 McDNN 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: MCDNN_CONVOLUTION + 1: MCDNN_CROSS_CORRELATION + tensor_format: int + 0: MCDNN_TENSOR_NCHW + 1: MCDNN_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.tir.expr.IntImm + ), "Dynamic batch is not supported for mcdnn conv2d backwad data yet." + + dx_shape = conv_dgrad_shape( + tensor_format, pad, stride, dilation, dy.shape, w.shape, output_padding, groups + ) + + if exists(): + # When mcdnn 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.tir.call_packed( + "tvm.contrib.mcdnn.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 McDNN 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: MCDNN_CONVOLUTION + 1: MCDNN_CROSS_CORRELATION + tensor_format: int + 0: MCDNN_TENSOR_NCHW + 1: MCDNN_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.tir.expr.IntImm + ), "Dynamic batch is not supported for mcdnn 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.tir.call_packed( + "tvm.contrib.mcdnn.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 McDNN + + 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.tir.call_packed( + "tvm.contrib.mcdnn.softmax.forward", ins[0], outs[0], axis + ), + name="y", + ) + + +def log_softmax(x, axis=-1): + """Compute log_softmax using McDNN + + 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.tir.call_packed( + "tvm.contrib.mcdnn.log_softmax.forward", ins[0], outs[0], axis + ), + name="y", + ) diff --git a/python/tvm/contrib/mxcc.py b/python/tvm/contrib/mxcc.py index ca59b3dad583..b9aec4179d87 100644 --- a/python/tvm/contrib/mxcc.py +++ b/python/tvm/contrib/mxcc.py @@ -160,20 +160,19 @@ def have_matrixcore(compute_version=None): else: raise RuntimeError("No MACA runtime found") major, _ = parse_compute_version(compute_version) - # matrix core first introduced in 8.0 - if major >= 8: + # matrix core first introduced in 10.0 + if major >= 10: return True return False - def have_fp16(compute_version): """Either fp16 support is provided in the compute capability or not Parameters ---------- compute_version: str - compute capability of a GPU (e.g. "6.0") + compute capability of a GPU (e.g. "10.0") """ major, minor = parse_compute_version(compute_version) if major >= 10: @@ -181,7 +180,6 @@ def have_fp16(compute_version): return False - @tvm._ffi.register_func("tvm_callback_maca_get_arch") def get_maca_arch(maca_path="/opt/maca"): """Utility function to get the MetaX GPU architecture diff --git a/python/tvm/relax/backend/contrib/mcblas.py b/python/tvm/relax/backend/contrib/mcblas.py new file mode 100644 index 000000000000..a8e8ac770f7a --- /dev/null +++ b/python/tvm/relax/backend/contrib/mcblas.py @@ -0,0 +1,243 @@ +# 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. + +"""Pattern table for mcBLAS backend""" +import operator +from functools import reduce + +import tvm +from tvm import DataType +from tvm.arith import Analyzer +from tvm.relax import transform +from tvm.relax.transform import PatternCheckContext + +from ..pattern_registry import get_patterns_with_prefix, register_patterns +from ..patterns import ( + make_matmul_dequantize_pattern, + make_matmul_multiply_pattern, + make_matmul_pattern, +) +from ..utils import has_leaking_intermediate_variables + + +def _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype): + """Check if dtypes in the given workload are supported by mcBLAS BYOC.""" + if lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn": + # The output cannot be 'float8_e5m2' if inputs are 'float8_e4m3fn' + return out_dtype != "float8_e5m2" + return ( + (lhs_dtype == "float16" and rhs_dtype == "float16") + or (lhs_dtype == "float32" and rhs_dtype == "float32") + or (lhs_dtype == "int8" and rhs_dtype == "int8") + ) + + +def _check_matmul(context: PatternCheckContext) -> bool: + if has_leaking_intermediate_variables(context): + return False + lhs = context.annotated_expr["lhs"] + rhs = context.annotated_expr["rhs"] + matmul_call = context.annotated_expr["root"] + + if "scale" in context.annotated_expr and "zp" in context.annotated_expr: + scale = context.annotated_expr["scale"] + zero_point = context.annotated_expr["zp"] + # Only scalar values for scale and zero_point are supported. + if scale.struct_info.ndim != 0 or zero_point.struct_info.ndim != 0: + return False + # Only zero_point == 0.0 is supported. + if zero_point.data.numpy()[()].item() != 0.0: + return False + + lhs_dtype = lhs.struct_info.dtype + rhs_dtype = rhs.struct_info.dtype + out_dtype = matmul_call.struct_info.dtype + if not _is_supported_dtype(lhs_dtype, rhs_dtype, out_dtype): + return False + + lhs_shape = lhs.struct_info.shape.values + rhs_shape = rhs.struct_info.shape.values + + if not isinstance(lhs_shape[-1], (tvm.tir.expr.IntImm, int)): + # Reduction axis must be constant + return False + + if lhs_dtype == "int8" and rhs_dtype == "int8": + if lhs_shape[-1] % 4 != 0: + # Reduction axis must be multiples of 4 for IGEMM + return False + if not isinstance(rhs_shape[-1], (tvm.tir.expr.IntImm, int)) or rhs_shape[-1] % 4 != 0: + # Rows number must be multiples of 4 for IGEMM + return False + elif lhs_dtype == "float8_e4m3fn" and rhs_dtype == "float8_e4m3fn": + matmul_rhs_var = matmul_call.args[1] + rhs_transposed = False + if matmul_rhs_var in context.matched_bindings: + matmul_rhs_call = context.matched_bindings[matmul_rhs_var] + assert ( + isinstance(matmul_rhs_call, tvm.relax.Call) + and matmul_rhs_call.op.name == "relax.permute_dims" + ) + rhs_transposed = True + + if not rhs_transposed: + # mcBLAS FP8 operations require rhs being transposed + return False + + # mcBLAS FP8 operations require all tensors being aligned to 16 bytes. + if ( + not isinstance(rhs_shape[-1], (tvm.tir.expr.IntImm, int)) + or rhs_shape[-1] % (16 // DataType(lhs_dtype).itemsize()) != 0 + ): + return False + if ( + not isinstance(rhs_shape[-2], (tvm.tir.expr.IntImm, int)) + or rhs_shape[-2] % (16 // DataType(out_dtype).itemsize()) != 0 + ): + return False + + lhs_batches = reduce(operator.mul, lhs_shape[:-2], 1) + rhs_batches = reduce(operator.mul, rhs_shape[:-2], 1) + + if "bias" in context.annotated_expr: + if lhs_dtype == "int8" and rhs_dtype == "int8": + # Non-default epilogue not supported for IGEMM + return False + bias = context.annotated_expr["bias"] + bias_shape = bias.struct_info.shape.values + bias_batches = reduce(operator.mul, bias_shape[:-1], 1) + if not isinstance(bias_batches, (tvm.tir.expr.IntImm, int)) or int(bias_batches) > 1: + # mcBLAS only supports bias vector + return False + + analyzer = Analyzer() + + # mcBLASLt does not seem to support batched GEMM with one of matrices having + # one batch (with batch_stride 0). So for batched GEMM, the two batch counts + # must be equal. If lhs is batched but rhs is not, we can use the regular GEMM by + # flattening all batch axes into the M axis. + return ( + isinstance(lhs_batches, tvm.tir.Var) + or isinstance(rhs_batches, tvm.tir.Var) + or (analyzer.can_prove_equal(lhs_batches, rhs_batches)) + or (analyzer.can_prove(lhs_batches >= 1) and analyzer.can_prove(rhs_batches == 1)) + ) + + +register_patterns( + [ + ( + "mcblas.matmul", + *make_matmul_pattern( + with_bias=False, + ), + _check_matmul, + ), + ( + "mcblas.matmul_bias", + *make_matmul_pattern( + with_bias=True, + ), + _check_matmul, + ), + ( + "mcblas.matmul_bias_relu", + *make_matmul_pattern( + with_bias=True, + activation="relax.nn.relu", + ), + _check_matmul, + ), + ( + "mcblas.matmul_bias_gelu", + *make_matmul_pattern( + with_bias=True, + activation="relax.nn.gelu", + ), + _check_matmul, + ), + ( + "mcblas.matmul_transposed", + *make_matmul_pattern( + with_bias=False, + transposed_rhs=True, + ), + _check_matmul, + ), + ( + "mcblas.matmul_transposed_bias", + *make_matmul_pattern( + with_bias=True, + transposed_rhs=True, + ), + _check_matmul, + ), + ( + "mcblas.matmul_transposed_bias_relu", + *make_matmul_pattern( + with_bias=True, + activation="relax.nn.relu", + transposed_rhs=True, + ), + _check_matmul, + ), + ( + "mcblas.matmul_transposed_bias_gelu", + *make_matmul_pattern( + with_bias=True, + activation="relax.nn.gelu", + transposed_rhs=True, + ), + _check_matmul, + ), + ( + "mcblas.matmul_transposed_dequantize", + *make_matmul_dequantize_pattern(transposed_rhs=True), + _check_matmul, + ), + ( + "mcblas.matmul_transposed_multiply", + *make_matmul_multiply_pattern(transposed_rhs=True), + _check_matmul, + ), + ] +) + + +def partition_for_mcblas(mod, bind_constants=False): + """ + Partition the input module into mcBLAS-supported subgraphs. + + Parameters + ---------- + mod: tvm.IRModule + The IRModule to be partitioned. + + bind_constants : bool + Whether or not to keep bound constants in the grouped function. + + Returns + ------- + mod: tvm.IRModule + The resulting IRModule, containing partitioned subgraphs to be + offloaded to the mcBLAS backend. + """ + + patterns = get_patterns_with_prefix("mcblas") + return transform.FuseOpsByPattern( + patterns, bind_constants=bind_constants, annotate_codegen=True + )(mod) diff --git a/python/tvm/relay/backend/te_compiler.py b/python/tvm/relay/backend/te_compiler.py index 84e4ecbaecfb..202ac9a9afaf 100644 --- a/python/tvm/relay/backend/te_compiler.py +++ b/python/tvm/relay/backend/te_compiler.py @@ -137,7 +137,7 @@ def get_valid_implementations(op, attrs, inputs, out_type, target): return ret -def select_implementation(op, attrs, inputs, out_type, target, use_autotvm=True): +def select_implementation(op, attrs, inputs, out_type, target, use_autotvm=False): """Select the best implementation from the op strategy. If use_autotvm is True, it'll first try to find the best implementation diff --git a/python/tvm/relay/op/contrib/mcblas.py b/python/tvm/relay/op/contrib/mcblas.py new file mode 100644 index 000000000000..22d40b14c452 --- /dev/null +++ b/python/tvm/relay/op/contrib/mcblas.py @@ -0,0 +1,147 @@ +# 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. +# pylint: disable=unused-argument +"""mcBLAS Relay integration.""" +from typing import Callable, List, Tuple, Dict, Optional + +import tvm +import tvm.ir +from tvm import relay +from tvm import te +from tvm.relay import transform +from tvm.contrib import mcblas + +from ...dataflow_pattern import is_op, wildcard +from .te_target import lower_composite, relay_to_runtime +from .register import register_pattern_table + + +tvm._ffi.register_func("relay.ext.mcblas", relay_to_runtime(tvm.target.maca())) + + +def partition_for_mcblas( + mod: tvm.IRModule, params: Optional[Dict[str, tvm.runtime.NDArray]] = None +) -> tvm.IRModule: + """Partition the graph to offload for mcBLAS. + + Parameters + ---------- + mod : tvm.IRModule + The module to partition. + params : Optional[Dict[str, tvm.runtime.NDArray]] + Constant input parameters. + + Returns + ------- + tvm.IRModule + The partitioned module. + """ + + seq = tvm.transform.Sequential( + [ + transform.InferType(), + transform.MergeComposite(pattern_table()), + transform.AnnotateTarget("mcblas"), + transform.PartitionGraph(), + transform.InferType(), + ] + ) + return seq(mod) + + +@register_pattern_table("mcblas") +def pattern_table() -> List[Tuple[str, relay.Pattern, Callable[[relay.Call], bool]]]: + """Get the mcBLAS pattern table.""" + + def matmul_pattern() -> relay.Pattern: + """Create pattern for matmul.""" + return is_op("nn.matmul")(wildcard(), wildcard()) + + def batch_matmul_pattern() -> relay.Pattern: + """Create pattern for batch_matmul.""" + return is_op("nn.batch_matmul")(wildcard(), wildcard()) + + def dense_pattern() -> relay.Pattern: + """Create pattern for dense.""" + return is_op("nn.dense")(wildcard(), wildcard()) + + def check_matmul_like(matched: relay.Call) -> bool: + """Check if matmul is supported by mcBLAS.""" + # Input data types can't be mixed + if matched.args[0].checked_type.dtype != matched.args[1].checked_type.dtype: + return False + + in_dtype = matched.args[0].checked_type.dtype + out_dtype = matched.checked_type.dtype + # Only the following data type combinations are supported + if (in_dtype, out_dtype) not in [ + ("float32", "float32"), + ("float16", "float16"), + ("float16", "float32"), + ("int8", "int32"), + ("float64", "float64"), + ("int8", "float32"), + ]: + return False + + # If inputs are int8, input column strides must be a multiple of 4 + if in_dtype == "int8": + if ( + matched.args[0].checked_type.shape[-1] % 4 != 0 + or matched.args[1].checked_type.shape[-1] % 4 != 0 + ): + return False + + return True + + return [ + ("mcblas.matmul", matmul_pattern(), check_matmul_like), + ("mcblas.batch_matmul", batch_matmul_pattern(), check_matmul_like), + ("mcblas.dense", dense_pattern(), check_matmul_like), + ] + + +@lower_composite("mcblas.matmul") +def _lower_matmul(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a matmul using mcBLAS.""" + return mcblas.matmul( + inputs[0], + inputs[1], + transa=op.attrs["transpose_a"], + transb=op.attrs["transpose_b"], + dtype=op.checked_type.dtype, + ) + + +@lower_composite("mcblas.batch_matmul") +def _lower_batch_matmul(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a batch_matmul using mcBLAS.""" + return mcblas.batch_matmul( + inputs[0], + inputs[1], + transa=op.attrs["transpose_a"], + transb=op.attrs["transpose_b"], + dtype=op.checked_type.dtype, + ) + + +@lower_composite("mcblas.dense") +def _lower_dense(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a dense using mcBLAS.""" + return mcblas.matmul( + inputs[0], inputs[1], transa=False, transb=True, dtype=op.checked_type.dtype + ) diff --git a/python/tvm/relay/op/contrib/mcdnn.py b/python/tvm/relay/op/contrib/mcdnn.py new file mode 100644 index 000000000000..d050017135cc --- /dev/null +++ b/python/tvm/relay/op/contrib/mcdnn.py @@ -0,0 +1,212 @@ +# 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. +# pylint: disable=unused-argument +"""mcDNN Relay integration.""" +from typing import Callable, List, Tuple + +import tvm +import tvm.ir +from tvm import relay +from tvm import te +from tvm.relay import transform +from tvm.contrib import mcdnn + +from ...dataflow_pattern import is_op, wildcard +from .te_target import lower_composite, relay_to_runtime +from .register import register_pattern_table + + +tvm._ffi.register_func("relay.ext.mcdnn", relay_to_runtime(tvm.target.maca())) + + +def partition_for_mcdnn(mod: tvm.IRModule) -> tvm.IRModule: + """Partition the graph to offload for mcDNN. + + Parameters + ---------- + mod : tvm.IRModule + The module to partition. + + Returns + ------- + tvm.IRModule + The partitioned module. + """ + + seq = tvm.transform.Sequential( + [ + transform.InferType(), + transform.MergeComposite(pattern_table()), + transform.AnnotateTarget("mcdnn"), + transform.PartitionGraph(), + transform.InferType(), + ] + ) + return seq(mod) + + +@register_pattern_table("mcdnn") +def pattern_table() -> List[Tuple[str, relay.Pattern, Callable[[relay.Call], bool]]]: + """Get the mcDNN pattern table.""" + + def softmax_pattern() -> relay.Pattern: + """Create pattern for softmax.""" + return is_op("nn.softmax")(wildcard()) + + def log_softmax_pattern() -> relay.Pattern: + """Create pattern for log_softmax.""" + return is_op("nn.log_softmax")(wildcard()) + + def conv2d_pattern() -> relay.Pattern: + """Create pattern for conv2d.""" + return is_op("nn.conv2d")(wildcard(), wildcard()) + + def conv2d_bias_act_pattern() -> relay.Pattern: + """Create pattern for fused conv2d+bias+activation.""" + conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) + bias = is_op("nn.bias_add")(conv2d, wildcard()) + return bias.optional(is_op("nn.relu")) + + def check_softmax(matched: relay.Call) -> bool: + """Check if softmax is supported by mcDNN.""" + if matched.args[0].checked_type.dtype not in ["float64", "float32", "float16"]: + return False + + return True + + def check_log_softmax(matched: relay.Call) -> bool: + """Check if log_softmax is supported by mcDNN.""" + if matched.args[0].checked_type.dtype not in ["float64", "float32", "float16"]: + return False + + if len(matched.args[0].checked_type.shape) != 2: + return False + + if matched.attrs["axis"] not in (1, -1): + return False + + return True + + def check_conv2d(matched: relay.Call) -> bool: + if matched.args[0].checked_type.dtype not in ["float64", "float32", "float16"]: + return False + + if matched.attrs["data_layout"] != "NCHW" or matched.attrs["kernel_layout"] != "OIHW": + return False + + padding = matched.attrs["padding"] + if padding[0] != padding[2] or padding[1] != padding[3]: + return False + + return True + + def check_conv2d_bias_act(matched: relay.Call) -> bool: + return True + + return [ + ("mcdnn.softmax", softmax_pattern(), check_softmax), + ("mcdnn.log_softmax", log_softmax_pattern(), check_log_softmax), + ("mcdnn.conv2d_bias_act", conv2d_bias_act_pattern(), check_conv2d_bias_act), + ("mcdnn.conv2d", conv2d_pattern(), check_conv2d), + ] + + +@lower_composite("mcdnn.softmax") +def _lower_softmax(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a softmax using mcDNN.""" + return mcdnn.softmax(inputs[0], axis=op.attrs["axis"]) + + +@lower_composite("mcdnn.log_softmax") +def _lower_log_softmax(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a log_softmax using mcDNN.""" + return mcdnn.log_softmax(inputs[0], axis=op.attrs["axis"]) + + +@lower_composite("mcdnn.conv2d_bias_act") +def _lower_conv2d_bias_act(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a fused conv2d+bias+activation using mcDNN.""" + conv_dtype = op.checked_type.dtype + if op.op.name == "nn.relu": + activation_mode = 1 # Relu + conv2d = op.args[0].args[0] + else: + activation_mode = 5 # Identity + conv2d = op.args[0] + + conv_mode = 1 + tensor_format = 0 + algo = 1 + pad = conv2d.attrs["padding"] + strides = conv2d.attrs["strides"] + dilation = conv2d.attrs["dilation"] + groups = conv2d.attrs["groups"] + + oshape = mcdnn.conv_output_shape( + tensor_format, + pad, + strides, + dilation, + inputs[0].shape, + inputs[1].shape, + inputs[0].dtype, + conv_dtype, + groups, + ) + + return te.extern( + oshape, + inputs, + lambda ins, outs: tvm.tir.call_packed( + "tvm.contrib.mcdnn.conv2d+bias+act.forward", + conv_mode, + tensor_format, + algo, + pad[0], + pad[1], + strides[0], + strides[1], + dilation[0], + dilation[1], + activation_mode, + 0, + ins[0], + ins[1], + ins[2], + outs[0], + conv_dtype, + groups, + ), + name="y", + ) + + +@lower_composite("mcdnn.conv2d") +def _lower_conv2d(op: relay.Call, inputs: List[te.Tensor]) -> te.Tensor: + """Lower a conv2d using mcDNN.""" + return mcdnn.conv_forward( + inputs[0], + inputs[1], + pad=op.attrs["padding"], + stride=op.attrs["strides"], + dilation=op.attrs["dilation"], + conv_mode=1, + tensor_format=0, + algo=1, + conv_dtype=op.checked_type.dtype, + groups=op.attrs["groups"], + ) diff --git a/python/tvm/relay/op/strategy/__init__.py b/python/tvm/relay/op/strategy/__init__.py index 1be5425e702c..40dc75d72431 100644 --- a/python/tvm/relay/op/strategy/__init__.py +++ b/python/tvm/relay/op/strategy/__init__.py @@ -30,3 +30,4 @@ from . import intel_graphics from . import hexagon from . import adreno +from . import maca diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 1fd806b7cf5c..e7a9eff3abd1 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -29,21 +29,21 @@ from .generic import * -@schedule_injective.register(["cuda", "gpu"]) +@schedule_injective.register(["cuda", "gpu", "maca"]) def schedule_injective_cuda(attrs, outs, target): """schedule injective ops for cuda""" with target: return topi.cuda.schedule_injective(outs) -@schedule_reduce.register(["cuda", "gpu"]) +@schedule_reduce.register(["cuda", "gpu", "maca"]) def schedule_reduce_cuda(attrs, outs, target): """schedule reduction ops for cuda""" with target: return topi.cuda.schedule_reduce(outs) -@concatenate_strategy.register(["cuda", "gpu"]) +@concatenate_strategy.register(["cuda", "gpu", "maca"]) def concatenate_strategy_cuda(attrs, inputs, out_type, target): strategy = _op.OpStrategy() strategy.add_implementation( @@ -54,14 +54,14 @@ def concatenate_strategy_cuda(attrs, inputs, out_type, target): return strategy -@schedule_pool.register(["cuda", "gpu"]) +@schedule_pool.register(["cuda", "gpu", "maca"]) def schedule_pool_cuda(attrs, outs, target): """schedule pooling ops for cuda""" with target: return topi.cuda.schedule_pool(outs, attrs.layout) -@schedule_pool_grad.register(["cuda", "gpu"]) +@schedule_pool_grad.register(["cuda", "gpu", "maca"]) def schedule_pool_grad_cuda(attrs, outs, target): """schedule pooling gradient ops for cuda""" with target: @@ -94,7 +94,7 @@ def softmax_strategy_cuda(attrs, inputs, out_type, target): return strategy -@fast_softmax_strategy.register(["cuda", "gpu"]) +@fast_softmax_strategy.register(["cuda", "gpu", "maca"]) def fast_softmax_strategy_cuda(attrs, inputs, out_type, target): """fast_softmax cuda strategy""" strategy = _op.OpStrategy() @@ -125,7 +125,7 @@ def log_softmax_strategy_cuda(attrs, inputs, out_type, target): return strategy -@schedule_lrn.register(["cuda", "gpu"]) +@schedule_lrn.register(["cuda", "gpu", "maca"]) def schedule_lrn_cuda(attrs, outs, target): """schedule LRN for cuda""" with target: @@ -516,7 +516,7 @@ def judge_winograd( return judge_winograd_tensorcore, judge_winograd_autotvm, judge_winograd_auto_scheduler -@conv2d_winograd_without_weight_transform_strategy.register(["cuda", "gpu"]) +@conv2d_winograd_without_weight_transform_strategy.register(["cuda", "gpu", "maca"]) def conv2d_winograd_without_weight_transform_strategy_cuda(attrs, inputs, out_type, target): """conv2d_winograd_without_weight_transform cuda strategy""" dilation = attrs.get_int_tuple("dilation") @@ -607,7 +607,7 @@ def conv2d_winograd_without_weight_transform_strategy_cuda(attrs, inputs, out_ty return strategy -@deformable_conv2d_strategy.register(["cuda", "gpu"]) +@deformable_conv2d_strategy.register(["cuda", "gpu", "maca"]) def deformable_conv2d_strategy_cuda(attrs, inputs, out_type, target): """deformable_conv2d cuda strategy""" layout = attrs.data_layout @@ -693,7 +693,7 @@ def conv2d_transpose_strategy_cuda(attrs, inputs, out_type, target): return strategy -@conv3d_transpose_strategy.register(["cuda", "gpu"]) +@conv3d_transpose_strategy.register(["cuda", "gpu", "maca"]) def conv3d_transpose_strategy_cuda(attrs, inputs, out_type, target): """conv3d_transpose cuda strategy""" layout = attrs.data_layout @@ -777,7 +777,7 @@ def conv3d_strategy_cuda(attrs, inputs, out_type, target): return strategy -@conv3d_winograd_without_weight_transform_strategy.register(["cuda", "gpu"]) +@conv3d_winograd_without_weight_transform_strategy.register(["cuda", "gpu", "maca"]) def conv3d_winograd_without_weight_transform_strategy_cuda(attrs, inputs, out_type, target): """conv3d_winograd_without_weight_transform cuda strategy""" dilation = attrs.get_int_tuple("dilation") @@ -797,7 +797,7 @@ def conv3d_winograd_without_weight_transform_strategy_cuda(attrs, inputs, out_ty return strategy -@conv1d_strategy.register(["cuda", "gpu"]) +@conv1d_strategy.register(["cuda", "gpu", "maca"]) def conv1d_strategy_cuda(attrs, inputs, out_type, target): """conv1d cuda strategy""" layout = attrs.data_layout @@ -838,7 +838,7 @@ def conv1d_strategy_cuda(attrs, inputs, out_type, target): return strategy -@conv1d_transpose_strategy.register(["cuda", "gpu"]) +@conv1d_transpose_strategy.register(["cuda", "gpu", "maca"]) def conv1d_transpose_strategy_cuda(attrs, inputs, out_type, target): """conv1d_transpose cuda strategy""" strategy = _op.OpStrategy() @@ -1016,7 +1016,7 @@ def batch_matmul_strategy_cuda(attrs, inputs, out_type, target): return strategy -@sparse_dense_strategy.register(["cuda", "gpu"]) +@sparse_dense_strategy.register(["cuda", "gpu", "maca"]) def sparse_dense_strategy_cuda(attrs, inputs, out_type, target): """sparse dense cuda strategy""" strategy = _op.OpStrategy() @@ -1029,7 +1029,7 @@ def sparse_dense_strategy_cuda(attrs, inputs, out_type, target): return strategy -@sparse_reshape_strategy.register(["cuda", "gpu"]) +@sparse_reshape_strategy.register(["cuda", "gpu", "maca"]) def sparse_reshape_strategy_cuda(attrs, inputs, out_type, target): strategy = _op.OpStrategy() strategy.add_implementation( @@ -1040,7 +1040,7 @@ def sparse_reshape_strategy_cuda(attrs, inputs, out_type, target): return strategy -@sparse_dense_padded_strategy.register(["cuda", "gpu", "rocm"]) +@sparse_dense_padded_strategy.register(["cuda", "gpu", "rocm", "maca"]) def sparse_dense_padded_strategy_cuda(attrs, inputs, out_type, target): """sparse dense cuda strategy""" strategy = _op.OpStrategy() @@ -1053,7 +1053,7 @@ def sparse_dense_padded_strategy_cuda(attrs, inputs, out_type, target): return strategy -@scatter_elements_strategy.register(["cuda", "gpu"]) +@scatter_elements_strategy.register(["cuda", "gpu", "maca"]) def scatter_elements_cuda(attrs, inputs, out_type, target): """scatter elements cuda strategy""" strategy = _op.OpStrategy() @@ -1077,7 +1077,7 @@ def scatter_elements_cuda(attrs, inputs, out_type, target): return strategy -@scatter_nd_strategy.register(["cuda", "gpu"]) +@scatter_nd_strategy.register(["cuda", "gpu", "maca"]) def scatter_nd_cuda(attrs, inputs, out_type, target): """scatter_nd cuda strategy""" strategy = _op.OpStrategy() @@ -1090,7 +1090,7 @@ def scatter_nd_cuda(attrs, inputs, out_type, target): return strategy -@sort_strategy.register(["cuda", "gpu"]) +@sort_strategy.register(["cuda", "gpu", "maca"]) def sort_strategy_cuda(attrs, inputs, out_type, target): """sort cuda strategy""" strategy = _op.OpStrategy() @@ -1109,7 +1109,7 @@ def sort_strategy_cuda(attrs, inputs, out_type, target): return strategy -@argsort_strategy.register(["cuda", "gpu"]) +@argsort_strategy.register(["cuda", "gpu", "maca"]) def argsort_strategy_cuda(attrs, inputs, out_type, target): """argsort cuda strategy""" strategy = _op.OpStrategy() @@ -1128,7 +1128,7 @@ def argsort_strategy_cuda(attrs, inputs, out_type, target): return strategy -@topk_strategy.register(["cuda", "gpu"]) +@topk_strategy.register(["cuda", "gpu", "maca"]) def topk_strategy_cuda(attrs, inputs, out_type, target): """topk cuda strategy""" strategy = _op.OpStrategy() @@ -1147,7 +1147,7 @@ def topk_strategy_cuda(attrs, inputs, out_type, target): return strategy -@searchsorted_strategy.register(["cuda", "gpu"]) +@searchsorted_strategy.register(["cuda", "gpu", "maca"]) def searchsorted_strategy_cuda(attrs, inputs, out_type, target): """searchsorted cuda strategy""" strategy = _op.OpStrategy() @@ -1159,7 +1159,7 @@ def searchsorted_strategy_cuda(attrs, inputs, out_type, target): return strategy -@multibox_prior_strategy.register(["cuda", "gpu"]) +@multibox_prior_strategy.register(["cuda", "gpu", "maca"]) def multibox_prior_strategy_cuda(attrs, inputs, out_type, target): """multibox_prior cuda strategy""" strategy = _op.OpStrategy() @@ -1171,7 +1171,7 @@ def multibox_prior_strategy_cuda(attrs, inputs, out_type, target): return strategy -@multibox_transform_loc_strategy.register(["cuda", "gpu"]) +@multibox_transform_loc_strategy.register(["cuda", "gpu", "maca"]) def multibox_transform_loc_strategy_cuda(attrs, inputs, out_type, target): """multibox_transform_loc cuda strategy""" strategy = _op.OpStrategy() @@ -1183,7 +1183,7 @@ def multibox_transform_loc_strategy_cuda(attrs, inputs, out_type, target): return strategy -@get_valid_counts_strategy.register(["cuda", "gpu"]) +@get_valid_counts_strategy.register(["cuda", "gpu", "maca"]) def get_valid_counts_strategy_cuda(attrs, inputs, out_type, target): """get_valid_counts cuda strategy""" strategy = _op.OpStrategy() @@ -1195,7 +1195,7 @@ def get_valid_counts_strategy_cuda(attrs, inputs, out_type, target): return strategy -@nms_strategy.register(["cuda", "gpu"]) +@nms_strategy.register(["cuda", "gpu", "maca"]) def nms_strategy_cuda(attrs, inputs, out_type, target): """nms cuda strategy""" strategy = _op.OpStrategy() @@ -1207,7 +1207,7 @@ def nms_strategy_cuda(attrs, inputs, out_type, target): return strategy -@all_class_nms_strategy.register(["cuda", "gpu"]) +@all_class_nms_strategy.register(["cuda", "gpu", "maca"]) def all_class_nms_strategy_cuda(attrs, inputs, out_type, target): """all class nms cuda strategy""" strategy = _op.OpStrategy() @@ -1219,7 +1219,7 @@ def all_class_nms_strategy_cuda(attrs, inputs, out_type, target): return strategy -@roi_align_strategy.register(["cuda", "gpu"]) +@roi_align_strategy.register(["cuda", "gpu", "maca"]) def roi_align_strategy_cuda(attrs, inputs, out_type, target): """roi_align cuda strategy""" strategy = _op.OpStrategy() @@ -1241,14 +1241,14 @@ def roi_align_strategy_cuda(attrs, inputs, out_type, target): return strategy -@schedule_roi_pool.register(["cuda", "gpu"]) +@schedule_roi_pool.register(["cuda", "gpu", "maca"]) def schedule_roi_pool_cuda(attrs, outs, target): """schedule roi_pool for cuda""" with target: return topi.cuda.schedule_roi_pool(outs) -@proposal_strategy.register(["cuda", "gpu"]) +@proposal_strategy.register(["cuda", "gpu", "maca"]) def proposal_strategy_cuda(attrs, inputs, out_type, target): """proposal cuda strategy""" strategy = _op.OpStrategy() @@ -1260,7 +1260,7 @@ def proposal_strategy_cuda(attrs, inputs, out_type, target): return strategy -@correlation_strategy.register(["cuda", "gpu"]) +@correlation_strategy.register(["cuda", "gpu", "maca"]) def correlation_strategy_cuda(attrs, inputs, out_type, target): """correlation cuda strategy""" layout = attrs.layout @@ -1274,7 +1274,7 @@ def correlation_strategy_cuda(attrs, inputs, out_type, target): return strategy -@argwhere_strategy.register(["cuda", "gpu"]) +@argwhere_strategy.register(["cuda", "gpu", "maca"]) def argwhere_strategy_cuda(attrs, inputs, out_type, target): """argwhere cuda strategy""" strategy = _op.OpStrategy() @@ -1286,7 +1286,7 @@ def argwhere_strategy_cuda(attrs, inputs, out_type, target): return strategy -@cumsum_strategy.register(["cuda", "gpu"]) +@cumsum_strategy.register(["cuda", "gpu", "maca"]) def cumsum_strategy_cuda(attrs, inputs, out_type, target): """cumsum cuda strategy""" strategy = _op.OpStrategy() @@ -1298,7 +1298,7 @@ def cumsum_strategy_cuda(attrs, inputs, out_type, target): return strategy -@cumprod_strategy.register(["cuda", "gpu"]) +@cumprod_strategy.register(["cuda", "gpu", "maca"]) def cumprod_strategy_cuda(attrs, inputs, out_type, target): """cumprod cuda strategy""" strategy = _op.OpStrategy() @@ -1310,7 +1310,7 @@ def cumprod_strategy_cuda(attrs, inputs, out_type, target): return strategy -@unique_strategy.register(["cuda", "gpu"]) +@unique_strategy.register(["cuda", "gpu", "maca"]) def unique_strategy_cuda(attrs, inputs, out_type, target): """unique cuda strategy""" strategy = _op.OpStrategy() @@ -1322,7 +1322,7 @@ def unique_strategy_cuda(attrs, inputs, out_type, target): return strategy -@schedule_transpose.register(["cuda", "gpu", "rocm"]) +@schedule_transpose.register(["cuda", "gpu", "rocm", "maca"]) def schedule_transpose_cuda(attrs, outs, target): """ Transpose cuda strategy @@ -1342,7 +1342,7 @@ def schedule_transpose_cuda(attrs, outs, target): return schedule_injective(attrs, outs, target) -@invert_permutation_strategy.register(["cuda", "gpu"]) +@invert_permutation_strategy.register(["cuda", "gpu", "maca"]) def invert_permutation_strategy_cuda(attrs, inputs, out_type, target): """invert_permutation cuda strategy""" strategy = _op.OpStrategy() @@ -1354,7 +1354,7 @@ def invert_permutation_strategy_cuda(attrs, inputs, out_type, target): return strategy -@einsum_strategy.register(["cuda", "gpu"]) +@einsum_strategy.register(["cuda", "gpu", "maca"]) def einsum_strategy_cuda(attrs, inputs, out_type, target): """einsum cuda strategy""" strategy = _op.OpStrategy() @@ -1367,7 +1367,7 @@ def einsum_strategy_cuda(attrs, inputs, out_type, target): return strategy -@stft_strategy.register(["cuda", "gpu"]) +@stft_strategy.register(["cuda", "gpu", "maca"]) def stft_strategy_cuda(attrs, inputs, out_type, target): strategy = _op.OpStrategy() strategy.add_implementation( @@ -1378,7 +1378,7 @@ def stft_strategy_cuda(attrs, inputs, out_type, target): return strategy -@dft_strategy.register(["cuda", "gpu"]) +@dft_strategy.register(["cuda", "gpu", "maca"]) def dft_strategy_cuda(attrs, inputs, out_type, target): strategy = _op.OpStrategy() strategy.add_implementation( @@ -1389,7 +1389,7 @@ def dft_strategy_cuda(attrs, inputs, out_type, target): return strategy -@layout_transform_strategy.register(["cuda", "gpu"]) +@layout_transform_strategy.register(["cuda", "gpu", "maca"]) def layout_transform_strategy_cuda(attrs, inputs, out_type, target): strategy = _op.OpStrategy() strategy.add_implementation( diff --git a/python/tvm/relay/op/strategy/maca.py b/python/tvm/relay/op/strategy/maca.py new file mode 100644 index 000000000000..b6937ee72e16 --- /dev/null +++ b/python/tvm/relay/op/strategy/maca.py @@ -0,0 +1,511 @@ +# 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. +"""Definition of MACA/GPU operator strategy.""" +# pylint: disable=invalid-name,unused-argument,wildcard-import,unused-wildcard-import +from tvm import topi +from tvm.auto_scheduler import is_auto_scheduler_enabled +from tvm.contrib import mxcc, nvcc +# from tvm.contrib.thrust import can_use_thrust +from tvm.meta_schedule import is_meta_schedule_enabled +from tvm.te import SpecializedCondition +from .cuda import batch_matmul_strategy_cuda, conv2d_strategy_cuda, dense_strategy_cuda, judge_winograd +from ....target import Target +from ....tir import IntImm +from .. import op as _op +from .generic import * + +@matmul_strategy.register(["maca"]) +def matmul_strategy_maca(attrs, inputs, out_type, target): + """Matmul maca strategy.""" + strategy = _op.OpStrategy() + + if is_auto_scheduler_enabled(): + strategy.add_implementation( + wrap_compute_matmul(topi.nn.matmul), naive_schedule, name="matmul.maca" + ) + elif is_meta_schedule_enabled(): + strategy.add_implementation( + wrap_compute_matmul(topi.nn.matmul), naive_schedule, name="matmul.maca" + ) + else: + logger.warning( + "Matmul is not optimized for maca. Recommend to use mcblas for better performance." + ) + # Temporary use this as a basic schedule + strategy.add_implementation( + wrap_compute_matmul(topi.gpu.matmul_default), + wrap_topi_schedule(topi.gpu.schedule_matmul_default), + name="matmul_default.gpu", + ) + + if target.kind.name == "maca" and "mcblas" in target.libs: + strategy.add_implementation( + wrap_compute_matmul(topi.maca.matmul_mcblas), + wrap_topi_schedule(topi.maca.schedule_matmul_mcblas), + name="matmul_mcblas.maca", + plevel=25, + ) + return strategy + +@dense_strategy.register(["maca"]) +def dense_strategy_maca(attrs, inputs, out_type, target): + """dense maca strategy""" + strategy = _op.OpStrategy() + data, weights = inputs + b, i = get_const_tuple(data.shape) + o, _ = get_const_tuple(weights.shape) + if ( + target.kind.name in ["maca"] + and data.dtype == "int8" + and weights.dtype == "int8" + and out_type.dtype == "int32" + ): + strategy.add_implementation( + wrap_compute_dense(topi.cuda.dense_int8), + wrap_topi_schedule(topi.cuda.schedule_dense_int8), + name="dense_int8.maca", + ) + + if target.kind.name == "maca" and "mcblas" in target.libs: + strategy.add_implementation( + wrap_compute_dense(topi.maca.dense_mcblas), + wrap_topi_schedule(topi.maca.schedule_dense_mcblas), + name="dense_mcblas.maca", + plevel=25, + ) + return strategy + +@batch_matmul_strategy.register(["maca"]) +def batch_matmul_strategy_maca(attrs, inputs, out_type, target): + """batch_matmul maca strategy""" + strategy = _op.OpStrategy() + x, y = inputs + + if target.kind.name == "maca" and "mcblas" in target.libs: + strategy.add_implementation( + wrap_compute_batch_matmul(topi.maca.batch_matmul_mcblas, need_out_dtype=True), + wrap_topi_schedule(topi.generic.schedule_extern), + name="batch_matmul_mcblas.maca", + plevel=30, + ) + + return strategy + +@conv2d_strategy.register(["maca"]) +def conv2d_strategy_maca(attrs, inputs, out_type, target): + """conv2d maca strategy""" + strategy = _op.OpStrategy() + data, kernel = inputs + stride_h, stride_w = attrs.get_int_tuple("strides") + dilation_h, dilation_w = attrs.get_int_tuple("dilation") + padding = attrs.get_int_tuple("padding") + groups = attrs.groups + layout = attrs.data_layout + kernel_layout = attrs.kernel_layout + if dilation_h < 1 or dilation_w < 1: + raise ValueError("dilation should be positive value") + if groups == 1: + if layout == "NCHW": + assert kernel_layout == "OIHW" + if ( + (target.kind.name in ["cuda", "vulkan", "rocm"]) + and data.dtype in ("int8", "uint8") + and kernel.dtype in ("int8", "uint8") + ): + assert data.dtype == kernel.dtype + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_nchw_int8), + wrap_topi_schedule(topi.cuda.schedule_conv2d_nchw_int8), + name="conv2d_nchw_int8.cuda", + ) + else: + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_nchw), + wrap_topi_schedule(topi.cuda.schedule_conv2d_nchw), + name="conv2d_nchw.cuda", + ) + N, _, H, W = get_const_tuple(data.shape) + CO, CI, KH, KW = get_const_tuple(kernel.shape) + (_, _, judge_winograd_auto_scheduler) = judge_winograd( + N, + H, + W, + KH, + KW, + CI, + CO, + padding, + stride_h, + stride_w, + dilation_h, + dilation_w, + data.dtype, + kernel.dtype, + pre_flag=False, + ) + if is_meta_schedule_enabled() and judge_winograd_auto_scheduler: + strategy.add_implementation( + wrap_compute_conv2d(topi.nn.conv2d_winograd_nchw), + naive_schedule, # this implementation should never be picked by autotvm + name="conv2d_nchw_winograd.cuda", + plevel=15, + ) + elif ( + (2 < KH < 8 and 2 < KW < 8 and KH == KW) + and (stride_h == 1 and stride_w == 1) + and (dilation_h == 1 and dilation_w == 1) + ): + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_nchw_winograd), + wrap_topi_schedule(topi.cuda.schedule_conv2d_nchw_winograd), + name="conv2d_nchw_winograd.cuda", + plevel=5, + ) + elif layout == "HWCN": + assert kernel_layout == "HWIO" + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_hwcn), + wrap_topi_schedule(topi.cuda.schedule_conv2d_hwcn), + name="conv2d_hwcn.cuda", + ) + elif layout == "NHWC" and kernel_layout == "HWIO": + strategy.add_implementation( + wrap_compute_conv2d(topi.gpu.conv2d_nhwc), + wrap_topi_schedule(topi.gpu.schedule_conv2d_nhwc), + name="conv2d_nhwc.gpu", + ) + + N, H, W, _ = get_const_tuple(data.shape) + KH, KW, CI, CO = get_const_tuple(kernel.shape) + # Winograd shape related judgment + ( + judge_winograd_tensorcore, + judge_winograd_autotvm, + judge_winograd_auto_scheduler, + ) = judge_winograd( + N, + H, + W, + KH, + KW, + CI, + CO, + padding, + stride_h, + stride_w, + dilation_h, + dilation_w, + data.dtype, + kernel.dtype, + pre_flag=False, + ) + if judge_winograd_autotvm: + if ( + target.kind.name == "cuda" + and nvcc.have_tensorcore(target=target) + and judge_winograd_tensorcore + ): + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_nhwc_winograd_tensorcore), + wrap_topi_schedule(topi.cuda.schedule_conv2d_nhwc_winograd_tensorcore), + name="conv2d_nhwc_winograd_tensorcore.cuda", + plevel=5, + ) + else: + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_nhwc_winograd_direct), + wrap_topi_schedule(topi.cuda.schedule_conv2d_nhwc_winograd_direct), + name="conv2d_nhwc_winograd_direct.cuda", + plevel=5, + ) + if ( + target.kind.name in ["maca"] + and not is_auto_scheduler_enabled() + and not is_meta_schedule_enabled() + and mxcc.have_matrixcore() + and ( + (N % 16 == 0 and CI % 16 == 0 and CO % 16 == 0) + or (N % 8 == 0 and CI % 16 == 0 and CO % 32 == 0) + or (N % 32 == 0 and CI % 16 == 0 and CO % 8 == 0) + ) + ): + strategy.add_implementation( + wrap_compute_conv2d(topi.maca.conv2d_nhwc_tensorcore), + wrap_topi_schedule(topi.maca.schedule_conv2d_nhwc_tensorcore), + name="conv2d_nhwc_tensorcore.maca", + plevel=20, + ) + + # register auto-scheduler implementations + if is_auto_scheduler_enabled() and judge_winograd_auto_scheduler: + strategy.add_implementation( + wrap_compute_conv2d(topi.nn.conv2d_winograd_nhwc), + naive_schedule, # this implementation should never be picked by autotvm + name="conv2d_nhwc.winograd", + plevel=15, + ) + # register meta-schedule implementations + if is_meta_schedule_enabled() and judge_winograd_auto_scheduler: + strategy.add_implementation( + wrap_compute_conv2d(topi.nn.conv2d_winograd_nhwc), + naive_schedule, # this implementation should never be picked by autotvm + name="conv2d_nhwc.winograd", + plevel=15, + ) + + elif layout == "HWNC": + assert kernel_layout in ["HWOI", "HWOI16o16i", "HWOI8o32i", "HWOI32o16i"] + _, _, N, in_channels = get_const_tuple(data.shape) + pre_computed = len(kernel.shape) == 6 + if pre_computed: + _, _, oc_chunk, _, oc_block_factor, _ = get_const_tuple(kernel.shape) + out_channels = oc_chunk * oc_block_factor + else: + _, _, out_channels, _ = get_const_tuple(kernel.shape) + + tensorcore_dtypes = ["int4", "uint4", "int8", "uint8"] + if ( + target.kind.name == "cuda" + and nvcc.have_tensorcore(target=target) + and kernel.dtype in tensorcore_dtypes + and ( + ( + data.dtype in ["int4", "uint4"] + and N % 8 == 0 + and in_channels % 32 == 0 + and out_channels % 8 == 0 + ) + or ( + data.dtype in ["int8", "uint8"] + and N % 8 == 0 + and in_channels % 16 == 0 + and out_channels % 32 == 0 + ) + ) + ): + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_hwnc_tensorcore), + wrap_topi_schedule(topi.cuda.schedule_conv2d_hwnc_tensorcore), + name="conv2d_hwnc_tensorcore_direct.cuda", + plevel=20, + ) + else: + raise RuntimeError( + "Unsupported shape for conv2d HWNC.\ + Need to satisfy tensor core schedule." + ) + elif ( + (target.kind.name in ["cuda", "vulkan", "rocm"]) + and layout == "NCHW4c" + and data.dtype in ["int8", "uint8"] + ): + assert kernel_layout == "OIHW4o4i" + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.conv2d_NCHWc_int8, need_data_layout=True), + wrap_topi_schedule(topi.cuda.schedule_conv2d_NCHWc_int8), + name="conv2d_NCHWc_int8.cuda", + ) + elif is_auto_scheduler_enabled() or is_meta_schedule_enabled(): + strategy.add_implementation( + wrap_compute_conv2d( + topi.nn.conv, need_data_layout=True, need_kernel_layout=True, has_groups=True + ), + naive_schedule, + name="conv2d.cuda", + plevel=15, + ) + elif target.kind.name == "maca" and "mcdnn" not in target.libs: + # No TVM native kernel applicable + raise RuntimeError(f"Unsupported conv2d layout {layout} for MACA") + + if ( + target.kind.name == "maca" + and "mcdnn" in target.libs + and layout in ["NCHW", "NHWC"] + and padding[0] == padding[2] + and padding[1] == padding[3] + and not (data.dtype in ["uint8", "int8"] or kernel.dtype in ["uint8", "int8"]) + ): + # add mcdnn implementation + if layout == "NHWC": + assert kernel_layout == "OHWI" + strategy.add_implementation( + wrap_compute_conv2d(topi.maca.conv2d_mcdnn, need_data_layout=True, has_groups=True), + wrap_topi_schedule(topi.maca.schedule_conv2d_mcdnn), + name="conv2d_mcdnn.maca", + plevel=25, + ) + + elif is_depthwise_conv2d(data.shape, layout, kernel.shape, kernel_layout, groups) and ( + layout == "NCHW" or "mcdnn" not in target.libs + ): # mcDNN requires a different kernel layout for NHWC inputs. + if layout == "NCHW": + assert kernel_layout == "OIHW" + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.depthwise_conv2d_nchw), + wrap_topi_schedule(topi.cuda.schedule_depthwise_conv2d_nchw), + name="depthwise_conv2d_nchw.cuda", + ) + elif layout == "NHWC": + assert kernel_layout == "HWOI" + strategy.add_implementation( + wrap_compute_conv2d(topi.nn.depthwise_conv2d_nhwc), + wrap_topi_schedule(topi.cuda.schedule_depthwise_conv2d_nhwc), + name="depthwise_conv2d_nhwc.cuda", + ) + else: + raise RuntimeError(f"Unsupported depthwise_conv2d layout {layout}") + else: # group_conv2d + # add mcdnn implementation, if any + mcdnn_impl = False + if target.kind.name == "maca" and "mcdnn" in target.libs: + if ( + layout in ["NCHW", "NHWC"] + and padding[0] == padding[2] + and padding[1] == padding[3] + and not (data.dtype in ["uint8", "int8"] or kernel.dtype in ["uint8", "int8"]) + ): + strategy.add_implementation( + wrap_compute_conv2d( + topi.maca.conv2d_mcdnn, need_data_layout=True, has_groups=True + ), + wrap_topi_schedule(topi.maca.schedule_conv2d_mcdnn), + name="conv2d_mcdnn.maca", + plevel=25, + ) + mcdnn_impl = True + + if layout == "NCHW": + assert kernel_layout == "OIHW" + _, channels, _, _ = get_const_tuple(data.shape) + out_channels, in_channels, _, _ = get_const_tuple(kernel.shape) + oc_chunk = out_channels // 4 + ic_chunk = in_channels // 4 + + if ( + (target.kind.name in ["cuda", "vulkan", "rocm"]) + and data.dtype in ["int8", "uint8"] + and kernel.dtype in ["int8", "uint8"] + and channels % groups == 0 + and out_channels % groups == 0 + and channels % 4 == 0 + and out_channels % 4 == 0 + and groups <= oc_chunk + and groups <= ic_chunk + ): + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.group_conv2d_nchw_int8, has_groups=True), + wrap_topi_schedule(topi.cuda.schedule_group_conv2d_nchw_int8), + name="group_conv2d_nchw_int8.cuda", + ) + else: + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.group_conv2d_nchw, has_groups=True), + wrap_topi_schedule(topi.cuda.schedule_group_conv2d_nchw), + name="group_conv2d_nchw.cuda", + ) + elif layout == "NCHW4c" and data.dtype in ["int8", "uint8"]: + assert kernel_layout == "OIHW4o4i" + strategy.add_implementation( + wrap_compute_conv2d(topi.cuda.group_conv2d_NCHWc_int8, has_groups=True), + wrap_topi_schedule(topi.cuda.schedule_group_conv2d_NCHWc_int8), + name="group_conv2d_NCHWc_int8.cuda", + ) + elif not mcdnn_impl: + raise RuntimeError(f"Unsupported group_conv2d layout {layout}") + return strategy + +@softmax_strategy.register(["maca"]) +def softmax_strategy_maca(attrs, inputs, out_type, target): + """softmax maca strategy""" + strategy = _op.OpStrategy() + if target.kind.name == "maca" and "mcdnn" in target.libs: + strategy.add_implementation( + wrap_compute_softmax(topi.maca.softmax_mcdnn), + wrap_topi_schedule(topi.maca.schedule_softmax_mcdnn), + name="softmax.mcdnn", + plevel=15, + ) + return strategy + +@log_softmax_strategy.register(["maca"]) +def log_softmax_strategy_maca(attrs, inputs, out_type, target): + """log_softmax maca strategy""" + strategy = _op.OpStrategy() + if target.kind.name == "maca" and "mcdnn" in target.libs: + strategy.add_implementation( + wrap_compute_softmax(topi.maca.log_softmax_mcdnn), + wrap_topi_schedule(topi.maca.schedule_log_softmax_mcdnn), + name="log_softmax.mcdnn", + plevel=15, + ) + return strategy + +@conv2d_transpose_strategy.register(["maca"]) +def conv2d_transpose_strategy_maca(attrs, inputs, out_type, target): + """conv2d_transpose maca strategy""" + layout = attrs.data_layout + dilation = get_const_tuple(attrs.dilation) + groups = attrs.groups + assert dilation == (1, 1), "not support dilate now" + strategy = _op.OpStrategy() + num_strategies = 0 + + if layout == "NCHW": + strategy.add_implementation( + wrap_compute_conv2d_transpose(topi.cuda.conv2d_transpose_nchw, has_groups=True), + wrap_topi_schedule(topi.cuda.schedule_conv2d_transpose_nchw), + name="conv2d_transpose_nchw.cuda", + ) + num_strategies += 1 + + if ( + target.kind.name == "maca" + and "mcdnn" in target.libs + and ( + (layout == "NCHW" and attrs.kernel_layout == "IOHW") + or (layout == "NHWC" and attrs.kernel_layout == "IHWO") + ) + ): + strategy.add_implementation( + wrap_compute_conv2d_transpose( + topi.maca.conv2d_transpose_mcdnn, add_layout=True, has_groups=True + ), + wrap_topi_schedule(topi.generic.schedule_extern), + name="conv2d_transpose.mcdnn.maca", + plevel=25, + ) + num_strategies += 1 + + # TODO(masahi): Support conv2d_transpose NHWC for non-cudnn path. + assert ( + num_strategies > 0 + ), f"Unsupported conv2d_transpose workload, layout = {layout}, groups = {groups}" + return strategy + +@conv3d_strategy.register(["maca"]) +def conv3d_strategy_maca(attrs, inputs, out_type, target): + """conv3d maca strategy""" + strategy = _op.OpStrategy() + if target.kind.name == "maca" and "mcdnn" in target.libs: + strategy.add_implementation( + wrap_compute_conv3d(topi.maca.conv3d_mcdnn, True), + wrap_topi_schedule(topi.maca.schedule_conv3d_mcdnn), + name="conv3d_mcdnn.maca", + plevel=25, + ) + return strategy \ No newline at end of file diff --git a/python/tvm/relay/transform/mixed_precision.py b/python/tvm/relay/transform/mixed_precision.py index f6bb8b815085..97c866019df8 100644 --- a/python/tvm/relay/transform/mixed_precision.py +++ b/python/tvm/relay/transform/mixed_precision.py @@ -98,7 +98,6 @@ "nn.prelu", "nn.dropout", # Complicated activations which saturate in a narrow range - "sigmoid", "tanh", "fast_tanh", # Some coefficients outside of representable range, but probably ok "fast_exp", @@ -118,14 +117,16 @@ "nn.adaptive_max_pool2d", "nn.adaptive_max_pool3d", "image.resize2d", + "nn.softmax", + "sum", ] DEFAULT_NEVER_LIST = [ # In general if |f(x)| >> |x| for expected inputs then put the op here. + "sigmoid", "exp", "power", "nn.cross_entropy", "nn.cross_entropy_with_logits", - "nn.softmax", "nn.l2_normalize", # Error function doesn't seem to be able to be lowered into fp16 version in llvm. # Move to follow list when it does. @@ -138,7 +139,6 @@ "nn.adaptive_avg_pool1d", "nn.adaptive_avg_pool2d", "nn.adaptive_avg_pool3d", - "sum", "mean", "variance", "nn.layer_norm", diff --git a/python/tvm/testing/plugin.py b/python/tvm/testing/plugin.py index c9aaab57d2f5..43a71dc05e9e 100644 --- a/python/tvm/testing/plugin.py +++ b/python/tvm/testing/plugin.py @@ -296,6 +296,8 @@ def _target_to_requirement(target): return utils.requires_rocm.marks() if target.kind.name == "maca": return utils.requires_maca.marks() + if target.kind.name == "maca" and "mcblas" in target.attrs.get("libs", []): + return utils.requires_mcblas.marks() if target.kind.name == "vulkan": return utils.requires_vulkan.marks() if target.kind.name == "nvptx": diff --git a/python/tvm/testing/utils.py b/python/tvm/testing/utils.py index c761dc17d187..9a06259236fd 100644 --- a/python/tvm/testing/utils.py +++ b/python/tvm/testing/utils.py @@ -953,6 +953,12 @@ def _multi_gpu_exists(): parent_features="gpu", ) +# Mark a test as requiring the mcBLAS library. +requires_mcblas = Feature("mcblas", "mcBLAS", cmake_flag="USE_MCBLAS", parent_features="maca") + +# Mark a test as requiring the mcDNN library. +requires_mcdnn = Feature("mcdnn", "mcDNN", cmake_flag="USE_MCDNN", parent_features="maca") + # Mark a test as requiring a matrixcore to run requires_matrixcore = Feature( "matrixcore", diff --git a/python/tvm/topi/__init__.py b/python/tvm/topi/__init__.py index fc316fd19307..5704c9a8d5b2 100644 --- a/python/tvm/topi/__init__.py +++ b/python/tvm/topi/__init__.py @@ -65,7 +65,7 @@ from . import random from . import hexagon from . import adreno - +from . import maca # error reporting from .utils import InvalidShapeError diff --git a/python/tvm/topi/maca/__init__.py b/python/tvm/topi/maca/__init__.py new file mode 100644 index 000000000000..54a7eef40085 --- /dev/null +++ b/python/tvm/topi/maca/__init__.py @@ -0,0 +1,22 @@ +# 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. + +# pylint: disable=redefined-builtin, wildcard-import +"""MACA specific declaration and schedules.""" +from .conv2d_nhwc_tensorcore import * +from .conv2d_hwnc_tensorcore import * +from .conv2d import * diff --git a/python/tvm/topi/maca/batch_matmul.py b/python/tvm/topi/maca/batch_matmul.py new file mode 100644 index 000000000000..f6439faf303c --- /dev/null +++ b/python/tvm/topi/maca/batch_matmul.py @@ -0,0 +1,82 @@ +# 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. +# pylint: disable=invalid-name,too-many-locals,unused-variable,unused-argument +"""maca batch_matmul operators""" +import tvm +from tvm import autotvm +from tvm import te +from tvm.contrib import mcblas +from tvm.autotvm.task.space import SplitEntity, OtherOptionEntity +from .. import nn, generic +from ..utils import traverse_inline, get_const_tuple, get_max_power2_factor +from .tensor_intrin import dp4a + +@autotvm.register_topi_compute("batch_matmul_mcblas.maca") +def batch_matmul_mcblas( + cfg, x, y, out_shape=None, out_dtype=None, transpose_a=False, transpose_b=True +): + """Compute batch matrix multiplication of `x` and `y`. + + Both `x` and `y` can be transposed. For legacy reason, we use NT format + (transpose_a=False, transpose_b=True) by default. + + Parameters + ---------- + cfg : ConfigSpace + Autotvm tuning space config file. + + x : tvm.te.Tensor + 3-D with shape [batch, M, K] or [batch, K, M]. + + y : tvm.te.Tensor + 3-D with shape [batch, K, N] or [batch, N, K]. + + out_shape : List[Optional] + Explicit intended output shape of the computation. Can be useful in cases + with dynamic input shapes. + + out_dtype : Optional[str] + Specifies the output data type for mixed precision batch matmul. + + transpose_a : Optional[bool] = False + Whether the first tensor is in transposed format. + + transpose_b : Optional[bool] = True + Whether the second tensor is in transposed format. + + Returns + ------- + output : tvm.te.Tensor + 3-D with shape [batch, M, N] + """ + if transpose_a: + b, k, m = get_const_tuple(x.shape) + else: + b, m, k = get_const_tuple(x.shape) + if transpose_b: + b, n, k = get_const_tuple(y.shape) + else: + b, k, n = get_const_tuple(y.shape) + if all([isinstance(s, int) for s in [b, m, n, k]]): + cfg.add_flop(b * m * k * n * 2) + return mcblas.batch_matmul(x, y, transa=transpose_a, transb=transpose_b, dtype=out_dtype) + + +@autotvm.register_topi_schedule("batch_matmul_mcblas.maca") +def schedule_batch_matmul_mcblas(_, outs): + """Schedule batch_matmul operator using MCBLAS""" + return generic.schedule_extern(outs) diff --git a/python/tvm/topi/maca/conv2d.py b/python/tvm/topi/maca/conv2d.py new file mode 100644 index 000000000000..df11585aab05 --- /dev/null +++ b/python/tvm/topi/maca/conv2d.py @@ -0,0 +1,150 @@ +# 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. +# pylint: disable=invalid-name, unused-argument +"""Compute definition for conv2d with maca backend""" +from tvm import te +from tvm import autotvm +from tvm.autotvm.task.space import OtherOptionEntity +from tvm.contrib import mcdnn + +from .. import nn, generic +from ..nn.utils import get_pad_tuple +from ..utils import get_const_tuple, traverse_inline +from ..cuda.conv2d_direct import schedule_direct_cuda + + +@autotvm.register_topi_compute("conv2d_nchw.maca") +def conv2d_nchw(cfg, data, kernel, strides, padding, dilation, out_dtype="float32"): + """Compute conv2d with NCHW layout""" + return nn.conv2d_nchw(data, kernel, strides, padding, dilation, out_dtype) + + +@autotvm.register_topi_schedule("conv2d_nchw.maca") +def schedule_conv2d_nchw(cfg, outs): + """Create the schedule for conv2d_nchw""" + outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs + s = te.create_schedule([x.op for x in outs]) + + def _callback(op): + if op.tag == "conv2d_nchw": + schedule_direct_cuda(cfg, s, op.output(0)) + + traverse_inline(s, outs[0].op, _callback) + return s + + +@autotvm.register_topi_compute("conv2d_mcdnn.maca") +def conv2d_mcdnn( + cfg, data, kernel, strides, padding, dilation, groups=1, layout="NCHW", out_dtype="float32" +): + """Compute conv2d using McDNN library""" + if layout == "NCHW": + tensor_format = 0 # MCDNN_TENSOR_NCHW + N, _, H, W = get_const_tuple(data.shape) + elif layout == "NHWC": + tensor_format = 1 # MCDNN_TENSOR_NHWC + N, H, W, _ = get_const_tuple(data.shape) + else: + raise ValueError(f"Unsupported layout {layout} in mcdnn") + CO, CI, KH, KW = get_const_tuple(kernel.shape) + + # handle dilation + stride_h, stride_w = (strides, strides) if isinstance(strides, int) else strides + dilation_h, dilation_w = (dilation, dilation) if isinstance(dilation, int) else dilation + KH_dilated = (KH - 1) * dilation_h + 1 + KW_dilated = (KW - 1) * dilation_h + 1 + + pt, pl, pb, pr = get_pad_tuple(padding, (KH_dilated, KW_dilated)) + if (pt != pb) or (pl != pr): + raise ValueError("Mcdnn doesn't support asymmetric padding.") + + OH = (H + pt + pb - KH) // stride_h + 1 + OW = (W + pl + pr - KW) // stride_w + 1 + + if isinstance(N, int): + cfg.add_flop( + groups + * 2 + * N + * OH + * OW + * CO + * CI + * ((KH - 1) * dilation_h + 1) + * ((KW - 1) * dilation_w + 1) + ) + + if data.dtype == "int8" or kernel.dtype == "int8": + if layout == "NCHW": + raise ValueError("NCHW layout do not support int8 in mcdnn") + dtype = "int32" + else: + dtype = data.dtype + + cfg.define_knob("algo", range(mcdnn.algo_to_index("fwd", "MCDNN_CONVOLUTION_FWD_ALGO_COUNT"))) + if cfg.is_fallback: + if mcdnn.exists(): + # Let MCDNN choose the best algo, based on benchmarks run + # on the local machine. In the future, this should be + # based on parameters stored in the Target. + cfg["algo"] = OtherOptionEntity(-1) + else: + cfg["algo"] = OtherOptionEntity(0) + + return mcdnn.conv_forward( + data, + kernel, + [pt, pl], # mcdnn padding pt, pl on both sides of input + [stride_h, stride_w], + [dilation_h, dilation_w], + conv_mode=1, + tensor_format=tensor_format, + algo=cfg["algo"].val, + conv_dtype=dtype, + groups=groups, + ) + + +@autotvm.register_topi_schedule("conv2d_mcdnn.maca") +def schedule_conv2d_mcdnn(cfg, outs): + """Create the schedule for conv2d_mcdnn""" + return generic.schedule_extern(outs) + + +def conv2d_backward_weight_mcdnn( + dy, x, kernel_size, padding, stride, dilation, groups, layout, output_dtype +): + """Compute conv2d wgrad using McDNN library""" + assert layout in ["NCHW", "NHWC"] + + if dy.dtype == "float16": + # mcDNN does not seem to support other combination. + assert output_dtype == "float16", "Only supports fp16 output for mcDNN fp16 wgrad." + + conv_dtype = "float32" # Accumulation is always fp32 + return mcdnn.conv_backward_filter( + dy, + x, + kernel_size, + padding, + stride, + dilation, + conv_mode=1, + tensor_format=0 if layout == "NCHW" else 1, + conv_dtype=conv_dtype, + groups=groups, + ) diff --git a/python/tvm/topi/maca/conv2d_hwnc_tensorcore.py b/python/tvm/topi/maca/conv2d_hwnc_tensorcore.py new file mode 100644 index 000000000000..f4427d7361e8 --- /dev/null +++ b/python/tvm/topi/maca/conv2d_hwnc_tensorcore.py @@ -0,0 +1,428 @@ +# 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. +# pylint: disable=invalid-name, too-many-locals, too-many-function-args +# pylint: disable=too-many-statements, unused-argument, too-many-arguments +"""Tensorcore template for maca backend""" +import tvm +from tvm import te +from tvm import autotvm +from tvm.target import Target +from tvm.topi.cuda.injective import schedule_injective_from_existing +from ..utils import get_const_tuple, traverse_inline, simplify, tag +from ..nn.pad import pad +from ..nn.utils import get_pad_tuple +from .tensor_intrin import intrin_wmma_load_matrix_A +from .tensor_intrin import intrin_wmma_load_matrix_W +from .tensor_intrin import intrin_wmma_store_matrix +from .tensor_intrin import intrin_wmma_gemm + + +def unpack_HWNCnc_to_hwnc(packed_out, out_dtype): + """Unpack conv2d_hwnc output from layout hwncnc to hwnc + + Parameters + ----------- + packed_out : tvm.te.Tensor + The output tensor of conv2d_hwnc. + + out_dtype : str + The output dtype. + + Returns + ------- + unpacked_out : tvm.te.Tensor + The unpacked output tensor in hwnc layout. + """ + H, W, N, O, wmma_m, wmma_n = get_const_tuple(packed_out.shape) + + idxmod = tvm.tir.indexmod + idxdiv = tvm.tir.indexdiv + + oshape = (H, W, N * wmma_m, O * wmma_n) + unpacked_out = te.compute( + oshape, + lambda h, w, n, o: packed_out[ + h, w, idxdiv(n, wmma_m), idxdiv(o, wmma_n), idxmod(n, wmma_m), idxmod(o, wmma_n) + ].astype(out_dtype), + name="output_unpack", + tag=tag.INJECTIVE + ",unpack_hwncc", + ) + return unpacked_out + + +def conv2d_hwnc_tensorcore(data, kernel, strides, padding, dilation, in_dtype, out_dtype="int32"): + """ "Compute conv2d with tensorcore for HWNC layout with int8/int4""" + assert data.dtype in ("int4", "uint4", "int8", "uint8") + assert kernel.dtype in ("int4", "uint4", "int8", "uint8") + packed_out = hwnc_tensorcore_maca(data, kernel, strides, padding, dilation, out_dtype) + return unpack_HWNCnc_to_hwnc(packed_out, out_dtype) + + +@autotvm.register_topi_compute("conv2d_HWNCnc_tensorcore.maca") +def hwnc_tensorcore_maca(cfg, Input, Filter, stride, padding, dilation, out_dtype="int32"): + """Compute declaration for tensorcore""" + assert isinstance(stride, int) or len(stride) == 2 + assert isinstance(dilation, int) or len(dilation) == 2 + + if isinstance(stride, int): + stride_h = stride_w = stride + else: + stride_h, stride_w = stride + + if isinstance(dilation, int): + dilation_h = dilation_w = dilation + else: + dilation_h, dilation_w = dilation + + in_dtype = Input.dtype + + if in_dtype in ["int4", "uint4"]: + wmma_n = wmma_m = 8 + wmma_k = 32 + else: + wmma_m = 8 + wmma_n = 32 + wmma_k = 16 + + pre_computed = len(Filter.shape) == 6 + in_height, in_width, batch, in_channels = get_const_tuple(Input.shape) + if pre_computed: + kernel_h, kernel_w, oc_chunk, _, oc_block_factor, _ = get_const_tuple(Filter.shape) + num_filter = oc_block_factor * oc_chunk + else: + kernel_h, kernel_w, num_filter, _ = get_const_tuple(Filter.shape) + + if in_dtype in ["int4", "uint4"]: + assert batch % 8 == 0 and in_channels % 32 == 0 and num_filter % 8 == 0 + else: + assert batch % 8 == 0 and in_channels % 16 == 0 and num_filter % 32 == 0, ( + "The shape of (batch, in_channels, num_filter) " + "must be multiple of (8, 16, 32) for int8, " + "and (8, 32, 8) for int4" + ) + + # compute the output shape + dilated_kernel_h = (kernel_h - 1) * dilation_h + 1 + dilated_kernel_w = (kernel_w - 1) * dilation_w + 1 + + pad_top, pad_left, pad_down, pad_right = get_pad_tuple( + padding, (dilated_kernel_h, dilated_kernel_w) + ) + + out_channels = num_filter + out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1) + out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1) + + cfg.add_flop( + 2 * batch * out_height * out_width * out_channels * in_channels * kernel_h * kernel_w + ) + + # Input feature map: (H, W, N, IC, n, ic) + data_shape = (in_height, in_width, batch // wmma_m, in_channels // wmma_k, wmma_m, wmma_k) + + # Kernel: (H, W, OC, IC, oc, ic) + kernel_shape = ( + kernel_h, + kernel_w, + out_channels // wmma_n, + in_channels // wmma_k, + wmma_n, + wmma_k, + ) + + # Reduction axes + kh = te.reduce_axis((0, kernel_h), name="kh") + kw = te.reduce_axis((0, kernel_w), name="kw") + ic = te.reduce_axis((0, in_channels // wmma_k), name="ic") + ii = te.reduce_axis((0, wmma_k), name="ii") + + if pre_computed: + packed_kernel = Filter + else: + packed_kernel = te.compute( + kernel_shape, + lambda kh, kw, o, i, oo, ii: Filter[kh, kw, o * wmma_n + oo, i * wmma_k + ii], + name="packed_kernel", + ) + + packed_data = te.compute( + data_shape, lambda h, w, n, i, nn, ii: Input[h, w, n * wmma_m + nn, i * wmma_k + ii] + ) + + pad_before = [pad_top, pad_left, 0, 0, 0, 0] + pad_after = [pad_down, pad_right, 0, 0, 0, 0] + pad_data = pad(packed_data, pad_before, pad_after, name="pad_data") + + Conv = te.compute( + (out_height, out_width, batch // wmma_m, out_channels // wmma_n, wmma_m, wmma_n), + lambda h, w, n, o, nn, oo: te.sum( + ( + pad_data[h * stride_h + kh, w * stride_w + kw, n, ic, nn, ii].astype("int32") + * packed_kernel[kh, kw, o, ic, oo, ii].astype("int32") + ), + axis=[ic, kh, kw, ii], + ), + name="Conv", + tag="conv2d_HWNCnc_tensorcore", + ) + return Conv + + +def schedule_hwnc_tensorcore_maca(cfg, s, Conv): + """Schedule tensorcore template""" + pad_data, packed_kernel = s[Conv].op.input_tensors + ic, kh, kw, ii = s[Conv].op.reduce_axis + packed_data = s[pad_data].op.input_tensors[0] + + block_x = te.thread_axis("blockIdx.x") + block_y = te.thread_axis("blockIdx.y") + block_z = te.thread_axis("blockIdx.z") + thread_x = te.thread_axis("threadIdx.x") + thread_y = te.thread_axis("threadIdx.y") + thread_z = te.thread_axis("threadIdx.z") + + # Designate the memory hierarchy + AS = s.cache_read(pad_data, "shared", [Conv]) + WS = s.cache_read(packed_kernel, "shared", [Conv]) + AF = s.cache_read(AS, "wmma.matrix_a", [Conv]) + WF = s.cache_read(WS, "wmma.matrix_b", [Conv]) + ConvF = s.cache_write(Conv, "wmma.accumulator") + + if Conv.op in s.outputs: + output = Conv + ConvS = s.cache_read(ConvF, "shared", [Conv]) + OL = ConvS + else: + output = s.outputs[0].output(0) + s[Conv].set_scope("shared") + OL = Conv + + out_dtype = Conv.dtype + + if isinstance(packed_kernel.op, te.tensor.ComputeOp) and packed_kernel.name == "packed_kernel": + if autotvm.GLOBAL_SCOPE.in_tuning: + s[packed_kernel].pragma(s[packed_kernel].op.axis[0], "debug_skip_region") + else: + with Target("maca"): + schedule_injective_from_existing(s, packed_kernel) + + if isinstance(pad_data.op, te.tensor.ComputeOp) and "pad" in pad_data.op.tag: + s[pad_data].compute_inline() + data = pad_data.op.input_tensors[0] + + if autotvm.GLOBAL_SCOPE.in_tuning: + # skip this part during tuning to make recrods accurate + # this part will be pre-computed during NNVM's pre-compute optimization pass + s[pad_data].pragma(s[pad_data].op.axis[0], "debug_skip_region") + else: + data = pad_data + s[data].compute_inline() + + data_dtype = data.dtype + kernel_dtype = packed_kernel.dtype + + # Schedule for autotvm + cfg.define_knob("block_row_warps", [1, 2, 4]) + cfg.define_knob("block_col_warps", [1, 2, 4]) + cfg.define_knob("warp_row_tiles", [1, 2, 4, 8, 16]) + cfg.define_knob("warp_col_tiles", [1, 2, 4, 8, 16]) + cfg.define_knob("chunk", [1, 2, 4, 8]) + cfg.define_knob("split_block_k_nums", [1, 2, 4, 8, 16, 32]) + cfg.define_knob("vector_ws", [1, 8]) + cfg.define_knob("vector_as", [1, 8, 16]) + + block_row_warps = cfg["block_row_warps"].val + block_col_warps = cfg["block_col_warps"].val + warp_row_tiles = cfg["warp_row_tiles"].val + warp_col_tiles = cfg["warp_col_tiles"].val + chunk = cfg["chunk"].val + vector_as = cfg["vector_as"].val + vector_ws = cfg["vector_ws"].val + split_block_k_nums = cfg["split_block_k_nums"].val + + s[packed_data].compute_inline() + + if data_dtype in ["int4", "uint4"]: + wmma_m = wmma_n = 8 + wmma_k = 32 + else: + wmma_m = 8 + wmma_n = 32 + wmma_k = 16 + + warp_size = 64 + + # Schedule for output + if len(s[output].op.axis) == 4: + ( + hc, + wc, + nc, + oc, + ) = output.op.axis + nc, nnc = s[output].split(nc, factor=wmma_m) + oc, ooc = s[output].split(oc, factor=wmma_n) + else: + hc, wc, nc, oc, nnc, ooc = output.op.axis + + kernel_scope, hc = s[output].split(hc, nparts=1) + + block_k = s[output].fuse(hc, wc) + block_k, split_block_k = s[output].split(block_k, factor=split_block_k_nums) + nc, nci = s[output].split(nc, factor=warp_row_tiles) + block_i, nc = s[output].split(nc, factor=block_row_warps) + oc, oci = s[output].split(oc, factor=warp_col_tiles) + block_j, oc = s[output].split(oc, factor=block_col_warps) + s[output].reorder(block_k, split_block_k, block_i, block_j, nc, oc, nci, oci, nnc, ooc) + t = s[output].fuse(nnc, ooc) + _, tx = s[output].split(t, factor=warp_size) + s[output].bind(block_k, block_z) + s[output].bind(block_i, block_x) + s[output].bind(block_j, block_y) + s[output].bind(tx, thread_x) + s[output].bind(nc, thread_y) + s[output].bind(oc, thread_z) + + # Schedule wmma store + s[OL].compute_at(s[output], block_j) + hc, wc, nc, oc, nnc, ooc = OL.op.axis + oc, oci = s[OL].split(oc, factor=warp_col_tiles) + _, oc = s[OL].split(oc, factor=block_col_warps) + nc, nci = s[OL].split(nc, factor=warp_row_tiles) + _, nc = s[OL].split(nc, factor=block_row_warps) + s[OL].reorder(nc, oc, nci, oci, nnc, ooc) + s[OL].bind(nc, thread_y) + s[OL].bind(oc, thread_z) + + # Schedule local computation + s[ConvF].compute_at(s[OL], oc) + _, _, n, o, nnf, oof = ConvF.op.axis + ko, ki = s[ConvF].split(ic, factor=chunk) + s[ConvF].reorder(ko, kh, ki, kw, n, o, nnf, oof, ii) + + cfg.define_reorder("reorder_inner", [ko, kh], policy="all") + cfg["reorder_inner"].apply(s, ConvF, [ko, kh]) + cfg["reorder_inner"].apply(s, ConvF, [ki, kw]) + + # Move intermediate computation into each output compute tile + s[AF].compute_at(s[ConvF], kw) + s[WF].compute_at(s[ConvF], kw) + + # Schedule for A's share memory + s[AS].compute_at(s[ConvF], ko) + + _, _, n, _, nn, ii = AS.op.axis + tx, xo = s[AS].split(n, nparts=block_row_warps) + ty, _ = s[AS].split(xo, nparts=block_col_warps) + t = s[AS].fuse(nn, ii) + to, ti = s[AS].split(t, nparts=warp_size) + ti, _t = s[AS].split(ti, factor=vector_as) + s[AS].bind(tx, thread_y) + s[AS].bind(ty, thread_z) + s[AS].bind(to, thread_x) + s[AS].vectorize(_t) + + # Schedule for W's share memory + s[WS].compute_at(s[ConvF], kw) + kh, kw, ic, o, ii, oo = WS.op.axis + tx, xo = s[WS].split(o, nparts=block_row_warps) + ty, _ = s[WS].split(xo, nparts=block_col_warps) + t = s[WS].fuse(ii, oo) + to, ti = s[WS].split(t, nparts=warp_size) + ti, _t = s[WS].split(ti, factor=vector_ws) + s[WS].bind(tx, thread_y) + s[WS].bind(ty, thread_z) + s[WS].bind(to, thread_x) + s[WS].vectorize(ti) + + # double buffer + cfg.define_knob("AS_double_buffer", [0, 1]) + cfg.define_knob("WS_double_buffer", [0, 1]) + if cfg["AS_double_buffer"].val: + s[AS].double_buffer() + if cfg["WS_double_buffer"].val: + s[WS].double_buffer() + + # unroll + cfg.define_knob("auto_unroll_max_step", [0, 512, 1500]) + s[output].pragma(kernel_scope, "auto_unroll_max_step", cfg["auto_unroll_max_step"].val) + s[output].pragma(kernel_scope, "unroll_explicit", False) + + shape = (wmma_m, wmma_n, wmma_k) + + AS_shape = (wmma_m, wmma_k) + AL_shape = (wmma_m, wmma_k) + WS_shape = (wmma_n, wmma_k) + WL_shape = (wmma_n, wmma_k) + CL_shape = (wmma_m, wmma_n) + CS_shape = (wmma_m, wmma_n) + + AL_gemm = te.placeholder(AL_shape, name="A", dtype=data_dtype) + WL_gemm = te.placeholder(WL_shape, name="B", dtype=kernel_dtype) + k_gemm = te.reduce_axis((0, wmma_k), name="k") + CL_compute = te.compute( + CL_shape, + lambda ii, jj: te.sum( + (AL_gemm[ii, k_gemm].astype("int32") * WL_gemm[jj, k_gemm].astype("int32")), axis=k_gemm + ), + name="C", + ) + + AL_strides = [wmma_k, 1] + AS_strides = [wmma_k, 1] + WL_strides = [wmma_k, 1] + WS_strides = [wmma_k, 1] + CL_strides = [wmma_n, 1] + CS_strides = [wmma_n, 1] + + s[AF].tensorize( + AF.op.axis[-2], + intrin_wmma_load_matrix_A( + AL_strides, AS_strides, shape, "row_major", AS_shape, AL_shape, data_dtype + ), + ) + + s[WF].tensorize( + WF.op.axis[-2], + intrin_wmma_load_matrix_W( + WL_strides, WS_strides, shape, "col_major", WS_shape, WL_shape, kernel_dtype + ), + ) + + s[OL].tensorize( + nnc, intrin_wmma_store_matrix(CS_strides, CL_strides, shape, out_dtype, CL_shape, CS_shape) + ) + + s[ConvF].tensorize( + nnf, + intrin_wmma_gemm(AL_gemm, WL_gemm, CL_compute, AL_strides, WL_strides, CL_strides, shape), + ) + + return s + + +@autotvm.register_topi_schedule("conv2d_HWNCnc_tensorcore.maca") +def schedule_conv2d_hwnc_tensorcore(cfg, outs): + """TOPI schedule callback""" + s = te.create_schedule([x.op for x in outs]) + + def _callback(op): + if "conv2d_HWNCnc_tensorcore" in op.tag: + schedule_hwnc_tensorcore_maca(cfg, s, op.output(0)) + + traverse_inline(s, outs[0].op, _callback) + return s diff --git a/python/tvm/topi/maca/conv2d_nhwc_tensorcore.py b/python/tvm/topi/maca/conv2d_nhwc_tensorcore.py new file mode 100644 index 000000000000..2424b0cf9946 --- /dev/null +++ b/python/tvm/topi/maca/conv2d_nhwc_tensorcore.py @@ -0,0 +1,343 @@ +# 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. +# pylint: disable=invalid-name, too-many-locals, too-many-function-args +# pylint: disable=too-many-statements, unused-argument, too-many-arguments +"""Tensorcore template for maca backend""" +import numpy as np +import tvm +from tvm import te +from tvm import autotvm +from ..utils import get_const_tuple, traverse_inline, simplify +from ..nn.pad import pad +from ..nn.utils import get_pad_tuple +from ..cuda.tensor_intrin import intrin_wmma_load_matrix_A +from ..cuda.tensor_intrin import intrin_wmma_load_matrix_W +from ..cuda.tensor_intrin import intrin_wmma_store_matrix +from ..cuda.tensor_intrin import intrin_wmma_gemm + + +def nhwc_tensorcore_maca(cfg, Input, Filter, stride, padding, dilation, out_dtype): + """Compute declaration for tensorcore""" + assert isinstance(stride, int) or len(stride) == 2 + assert isinstance(dilation, int) or len(dilation) == 2 + + if isinstance(stride, int): + stride_h = stride_w = stride + else: + stride_h, stride_w = stride + + if isinstance(dilation, int): + dilation_h = dilation_w = dilation + else: + dilation_h, dilation_w = dilation + + batch, in_height, in_width, in_channel = get_const_tuple(Input.shape) + kernel_h, kernel_w, _, num_filter = get_const_tuple(Filter.shape) + assert ( + (batch % 16 == 0 and in_channel % 16 == 0 and num_filter % 16 == 0) + or (batch % 8 == 0 and in_channel % 16 == 0 and num_filter % 32 == 0) + or (batch % 32 == 0 and in_channel % 16 == 0 and num_filter % 8 == 0) + ), ( + "The shape of (batch, in_channel, num_filter) " + "must be multiple of (16, 16, 16) or (32, 16, 8) or (8, 16, 32) for now" + ) + + # compute the output shape + dilated_kernel_h = (kernel_h - 1) * dilation_h + 1 + dilated_kernel_w = (kernel_w - 1) * dilation_w + 1 + pad_top, pad_left, pad_down, pad_right = get_pad_tuple( + padding, (dilated_kernel_h, dilated_kernel_w) + ) + out_channel = num_filter + out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1) + out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1) + pad_before = [0, pad_top, pad_left, 0] + pad_after = [0, pad_down, pad_right, 0] + PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput") + rc = te.reduce_axis((0, in_channel), name="rc") + ry = te.reduce_axis((0, kernel_h), name="ry") + rx = te.reduce_axis((0, kernel_w), name="rx") + # convert data type of input feature maps and weights + # TODO: add checking here, datatype casting may cause precision loss + TransPaddedInput = te.compute( + PaddedInput.shape, lambda n, h, w, c: PaddedInput[n, h, w, c].astype("float16") + ) + TransFilter = te.compute(Filter.shape, lambda h, w, i, o: Filter[h, w, i, o].astype("float16")) + Output = te.compute( + (batch, out_height, out_width, out_channel), + lambda nn, yy, xx, ff: te.sum( + TransPaddedInput[ + nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rc + ].astype(out_dtype) + * TransFilter[ry, rx, rc, ff].astype(out_dtype), + axis=[ry, rx, rc], + ), + name="Conv2dOutput", + tag="conv2d_nhwc_tensorcore", + ) + return Output + + +def schedule_nhwc_tensorcore_maca(cfg, s, Conv): + """Schedule tensorcore template""" + kh, kw, ic = s[Conv].op.reduce_axis + out_dtype = Conv.dtype + trans_paddata, kernel = s[Conv].op.input_tensors + in_dtype = trans_paddata.dtype + batch, _, _, _ = get_const_tuple(Conv.shape) + _, _, _, out_channels = get_const_tuple(kernel.shape) + paddata = s[trans_paddata].op.input_tensors + + # inline the pad and dtype transform + s[trans_paddata].compute_inline() + s[kernel].compute_inline() + s[paddata[0]].compute_inline() + + # Designate the memory hierarchy + AS = s.cache_read(trans_paddata, "shared", [Conv]) + WS = s.cache_read(kernel, "shared", [Conv]) + AF = s.cache_read(AS, "wmma.matrix_a", [Conv]) + WF = s.cache_read(WS, "wmma.matrix_b", [Conv]) + ConvF = s.cache_write(Conv, "wmma.accumulator") + + if Conv.op in s.outputs: + output = Conv + ConvS = s.cache_read(ConvF, "shared", [Conv]) + OL = ConvS + else: + output = s.outputs[0].output(0) + s[Conv].set_scope("shared") + OL = Conv + + # Schedule for autotvm + cfg.define_knob("block_row_warps", [1, 2, 4]) + cfg.define_knob("block_col_warps", [1, 2, 4]) + cfg.define_knob("warp_row_tiles", [1, 2, 4]) + cfg.define_knob("warp_col_tiles", [1, 2, 4]) + cfg.define_knob("chunk", [1, 2, 4, 8]) + cfg.define_knob("offset", [0, 8]) + cfg.define_knob("vector_width", [1, 2, 4, 8]) + + if batch % 16 == 0 and out_channels % 16 == 0: + cfg.define_knob("wmma_m", [16, 8, 32]) + elif batch % 8 == 0 and out_channels % 32 == 0: + cfg.define_knob("wmma_m", [8, 16, 32]) + elif batch % 32 == 0 and out_channels % 8 == 0: + cfg.define_knob("wmma_m", [32, 16, 8]) + + # fallback support + target = tvm.target.Target.current() + if cfg.is_fallback: + ref_log = autotvm.tophub.load_reference_log( + target.kind.name, target.model, "conv2d_nhwc_tensorcore.cuda" + ) + cfg.fallback_with_reference_log(ref_log) + + block_row_warps = cfg["block_row_warps"].val + block_col_warps = cfg["block_col_warps"].val + warp_row_tiles = cfg["warp_row_tiles"].val + warp_col_tiles = cfg["warp_col_tiles"].val + chunk = cfg["chunk"].val + offset = cfg["offset"].val + wmma_m = cfg["wmma_m"].val + vector_width = cfg["vector_width"].val + + wmma_k = 16 + if wmma_m == 16: + wmma_n = 16 + elif wmma_m == 8: + wmma_n = 32 + elif wmma_m == 32: + wmma_n = 8 + + warp_size = 64 + + block_x = te.thread_axis("blockIdx.x") + block_y = te.thread_axis("blockIdx.y") + block_z = te.thread_axis("blockIdx.z") + thread_x = te.thread_axis("threadIdx.x") + thread_y = te.thread_axis("threadIdx.y") + thread_z = te.thread_axis("threadIdx.z") + + # Define the intrin strides + def get_strides(extents): + return [np.prod(extents[i:]).tolist() for i in range(len(extents))] + + AS_align = chunk * wmma_k + offset + WS_align = warp_col_tiles * block_col_warps * wmma_n + offset + block_factor_n = wmma_m * warp_row_tiles * block_row_warps + block_factor_o = wmma_n * warp_col_tiles * block_col_warps + CS_align = block_factor_o + offset + AS_strides = get_strides([1, 1, AS_align, 1]) + AL_strides = get_strides([1, 1, wmma_k, 1]) + WS_strides = get_strides([WS_align, 1]) + WL_strides = get_strides([wmma_n * warp_col_tiles, 1]) + CL_strides = get_strides([1, 1, wmma_n * warp_col_tiles, 1]) + CS_strides = get_strides([1, 1, CS_align, 1]) + + # Schedule for output + nc, hc, wc, oc = output.op.axis + block_k = s[output].fuse(hc, wc) + s[output].bind(block_k, block_z) + block_i, nc = s[output].split(nc, factor=block_factor_n) + block_j, oc = s[output].split(oc, factor=block_factor_o) + s[output].reorder(block_k, block_i, block_j, nc, oc) + t = s[output].fuse(nc, oc) + t, ti = s[output].split(t, factor=vector_width) + t, tx = s[output].split(t, factor=warp_size) + t, ty = s[output].split(t, factor=block_row_warps) + t, tz = s[output].split(t, factor=block_col_warps) + s[output].bind(block_i, block_x) + s[output].bind(block_j, block_y) + s[output].bind(tz, thread_z) + s[output].bind(ty, thread_y) + s[output].bind(tx, thread_x) + s[output].vectorize(ti) + + # Schedule wmma store + s[OL].compute_at(s[output], block_j) + nc, hc, wc, oc = OL.op.axis + s[OL].reorder(hc, wc, nc, oc) + s[OL].storage_align(wc, CS_align - 1, CS_align) + oc, ooc = s[OL].split(oc, factor=wmma_n) + oc, oci = s[OL].split(oc, factor=warp_col_tiles) + _, oc = s[OL].split(oc, factor=block_col_warps) + nc, nnc = s[OL].split(nc, factor=wmma_m) + nc, nci = s[OL].split(nc, factor=warp_row_tiles) + _, nc = s[OL].split(nc, factor=block_row_warps) + s[OL].reorder(nc, oc, nci, oci, nnc, ooc) + s[OL].bind(nc, thread_y) + s[OL].bind(oc, thread_z) + + # Schedule wmma computation + s[ConvF].compute_at(s[OL], oc) + n, h, w, o = ConvF.op.axis + n, nnf = s[ConvF].split(n, factor=wmma_m) + o, oof = s[ConvF].split(o, factor=wmma_n) + ic, ii = s[ConvF].split(ic, factor=wmma_k) + ko, ki = s[ConvF].split(ic, factor=chunk) + s[ConvF].reorder(kh, kw, ko, ki, n, o, nnf, oof, ii) + + s[AF].compute_at(s[ConvF], ki) + s[WF].compute_at(s[ConvF], ki) + + # Schedule wmma load + n, h, w, i = AF.op.axis + n, nn = s[AF].split(n, factor=wmma_m) + i, ii = s[AF].split(i, factor=wmma_k) + s[AF].reorder(n, i, nn, ii) + + kh, kw, i, o = WF.op.axis + i, ii = s[WF].split(i, factor=wmma_k) + o, oo = s[WF].split(o, factor=wmma_n) + s[WF].reorder(o, i, oo) + s[WF].reorder(i, o, ii, oo) + + s[WS].compute_at(s[ConvF], ko) + s[AS].compute_at(s[ConvF], ko) + + # Schedule for data's share memory + n, h, w, i = AS.op.axis + s[AS].reorder(h, w, n, i) + s[AS].storage_align(w, AS_align - 1, AS_align) + t = s[AS].fuse(n, i) + t, ti = s[AS].split(t, factor=vector_width) + t, tx = s[AS].split(t, factor=warp_size) + t, ty = s[AS].split(t, factor=block_row_warps) + _, tz = s[AS].split(t, factor=block_col_warps) + s[AS].bind(ty, thread_y) + s[AS].bind(tz, thread_z) + s[AS].bind(tx, thread_x) + s[AS].vectorize(ti) + + # Schedule for kernel's share memory + kh, kw, ic, o = WS.op.axis + t = s[WS].fuse(ic, o) + s[WS].storage_align(ic, WS_align - 1, WS_align) + t, ti = s[WS].split(t, factor=vector_width) + t, tx = s[WS].split(t, factor=warp_size) + t, ty = s[WS].split(t, factor=block_row_warps) + _, tz = s[WS].split(t, factor=block_col_warps) + s[WS].bind(ty, thread_y) + s[WS].bind(tz, thread_z) + s[WS].bind(tx, thread_x) + s[WS].vectorize(ti) + + shape = (wmma_m, wmma_n, wmma_k) + + # tensorize the wmma process + AS_shape = (wmma_m, 1, 1, wmma_k) + AL_shape = (wmma_m, 1, 1, wmma_k) + WS_shape = (wmma_k, wmma_n) + WL_shape = (wmma_k, wmma_n) + CL_shape = (wmma_m, 1, 1, wmma_n) + CS_shape = (wmma_m, 1, 1, wmma_n) + + AL_gemm = te.placeholder(AL_shape, name="A", dtype=in_dtype) + WL_gemm = te.placeholder(WL_shape, name="B", dtype=in_dtype) + k_gemm = te.reduce_axis((0, wmma_k), name="k") + CL_compute = te.compute( + CL_shape, + lambda ii, t0, t1, jj: te.sum( + AL_gemm[ii, t0, t1, k_gemm].astype(out_dtype) * WL_gemm[k_gemm, jj].astype(out_dtype), + axis=k_gemm, + ), + name="C", + ) + + s[AF].tensorize( + nn, + intrin_wmma_load_matrix_A( + AL_strides, AS_strides, shape, "row_major", AS_shape, AL_shape, in_dtype + ), + ) + s[WF].tensorize( + ii, + intrin_wmma_load_matrix_W( + WL_strides, WS_strides, shape, "row_major", WS_shape, WL_shape, in_dtype + ), + ) + s[OL].tensorize( + nnc, intrin_wmma_store_matrix(CS_strides, CL_strides, shape, out_dtype, CL_shape, CS_shape) + ) + s[ConvF].tensorize( + nnf, + intrin_wmma_gemm(AL_gemm, WL_gemm, CL_compute, AL_strides, WL_strides, CL_strides, shape), + ) + + N, OH, OW, CO = get_const_tuple(output.shape) + KH, KW, CI, _ = get_const_tuple(kernel.shape) + cfg.add_flop(2 * N * OH * OW * CO * CI * KH * KW) + + +@autotvm.register_topi_compute("conv2d_nhwc_tensorcore.maca") +def conv2d_nhwc_tensorcore(cfg, data, kernel, strides, padding, dilation, out_dtype): + """Compute conv2d with tensorcore for NCHW layout""" + return nhwc_tensorcore_maca(cfg, data, kernel, strides, padding, dilation, out_dtype) + + +@autotvm.register_topi_schedule("conv2d_nhwc_tensorcore.maca") +def schedule_conv2d_nhwc_tensorcore(cfg, outs): + """TOPI schedule callback""" + s = te.create_schedule([x.op for x in outs]) + + def _callback(op): + if "conv2d_nhwc_tensorcore" in op.tag: + schedule_nhwc_tensorcore_maca(cfg, s, op.output(0)) + + traverse_inline(s, outs[0].op, _callback) + return s diff --git a/python/tvm/topi/maca/conv2d_transpose.py b/python/tvm/topi/maca/conv2d_transpose.py new file mode 100644 index 000000000000..a1f8a43ac806 --- /dev/null +++ b/python/tvm/topi/maca/conv2d_transpose.py @@ -0,0 +1,44 @@ +# 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. +# pylint: disable=invalid-name +"""Conv2d transpose template for cuda backend""" + +import tvm +from tvm import te +from tvm.contrib import mcdnn +from tvm import autotvm +from tvm.autotvm.task.space import SplitEntity, OtherOptionEntity +from .. import nn +from ..utils import get_const_tuple, traverse_inline + +def conv2d_transpose_mcdnn( + x, w, stride, padding, out_dtype, output_padding=(0, 0), layout="NCHW", groups=1 +): + """Compute conv2d_tranpose using mcdnn dgrad kernel""" + tensor_format = 0 if layout == "NCHW" else 1 + return mcdnn.conv_backward_data( + x, + w, + padding, + stride, + (1, 1), + 1, + tensor_format, + out_dtype, + groups=groups, + output_padding=output_padding, + ) diff --git a/python/tvm/topi/maca/conv3d.py b/python/tvm/topi/maca/conv3d.py new file mode 100644 index 000000000000..bb03780bbfaa --- /dev/null +++ b/python/tvm/topi/maca/conv3d.py @@ -0,0 +1,144 @@ +# 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. +# pylint: disable=invalid-name, unused-argument +"""Compute definition for conv3d with cuda backend""" +from tvm import te +from tvm import autotvm +from tvm.contrib import mcdnn + +from .. import nn, generic +from ..utils import get_const_tuple, traverse_inline +from .conv3d_direct import schedule_direct_conv3d_cuda + +@autotvm.register_topi_compute("conv3d_mcdnn.maca") +def conv3d_mcdnn( + cfg, data, kernel, strides, padding, dilation, groups, layout="NCDHW", out_dtype="float32" +): + """Conv3D operator for maca backend. + + Parameters + ---------- + cfg: ConfigEntity + The config for this template + + data : tvm.te.Tensor + 5-D with shape [batch, in_channel, in_depth, in_height, in_width] + + kernel : tvm.te.Tensor + 5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width] + + strides : int or a list/tuple of three ints + stride size, or [stride_depth, stride_height, stride_width] + + padding : int or a list/tuple of three ints + padding size, or [pad_depth, pad_height, pad_width] + + dilation: int or a list/tuple of three ints + dilation size, or [dilation_depth, dilation_height, dilation_width] + + layout : str + layout of data + + out_dtype: str + The output type. This is used for mixed precision. + + Returns + ------- + output : tvm.te.Tensor + 5-D with shape [batch, out_channel, out_depth, out_height, out_width] + """ + if layout == "NCDHW": + tensor_format = 0 # MCDNN_TENSOR_NCHW + N, _, D, H, W = get_const_tuple(data.shape) + elif layout == "NDHWC": + tensor_format = 1 # MCDNN_TENSOR_NHWC + N, D, H, W, _ = get_const_tuple(data.shape) + else: + raise ValueError(f"Unsupported layout {layout} in mcdnn") + CO, CI, KD, KH, KW = get_const_tuple(kernel.shape) + + assert groups == 1, "conv3d_mcdnn does not support groups" + + # handle dilation + stride_d, stride_h, stride_w = ( + (strides, strides, strides) if isinstance(strides, int) else strides + ) + pad_d, pad_h, pad_w = (padding, padding, padding) if isinstance(padding, int) else padding + dilation_d, dilation_h, dilation_w = ( + (dilation, dilation, dilation) if isinstance(dilation, int) else dilation + ) + + OD = (D + 2 * pad_d - KD) // stride_d + 1 + OH = (H + 2 * pad_h - KH) // stride_h + 1 + OW = (W + 2 * pad_w - KW) // stride_w + 1 + + if isinstance(N, int): + cfg.add_flop( + 2 + * N + * OD + * OH + * OW + * CO + * CI + * ((KD - 1) * dilation_d + 1) + * ((KH - 1) * dilation_h + 1) + * ((KW - 1) * dilation_w + 1) + ) + + cfg.define_knob("algo", range(mcdnn.algo_to_index("fwd", "MCDNN_CONVOLUTION_FWD_ALGO_COUNT"))) + if cfg.is_fallback: + if mcdnn.exists(): + # Let MCDNN choose the best algo, based on benchmarks run + # on the local machine. In the future, this should be + # based on parameters stored in the Target. + cfg["algo"] = OtherOptionEntity(-1) + else: + cfg["algo"] = OtherOptionEntity(0) + + return mcdnn.conv_forward( + data, + kernel, + [pad_d, pad_h, pad_w], + [stride_d, stride_h, stride_w], + [dilation_d, dilation_h, dilation_w], + conv_mode=1, + tensor_format=tensor_format, + algo=cfg["algo"].val, + conv_dtype=dtype, + ) + + +@autotvm.register_topi_schedule("conv3d_mcdnn.maca") +def schedule_conv3d_mcdnn(_, outs): + """TOPI schedule callback of conv3d for maca gpu + + Parameters + ---------- + cfg: ConfigEntity + The config for this template + + outs: Array of Tensor + The computation graph description of conv2d + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for conv2d. + """ + return generic.schedule_extern(outs) diff --git a/python/tvm/topi/maca/dense.py b/python/tvm/topi/maca/dense.py new file mode 100644 index 000000000000..e73ab84c98d6 --- /dev/null +++ b/python/tvm/topi/maca/dense.py @@ -0,0 +1,76 @@ +# 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. +# pylint: disable=invalid-name, unused-argument +"""Schedule for dense operator""" +import logging +import tvm +from tvm import te, autotvm +from tvm.contrib import mcblas +from .tensor_intrin import dp4a +from .. import tag +from .. import generic +from ..utils import traverse_inline, get_const_tuple + +logger = logging.getLogger("topi") + + +def _matmul_mcblas_common( + cfg, tensor_a, tensor_b, bias=None, out_dtype=None, transpose_a=False, transpose_b=False +): + assert len(tensor_a.shape) == 2 and len(tensor_b.shape) == 2, "only support 2-dim matmul" + if bias is not None: + assert len(bias.shape) == 1 + if out_dtype is None: + out_dtype = tensor_a.dtype + if out_dtype not in [tensor_a.dtype, "int32"]: + assert out_dtype == tensor_a.dtype, "Mixed precision other than int8 + int32 not supported." + batch, in_dim = get_const_tuple(tensor_a.shape) + out_dim, _ = get_const_tuple(tensor_b.shape) + matmul = mcblas.matmul(tensor_a, tensor_b, transpose_a, transpose_b, dtype=out_dtype) + if all(isinstance(d, int) for d in [batch, in_dim, out_dim]): + cfg.add_flop(batch * in_dim * out_dim * 2) + if bias is not None: + matmul = te.compute( + (batch, out_dim), lambda i, j: matmul[i, j] + bias[j], tag=tag.BROADCAST + ) + return matmul + + +@autotvm.register_topi_compute("matmul_mcblas.maca") +def matmul_mcblas( + cfg, tensor_a, tensor_b, bias=None, out_dtype=None, transpose_a=False, transpose_b=False +): + """Matmul operator on MACA with MCBLAS""" + return _matmul_mcblas_common(cfg, tensor_a, tensor_b, bias, out_dtype, transpose_a, transpose_b) + + +@autotvm.register_topi_schedule("matmul_mcblas.maca") +def schedule_matmul_mcblas(_, outs): + """Schedule matmul operator using MCBLAS""" + return generic.schedule_extern(outs) + + +@autotvm.register_topi_compute("dense_mcblas.maca") +def dense_mcblas(cfg, data, weight, bias=None, out_dtype=None): + """Dense operator on MACA with MCBLAS. This is an alias of matmul_nt operator.""" + return _matmul_mcblas_common(cfg, data, weight, bias, out_dtype, False, True) + + +@autotvm.register_topi_schedule("dense_mcblas.maca") +def schedule_dense_mcblas(_, outs): + """Schedule dense operator using MCBLAS""" + return generic.schedule_extern(outs) diff --git a/python/tvm/topi/maca/softmax.py b/python/tvm/topi/maca/softmax.py new file mode 100644 index 000000000000..8f6a33096d6e --- /dev/null +++ b/python/tvm/topi/maca/softmax.py @@ -0,0 +1,43 @@ +# 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. +# pylint: disable=invalid-name, unused-variable, trailing-whitespace +"""Schedule for softmax operator""" +from tvm.target import Target +from tvm import te +from tvm.contrib import mcdnn +from .. import generic +from .injective import schedule_injective_from_existing +from ..utils import get_const_int, traverse_inline + +def softmax_mcdnn(x, axis=-1): + """Perform softmax on the data using mcdnn""" + return mcdnn.softmax(x, axis) + + +def schedule_softmax_mcdnn(outs): + """Schedule for softmax mcdnn op""" + return generic.schedule_extern(outs) + + +def log_softmax_mcdnn(x, axis=-1): + """Perform log_softmax on the data using mcdnn""" + return mcdnn.log_softmax(x, axis) + + +def schedule_log_softmax_mcdnn(outs): + """Schedule for log_softmax mcdnn op""" + return generic.schedule_extern(outs) diff --git a/python/tvm/topi/maca/tensor_intrin.py b/python/tvm/topi/maca/tensor_intrin.py new file mode 100644 index 000000000000..e5317dcc89b1 --- /dev/null +++ b/python/tvm/topi/maca/tensor_intrin.py @@ -0,0 +1,247 @@ +# 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. +# pylint: disable=invalid-name, unnecessary-lambda, too-many-arguments +"""Tensor intrinsics on MACA.""" +import tvm +from tvm import te +from ..utils import is_target + +def intrin_wmma_load_matrix_A(strides_dst, strides_from, shape, layout, A_shape, C_shape, in_dtype): + """Intrin function for loading data from shared memory to wmma.matrix_a""" + wmma_m, wmma_n, wmma_k = shape + + A = te.placeholder(A_shape, name="A", dtype=in_dtype) + BA = tvm.tir.decl_buffer( + A.shape, A.dtype, scope="shared", strides=strides_from, data_alignment=32, offset_factor=8 + ) + C = te.compute(C_shape, lambda *i: A(*i), name="C") + BC = tvm.tir.decl_buffer( + C.shape, + C.dtype, + scope="wmma.matrix_a", + strides=strides_dst, + data_alignment=32, + offset_factor=8, + ) + + def intrin_func(ins, outs): + ib = tvm.tir.ir_builder.create() + + BA = ins[0] + BC = outs[0] + row = wmma_m * wmma_k + warp_index = BC.elem_offset // row + BC.elem_offset % row // wmma_k + ib.emit( + tvm.tir.call_intrin( + "handle", + "tir.tvm_load_matrix_sync", + BC.data, + wmma_m, + wmma_n, + wmma_k, + warp_index, + BA.access_ptr("r"), + strides_from[0], + layout, + ) + ) + return ib.get() + + return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, C: BC}) + + +def intrin_wmma_load_matrix_W(strides_dst, strides_from, shape, layout, A_shape, C_shape, in_dtype): + """Intrin function for loading data from shared memory to wmma.matrix_b""" + wmma_m, wmma_n, wmma_k = shape + + A = te.placeholder(A_shape, name="A", dtype=in_dtype) + BA = tvm.tir.decl_buffer( + A.shape, A.dtype, scope="shared", strides=strides_from, data_alignment=32, offset_factor=8 + ) + C = te.compute(C_shape, lambda *i: A(*i), name="C") + BC = tvm.tir.decl_buffer( + C.shape, + C.dtype, + scope="wmma.matrix_b", + strides=strides_dst, + data_alignment=32, + offset_factor=8, + ) + + def intrin_func(ins, outs): + ib = tvm.tir.ir_builder.create() + + BA = ins[0] + BC = outs[0] + row = wmma_n * wmma_k + warp_index = BC.elem_offset // row + BC.elem_offset % row // wmma_n + ib.emit( + tvm.tir.call_intrin( + "handle", + "tir.tvm_load_matrix_sync", + BC.data, + wmma_m, + wmma_n, + wmma_k, + warp_index, + BA.access_ptr("r"), + strides_from[0], + layout, + ) + ) + return ib.get() + + return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, C: BC}) + + +def intrin_wmma_store_matrix(strides_dst, strides_from, shape, out_dtype, A_shape, C_shape): + """Intrin function for storing the results from wmma.accumulator to shared""" + wmma_m, wmma_n, wmma_k = shape + A = te.placeholder(A_shape, name="A", dtype=out_dtype) + BA = tvm.tir.decl_buffer( + A.shape, + A.dtype, + scope="wmma.accumulator", + strides=strides_from, + data_alignment=32, + offset_factor=8, + ) + C = te.compute(C_shape, lambda *i: A(*i), name="C") + BC = tvm.tir.decl_buffer( + C.shape, C.dtype, scope="shared", strides=strides_dst, data_alignment=32, offset_factor=8 + ) + + def intrin_func(ins, outs): + ib = tvm.tir.ir_builder.create() + + BA = ins[0] + BC = outs[0] + row = wmma_m * wmma_n + warp_index = BA.elem_offset // row + BA.elem_offset % row // wmma_n + ib.emit( + tvm.tir.call_intrin( + "handle", + "tir.tvm_store_matrix_sync", + BA.data, + wmma_m, + wmma_n, + wmma_k, + warp_index, + BC.access_ptr("w"), + strides_dst[0], + "row_major", + ) + ) + return ib.get() + + return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, C: BC}) + + +def intrin_wmma_gemm(AL_gemm, WL_gemm, CL_compute, strides_A, strides_W, strides_Conv, shape): + """Intrin for wmma fill_fragment and mma_sync + + Parameters + ---------- + AL_gemm : tvm.te.placeholder + wmma matrix A + WL_gemm : tvm.te.placeholder + wmma matrix B + CL_compute : tvm.te.compute + The definition of wmma gemm + """ + wmma_m, wmma_n, wmma_k = shape + A = AL_gemm + B = WL_gemm + C = CL_compute + + BA = tvm.tir.decl_buffer( + A.shape, + A.dtype, + name="BA", + scope="wmma.matrix_a", + data_alignment=32, + offset_factor=8, + strides=strides_A, + ) + BB = tvm.tir.decl_buffer( + B.shape, + B.dtype, + name="BB", + scope="wmma.matrix_b", + data_alignment=32, + offset_factor=8, + strides=strides_W, + ) + BC = tvm.tir.decl_buffer( + C.shape, + C.dtype, + name="BC", + scope="wmma.accumulator", + data_alignment=32, + offset_factor=8, + strides=strides_Conv, + ) + + def intrin_func(ins, outs): + BA, BB = ins + (BC,) = outs + + def warp_idnex(offset, row, col): + row = row * col + return offset // row + offset % row // col + + warp_index_A = warp_idnex(BA.elem_offset, wmma_m, wmma_k) + warp_index_B = warp_idnex(BB.elem_offset, wmma_k, wmma_n) + warp_index_C = warp_idnex(BC.elem_offset, wmma_m, wmma_n) + + def init(): + ib = tvm.tir.ir_builder.create() + ib.emit( + tvm.tir.call_intrin( + "handle", + "tir.tvm_fill_fragment", + BC.data, + wmma_m, + wmma_n, + wmma_k, + warp_index_C, + 0.0, + ) + ) + return ib.get() + + def update(): + ib = tvm.tir.ir_builder.create() + ib.emit( + tvm.tir.call_intrin( + "handle", + "tir.tvm_mma_sync", + BC.data, + warp_index_C, + BA.data, + warp_index_A, + BB.data, + warp_index_B, + BC.data, + warp_index_C, + ) + ) + return ib.get() + + return update(), init(), update() + + return te.decl_tensor_intrin(C.op, intrin_func, binds={A: BA, B: BB, C: BC}) diff --git a/src/relax/backend/contrib/mcblas/codegen.cc b/src/relax/backend/contrib/mcblas/codegen.cc new file mode 100644 index 000000000000..e3e371ce5f4f --- /dev/null +++ b/src/relax/backend/contrib/mcblas/codegen.cc @@ -0,0 +1,133 @@ +/* + * 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. + */ + +/*! + * \file src/relax/backend/contrib/mcblas/codegen.cc + * \brief Implementation of the MCBLAS JSON serializer. + */ +#include +#include + +#include + +#include "../codegen_json/codegen_json.h" +#include "../utils.h" + +namespace tvm { +namespace relax { +namespace contrib { + +using JSONGraphNode = tvm::runtime::json::JSONGraphNode; +using JSONGraphNodeEntry = tvm::runtime::json::JSONGraphNodeEntry; +using JSONSerializer = backend::contrib::JSONSerializer; +using backend::contrib::NodeEntries; + +class McblasJSONSerializer : public JSONSerializer { + public: + McblasJSONSerializer(Map constant_names, Map bindings) + : JSONSerializer(constant_names), bindings_(bindings) {} + + using JSONSerializer::VisitExpr_; + + NodeEntries VisitExpr_(const CallNode* call_node) final { + const auto* fn_var = call_node->op.as(); + ICHECK(fn_var); + const auto fn = Downcast(bindings_[GetRef(fn_var)]); + ICHECK(fn.defined()) << "Expects the callee to be a function."; + + auto composite_opt = fn->GetAttr(attr::kComposite); + ICHECK(composite_opt.defined()) << "Only composite functions are supported."; + + std::string composite_name = composite_opt.value(); + + NodeEntries inputs_tmp; + for (const auto& arg : call_node->args) { + auto res = VisitExpr(arg); + inputs_tmp.insert(inputs_tmp.end(), res.begin(), res.end()); + } + + ICHECK(inputs_tmp.size() <= 4); + NodeEntries inputs(inputs_tmp.size()); + + auto arg_idx = backend::ExtractArgIdx(composite_name, fn); + inputs[0] = inputs_tmp[arg_idx["lhs"]->value]; + inputs[1] = inputs_tmp[arg_idx["rhs"]->value]; + if (inputs_tmp.size() == 3) { + inputs[2] = inputs_tmp[arg_idx["bias"]->value]; + } else if (inputs_tmp.size() == 4) { + inputs[2] = inputs_tmp[arg_idx["scaleA"]->value]; + inputs[3] = inputs_tmp[arg_idx["scaleB"]->value]; + } + + auto node = std::make_shared(composite_name, /* name_ */ + "kernel", /* op_type_ */ + inputs, 1 /* num_outputs_ */); + if (composite_name.find("dequantize") != std::string::npos) { + const CallNode* dequantize_call = backend::GetOpInFunction(fn, "relax.dequantize"); + if (dequantize_call->args[1]->IsInstance()) { + const auto* const_expr = dequantize_call->args[1].as(); + auto sinfo = Downcast(const_expr->struct_info_); + float alpha = 1.0; + if (sinfo->dtype == DataType::Float(16)) { + alpha = __gnu_h2f_ieee(static_cast(const_expr->data->data)[0]); + } else { + ICHECK(sinfo->dtype == DataType::Float(32)); + alpha = static_cast(const_expr->data->data)[0]; + } + + std::vector dq_scale = {backend::to_str(alpha)}; + std::vector dq_scale_attr; + dq_scale_attr.emplace_back(dq_scale); + node->SetAttr("dq_scale", dq_scale_attr); + } + } + + const CallNode* root_call = backend::GetOpInFunction(fn, "relax.matmul"); + SetCallNodeAttribute(node, root_call); + return AddNode(node, GetRef(call_node)); + } + + private: + /*! \brief The bindings to look up composite functions. */ + Map bindings_; +}; + +Array McblasCompiler(Array functions, Map /*unused*/, + Map constant_names) { + Array compiled_functions; + + for (const auto& func : functions) { + McblasJSONSerializer serializer(constant_names, AnalyzeVar2Value(func)); + serializer.serialize(func); + auto graph_json = serializer.GetJSON(); + auto constant_names = serializer.GetConstantNames(); + const auto* pf = runtime::Registry::Get("runtime.McblasJSONRuntimeCreate"); + ICHECK(pf != nullptr) << "Cannot find MCBLAS runtime module create function."; + auto func_name = GetExtSymbol(func); + compiled_functions.push_back((*pf)(func_name, graph_json, constant_names)); + } + + return compiled_functions; +} + +TVM_REGISTER_GLOBAL("relax.ext.mcblas").set_body_typed(McblasCompiler); + +} // namespace contrib +} // namespace relax +} // namespace tvm diff --git a/src/relax/backend/contrib/mcdnn/codegen.cc b/src/relax/backend/contrib/mcdnn/codegen.cc new file mode 100644 index 000000000000..8635674f1e10 --- /dev/null +++ b/src/relax/backend/contrib/mcdnn/codegen.cc @@ -0,0 +1,157 @@ +/* + * 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. + */ + +/*! + * \file src/relax/backend/contrib/mcdnn/codegen.cc + * \brief Implementation of the mcDNN JSON serializer. + */ +#include + +#include + +#include "../codegen_json/codegen_json.h" +#include "../utils.h" + +namespace tvm { +namespace relax { +namespace contrib { + +using JSONGraphNode = tvm::runtime::json::JSONGraphNode; +using JSONGraphNodeEntry = tvm::runtime::json::JSONGraphNodeEntry; +using JSONSerializer = backend::contrib::JSONSerializer; +using backend::contrib::NodeEntries; + +class mcDNNJSONSerializer : public JSONSerializer { + public: + mcDNNJSONSerializer(Map constant_names, Map bindings) + : JSONSerializer(constant_names), bindings_(bindings) {} + + using JSONSerializer::VisitExpr_; + + NodeEntries VisitExpr_(const CallNode* call_node) final { + const auto* fn_var = call_node->op.as(); + ICHECK(fn_var); + const auto fn = Downcast(bindings_[GetRef(fn_var)]); + ICHECK(fn.defined()) << "Expects the callee to be a function."; + + auto composite_opt = fn->GetAttr(attr::kComposite); + ICHECK(composite_opt.defined()) << "Only composite functions are supported."; + + std::string composite_name = composite_opt.value(); + + if (composite_name.find("mcdnn.conv2d") != std::string::npos) { + return HandleConv2D(call_node, fn, composite_name); + } else if (composite_name.find("mcdnn.attention") != std::string::npos) { + return HandleAttention(call_node, fn, composite_name); + } else { + LOG(FATAL) << "Unsupported composite function: " << composite_name; + } + } + + NodeEntries HandleConv2D(const CallNode* call_node, const Function& fn, + const std::string& composite_name) { + NodeEntries inputs_tmp; + for (const auto& arg : call_node->args) { + auto res = VisitExpr(arg); + inputs_tmp.insert(inputs_tmp.end(), res.begin(), res.end()); + } + + ICHECK(inputs_tmp.size() <= 3); + NodeEntries inputs(inputs_tmp.size()); + + auto arg_idx = backend::ExtractArgIdx(composite_name, fn); + inputs[0] = inputs_tmp[arg_idx["input"]->value]; + inputs[1] = inputs_tmp[arg_idx["weight"]->value]; + if (inputs_tmp.size() == 3) { + inputs[2] = inputs_tmp[arg_idx["bias"]->value]; + } + + auto node = std::make_shared(composite_name, /* name_ */ + "kernel", /* op_type_ */ + inputs, 1 /* num_outputs_ */); + + const CallNode* root_call = backend::GetOpInFunction(fn, "relax.nn.conv2d"); + SetCallNodeAttribute(node, root_call); + return AddNode(node, GetRef(call_node)); + } + + NodeEntries HandleAttention(const CallNode* call_node, const Function& fn, + const std::string& composite_name) { + std::string layout = composite_name.substr(composite_name.find_last_of(".") + 1); + NodeEntries inputs; + for (const auto& arg : call_node->args) { + auto res = VisitExpr(arg); + inputs.insert(inputs.end(), res.begin(), res.end()); + } + ICHECK_EQ(inputs.size(), 2); + auto node = std::make_shared(composite_name, /* name_ */ + "kernel", /* op_type_ */ + inputs, 1 /* num_outputs_ */); + const CallNode* root_call = backend::GetOpInFunction(fn, "relax.nn.attention"); + auto q_shape = Downcast( + Downcast(root_call->args[0]->struct_info_.value())->shape.value()); + auto k_shape = Downcast( + Downcast(root_call->args[1]->struct_info_.value())->shape.value()); + auto v_shape = Downcast( + Downcast(root_call->args[2]->struct_info_.value())->shape.value()); + int num_heads = q_shape->values[2].as()->value; + int num_kv_heads = k_shape->values[2].as()->value; + int head_size = q_shape->values[3].as()->value; + int head_size_v = v_shape->values[3].as()->value; + SetCallNodeAttribute(node, root_call); + + auto to_str_array = [](int val) { + return std::vector{std::vector{std::to_string(val)}}; + }; + node->SetAttr("num_heads", to_str_array(num_heads)); + node->SetAttr("num_kv_heads", to_str_array(num_kv_heads)); + node->SetAttr("head_size", to_str_array(head_size)); + node->SetAttr("head_size_v", to_str_array(head_size_v)); + node->SetAttr("layout", std::vector{std::vector{layout}}); + return AddNode(node, GetRef(call_node)); + } + + private: + /*! \brief The bindings to look up composite functions. */ + Map bindings_; +}; + +Array mcDNNCompiler(Array functions, Map /*unused*/, + Map constant_names) { + Array compiled_functions; + + for (const auto& func : functions) { + mcDNNJSONSerializer serializer(constant_names, AnalyzeVar2Value(func)); + serializer.serialize(func); + auto graph_json = serializer.GetJSON(); + auto constant_names = serializer.GetConstantNames(); + const auto* pf = runtime::Registry::Get("runtime.mcDNNJSONRuntimeCreate"); + ICHECK(pf != nullptr) << "Cannot find mcDNN runtime module create function."; + auto func_name = GetExtSymbol(func); + compiled_functions.push_back((*pf)(func_name, graph_json, constant_names)); + } + + return compiled_functions; +} + +TVM_REGISTER_GLOBAL("relax.ext.mcdnn").set_body_typed(mcDNNCompiler); + +} // namespace contrib +} // namespace relax +} // namespace tvm diff --git a/src/relay/backend/contrib/mcblas/target.cc b/src/relay/backend/contrib/mcblas/target.cc new file mode 100644 index 000000000000..a30efde3cb04 --- /dev/null +++ b/src/relay/backend/contrib/mcblas/target.cc @@ -0,0 +1,44 @@ +/* + * 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. + */ + +/*! + * \file src/relay/backend/contrib/cudnn/target.cc + * \brief Registers the "mcblas" external codegen TargetKind. + */ + +#include + +namespace tvm { +namespace relay { +namespace contrib { + +/*! + * \brief This external codegen target can use the McBLAS library linked into the TVM runtime. + * - Patterns and custom compiler: python/tvm/relay/op/contrib/mcblas.py + * - Custom schedules: python/tvm/contrib/mcblas.py + * - Runtime: src/runtime/contrib/mcblas/mcblas.cc + * + * McBLAS can also be used via the "-libs=mcblas" Target option. + */ +TVM_REGISTER_TARGET_KIND("mcblas", kDLMACA) + .set_attr(tvm::attr::kIsExternalCodegen, Bool(true)); + +} // namespace contrib +} // namespace relay +} // namespace tvm diff --git a/src/relay/backend/contrib/mcdnn/target.cc b/src/relay/backend/contrib/mcdnn/target.cc new file mode 100644 index 000000000000..5134322cbc4c --- /dev/null +++ b/src/relay/backend/contrib/mcdnn/target.cc @@ -0,0 +1,42 @@ +/* + * 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. + */ + +/*! + * \file src/relay/backend/contrib/cudnn/target.cc + * \brief Registers the "cudnn" external codegen TargetKind. + */ + +#include + +namespace tvm { +namespace relay { +namespace contrib { + +/*! + * \brief This external codegen target can use the CuDNN library linked into the TVM runtime. + * - Patterns and custom compiler: python/tvm/relay/op/contrib/cudnn.py + * - Custom schedules: python/tvm/contrib/cudnn.py + * - Runtime: src/runtime/contrib/cudnn/ *.cc + */ +TVM_REGISTER_TARGET_KIND("mcdnn", kDLMACA) + .set_attr(tvm::attr::kIsExternalCodegen, Bool(true)); + +} // namespace contrib +} // namespace relay +} // namespace tvm diff --git a/src/relay/backend/contrib/mctlass/codegen.cc b/src/relay/backend/contrib/mctlass/codegen.cc new file mode 100644 index 000000000000..94cfc9f17391 --- /dev/null +++ b/src/relay/backend/contrib/mctlass/codegen.cc @@ -0,0 +1,442 @@ +/* + * 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. + */ + +/*! + * \file src/relay/backend/contrib/mctlass/codegen.cc + * \brief The 'custom' compilation pass for MCTLASS (invoked by the RelayToTIRTargetHook pass). + */ + +#include "codegen.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "../../../transforms/compiler_function_utils.h" +#include "../../utils.h" +#include "../codegen_c/codegen_c.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace mctlass { + +std::string EmitSignature(const std::vector& out, const std::string& func_id, + const std::vector& arg_names) { + std::ostringstream code_stream_; + code_stream_ << "void " << func_id << "_("; + for (const auto& arg_name : arg_names) { + code_stream_ << "DLTensor* " << arg_name << ", "; + } + for (size_t i = 0; i < out.size() - 1; ++i) { + code_stream_ << "DLTensor* out" << i << ", "; + } + code_stream_ << "DLTensor* out" << out.size() - 1 << ")"; + return code_stream_.str(); +} + +runtime::Module Finalize(const std::string& code, const Array& func_names) { + ICHECK(!func_names.empty()) + << "Should only create MCTLASS CSourceModule if there is at least one MCTLASS partition"; + + std::ostringstream default_headers; + default_headers << "#include \n"; + default_headers << "#include \n"; + default_headers << "#include \n"; + default_headers << "#include \n"; + default_headers << "#include \n"; + default_headers << "#include \n"; + default_headers << "#include \n"; + + const auto* pf = runtime::Registry::Get("runtime.CSourceModuleCreate"); + ICHECK(pf != nullptr) << "Cannot find CSource module to create the external runtime module"; + VLOG(1) << "Generated MCTLASS code:" << std::endl << code; + return (*pf)(default_headers.str() + code, "cu", func_names, /*const_vars=*/Array()); +} + +class CodegenResultNode : public Object { + public: + String code; + Array headers; + + void VisitAttrs(AttrVisitor* v) { + v->Visit("code", &code); + v->Visit("headers", &headers); + } + static constexpr const char* _type_key = "contrib.mctlass.CodegenResult"; + TVM_DECLARE_FINAL_OBJECT_INFO(CodegenResultNode, Object); +}; + +class CodegenResult : public ObjectRef { + public: + CodegenResult(String code, Array headers) { + auto n = make_object(); + n->code = std::move(code); + n->headers = std::move(headers); + data_ = std::move(n); + } + + TVM_DEFINE_OBJECT_REF_METHODS(CodegenResult, ObjectRef, CodegenResultNode) +}; + +TVM_REGISTER_NODE_TYPE(CodegenResultNode); + +TVM_REGISTER_GLOBAL("contrib.mctlass.CodegenResult") + .set_body_typed([](String code, Array headers) { + return CodegenResult(code, headers); + }); + +GenerateBodyOutput GenerateBody(const std::string& func_name, const std::string& ext_func_id, + const std::vector& output_types, + const Array& func_args, const Map& attrs, + int* buf_idx) { + // Make function call with input buffers when visiting arguements + ICHECK_GT(func_args.size(), 0); + std::ostringstream decl_stream; + decl_stream << "(" << func_args[0]; + for (size_t i = 1; i < func_args.size(); ++i) { + decl_stream << ", " << func_args[i]; + } + GenerateBodyOutput ret; + for (const auto& out_type : output_types) { + const std::string out = "out" + std::to_string(*buf_idx++); + decl_stream << ", " << out; + Output output; + output.name = out; + output.dtype = out_type; + output.need_copy = false; + ret.outputs.push_back(output); + } + decl_stream << ");"; + + const auto* instantiate_template_func = + runtime::Registry::Get("contrib.mctlass.instantiate_template"); + ICHECK(instantiate_template_func); + + CodegenResult codegen_res = (*instantiate_template_func)(func_name, attrs, func_args); + ret.decl = codegen_res->code; + ret.headers = codegen_res->headers; + + return ret; +} + +namespace { + +/*! \brief Return the "mctlass" Target instance to use to guide compilation. */ +Target GetCutlassTarget() { + Target target = Target::Current(/*allow_not_defined=*/true); + if (!target.defined() || target->kind->name != "mctlass") { + // Use the default MCTLASS compilation options if no specific "mctlass" target was given + // in the overall targets list. In that case target_hooks.cc will invoke the custom pass + // without pushing any target instance onto the implicit target stack. + target = Target("mctlass"); + } + return target; +} + +class CodegenCutlass : public backend::MemoizedExprTranslator>, + public CodegenCBase { + public: + CodegenCutlass(const std::string& id, const Map& attrs) { + this->ext_func_id_ = id; + this->attrs_ = attrs; + } + + std::vector VisitExprDefault_(const Object* op) final { + LOG(FATAL) << "Cutlass codegen doesn't support: " << op->GetTypeKey(); + } + + std::vector VisitExpr_(const VarNode* node) final { + ext_func_args_.push_back(GetRef(node)); + Output output; + output.name = node->name_hint(); + return {output}; + } + + std::vector VisitExpr_(const CallNode* call) final { + const auto* func = call->op.as(); + ICHECK(func) << "Only composite function is supported for MCTLASS."; + GenerateBodyOutput ret = GenerateCompositeFunctionCall(func, call); + ext_func_body_.push_back(ret.decl); + headers_ = ret.headers; + return ret.outputs; + } + + std::string JIT(const std::vector& out) { + std::vector arg_names; + for (const auto& arg : ext_func_args_) { + arg_names.push_back(arg->name_hint()); + } + + code_stream_ << EmitSignature(out, ext_func_id_, arg_names) << "{\n"; + + this->EnterScope(); + + // Function body + for (auto decl : buf_decl_) { + this->PrintIndents(); + code_stream_ << decl << "\n"; + } + code_stream_ << "\n"; + for (auto stmt : ext_func_body_) { + this->PrintIndents(); + code_stream_ << stmt << "\n"; + } + + this->ExitScope(); + code_stream_ << "}\n"; + + this->GenerateBackendCFunc(ext_func_id_, ext_func_args_, /*const_arr_name=*/"", out, true); + return code_stream_.str(); + } + + Array GetHeaders() { return headers_; } + + private: + Array GetArgumentNames(const CallNode* call) { + Array arg_names; + for (size_t i = 0; i < call->args.size(); ++i) { + auto res = VisitExpr(call->args[i]); + for (const auto& out : res) { + arg_names.push_back(out.name); + } + } + return arg_names; + } + + // Is node `x` an ancestor of `y`? + bool IsAncestor(const CallNode* x, const CallNode* y) { + if (x == y) return true; + for (auto arg : y->args) { + const CallNode* arg_ptr = arg.as(); + if (arg_ptr && IsAncestor(x, arg_ptr)) return true; + } + return false; + } + + GenerateBodyOutput GenerateCompositeFunctionCall(const FunctionNode* callee, + const CallNode* caller) { + const auto pattern_name_opt = callee->GetAttr(attr::kComposite); + ICHECK(pattern_name_opt.defined()) << "Only functions with composite attribute are supported."; + const std::string pattern_name = pattern_name_opt.value(); + + if (pattern_name.find("conv2d") != std::string::npos && + pattern_name.find("residual") != std::string::npos) { + const CallNode* current_call = callee->body.as(); + bool has_relu = current_call->args.size() == 1; + const CallNode* binop = has_relu ? current_call->args[0].as() : current_call; + ICHECK(binop->args.size() == 2); + // Figure out which of the first or second argument corresponds to the residual input + // The root conv2d call can be reached via the other input of the binary op + int residual_index; + if (binop->args[1].as()) { + residual_index = 1; + } else if (binop->args[0].as()) { + residual_index = 0; + } else { + const CallNode* lhs = binop->args[0].as(); + const CallNode* rhs = binop->args[1].as(); + ICHECK(lhs && rhs); + // The residual input should be an ancestor of the non-residual input + residual_index = IsAncestor(rhs, lhs) ? 1 : 0; + } + const auto residual_input = binop->args[residual_index]; + auto call_args = GetArgumentNames(caller); + auto func_args = call_args; + if (call_args.size() == 3) { + // TODO(masahi): This code assumes that there is always a bias_add in a residual block. + for (size_t i = 0; i < call_args.size(); ++i) { + if (callee->params[i] == residual_input) { + auto residual_input_name = call_args[i]; + func_args.push_back(residual_input_name); + } + } + } else { + ICHECK_EQ(func_args.size(), 4) << "Residual block fusion expects 4 input tensors: data, " + "weight, bias, and residual tensor."; + } + return GenerateBody(caller, pattern_name, func_args, attrs_); + } else { + return GenerateBody(caller, pattern_name, attrs_); + } + + LOG(FATAL) << "Unknown composite function: " << pattern_name; + } + + GenerateBodyOutput GenerateBody(const CallNode* call, const std::string& func_name, + const Array& func_args, + const Map& attrs) { + std::vector out_types; + if (call->checked_type()->IsInstance()) { + auto type_node = call->checked_type().as(); + for (auto field : type_node->fields) { + ICHECK(field->IsInstance()); + out_types.push_back(field); + } + } else if (call->checked_type()->IsInstance()) { + ICHECK(call->checked_type()->IsInstance()); + out_types.push_back(call->checked_type()); + } else { + LOG(FATAL) << "Unrecognized type node: " << AsText(call->checked_type(), false); + } + + std::vector out_types_str; + for (const auto& out_type : out_types) { + out_types_str.push_back(GetDtypeString(out_type.as())); + } + + return mctlass::GenerateBody(func_name, ext_func_id_, out_types_str, func_args, attrs, + &buf_idx_); + } + + GenerateBodyOutput GenerateBody(const CallNode* call, const std::string& func_name, + const Map& attrs) { + auto func_args = GetArgumentNames(call); + return GenerateBody(call, func_name, func_args, attrs); + } + + /*! \brief The id of the external mctlass ext_func. */ + std::string ext_func_id_; + /*! \brief The attrs of the external mctlass ext_func. */ + Map attrs_; + /*! + * \brief The index to track the output buffer. Each kernel will redirect the + * output to a buffer that may be consumed by other kernels. + */ + int buf_idx_{0}; + /*! \brief The arguments used by a wrapped function that calls MCTLASS kernels. */ + Array ext_func_args_; + /*! \brief Statement of the function that will be compiled using MCTLASS kernels. */ + std::vector ext_func_body_; + /*! \brief The declaration of intermediate buffers. */ + std::vector buf_decl_; + /*! \brief Required header-file names. */ + Array headers_; +}; // class CodegenCutlass + +class CutlassModuleCodegen { + public: + explicit CutlassModuleCodegen(IRModule mod) : mod_(std::move(mod)) {} + + runtime::Module CreateCSourceModule() { + for (const auto& entry : mod_->functions) { + if (const auto* function_node = GetCutlassFunctionNode(entry.second)) { + GenCutlassFunc(GetRef(function_node)); + } + } + return Finalize(code_stream_.str(), func_names_); + } + + private: + void GenCutlassFunc(const Function& function) { + ICHECK(function.defined()) << "Input error: expect a Relay function."; + + // Record the external symbol for runtime lookup. + Optional opt_global_symbol = function->GetAttr(tvm::attr::kGlobalSymbol); + ICHECK(opt_global_symbol.defined()) + << "MCTLASS functions must have a " << tvm::attr::kGlobalSymbol << " attribute"; + std::string sid = opt_global_symbol.value(); + if (std::find(func_names_.begin(), func_names_.end(), sid) != func_names_.end()) { + // Already emitted. + return; + } + func_names_.push_back(sid); + + const auto* attrs = function->attrs.as(); + ICHECK(attrs != nullptr); + const auto dict = attrs->dict; + CodegenCutlass builder(sid, dict); + VLOG(1) << "Creating mctlass C code for '" << sid << "' from:\n" << PrettyPrint(function); + auto out = builder.VisitExpr(function->body); + auto code = builder.JIT(out); + for (const auto& header : builder.GetHeaders()) { + code_stream_ << "#include <" << header << ">\n"; + } + code_stream_ << "\n" + code; + } + + /*! + * \brief Returns \p expr as function if it is a \p Function with "Compiler" attribute + * value "mctlass". + */ + static const FunctionNode* GetCutlassFunctionNode(const Expr& expr) { + if (const auto* function_node = expr.as()) { + Optional opt_compiler = function_node->GetAttr(attr::kCompiler); + if (opt_compiler.defined() && opt_compiler.value() == "mctlass") { + return function_node; + } + } + return nullptr; + } + + /*! \brief Module we are compiling. */ + IRModule mod_; + /*! \brief The accumulated code stream that will be compiled by NVCC */ + std::ostringstream code_stream_; + /*! \brief The accumulated function names. */ + Array func_names_; +}; // CutlassModuleCodegen + +/*! + * \brief A small shim to redirect to the 'relay.ext.mctlass.compile_for_mctlass' Python + * function which does the main MCTLASS training, c-code generation and compilation steps. + */ +tvm::transform::Pass CompileForCutlassImpl() { + auto pass_func = [=](IRModule mod, const tvm::transform::PassContext& pass_ctx) { + VLOG(1) << "CompileForCutlass input:" << std::endl << PrettyPrint(mod); + const auto* pf = runtime::Registry::Get("relay.ext.mctlass.compile_for_mctlass"); + ICHECK(pf != nullptr) << "Cannot find compile_for_mctlass function"; + Target target = GetCutlassTarget(); + runtime::Module runtime_mod = (*pf)(mod, target); + Array external_mods = + mod->GetAttr>(tvm::attr::kExternalMods).value_or({}); + external_mods.push_back(runtime_mod); + return WithAttr(mod, tvm::attr::kExternalMods, external_mods); + }; + return tvm::transform::CreateModulePass(pass_func, 0, "CompileForCutlass", {}); +} + +runtime::Module CreateCSourceModule(const IRModule& mod) { + VLOG(1) << "Creating MCTLASS CSource module from:" << std::endl << PrettyPrint(mod); + return CutlassModuleCodegen(mod).CreateCSourceModule(); +} + +} // namespace + +TVM_REGISTER_GLOBAL("relay.ext.mctlass.create_c_source_module").set_body_typed(CreateCSourceModule); + +tvm::transform::Pass CompileForCutlass() { + return transform::Sequential( + {transform::OutlineCompilerFunctionsWithExistingGlobalSymbols("mctlass"), + CompileForCutlassImpl(), transform::MarkCompilerFunctionsAsExtern("mctlass")}); +} + +} // namespace mctlass +} // namespace contrib +} // namespace relay +} // namespace tvm diff --git a/src/relay/backend/contrib/mctlass/codegen.h b/src/relay/backend/contrib/mctlass/codegen.h new file mode 100644 index 000000000000..074b36a69423 --- /dev/null +++ b/src/relay/backend/contrib/mctlass/codegen.h @@ -0,0 +1,69 @@ +/* + * 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. + */ + +/*! + * \file src/relay/backend/contrib/mctlass/codegen.h + * \brief The 'custom' compilation pass for MCTLASS (invoked by the RelayToTIRTargetHook pass). + */ + +#ifndef TVM_RELAY_BACKEND_CONTRIB_MCTLASS_CODEGEN_H_ +#define TVM_RELAY_BACKEND_CONTRIB_MCTLASS_CODEGEN_H_ + +#include + +#include +#include + +#include "../codegen_c/codegen_c.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace mctlass { + +/*! + * \brief Returns the pass which replaces all calls to "Primitive" functions with "Compiler" + * attribute of "mctlass" with an call to an extern, and binds a \p runtime::StaticLibrary + * to the IRModule's "external_mods" attribute containing compiled implementations of + * those functions using the MCTLASS C++ template library. + */ +transform::Pass CompileForCutlass(); + +// The rest is sparsely documented since they are exposed only for code sharing between Relay +// and Relax backend implementations. + +/*! \brief Emit the function signature for a kernel */ +std::string EmitSignature(const std::vector& out, + const std::string& func_id, const std::vector& arg_names); + +/*! \brief Generate the body of the kernel */ +GenerateBodyOutput GenerateBody(const std::string& func_name, const std::string& ext_func_id, + const std::vector& output_types, + const Array& func_args, const Map& attrs, + int* buf_idx); + +/*! \brief Create a C-source module from the given kernel string */ +runtime::Module Finalize(const std::string& code, const Array& func_names); + +} // namespace mctlass +} // namespace contrib +} // namespace relay +} // namespace tvm + +#endif // TVM_RELAY_BACKEND_CONTRIB_MCTLASS_CODEGEN_H_ diff --git a/src/relay/backend/contrib/mctlass/target.cc b/src/relay/backend/contrib/mctlass/target.cc new file mode 100644 index 000000000000..0d635e868d4d --- /dev/null +++ b/src/relay/backend/contrib/mctlass/target.cc @@ -0,0 +1,74 @@ +/* + * 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. + */ + +/*! + * \file src/relay/backend/contrib/mctlass/target.cc + * \brief Registers the "mctlass" external codegen TargetKind. + */ + +#include + +#include "./codegen.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace mctlass { + +/*! + * \brief This external codegen target can use the CUTLASS template library included in + * TVM's 3rdparty/mctlass. + * - Patterns: python/tvm/relay/op/contrib/mctlass.py + * - Custom compiler: python/tvm/contrib/mctlass/build.py, + * src/relay/backend/contrib/mctlass/codegen.cc + */ +TVM_REGISTER_TARGET_KIND("mctlass", kDLMACA) + .set_attr(tvm::attr::kIsExternalCodegen, runtime::Bool(true)) + .set_attr("RelayToTIR", CompileForCutlass()) + // An integer specifying the compute capability. For example, 75 for Turing and + // 80 or 86 for Ampere. + .add_attr_option("sm", runtime::Int(80)) + // Whether to use slower but very accurate (compared to tf32) 3xtf32 mode for + // fp32 inputs on tensorcore. + .add_attr_option("use_3xtf32", runtime::Bool(true)) + // Split factor candidates for split-K GEMM. If split-K > 1, the GEMM K-loop is computed in + // parallel across split-K blocks, and a separate global reduction kernel is launched to + // accumulate partial reductions. The profiler will pick the best split-k factor from the + // given candidate list. Note that the larger split-K factor requires a larger workspace. + // Currently, parallel split-k has been tested only for wgrad. For GEMM and other conv2d + // kinds, split_k_slices is ignored. + .add_attr_option>("split_k_slices", Array{runtime::Int(1)}) + // When True, profile all kernel variants with smaller alignments than the largest possible. + .add_attr_option("profile_all_alignments", runtime::Bool(false)) + // Whether to profile all candidate kernels, or stop profiling after the first applicable kernel + // is found. + .add_attr_option("find_first_valid", runtime::Bool(false)) + // Whether to compile profiler executables for different kernels in parallel. + .add_attr_option("use_multiprocessing", runtime::Bool(false)) + // Number of threads to use during compilation, or -1 to use number of cpus. + .add_attr_option("threads", runtime::Int(-1)) + // Whether to replace sigmoid with tanh. + .add_attr_option("use_fast_math", runtime::Bool(false)) + // A temporary directory where intermediate compiled artifacts will be stored. + .add_attr_option("tmp_dir", String("./tmp")); + +} // namespace mctlass +} // namespace contrib +} // namespace relay +} // namespace tvm diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc index 79a41ae050c6..57b492e7948e 100644 --- a/src/relay/backend/te_compiler_cache.cc +++ b/src/relay/backend/te_compiler_cache.cc @@ -320,11 +320,11 @@ class LowerToTECompute : public backend::MemoizedExprTranslator outputs; - if (pattern_matcher_.find(op)) { if (pattern_matcher_.IsLeafOp(op)) { // Lower anchor op when pattern leaf op was reached auto anchor_op = pattern_matcher_.GetAnchorOp(); + LoweredOutput lowered_out = (*flower_call)(GetRef(anchor_op), inputs, target_, call_node->checked_type()); outputs = lowered_out->outputs; diff --git a/src/runtime/contrib/mcblas/mcblas.cc b/src/runtime/contrib/mcblas/mcblas.cc new file mode 100644 index 000000000000..317b6a5e46ec --- /dev/null +++ b/src/runtime/contrib/mcblas/mcblas.cc @@ -0,0 +1,570 @@ +/* + * 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. + */ + +/*! + * \file Use external cblas library call. + */ +#include +#include +#include + +#include "../../3rdparty/compiler-rt/builtin_fp16.h" +#include "../cblas/gemm_common.h" +#include "mcblas_utils.h" + +namespace tvm { +namespace contrib { + +using namespace runtime; +inline mcblasOperation_t MCBLASBooleanToTranspose(bool item) { + return item ? MCBLAS_OP_T : MCBLAS_OP_N; +} + +inline void MCBLASTryEnableTensorCore(mcblasHandle_t hdl) { + // TensorCores are only supported in mcblas 9.0 or higher + int version; + CHECK_MCBLAS_ERROR(mcblasGetVersion(hdl, &version)); + if (version >= 9000) CHECK_MCBLAS_ERROR(mcblasSetMathMode(hdl, MCBLAS_DEFAULT_MATH)); +} + +struct McblasHgemmOp { + typedef half TDatatype; + mcblasHandle_t handle; + explicit McblasHgemmOp(mcblasHandle_t hdl) : handle(hdl) {} + + void operator()(bool ta, bool tb, int M, int N, int K, half alpha, half* A, int lda, half* B, + int ldb, half beta, half* C, int ldc) { + CHECK_MCBLAS_ERROR(mcblasHgemm(handle, MCBLASBooleanToTranspose(ta), + MCBLASBooleanToTranspose(tb), M, N, K, &alpha, A, lda, B, ldb, + &beta, C, ldc)); + } +}; + +struct McblasSgemmOp { + typedef float TDatatype; + mcblasHandle_t handle; + explicit McblasSgemmOp(mcblasHandle_t hdl) : handle(hdl) {} + + void operator()(bool ta, bool tb, int M, int N, int K, float alpha, float* A, int lda, float* B, + int ldb, float beta, float* C, int ldc) { + CHECK_MCBLAS_ERROR(mcblasSgemm(handle, MCBLASBooleanToTranspose(ta), + MCBLASBooleanToTranspose(tb), M, N, K, &alpha, A, lda, B, ldb, + &beta, C, ldc)); + } +}; + +struct McblasDgemmOp { + typedef double TDatatype; + mcblasHandle_t handle; + explicit McblasDgemmOp(mcblasHandle_t hdl) : handle(hdl) {} + void operator()(bool ta, bool tb, int M, int N, int K, double alpha, double* A, int lda, + double* B, int ldb, double beta, double* C, int ldc) { + CHECK_MCBLAS_ERROR(mcblasDgemm(handle, MCBLASBooleanToTranspose(ta), + MCBLASBooleanToTranspose(tb), M, N, K, &alpha, A, lda, B, ldb, + &beta, C, ldc)); + } +}; + +struct McblasHgemmBatchOp { + typedef half TDatatype; + mcblasHandle_t handle; + explicit McblasHgemmBatchOp(mcblasHandle_t hdl) : handle(hdl) {} + void operator()(int batch_size, bool ta, bool tb, int M, int N, int K, half alpha, half* A, + int a_stride, int lda, half* B, int b_stride, int ldb, half beta, half* C, + int c_stride, int ldc) { + CHECK_MCBLAS_ERROR(mcblasHgemmStridedBatched( + handle, MCBLASBooleanToTranspose(ta), MCBLASBooleanToTranspose(tb), M, N, K, &alpha, A, lda, + a_stride, B, ldb, b_stride, &beta, C, ldc, c_stride, batch_size)); + } +}; + +struct McblasSgemmBatchOp { + typedef float TDatatype; + mcblasHandle_t handle; + explicit McblasSgemmBatchOp(mcblasHandle_t hdl) : handle(hdl) {} + void operator()(int batch_size, bool ta, bool tb, int M, int N, int K, float alpha, float* A, + int a_stride, int lda, float* B, int b_stride, int ldb, float beta, float* C, + int c_stride, int ldc) { + CHECK_MCBLAS_ERROR(mcblasSgemmStridedBatched( + handle, MCBLASBooleanToTranspose(ta), MCBLASBooleanToTranspose(tb), M, N, K, &alpha, A, lda, + a_stride, B, ldb, b_stride, &beta, C, ldc, c_stride, batch_size)); + } +}; + +struct McblasDgemmBatchOp { + typedef double TDatatype; + mcblasHandle_t handle; + explicit McblasDgemmBatchOp(mcblasHandle_t hdl) : handle(hdl) {} + void operator()(int batch_size, bool ta, bool tb, int M, int N, int K, double alpha, double* A, + int a_stride, int lda, double* B, int b_stride, int ldb, double beta, double* C, + int c_stride, int ldc) { + CHECK_MCBLAS_ERROR(mcblasDgemmStridedBatched( + handle, MCBLASBooleanToTranspose(ta), MCBLASBooleanToTranspose(tb), M, N, K, &alpha, A, lda, + a_stride, B, ldb, b_stride, &beta, C, ldc, c_stride, batch_size)); + } +}; + +// Check mcblas supported mix-precision computation type and return computeType +bool CheckMixPrecisionType(DLDataType in_dtype, DLDataType out_dtype, bool int_support = true) { + if (int_support && TypeMatch(out_dtype, kDLInt, 32)) { + return TypeMatch(in_dtype, kDLInt, 8); + } else if (TypeMatch(out_dtype, kDLFloat, 32)) { + return TypeMatch(in_dtype, kDLInt, 8) || TypeMatch(in_dtype, kDLFloat, 16); + } else { + return false; + } +} + +int roundoff(int v, int d) { return (v + d - 1) / d * d; } + +void CallMcblasLt(mcblasLtHandle_t hdl, mcStream_t stream, + mcblasLtMatmulPreference_t matmul_pref_desc, const DLTensor* A, const DLTensor* B, + const DLTensor* bias, const DLTensor* scaleA, const DLTensor* scaleB, + const DLTensor* C, bool transa, bool transb, void* workspace_ptr, + size_t workspace_size, mcblasLtEpilogue_t epilogue, + std::optional dq_scale) { + ICHECK(TypeEqual(A->dtype, B->dtype)); + // Reversed strides indicates an in-place transpose operation. + transa = IsInPlaceTransposed(A) ? !transa : transa; + transb = IsInPlaceTransposed(B) ? !transb : transb; + + auto compute_type = MCBLAS_COMPUTE_32F; + auto scale_type = MACA_R_32F; + macaDataType_t ab_type = MACA_R_32F; + macaDataType_t c_type = MACA_R_32F; + float one_fp32 = 1.0; + float zero_fp32 = 0.0; + int32_t one_i32 = 1; + int32_t zero_i32 = 0; + // Pass dequantization scale through the "alpha" parameter. If there is no dequantization after + // matmul, then alpha == 1.0 + float alpha_value = dq_scale.value_or(one_fp32); + void* alpha = &alpha_value; + void* beta = &zero_fp32; + + if (TypeMatch(A->dtype, kDLFloat, 16)) { + ab_type = MACA_R_16F; + } else if (TypeMatch(A->dtype, kDLInt, 8)) { + ab_type = MACA_R_8I; + } + // TODO: Support MACA_R_8F_E4M3 in mcblas + // else if (TypeMatch(A->dtype, DataType::TypeCode::kFloat8_e4m3fn, 8)) { + // ICHECK(TypeMatch(B->dtype, DataType::TypeCode::kFloat8_e4m3fn, 8)); + // ab_type = MACA_R_8F_E4M3; + // } + + if (TypeMatch(C->dtype, kDLFloat, 16)) { + c_type = MACA_R_16F; + } else if (TypeMatch(C->dtype, kDLInt, 32)) { + c_type = MACA_R_32I; + compute_type = MCBLAS_COMPUTE_32I; + scale_type = MACA_R_32I; + alpha = &one_i32; + beta = &zero_i32; + } + + mcblasLtMatmulDesc_t op_desc; + mcblasOperation_t op_transa = MCBLASBooleanToTranspose(transa); + mcblasOperation_t op_transb = MCBLASBooleanToTranspose(transb); + + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescCreate(&op_desc, compute_type, scale_type)); + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(op_desc, MCBLASLT_MATMUL_DESC_TRANSA, + &op_transb, sizeof(op_transb))); + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(op_desc, MCBLASLT_MATMUL_DESC_TRANSB, + &op_transa, sizeof(op_transa))); + + if (bias != nullptr) { + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(op_desc, MCBLASLT_MATMUL_DESC_BIAS_POINTER, + &bias->data, sizeof(float*))); + } + + if (scaleA != nullptr) { + auto scaleA_data = static_cast(scaleA->data) + scaleA->byte_offset; + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(op_desc, MCBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &scaleA_data, sizeof(float*))); + } + if (scaleB != nullptr) { + auto scaleB_data = static_cast(scaleB->data) + scaleB->byte_offset; + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(op_desc, MCBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &scaleB_data, sizeof(float*))); + } + + if (epilogue != MCBLASLT_EPILOGUE_DEFAULT) { + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(op_desc, MCBLASLT_MATMUL_DESC_EPILOGUE, + &epilogue, sizeof(epilogue))); + } + + int batch_offset_A = A->ndim - 2; + int batch_offset_B = B->ndim - 2; + + int M = ColumnCount(B, transb, batch_offset_B); + int N = RowCount(A, transa, batch_offset_A); + int K = ColumnCount(A, transa, batch_offset_A); + bool use_batched_gemm = A->ndim > 2 || B->ndim > 2; + + // If A is batched but B is not, flatten all non-reduction axes of A to use the regular GEMM. + // This trick is only applicable if batch axes and the other spatial axis (M or N) are + // adjacent in both the input and the output matrix. In particular, if A is of shape (M, K) + // and B matrix is of shape (Batch, N, K) with transb = true, the output shape + // is (Batch, M, N). Since the Batch and the N axes are not adjacent in the output, we cannot + // use the regular GEMM if only B is batched. + if (A->ndim > 2 && B->ndim == 2 && transa == false) { + N = 1; + for (int i = 0; i < A->ndim - 1; ++i) { + N *= A->shape[i]; + } + use_batched_gemm = false; + } + + int lda = transb ? K : M; + int ldb = transa ? N : K; + int ldc = M; + + mcblasLtMatrixLayout_t A_desc, B_desc, C_desc; + CHECK_MCBLAS_ERROR( + mcblasLtMatrixLayoutCreate(&A_desc, ab_type, !transb ? M : K, !transb ? K : M, lda)); + CHECK_MCBLAS_ERROR( + mcblasLtMatrixLayoutCreate(&B_desc, ab_type, !transa ? K : N, !transa ? N : K, ldb)); + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutCreate(&C_desc, c_type, M, N, ldc)); + + if (use_batched_gemm) { + auto get_batch_count = [](int64_t* shape, int batch_offset) { + int64_t count = 1; + for (int i = 0; i < batch_offset; ++i) { + count *= shape[i]; + } + return count; + }; + auto set_batch = [](mcblasLtMatrixLayout_t mat_desc, int batch_count, int64_t batch_stride) { + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutSetAttribute( + mat_desc, MCBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count))); + CHECK_MCBLAS_ERROR( + mcblasLtMatrixLayoutSetAttribute(mat_desc, MCBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, + &batch_stride, sizeof(batch_stride))); + }; + + int batch_count_A = get_batch_count(A->shape, batch_offset_A); + int batch_count_B = get_batch_count(B->shape, batch_offset_B); + int batch_count_C = get_batch_count(C->shape, C->ndim - 2); + int64_t batch_stride_A = M * K; + int64_t batch_stride_B = K * N; + int64_t batch_stride_C = M * N; + + // McblasLt does not seem to support batched GEMM with one of matrices having + // one batch (with batch_stride 0). + ICHECK_EQ(batch_count_A, batch_count_B); + + set_batch(A_desc, batch_count_A, batch_stride_A); + set_batch(B_desc, batch_count_B, batch_stride_B); + set_batch(C_desc, batch_count_C, batch_stride_C); + } + + auto A_data = static_cast(A->data) + A->byte_offset; + auto B_data = static_cast(B->data) + B->byte_offset; + auto C_data = static_cast(C->data) + C->byte_offset; + + mcblasLtMatmulPreferenceSetAttribute(matmul_pref_desc, MCBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &workspace_size, sizeof(size_t)); + + mcblasLtMatmulHeuristicResult_t heuristic_result = {}; + int returned_result = 0; + CHECK_MCBLAS_ERROR(mcblasLtMatmulAlgoGetHeuristic(hdl, op_desc, A_desc, B_desc, C_desc, C_desc, + matmul_pref_desc, 1, &heuristic_result, + &returned_result)); + if (returned_result == 0) { + CHECK_MCBLAS_ERROR(MCBLAS_STATUS_NOT_SUPPORTED); + } + + CHECK_MCBLAS_ERROR(mcblasLtMatmul(hdl, op_desc, alpha, B_data, A_desc, A_data, B_desc, beta, + C_data, C_desc, C_data, C_desc, &heuristic_result.algo, + workspace_ptr, workspace_size, stream)); + + mcblasLtMatmulDescDestroy(op_desc); + mcblasLtMatrixLayoutDestroy(A_desc); + mcblasLtMatrixLayoutDestroy(B_desc); + mcblasLtMatrixLayoutDestroy(C_desc); +} + +inline void CallLtIgemm(TVMArgs args, TVMRetValue* ret, mcblasLtHandle_t hdl, mcStream_t stream) { + DLTensor* A = args[0]; + DLTensor* B = args[1]; + DLTensor* C = args[2]; + bool transa = args[3]; + bool transb = args[4]; + // Reversed strides indicates an in-place transpose operation. + transa = IsInPlaceTransposed(A) ? !transa : transa; + transb = IsInPlaceTransposed(B) ? !transb : transb; + int M = ColumnCount(B, transb); + int N = RowCount(A, transa); + int K = ColumnCount(A, transa); + int N_out = ColumnCount(C, false); + int m = M; + int n = m; + int k = m; + int lda = M * K / (roundoff(K, 32) / 32); + int ldb = K * N / (roundoff(K, 32) / 32); + int ldc = M * N_out / (roundoff(N_out, 32) / 32); + ICHECK_EQ(A->ndim, 2); + ICHECK_EQ(B->ndim, 2); + ICHECK_EQ(C->ndim, 2); + + ICHECK_EQ(ElementStride(A), 1); + ICHECK_EQ(ElementStride(B), 1); + ICHECK_EQ(ElementStride(C), 1); + + ICHECK(TypeEqual(A->dtype, B->dtype)); + ICHECK(TypeMatch(A->dtype, kDLInt, 8)); + ICHECK(TypeMatch(C->dtype, kDLInt, 32)); + + ICHECK(CheckMixPrecisionType(A->dtype, C->dtype)) << "Unsupported data type"; + int32_t alpha = args.size() > 5 ? args[5] : 1; + int32_t beta = args.size() > 6 ? args[6] : 0; + mcblasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr; + auto A_data = reinterpret_cast(static_cast(A->data) + A->byte_offset); + auto B_data = reinterpret_cast(static_cast(B->data) + B->byte_offset); + auto C_data = reinterpret_cast(static_cast(C->data) + C->byte_offset); + + mcblasLtOrder_t order_COL32 = MCBLASLT_ORDER_COL32; + mcblasLtOrder_t order_COL4_4R2_8C = MCBLASLT_ORDER_COL4_4R2_8C; + mcblasLtMatmulDesc_t operationDesc = nullptr; + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescCreate(&operationDesc, MCBLAS_COMPUTE_32I, MACA_R_32I)); + mcblasOperation_t opTransA = MCBLASBooleanToTranspose(transa); + mcblasOperation_t opTransB = MCBLASBooleanToTranspose(transb); + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(operationDesc, MCBLASLT_MATMUL_DESC_TRANSA, + &opTransA, sizeof(opTransA))); + CHECK_MCBLAS_ERROR(mcblasLtMatmulDescSetAttribute(operationDesc, MCBLASLT_MATMUL_DESC_TRANSB, + &opTransB, sizeof(opTransB))); + // Create descriptors for the original matrices + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutCreate(&Adesc, MACA_R_8I, opTransA == MCBLAS_OP_N ? m : k, + opTransA == MCBLAS_OP_N ? k : m, lda)); + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutCreate(&Bdesc, MACA_R_8I, opTransB == MCBLAS_OP_N ? k : n, + opTransB == MCBLAS_OP_N ? n : k, ldb)); + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutCreate(&Cdesc, MACA_R_32I, m, n, ldc)); + + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutSetAttribute(Adesc, MCBLASLT_MATRIX_LAYOUT_ORDER, + &order_COL32, sizeof(order_COL32))); + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutSetAttribute( + Bdesc, MCBLASLT_MATRIX_LAYOUT_ORDER, &order_COL4_4R2_8C, sizeof(order_COL4_4R2_8C))); + CHECK_MCBLAS_ERROR(mcblasLtMatrixLayoutSetAttribute(Cdesc, MCBLASLT_MATRIX_LAYOUT_ORDER, + &order_COL32, sizeof(order_COL32))); + + CHECK_MCBLAS_ERROR(mcblasLtMatmul(hdl, operationDesc, &alpha, B_data, Adesc, A_data, Bdesc, &beta, + C_data, Cdesc, C_data, Cdesc, nullptr, nullptr, 0, stream)); +} + +inline void CallGemmEx(TVMArgs args, TVMRetValue* ret, mcblasHandle_t hdl) { + DLTensor* A = args[0]; + DLTensor* B = args[1]; + DLTensor* C = args[2]; + bool transa = args[3]; + bool transb = args[4]; + ICHECK_EQ(A->ndim, 2); + ICHECK_EQ(B->ndim, 2); + ICHECK_EQ(C->ndim, 2); + + ICHECK_EQ(ElementStride(A), 1); + ICHECK_EQ(ElementStride(B), 1); + ICHECK_EQ(ElementStride(C), 1); + + ICHECK(TypeEqual(A->dtype, B->dtype)); + + // C can never be transposed. + ICHECK(!IsInPlaceTransposed(C)); + + // Reversed strides indicates an in-place transpose operation. + transa = IsInPlaceTransposed(A) ? !transa : transa; + transb = IsInPlaceTransposed(B) ? !transb : transb; + + ICHECK(CheckMixPrecisionType(A->dtype, C->dtype)) << "Unsupported data type"; + ICHECK(!TypeMatch(A->dtype, kDLInt, 8) || ColumnStride(A) % 4 == 0) + << "leading dimension must divide 4 for int8 gemm"; + ICHECK(!TypeMatch(B->dtype, kDLInt, 8) || ColumnStride(B) % 4 == 0) + << "leading dimension must divide 4 for int8 gemm"; + double alpha = args.size() > 5 ? args[5] : 1.0; + double beta = args.size() > 6 ? args[6] : 0.0; + + macaDataType_t maca_in_type = GetMacaDataType(A->dtype); + macaDataType_t maca_out_type = GetMacaDataType(C->dtype); + mcblasGemmAlgo_t algo = MCBLAS_GEMM_DEFAULT; + void *alpha_ptr = nullptr, *beta_ptr = nullptr; + auto alpha_int = static_cast(alpha); + auto beta_int = static_cast(beta); + auto alpha_float = static_cast(alpha); + auto beta_float = static_cast(beta); + if (C->dtype.code == kDLInt) { + alpha_ptr = &alpha_int; + beta_ptr = &beta_int; + } else if (C->dtype.code == kDLFloat) { + alpha_ptr = &alpha_float; + beta_ptr = &beta_float; + } + + auto A_data = reinterpret_cast(static_cast(A->data) + A->byte_offset); + auto B_data = reinterpret_cast(static_cast(B->data) + B->byte_offset); + auto C_data = reinterpret_cast(static_cast(C->data) + C->byte_offset); + + CHECK_MCBLAS_ERROR( + mcblasGemmEx(hdl, MCBLASBooleanToTranspose(transb), MCBLASBooleanToTranspose(transa), + ColumnCount(B, transb), RowCount(A, transa), ColumnCount(A, transa), alpha_ptr, + B_data, maca_in_type, ColumnStride(B), A_data, maca_in_type, ColumnStride(A), + beta_ptr, C_data, maca_out_type, ColumnStride(C), maca_out_type, algo)); +} + +inline void CallBatchGemmEx(TVMArgs args, TVMRetValue* ret, mcblasHandle_t hdl) { + DLTensor* A = args[0]; + DLTensor* B = args[1]; + DLTensor* C = args[2]; + bool transa = args[3]; + bool transb = args[4]; + ICHECK_EQ(A->ndim, 3); + ICHECK_EQ(B->ndim, 3); + ICHECK_EQ(C->ndim, 3); + + int batch_size = BatchCount3D(C); + ICHECK_EQ(ElementStride3D(A), 1); + ICHECK_EQ(ElementStride3D(B), 1); + ICHECK_EQ(ElementStride3D(C), 1); + + ICHECK(TypeEqual(A->dtype, B->dtype)); + + // C can never be transposed. + ICHECK(!IsInPlaceTransposed3D(C)); + + // Reversed strides indicates an in-place transpose operation. + transa = IsInPlaceTransposed3D(A) ? !transa : transa; + transb = IsInPlaceTransposed3D(B) ? !transb : transb; + + ICHECK(CheckMixPrecisionType(A->dtype, C->dtype, true)) << "Unsupported data type"; + ICHECK(!TypeMatch(A->dtype, kDLInt, 8) || ColumnStride3D(A) % 4 == 0) + << "leading dimension must divide 4 for int8 gemm"; + ICHECK(!TypeMatch(B->dtype, kDLInt, 8) || ColumnStride3D(B) % 4 == 0) + << "leading dimension must divide 4 for int8 gemm"; + double alpha = args.size() > 5 ? args[5] : 1.0; + double beta = args.size() > 6 ? args[6] : 0.0; + + int A_stride = A->shape[1] * A->shape[2]; + int B_stride = B->shape[1] * B->shape[2]; + int C_stride = C->shape[1] * C->shape[2]; + + // Broadcast A or B by changing its stride. + int batch_size_a = BatchCount3D(A); + int batch_size_b = BatchCount3D(B); + if (batch_size_a != batch_size_b) { + if (batch_size_a == 1) { + A_stride = 0; + } else if (batch_size_b == 1) { + B_stride = 0; + } + } else { + ICHECK_EQ(batch_size_a, batch_size); + ICHECK_EQ(batch_size_b, batch_size); + } + + macaDataType_t maca_in_type = GetMacaDataType(A->dtype); + macaDataType_t maca_out_type = GetMacaDataType(C->dtype); + mcblasGemmAlgo_t algo = MCBLAS_GEMM_DEFAULT; + void *alpha_ptr = nullptr, *beta_ptr = nullptr; + auto alpha_int = static_cast(alpha); + auto beta_int = static_cast(beta); + auto alpha_float = static_cast(alpha); + auto beta_float = static_cast(beta); + if (C->dtype.code == kDLInt) { + alpha_ptr = &alpha_int; + beta_ptr = &beta_int; + } else if (C->dtype.code == kDLFloat) { + alpha_ptr = &alpha_float; + beta_ptr = &beta_float; + } + + auto A_data = reinterpret_cast(static_cast(A->data) + A->byte_offset); + auto B_data = reinterpret_cast(static_cast(B->data) + B->byte_offset); + auto C_data = reinterpret_cast(static_cast(C->data) + C->byte_offset); + CHECK_MCBLAS_ERROR(mcblasGemmStridedBatchedEx( + hdl, MCBLASBooleanToTranspose(transb), MCBLASBooleanToTranspose(transa), + ColumnCount3D(B, transb), RowCount3D(A, transa), ColumnCount3D(A, transa), alpha_ptr, B_data, + maca_in_type, ColumnStride3D(B), B_stride, A_data, maca_in_type, ColumnStride3D(A), A_stride, + beta_ptr, C_data, maca_out_type, ColumnStride3D(C), C_stride, batch_size, maca_out_type, + algo)); +} + +// matrix multiplication for row major +TVM_REGISTER_GLOBAL("tvm.contrib.mcblas.matmul").set_body([](TVMArgs args, TVMRetValue* ret) { + DLTensor* A = args[0]; + DLTensor* C = args[2]; + + McBlasThreadEntry* entry_ptr = McBlasThreadEntry::ThreadLocal(); + + MCBLASTryEnableTensorCore(entry_ptr->handle); + + if (TypeEqual(A->dtype, C->dtype)) { + ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) || + TypeMatch(A->dtype, kDLFloat, 64)); + + if (TypeMatch(A->dtype, kDLFloat, 16)) + CallGemm(args, ret, McblasHgemmOp(entry_ptr->handle)); + else if (TypeMatch(A->dtype, kDLFloat, 32)) + CallGemm(args, ret, McblasSgemmOp(entry_ptr->handle)); + else + CallGemm(args, ret, McblasDgemmOp(entry_ptr->handle)); + } else { + CallGemmEx(args, ret, entry_ptr->handle); + } +}); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcblaslt.matmul").set_body([](TVMArgs args, TVMRetValue* ret) { + DLTensor* A = args[0]; + + McBlasThreadEntry* entry_ptr = McBlasThreadEntry::ThreadLocal(); + + MCBLASTryEnableTensorCore(entry_ptr->handle); + + ICHECK(TypeMatch(A->dtype, kDLInt, 8)) << "Expects dtype to be int8\n"; + mcblasLtHandle_t ltHandle; + CHECK_MCBLAS_ERROR(mcblasLtCreate(<Handle)); + auto func = tvm::runtime::Registry::Get("runtime.get_maca_stream"); + ICHECK(func != nullptr); + mcStream_t stream = static_cast((*func)().operator void*()); + CallLtIgemm(args, ret, ltHandle, stream); + CHECK_MCBLAS_ERROR(mcblasLtDestroy(ltHandle)); +}); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcblas.batch_matmul").set_body([](TVMArgs args, TVMRetValue* ret) { + DLTensor* A = args[0]; + DLTensor* C = args[2]; + + McBlasThreadEntry* entry_ptr = McBlasThreadEntry::ThreadLocal(); + + MCBLASTryEnableTensorCore(entry_ptr->handle); + if (TypeEqual(A->dtype, C->dtype)) { + ICHECK(TypeMatch(A->dtype, kDLFloat, 16) || TypeMatch(A->dtype, kDLFloat, 32) || + TypeMatch(A->dtype, kDLFloat, 64)); + + if (TypeMatch(A->dtype, kDLFloat, 16)) + CallBatchGemm(args, ret, McblasHgemmBatchOp(entry_ptr->handle)); + else if (TypeMatch(A->dtype, kDLFloat, 32)) + CallBatchGemm(args, ret, McblasSgemmBatchOp(entry_ptr->handle)); + else + CallBatchGemm(args, ret, McblasDgemmBatchOp(entry_ptr->handle)); + } else { + CallBatchGemmEx(args, ret, entry_ptr->handle); + } +}); + +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/contrib/mcblas/mcblas_json_runtime.cc b/src/runtime/contrib/mcblas/mcblas_json_runtime.cc new file mode 100644 index 000000000000..13596de4b985 --- /dev/null +++ b/src/runtime/contrib/mcblas/mcblas_json_runtime.cc @@ -0,0 +1,166 @@ +/* + * 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. + */ + +/*! + * \file src/runtime/contrib/mcblas/mcblas_json_runtime.cc + * \brief A simple JSON runtime for MCBLAS. + */ + +#include +#include + +#include +#include +#include + +#include "../json/json_node.h" +#include "../json/json_runtime.h" +#include "mcblas_utils.h" + +namespace tvm { +namespace runtime { +namespace contrib { + +using namespace tvm::runtime; +using namespace tvm::runtime::json; + +class McblasJSONRuntime : public JSONRuntimeBase { + public: + McblasJSONRuntime(const std::string& symbol_name, const std::string& graph_json, + const Array const_names) + : JSONRuntimeBase(symbol_name, graph_json, const_names) {} + + void Init(const Array& consts) override {} + + PackedFunc GetFunction(const String& name, const ObjectPtr& sptr_to_self) override { + // JSONRuntimeBase::SetInputOutputBuffers(...) is not thread safe. Since McblasJSONRuntime + // can be used by multiple GPUs running on different threads, we avoid using that function + // and directly call mcBLAS on the inputs from TVMArgs. + if (this->symbol_name_ == name) { + return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { + ICHECK(this->initialized_) << "The module has not been initialized"; + this->Run(args); + }); + } else { + return JSONRuntimeBase::GetFunction(name, sptr_to_self); + } + } + + const char* type_key() const override { return "mcblas_json"; } // May be overridden + + void Run(TVMArgs args) { + auto* entry_ptr = tvm::contrib::McBlasLtThreadEntry::ThreadLocal(); + + auto func = tvm::runtime::Registry::Get("runtime.get_maca_stream"); + ICHECK(func != nullptr); + mcStream_t stream = static_cast((*func)().operator void*()); + + std::vector dl_tensors(NumEntries()); + + for (size_t i = 0; i < static_cast(args.size()); i++) { + auto eid = i < input_var_eid_.size() ? input_var_eid_[i] + : EntryID(outputs_[i - input_var_eid_.size()]); + ICHECK(args[i].type_code() == kTVMNDArrayHandle || args[i].type_code() == kTVMDLTensorHandle) + << "Expect NDArray or DLTensor as inputs"; + + const DLTensor* arg; + if (args[i].IsObjectRef()) { + NDArray arr = args[i]; + arg = arr.operator->(); + } else { + arg = args[i].operator DLTensor*(); + } + + dl_tensors[eid] = arg; + } + + auto get_input = [this, &dl_tensors](const JSONGraphNode& node, int idx) { + ICHECK_LT(idx, node.GetInputs().size()); + auto eid = EntryID(node.GetInputs()[idx]); + ICHECK(eid < dl_tensors.size()); + return dl_tensors[eid]; + }; + + auto get_inputs = [=](const JSONGraphNode& node, bool has_bias, bool has_scale) { + const DLTensor *bias = nullptr, *scaleA = nullptr, *scaleB = nullptr; + if (has_bias) { + bias = get_input(node, 2); + } else if (has_scale) { + scaleA = get_input(node, 2); + scaleB = get_input(node, 3); + } + return std::make_tuple(get_input(node, 0), get_input(node, 1), bias, scaleA, scaleB); + }; + + for (size_t i = 0; i < nodes_.size(); ++i) { + const auto& node = nodes_[i]; + if (node.GetOpType() == "kernel") { + auto op_name = node.GetOpName(); + uint32_t output_eid = EntryID(outputs_[0]); + auto out_ptr = dl_tensors[output_eid]; + bool transa = false; + bool transb = false; + mcblasLtEpilogue_t epilogue = MCBLASLT_EPILOGUE_DEFAULT; + + if (op_name.find("transposed") != std::string::npos) { + transb = true; + } + + if (op_name.find("relu") != std::string::npos) { + epilogue = MCBLASLT_EPILOGUE_RELU_BIAS; + } else if (op_name.find("gelu") != std::string::npos) { + epilogue = MCBLASLT_EPILOGUE_GELU_BIAS; + } else if (op_name.find("bias") != std::string::npos) { + epilogue = MCBLASLT_EPILOGUE_BIAS; + } + + bool has_scale = op_name.find("multiply") != std::string::npos; + auto [a_ptr, b_ptr, bias_ptr, scaleA_ptr, scaleB_ptr] = + get_inputs(node, epilogue != MCBLASLT_EPILOGUE_DEFAULT, has_scale); + + std::optional dq_scale = std::nullopt; + if (op_name.find("dequantize") != std::string::npos) { + dq_scale = std::stof(node.GetAttr>("dq_scale")[0]); + } + + tvm::contrib::CallMcblasLt(entry_ptr->handle, stream, entry_ptr->matmul_pref_desc, a_ptr, + b_ptr, bias_ptr, scaleA_ptr, scaleB_ptr, out_ptr, transa, transb, + entry_ptr->workspace_ptr, entry_ptr->workspace_size, epilogue, + dq_scale); + } + } + } + + void Run() override { LOG(FATAL) << "Unreachable"; } +}; + +runtime::Module McblasJSONRuntimeCreate(String symbol_name, String graph_json, + const Array& const_names) { + auto n = make_object(symbol_name, graph_json, const_names); + return runtime::Module(n); +} + +TVM_REGISTER_GLOBAL("runtime.McblasJSONRuntimeCreate").set_body_typed(McblasJSONRuntimeCreate); + +TVM_REGISTER_GLOBAL("runtime.module.loadbinary_mcblas_json") + .set_body_typed(JSONRuntimeBase::LoadFromBinary); + +} // namespace contrib +} // namespace runtime +} // namespace tvm diff --git a/src/runtime/contrib/mcblas/mcblas_utils.cc b/src/runtime/contrib/mcblas/mcblas_utils.cc new file mode 100644 index 000000000000..c51ae111eece --- /dev/null +++ b/src/runtime/contrib/mcblas/mcblas_utils.cc @@ -0,0 +1,77 @@ +/* + * 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. + */ + +/*! + * \file Use external mcdnn utils function + */ +#include "mcblas_utils.h" + +#include +#include + +#include "../../maca/maca_common.h" + +namespace tvm { +namespace contrib { + +McBlasThreadEntry::McBlasThreadEntry() { CHECK_MCBLAS_ERROR(mcblasCreate(&handle)); } + +McBlasThreadEntry::~McBlasThreadEntry() { + if (handle) { + mcblasDestroy(handle); + handle = nullptr; + } +} + +typedef dmlc::ThreadLocalStore McBlasThreadStore; + +McBlasThreadEntry* McBlasThreadEntry::ThreadLocal() { + auto stream = runtime::MACAThreadEntry::ThreadLocal()->stream; + McBlasThreadEntry* retval = McBlasThreadStore::Get(); + CHECK_MCBLAS_ERROR(mcblasSetStream(retval->handle, static_cast(stream))); + return retval; +} + +McBlasLtThreadEntry::McBlasLtThreadEntry() { + CHECK_MCBLAS_ERROR(mcblasLtCreate(&handle)); + CHECK_MCBLAS_ERROR(mcblasLtMatmulPreferenceCreate(&matmul_pref_desc)); + MACA_CALL(mcMalloc(&workspace_ptr, workspace_size)); +} + +McBlasLtThreadEntry::~McBlasLtThreadEntry() { + if (handle) { + mcblasLtDestroy(handle); + handle = nullptr; + } + if (matmul_pref_desc) { + mcblasLtMatmulPreferenceDestroy(matmul_pref_desc); + matmul_pref_desc = nullptr; + } + if (workspace_ptr != nullptr) { + mcFree(workspace_ptr); + workspace_ptr = nullptr; + } +} + +typedef dmlc::ThreadLocalStore McBlasLtThreadStore; + +McBlasLtThreadEntry* McBlasLtThreadEntry::ThreadLocal() { return McBlasLtThreadStore::Get(); } + +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/contrib/mcblas/mcblas_utils.h b/src/runtime/contrib/mcblas/mcblas_utils.h new file mode 100644 index 000000000000..1577c0de24d3 --- /dev/null +++ b/src/runtime/contrib/mcblas/mcblas_utils.h @@ -0,0 +1,131 @@ +/* + * 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. + */ + +/*! + * \file Use external mcdnn utils function + */ + +#ifndef TVM_RUNTIME_CONTRIB_MCBLAS_MCBLAS_UTILS_H_ +#define TVM_RUNTIME_CONTRIB_MCBLAS_MCBLAS_UTILS_H_ + +// #include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace tvm { +namespace contrib { + +inline const char* GetMcblasErrorString(int error) { + switch (error) { + case MCBLAS_STATUS_NOT_INITIALIZED: + return "MCBLAS_STATUS_NOT_INITIALIZED"; + case MCBLAS_STATUS_ALLOC_FAILED: + return "MCBLAS_STATUS_ALLOC_FAILED"; + case MCBLAS_STATUS_INVALID_VALUE: + return "MCBLAS_STATUS_INVALID_VALUE"; + case MCBLAS_STATUS_ARCH_MISMATCH: + return "MCBLAS_STATUS_ARCH_MISMATCH"; + case MCBLAS_STATUS_MAPPING_ERROR: + return "MCBLAS_STATUS_MAPPING_ERROR"; + case MCBLAS_STATUS_EXECUTION_FAILED: + return "MCBLAS_STATUS_EXECUTION_FAILED"; + case MCBLAS_STATUS_INTERNAL_ERROR: + return "MCBLAS_STATUS_INTERNAL_ERROR"; + case MCBLAS_STATUS_NOT_SUPPORTED: + return "MCBLAS_STATUS_NOT_SUPPORTED"; + case MCBLAS_STATUS_LICENSE_ERROR: + return "MCBLAS_STATUS_LICENSE_ERROR"; + } + return "Unrecognized error"; +} + +#ifndef CHECK_MCBLAS_ERROR +#define CHECK_MCBLAS_ERROR(fn) \ + do { \ + int error = static_cast(fn); \ + ICHECK_EQ(error, MCBLAS_STATUS_SUCCESS) << "MCBLAS: " << GetMcblasErrorString(error); \ + } while (0) // ; intentionally left off. +#endif // CHECK_MCBLAS_ERROR + +struct McBlasThreadEntry { + McBlasThreadEntry(); + ~McBlasThreadEntry(); + mcblasHandle_t handle{nullptr}; + static McBlasThreadEntry* ThreadLocal(); +}; // McBlasThreadEntry + +struct McBlasLtThreadEntry { + McBlasLtThreadEntry(); + ~McBlasLtThreadEntry(); + + mcblasLtHandle_t handle{nullptr}; + mcblasLtMatmulPreference_t matmul_pref_desc{nullptr}; + void* workspace_ptr{nullptr}; + static constexpr const size_t workspace_size = 33554432; + + static McBlasLtThreadEntry* ThreadLocal(); +}; // McBlasLtThreadEntry + +inline macaDataType_t GetMacaDataType(DLDataType type) { + if (type.code == kDLInt) { + switch (type.bits) { + case 8: + return MACA_R_8I; + case 32: + return MACA_R_32I; + } + } else if (type.code == kDLUInt) { + switch (type.bits) { + case 8: + return MACA_R_8U; + case 32: + return MACA_R_32U; + } + } else if (type.code == kDLFloat) { + switch (type.bits) { + case 16: + return MACA_R_16F; + case 32: + return MACA_R_32F; + case 64: + return MACA_R_64F; + } + } + LOG(FATAL) << "Unsupported maca type"; +} + +/*! \brief Execute matrix multiply followed by the specified epilogue, using mcBLASLt. */ +void CallMcblasLt(mcblasLtHandle_t hdl, mcStream_t stream, + mcblasLtMatmulPreference_t matmul_pref_desc, const DLTensor* A, const DLTensor* B, + const DLTensor* bias, const DLTensor* scaleA, const DLTensor* scaleB, + const DLTensor* C, bool transa, bool transb, void* workspace_ptr, + size_t workspace_size, mcblasLtEpilogue_t epilogue = MCBLASLT_EPILOGUE_DEFAULT, + std::optional dq_scale = std::nullopt); + +} // namespace contrib +} // namespace tvm + +#endif // TVM_RUNTIME_CONTRIB_MCBLAS_MCBLAS_UTILS_H_ diff --git a/src/runtime/contrib/mcdnn/conv_backward.cc b/src/runtime/contrib/mcdnn/conv_backward.cc new file mode 100644 index 000000000000..f1c91dc10a4f --- /dev/null +++ b/src/runtime/contrib/mcdnn/conv_backward.cc @@ -0,0 +1,271 @@ +/* + * 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. + */ + +/*! + * \file mcDNN kernel calls for backward algorithms. + */ +#include +#include +#include + +#include "mcdnn_utils.h" + +namespace tvm { +namespace contrib { +namespace maca { +using namespace runtime; + +void ConvolutionBackwardData(int mode, int format, int algo, int dims, int groups, const int pad[], + const int stride[], const int dilation[], DLTensor* dy, DLTensor* w, + DLTensor* dx, const std::string& conv_dtype) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + // Set Mode + entry_ptr->conv_entry.mode = static_cast(mode); + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, dx->shape, w->shape, + dy->shape, dy->dtype, conv_dtype); + // Set Device + entry_ptr->conv_entry.device = dy->device; + // Set Algo + entry_ptr->conv_entry.bwd_data_algo = static_cast(algo); + + // Set workspace + size_t workspace_size = 0; + MCDNN_CALL(mcdnnGetConvolutionBackwardDataWorkspaceSize( + entry_ptr->handle, entry_ptr->conv_entry.filter_desc, entry_ptr->conv_entry.output_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.input_desc, + entry_ptr->conv_entry.bwd_data_algo, &workspace_size)); + entry_ptr->conv_entry.UpdateWorkspace(workspace_size); + MCDNN_CALL(mcdnnConvolutionBackwardData( + entry_ptr->handle, McDNNDataType::GetConst<1>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.filter_desc, w->data, entry_ptr->conv_entry.output_desc, dy->data, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.bwd_data_algo, + entry_ptr->conv_entry.workspace, workspace_size, + McDNNDataType::GetConst<0>(entry_ptr->conv_entry.data_type), entry_ptr->conv_entry.input_desc, + dx->data)); +} + +void BackwardDataFindAlgo(int format, int dims, int groups, const int pad[], const int stride[], + const int dilation[], const int dy_dim[], const int w_dim[], + const int dx_dim[], const std::string& data_dtype, + const std::string& conv_dtype, bool verbose, TVMRetValue* ret) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + const int full_dims = dims + 2; + std::vector dy_dim_int64(full_dims); + std::vector w_dim_int64(full_dims); + std::vector dx_dim_int64(full_dims); + for (int i = 0; i < full_dims; ++i) { + dy_dim_int64[i] = dy_dim[i]; + w_dim_int64[i] = w_dim[i]; + dx_dim_int64[i] = dx_dim[i]; + } + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, dx_dim_int64.data(), + w_dim_int64.data(), dy_dim_int64.data(), String2DLDataType(data_dtype), + conv_dtype); + + int returned_algo_count = 0; + + mcdnnConvolutionBwdDataAlgoPerf_t perf_results[MCDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT]; + MCDNN_CALL(mcdnnFindConvolutionBackwardDataAlgorithm( + entry_ptr->handle, entry_ptr->conv_entry.filter_desc, entry_ptr->conv_entry.output_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.input_desc, + MCDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT, &returned_algo_count, perf_results)); + + const std::vector bwd_data_algo_names{ + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_0", // non-deterministic + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_1", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_FFT", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_BWD_DATA_ALGO_COUNT"}; + + auto best_algo = perf_results[0].algo; + if (verbose) { + LOG(INFO) << "\tMCDNN Found " << returned_algo_count << " bwd data algorithms, choosing " + << bwd_data_algo_names[best_algo]; + for (int i = 0; i < returned_algo_count; ++i) { + LOG(INFO) << "\t\t" << i << ") " << bwd_data_algo_names[perf_results[i].algo] + << " - time: " << perf_results[i].time << " ms" + << ", Memory: " << perf_results[i].memory; + } + } + ret[0] = best_algo; +} + +void ConvolutionBackwardFilter(int mode, int format, int algo, int dims, int groups, + const int pad[], const int stride[], const int dilation[], + DLTensor* dy, DLTensor* x, DLTensor* dw, + const std::string& conv_dtype) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + // Set Mode + entry_ptr->conv_entry.mode = static_cast(mode); + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, x->shape, dw->shape, + dy->shape, x->dtype, conv_dtype); + // Set Device + entry_ptr->conv_entry.device = x->device; + // Set Algo + entry_ptr->conv_entry.bwd_filter_algo = static_cast(algo); + + // Set workspace + size_t workspace_size = 0; + MCDNN_CALL(mcdnnGetConvolutionBackwardFilterWorkspaceSize( + entry_ptr->handle, entry_ptr->conv_entry.input_desc, entry_ptr->conv_entry.output_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.filter_desc, + entry_ptr->conv_entry.bwd_filter_algo, &workspace_size)); + entry_ptr->conv_entry.UpdateWorkspace(workspace_size); + MCDNN_CALL(mcdnnConvolutionBackwardFilter( + entry_ptr->handle, McDNNDataType::GetConst<1>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.input_desc, x->data, entry_ptr->conv_entry.output_desc, dy->data, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.bwd_filter_algo, + entry_ptr->conv_entry.workspace, workspace_size, + McDNNDataType::GetConst<0>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.filter_desc, dw->data)); +} + +void BackwardFilterFindAlgo(int format, int dims, int groups, const int pad[], const int stride[], + const int dilation[], const int dy_dim[], const int x_dim[], + const int dw_dim[], const std::string& data_dtype, + const std::string& conv_dtype, bool verbose, TVMRetValue* ret) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + const int full_dims = dims + 2; + std::vector x_dim_int64(full_dims); + std::vector dy_dim_int64(full_dims); + std::vector dw_dim_int64(full_dims); + for (int i = 0; i < full_dims; ++i) { + x_dim_int64[i] = x_dim[i]; + dy_dim_int64[i] = dy_dim[i]; + dw_dim_int64[i] = dw_dim[i]; + } + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, x_dim_int64.data(), + dw_dim_int64.data(), dy_dim_int64.data(), String2DLDataType(data_dtype), + conv_dtype); + + int returned_algo_count = 0; + + mcdnnConvolutionBwdFilterAlgoPerf_t perf_results[MCDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT]; + MCDNN_CALL(mcdnnFindConvolutionBackwardFilterAlgorithm( + entry_ptr->handle, entry_ptr->conv_entry.input_desc, entry_ptr->conv_entry.output_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.filter_desc, + MCDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT, &returned_algo_count, perf_results)); + + const std::vector bwd_filter_algo_names{ + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_0", // non-deterministic + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_1", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_3", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_BWD_FILTER_ALGO_COUNT", + }; + auto best_algo = perf_results[0].algo; + if (verbose) { + LOG(INFO) << "\tMCDNN Found " << returned_algo_count << " bwd filter algorithms, choosing " + << bwd_filter_algo_names[best_algo]; + for (int i = 0; i < returned_algo_count; ++i) { + LOG(INFO) << "\t\t" << i << ") " << bwd_filter_algo_names[perf_results[i].algo] + << " - time: " << perf_results[i].time << " ms" + << ", Memory: " << perf_results[i].memory; + } + } + ret[0] = best_algo; +} + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv2d.backward_data") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int mode = args[0]; + int format = args[1]; + int algo = args[2]; + int pad_v[2], stride_v[2], dilation_v[2]; + for (int i = 0; i < 2; i++) { + pad_v[i] = args[3 + i]; + stride_v[i] = args[5 + i]; + dilation_v[i] = args[7 + i]; + } + DLTensor* dy = args[9]; + DLTensor* w = args[10]; + DLTensor* dx = args[11]; + std::string conv_dtype = args[12]; + int groups = args[13]; + + ConvolutionBackwardData(mode, format, algo, 2, groups, pad_v, stride_v, dilation_v, dy, w, dx, + conv_dtype); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv.backward_data_find_algo") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int format = args[0]; + int dims = args[1]; + int* pad = static_cast(static_cast(args[2])); + int* stride = static_cast(static_cast(args[3])); + int* dilation = static_cast(static_cast(args[4])); + int* dy_dim = static_cast(static_cast(args[5])); + int* w_dim = static_cast(static_cast(args[6])); + int* dx_dim = static_cast(static_cast(args[7])); + std::string data_dtype = args[8]; + std::string conv_dtype = args[9]; + int groups = args[10]; + bool verbose = args[11]; + + BackwardDataFindAlgo(format, dims, groups, pad, stride, dilation, dy_dim, w_dim, dx_dim, + data_dtype, conv_dtype, verbose, ret); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv2d.backward_filter") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int mode = args[0]; + int format = args[1]; + int algo = args[2]; + int pad_v[2], stride_v[2], dilation_v[2]; + for (int i = 0; i < 2; i++) { + pad_v[i] = args[3 + i]; + stride_v[i] = args[5 + i]; + dilation_v[i] = args[7 + i]; + } + DLTensor* dy = args[9]; + DLTensor* x = args[10]; + DLTensor* dw = args[11]; + std::string conv_dtype = args[12]; + int groups = args[13]; + + ConvolutionBackwardFilter(mode, format, algo, 2, groups, pad_v, stride_v, dilation_v, dy, x, + dw, conv_dtype); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv.backward_filter_find_algo") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int format = args[0]; + int dims = args[1]; + int* pad = static_cast(static_cast(args[2])); + int* stride = static_cast(static_cast(args[3])); + int* dilation = static_cast(static_cast(args[4])); + int* dy_dim = static_cast(static_cast(args[5])); + int* x_dim = static_cast(static_cast(args[6])); + int* dw_dim = static_cast(static_cast(args[7])); + std::string data_dtype = args[8]; + std::string conv_dtype = args[9]; + int groups = args[10]; + bool verbose = args[11]; + + BackwardFilterFindAlgo(format, dims, groups, pad, stride, dilation, dy_dim, x_dim, dw_dim, + data_dtype, conv_dtype, verbose, ret); + }); +} // namespace maca +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/contrib/mcdnn/conv_forward.cc b/src/runtime/contrib/mcdnn/conv_forward.cc new file mode 100644 index 000000000000..aaeb1792391d --- /dev/null +++ b/src/runtime/contrib/mcdnn/conv_forward.cc @@ -0,0 +1,242 @@ +/* + * 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. + */ + +/*! + * \file mcDNN kernel calls for the forward algorithm. + */ +#include +#include +#include + +#include "mcdnn_utils.h" + +namespace tvm { +namespace contrib { +namespace maca { +using namespace runtime; + +void ConvolutionForward(int mode, int format, int algo, int dims, int groups, const int pad[], + const int stride[], const int dilation[], const DLTensor* x, + const DLTensor* w, const DLTensor* y, const std::string& conv_dtype) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + // Set Mode + entry_ptr->conv_entry.mode = static_cast(mode); + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, x->shape, w->shape, + y->shape, x->dtype, conv_dtype); + // Set Device + entry_ptr->conv_entry.device = x->device; + // Set Algo + entry_ptr->conv_entry.fwd_algo = static_cast(algo); + + // Set workspace + size_t workspace_size = 0; + // auto desc = static_cast(entry_ptr->conv_entry.conv_desc); + MCDNN_CALL(mcdnnGetConvolutionForwardWorkspaceSize( + entry_ptr->handle, entry_ptr->conv_entry.input_desc, entry_ptr->conv_entry.filter_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.output_desc, + entry_ptr->conv_entry.fwd_algo, &workspace_size)); + entry_ptr->conv_entry.UpdateWorkspace(workspace_size); + // Compute convolution + // auto &x_desc = static_cast(*(entry_ptr->conv_entry.input_desc)); + MCDNN_CALL(mcdnnConvolutionForward( + entry_ptr->handle, McDNNDataType::GetConst<1>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.input_desc, x->data, entry_ptr->conv_entry.filter_desc, w->data, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.fwd_algo, + entry_ptr->conv_entry.workspace, workspace_size, + McDNNDataType::GetConst<0>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.output_desc, y->data)); +} + +void ConvolutionBiasActivationForward(int mode, int format, int algo, int dims, int groups, int act, + double coef, const int pad[], const int stride[], + const int dilation[], const DLTensor* x, const DLTensor* w, + const DLTensor* y, const DLTensor* bias, + const std::string& conv_dtype) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + // Set Mode + entry_ptr->conv_entry.mode = static_cast(mode); + MCDNN_CALL(mcdnnSetActivationDescriptor(entry_ptr->conv_entry.activation_desc, + static_cast(act), + mcdnnNanPropagation_t::MCDNN_NOT_PROPAGATE_NAN, coef)); + MCDNN_CALL(mcdnnSetTensor4dDescriptor( + entry_ptr->conv_entry.bias_desc, entry_ptr->conv_entry.tensor_format, + McDNNDataType::DLTypeToMcDNNType(bias->dtype), 1, static_cast(w->shape[0]), 1, 1)); + + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, x->shape, w->shape, + y->shape, x->dtype, conv_dtype); + // Set Device + entry_ptr->conv_entry.device = x->device; + // Set Algo + entry_ptr->conv_entry.fwd_algo = static_cast(algo); + + // Set workspace + size_t workspace_size = 0; + MCDNN_CALL(mcdnnGetConvolutionForwardWorkspaceSize( + entry_ptr->handle, entry_ptr->conv_entry.input_desc, entry_ptr->conv_entry.filter_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.output_desc, + entry_ptr->conv_entry.fwd_algo, &workspace_size)); + + entry_ptr->conv_entry.UpdateWorkspace(workspace_size); + + // Compute convolution, add bias and apply activation + MCDNN_CALL(mcdnnConvolutionBiasActivationForward( + entry_ptr->handle, McDNNDataType::GetConst<1>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.input_desc, x->data, entry_ptr->conv_entry.filter_desc, w->data, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.fwd_algo, + entry_ptr->conv_entry.workspace, workspace_size, + McDNNDataType::GetConst<0>(entry_ptr->conv_entry.data_type), + entry_ptr->conv_entry.output_desc, y->data, entry_ptr->conv_entry.bias_desc, bias->data, + entry_ptr->conv_entry.activation_desc, entry_ptr->conv_entry.output_desc, y->data)); +} + +void FindAlgo(int format, int dims, int groups, const int pad[], const int stride[], + const int dilation[], const int x_dim[], const int w_dim[], const int y_dim[], + const std::string& data_dtype, const std::string& conv_dtype, bool verbose, + TVMRetValue* ret) { + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + const int full_dims = dims + 2; + std::vector x_dim_int64(full_dims); + std::vector w_dim_int64(full_dims); + std::vector y_dim_int64(full_dims); + for (int i = 0; i < full_dims; ++i) { + x_dim_int64[i] = x_dim[i]; + w_dim_int64[i] = w_dim[i]; + y_dim_int64[i] = y_dim[i]; + } + SetConvDescriptors(entry_ptr, format, dims, groups, pad, stride, dilation, x_dim_int64.data(), + w_dim_int64.data(), y_dim_int64.data(), String2DLDataType(data_dtype), + conv_dtype); + + int returned_algo_count = 0; + mcdnnConvolutionFwdAlgoPerf_t perf_results[MCDNN_CONVOLUTION_FWD_ALGO_COUNT]; + MCDNN_CALL(mcdnnFindConvolutionForwardAlgorithm( + entry_ptr->handle, entry_ptr->conv_entry.input_desc, entry_ptr->conv_entry.filter_desc, + entry_ptr->conv_entry.conv_desc, entry_ptr->conv_entry.output_desc, + MCDNN_CONVOLUTION_FWD_ALGO_COUNT, &returned_algo_count, perf_results)); + + const std::vector fwd_algo_names{"MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_GEMM", + "MCDNN_CONVOLUTION_FWD_ALGO_DIRECT", + "MCDNN_CONVOLUTION_FWD_ALGO_FFT", + "MCDNN_CONVOLUTION_FWD_ALGO_FFT_TILING", + "MCDNN_CONVOLUTION_FWD_ALGO_WINOGRAD", + "MCDNN_CONVOLUTION_FWD_ALGO_WINOGRAD_NONFUSED", + "MCDNN_CONVOLUTION_FWD_ALGO_COUNT"}; + + auto best_algo = perf_results[0].algo; + if (verbose) { + LOG(INFO) << "\tMCDNN Found " << returned_algo_count << " fwd algorithms, choosing " + << fwd_algo_names[best_algo]; + for (int i = 0; i < returned_algo_count; ++i) { + LOG(INFO) << "\t\t" << i << ") " << fwd_algo_names[perf_results[i].algo] + << " - time: " << perf_results[i].time << " ms" + << ", Memory: " << perf_results[i].memory; + } + } + + ret[0] = best_algo; +} + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv2d.forward") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int mode = args[0]; + int format = args[1]; + int algo = args[2]; + int pad_v[2], stride_v[2], dilation_v[2]; + for (int i = 0; i < 2; i++) { + pad_v[i] = args[3 + i]; + stride_v[i] = args[5 + i]; + dilation_v[i] = args[7 + i]; + } + DLTensor* x = args[9]; + DLTensor* w = args[10]; + DLTensor* y = args[11]; + std::string conv_dtype = args[12]; + int groups = args[13]; + + ConvolutionForward(mode, format, algo, 2, groups, pad_v, stride_v, dilation_v, x, w, y, + conv_dtype); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv2d+bias+act.forward") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int mode = args[0]; + int format = args[1]; + int algo = args[2]; + int pad_v[2], stride_v[2], dilation_v[2]; + for (int i = 0; i < 2; i++) { + pad_v[i] = args[3 + i]; + stride_v[i] = args[5 + i]; + dilation_v[i] = args[7 + i]; + } + int act = args[9]; + double coef = args[10]; + DLTensor* x = args[11]; + DLTensor* w = args[12]; + DLTensor* bias = args[13]; + DLTensor* y = args[14]; + std::string conv_dtype = args[15]; + int groups = args[16]; + + ConvolutionBiasActivationForward(mode, format, algo, 2, groups, act, coef, pad_v, stride_v, + dilation_v, x, w, y, bias, conv_dtype); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv3d.forward") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int mode = args[0]; + int format = args[1]; + int algo = args[2]; + int pad_v[3], stride_v[3], dilation_v[3]; + for (int i = 0; i < 3; i++) { + pad_v[i] = args[3 + i]; + stride_v[i] = args[6 + i]; + dilation_v[i] = args[9 + i]; + } + DLTensor* x = args[12]; + DLTensor* w = args[13]; + DLTensor* y = args[14]; + std::string conv_dtype = args[15]; + int groups = args[16]; + + ConvolutionForward(mode, format, algo, 3, groups, pad_v, stride_v, dilation_v, x, w, y, + conv_dtype); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.conv.forward_find_algo") + .set_body([](TVMArgs args, TVMRetValue* ret) { + int format = args[0]; + int dims = args[1]; + int* pad = static_cast(static_cast(args[2])); + int* stride = static_cast(static_cast(args[3])); + int* dilation = static_cast(static_cast(args[4])); + int* x_dim = static_cast(static_cast(args[5])); + int* w_dim = static_cast(static_cast(args[6])); + int* y_dim = static_cast(static_cast(args[7])); + std::string data_dtype = args[8]; + std::string conv_dtype = args[9]; + int groups = args[10]; + bool verbose = args[11]; + FindAlgo(format, dims, groups, pad, stride, dilation, x_dim, w_dim, y_dim, data_dtype, + conv_dtype, verbose, ret); + }); +} // namespace maca +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/contrib/mcdnn/mcdnn_frontend/attention.cc b/src/runtime/contrib/mcdnn/mcdnn_frontend/attention.cc new file mode 100644 index 000000000000..4d526b4e5ef5 --- /dev/null +++ b/src/runtime/contrib/mcdnn/mcdnn_frontend/attention.cc @@ -0,0 +1,124 @@ +/* + * 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. + */ + +/*! + * \file src/runtime/contrib/mcdnn/mcdnn_frontend/attention.cc + * \brief mcDNN scale dot product attention implementation + */ + +#include "./attention.h" + +#include +#include + +#include "../../../maca/maca_common.h" +#include "../mcdnn_utils.h" + +namespace tvm { +namespace contrib { +namespace maca { +void McDNNSDPARunnerNode::Init(int64_t batch, int64_t seq_len, int64_t num_heads, + int64_t num_kv_heads, int64_t head_size, int64_t head_size_v, + double scale, const DLDataType& data_type, + const std::string& layout) { + graph_ = std::make_unique(); + + CHECK(data_type.code == DLDataTypeCode::kDLFloat && data_type.bits == 16) + << "Only float16 is supported"; + + graph_->set_io_data_type(mcdnn_frontend::DataType_t::HALF) + .set_intermediate_data_type(mcdnn_frontend::DataType_t::FLOAT) + .set_compute_data_type(mcdnn_frontend::DataType_t::FLOAT); + + auto q_desc = mcdnn_frontend::graph::Tensor_attributes().set_name("Q").set_uid(kTensorIDQ); + auto k_desc = mcdnn_frontend::graph::Tensor_attributes().set_name("K").set_uid(kTensorIDK); + auto v_desc = mcdnn_frontend::graph::Tensor_attributes().set_name("V").set_uid(kTensorIDV); + auto o_desc = mcdnn_frontend::graph::Tensor_attributes().set_name("Out").set_uid(kTensorIDOut); + + std::vector q_stride, k_stride, v_stride, + o_stride; // stride in the order of (batch, num_heads, seq_len, head_size) + + if (layout == "BS3NH") { + int64_t stride_H = 1; + int64_t q_stride_N = head_size; + int64_t k_stride_N = head_size; + int64_t v_stride_N = head_size_v; + int64_t stride_S = + num_heads * q_stride_N + num_kv_heads * k_stride_N + num_kv_heads * v_stride_N; + int64_t stride_B = stride_S * seq_len; + q_stride = {stride_B, q_stride_N, stride_S, stride_H}; + k_stride = {stride_B, k_stride_N, stride_S, stride_H}; + v_stride = {stride_B, v_stride_N, stride_S, stride_H}; + o_stride = {seq_len * num_heads * head_size_v, head_size_v, num_heads * head_size_v, 1}; + offset_k_ = num_heads * head_size; + offset_v_ = offset_k_ + num_kv_heads * head_size; + } else if (layout == "SBN3H") { + CHECK_EQ(num_kv_heads, num_heads); + int64_t stride_H = 1; + int64_t stride_N = head_size + head_size + head_size_v; + int64_t stride_B = num_heads * stride_N; + int64_t stride_S = stride_B * batch; + q_stride = k_stride = v_stride = {stride_B, stride_N, stride_S, stride_H}; + o_stride = {num_heads * head_size_v, head_size_v, num_heads * head_size_v * batch, 1}; + offset_k_ = head_size; + offset_v_ = offset_k_ * 2; + } else { + LOG(FATAL) << "Unsupported layout: " << layout; + } + + q_desc = q_desc.set_dim({batch, num_heads, seq_len, head_size}).set_stride(q_stride); + k_desc = k_desc.set_dim({batch, num_kv_heads, seq_len, head_size}).set_stride(k_stride); + v_desc = v_desc.set_dim({batch, num_kv_heads, seq_len, head_size_v}).set_stride(v_stride); + auto sdpa_options = mcdnn_frontend::graph::SDPA_attributes() + .set_name("flash_attention") + .set_is_inference(true) + .set_alibi_mask(false) + .set_causal_mask(false) + .set_attn_scale(scale); + + auto q = graph_->tensor(q_desc); + auto k = graph_->tensor(k_desc); + auto v = graph_->tensor(v_desc); + auto [o, stats] = graph_->sdpa(q, k, v, sdpa_options); + CHECK(stats == nullptr); + o->set_output(true).set_dim({batch, num_heads, seq_len, head_size_v}).set_stride(o_stride); + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + MCDNN_FRONTEND_CALL(graph_->build(entry_ptr->handle, {mcdnn_frontend::HeurMode_t::A})); +} + +void McDNNSDPARunnerNode::Run(const DLTensor* qkv, DLTensor* workspace, DLTensor* out) { + MCDNN_CALL( + mcdnnSetStream(McDNNThreadEntry::ThreadLocal()->handle, tvm::runtime::GetMCDAStream())); + auto* qkv_base = reinterpret_cast(qkv->data) + qkv->byte_offset; + auto* q_ptr = reinterpret_cast(qkv_base) + offset_q_; + auto* k_ptr = reinterpret_cast(qkv_base) + offset_k_; + auto* v_ptr = reinterpret_cast(qkv_base) + offset_v_; + auto* out_ptr = reinterpret_cast(out->data) + out->byte_offset; + + size_t workspace_size = graph_->get_workspace_size(); + CHECK_LE(workspace_size, workspace->shape[0]) << "Workspace size too small"; + std::unordered_map inputs = { + {kTensorIDQ, q_ptr}, {kTensorIDK, k_ptr}, {kTensorIDV, v_ptr}, {kTensorIDOut, out_ptr}}; + + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + MCDNN_FRONTEND_CALL(graph_->execute(entry_ptr->handle, inputs, workspace->data)); +} +} // namespace maca +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/contrib/mcdnn/mcdnn_frontend/attention.h b/src/runtime/contrib/mcdnn/mcdnn_frontend/attention.h new file mode 100644 index 000000000000..4cfbe331f191 --- /dev/null +++ b/src/runtime/contrib/mcdnn/mcdnn_frontend/attention.h @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICUCE 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. + */ + +/*! + * \file src/runtime/contrib/mcdnn/mcdnn_frontend/attention.h + * \brief mcDNN scale dot product attention implementation + */ + +#ifndef TVM_RUNTIME_CONTRIB_MCDNN_MCDNN_FRONTEND_ATTENTION_H_ +#define TVM_RUNTIME_CONTRIB_MCDNN_MCDNN_FRONTEND_ATTENTION_H_ + +#include +#include + +#include +#include + +#define MCDNN_FRONTEND_CALL(func) \ + do { \ + auto status = (func); \ + CHECK(status.is_good()) << status.get_message(); \ + } while (0) + +namespace tvm { +namespace contrib { +namespace maca { +class McDNNSDPARunnerNode : public tvm::runtime::Object { + public: + McDNNSDPARunnerNode() {} + + ~McDNNSDPARunnerNode() {} + + static constexpr const char* _type_key = "contrib.mcdnn.SDPARunner"; + + void Init(int64_t batch, int64_t seq_len, int64_t num_heads, int64_t num_kv_heads, + int64_t head_size, int64_t head_size_v, double scale, const DLDataType& data_type, + const std::string& layout); + + void Run(const DLTensor* qkv, DLTensor* workspace, DLTensor* out); + + static constexpr int kTensorIDQ = 0; + static constexpr int kTensorIDK = 1; + static constexpr int kTensorIDV = 2; + static constexpr int kTensorIDOut = 4; + + private: + std::unique_ptr graph_{nullptr}; + int64_t offset_q_{0}; + int64_t offset_k_{0}; + int64_t offset_v_{0}; +}; + +class McDNNSDPARunner : public tvm::runtime::ObjectRef { + public: + static McDNNSDPARunner Create() { + auto n = make_object(); + return McDNNSDPARunner(n); + } + + TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(McDNNSDPARunner, tvm::runtime::ObjectRef, + McDNNSDPARunnerNode); +}; +} // namespace maca +} // namespace contrib +} // namespace tvm + +#endif // TVM_RUNTIME_CONTRIB_MCDNN_MCDNN_FRONTEND_ATTENTION_H_ diff --git a/src/runtime/contrib/mcdnn/mcdnn_utils.cc b/src/runtime/contrib/mcdnn/mcdnn_utils.cc new file mode 100644 index 000000000000..ec2cdf12aca6 --- /dev/null +++ b/src/runtime/contrib/mcdnn/mcdnn_utils.cc @@ -0,0 +1,273 @@ +/* + * 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. + */ + +/*! + * \file Use external mcdnn utils function + */ + +#include "mcdnn_utils.h" + +#include +#include +#include + +#include +#include + +namespace tvm { +namespace contrib { +namespace maca { +// McDNN Data Type +mcdnnDataType_t McDNNDataType::DLTypeToMcDNNType(const DLDataType& dtype) { + switch (dtype.code) { + case kDLInt: + if (dtype.bits == 8 && dtype.lanes == 1) + return MCDNN_DATA_INT8; + else if (dtype.bits == 32 && dtype.lanes == 1) + return MCDNN_DATA_INT32; + else if (dtype.bits == 8 && dtype.lanes == 4) + return MCDNN_DATA_INT8x4; + else + LOG(FATAL) << "Unsupported type"; + break; + case kDLUInt: + LOG(FATAL) << "Unsupported type"; + break; + case kDLFloat: + if (dtype.bits == 32 && dtype.lanes == 1) + return MCDNN_DATA_FLOAT; + else if (dtype.bits == 64 && dtype.lanes == 1) + return MCDNN_DATA_DOUBLE; + else if (dtype.bits == 16 && dtype.lanes == 1) + return MCDNN_DATA_HALF; + else + LOG(FATAL) << "Unsupported type"; + break; + } + return MCDNN_DATA_FLOAT; +} + +template <> +const void* McDNNDataType::GetConst<0>(mcdnnDataType_t type) { + static const int int_v = 0; + static const float float_v = 0; + static const double double_v = 0; + if (type == MCDNN_DATA_FLOAT || type == MCDNN_DATA_HALF) { + return static_cast(&float_v); + } + if (type == MCDNN_DATA_DOUBLE) { + return static_cast(&double_v); + } + if (type == MCDNN_DATA_INT8 || type == MCDNN_DATA_INT32 || type == MCDNN_DATA_INT8x4) { + return static_cast(&int_v); + } + return nullptr; +} + +template <> +const void* McDNNDataType::GetConst<1>(mcdnnDataType_t type) { + static const int int_v = 1; + static const float float_v = 1.f; + static const double double_v = 1.f; + if (type == MCDNN_DATA_FLOAT || type == MCDNN_DATA_HALF) { + return static_cast(&float_v); + } + if (type == MCDNN_DATA_DOUBLE) { + return static_cast(&double_v); + } + if (type == MCDNN_DATA_INT8 || type == MCDNN_DATA_INT32 || type == MCDNN_DATA_INT8x4) { + return static_cast(&int_v); + } + return nullptr; +} + +// McDNNThreadEntry + +McDNNThreadEntry::McDNNThreadEntry() { + auto stream = runtime::MACAThreadEntry::ThreadLocal()->stream; + auto func = runtime::Registry::Get("device_api.maca"); + void* ret = (*func)(); + maca_api = static_cast(ret); + + // If no McDNN-capable device is present, allow the McDNNThreadEntry + // object to be created. This is needed for + // McDNNThreadEntry::exists. + { + mcdnnStatus_t create_res = mcdnnCreate(&handle); + if (create_res == MCDNN_STATUS_NOT_INITIALIZED) { + return; + } + MCDNN_CALL(create_res); + } + + MCDNN_CALL(mcdnnSetStream(handle, stream)); + conv_entry.maca_api = maca_api; +} + +McDNNThreadEntry::~McDNNThreadEntry() {} + +typedef dmlc::ThreadLocalStore McDNNThreadStore; + +McDNNThreadEntry* McDNNThreadEntry::ThreadLocal(bool check_exists) { + auto* res = McDNNThreadStore::Get(); + if (check_exists) { + ICHECK(res->exists()) << "MCDNN_STATUS_NOT_INITIALIZED"; + } + + return res; +} + +// ConvEntry + +ConvEntry::ConvEntry() { + MCDNN_CALL(mcdnnCreateConvolutionDescriptor(&conv_desc)); + MCDNN_CALL(mcdnnCreateFilterDescriptor(&filter_desc)); + MCDNN_CALL(mcdnnCreateTensorDescriptor(&input_desc)); + MCDNN_CALL(mcdnnCreateTensorDescriptor(&output_desc)); + MCDNN_CALL(mcdnnCreateTensorDescriptor(&bias_desc)); + MCDNN_CALL(mcdnnCreateActivationDescriptor(&activation_desc)); +} + +ConvEntry::~ConvEntry() { + MCDNN_CALL(mcdnnDestroyFilterDescriptor(filter_desc)); + MCDNN_CALL(mcdnnDestroyConvolutionDescriptor(conv_desc)); + MCDNN_CALL(mcdnnDestroyTensorDescriptor(input_desc)); + MCDNN_CALL(mcdnnDestroyTensorDescriptor(output_desc)); + MCDNN_CALL(mcdnnDestroyTensorDescriptor(bias_desc)); + MCDNN_CALL(mcdnnDestroyActivationDescriptor(activation_desc)); + CleanWorkspace(); +} + +void ConvEntry::UpdateWorkspace(const size_t wsize) { + if (workspace_size < wsize) { + if (workspace != nullptr) { + CleanWorkspace(); + } + workspace_size = wsize; + workspace = maca_api->AllocWorkspace(device, workspace_size); + } +} + +void ConvEntry::CleanWorkspace() { + if (workspace) maca_api->FreeWorkspace(device, workspace); + workspace_size = 0; +} + +void SetConvDescriptors(McDNNThreadEntry* entry_ptr, int format, int dims, int groups, + const int pad[], const int stride[], const int dilation[], int64_t x_dim[], + int64_t w_dim[], int64_t y_dim[], DLDataType data_dtype, + const std::string& conv_dtype) { + // Set Format + entry_ptr->conv_entry.tensor_format = static_cast(format); + // Set Data Type + entry_ptr->conv_entry.data_type = + McDNNDataType::DLTypeToMcDNNType(runtime::String2DLDataType(conv_dtype)); + + mcdnnDataType_t mcdnn_data_type = McDNNDataType::DLTypeToMcDNNType(data_dtype); + + // Dims includes N and C + int full_dims = dims + 2; + + std::vector dim(full_dims); + std::vector tensor_stride(full_dims); + + // Note: For 2D tenor, using ND setters causes MCDNN_STATUS_NOT_SUPPORTED error + // in following mcdnnGetConvolutionForwardWorkspaceSize() when data type is fp16, int + + // MCDNN_CALL(mcdnnSetConvolutionGroupCount(entry_ptr->conv_entry.conv_desc, groups)); + if (dims == 2) { + // Set Desc + MCDNN_CALL(mcdnnSetConvolution2dDescriptor( + entry_ptr->conv_entry.conv_desc, pad[0], pad[1], stride[0], stride[1], dilation[0], + dilation[1], entry_ptr->conv_entry.mode, entry_ptr->conv_entry.data_type)); + int ni, ci, hi, wi; + if (entry_ptr->conv_entry.tensor_format == MCDNN_TENSOR_NHWC) { + ni = 0; + ci = 3; + hi = 1; + wi = 2; + } else { + ni = 0; + ci = 1; + hi = 2; + wi = 3; + } + + // Set Input + MCDNN_CALL(mcdnnSetTensor4dDescriptor( + entry_ptr->conv_entry.input_desc, entry_ptr->conv_entry.tensor_format, mcdnn_data_type, + static_cast(x_dim[ni]), static_cast(x_dim[ci]), static_cast(x_dim[hi]), + static_cast(x_dim[wi]))); + // Set Filter + MCDNN_CALL(mcdnnSetFilter4dDescriptor( + entry_ptr->conv_entry.filter_desc, mcdnn_data_type, entry_ptr->conv_entry.tensor_format, + static_cast(w_dim[ni]), static_cast(w_dim[ci]), static_cast(w_dim[hi]), + static_cast(w_dim[wi]))); + // Set Output + MCDNN_CALL(mcdnnSetTensor4dDescriptor( + entry_ptr->conv_entry.output_desc, entry_ptr->conv_entry.tensor_format, mcdnn_data_type, + static_cast(y_dim[ni]), static_cast(y_dim[ci]), static_cast(y_dim[hi]), + static_cast(y_dim[wi]))); + } else { + ICHECK_EQ(format, 0) << "Use of layout MCDNN_TENSOR_NHWC is supported only for 4-D tensors."; + + MCDNN_CALL(mcdnnSetConvolutionNdDescriptor(entry_ptr->conv_entry.conv_desc, dims, pad, stride, + dilation, entry_ptr->conv_entry.mode, + entry_ptr->conv_entry.data_type)); + + // Set Filter + for (int i = 0; i < full_dims; i++) { + dim[i] = static_cast(w_dim[i]); + } + MCDNN_CALL(mcdnnSetFilterNdDescriptor(entry_ptr->conv_entry.filter_desc, mcdnn_data_type, + entry_ptr->conv_entry.tensor_format, full_dims, + dim.data())); + // Set Input + for (int i = 0; i < full_dims; i++) { + dim[i] = static_cast(x_dim[i]); + } + GetMcdnnStride(full_dims, dim.data(), tensor_stride.data()); + MCDNN_CALL(mcdnnSetTensorNdDescriptor(entry_ptr->conv_entry.input_desc, mcdnn_data_type, + full_dims, dim.data(), tensor_stride.data())); + // Set Output + for (int i = 0; i < full_dims; i++) { + dim[i] = static_cast(y_dim[i]); + } + GetMcdnnStride(full_dims, dim.data(), tensor_stride.data()); + MCDNN_CALL(mcdnnSetTensorNdDescriptor(entry_ptr->conv_entry.output_desc, mcdnn_data_type, + full_dims, dim.data(), tensor_stride.data())); + } + MCDNN_CALL(mcdnnSetConvolutionGroupCount(entry_ptr->conv_entry.conv_desc, groups)); + if (mcdnnGetVersion() > 7000) { + MCDNN_CALL(mcdnnSetConvolutionMathType(entry_ptr->conv_entry.conv_desc, MCDNN_TENSOR_OP_MATH)) + } +} + +// SoftmaxEntry + +SoftmaxEntry::SoftmaxEntry() { MCDNN_CALL(mcdnnCreateTensorDescriptor(&shape_desc)); } + +SoftmaxEntry::~SoftmaxEntry() { MCDNN_CALL(mcdnnDestroyTensorDescriptor(shape_desc)); } + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.exists").set_body_typed([]() -> bool { + return McDNNThreadEntry::ThreadLocal(false)->exists(); +}); +} // namespace maca +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/contrib/mcdnn/softmax.cc b/src/runtime/contrib/mcdnn/softmax.cc new file mode 100644 index 000000000000..d2ac28d64aee --- /dev/null +++ b/src/runtime/contrib/mcdnn/softmax.cc @@ -0,0 +1,89 @@ +/* + * 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. + */ + +/*! + * \file src/runtime/contrib/mcdnn/softmax.cc + * \brief Use external mcdnn softmax function + */ +#include +#include + +#include "mcdnn_utils.h" + +namespace tvm { +namespace contrib { +namespace maca { +using namespace runtime; + +void softmax_impl(mcdnnSoftmaxAlgorithm_t alg, TVMArgs args, TVMRetValue* ret) { + DLTensor* x = args[0]; + DLTensor* y = args[1]; + int axis = args[2]; + int ndim = x->ndim; + int64_t* shape = x->shape; + if (axis < 0) axis += ndim; + ICHECK(axis >= 0 && axis < ndim); + + McDNNThreadEntry* entry_ptr = McDNNThreadEntry::ThreadLocal(); + entry_ptr->softmax_entry.data_type = McDNNDataType::DLTypeToMcDNNType(x->dtype); + + // Set mode and shape descriptor + if (axis == ndim - 1) { + int64_t N = 1; + for (int i = 0; i < ndim - 1; ++i) { + N *= shape[i]; + } + entry_ptr->softmax_entry.mode = MCDNN_SOFTMAX_MODE_INSTANCE; + MCDNN_CALL(mcdnnSetTensor4dDescriptor(entry_ptr->softmax_entry.shape_desc, MCDNN_TENSOR_NCHW, + entry_ptr->softmax_entry.data_type, static_cast(N), + static_cast(shape[ndim - 1]), 1, 1)); + } else { + int64_t pre_axis_dim = 1; + int64_t post_axis_dim = 1; + for (int i = 0; i < ndim; ++i) { + if (i < axis) { + pre_axis_dim *= shape[i]; + } else if (i > axis) { + post_axis_dim *= shape[i]; + } + } + entry_ptr->softmax_entry.mode = MCDNN_SOFTMAX_MODE_CHANNEL; + MCDNN_CALL(mcdnnSetTensor4dDescriptor( + entry_ptr->softmax_entry.shape_desc, MCDNN_TENSOR_NCHW, entry_ptr->softmax_entry.data_type, + static_cast(pre_axis_dim), static_cast(shape[axis]), + static_cast(post_axis_dim), 1)); + } + + auto alpha = McDNNDataType::GetConst<1>(entry_ptr->softmax_entry.data_type); + auto beta = McDNNDataType::GetConst<0>(entry_ptr->softmax_entry.data_type); + MCDNN_CALL(mcdnnSoftmaxForward(entry_ptr->handle, alg, entry_ptr->softmax_entry.mode, alpha, + entry_ptr->softmax_entry.shape_desc, x->data, beta, + entry_ptr->softmax_entry.shape_desc, y->data)); +} + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.softmax.forward") + .set_body([](TVMArgs args, TVMRetValue* ret) { + softmax_impl(MCDNN_SOFTMAX_ACCURATE, args, ret); + }); + +TVM_REGISTER_GLOBAL("tvm.contrib.mcdnn.log_softmax.forward") + .set_body([](TVMArgs args, TVMRetValue* ret) { softmax_impl(MCDNN_SOFTMAX_LOG, args, ret); }); +} // namespace maca +} // namespace contrib +} // namespace tvm diff --git a/src/runtime/maca/maca_common.h b/src/runtime/maca/maca_common.h index 6e28dcc145ba..f8cac0b4e85b 100644 --- a/src/runtime/maca/maca_common.h +++ b/src/runtime/maca/maca_common.h @@ -61,6 +61,8 @@ class MACAThreadEntry { // get the threadlocal workspace static MACAThreadEntry* ThreadLocal(); }; + +inline mcStream_t GetMACAStream() { return MACAThreadEntry::ThreadLocal()->stream; } } // namespace runtime } // namespace tvm #endif // TVM_RUNTIME_MACA_MACA_COMMON_H_ diff --git a/src/runtime/maca/maca_device_api.cc b/src/runtime/maca/maca_device_api.cc index 83b250462b5e..2966d333ef99 100644 --- a/src/runtime/maca/maca_device_api.cc +++ b/src/runtime/maca/maca_device_api.cc @@ -279,5 +279,9 @@ TVM_REGISTER_GLOBAL("profiling.timer.maca").set_body_typed([](Device dev) { return Timer(make_object()); }); +TVM_REGISTER_GLOBAL("runtime.get_maca_stream").set_body_typed([]() { + return static_cast(MACAThreadEntry::ThreadLocal()->stream); +}); + } // namespace runtime } // namespace tvm diff --git a/src/support/libinfo.cc b/src/support/libinfo.cc index 86f0c300bb8e..12fd07b845ae 100644 --- a/src/support/libinfo.cc +++ b/src/support/libinfo.cc @@ -83,6 +83,14 @@ #define TVM_INFO_USE_MACA "NOT-FOUND" #endif +#ifndef TVM_INFO_USE_MCBLAS +#define TVM_INFO_USE_MCBLAS "NOT-FOUND" +#endif + +#ifndef TVM_INFO_USE_MCDNN +#define TVM_INFO_USE_MCDNN "NOT-FOUND" +#endif + #ifndef TVM_INFO_ROCM_PATH #define TVM_INFO_ROCM_PATH "NOT-FOUND" #endif @@ -379,6 +387,8 @@ TVM_DLL Map GetLibInfo() { {"USE_HIPBLAS", TVM_INFO_USE_HIPBLAS}, {"USE_ROCM", TVM_INFO_USE_ROCM}, {"USE_MACA", TVM_INFO_USE_MACA}, + {"USE_MCBLAS", TVM_INFO_USE_MCBLAS}, + {"USE_MCDNN", TVM_INFO_USE_MCDNN}, {"USE_RCCL", TVM_INFO_USE_RCCL}, {"USE_RPC", TVM_INFO_USE_RPC}, {"USE_RTTI", TVM_INFO_USE_RTTI}, diff --git a/src/topi/schedule.cc b/src/topi/schedule.cc index 0999f00ffd11..73a1d18905e0 100644 --- a/src/topi/schedule.cc +++ b/src/topi/schedule.cc @@ -44,6 +44,7 @@ #include #include #include +#include namespace tvm { namespace topi { @@ -315,5 +316,14 @@ TVM_REGISTER_GENERIC_FUNC(dense) .register_func({"cuda", "gpu"}, WrapDenseOp(topi::cuda::dense_cuda)) .register_func({"rocm"}, WrapDenseOp(topi::rocm::dense_rocm)); +/* MACA schedules */ +TVM_REGISTER_GLOBAL("topi.maca.dense_maca").set_body([](TVMArgs args, TVMRetValue* rv) { + *rv = maca::dense_maca(args[0], args[1], args[2], args[3], args[4]); +}); + +TVM_REGISTER_GLOBAL("topi.maca.schedule_dense").set_body([](TVMArgs args, TVMRetValue* rv) { + *rv = topi::maca::schedule_dense(args[0], args[1]); +}); + } // namespace topi } // namespace tvm diff --git a/tests/python/contrib/test_mcblas.py b/tests/python/contrib/test_mcblas.py new file mode 100644 index 000000000000..e8e33edbcf2e --- /dev/null +++ b/tests/python/contrib/test_mcblas.py @@ -0,0 +1,381 @@ +# 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. +import pytest + +import tvm +from tvm import te +from tvm import relay +import numpy as np +from tvm.contrib import mcblas +from tvm.contrib import mcblaslt +from tvm.contrib import graph_executor +import tvm.testing +from tvm.relay.op.contrib import get_pattern_table +from tvm.relay.op.contrib.mcblas import partition_for_mcblas + + +def verify_matmul_add(in_dtype, out_dtype, rtol=1e-5): + n = 1024 + l = 128 + m = 236 + A = te.placeholder((n, l), name="A", dtype=in_dtype) + B = te.placeholder((l, m), name="B", dtype=in_dtype) + C = mcblas.matmul(A, B, dtype=out_dtype) + s = te.create_schedule(C.op) + + def verify(target="maca"): + if not tvm.get_global_func("tvm.contrib.mcblas.matmul", True): + print("skip because extern function is not available") + return + dev = tvm.maca(0) + f = tvm.build(s, [A, B, C], target) + a = tvm.nd.array(np.random.uniform(0, 128, size=(n, l)).astype(A.dtype), dev) + b = tvm.nd.array(np.random.uniform(0, 128, size=(l, m)).astype(B.dtype), dev) + c = tvm.nd.array(np.zeros((n, m), dtype=C.dtype), dev) + f(a, b, c) + tvm.testing.assert_allclose( + c.numpy(), np.dot(a.numpy().astype(C.dtype), b.numpy().astype(C.dtype)), rtol=rtol + ) + + verify() + + +def roundoff(v, d): + return int(np.floor((v + d - 1) / d) * d) + + +def verify_matmul_add_igemm(in_dtype, out_dtype, rtol=1e-5): + n = 1024 + l = 1024 + m = 1024 + L = roundoff(l, 32) + N = roundoff(n, 8) + N_out = roundoff(n, 32) + + A = te.placeholder((N, L), name="A", dtype=in_dtype) + B = te.placeholder((m, L), name="B", dtype=in_dtype) + # C has MCBLASLT_ORDER_COL32 layout, thus a different shape + C = mcblaslt.matmul(A, B, False, True, m, N_out, dtype=out_dtype) + s = te.create_schedule(C.op) + + def verify(target="maca"): + if not tvm.get_global_func("tvm.contrib.mcblaslt.matmul", True): + print("skip because extern function is not available") + return + dev = tvm.maca(0) + f = tvm.build(s, [A, B, C], target) + a_old = np.random.uniform(0, 128, size=(n, l)) + b_old = np.random.uniform(0, 128, size=(l, m)) + + # Transform a to become MCBLASLT_ORDER_COL4_4R2_8C layout + a_new = np.hstack((a_old.astype(A.dtype), np.zeros([n, L - l]))) + a_new = np.vstack((a_new.astype(A.dtype), np.zeros([N - n, L]))) + a_even = np.vsplit(a_new[::2], N / 8) + a_odd = np.vsplit(a_new[1::2], N / 8) + a_new = [None] * (len(a_even) + len(a_odd)) + a_new[::2] = a_even + a_new[1::2] = a_odd + a_new = np.vstack(a_new) + a_new = np.vstack( + np.vstack(np.vstack(np.hsplit(i, 8)).reshape([4, 32]) for i in np.vsplit(j, N / 4)) + for j in np.hsplit(a_new, L / 32) + ) + a_new = a_new.reshape([N, L]) + # Transform b to become MCBLASLT_ORDER_COL32 layout + b_new = np.vstack( + np.hsplit(np.hstack((b_old.T.astype(B.dtype), np.zeros([m, L - l]))), L / 32) + ) + b_new = b_new.reshape([m, L]) + + a = tvm.nd.array(a_new.astype(A.dtype), dev) + b = tvm.nd.array(b_new.astype(B.dtype), dev) + c = tvm.nd.array(np.zeros((m, N_out), dtype=C.dtype), dev) + f(a, b, c) + # Transform output c from layout MCBLASLT_ORDER_COL32 to row major layout + c_out = c.numpy() + c_out = c_out.reshape([int(m * N_out / 32), 32]) + c_out = np.hstack(np.vsplit(c_out, int(N_out / 32))) + c_out = c_out[:, :n] + c_out = c_out.T + tvm.testing.assert_allclose( + c_out, np.dot(a_old.astype(C.dtype), b_old.astype(C.dtype)), rtol=rtol + ) + + verify() + + +def verify_batch_matmul(Ashape, Bshape, Cshape, in_dtype, out_dtype, rtol=1e-5): + A = te.placeholder(Ashape, name="A", dtype=in_dtype) + B = te.placeholder(Bshape, name="B", dtype=in_dtype) + C = mcblas.batch_matmul(A, B, dtype=out_dtype) + s = te.create_schedule(C.op) + + dev = tvm.maca(0) + f = tvm.build(s, [A, B, C], "maca") + + if "int" in in_dtype: + a = tvm.nd.array(np.random.uniform(1, 10, size=Ashape).astype(in_dtype), dev) + b = tvm.nd.array(np.random.uniform(1, 10, size=Bshape).astype(in_dtype), dev) + else: + a = tvm.nd.array(np.random.uniform(size=Ashape).astype(A.dtype), dev) + b = tvm.nd.array(np.random.uniform(size=Bshape).astype(B.dtype), dev) + + c = tvm.nd.array(np.zeros(Cshape, dtype=C.dtype), dev) + f(a, b, c) + tvm.testing.assert_allclose( + c.numpy(), + np.matmul(a.numpy().astype(C.dtype), b.numpy().astype(C.dtype)).astype(C.dtype), + rtol=rtol, + ) + + +@tvm.testing.requires_maca +def test_matmul_add(): + verify_matmul_add("float", "float", rtol=1e-3) + verify_matmul_add("float16", "float") + verify_matmul_add("float16", "float16", rtol=1e-2) + verify_matmul_add("int8", "int32") + +@pytest.mark.skip("fail") +@tvm.testing.requires_maca +def test_matmul_add_igemm(): + verify_matmul_add_igemm("int8", "int32") + + +@tvm.testing.requires_maca +def test_batch_matmul(): + if not tvm.get_global_func("tvm.contrib.mcblas.matmul", True): + print("skip because extern function is not available") + return + + verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float", "float") + verify_batch_matmul((16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float", "float") + verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float16", "float") + verify_batch_matmul((16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float16", "float") + verify_batch_matmul( + (16, 1024, 128), (16, 128, 236), (16, 1024, 236), "float16", "float16", rtol=1e-2 + ) + verify_batch_matmul( + (16, 1024, 128), (1, 128, 236), (16, 1024, 236), "float16", "float16", rtol=1e-2 + ) + + verify_batch_matmul((16, 1024, 128), (16, 128, 236), (16, 1024, 236), "int8", "int32") + + +def _verify_mcblas_relay(expr): + np.random.seed(42) + + mod = tvm.IRModule.from_expr(expr) + mod = relay.transform.InferType()(mod) + func = mod["main"] + mcblas_mod = partition_for_mcblas(mod) + assert len(mcblas_mod.get_global_vars()) == 2 + + input_data = [] + for param in func.params: + shape = [int(x) for x in param.checked_type.shape] + input_data.append( + (param.name_hint, np.random.uniform(0, 32, size=shape).astype(param.checked_type.dtype)) + ) + + # Test against CPU reference + maca_config = (tvm.target.maca(), tvm.maca(), mcblas_mod) + cpu_config = (tvm.target.Target("llvm"), tvm.cpu(), mod) + outputs = [] + for target, dev, test_mod in [maca_config, cpu_config]: + with tvm.transform.PassContext(opt_level=3): + lib = relay.build(test_mod, target=target, target_host=cpu_config[0]) + module = graph_executor.GraphModule(lib["default"](dev)) + for name, data in input_data: + module.set_input(name, tvm.nd.array(data, dev)) + + module.run() + out_type = func.body.checked_type + outputs.append( + module.get_output(0, tvm.nd.empty(out_type.shape, dtype=out_type.dtype)).numpy() + ) + + tvm.testing.assert_allclose( + outputs[0], + outputs[1], + rtol=1e-2, + ) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "n,m,k,transpose_a,transpose_b", + [ + (64, 128, 32, False, False), + (17, 32, 16, True, False), + (24, 17, 12, False, True), + (96, 4, 17, True, True), + ], +) +@pytest.mark.parametrize( + "in_dtype,out_dtype", + [ + ("float32", "float32"), + ("float16", "float16"), + ("float16", "float32"), + ("int8", "int32"), + ("float64", "float64"), + ("int8", "float32"), + ], +) +def test_relay_mcblas_matmul(n, m, k, in_dtype, out_dtype, transpose_a, transpose_b): + unsupported_configs = [ + (17, 32, 16, "int8", "float32", True, False), + (96, 4, 17, "int8", "float32", True, True), + (17, 32, 16, "int8", "int32", True, False), + (96, 4, 17, "int8", "int32", True, True), + ] + if (n, m, k, in_dtype, out_dtype, transpose_a, transpose_b) in unsupported_configs: + pytest.skip("Unsupported parameters.") + + a_shape = (k, n) if transpose_a else (n, k) + b_shape = (m, k) if transpose_b else (k, m) + a = tvm.relay.var("A", tvm.relay.TensorType(a_shape, in_dtype)) + b = tvm.relay.var("B", tvm.relay.TensorType(b_shape, in_dtype)) + # Directly use matmul because nn.matmul sometimes defers to nn.dense + matmul = relay.op.nn._make.matmul(a, b, None, out_dtype, transpose_a, transpose_b) + _verify_mcblas_relay(matmul) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "n,m,k", + [ + (64, 128, 32), + (17, 32, 16), + (24, 17, 12), + (96, 4, 17), + ], +) +@pytest.mark.parametrize( + "in_dtype,out_dtype", + [ + ("float32", "float32"), + ("float16", "float16"), + ("float16", "float32"), + ("int8", "int32"), + ("float64", "float64"), + ("int8", "float32"), + ], +) +def test_relay_mcblas_dense(n, m, k, in_dtype, out_dtype): + unsupported_configs = [ + (96, 4, 17, "int8", "float32"), + (96, 4, 17, "int8", "int32"), + ] + if (n, m, k, in_dtype, out_dtype) in unsupported_configs: + pytest.skip("Unsupported parameters.") + + data = tvm.relay.var("data", tvm.relay.TensorType((n, k), in_dtype)) + weight = tvm.relay.var("weight", tvm.relay.TensorType((m, k), in_dtype)) + dense = relay.op.nn.dense(data, weight, out_dtype=out_dtype) + _verify_mcblas_relay(dense) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "n,m,k,batch_a,batch_b,transpose_a,transpose_b", + [ + (64, 128, 32, 16, 16, False, False), + (17, 32, 16, 16, 1, True, False), + (24, 17, 12, 17, 17, False, True), + (96, 4, 17, 53, 1, True, True), + ], +) +@pytest.mark.parametrize( + "in_dtype,out_dtype", + [ + ("float32", "float32"), + ("float16", "float16"), + ("float16", "float32"), + ("int8", "int32"), + ("float64", "float64"), + ("int8", "float32"), + ], +) +def test_relay_mcblas_batch_matmul( + n, m, k, batch_a, batch_b, in_dtype, out_dtype, transpose_a, transpose_b +): + unsupported_configs = [ + (17, 32, 16, 16, 1, "int8", "float32", True, False), + (96, 4, 17, 53, 1, "int8", "float32", True, True), + (17, 32, 16, 16, 1, "int8", "int32", True, False), + (96, 4, 17, 53, 1, "int8", "int32", True, True), + ] + if ( + n, + m, + k, + batch_a, + batch_b, + in_dtype, + out_dtype, + transpose_a, + transpose_b, + ) in unsupported_configs: + pytest.skip("Unsupported parameters.") + + a_shape = (batch_a, k, n) if transpose_a else (batch_a, n, k) + b_shape = (batch_b, m, k) if transpose_b else (batch_b, k, m) + a = tvm.relay.var("A", tvm.relay.TensorType(a_shape, in_dtype)) + b = tvm.relay.var("B", tvm.relay.TensorType(b_shape, in_dtype)) + batch_matmul = relay.op.nn.batch_matmul(a, b, out_dtype, transpose_a, transpose_b) + _verify_mcblas_relay(batch_matmul) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "n,m,k", + [ + (64, 128, 32), + (17, 32, 16), + (24, 17, 12), + (96, 4, 17), + ], +) +@pytest.mark.parametrize( + "in_dtype,out_dtype", + [ + ("float32", "float32"), + ("float16", "float16"), + ("float16", "float32"), + ("int8", "int32"), + ("float64", "float64"), + ("int8", "float32"), + ], +) +def test_relay_mcblas_dense(n, m, k, in_dtype, out_dtype): + unsupported_configs = [ + (96, 4, 17, "int8", "float32"), + (96, 4, 17, "int8", "int32"), + ] + if (n, m, k, in_dtype, out_dtype) in unsupported_configs: + pytest.skip("Unsupported parameters.") + + data = tvm.relay.var("data", tvm.relay.TensorType((n, k), in_dtype)) + weight = tvm.relay.var("weight", tvm.relay.TensorType((m, k), in_dtype)) + dense = relay.op.nn.dense(data, weight, out_dtype=out_dtype) + _verify_mcblas_relay(dense) + + +if __name__ == "__main__": + tvm.testing.main() \ No newline at end of file diff --git a/tests/python/contrib/test_mcdnn.py b/tests/python/contrib/test_mcdnn.py new file mode 100644 index 000000000000..4abcca46229b --- /dev/null +++ b/tests/python/contrib/test_mcdnn.py @@ -0,0 +1,630 @@ +# 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. + +import sys + +import pytest + +import tvm +import tvm.testing +from tvm import te +from tvm import relay +from tvm.contrib import mcdnn +from tvm.contrib.mxcc import have_fp16 +from tvm.contrib import graph_executor +import numpy as np +import tvm.topi.testing +import tvm.testing +from tvm.relay.op.contrib.mcdnn import partition_for_mcdnn + + +requires_mcdnn = pytest.mark.skipif( + tvm.get_global_func("tvm.contrib.mcdnn.conv2d.forward", True) is None, + reason="McDNN is not enabled", +) + + +def verify_conv2d(data_dtype, conv_dtype, tensor_format=0, groups=1): + in_channel = 4 + out_channel = 16 + filter_h = 3 + filter_w = 3 + pad_h = 1 + pad_w = 1 + stride_h = 1 + stride_w = 1 + dilation_h = 1 + dilation_w = 1 + batch = 3 + height = 32 + width = 32 + + if data_dtype == "float16" and not have_fp16(tvm.maca(0).compute_version): + print("Skip because gpu does not have fp16 support") + return + + # schedule + if tensor_format == 0: + xshape = [batch, in_channel, height, width] + wshape = [out_channel, in_channel // groups, filter_h, filter_w] + else: + xshape = [batch, height, width, in_channel] + wshape = [out_channel, filter_h, filter_w, in_channel // groups] + + X = te.placeholder(xshape, name="X", dtype=data_dtype) + W = te.placeholder(wshape, name="W", dtype=data_dtype) + Y = mcdnn.conv_forward( + X, + W, + [pad_h, pad_w], + [stride_h, stride_w], + [dilation_h, dilation_w], + conv_mode=1, + tensor_format=tensor_format, + conv_dtype=conv_dtype, + algo=-1, + groups=groups, + ) + yshape = [x.value for x in Y.shape] + s = te.create_schedule(Y.op) + + # validation + dev = tvm.maca(0) + f = tvm.build(s, [X, W, Y], "maca --host=llvm", name="conv2d") + x_np = np.random.uniform(-1, 1, xshape).astype(data_dtype) + w_np = np.random.uniform(-1, 1, wshape).astype(data_dtype) + y_np = np.zeros(yshape).astype(data_dtype) + x = tvm.nd.array(x_np, dev) + w = tvm.nd.array(w_np, dev) + y = tvm.nd.array(y_np, dev) + if tensor_format == 0: + c_np = tvm.topi.testing.conv2d_nchw_python(x_np, w_np, 1, 1, groups=groups) + elif tensor_format == 1: + wt = w_np.transpose((1, 2, 3, 0)) # OHWI => HWIO + c_np = tvm.topi.testing.conv2d_nhwc_python(x_np, wt, 1, 1, groups=groups) + + f(x, w, y) + tvm.testing.assert_allclose(y.numpy(), c_np, atol=1e-2, rtol=1e-2) + + +@tvm.testing.requires_gpu +@requires_mcdnn +def test_conv2d(): + verify_conv2d("float32", "float32", tensor_format=0) + verify_conv2d("float16", "float32", tensor_format=1) + verify_conv2d("float16", "float16", tensor_format=0) + verify_conv2d("float16", "float16", tensor_format=1) + verify_conv2d("int8", "int32", tensor_format=1) + + verify_conv2d("float32", "float32", tensor_format=0, groups=2) + verify_conv2d("float16", "float32", tensor_format=1, groups=2) + verify_conv2d("float16", "float16", tensor_format=0, groups=2) + verify_conv2d("int8", "int32", tensor_format=1, groups=2) + + +def verify_conv3d(data_dtype, conv_dtype, tensor_format=0, groups=1): + in_channel = 4 + out_channel = 16 + filter_d = 3 + filter_h = 3 + filter_w = 3 + pad_d = 1 + pad_h = 1 + pad_w = 1 + stride_d = 1 + stride_h = 1 + stride_w = 1 + dilation_d = 1 + dilation_h = 1 + dilation_w = 1 + batch = 3 + depth = 32 + height = 32 + width = 32 + + # schedule + xshape = [batch, in_channel, depth, height, width] + wshape = [out_channel, in_channel // groups, filter_d, filter_h, filter_w] + + X = te.placeholder(xshape, name="X", dtype=data_dtype) + W = te.placeholder(wshape, name="W", dtype=data_dtype) + Y = mcdnn.conv_forward( + X, + W, + [pad_d, pad_h, pad_w], + [stride_d, stride_h, stride_w], + [dilation_d, dilation_h, dilation_w], + conv_mode=1, + tensor_format=tensor_format, + algo=-1, + conv_dtype=conv_dtype, + groups=groups, + ) + yshape = [x.value for x in Y.shape] + s = te.create_schedule(Y.op) + + # validation + dev = tvm.maca(0) + f = tvm.build(s, [X, W, Y], target="maca --host=llvm", name="conv3d") + x_np = np.random.uniform(-1, 1, xshape).astype(data_dtype) + w_np = np.random.uniform(-1, 1, wshape).astype(data_dtype) + y_np = np.zeros(yshape).astype(data_dtype) + x = tvm.nd.array(x_np, dev) + w = tvm.nd.array(w_np, dev) + y = tvm.nd.array(y_np, dev) + if tensor_format == 0: + c_np = tvm.topi.testing.conv3d_ncdhw_python(x_np, w_np, 1, 1, groups) + else: + raise AssertionError("For now, conv3d tensor format only support: 0(NCHW)") + + f(x, w, y) + tvm.testing.assert_allclose(y.numpy(), c_np, atol=3e-4, rtol=1e-4) + + +@tvm.testing.requires_gpu +@requires_mcdnn +@pytest.mark.skip("fail, Max absolute difference: 0.00356224") +def test_conv3d(): + # verify_conv3d("float32", "float32", tensor_format=0) + verify_conv3d("float32", "float32", tensor_format=0, groups=2) + + +def verify_softmax(shape, axis, dtype="float32", log_softmax=False): + mcdnn_op = mcdnn.log_softmax if log_softmax else mcdnn.softmax + testing_op = ( + tvm.topi.testing.log_softmax_python if log_softmax else tvm.topi.testing.softmax_python + ) + + A = te.placeholder(shape, dtype=dtype, name="A") + B = mcdnn_op(A, axis) + s = te.create_schedule([B.op]) + + dev = tvm.maca(0) + a_np = np.random.uniform(size=shape).astype(dtype) + b_np = testing_op(a_np) + a = tvm.nd.array(a_np, dev) + b = tvm.nd.array(b_np, dev) + f = tvm.build(s, [A, B], target="maca --host=llvm", name="softmax") + f(a, b) + tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-3) + + +def verify_softmax_4d(shape, dtype="float32", log_softmax=False): + mcdnn_op = mcdnn.log_softmax if log_softmax else mcdnn.softmax + testing_op = ( + tvm.topi.testing.log_softmax_python if log_softmax else tvm.topi.testing.softmax_python + ) + + A = te.placeholder(shape, dtype=dtype, name="A") + B = mcdnn_op(A, axis=1) + s = te.create_schedule([B.op]) + + dev = tvm.maca(0) + n, c, h, w = shape + a_np = np.random.uniform(size=shape).astype(dtype) + b_np = testing_op(a_np.transpose(0, 2, 3, 1).reshape(h * w, c)) + b_np = b_np.reshape(n, h, w, c).transpose(0, 3, 1, 2) + a = tvm.nd.array(a_np, dev) + b = tvm.nd.array(b_np, dev) + f = tvm.build(s, [A, B], target="maca --host=llvm", name="softmax") + f(a, b) + tvm.testing.assert_allclose(b.numpy(), b_np, rtol=1e-3) + + +@tvm.testing.requires_gpu +@requires_mcdnn +def test_softmax(): + verify_softmax((32, 10), -1) + verify_softmax((3, 4), -1) + verify_softmax((1, 5), -1, "float64") + verify_softmax_4d((1, 16, 256, 256)) + verify_softmax_4d((1, 16, 256, 256), "float64") + + verify_softmax((32, 10), -1, log_softmax=True) + verify_softmax((3, 4), -1, log_softmax=True) + verify_softmax((1, 5), -1, "float64", log_softmax=True) + verify_softmax_4d((1, 16, 256, 256), log_softmax=True) + verify_softmax_4d((1, 16, 256, 256), "float64", log_softmax=True) + + +def verify_conv2d_backward_data(data_dtype, conv_dtype, tensor_format=0, tol=1e-5): + batch = 3 + in_channel = 4 + out_channel = 16 + filter_h, filter_w = 3, 3 + pad_h, pad_w = 1, 1 + stride_h, stride_w = 1, 1 + height, width = 32, 32 + + if tensor_format == 0: + xshape = [batch, in_channel, height, width] + wshape = [out_channel, in_channel, filter_h, filter_w] + oshape = xshape + oshape[1] = out_channel + ref_func = tvm.topi.testing.conv2d_transpose_nchw_python + else: + xshape = [batch, height, width, in_channel] + wshape = [out_channel, filter_h, filter_w, in_channel] + oshape = xshape + oshape[3] = out_channel + ref_func = lambda dy_np, w_np, strides, padding, out_pad: tvm.topi.testing.conv2d_transpose_nhwc_python( + dy_np, np.transpose(w_np, [1, 2, 3, 0]), "HWOI", strides, padding, out_pad + ) + + dy_np = np.random.uniform(-1, 1, oshape).astype(data_dtype) + w_np = np.random.uniform(-1, 1, wshape).astype(data_dtype) + + if data_dtype == "float16": + dx_np = ref_func( + dy_np.astype("float32"), + w_np.astype("float32"), + (stride_h, stride_w), + (pad_h, pad_w), + (0, 0), + ) + dx_np = dx_np.astype("float16") + else: + dx_np = ref_func(dy_np, w_np, (stride_h, stride_w), (pad_h, pad_w), (0, 0)) + + dy = te.placeholder(oshape, name="dy", dtype=data_dtype) + w = te.placeholder(wshape, name="dw", dtype=data_dtype) + dx = mcdnn.conv_backward_data( + dy, + w, + [pad_h, pad_w], + [stride_h, stride_w], + [1, 1], + conv_mode=1, + tensor_format=tensor_format, + conv_dtype=conv_dtype, + groups=1, + ) + + s = te.create_schedule(dx.op) + + dev = tvm.maca(0) + f = tvm.build(s, [dy, w, dx], "maca --host=llvm", name="conv2d_backward_data") + + dy = tvm.nd.array(dy_np, dev) + w = tvm.nd.array(w_np, dev) + dx = tvm.nd.array(dx_np, dev) + + f(dy, w, dx) + tvm.testing.assert_allclose(dx.numpy(), dx_np, atol=tol, rtol=tol) + + +@tvm.testing.requires_gpu +@requires_mcdnn +def test_conv2d_backward_data(): + # FIXME: Max absolute difference: 0.00397015 + # verify_conv2d_backward_data("float32", "float32", tensor_format=0, tol=1e-5) + verify_conv2d_backward_data("float32", "float32", tensor_format=1, tol=1e-2) + # The scipy convolve function does not support fp16, so the reference will be computed with + # fp32. Use larger tolerance to be on the safe side (1e-2 also seems mostly ok). + verify_conv2d_backward_data("float16", "float16", tensor_format=1, tol=1e-1) + + +def verify_conv2d_backward_filter(data_dtype, conv_dtype, tensor_format=0, tol=1e-5): + batch = 3 + in_channel = 4 + out_channel = 16 + filter_h, filter_w = 3, 3 + pad_h, pad_w = 1, 1 + stride_h, stride_w = 1, 1 + height, width = 32, 32 + + if tensor_format == 0: + x_shape = [batch, in_channel, height, width] + dy_shape = [batch, out_channel, height, width] + else: + x_shape = [batch, height, width, in_channel] + dy_shape = [batch, height, width, out_channel] + + x_np = np.random.uniform(-1, 1, x_shape).astype(data_dtype) + dy_np = np.random.uniform(-1, 1, dy_shape).astype(data_dtype) + + dw_np = tvm.topi.testing.conv2d_backward_weight_python( + dy_np, + x_np, + (filter_h, filter_w), + (stride_h, stride_w), + (pad_h, pad_w), + "NCHW" if tensor_format == 0 else "NHWC", + ) + + x = te.placeholder(x_shape, name="x", dtype=data_dtype) + dy = te.placeholder(dy_shape, name="dy", dtype=data_dtype) + dw = mcdnn.conv_backward_filter( + dy, + x, + (filter_h, filter_w), + [pad_h, pad_w], + [stride_h, stride_w], + [1, 1], + conv_mode=1, + tensor_format=tensor_format, + conv_dtype=conv_dtype, + ) + + s = te.create_schedule(dw.op) + + dev = tvm.maca(0) + f = tvm.build(s, [dy, x, dw], "maca --host=llvm", name="conv2d_backward_filter") + + x = tvm.nd.array(x_np, dev) + dy = tvm.nd.array(dy_np, dev) + dw = tvm.nd.array(dw_np, dev) + + f(dy, x, dw) + tvm.testing.assert_allclose(dw.numpy(), dw_np, atol=tol, rtol=tol) + + +@tvm.testing.requires_gpu +@requires_mcdnn +@pytest.mark.skip("fail: Max absolute difference: 0.01865146") +def test_conv2d_backward_filter(): + verify_conv2d_backward_filter("float32", "float32", tensor_format=0, tol=1e-2) + verify_conv2d_backward_filter("float32", "float32", tensor_format=1, tol=1e-2) + + +test_kwargs_default_2d = { + "tensor_format": 0, + "pad": [1, 1], + "stride": [1, 1], + "dilation": [1, 1], + "x_shape": [16, 4, 32, 32], + "w_shape": [8, 4, 3, 3], + "groups": 1, + "conv_dtype": "float32", + "data_dtype": "float32", +} +test_kwargs_default_3d = { + "tensor_format": 0, + "pad": [1, 1, 1], + "stride": [1, 1, 1], + "dilation": [1, 1, 1], + "x_shape": [16, 4, 32, 32, 32], + "w_shape": [8, 4, 3, 3, 3], + "groups": 1, + "conv_dtype": "float32", + "data_dtype": "float32", +} +conv_output_shape_conditions = { + "2d_small": test_kwargs_default_2d, + "2d_large": { + **test_kwargs_default_2d, + "x_shape": [16, 32, 512, 1024], + "w_shape": [8, 32, 5, 5], + }, + "2d_pad": {**test_kwargs_default_2d, "pad": [2, 3]}, + "2d_stride": {**test_kwargs_default_2d, "stride": [2, 3]}, + "2d_dilation": {**test_kwargs_default_2d, "dilation": [2, 3]}, + "2d_groups": {**test_kwargs_default_2d, "groups": 4, "w_shape": [8, 1, 3, 3]}, + "2d_NHWC": { + **test_kwargs_default_2d, + "tensor_format": 1, + "x_shape": [16, 32, 32, 4], + "w_shape": [8, 3, 3, 4], + }, + "2d_NCHW_VECT_C": { + **test_kwargs_default_2d, + "tensor_format": 2, + "w_shape": [8, 16, 3, 3], + "data_dtype": "int8x4", + }, + "3d_small": test_kwargs_default_3d, + "3d_large": { + **test_kwargs_default_3d, + "x_shape": [16, 32, 64, 128, 256], + "w_shape": [8, 32, 5, 5, 5], + }, + "3d_pad": {**test_kwargs_default_3d, "pad": [2, 3, 4]}, + "3d_stride": {**test_kwargs_default_3d, "stride": [2, 3, 4]}, + "3d_dilation": {**test_kwargs_default_3d, "dilation": [2, 3, 4]}, + "3d_groups": {**test_kwargs_default_3d, "groups": 4, "w_shape": [8, 1, 3, 3, 3]}, + "3d_NCHW_VECT_C": { + **test_kwargs_default_3d, + "tensor_format": 2, + "w_shape": [8, 16, 3, 3, 3], + "data_dtype": "int8x4", + }, +} + + +@pytest.fixture( + params=[pytest.param(kwargs, id=name) for name, kwargs in conv_output_shape_conditions.items()] +) +def conv_output_shape_kwargs(request): + return request.param + + +def _verify_mcdnn_relay(expr): + np.random.seed(42) + + mod = tvm.IRModule.from_expr(expr) + mod = relay.transform.InferType()(mod) + func = mod["main"] + mcdnn_mod = partition_for_mcdnn(mod) + assert len(mcdnn_mod.get_global_vars()) == 2 + + input_data = [] + for param in func.params: + shape = [int(x) for x in param.checked_type.shape] + input_data.append( + ( + param.name_hint, + np.random.uniform(-32, 32, size=shape).astype(param.checked_type.dtype), + ) + ) + + maca_config = (tvm.target.maca(), tvm.maca(), mcdnn_mod) + cpu_config = (tvm.target.Target("llvm"), tvm.cpu(), mod) + outputs = [] + for target, dev, test_mod in [maca_config, cpu_config]: + with tvm.transform.PassContext(opt_level=3): + lib = relay.build(test_mod, target=target, target_host=cpu_config[0]) + module = graph_executor.GraphModule(lib["default"](dev)) + for name, data in input_data: + module.set_input(name, tvm.nd.array(data, dev)) + + module.run() + out_type = func.body.checked_type + outputs.append( + module.get_output(0, tvm.nd.empty(out_type.shape, dtype=out_type.dtype)).numpy() + ) + + tvm.testing.assert_allclose( + outputs[0], + outputs[1], + rtol=1e-2, + atol=30, + ) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "shape,axis", + [ + ((200,), 0), + ((13, 27), 0), + ((44, 12, 67), 1), + ((1, 16, 16, 8), 2), + ((2, 4, 6, 8, 10), 3), + ], +) +@pytest.mark.parametrize( + "dtype", + [ + "float32", + "float16", + "float64", + ], +) +def test_relay_mcdnn_softmax(shape, axis, dtype): + x = tvm.relay.var("x", tvm.relay.TensorType(shape, dtype)) + softmax = relay.op.nn.softmax(x, axis=axis) + _verify_mcdnn_relay(softmax) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "shape,axis", + [ + ((32, 16), -1), + ((13, 27), 1), + ], +) +@pytest.mark.parametrize( + "dtype", + [ + "float32", + "float16", + "float64", + ], +) +def test_relay_mcdnn_log_softmax(shape, axis, dtype): + x = tvm.relay.var("x", tvm.relay.TensorType(shape, dtype)) + log_softmax = relay.op.nn.log_softmax(x, axis=axis) + _verify_mcdnn_relay(log_softmax) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "n,h,w,ci,co,groups", + [ + (1, 16, 20, 8, 16, 1), + (10, 17, 19, 16, 8, 4), + ], +) +@pytest.mark.parametrize( + "kh,kw,padding", + [ + (1, 1, (3, 1, 3, 1)), + (3, 3, (1, 2)), + (7, 2, (0, 0)), + ], +) +@pytest.mark.parametrize( + "strides,dilation,dtype", + [ + ((1, 1), (1, 1), "float32"), + ((2, 1), (2, 2), "float16"), + ((3, 3), (1, 2), "float64"), + ], +) +def test_relay_mcdnn_conv2d(n, h, w, ci, co, kh, kw, strides, dilation, padding, groups, dtype): + data = tvm.relay.var("data", tvm.relay.TensorType((n, ci, h, w), dtype)) + weight = tvm.relay.var("weight", tvm.relay.TensorType((co, ci // groups, kh, kw), dtype)) + conv2d = relay.op.nn.conv2d( + data, + weight, + groups=groups, + channels=co, + kernel_size=(kh, kw), + strides=strides, + dilation=dilation, + padding=padding, + data_layout="NCHW", + kernel_layout="OIHW", + ) + _verify_mcdnn_relay(conv2d) + + +@tvm.testing.requires_maca +@pytest.mark.parametrize( + "n,h,w,ci,co,groups", + [ + (1, 16, 20, 8, 16, 1), + (10, 17, 19, 16, 8, 4), + ], +) +@pytest.mark.parametrize( + "kh,kw,padding,strides,dilation,dtype", + [ + (1, 1, (3, 1, 3, 1), (1, 1), (1, 1), "float32"), + (3, 3, (1, 2), (2, 1), (2, 2), "float16"), + (7, 2, (0, 0), (3, 3), (1, 2), "float64"), + ], +) +@pytest.mark.parametrize("activation", [True, False]) +def test_relay_mcdnn_conv2d_bias_act( + n, h, w, ci, co, kh, kw, strides, dilation, padding, groups, dtype, activation +): + data = tvm.relay.var("data", tvm.relay.TensorType((n, ci, h, w), dtype)) + weight = tvm.relay.var("weight", tvm.relay.TensorType((co, ci // groups, kh, kw), dtype)) + bias = relay.var("bias", relay.TensorType((co,), dtype)) + conv2d = relay.op.nn.conv2d( + data, + weight, + groups=groups, + channels=co, + kernel_size=(kh, kw), + strides=strides, + dilation=dilation, + padding=padding, + data_layout="NCHW", + kernel_layout="OIHW", + ) + out = relay.op.nn.bias_add(conv2d, bias) + if activation: + out = relay.op.nn.relu(out) + + _verify_mcdnn_relay(out) + + +if __name__ == "__main__": + tvm.testing.main()