################################################################################################### 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.##################################################################################################"""Utility functions for checking constraints on kernels and calculating kernel attributes"""importctypesimportcutlass\_bindingsimportcutlassfromcutlass.backend.libraryimportDataTypeSize,TileDescription
[[docs]](../../../cutlass.utils.html#cutlass.utils.check.calculate_smem_usage_per_stage)def calculate\_smem\_usage\_per\_stage(tile\_description, operation\_kind):""" Returns the amount of shared memory in bytes consumed in a single stage of a kernel. :return: number of bytes of shared memory consumed by a single stage :rtype: int """m, n, k = tile\_description.threadblock\_shapeif operation\_kind == cutlass.OperationKind.Gemm:stage\_barrier\_bytes = 32return ((DataTypeSize[tile\_description.math\_instruction.element\_a] \* m \* k // 8)+ (DataTypeSize[tile\_description.math\_instruction.element\_b] \* k \* n // 8)+ stage\_barrier\_bytes)else:raise Exception(f"No available shared memory calculation for operation kind {operation.operation\_kind}")

[[docs]](../../../cutlass.utils.html#cutlass.utils.check.calculate_smem_usage)def calculate\_smem\_usage(operation):""" Returns the amount of shared memory in bytes consumed by a kernel. :return: number of bytes of shared memory consumed by the operation :return: int """\_per\_stage = calculate\_smem\_usage\_per\_stage(operation.tile\_description, operation.operation\_kind)return \_per\_stage \* operation.tile\_description.stages

[[docs]](../../../cutlass.utils.html#cutlass.utils.check.valid_stage_count)def valid\_stage\_count(cc: int, td: TileDescription) -\> tuple:""" Checks whether a device with `cc` supports the number of stages within `tile_description`, both based on raw limits on the number of stages and based on shared memory capacity :param cc: compute capability of device in question :type cc: int :param td: tile description to check :type td: TileDescription :return: tuple with the first element indicating whether the provided tile description is valid for the provided device and the second element being an error message :rtype: tuple """if cc == 90 and (td.stages is None or td.stages == 0):# Stage count of None or 0 for SM90 indicates that the CollectiveBuilder automatically# determines the stage count to use. Thus, all settings are valid in these scenarios.return (True, "")if td.stages \<= 0:return (False, f"Stage counts must be positive integers. Tile description has stage count of {td.stages}.")if cc \< 80 and td.stages != 2:return (False, f"Tile description has stage count of {td.stages}, "f"but only 2 stages are supported on SM{cc}.")smem\_per\_stage = calculate\_smem\_usage\_per\_stage(td, cutlass.OperationKind.Gemm)smem\_arch = cutlass.SharedMemPerCC[cc] \<\< 10if (smem\_per\_stage \* td.stages) \> smem\_arch:return ( False,"Configuration uses too much shared memory. Consider reducing stage count or tile shape.\n"f"Details: configuration uses {smem\_per\_stage} bytes of shared memory per stage, and "f"{td.stages} stages for a total of {smem\_per\_stage \* td.stages} bytes.\n"f"The maxmium amoung of shared memory that can be used per block on CC {cc} is {smem\_arch}.")return (True, "")

[[docs]](../../../cutlass.utils.html#cutlass.utils.check.valid_cluster_shape)def valid\_cluster\_shape(cc: int, cluster\_shape: list) -\> tuple:""" Checks whether a device with `cc` supports a thread block cluster of shape `cluster_shape`. :param cc: compute capability of device in question :type cc: int :param cluster\_shape: dimensions of thread block cluster shape to check :type cluster\_shape: list :return: tuple with the first element indicating whether the provided cluster shape is valid for the provided device and the second element being an error message :rtype: tuple """if cc \< 90:if cluster\_shape != [1, 1, 1]:return (False,f"Cluster shape for pre-SM90 architectures must be [1, 1, 1]. Received cluster shape of "f"{cluster\_shape} for SM{cc}.")else:return (True, "")if len(cluster\_shape) != 3:return (False,f"Cluster shapes must be rank-3. Received {cluster\_shape} (rank {len(cluster\_shape)}")if cluster\_shape[2] != 1:return (False,"CUTLASS kernels currently require the third dimension of cluster shape to be 1. "f"Received cluster shape of {cluster\_shape}.")# The CUDA programming guide currently defines a maximum of 8 thread blocks per cluster# as being portably supported (https://docs.nvidia.com/cuda/cuda-c-programming-guide/#thread-block-clusters).# Current CUTLASS kernels only have non-unit cluster dimensions within the first two dimensions,# so we check that the first two dimensions of the cluster shape do not exceed 8 thread blocks in total.blocks\_in\_2d = cluster\_shape[0] \* cluster\_shape[1]if blocks\_in\_2d \> 8:return (False,f"Thread block clusters with more than 8 thread blocks are currently unsupported on SM{cc}. "f"Received cluster shape {cluster\_shape}, which has {blocks\_in\_2d} thread blocks.")return (True, "")

[[docs]](../../../cutlass.utils.html#cutlass.utils.check.valid_kernel_schedule)def valid\_kernel\_schedule(cc: int, kernel\_schedule: cutlass.KernelScheduleType) -\> tuple:""" Checks whether a device with ``cc`` supports ``kernel\_schedule``. :param cc: compute capability of device in question :type cc: int :param kernel\_schedule: kernel schedule type :type KernelScheduleType: cutlass.KernelScheduleType :return: tuple with the first element indicating whether the provided kernel schedule is valid for the provided device and the second element being an error message :rtype: tuple """if kernel\_schedule != cutlass.KernelScheduleType.ScheduleAuto and cc \< 90:return (False, "Non-default kernel schedules are only supported on SM90 and beyond")return (True, "")

[[docs]](../../../cutlass.utils.html#cutlass.utils.check.alignment_or_default)def alignment\_or\_default(alignment\_provided: int, default\_alignment: int) -\> int:""" Returns `alignment_provided` if it is set, otherwise `default_alignment` and checks that `alignment_provided` does not exceed `default_alignment`. :param alignment\_provided: alignment preference specified. Can be None. :type alignment\_provided: int :param default\_alignment: alignment to use if `alignment_provided` is None :type default\_alignment: int :return: alignment to use :rtype: int """if alignment\_provided is not None:if alignment\_provided \> default\_alignment:raise Exception(f"Alignment {alignment\_provided} exceeds the maximum supported of {default\_alignment}.")return alignment\_providedreturn default\_alignment