Back to Cutlass

Source code for cutlass.epilogue

python/docs/_modules/cutlass/epilogue.html

4.5.03.7 KB
Original Source

Source code for cutlass.epilogue

################################################################################################### Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.# SPDX-License-Identifier: BSD-3-Clause## Redistribution and use in source and binary forms, with or without# modification, are permitted provided that the following conditions are met:## 1. Redistributions of source code must retain the above copyright notice, this# list of conditions and the following disclaimer.## 2. Redistributions in binary form must reproduce the above copyright notice,# this list of conditions and the following disclaimer in the documentation# and/or other materials provided with the distribution.## 3. Neither the name of the copyright holder nor the names of its# contributors may be used to endorse or promote products derived from# this software without specific prior written permission.## THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.##################################################################################################"""Registry of elementwise epiloguesElementwise epilogues can be added to many CUTLASS kernels in the CUTLAS Python interface viacode like the following for GEMM:.. highlight:: python.. code-block:: python plan = cutlass.op.Gemm(element=cutlass.DataType.f32, layout=cutlass.LayoutType.RowMajor) plan.activation = cutlass.epilogue.relu"""fromcutlass.backendimportepiloguegelu=epilogue.geluhardswish=epilogue.hardswishidentity=epilogue.identityleaky\_relu=epilogue.leaky\_relurelu=epilogue.relusigmoid=epilogue.sigmoidsilu=epilogue.silutanh=epilogue.tanh\_activations=[gelu,hardswish,identity,leaky\_relu,relu,sigmoid,silu,tanh]
[[docs]](../../cutlass.html#cutlass.epilogue.get_activations)def get\_activations() -\> list:""" Returns a list of available activation functions :return: list of available activation functions :rtype: list """return \_activations

[[docs]](../../cutlass.html#cutlass.epilogue.get_activation_epilogue)def get\_activation\_epilogue(activation,element\_output,elements\_per\_access,element\_accumulator,element\_compute,):""" Return an epilogue corresponding to the activation function, data types, and alignment used in the kernel :param activation: elementwise activation function to use :param element\_output: data type of the output :param elements\_per\_access: alignment of operand C of the kernel :type elements\_per\_access: int :param element\_accumulator: data type of the accumulated output C :param element\_compute: data type in which compute operations should be performed :return: epilogue functor """if activation not in \_activations:raise Exception(f"Unsupported activation type {activation}. Available activations are: {\_activations}")if activation == identity:return epilogue.LinearCombination(element\_output, elements\_per\_access, element\_accumulator, element\_compute)else:return epilogue.LinearCombinationGeneric(activation(element\_compute),element\_output,elements\_per\_access,element\_accumulator,element\_compute,)