################################################################################################# # # Copyright (c) 2017 - 2025 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. # ################################################################################################# """ Utilities for enumerating CUTLASS library kernels """ import argparse import enum from itertools import chain, product import logging import os.path import shutil import sys import copy from typing import Any, Dict, Optional, Sequence, Tuple _LOGGER = logging.getLogger(__name__) def logging_prefix(indent_level: int = 0) -> str: """String prefix for start of each debug log entry""" prefix = '*** ' indent = ' ' return f"{prefix}{indent_level * indent}" def log_debug_line(line: str, indent_level: int = 0) -> None: """Log one line of debug output""" prefix = logging_prefix(indent_level) _LOGGER.debug(prefix + line) # Certain usecases of cutlass_library nearly always prefer to run as scripts with # relative imports, rather than via an installed Python package. An example of this # is using CUTLASS's CMake system to generate a library of kernels to be profiled. # To make it easy to use these use cases when an existing installation of cutlass_library # exists, this global flag can be set to true (via command-line arguments) to ensure # that package-based installations are not used. # Create a temporary argument parser to check only for the availability of the # --disable-cutlass-package-imports argument, which controls whether package-based # imports are disabled. def _add_package_disablement_flag(argparser): argparser.add_argument("--disable-cutlass-package-imports", action='store_true', required=False, help="Disable use of cutlass_library from Python package") _parser = argparse.ArgumentParser() _add_package_disablement_flag(_parser) _args, _ = _parser.parse_known_args() # Add `CUTLASS_IGNORE_PACKAGE` to `builtins` so that it is visible for gating future # imports without requiring importing another module. Ideally, we would just place this # as a global variable in a module to that could be imported and checked (e.g., # utils.CUTLASS_IGNORE_PACKAGE). However, this raises the issue of determining # where this module should be sourced (from the cutlass_library package or from # a relative import), which is the problem this variable is being used to solve in the # first place. import builtins builtins.CUTLASS_IGNORE_PACKAGE = _args.disable_cutlass_package_imports try: if CUTLASS_IGNORE_PACKAGE: raise ImportError("Disabling attempt to import cutlass_library") from cutlass_library.library import * from cutlass_library.manifest import * from cutlass_library.heuristics import * from cutlass_library.emit_kernel_listing import emit_gemm_kernel_testlist except ImportError: from library import * from manifest import * from heuristics import * from emit_kernel_listing import emit_gemm_kernel_testlist ################################################################################################### # def CudaToolkitVersionSatisfies(semantic_ver_string, major, minor, patch = 0): # by default, use the latest CUDA Toolkit version cuda_version = [11, 0, 132] # Update cuda_version based on parsed string if semantic_ver_string != '': for i, x in enumerate([int(x) for x in semantic_ver_string.split('.')[:3]]): if i < len(cuda_version): cuda_version[i] = x else: cuda_version.append(x) return cuda_version >= [major, minor, patch] # From cuda 13.0, Thor SM is renumbered from 101 to 110 def ThorSMRenumbering(cuda_version): return 110 if CudaToolkitVersionSatisfies(cuda_version, 13, 0) else 101 ################################################################################################### ################################################################################################### # def EpilogueAlignment(max_alignment, tile, epilogue_steps = 8): ''' Helper to compute the maximum alignment of the epilogue ''' def product(X, identity = 1): result = identity for item in X: result *= item return result elements_per_thread = product(tile.threadblock_shape[:-1]) // product(tile.warp_count) // 32 // epilogue_steps return min(max_alignment, elements_per_thread) def DefaultSwizzlingFunctor(): return SwizzlingFunctor.Identity8 # To use StreamK decomposition for basic GEMMs, set `swizzling_functor = SwizzlingFunctor.StreamK` # def CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, \ alignment_constraints, complex_transforms = None, epilogue_functor = EpilogueFunctor.LinearCombination, \ swizzling_functor = DefaultSwizzlingFunctor()): if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none),] element_a, element_b, element_c, element_epilogue = data_type operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for layout in layouts: for tile_description in tile_descriptions: for alignment in alignment_constraints: for complex_transform in complex_transforms: # If alignment is a tuple or a list, then we have different alignments for A and B alignment_a = alignment if isinstance(alignment, int) else alignment[0] alignment_b = alignment if isinstance(alignment, int) else alignment[1] alignment_c = min(8, alignment_a) if isinstance(alignment, int) else alignment[2] A = TensorDescription(element_a, layout[0], alignment_a, complex_transform[0]) B = TensorDescription(element_b, layout[1], alignment_b, complex_transform[1]) C = TensorDescription(element_c, layout[2], alignment_c) new_operation = GemmOperation(GemmKind.Universal, tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor) manifest.append(new_operation) operations.append(new_operation) return operations # Generates 3.0 API based GemmUniversal API kernels. Alignment constraints are folded in with layouts def CreateGemmUniversal3xOperator( manifest, layouts, tile_descriptions, data_types, schedules = [[KernelScheduleType.ScheduleAuto, EpilogueScheduleType.ScheduleAuto]], complex_transforms=None, epilogue_functor=EpilogueFunctor.LinearCombination, swizzling_functor=SwizzlingFunctor.Identity1, tile_schedulers=[TileSchedulerType.Default], gemm_kind=GemmKind.Universal3x): if type(data_types) is dict: data_types = [data_types] for s in schedules: assert(len(s) == 2) if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none), ] operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': if len(tile_descriptions) == 0: return operations tile_descriptions = [tile_descriptions[0]] combinations = product(layouts, tile_descriptions, data_types, complex_transforms, schedules, tile_schedulers) for layout, tile_description, data_type, complex_transform, schedules, tile_scheduler in combinations: kernel_schedule, epilogue_schedule = schedules A = TensorDescription( data_type["a_type"], layout[0][0], layout[0][1], complex_transform[0]) B = TensorDescription( data_type["b_type"], layout[1][0], layout[1][1], complex_transform[1]) C = TensorDescription(data_type["c_type"], layout[2][0], layout[2][1]) D = TensorDescription(data_type["d_type"], layout[2][0], layout[2][1]) gemm_op_extra_args = {} element_compute = data_type.get("epi_type", data_type["acc_type"]) if "sf_type" in data_type: gemm_op_extra_args["ScaleFactorA"] = data_type["sf_type"] gemm_op_extra_args["ScaleFactorB"] = data_type["sf_type"] gemm_op_extra_args["ScaleFactorD"] = { "tensor": TensorDescription(data_type["sfd_type"]["type"], data_type["sfd_type"]["layout"]), "vector_size" : data_type["sfd_type"]["vector_size"]} assert is_block_scaled(gemm_kind) if tile_description.explicit_vector_sizes != None: assert len(tile_description.explicit_vector_sizes) == 3 gemm_op_extra_args["ScaleFactorMVecSize"] = tile_description.explicit_vector_sizes[0] gemm_op_extra_args["ScaleFactorNVecSize"] = tile_description.explicit_vector_sizes[1] gemm_op_extra_args["ScaleFactorKVecSize"] = tile_description.explicit_vector_sizes[2] assert is_blockwise(gemm_kind) else: assert not is_blockwise(gemm_kind) A_dtype = data_type["a_type"] B_dtype = data_type["b_type"] A_dtype_bits = DataTypeSize[A_dtype] B_dtype_bits = DataTypeSize[B_dtype] is_A_dtype_narrow = A_dtype_bits < B_dtype_bits if is_A_dtype_narrow: narrow_dtype, wide_dtype = (A_dtype, B_dtype) narrow_dtype_bits, wide_dtype_bits = (A_dtype_bits, B_dtype_bits) else: narrow_dtype, wide_dtype = (B_dtype, A_dtype) narrow_dtype_bits, wide_dtype_bits = (B_dtype_bits, A_dtype_bits) mixed_input_modes = [None] if narrow_dtype_bits != wide_dtype_bits: if narrow_dtype == DataType.s4 and (wide_dtype == DataType.e4m3 or wide_dtype == DataType.e5m2): mixed_input_modes = [MixedInputMode.ScaleOnly] else: mixed_input_modes = [MixedInputMode.ConvertOnly, MixedInputMode.ScaleOnly, MixedInputMode.ScaleWithZeroPoint] mixed_input_shuffle_options = [False] if (mixed_input_modes[0] is not None) and (wide_dtype_bits == 16) and (narrow_dtype_bits == 4 or narrow_dtype_bits == 8): mixed_input_shuffle_options = [False, True] for mixed_input_mode, mixed_input_shuffle in product(mixed_input_modes, mixed_input_shuffle_options): operation = GemmOperation( gemm_kind, tile_description.minimum_compute_capability, tile_description, A, B, C, element_compute, epilogue_functor, swizzling_functor, D, kernel_schedule, epilogue_schedule, tile_scheduler, mixed_input_mode=mixed_input_mode, mixed_input_shuffle=mixed_input_shuffle, **gemm_op_extra_args) manifest.append(operation) operations.append(operation) return operations # Generates 3.0 API based GemmUniversal API kernels. Alignment constraints are folded in with layouts def CreateSparseGemmUniversal3xOperator( manifest, layouts, tile_descriptions, data_types, schedules = [[KernelScheduleType.ScheduleAuto, EpilogueScheduleType.ScheduleAuto]], complex_transforms=None, epilogue_functor=EpilogueFunctor.LinearCombination, swizzling_functor=SwizzlingFunctor.Identity1, tile_schedulers=[TileSchedulerType.Default]): if type(data_types) is dict: data_types = [data_types] for s in schedules: assert(len(s) == 2) if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none), ] operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0]] combinations = product(layouts, tile_descriptions, data_types, complex_transforms, schedules, tile_schedulers) for layout, tile_description, data_type, complex_transform, schedules, tile_scheduler in combinations: kernel_schedule, epilogue_schedule = schedules A = TensorDescription( data_type["a_type"], layout[0][0], layout[0][1], complex_transform[0]) B = TensorDescription( data_type["b_type"], layout[1][0], layout[1][1], complex_transform[1]) # Currently assume tensor C/D have same layout requirement. C = TensorDescription(data_type["c_type"], layout[2][0], layout[2][1]) D = TensorDescription(data_type["d_type"], layout[2][0], layout[2][1]) element_compute = data_type.get("epi_type", data_type["acc_type"]) operation = GemmOperation( GemmKind.SparseUniversal3x, tile_description.minimum_compute_capability, tile_description, A, B, C, element_compute, epilogue_functor, swizzling_functor, D, kernel_schedule, epilogue_schedule, tile_scheduler) manifest.append(operation) operations.append(operation) return operations # def CreateSparseGemmOperator(manifest, layouts, tile_descriptions, data_type, \ alignment_constraints, complex_transforms = None, epilogue_functor = EpilogueFunctor.LinearCombination, \ swizzling_functor = SwizzlingFunctor.Identity8): if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none),] element_a, element_b, element_c, element_epilogue = data_type gemm_kinds = [GemmKind.Sparse] operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for layout in layouts: for tile_description in tile_descriptions: for alignment in alignment_constraints: for complex_transform in complex_transforms: alignment_c = min(8, alignment) A = TensorDescription(element_a, layout[0], alignment, complex_transform[0]) B = TensorDescription(element_b, layout[1], alignment, complex_transform[1]) C = TensorDescription(element_c, layout[2], alignment_c) new_operation = GemmOperation(GemmKind.Sparse, tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor) manifest.append(new_operation) operations.append(new_operation) return operations # def CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, data_type, \ alignment_constraints, complex_transforms): if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none),] element_a, element_b, element_c, element_epilogue = data_type gemm_kinds = [GemmKind.PlanarComplex, GemmKind.PlanarComplexArray] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for gemm_kind in gemm_kinds: for layout in layouts: for tile_description in tile_descriptions: for alignment in alignment_constraints: for complex_transform in complex_transforms: alignment_c = min(8, alignment) A = TensorDescription(element_a, layout[0], alignment, complex_transform[0]) B = TensorDescription(element_b, layout[1], alignment, complex_transform[1]) C = TensorDescription(element_c, layout[2], alignment_c) manifest.append(GemmOperation(gemm_kind, \ tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue)) return # def CreateGemmGroupedOperator(manifest, layouts, tile_descriptions, data_type, \ alignment_constraints, complex_transforms = None, epilogue_functor = EpilogueFunctor.LinearCombination, \ swizzling_functor = SwizzlingFunctor.Identity8): if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none),] element_a, element_b, element_c, element_epilogue = data_type operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for layout in layouts: for tile_description in tile_descriptions: for alignment in alignment_constraints: for complex_transform in complex_transforms: alignment_c = min(8, alignment) A = TensorDescription(element_a, layout[0], alignment, complex_transform[0]) B = TensorDescription(element_b, layout[1], alignment, complex_transform[1]) C = TensorDescription(element_c, layout[2], alignment_c) new_operation = GroupedGemmOperation(GemmKind.Grouped, tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor) manifest.append(new_operation) operations.append(new_operation) return operations # def CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, data_type, \ alignment_constraints, blas_mode, epilogue_functor = EpilogueFunctor.LinearCombination, \ swizzling_functor = SwizzlingFunctor.Identity8): element_a, element_c, element_epilogue = data_type operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for layout in layouts: for fill_mode in fill_modes: for tile_description in tile_descriptions: for alignment in alignment_constraints: # SERK supported layouts (RowMajor, ColumnMajor) with no conjugation complex_transform = ComplexTransform.none # HERK supported layouts (RowMajor + conj, ColumnMajor) if blas_mode == BlasMode.hermitian and layout[0] == LayoutType.RowMajor: complex_transform = ComplexTransform.conj alignment_c = 1 # Alignment only applies to A in SYRK A = TensorDescription(element_a, layout[0], alignment, complex_transform) C = SymmetricTensorDescription(element_c, layout[1], fill_mode, alignment_c) # Rank-K update new_operation = RankKOperation(RankKKind.Universal, tile_description.minimum_compute_capability, \ tile_description, A, C, element_epilogue, epilogue_functor, swizzling_functor, blas_mode) manifest.append(new_operation) operations.append(new_operation) # Rank-2K update new_operation = Rank2KOperation(RankKKind.Universal, tile_description.minimum_compute_capability, \ tile_description, A, C, element_epilogue, epilogue_functor, swizzling_functor, blas_mode) manifest.append(new_operation) operations.append(new_operation) return operations # def CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, data_type, \ alignment_constraints, complex_transforms = None, epilogue_functor = EpilogueFunctor.LinearCombination, \ swizzling_functor = SwizzlingFunctor.Identity8): if complex_transforms is None: complex_transforms = [(ComplexTransform.none),] element_a, element_b, element_c, element_epilogue = data_type operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for layout in layouts: for side_mode in side_modes: for fill_mode in fill_modes: for diag_type in diag_types: for tile_description in tile_descriptions: for alignment in alignment_constraints: for complex_transform in complex_transforms: alignment_c = min(8, alignment) A = TriangularTensorDescription(element_a, layout[0], side_mode, fill_mode, diag_type, alignment, complex_transform) B = TensorDescription(element_b, layout[1], alignment) C = TensorDescription(element_c, layout[2], alignment_c) new_operation = TrmmOperation(TrmmKind.Universal, tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor) manifest.append(new_operation) operations.append(new_operation) return operations # def CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, data_type, \ alignment_constraints, blas_mode, epilogue_functor = EpilogueFunctor.LinearCombination, \ swizzling_functor = SwizzlingFunctor.Identity8): element_a, element_b, element_c, element_epilogue = data_type operations = [] # by default, only generate the largest tile and largest alignment if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] for layout in layouts: for side_mode in side_modes: for fill_mode in fill_modes: for tile_description in tile_descriptions: for alignment in alignment_constraints: # SYMM supported layouts (RowMajor, ColumnMajor) with no conjugation complex_transform = ComplexTransform.none alignment_a = 1 # No vectorized access for the triangular matrix alignment_c = min(8, alignment) A = SymmetricTensorDescription(element_a, layout[0], fill_mode, alignment_a, complex_transform, side_mode) # tensor A and B have same data type and layout B = TensorDescription(element_b, layout[0], alignment) C = TensorDescription(element_c, layout[1], alignment_c) # SYMM/HEMM update new_operation = SymmOperation(SymmKind.Universal, tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor, blas_mode) manifest.append(new_operation) operations.append(new_operation) # SYMM/HEMM update new_operation = SymmOperation(SymmKind.Universal, tile_description.minimum_compute_capability, \ tile_description, A, B, C, element_epilogue, epilogue_functor, swizzling_functor, blas_mode) manifest.append(new_operation) operations.append(new_operation) return operations ########################################################################################################### # ConvolutionOperator support variations # ____________________________________________________________________ # ConvolutionalOperator | Analytic | Optimized # ____________________________________________________________________ # | Fprop | (strided) | (strided) # | Dgrad | (strided, unity*) | (strided, unity) # | Wgrad | (strided) | (strided) # ____________________________________________________________________ # # Note : Operator marked (*) are supported but not generated to keep the instantiated kernel count low ########################################################################################################### # Convolution for 2D operations def CreateConv2dOperator(manifest, layout, tile_descriptions, data_type, alignment_constraints, \ conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], \ epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity4): element_a, element_b, element_c, element_epilogue = data_type # one exceptional case # iterator algorithm (analytic and optimized) iterator_algorithms = [IteratorAlgorithm.Analytic, IteratorAlgorithm.Optimized] # by default, only generate the largest tile size, largest alignment, and optimized iterator if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] iterator_algorithms = [IteratorAlgorithm.Optimized] operations = [] for tile in tile_descriptions: for alignment in alignment_constraints: alignment_c = min(8, alignment) A = TensorDescription(element_a, layout[0], alignment) B = TensorDescription(element_b, layout[1], alignment) C = TensorDescription(element_c, layout[2], alignment_c) swizzling_functor_ = swizzling_functor # # Conv2d Fprop # if ConvKind.Fprop in conv_kinds: # Strided support for Analytic and Optimized Fprop for iterator_algorithm in iterator_algorithms: new_operations = [ # None grouped kernel Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor, swizzling_functor_), ] # Instance group conv kernel if tile.math_instruction.opcode_class == OpcodeClass.TensorOp and A.layout == LayoutType.TensorNHWC and \ tile.minimum_compute_capability >= 80: # SingleGroup kernel new_operations.append(Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor, swizzling_functor_, group_mode=GroupMode.SingleGroup)) # Analytic iterator supports MultipleGroup mode if iterator_algorithm == IteratorAlgorithm.Analytic: new_operations.append(Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor, swizzling_functor_, group_mode=GroupMode.MultipleGroup)) for new_operation in new_operations: manifest.append(new_operation) operations.append(new_operation) # # Conv2d Dgrad # if ConvKind.Dgrad in conv_kinds: # Unity stride for Analytic and Optimized Dgrad for iterator_algorithm in iterator_algorithms: new_operation = Conv2dOperation(ConvKind.Dgrad, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor, swizzling_functor_) manifest.append(new_operation) operations.append(new_operation) # Strided support for Analytic Dgrad # strided dgrad uses a special threadblock swizzle # note that SwizzlingFunctor.StridedDgradHorizontal might be # better for problem sizes with large activation channel count swizzling_functor_strided_dgrad_ = SwizzlingFunctor.StridedDgradIdentity1 if IteratorAlgorithm.Analytic in iterator_algorithms: new_operation = Conv2dOperation(ConvKind.Dgrad, IteratorAlgorithm.Analytic, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_strided_dgrad_) manifest.append(new_operation) operations.append(new_operation) # Strided support for Optimized Dgrad if IteratorAlgorithm.Optimized in iterator_algorithms: new_operation = Conv2dOperation(ConvKind.Dgrad, IteratorAlgorithm.Optimized, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_strided_dgrad_) manifest.append(new_operation) operations.append(new_operation) # # Conv2d Wgrad # if ConvKind.Wgrad in conv_kinds: # Strided support for Analytic and Optimized Wgrad for iterator_algorithm in iterator_algorithms: new_operation = Conv2dOperation(ConvKind.Wgrad, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_) manifest.append(new_operation) operations.append(new_operation) return operations # Convolution for 2D operations specialized for few channels def CreateConv2dFixedChannelsOperator(manifest, layout, tile_descriptions, data_type, channel_counts, \ conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], \ epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity4): element_a, element_b, element_c, element_epilogue = data_type # one exceptional case # iterator algorithm (analytic and optimized) iterator_algorithms = [IteratorAlgorithm.FixedChannels,] # by default, only generate the largest tile size, largest alignment, and optimized iterator if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] channel_counts = [channel_counts[0],] operations = [] for tile in tile_descriptions: for channel_count in channel_counts: alignment_c = EpilogueAlignment(channel_count, tile) A = TensorDescription(element_a, layout[0], channel_count) B = TensorDescription(element_b, layout[1], channel_count) C = TensorDescription(element_c, layout[2], alignment_c) swizzling_functor_ = swizzling_functor # # Conv2d Fprop # if ConvKind.Fprop in conv_kinds: # Strided support for Analytic and Optimized Fprop for iterator_algorithm in iterator_algorithms: new_operation = Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_) manifest.append(new_operation) operations.append(new_operation) return operations # Convolution for 2D operations specialized for few channels def CreateConv2dFewChannelsOperator(manifest, layout, tile_descriptions, data_type, channel_counts, \ conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], \ epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity4): element_a, element_b, element_c, element_epilogue = data_type # one exceptional case # iterator algorithm (analytic and optimized) iterator_algorithms = [IteratorAlgorithm.FewChannels,] # by default, only generate the largest tile size, largest alignment, and optimized iterator if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] channel_counts = [channel_counts[0],] operations = [] for tile in tile_descriptions: for channel_count in channel_counts: alignment_c = EpilogueAlignment(channel_count, tile) A = TensorDescription(element_a, layout[0], channel_count) B = TensorDescription(element_b, layout[1], channel_count) C = TensorDescription(element_c, layout[2], alignment_c) swizzling_functor_ = swizzling_functor # # Conv2d Fprop # if ConvKind.Fprop in conv_kinds: # Strided support for Analytic and Optimized Fprop for iterator_algorithm in iterator_algorithms: new_operation = Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor, swizzling_functor_) manifest.append(new_operation) operations.append(new_operation) return operations # Convolution for 3D operations def CreateConv3dOperator(manifest, layout, tile_descriptions, data_type, alignment, \ conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], epilogue_functor = EpilogueFunctor.LinearCombination): element_a, element_b, element_c, element_epilogue = data_type # one exceptional case alignment_c = min(8, alignment) # iterator algorithm (analytic and optimized) iterator_algorithms = [IteratorAlgorithm.Analytic, IteratorAlgorithm.Optimized] # by default, only generate the largest tile size and optimized iterators if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] iterator_algorithms = [IteratorAlgorithm.Optimized] operations = [] # All tile sizes for Conv3dFprop and Conv3dWgrad for tile in tile_descriptions: A = TensorDescription(element_a, layout, alignment) B = TensorDescription(element_b, layout, alignment) C = TensorDescription(element_c, layout, alignment_c) # # Conv3d Fprop # if ConvKind.Fprop in conv_kinds: # Strided support for Analytic and Optimized Fprop for iterator_algorithm in iterator_algorithms: new_operation = Conv3dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided) manifest.append(new_operation) operations.append(new_operation) # # Conv3d Wgrad # if ConvKind.Wgrad in conv_kinds: # Strided support for Analytic and Optimized Wgrad for iterator_algorithm in iterator_algorithms: new_operation = Conv3dOperation(ConvKind.Wgrad, iterator_algorithm, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor) manifest.append(new_operation) operations.append(new_operation) # All tile sizes for Conv3dDgrad for tile in tile_descriptions: A = TensorDescription(element_a, layout, alignment) B = TensorDescription(element_b, layout, alignment) C = TensorDescription(element_c, layout, alignment_c) # # Conv3d Dgrad # if ConvKind.Dgrad in conv_kinds: # Unity stride for Optimized Dgrad new_operation = Conv3dOperation(ConvKind.Dgrad, IteratorAlgorithm.Optimized, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Unity, epilogue_functor) manifest.append(new_operation) operations.append(new_operation) # Strided support for Analytic Dgrad # Conv3dDgrad has a naive strided support which does not cut down redundant MMAs new_operation = Conv3dOperation(ConvKind.Dgrad, IteratorAlgorithm.Analytic, tile.minimum_compute_capability, tile,\ A, B, C, element_epilogue, StrideSupport.Strided, epilogue_functor) manifest.append(new_operation) operations.append(new_operation) return operations # Convolution for Depthwise 2d conv def CreateDepthwiseConv2dOperator(manifest, layout, tile_descriptions, data_type, alignment_constraints, \ conv_kinds = [ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad], \ epilogue_functor = EpilogueFunctor.LinearCombination, swizzling_functor = SwizzlingFunctor.Identity4): element_a, element_b, element_c, element_epilogue = data_type # iterator algorithm (FixedStrideDilation, Optimized) iterator_algorithms = [IteratorAlgorithm.FixedStrideDilation, IteratorAlgorithm.Optimized] # by default, only generate the largest tile size, largest alignment, and optimized iterator if manifest.kernel_filter == '': tile_descriptions = [tile_descriptions[0],] alignment_constraints = [alignment_constraints[0],] operations = [] for tile in tile_descriptions: for alignment in alignment_constraints: alignment_c = min(8, alignment) A = TensorDescription(element_a, layout[0], alignment) B = TensorDescription(element_b, layout[1], alignment) C = TensorDescription(element_c, layout[2], alignment_c) swizzling_functor_ = swizzling_functor if ConvKind.Fprop in conv_kinds: # Strided support for Optimized and FixedStridedDilation Depthwise Conv for iterator_algorithm in iterator_algorithms: stride_support = StrideSupport.Strided if iterator_algorithm == IteratorAlgorithm.FixedStrideDilation: if tile.stride == [-1, -1] or tile.dilation == [-1,-1]: continue stride_support = StrideSupport.Fixed if iterator_algorithm == IteratorAlgorithm.Optimized: if tile.stride != [-1, -1] or tile.dilation != [-1,-1]: continue new_operation = Conv2dOperation(ConvKind.Fprop, iterator_algorithm, tile.minimum_compute_capability, tile, A, B, C, element_epilogue, stride_support, epilogue_functor, swizzling_functor_, group_mode=GroupMode.Depthwise) manifest.append(new_operation) operations.append(new_operation) return operations class ConvOperation3x: """All parameters of a CUTLASS 3 convolution operation. Unlike CUTLASS 2 convolutions, CUTLASS 3 convolutions do not distinguish between 2-D and 3-D convolutions by kernel class name. Instead, for CUTLASS 3 convolutions, the tensor layouts encode whether the convolution is 2-D or 3-D. Thus, this class deduces the OperationKind (either Conv2d or Conv3d) from the layouts, rather than taking it as a constructor parameter. """ def __init__(self, conv_kind: ConvKind, tile_description: TileDescription, A: TensorDescription, B: TensorDescription, C: TensorDescription, element_compute: Optional[DataType] = None, D: Optional[TensorDescription] = None, kernel_schedule: KernelScheduleType = KernelScheduleType.ScheduleAuto, epilogue_schedule: EpilogueScheduleType = EpilogueScheduleType.ScheduleAuto, tile_scheduler: TileSchedulerType = TileSchedulerType.Default, log_indent_level: int = 1): log_debug_line(f'ConvOperation3x::init: conv_kind: {conv_kind}', log_indent_level) log_indent_level = log_indent_level + 1 self.conv_kind = conv_kind self.tile_description = tile_description self.A = A self.B = B self.C = C self.element_compute = C.element if element_compute is None else element_compute self.kernel_schedule = kernel_schedule self.epilogue_schedule = epilogue_schedule self.arch = tile_description.minimum_compute_capability self.tile_scheduler = tile_scheduler if D == None: self.D = C else: self.D = D self.is_3x = True self.group_mode = GroupMode.NoneGroup # CUTLASS 3 convolutions currently aren't grouped operation_kind = None for layout in (A.layout, B.layout, C.layout): assert(isinstance(layout, LayoutType)) new_operation_kind = convolution_tensor_layout_type_to_operation_kind(layout) if operation_kind is None: operation_kind = new_operation_kind else: # CUTLASS 3 convolutions don't permit mixing 2-D and 3-D layouts. assert(operation_kind == new_operation_kind) assert(operation_kind is not None) self.operation_kind = operation_kind def __str__(self): return f"ConvOperation3x: operation_kind={self.operation_kind}, conv_kind={self.conv_kind}, tile_description={self.tile_description}" def is_complex(self): complex_operators = [ MathOperation.multiply_add_complex, MathOperation.multiply_add_complex_gaussian, MathOperation.multiply_add_complex_fast_f32 ] return self.tile_description.math_instruction.math_operation in complex_operators def is_mixed_input(self): return self.A.element != self.B.element def accumulator_type(self): accum = self.tile_description.math_instruction.element_accumulator if self.is_complex(): return get_complex_from_real(accum) return accum def short_math_name(self): if self.tile_description.math_instruction.math_operation == MathOperation.multiply_add_complex_gaussian: return "g%s" % ShortDataTypeNames[self.accumulator_type()] return ShortDataTypeNames[self.accumulator_type()] def core_name(self): ''' The basic operation kind is prefixed with a letter indicating the accumulation type. ''' inst_shape = '' inst_operation = '' intermediate_type = '' math_operations_map = { MathOperation.xor_popc: 'xor', MathOperation.and_popc: 'and', } tensor_ops = [ OpcodeClass.TensorOp, OpcodeClass.WmmaTensorOp, OpcodeClass.SparseTensorOp, OpcodeClass.BlockScaledTensorOp, ] is_tensor_op = self.tile_description.math_instruction.opcode_class in tensor_ops if is_tensor_op: math_op = self.tile_description.math_instruction.math_operation math_op_string = math_operations_map[math_op] if math_op in math_operations_map.keys() else '' if self.tile_description.math_instruction.element_a != self.A.element and \ self.tile_description.math_instruction.element_a != self.tile_description.math_instruction.element_accumulator: intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a] return "%s%s%s" % (math_op_string, intermediate_type, ConvKindNames[self.conv_kind]) def extended_name(self): '''Generates a string representing the MMA atom. Assumes accumulator type is C type.''' extended_name = "{core_name}_{element_a}{layout_a}_{element_b}{layout_b}_{element_acc}_{element_c}_{element_d}{layout_c}".format( element_a = DataTypeNames[self.A.element], layout_a = ShortLayoutTypeNames[self.A.layout], element_b = DataTypeNames[self.B.element], layout_b = ShortLayoutTypeNames[self.B.layout], element_acc = DataTypeNames[self.accumulator_type()], element_c = DataTypeNames[self.C.element], layout_c = ShortLayoutTypeNames[self.C.layout], element_d = DataTypeNames[self.D.element], core_name = self.core_name()) return extended_name # Generates a short string representing underlying kernel schedule type def kernel_schedule_name(self): return KernelScheduleSuffixes[self.kernel_schedule] # Generates a short string representing underlying epilogue schedule type def epilogue_schedule_name(self): return EpilogueScheduleSuffixes[self.epilogue_schedule] # Generate a short string representing the operation class def opcode_class_name(self): return OpcodeClassNames[self.tile_description.math_instruction.opcode_class] # Generates the full kernel function name def configuration_name(self): ''' The full function name indicates architecture, extended name, tile size, and layout. ''' kernel_name_template = "cutlass3x_sm{ar}_{op}_{ex}{ct}{cs}_{l}_align{al}{t}{k}{e}" return kernel_name_template.format( ar = self.arch, op = self.opcode_class_name(), ex = self.extended_name(), ct = '_' + 'x'.join([str(i) for i in self.tile_description.tile_shape]) if self.tile_description.tile_shape[0] > 0 else "", cs = '_' + 'x'.join([str(i) for i in self.tile_description.cluster_shape]), l = self.tile_description.stages, al = str(max(self.A.alignment, self.B.alignment)), t = TileSchedulerSuffixes[self.tile_scheduler], k = self.kernel_schedule_name(), e = self.epilogue_schedule_name()) def procedural_name(self): return self.configuration_name() def convolution_tensor_layout_type_to_operation_kind(layout: LayoutType) -> OperationKind: if layout == LayoutType.TensorNHWC or layout == LayoutType.TensorKCSR: return OperationKind.Conv2d elif layout == LayoutType.TensorNDHWC or layout == LayoutType.TensorKCSRT: return OperationKind.Conv3d else: raise RuntimeError(f'LayoutType {layout} does not have a corresponding OperationKind') def CreateConvOperator3x(manifest: Manifest, dims_and_alignments: Sequence[Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]], tile_descriptions: Sequence[Sequence[TileDescription]], data_types, schedule_pairs: Sequence[Tuple[KernelScheduleType, KernelScheduleType]] = \ [(KernelScheduleType.ScheduleAuto, EpilogueScheduleType.ScheduleAuto)], complex_transforms: Optional[Sequence[ComplexTransform]] = None, tile_schedulers: Sequence[TileSchedulerType] = [TileSchedulerType.Default], conv_kind: ConvKind = ConvKind.Fprop, log_indent_level: int = 1): """ Create zero or more CUTLASS 3 two-dimensional convolution operators. Create a CUTLASS 3 two-dimensional convolution operator for all feasible combinations of the input parameters. Add the operators to the manifest. dims_and_alignments: 3-level list. Each outer list term is a list [A, B, C]. Each inner list (A, B, or C) has the form [num_spatial_dimensions, alignment]. Both are integers; the first is the number of spatial dimensions (currently, only 2 or 3 are supported), and the second is the byte alignment. We deduce the operation_kind (either OperationKind.Conv2d or OperationKind.Conv3d) from num_spatial_dimensions. This function doesn't take layouts, unlike the GEMM functions. CUTLASS 3 convolutions currently support three input layouts: * TensorNWC for 1-D convolutions, * TensorNHWC for 2-D convolutions, and * TensorNDHWC for 3-D convolutions. Output (C and D) layouts are the same as input layouts, except for Wgrad convolutions, where the layouts are * TensorKCS for 1-D convolutions, * TensorKCSR for 2-D convolutions, and * TensorKCSRT for 3-D convolutions. The output layouts are completely constrained by the input layouts and the convolution kind. tile_descriptions: 2-level list. Outer level has one list per math instruction. Inner level has one TileDescription for each cluster shape. data_types: Either a single data_type dictionary, or a list of them. Keys: 'a_type', 'b_type', 'c_type', 'd_type', 'acc_type', 'epi_type' complex_transforms: Optional list of pairs. First element of each pair is the complex transform for A, and second element of each pair is the complex transform for B. schedule_pairs: [(kernel_schedule, epilogue_schedule), ...] conv_kind: Convolution kind (Fprop, Dgrad, or Wgrad). """ log_debug_line('CreateConvOperator3x', log_indent_level) log_indent_level = log_indent_level + 1 log_debug_line(f'conv_kind: {conv_kind}', log_indent_level) for triple in dims_and_alignments: assert(isinstance(triple, tuple) or isinstance(triple, list)) assert(len(triple) == 3) spatial_dimensionality = None # to be determined by loop below for entry in triple: # [A, B, C] assert(len(entry) == 2) [dim, alignment] = entry assert(type(dim) is int) assert(dim == 2 or dim == 3) assert(type(alignment) is int) assert(alignment > 0) if spatial_dimensionality is None: spatial_dimensionality = dim else: # A, B, and C need to have the same spatial dimensionality assert(spatial_dimensionality == dim) def input_and_output_layouts(spatial_dim: int, kind: ConvKind) -> Tuple[LayoutType, LayoutType]: if spatial_dim == 1: input_layout = LayoutType.TensorNWC if kind == ConvKind.Wgrad: output_layout = LayoutType.TensorKCS else: output_layout = input_layout elif spatial_dim == 2: input_layout = LayoutType.TensorNHWC if kind == ConvKind.Wgrad: output_layout = LayoutType.TensorKCSR else: output_layout = input_layout elif spatial_dim == 3: input_layout = LayoutType.TensorNDHWC if kind == ConvKind.Wgrad: output_layout = LayoutType.TensorKCSRT else: output_layout = input_layout else: assert(False) return (input_layout, output_layout) def dims_to_layouts(A_B_C: Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]) -> \ Tuple[Tuple[LayoutType, int], Tuple[LayoutType, int], Tuple[LayoutType, int]]: [A, B, C] = A_B_C [spatial_dim, alignment] = A [input_layout, output_layout] = input_and_output_layouts(spatial_dim, conv_kind) return ((input_layout, A[1]), (input_layout, B[1]), (output_layout, C[1])) # layouts: list of triples (A, B, C). # Each of A, B, and C has the form [layout, alignment]. layouts = [dims_to_layouts(A_B_C) for A_B_C in dims_and_alignments] if type(data_types) is dict: data_types = [data_types] for s in schedule_pairs: assert(len(s) == 2) if complex_transforms is None: complex_transforms = [(ComplexTransform.none, ComplexTransform.none)] # product produces a one-pass generator, so the loop must call it anew each time. def make_combinations(): return product( layouts, tile_descriptions, data_types, complex_transforms, schedule_pairs, tile_schedulers ) operations = [] for layout_triple, tile_description, data_type, complex_transform_pair, schedule_pair, tile_scheduler in make_combinations(): A_layout, A_alignment = layout_triple[0] A_xform = complex_transform_pair[0] B_layout, B_alignment = layout_triple[1] B_xform = complex_transform_pair[1] C_layout, C_alignment = layout_triple[2] D_layout = C_layout D_alignment = C_alignment A = TensorDescription(data_type["a_type"], A_layout, A_alignment, A_xform) B = TensorDescription(data_type["b_type"], B_layout, B_alignment, B_xform) C = TensorDescription(data_type["c_type"], C_layout, C_alignment) D = TensorDescription(data_type["d_type"], D_layout, D_alignment) element_compute = data_type.get("epi_type", data_type["acc_type"]) kernel_schedule, epilogue_schedule = schedule_pair operation = ConvOperation3x(conv_kind=conv_kind, tile_description=tile_description, A=A, B=B, C=C, element_compute=element_compute, D=D, kernel_schedule=kernel_schedule, epilogue_schedule=epilogue_schedule, tile_scheduler=tile_scheduler, log_indent_level=log_indent_level) log_debug_line(f'Created ConvOperation3x: {str(operation)}', log_indent_level) manifest.append(operation) operations.append(operation) return operations ################################################################################################### ################################################################################################### # def GenerateSM50_Simt(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.Simt, \ MathOperation.multiply_add), MathInstruction( \ [1, 1, 1], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.Simt, \ MathOperation.multiply_add), ] min_cc = 50 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 8], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 8], 2, [1, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) if math_inst.element_a == DataType.f32: conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # # def GenerateSM50_Simt_complex(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.Simt, \ MathOperation.multiply_add_complex), ] min_cc = 50 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 64, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 8], 2, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 8], 2, [4, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # # def GenerateSM50(manifest, cuda_version): GenerateSM50_Simt(manifest, cuda_version) GenerateSM50_Simt_complex(manifest, cuda_version) ################################################################################################### ################################################################################################### # def GenerateSM60_Simt(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.Simt, \ MathOperation.multiply_add), ] min_cc = 60 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 8], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 8], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 8], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 8], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 8], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 8], 2, [1, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # def GenerateSM60_Simt_DepthwiseConv2d(manifest, cuda_version): math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.Simt, \ MathOperation.multiply_add), ] min_cc = 60 max_cc = 1024 alignment_constraints = [8,] filter_3x3 = [3, 3] filter_5x5 = [5, 5] # [stride_h, stride_w] # [-1, -1] means all stride size. strides = [[-1,-1], [1, 1], [2, 2]] # [dilation_h, dilation_w] # [-1, -1] means all dilation size. dilations = [[-1,-1], [1, 1], [2, 2]] #groups per thread block g16 = 16 g32 = 32 g64 = 64 #output shape per thread block npq_1x4x4 = [1, 4, 4] npq_1x8x8 = [1, 8, 8] npq_1x10x10 = [1, 10, 10] tile_descriptions = [] for math_inst in math_instructions: for stride, dilation in product(strides, dilations): tile_descriptions.extend([ # filter3x3 ThreadBlock_output, filter, stage, warp Direct2dConvFixedStrideDilationTileDescription(npq_1x8x8+[g32], filter_3x3, 3, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x8x8+[g64], filter_3x3, 3, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x8x8+[g16], filter_3x3, 3, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x10x10+[g64], filter_3x3, 2, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x4x4+[g32], filter_3x3, 4, stride, dilation, [4, 1, 1], math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x4x4+[g64], filter_3x3, 4, stride, dilation,[4, 1, 1], math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x4x4+[g16], filter_3x3, 4, stride, dilation, [4, 1, 1], math_inst, min_cc, max_cc), # filter5x5 ThreadBlock_output, filter, stage, warp Direct2dConvFixedStrideDilationTileDescription(npq_1x8x8+[g32], filter_5x5, 3, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x8x8+[g64], filter_5x5, 3, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x8x8+[g16], filter_5x5, 3, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x10x10+[g64], filter_5x5, 2, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x4x4+[g32], filter_5x5, 4, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x4x4+[g64], filter_5x5, 4, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc), Direct2dConvFixedStrideDilationTileDescription(npq_1x4x4+[g16], filter_5x5, 4, stride, dilation,[4, 1, 1],math_inst, min_cc, max_cc) ]) data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateDepthwiseConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # # def GenerateSM60(manifest, cuda_version): GenerateSM60_Simt(manifest, cuda_version) GenerateSM60_Simt_DepthwiseConv2d(manifest, cuda_version) ################################################################################################### ################################################################################################### # def GenerateSM61_Simt(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [1, 1, 4], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.Simt, \ MathOperation.multiply_add), ] min_cc = 61 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 32], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 32], 2, [1, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) # # def GenerateSM61(manifest, cuda_version): GenerateSM61_Simt(manifest, cuda_version) ################################################################################################### ################################################################################################### # def GenerateSM70_TensorOp_884(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 1): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [8, 8, 4], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [8, 8, 4], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 70 max_cc = 75 alignment_constraints = [8, 4, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints) # def GenerateSM70_PlanarComplexTensorOp_884(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 1): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] math_instructions = [ MathInstruction( \ [8, 8, 4], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [8, 8, 4], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 70 max_cc = 75 alignment_constraints = [8, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, complex_transforms) # def GenerateSM70_WmmaTensorOp_161616(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 16, 16], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.WmmaTensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 16, 16], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.WmmaTensorOp, \ MathOperation.multiply_add), ] min_cc = 70 max_cc = 1024 alignment_constraints = [8,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) # ################################################################################################## # def GenerateSM70(manifest, cuda_version): GenerateSM70_TensorOp_884(manifest, cuda_version) GenerateSM70_PlanarComplexTensorOp_884(manifest, cuda_version) # To limit build size, WMMA GEMMs are disabled for now. # #GenerateSM70_WmmaTensorOp_161616(manifest, cuda_version) ################################################################################################### ################################################################################################### # def GenerateSM75_TensorOp_1688_FewChannels(manifest, cuda_version, math_inst): min_cc = 75 max_cc = 1024 tile_descriptions = [ TileDescription([128, 64, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 2, [2, 2, 2], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type, [4, 8]) CreateConv2dFewChannelsOperator(manifest, conv_layout, tile_descriptions, data_type, [1, 2, 4]) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, [4, 8]) CreateConv2dFewChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, [1, 2, 4]) # def GenerateSM75_TensorOp_1688(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 2): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 8], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 75 max_cc = 1024 alignment_constraints = [8, 4, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 2, [1, 2, 2], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints) # Separate generator for 'few channels' specializations GenerateSM75_TensorOp_1688_FewChannels(manifest, cuda_version, math_inst) # # def GenerateSM75_PlanarComplexTensorOp_1688(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 2): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 8], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 75 max_cc = 1024 alignment_constraints = [8, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([ 64, 128, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, complex_transforms) # def GenerateSM75_TensorOp_8816_TN(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 2): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [8, 8, 16], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [8, 8, 16], \ DataType.u8, DataType.u8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 75 max_cc = 90 alignment_constraints = [16,] alignment_constraints_small_channels = [16, 8, 4] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 64], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 64], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 64], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 32, 64], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 32], 2, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 32, 32], 2, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, DataType.s32, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombination) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32, ] operations = [] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints_small_channels, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dFewChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints_small_channels, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: op.C.alignment = 16 else: op.C.alignment = 8 # # def GenerateSM75_TensorOp_8816_Interleaved(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 2): return layouts = [ (LayoutType.ColumnMajorInterleaved32, LayoutType.RowMajorInterleaved32, LayoutType.ColumnMajorInterleaved32), ] math_instructions = [ MathInstruction( \ [8, 8, 16], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [8, 8, 16], \ DataType.u8, DataType.u8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 75 max_cc = 90 alignment_constraints = [16,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32, ] operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) conv_layout = (LayoutType.TensorNC32HW32, LayoutType.TensorC32RSK32, LayoutType.TensorNC32HW32) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: op.C.alignment = 8 # # def GenerateSM75_TensorOp_8832_TN(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 2): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [8, 8, 32], \ DataType.s4, DataType.s4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [8, 8, 32], \ DataType.u4, DataType.u4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 75 max_cc = 89 alignment_constraints = [32,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 128], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, DataType.s32, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombination) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32, ] operations = [] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: op.C.alignment = 16 elif op.tile_description.threadblock_shape[1] == 64: op.C.alignment = 8 else: op.C.alignment = 8 # # def GenerateSM75_TensorOp_8832_Interleaved(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 2): return layouts = [ (LayoutType.ColumnMajorInterleaved64, LayoutType.RowMajorInterleaved64, LayoutType.ColumnMajorInterleaved64), ] math_instructions = [ MathInstruction( \ [8, 8, 32], \ DataType.s4, DataType.s4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [8, 8, 32], \ DataType.u4, DataType.u4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 75 max_cc = 89 alignment_constraints = [32,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 128], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 2, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 2, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 2, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32, ] operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) conv_layout = (LayoutType.TensorNC64HW64, LayoutType.TensorC64RSK64, LayoutType.TensorNC64HW64) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: op.C.alignment = 16 # # def GenerateSM75_TensorOp_88128(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [8, 8, 128], \ DataType.b1, DataType.b1, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.xor_popc), ] min_cc = 75 max_cc = { MathOperation.xor_popc: 89, MathOperation.and_popc: 90 } alignment_constraints = [128,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 512], 2, [4, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 256, 512], 2, [2, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 128, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 256, 512], 2, [1, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([256, 64, 512], 2, [4, 1, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 128, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 64, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 64, 512], 2, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), ] data_type = [DataType.b1, DataType.b1, DataType.s32, DataType.s32] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM75_WmmaTensorOp_161616(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 10, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 16, 16], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.WmmaTensorOp, \ MathOperation.multiply_add), ] min_cc = 75 max_cc = 1024 alignment_constraints = [16,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 32], 2, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 2, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, DataType.f32, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) # # def GenerateSM75_Simt_complex(manifest, cuda_version): math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.Simt, \ MathOperation.multiply_add_complex), ] min_cc = 75 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc) ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 ] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # def GenerateSM75(manifest, cuda_version): GenerateSM75_TensorOp_1688(manifest, cuda_version) GenerateSM75_PlanarComplexTensorOp_1688(manifest, cuda_version) GenerateSM75_TensorOp_8816_TN(manifest, cuda_version) GenerateSM75_TensorOp_8816_Interleaved(manifest, cuda_version) GenerateSM75_TensorOp_8832_TN(manifest, cuda_version) GenerateSM75_TensorOp_8832_Interleaved(manifest, cuda_version) GenerateSM75_TensorOp_88128(manifest, cuda_version) #GenerateSM75_WmmaTensorOp_161616(manifest, cuda_version) GenerateSM75_Simt_complex(manifest, cuda_version) ################################################################################################### ################################################################################################### # def GenerateSM80_TensorOp_16816(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 16], \ DataType.bf16, DataType.bf16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 80 max_cc = 1024 alignment_constraints = [8, 4, 2] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 3, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) CreateGemmGroupedOperator(manifest, layouts, tile_descriptions, data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type, [4, 8]) CreateConv3dOperator(manifest, LayoutType.TensorNDHWC, tile_descriptions, data_type, 8) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints) CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, [4, 8]) CreateConv3dOperator(manifest, LayoutType.TensorNDHWC, tile_descriptions, data_type_mixed, 8) # # def GenerateSM80_SparseTensorOp_16832(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 1): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.RowMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 32], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 32], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 32], \ DataType.bf16, DataType.bf16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 80 max_cc = 1024 alignment_constraints = [8] for math_inst in math_instructions: tile_descriptions = [ TileDescription([ 64, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) # # def GenerateSM80_PlanarComplexTensorOp_16816(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] math_instructions = [ MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 16], \ DataType.bf16, DataType.bf16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 80 max_cc = 1024 alignment_constraints = [8, ] for math_inst in math_instructions: tile_descriptions = [ TileDescription([ 64, 128, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmPlanarComplexOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, complex_transforms) # def GenerateSM80_TensorOp_16816_mixed_input_upcast_a(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] # Upcast on Operand A math_instructions = [ MathInstruction( \ [16, 8, 16], \ DataType.s8, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.u8, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.s8, DataType.bf16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.u8, DataType.bf16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.s8, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.u8, DataType.f16, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), ] min_cc = 80 max_cc = 1024 # For mixed-input alignment constraints are a list of lists, where the # inner list contains the alignment constraints for operands/matrices # [[alignA, alignB, alignC],..] alignment_constraints = [[16, 8, 8],] for math_inst in math_instructions: tile_descriptions = [ # 128x128 TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), # 128x64 TileDescription([128, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), # 128x32 TileDescription([128, 32, 64], 9, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), # 128x16 TileDescription([128, 16, 64], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 16, 64], 3, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] # streamk uses more regs which can cause spill for the biggest warp tile size when the accumulators are 32bit. operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination, SwizzlingFunctor.Identity8) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_b != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_b, math_inst.element_accumulator, ] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombination, SwizzlingFunctor.Identity8) for op in operations: if (DataTypeSize[op.C.element] == 16) and \ (op.tile_description.threadblock_shape[1] <= 32): op.C.alignment = 4 # def GenerateSM80_TensorOp_16816_mixed_input_upcast_b(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.s8, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.u8, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.bf16, DataType.s8, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.bf16, DataType.u8, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.s8, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), MathInstruction( \ [16, 8, 16], \ DataType.f16, DataType.u8, DataType.f16, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), ] min_cc = 80 max_cc = 1024 # For mixed-input alignment constraints are a list of lists, where the # inner list contains the alignment constraints for operands/matrices # [[alignA, alignB, alignC],..] alignment_constraints = [[8, 16, 8],] for math_inst in math_instructions: tile_descriptions = [ # 128x128 TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), # 128x64 TileDescription([128, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), # 128x32 TileDescription([128, 32, 64], 9, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 32], 9, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), # 128x16 TileDescription([128, 16, 64], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 16, 64], 3, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 16, 32], 9, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 16, 32], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 16, 32], 3, [2, 1, 1], math_inst, min_cc, max_cc), # 256x16 TileDescription([256, 16, 32], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 16, 32], 3, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] # streamk uses more regs which can cause spill for the biggest warp tile size when the accumulators are 32bit. operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination, SwizzlingFunctor.Identity8) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombination, SwizzlingFunctor.Identity8) for op in operations: if op.tile_description.threadblock_shape[1] <= 32: op.C.alignment = 4 # def GenerateSM80_TensorOp_16832_TN(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 32], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [16, 8, 32], \ DataType.u8, DataType.u8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 80 max_cc = 1024 smem_usage = 164 alignment_constraints = [16,] alignment_constraints_small_channels = [16, 8, 4] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 64], 6, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 64], 6, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, DataType.s32] data_type_mixed = [math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombination) operations = [] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints_small_channels, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) operations += CreateConv2dFewChannelsOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints_small_channels, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: if op.tile_description.threadblock_shape[0] == 32: op.C.alignment = 8 else: op.C.alignment = 16 else: op.C.alignment = 8 # def GenerateSM80_TensorOp_16832_TN_mixed_input_upcast_a(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] # Upcast on Operand A math_instructions = [ MathInstruction( \ [16, 8, 32], \ DataType.s4, DataType.s8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), ] min_cc = 80 max_cc = 1024 # For mixed-input alignment constraints are a list of lists, where the # inner list contains the alignment constraints for operands/matrices # [[alignA, alignB, alignC],..] alignment_constraints = [[32, 16, 4],] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] # streamk uses more regs which can cause spill for the biggest warp tile size when the accumulators are 32bit. operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination, SwizzlingFunctor.Identity8) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. S8 accumulation) if math_inst.element_a != math_inst.element_accumulator: alignment_constraints = [[32, 16, 16],] data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_b, DataType.f32 ] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp, SwizzlingFunctor.Identity8) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: if op.tile_description.threadblock_shape[0] == 32: op.C.alignment = 8 else: op.C.alignment = 16 else: op.C.alignment = 8 # # def GenerateSM80_TensorOp_16832_TN_mixed_input_upcast_b(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] # Upcast on Operand B math_instructions = [ MathInstruction( \ [16, 8, 32], \ DataType.s8, DataType.s4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_mixed_input_upcast), ] min_cc = 80 max_cc = 1024 # For mixed-input alignment constraints are a list of lists, where the # inner list contains the alignment constraints for operands/matrices # [[alignA, alignB, alignC],..] alignment_constraints = [[16, 32, 4],] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 64], 6, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] # streamk uses more regs which can cause spill for the biggest warp tile size when the accumulators are 32bit. operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination, SwizzlingFunctor.Identity8) # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. S8 accumulation) if math_inst.element_a != math_inst.element_accumulator: alignment_constraints = [[16, 32, 16],] data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32, ] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp, SwizzlingFunctor.Identity8) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: if op.tile_description.threadblock_shape[0] == 32: op.C.alignment = 8 else: op.C.alignment = 16 else: op.C.alignment = 8 # # def GenerateSM80_SparseTensorOp_16864_TN(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 1): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor), ] math_inst = \ MathInstruction( \ [16, 8, 64], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate) min_cc = 80 max_cc = 1024 alignment_constraints = [16,] tile_descriptions = [ TileDescription([128, 64, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.s8, DataType.s8, DataType.s32, DataType.s32] data_type_mixed = [DataType.s8, DataType.s8, DataType.s8, DataType.f32] CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) operations = [] operations += CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: op.C.alignment = 16 else: op.C.alignment = 8 # # def GenerateSM80_TensorOp_16832_Interleaved(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajorInterleaved32, LayoutType.RowMajorInterleaved32, LayoutType.ColumnMajorInterleaved32), ] math_instructions = [ MathInstruction( \ [16, 8, 32], \ DataType.s8, DataType.s8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [16, 8, 32], \ DataType.u8, DataType.u8, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 80 max_cc = 1024 alignment_constraints = [16,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 10, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type_mixed = [math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32] operations = CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) conv_layout = (LayoutType.TensorNC32HW32, LayoutType.TensorC32RSK32, LayoutType.TensorNC32HW32) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: op.C.alignment = 8 # # def GenerateSM80_TensorOp_16864_TN(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 64], \ DataType.s4, DataType.s4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [16, 8, 64], \ DataType.u4, DataType.u4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 80 max_cc = 1024 alignment_constraints = [32,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 128], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 256], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 256], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 256], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 256], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 256], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, DataType.s32] data_type_mixed = [math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) operations = [] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombination) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: op.C.alignment = 16 elif op.tile_description.threadblock_shape[1] == 64: op.C.alignment = 8 else: op.C.alignment = 8 # # def GenerateSM80_SparseTensorOp_168128_TN(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 1): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor), ] math_inst = \ MathInstruction( \ [16, 8, 128], \ DataType.s4, DataType.s4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate) min_cc = 80 max_cc = 1024 alignment_constraints = [32,] tile_descriptions = [ TileDescription([ 64, 64, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 256], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 256], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 256], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 256], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 256], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 512], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 512], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 512], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 512], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.s4, DataType.s4, DataType.s32, DataType.s32] data_type_mixed = [DataType.s4, DataType.s4, DataType.s4, DataType.f32] CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) operations = [] operations += CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) for op in operations: if op.tile_description.threadblock_shape[1] > 128: op.C.alignment = 16 else: op.C.alignment = 8 # # def GenerateSM80_TensorOp_16864_Interleaved(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajorInterleaved64, LayoutType.RowMajorInterleaved64, LayoutType.ColumnMajorInterleaved64), ] math_instructions = [ MathInstruction( \ [16, 8, 64], \ DataType.s4, DataType.s4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), MathInstruction( \ [16, 8, 64], \ DataType.u4, DataType.u4, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_saturate), ] min_cc = 80 max_cc = 1024 alignment_constraints = [32,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 6, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type_mixed = [math_inst.element_a, math_inst.element_b, math_inst.element_a, DataType.f32] operations = [] operations += CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints, None, EpilogueFunctor.LinearCombinationClamp) conv_layout = (LayoutType.TensorNC64HW64, LayoutType.TensorC64RSK64, LayoutType.TensorNC64HW64) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombinationClamp) for op in operations: op.C.alignment = 16 # # def GenerateSM80_TensorOp_168256(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 256], \ DataType.b1, DataType.b1, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.xor_popc), MathInstruction( \ [16, 8, 256], \ DataType.b1, DataType.b1, DataType.s32, \ OpcodeClass.TensorOp, \ MathOperation.and_popc), ] min_cc = 80 max_cc = { MathOperation.xor_popc: 89, MathOperation.and_popc: 90 } alignment_constraints = [128,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 512], 3, [4, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 256, 512], 3, [2, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([256, 64, 512], 4, [4, 1, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 256, 512], 4, [1, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 128, 512], 5, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 64, 512], 6, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 128, 512], 6, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 64, 512], 10, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([256, 128, 1024], 3, [4, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 256, 1024], 3, [2, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([256, 64, 1024], 4, [4, 1, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 256, 1024], 4, [1, 4, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 128, 1024], 4, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([128, 64, 1024], 3, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 128, 1024], 3, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), TileDescription([ 64, 64, 1024], 5, [2, 2, 1], math_inst, min_cc, max_cc[math_inst.math_operation]), ] data_type = [DataType.b1, DataType.b1, DataType.s32, DataType.s32] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM80_TensorOp_1688(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) ] min_cc = 80 max_cc = 1024 alignment_constraints = [4, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] data_type_mixed = [ math_inst.element_a, math_inst.element_b, math_inst.element_a, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type_mixed, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type_mixed, alignment_constraints) # # def GenerateSM80_TensorOp_1688_fast_math(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 8], \ DataType.f16, DataType.f16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_fast_f16), MathInstruction( \ [16, 8, 8], \ DataType.bf16, DataType.bf16, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_fast_bf16), ] min_cc = 80 max_cc = 1024 alignment_constraints = [4, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # # def GenerateSM80_TensorOp_1688_fast_fp32_math(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_fast_f32), ] min_cc = 80 max_cc = 1024 alignment_constraints = [4, 2, 1] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # def GenerateSM80_TensorOp_1688_fast_fp32_math_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_fast_f32) min_cc = 80 max_cc = 1024 tile_descriptions = [ TileDescription([128, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 ] alignment_constraints = [1,] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # def GenerateSM80_SparseTensorOp_16816_fast_math(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 1): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.RowMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.RowMajor), ] math_instructions = [ MathInstruction( \ [16, 8, 16], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), ] min_cc = 80 max_cc = 1024 alignment_constraints = [4] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 32], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 32], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32] CreateSparseGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM80_TensorOp_1688_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 80 max_cc = 1024 tile_descriptions = [ TileDescription([128, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 ] alignment_constraints = [1,] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM80_TensorOp_1688_rank_k(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_fast_f32), ] min_cc = 80 max_cc = 1024 alignment_constraints = [1, 2, 4] # Alignment only applies to A in SYRK for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), #TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), #TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f32, DataType.f32, DataType.f32] CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # # def GenerateSM80_TensorOp_1688_rank_k_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex), MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_fast_f32), ] min_cc = 80 max_cc = 1024 for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32 ] alignment_constraints = [1,] # SYRK CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HERK CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM80_TensorOp_1688_trmm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_fast_f32), ] min_cc = 80 max_cc = 1024 alignment_constraints = [1, 2, 4] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), #TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM80_TensorOp_1688_trmm_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex), MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_fast_f32), ] min_cc = 80 max_cc = 1024 for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 ] alignment_constraints = [1,] complex_transforms = [ ComplexTransform.none, ComplexTransform.conj, ] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM80_TensorOp_1688_symm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return # A and B have same layouts layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add), MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_fast_f32), ] min_cc = 80 max_cc = 1024 alignment_constraints = [ 1, 2, 4 ] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), #TileDescription([256, 64, 16], 4, [4, 1, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 256, 16], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 64, 16], 10, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 32], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 32], 3, [2, 4, 1], math_inst, min_cc, max_cc), #TileDescription([256, 64, 32], 4, [4, 1, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 256, 32], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 32], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([128, 64, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 128, 32], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([ 64, 64, 32], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f32, DataType.f32, DataType.f32, DataType.f32] CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # # def GenerateSM80_TensorOp_1688_symm_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_instructions = [ MathInstruction( \ [16, 8, 8], \ DataType.tf32, DataType.tf32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex), MathInstruction( \ [16, 8, 8], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_fast_f32), ] min_cc = 80 max_cc = 1024 for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 64, 16], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([64, 32, 16], 4, [2, 1, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 ] alignment_constraints = [1,] # SYMM CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HEMM CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM80_TensorOp_884(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 16], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([32, 256, 16], 3, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM80_TensorOp_884_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8 ], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8 ], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8 ], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 8 ], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 8 ], 4, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 3, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # def GenerateSM80_TensorOp_884_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM80_TensorOp_884_rank_k(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64] CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # # def GenerateSM80_TensorOp_884_rank_k_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64] # SYRK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HERK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM80_TensorOp_884_rank_k_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ComplexTransform.none,] # SYRK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HERK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM80_TensorOp_884_trmm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM80_TensorOp_884_trmm_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ ComplexTransform.none, ComplexTransform.conj, ] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM80_TensorOp_884_trmm_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ ComplexTransform.none, ComplexTransform.conj, ] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM80_TensorOp_884_symm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64] CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # # def GenerateSM80_TensorOp_884_symm_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] # SYMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HEMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM80_TensorOp_884_symm_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 0): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [8, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 80 max_cc = 1024 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ComplexTransform.none,] # SYMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HEMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # ################################################################################################### # def GenerateSM80_Simt_f32(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.Simt, \ MathOperation.multiply_add), ] min_cc = 80 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 8], 5, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 8], 4, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 8], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 8], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 8], 5, [1, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # # def GenerateSM80_Simt_f64(manifest, cuda_version): layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.Simt, \ MathOperation.multiply_add), ] min_cc = 80 max_cc = 1024 alignment_constraints = [1,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 8], 5, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 8], 5, [1, 2, 1], math_inst, min_cc, max_cc), ] data_type = [ math_inst.element_a, math_inst.element_b, math_inst.element_accumulator, math_inst.element_accumulator, ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints) # ################################################################################################## # def GenerateSM80_Simt_complex(manifest, cuda_version): math_instructions = [ MathInstruction( \ [1, 1, 1], \ DataType.f32, DataType.f32, DataType.f32, \ OpcodeClass.Simt, \ MathOperation.multiply_add_complex), ] min_cc = 80 max_cc = 1024 alignment_constraints = [1,] data_type = [ DataType.cf32, DataType.cf32, DataType.cf32, DataType.cf32 ] layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 128, 8], 5, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 8], 4, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), ] CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, alignment_constraints, complex_transforms) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints) # ################################################################################################### # def GenerateSM80(manifest, cuda_version): GenerateSM80_TensorOp_16816(manifest, cuda_version) GenerateSM80_SparseTensorOp_16832(manifest, cuda_version) GenerateSM80_PlanarComplexTensorOp_16816(manifest, cuda_version) GenerateSM80_TensorOp_1688(manifest, cuda_version) GenerateSM80_TensorOp_1688_fast_math(manifest, cuda_version) GenerateSM80_SparseTensorOp_16816_fast_math(manifest, cuda_version) GenerateSM80_TensorOp_1688_complex(manifest, cuda_version) # 3xTF32 GenerateSM80_TensorOp_1688_fast_fp32_math(manifest, cuda_version) GenerateSM80_TensorOp_1688_fast_fp32_math_complex(manifest, cuda_version) GenerateSM80_TensorOp_1688_rank_k(manifest, cuda_version) GenerateSM80_TensorOp_1688_rank_k_complex(manifest, cuda_version) GenerateSM80_TensorOp_1688_trmm(manifest, cuda_version) GenerateSM80_TensorOp_1688_trmm_complex(manifest, cuda_version) GenerateSM80_TensorOp_1688_symm(manifest, cuda_version) GenerateSM80_TensorOp_1688_symm_complex(manifest, cuda_version) GenerateSM80_TensorOp_884(manifest, cuda_version) GenerateSM80_TensorOp_884_complex(manifest, cuda_version) GenerateSM80_TensorOp_884_complex_gaussian(manifest, cuda_version) GenerateSM80_TensorOp_884_rank_k(manifest, cuda_version) GenerateSM80_TensorOp_884_rank_k_complex(manifest, cuda_version) GenerateSM80_TensorOp_884_rank_k_complex_gaussian(manifest, cuda_version) GenerateSM80_TensorOp_884_trmm(manifest, cuda_version) GenerateSM80_TensorOp_884_trmm_complex(manifest, cuda_version) GenerateSM80_TensorOp_884_trmm_complex_gaussian(manifest, cuda_version) GenerateSM80_TensorOp_884_symm(manifest, cuda_version) GenerateSM80_TensorOp_884_symm_complex(manifest, cuda_version) GenerateSM80_TensorOp_884_symm_complex_gaussian(manifest, cuda_version) GenerateSM80_TensorOp_16816_mixed_input_upcast_a(manifest, cuda_version) GenerateSM80_TensorOp_16816_mixed_input_upcast_b(manifest, cuda_version) GenerateSM80_TensorOp_16832_TN(manifest, cuda_version) GenerateSM80_TensorOp_16832_TN_mixed_input_upcast_a(manifest, cuda_version) GenerateSM80_TensorOp_16832_TN_mixed_input_upcast_b(manifest, cuda_version) GenerateSM80_SparseTensorOp_16864_TN(manifest, cuda_version) GenerateSM80_TensorOp_16832_Interleaved(manifest, cuda_version) GenerateSM80_TensorOp_16864_TN(manifest, cuda_version) GenerateSM80_SparseTensorOp_168128_TN(manifest, cuda_version) GenerateSM80_TensorOp_16864_Interleaved(manifest, cuda_version) GenerateSM80_TensorOp_168256(manifest, cuda_version) GenerateSM80_Simt_f32(manifest, cuda_version) GenerateSM80_Simt_f64(manifest, cuda_version) GenerateSM80_Simt_complex(manifest, cuda_version) ################################################################################################### def GenerateSM89_TensorOp_16832_fp8(manifest, element_acc): layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor) ] math_instructions = [ MathInstruction( [16, 8, 32], DataType.e4m3, DataType.e4m3, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 32], DataType.e4m3, DataType.e5m2, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 32], DataType.e5m2, DataType.e4m3, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 32], DataType.e5m2, DataType.e5m2, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 32], DataType.e4m3, DataType.e4m3, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), MathInstruction( [16, 8, 32], DataType.e4m3, DataType.e5m2, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), MathInstruction( [16, 8, 32], DataType.e5m2, DataType.e4m3, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), MathInstruction( [16, 8, 32], DataType.e5m2, DataType.e5m2, element_acc, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), ] min_cc = 89 max_cc = 100 alignment_constraints = [16,] alignment_constraints_small_channels = [16, 8, 4] for math_inst in math_instructions: tile_descriptions = [ TileDescription([256, 128, 64], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 64], 6, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 64], 6, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 3, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 64], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 256, 64], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 128], 4, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 32, 64], 6, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 32, 128, 64], 6, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 64], 10, [2, 2, 1], math_inst, min_cc, max_cc), ] data_types = [ [ math_inst.element_a, math_inst.element_b, DataType.f32, math_inst.element_accumulator ], [ math_inst.element_a, math_inst.element_b, DataType.bf16, math_inst.element_accumulator ], ] operations = [] for data_type in data_types: operations += CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) conv_layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC) operations += CreateConv2dOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints, [ConvKind.Fprop], EpilogueFunctor.LinearCombination) operations += CreateConv2dFixedChannelsOperator(manifest, conv_layout, tile_descriptions, data_type, alignment_constraints_small_channels, [ConvKind.Fprop], EpilogueFunctor.LinearCombination) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: if op.tile_description.threadblock_shape[0] == 32: op.C.alignment = 8 else: op.C.alignment = 16 else: op.C.alignment = 8 def GenerateSM89_TensorOp_16832_fp8_fp32acc(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 4): return GenerateSM89_TensorOp_16832_fp8(manifest, DataType.f32) def GenerateSM89_TensorOp_16832_fp8_fp16acc(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return GenerateSM89_TensorOp_16832_fp8(manifest, DataType.f16) # def GenerateSM89_SparseTensorOp_16864_fp8(manifest, cuda_version): if ( not CudaToolkitVersionSatisfies(cuda_version, 12, 4) ): return layouts = [ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor) ] math_instructions = [ MathInstruction( [16, 8, 64], DataType.e4m3, DataType.e4m3, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 64], DataType.e4m3, DataType.e5m2, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 64], DataType.e5m2, DataType.e4m3, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 64], DataType.e5m2, DataType.e5m2, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [16, 8, 64], DataType.e4m3, DataType.e4m3, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), MathInstruction( [16, 8, 64], DataType.e4m3, DataType.e5m2, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), MathInstruction( [16, 8, 64], DataType.e5m2, DataType.e4m3, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), MathInstruction( [16, 8, 64], DataType.e5m2, DataType.e5m2, DataType.f32, OpcodeClass.TensorOp, MathOperation.multiply_add_fast_accum), ] min_cc = 89 max_cc = 89 alignment_constraints = [16,] for math_inst in math_instructions: tile_descriptions = [ TileDescription([128, 64, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 128, 128], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 256, 128], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 128], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 128], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 256, 128], 4, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 128], 6, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 128], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 256], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 128, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([ 64, 64, 256], 3, [2, 2, 1], math_inst, min_cc, max_cc), ] data_types = [ [ math_inst.element_a, math_inst.element_b, DataType.f32, math_inst.element_accumulator ], ] operations = [] for data_type in data_types: operations += CreateSparseGemmOperator(manifest, layouts, tile_descriptions, data_type, alignment_constraints, None, EpilogueFunctor.LinearCombination) for op in operations: if op.tile_description.threadblock_shape[1] >= 128: op.C.alignment = 16 else: op.C.alignment = 8 ################################################################################################### # def GenerateSM89(manifest, cuda_version): GenerateSM89_TensorOp_16832_fp8_fp32acc(manifest, cuda_version) GenerateSM89_TensorOp_16832_fp8_fp16acc(manifest, cuda_version) GenerateSM89_SparseTensorOp_16864_fp8(manifest, cuda_version) ################################################################################################### try: from .sm90_utils import ( generate_fp16_bf16_math_instructions_sm90, generate_tf32_math_instructions_sm90, generate_int8_math_instructions_sm90, generate_fp8_math_instructions_sm90, generate_mixed_dtype_math_instructions_sm90, make_sparse_math_instructions, generate_tile_descriptions_sm90, get_valid_schedules, generate_data_types_from_math_instruction, fix_alignments, ) except ImportError: from sm90_utils import ( generate_fp16_bf16_math_instructions_sm90, generate_tf32_math_instructions_sm90, generate_int8_math_instructions_sm90, generate_fp8_math_instructions_sm90, generate_mixed_dtype_math_instructions_sm90, make_sparse_math_instructions, generate_tile_descriptions_sm90, get_valid_schedules, generate_data_types_from_math_instruction, fix_alignments, ) def GenerateSM90_TensorOp_16b_WGMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 3 if is_grouped(gemm_kind) else 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=131, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments. layouts = [ [[LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 1]], ] math_instructions = generate_fp16_bf16_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_type_w_source = generate_data_types_from_math_instruction(math_inst) data_type_wo_source = generate_data_types_from_math_instruction(math_inst, element_source=DataType.void) data_types = [data_type_w_source, data_type_wo_source] # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed_w_source = generate_data_types_from_math_instruction( math_inst, element_source=math_inst.element_a, element_dest=math_inst.element_a ) data_type_mixed_wo_source = generate_data_types_from_math_instruction( math_inst, element_source=DataType.void, element_dest=math_inst.element_a ) data_types.append(data_type_mixed_w_source) data_types.append(data_type_mixed_wo_source) for layout in layouts: for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, gemm_kind=gemm_kind, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules, gemm_kind=gemm_kind) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_16b_WGMMA_alignx_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=101, exhaustive_level=9992) is_aligned = False # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]], ] math_instructions = generate_fp16_bf16_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_type_w_source = generate_data_types_from_math_instruction(math_inst) data_types = [data_type_w_source] # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed_w_source = generate_data_types_from_math_instruction( math_inst, element_source=math_inst.element_a, element_dest=math_inst.element_a ) data_types.append(data_type_mixed_w_source) for layout in layouts: for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_SparseTensorOp_16b_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 2): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=131, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments. layouts = [ [[LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 1]], ] math_instructions = make_sparse_math_instructions(generate_fp16_bf16_math_instructions_sm90(instantiation_level)) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_type_w_source = generate_data_types_from_math_instruction(math_inst) data_type_wo_source = generate_data_types_from_math_instruction(math_inst, element_source=DataType.void) data_types = [data_type_w_source, data_type_wo_source] # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_type_mixed_w_source = generate_data_types_from_math_instruction( math_inst, element_source=math_inst.element_a, element_dest=math_inst.element_a ) data_type_mixed_wo_source = generate_data_types_from_math_instruction( math_inst, element_source=DataType.void, element_dest=math_inst.element_a ) data_types.append(data_type_mixed_w_source) data_types.append(data_type_mixed_wo_source) for layout in layouts: for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_tf32_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=120, default_level=121, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4]], ] math_instructions = generate_tf32_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction for layout in layouts: data_type_tf32 = generate_data_types_from_math_instruction(math_inst) data_type_tf32_wo_source = generate_data_types_from_math_instruction(math_inst, element_source=DataType.void) data_type_f32 = copy.deepcopy(data_type_tf32) data_type_f32_wo_source = copy.deepcopy(data_type_tf32_wo_source) data_type_f32["a_type"] = DataType.f32 data_type_f32["b_type"] = DataType.f32 data_type_f32["epi_type"] = DataType.f32 data_type_f32_wo_source["a_type"] = DataType.f32 data_type_f32_wo_source["b_type"] = DataType.f32 data_type_f32_wo_source["epi_type"] = DataType.f32 data_types = [data_type_tf32, data_type_f32, data_type_tf32_wo_source, data_type_f32_wo_source] for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_tf32_WGMMA_alignx_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=101, exhaustive_level=9992) is_aligned = False # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 2], [LayoutType.RowMajor, 2], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 1], [LayoutType.ColumnMajor, 1], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 1], [LayoutType.RowMajor, 1], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 1], [LayoutType.ColumnMajor, 1], [LayoutType.ColumnMajor, 1]], [[LayoutType.ColumnMajor, 1], [LayoutType.RowMajor, 1], [LayoutType.ColumnMajor, 1]], ] math_instructions = generate_tf32_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction for layout in layouts: # Inconsistency: TF32 does not stamp out void-C data_type_tf32 = generate_data_types_from_math_instruction(math_inst) data_type_f32 = copy.deepcopy(data_type_tf32) data_type_f32["a_type"] = DataType.f32 data_type_f32["b_type"] = DataType.f32 data_type_f32["epi_type"] = DataType.f32 for data_type in [data_type_tf32, data_type_f32]: # Inconsistency: alignments aren't fixed in TF32 / alignx # layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_SparseTensorOp_tf32_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 2): return instantiation_level = manifest.get_instantiation_level(pruned_level=120, default_level=121, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4]], ] math_instructions = make_sparse_math_instructions(generate_tf32_math_instructions_sm90(instantiation_level)) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction for layout in layouts: data_type_tf32 = generate_data_types_from_math_instruction(math_inst) data_type_tf32_wo_source = generate_data_types_from_math_instruction(math_inst, element_source=DataType.void) data_type_f32 = copy.deepcopy(data_type_tf32) data_type_f32_wo_source = copy.deepcopy(data_type_tf32_wo_source) data_type_f32["a_type"] = DataType.f32 data_type_f32["b_type"] = DataType.f32 data_type_f32["epi_type"] = DataType.f32 data_type_f32_wo_source["a_type"] = DataType.f32 data_type_f32_wo_source["b_type"] = DataType.f32 data_type_f32_wo_source["epi_type"] = DataType.f32 data_types = [data_type_tf32, data_type_f32, data_type_tf32_wo_source, data_type_f32_wo_source] for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_int8_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=111, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16]], ] math_instructions = generate_int8_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_type_w_source = generate_data_types_from_math_instruction(math_inst) data_type_wo_source = generate_data_types_from_math_instruction(math_inst, element_source=DataType.void) data_type_int8_output = generate_data_types_from_math_instruction( math_inst, element_source=DataType.s8, element_dest=math_inst.element_a, element_epilogue=DataType.f32 ) data_types = [data_type_w_source, data_type_wo_source, data_type_int8_output] for layout in layouts: for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_int8_WGMMA_alignx_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=111, exhaustive_level=9992) is_aligned = False # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], ] math_instructions = generate_int8_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_type_w_source = generate_data_types_from_math_instruction(math_inst) data_type_int8_output = generate_data_types_from_math_instruction( math_inst, element_source=DataType.s8, element_dest=math_inst.element_a, element_epilogue=DataType.f32 ) data_types = [data_type_w_source, data_type_int8_output] for layout in layouts: for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_SparseTensorOp_int8_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 2): return instantiation_level = manifest.get_instantiation_level(pruned_level=100, default_level=111, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16]], ] math_instructions = make_sparse_math_instructions(generate_int8_math_instructions_sm90(instantiation_level)) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction # s8.u8 and u8.s8 wgmma variants require PTX 8.4 if math_inst.element_a != math_inst.element_b and not CudaToolkitVersionSatisfies(cuda_version, 12, 4): continue data_type_w_source = generate_data_types_from_math_instruction(math_inst) data_type_wo_source = generate_data_types_from_math_instruction(math_inst, element_source=DataType.void) data_type_int8_output = generate_data_types_from_math_instruction( math_inst, element_source=DataType.s8, element_dest=math_inst.element_a, element_epilogue=DataType.f32 ) data_types = [data_type_w_source, data_type_wo_source, data_type_int8_output] for layout in layouts: for data_type in data_types: layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_fp8_WGMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 3 if is_grouped(gemm_kind) else 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=20, default_level=121, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 1]], # TN Layout ] math_instructions = generate_fp8_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_types = [] fp8_types = [DataType.e4m3, DataType.e5m2] valid_types_for_d = [DataType.f32, DataType.bf16, DataType.f16, DataType.e4m3, DataType.e5m2] valid_types_for_c = copy.deepcopy(valid_types_for_d) valid_types_for_c.append(DataType.void) for c_type, d_type in product(valid_types_for_c, valid_types_for_d): data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=c_type, element_dest=d_type, ) ) else: for d_type in valid_types_for_d: data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=DataType.void, element_dest=d_type, ) ) for layout in layouts: for data_type in data_types: # Inconsistency: alignments aren't fixed in FP8 # layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, gemm_kind=gemm_kind, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules, gemm_kind=gemm_kind) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_fp8_WGMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 3 if is_grouped(gemm_kind) else 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=20, default_level=121, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 1]], # TN Layout ] math_instructions = generate_fp8_math_instructions_sm90(instantiation_level) tile_descriptions_ = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) tile_descriptions = list() for desc in tile_descriptions_: desc.explicit_vector_sizes = [1, desc.tile_shape[1], desc.tile_shape[2]] tile_descriptions.append(copy.deepcopy(desc)) desc.explicit_vector_sizes = [desc.tile_shape[0], desc.tile_shape[1], desc.tile_shape[2]] tile_descriptions.append(copy.deepcopy(desc)) desc.explicit_vector_sizes = [desc.tile_shape[0], desc.tile_shape[1], desc.tile_shape[2]] tile_descriptions.append(copy.deepcopy(desc)) desc.explicit_vector_sizes = [1, 1, desc.tile_shape[2]] tile_descriptions.append(copy.deepcopy(desc)) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_types = [] fp8_types = [DataType.e4m3, DataType.e5m2] valid_types_for_d = [DataType.f32, DataType.bf16, DataType.f16, DataType.e4m3, DataType.e5m2] valid_types_for_c = copy.deepcopy(valid_types_for_d) valid_types_for_c.append(DataType.void) for c_type, d_type in product(valid_types_for_c, valid_types_for_d): data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=c_type, element_dest=d_type, ) ) else: for d_type in valid_types_for_d: data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=DataType.void, element_dest=d_type, ) ) for layout in layouts: for data_type in data_types: # Inconsistency: alignments aren't fixed in FP8 # layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, gemm_kind=gemm_kind, enable_fp8_fast_acc=False, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules, gemm_kind=gemm_kind) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK], gemm_kind=gemm_kind) def GenerateSM90_TensorOp_fp8_WGMMA_alignx_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=0, default_level=101, exhaustive_level=9992) is_aligned = False # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 1]], # TN Layout [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 1]], # TN Layout ] math_instructions = generate_fp8_math_instructions_sm90(instantiation_level) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_types = [generate_data_types_from_math_instruction(math_inst)] fp8_types = [DataType.e4m3, DataType.e5m2] valid_types_for_d = [DataType.f32, DataType.bf16, DataType.f16, DataType.e4m3, DataType.e5m2] valid_types_for_c = copy.deepcopy(valid_types_for_d) valid_types_for_c.append(DataType.void) for c_type, d_type in product(valid_types_for_c, valid_types_for_d): data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=c_type, element_dest=d_type, ) ) for layout in layouts: for data_type in data_types: # Inconsistency: alignments aren't fixed in FP8 # layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_mixed_dtype_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 1): return instantiation_level = manifest.get_instantiation_level(pruned_level=20, default_level=121, exhaustive_level=9999) is_aligned = True # layouts for ABC, their alignments will be fixed later based on the data type layouts = [ [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16]], ] valid_types_for_a_b_acc = [ (DataType.e4m3, DataType.f16, DataType.f32), (DataType.e4m3, DataType.bf16, DataType.f32), (DataType.e5m2, DataType.f16, DataType.f32), (DataType.e5m2, DataType.bf16, DataType.f32), (DataType.s8, DataType.f16, DataType.f32), (DataType.s8, DataType.bf16, DataType.f32), (DataType.u8, DataType.f16, DataType.f32), (DataType.u8, DataType.bf16, DataType.f32), (DataType.s4, DataType.f16, DataType.f32), (DataType.s4, DataType.bf16, DataType.f32), (DataType.s4, DataType.e4m3, DataType.f32), (DataType.s4, DataType.e5m2, DataType.f32), (DataType.u4, DataType.f16, DataType.f32), (DataType.u4, DataType.bf16, DataType.f32), (DataType.u2, DataType.f16, DataType.f32), (DataType.u2, DataType.bf16, DataType.f32), (DataType.s2, DataType.f16, DataType.f32), (DataType.s2, DataType.bf16, DataType.f32), ] # Note: For sizeof(a_type) > sizeof(b_type), some generated kernels might crash due to a compiler bug. Disable it for now. #swapped_valid_types_for_a_b_acc = [(b_type, a_type, acc_type) for a_type, b_type, acc_type in valid_types_for_a_b_acc] #valid_types_for_a_b_acc = valid_types_for_a_b_acc + swapped_valid_types_for_a_b_acc math_instructions = generate_mixed_dtype_math_instructions_sm90(instantiation_level, valid_types_for_a_b_acc) valid_types_for_d = [DataType.f32, DataType.bf16, DataType.f16, DataType.e4m3, DataType.e5m2] valid_types_for_c = copy.deepcopy(valid_types_for_d) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_types = [] # Limit C/D types to avoid a giant number of instantiations. # A typical use case for mixed dtype in DL is weight quantization (tensor A), # therefore we can limit the output type to that of activation (tensor B). valid_types_for_c = [math_inst.element_b] valid_types_for_d = [math_inst.element_b] for c_type, d_type in product(valid_types_for_c, valid_types_for_d): data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=c_type, element_dest=d_type, ) ) for layout in layouts: for data_type in data_types: # Fix alignments, DataTypeSize are in the unit of bits alignment_bits = 128 layout[0][1] = alignment_bits // DataTypeSize[data_type['a_type']] layout[1][1] = alignment_bits // DataTypeSize[data_type['b_type']] layout[2][1] = alignment_bits // DataTypeSize[data_type['c_type']] schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_SparseTensorOp_fp8_WGMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 2): return instantiation_level = manifest.get_instantiation_level(pruned_level=20, default_level=121, exhaustive_level=9992) is_aligned = True # layouts for ABC and their alignments layouts = [ [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 1]], # TN Layout ] math_instructions = make_sparse_math_instructions(generate_fp8_math_instructions_sm90(instantiation_level)) tile_descriptions = generate_tile_descriptions_sm90( math_instructions=math_instructions, is_aligned=is_aligned, level=instantiation_level) for tile_desc in tile_descriptions: math_inst = tile_desc.math_instruction data_types = [] fp8_types = [DataType.e4m3, DataType.e5m2] valid_types_for_d = [DataType.f32, DataType.bf16, DataType.f16, DataType.e4m3, DataType.e5m2] valid_types_for_c = copy.deepcopy(valid_types_for_d) valid_types_for_c.append(DataType.void) for c_type, d_type in product(valid_types_for_c, valid_types_for_d): data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=c_type, element_dest=d_type, ) ) else: for d_type in valid_types_for_d: data_types.append( generate_data_types_from_math_instruction( math_inst, element_source=DataType.void, element_dest=d_type, ) ) for layout in layouts: for data_type in data_types: # Inconsistency: alignments aren't fixed in FP8 # layout = fix_alignments(data_type, layout, alignment_bits=128) schedules, stream_k_schedules = get_valid_schedules( tile_description=tile_desc, cuda_version=cuda_version, is_aligned=is_aligned, data_types=data_type, instantiation_level=instantiation_level, layout=layout, ) if len(schedules): CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, schedules) if len(stream_k_schedules): assert CudaToolkitVersionSatisfies(cuda_version, 12, 1) CreateSparseGemmUniversal3xOperator(manifest, [layout], [tile_desc], data_type, stream_k_schedules, tile_schedulers=[TileSchedulerType.StreamK]) def GenerateSM90_TensorOp_1684(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = MathInstruction( [16, 8, 4], DataType.f64, DataType.f64, DataType.f64, OpcodeClass.TensorOp, MathOperation.multiply_add) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([256, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 256, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([256, 32, 16], 3, [4, 1, 1], math_inst, min_cc, max_cc), TileDescription([32, 256, 16], 3, [1, 4, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64] CreateGemmOperator(manifest, layouts, tile_descriptions, data_type, alignment_constraints) # # def GenerateSM90_TensorOp_1684_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8 ], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8 ], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8 ], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 8 ], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 8 ], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 8 ], 4, [2, 1, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 3, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM90_TensorOp_1684_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.RowMajor, LayoutType.ColumnMajor), ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ (ComplexTransform.none, ComplexTransform.none), (ComplexTransform.conj, ComplexTransform.none), (ComplexTransform.none, ComplexTransform.conj), (ComplexTransform.conj, ComplexTransform.conj) ] CreateGemmOperator(manifest, layouts, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM90_TensorOp_1684_rank_k(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64] CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # # def GenerateSM90_TensorOp_1684_rank_k_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64] # SYRK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HERK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM90_TensorOp_1684_rank_k_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor), ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ComplexTransform.none,] # SYRK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HERK computation CreateRankKOperator(manifest, layouts, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM90_TensorOp_1684_trmm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints) # # def GenerateSM90_TensorOp_1684_trmm_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ ComplexTransform.none, ComplexTransform.conj, ] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM90_TensorOp_1684_trmm_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] diag_types = [ DiagType.NonUnit, DiagType.Unit, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ ComplexTransform.none, ComplexTransform.conj, ] CreateTrmmOperator(manifest, layouts, side_modes, fill_modes, diag_types, tile_descriptions, \ data_type, alignment_constraints, complex_transforms) # # def GenerateSM90_TensorOp_1684_symm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 128, 16], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([128, 64, 16], 3, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 16], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 32, 16], 5, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([16, 32, 16], 5, [1, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 16, 16], 5, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.f64, DataType.f64, DataType.f64, DataType.f64] CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # # def GenerateSM90_TensorOp_1684_symm_complex(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([128, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 128, 8], 3, [2, 4, 1], math_inst, min_cc, max_cc), TileDescription([64, 64, 8], 3, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] # SYMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HEMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # def GenerateSM90_TensorOp_1684_symm_complex_gaussian(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 11, 8): return layouts = [ (LayoutType.ColumnMajor, LayoutType.ColumnMajor), ] side_modes = [ SideMode.Left, SideMode.Right, ] fill_modes = [ FillMode.Lower, FillMode.Upper, ] math_inst = \ MathInstruction( \ [16, 8, 4], \ DataType.f64, DataType.f64, DataType.f64, \ OpcodeClass.TensorOp, \ MathOperation.multiply_add_complex_gaussian) min_cc = 90 max_cc = 90 alignment_constraints = [1,] tile_descriptions = [ TileDescription([64, 64, 8], 3, [4, 2, 1], math_inst, min_cc, max_cc), TileDescription([64, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), TileDescription([32, 64, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 32, 8], 4, [2, 2, 1], math_inst, min_cc, max_cc), #TileDescription([16, 32, 8], 4, [1, 2, 1], math_inst, min_cc, max_cc), #TileDescription([32, 16, 8], 4, [2, 1, 1], math_inst, min_cc, max_cc), ] data_type = [DataType.cf64, DataType.cf64, DataType.cf64, DataType.cf64] complex_transforms = [ComplexTransform.none,] # SYMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.symmetric) # HEMM computation CreateSymmOperator(manifest, layouts, side_modes, fill_modes, tile_descriptions, \ data_type, alignment_constraints, BlasMode.hermitian) # # Blackwell SM 100 generators try: import cutlass_library.sm100_utils from cutlass_library.sm100_utils import ( generate_tf32_math_instructions_sm100, generate_16b_math_instructions_sm100, generate_f8f6f4_math_instructions_sm100, generate_mxf8f6f4_math_instructions_sm100, generate_mxf4nvf4_math_instructions_sm100, generate_sparse_mxf4nvf4_math_instructions_sm100, generate_sparse_mxf8f6f4_math_instructions_sm100, generate_fp8_math_instructions_sm100, generate_cluster_shapes_sm100, get_pruning_level_from_global_level ) except ImportError: import sm100_utils from sm100_utils import ( generate_tf32_math_instructions_sm100, generate_16b_math_instructions_sm100, generate_f8f6f4_math_instructions_sm100, generate_mxf8f6f4_math_instructions_sm100, generate_mxf4nvf4_math_instructions_sm100, generate_sparse_mxf4nvf4_math_instructions_sm100, generate_sparse_mxf8f6f4_math_instructions_sm100, generate_fp8_math_instructions_sm100, generate_cluster_shapes_sm100, get_pruning_level_from_global_level ) ################################################################################################### def get_tma_alignment_elt(data_type : DataType, is_f8f6f4 : bool = True ) -> int: if DataTypeSize[data_type] < 8 and is_f8f6f4: return int(128) return int(16 * 8 / DataTypeSize[data_type]) sm100_cluster_shape_1sm = [ [4,4,1] , DynamicClusterShape ] sm100_cluster_shape_2sm = [ # cluster_m % 2 == 0 for 2sm [4,4,1] , DynamicClusterShape ] def GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999) # layouts for ABC and their alignments. layouts = [ [[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4]], [[LayoutType.ColumnMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4]], [[LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4]], [[LayoutType.RowMajor, 4], [LayoutType.ColumnMajor, 4], [LayoutType.RowMajor, 4]], [[LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4], [LayoutType.RowMajor, 4]], ] data_types = [ { "a_type" : DataType.f32, "b_type" : DataType.f32, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, { "a_type" : DataType.f32, "b_type" : DataType.f32, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) math_instructions_1sm, math_instructions_2sm = generate_tf32_math_instructions_sm100(instantiation_level) cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level) if thor_sm in manifest.compute_capabilities_baseline : if [4,4,1] in cluster_shapes_1sm : cluster_shapes_1sm.remove([4,4,1]) if [4,4,1] in cluster_shapes_2sm : cluster_shapes_2sm.remove([4,4,1]) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, [[KernelScheduleType.TmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) if math_inst.instruction_shape[0] == 128: epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2Sm else: epi_schedule = EpilogueScheduleType.ScheduleAuto CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, [[KernelScheduleType.TmaWarpSpecialized2SmSm100, epi_schedule]], tile_schedulers=tile_schedulers) def GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=490, default_level=490, exhaustive_level=9999) # layouts for ABC and their alignments. C alignment will be set later based on output type layouts = [ [[LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 0]], [[LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 8], [LayoutType.ColumnMajor, 8], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 8], [LayoutType.RowMajor, 8], [LayoutType.RowMajor, 0]], ] thor_sm = ThorSMRenumbering(cuda_version) math_instructions_1sm, math_instructions_2sm = generate_16b_math_instructions_sm100(instantiation_level) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) grouped = is_grouped(gemm_kind) cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level) if thor_sm in manifest.compute_capabilities_baseline : if [4,4,1] in cluster_shapes_1sm : cluster_shapes_1sm.remove([4,4,1]) if [4,4,1] in cluster_shapes_2sm : cluster_shapes_2sm.remove([4,4,1]) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_accumulator, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] kernel_schedule = KernelScheduleType.TmaWarpSpecialized1SmSm100 if not grouped else KernelScheduleType.PtrArrayTmaWarpSpecialized1SmSm100 epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1Sm if not grouped else EpilogueScheduleType.PtrArrayTmaWarpSpecialized1Sm CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_types_mixed = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_a, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types_mixed[0]["d_type"]] CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_accumulator, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] if grouped: epi_schedule = EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm elif math_inst.instruction_shape[0] == 128: epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2Sm else: epi_schedule = EpilogueScheduleType.ScheduleAuto kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized2SmSm100, grouped) CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_types_mixed = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_a, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types_mixed[0]["d_type"]] CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) def GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=591 , default_level=591 , exhaustive_level=9999) # layouts for ABC and their alignments. layouts = [ [[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) epi_type = DataType.f32 grouped = is_grouped(gemm_kind) math_instructions_1sm, math_instructions_2sm = generate_fp8_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped) cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level) if thor_sm in manifest.compute_capabilities_baseline : if [4,4,1] in cluster_shapes_1sm : cluster_shapes_1sm.remove([4,4,1]) if [4,4,1] in cluster_shapes_2sm : cluster_shapes_2sm.remove([4,4,1]) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type in data_types: if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ ( data_type["d_type"] == DataType.e5m2 ): continue kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized1SmSm100, grouped) epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped) CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type in data_types: if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ ( data_type["d_type"] == DataType.e5m2 ): continue if grouped: epi_schedule = EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm elif math_inst.instruction_shape[0] == 128: epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2Sm else: epi_schedule = EpilogueScheduleType.ScheduleAuto kernel_schedule = to_grouped_schedule(KernelScheduleType.TmaWarpSpecialized2SmSm100, grouped) CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, epi_schedule]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) def GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=593, default_level=593, exhaustive_level=9999) grouped = is_grouped(gemm_kind) # layouts for ABC and their alignments. layouts = [ [[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]], ] min_cc = 100 max_cc = 100 epi_type = DataType.f32 pruning_level = get_pruning_level_from_global_level(instantiation_level) math_instructions_1sm, math_instructions_2sm = generate_fp8_math_instructions_sm100(instantiation_level, enable_compile_time_dtype=grouped or pruning_level >= 1, enable_runtime_dtype=not grouped) cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level) tile_schedulers = [ TileSchedulerType.Default, ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape, [math_inst.instruction_shape[0], math_inst.instruction_shape[1], math_inst.instruction_shape[2] * 4])) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape, [1, math_inst.instruction_shape[1], math_inst.instruction_shape[2] * 4])) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape, [math_inst.instruction_shape[0], 1, math_inst.instruction_shape[2] * 4])) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) for data_type in data_types: if ( data_type["a_type"] == DataType.e4m3 ) and ( data_type["b_type"] == DataType.e4m3 ) and\ ( data_type["d_type"] == DataType.e5m2 ): continue is_runtime_datatype_a = is_runtime_datatype(data_type["a_type"]) is_runtime_datatype_b = is_runtime_datatype(data_type["d_type"]) # A/B datatypes should be both static or dynamic if (is_runtime_datatype_a != is_runtime_datatype_b): continue kernel_schedule = to_grouped_schedule(KernelScheduleType.BlockwiseTmaWarpSpecialized1SmSm100, grouped) epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped) epi_schedule_nosmem = to_grouped_schedule(EpilogueScheduleType.BlockwiseNoSmemWarpSpecialized1Sm, grouped) CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, epi_schedule], [kernel_schedule, epi_schedule_nosmem]], tile_schedulers=tile_schedulers, gemm_kind=gemm_kind) def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.Universal3x): # SM100 MMA with mixed F4/F6/F8 inputs + without block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return instantiation_level = manifest.get_instantiation_level(pruned_level=590, default_level=590, exhaustive_level=9999) grouped = is_grouped(gemm_kind) # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]], ] math_instructions_1sm, math_instructions_2sm = generate_f8f6f4_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped) def change_priority_func(shapes_1sm, shapes_2sm): shapes_1sm[(1,2,1)] = 6 shapes_1sm[(1,4,1)] = 6 shapes_2sm[(2,2,1)] = 6 shapes_2sm[(2,4,1)] = 6 shapes_2sm[(4,2,1)] = 6 cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level, change_priority_func) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) epi_type = DataType.f32 is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) if thor_sm in manifest.compute_capabilities_baseline : if [4,4,1] in cluster_shapes_1sm : cluster_shapes_1sm.remove([4,4,1]) if [4,4,1] in cluster_shapes_2sm : cluster_shapes_2sm.remove([4,4,1]) # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) kernel_data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, } ] for kernel_data_type in kernel_data_types: # Filter out some kernel if ( kernel_data_type["a_type"] == DataType.e4m3 ) and ( kernel_data_type["b_type"] == DataType.e4m3 ) and\ ( kernel_data_type["d_type"] == DataType.e5m2 ): continue # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.TmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) kernel_data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, } ] for kernel_data_type in kernel_data_types: # Filter some kernel if ( kernel_data_type["a_type"] == DataType.e4m3 ) and ( kernel_data_type["b_type"] == DataType.e4m3 ) and\ ( kernel_data_type["d_type"] == DataType.e5m2 ): continue # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) if math_inst.instruction_shape[0] == 128: CreateGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.TmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) else: CreateGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.TmaWarpSpecialized2SmSm100, EpilogueScheduleType.ScheduleAuto]], tile_schedulers=tile_schedulers) def GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x): # SM100 MMA with mixed F4/F6/F8 inputs + block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=590, default_level=590, exhaustive_level=9999) grouped = is_grouped(gemm_kind) layouts = [ [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 128], [LayoutType.RowMajor, 0]], ] math_instructions_1sm, math_instructions_2sm = generate_mxf8f6f4_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped) def change_priority_func(shapes_1sm, shapes_2sm): shapes_1sm[(1,2,1)] = 6 shapes_1sm[(1,4,1)] = 6 shapes_2sm[(2,2,1)] = 6 shapes_2sm[(2,4,1)] = 6 shapes_2sm[(4,2,1)] = 6 cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level, change_priority_func) ab_types = [ DataType.f4, DataType.f6, DataType.e2m1, DataType.e2m3, DataType.e3m2, DataType.e5m2, DataType.e4m3, ] acc_types = [ DataType.f32 ] def tile_schedulers(sfdtype): # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K. if sfdtype["type"] == DataType.void or grouped: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) epi_type = DataType.f32 is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) if thor_sm in manifest.compute_capabilities_baseline : if [4,4,1] in cluster_shapes_1sm : cluster_shapes_1sm.remove([4,4,1]) if [4,4,1] in cluster_shapes_2sm : cluster_shapes_2sm.remove([4,4,1]) # 1xSM MMA kernels for math_inst in math_instructions_1sm: assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e3m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type in data_types: CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[to_grouped_schedule(KernelScheduleType.Mxf8f6f4TmaWarpSpecialized1SmSm100, grouped), to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped)]] , tile_schedulers = tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) for math_inst in math_instructions_2sm: assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e3m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, ] # Set alignment d based on Destination format. for data_type in data_types: for layout in layouts: # alignment for a layout[0][1] = get_tma_alignment_elt(data_type["a_type"]) # alignment for b layout[1][1] = get_tma_alignment_elt(data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(data_type["d_type"]) for tile in tile_descriptions: math_inst = tile.math_instruction # Filter some kernels that does not meet the alignment requirements. if layout[0][0] == LayoutType.ColumnMajor: if math_inst.instruction_shape[0] // 2 % layout[0][1] != 0: continue else: if tile.threadblock_shape[2] // tile.cluster_shape[2] % layout[0][1] != 0: continue if layout[1][0] == LayoutType.RowMajor: if math_inst.instruction_shape[1] // 2 % layout[1][1] != 0: continue else: if tile.threadblock_shape[2] // tile.cluster_shape[2] % layout[1][1] != 0: continue if grouped: CreateGemmUniversal3xOperator(manifest, [layout], [tile], [data_type], [[to_grouped_schedule(KernelScheduleType.Mxf8f6f4TmaWarpSpecialized2SmSm100, grouped), to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized2Sm, grouped)]] , tile_schedulers = tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) elif math_inst.instruction_shape[0] == 128: CreateGemmUniversal3xOperator(manifest, [layout], [tile], [data_type], [[KernelScheduleType.Mxf8f6f4TmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]] , tile_schedulers = tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) else: CreateGemmUniversal3xOperator(manifest, [layout], [tile], [data_type], [[KernelScheduleType.Mxf8f6f4TmaWarpSpecialized2SmSm100, EpilogueScheduleType.ScheduleAuto]] , tile_schedulers = tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) def GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x): # SM100 MMA with F4 + block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999) grouped = is_grouped(gemm_kind) # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]], ] math_instructions_1sm, math_instructions_2sm = generate_mxf4nvf4_math_instructions_sm100(instantiation_level, enable_runtime_dtype=not grouped) def change_priority_func(shapes_1sm, shapes_2sm): shapes_1sm[(1,2,1)] = 6 shapes_1sm[(1,4,1)] = 6 shapes_2sm[(2,2,1)] = 6 shapes_2sm[(2,4,1)] = 6 shapes_2sm[(4,2,1)] = 6 cluster_shapes_1sm, cluster_shapes_2sm = generate_cluster_shapes_sm100(instantiation_level, change_priority_func=change_priority_func) acc_types = [ DataType.f32 ] # Accumulator is always 32 bits for block scaled MMA instructions def tile_schedulers(sfdtype): # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K. if sfdtype["type"] == DataType.void or grouped: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) epi_type = DataType.f32 is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) if thor_sm in manifest.compute_capabilities_baseline : if [4,4,1] in cluster_shapes_1sm : cluster_shapes_1sm.remove([4,4,1]) if [4,4,1] in cluster_shapes_2sm : cluster_shapes_2sm.remove([4,4,1]) # 1xSM MMA kernels for math_inst in math_instructions_1sm: assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) assert math_inst.instruction_shape[2] * 4 == 256 data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for layout in layouts: for data_type in data_types: if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue # E2M1 x E2M1, vector size 32, E8 # E2M1 x E2M1, vector size 16, UE4M3 isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 epi_schedule = to_grouped_schedule(EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped) epi_nosmem_schedule = to_grouped_schedule(EpilogueScheduleType.NoSmemWarpSpecialized1Sm, grouped) nvfp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100, grouped) fp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Mxf4TmaWarpSpecialized1SmSm100, grouped) nvfp4_schedules = [[nvfp4_kernel_schedule, epi_schedule]] fp4_schedules = [[fp4_kernel_schedule, epi_schedule]] if (data_type["sfd_type"]["type"] == DataType.void): nvfp4_schedules.append([nvfp4_kernel_schedule, epi_nosmem_schedule]) fp4_schedules.append([fp4_kernel_schedule, epi_nosmem_schedule]) CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) if isFp4: CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) for math_inst in math_instructions_2sm: assert math_inst.opcode_class == OpcodeClass.BlockScaledTensorOp tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for layout in layouts: for data_type in data_types: if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue # E2M1 x E2M1, vector size 32, E8 isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 epi_schedule = EpilogueScheduleType.ScheduleAuto if not grouped else EpilogueScheduleType.PtrArrayTmaWarpSpecialized2Sm epi_nosmem_schedule = to_grouped_schedule(EpilogueScheduleType.NoSmemWarpSpecialized2Sm, grouped) nvfp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100, grouped) fp4_kernel_schedule = to_grouped_schedule(KernelScheduleType.Mxf4TmaWarpSpecialized2SmSm100, grouped) nvfp4_schedules = [[nvfp4_kernel_schedule, epi_schedule]] fp4_schedules = [[fp4_kernel_schedule, epi_schedule]] if (data_type["sfd_type"]["type"] == DataType.void): nvfp4_schedules.append([nvfp4_kernel_schedule, epi_nosmem_schedule]) fp4_schedules.append([fp4_kernel_schedule, epi_nosmem_schedule]) CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) if isFp4: CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) def GenerateSM100_SparseTensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledSparseUniversal3x): # SM100 MMA with F4 inputs + block scale + sparse if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999) grouped = is_grouped(gemm_kind) if grouped: return # not support for grouped sparse kernels # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 64], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 64], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) def tile_schedulers(sfdtype): # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K. if sfdtype["type"] == DataType.void or grouped: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] math_instructions_1sm, math_instructions_2sm = generate_sparse_mxf4nvf4_math_instructions_sm100(instantiation_level) # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.e4m3, "vector_size": 32, "layout" : LayoutType.RowMajor} }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for layout in layouts: for data_type in data_types: if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue # E2M1 x E2M1, vector size 32, E8 # E2M1 x E2M1, vector size 16, UE4M3 isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 nvfp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1SmNvf4 fp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1SmMxf4 nvfp4_kernel_schedule = KernelScheduleType.SparseNvf4TmaWarpSpecialized1SmSm100 fp4_kernel_schedule = KernelScheduleType.SparseMxf4TmaWarpSpecialized1SmSm100 nvfp4_schedules = [[nvfp4_kernel_schedule, nvfp4_epi_schedule]] fp4_schedules = [[fp4_kernel_schedule, fp4_epi_schedule]] CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) if isFp4: CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.e4m3, "vector_size": 32, "layout" : LayoutType.RowMajor} }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for layout in layouts: for data_type in data_types: if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue # E2M1 x E2M1, vector size 32, E8 # E2M1 x E2M1, vector size 16, UE4M3 isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 nvfp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2SmNvf4 fp4_epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2SmMxf4 nvfp4_kernel_schedule = KernelScheduleType.SparseNvf4TmaWarpSpecialized2SmSm100 fp4_kernel_schedule = KernelScheduleType.SparseMxf4TmaWarpSpecialized2SmSm100 nvfp4_schedules = [[nvfp4_kernel_schedule, nvfp4_epi_schedule]] fp4_schedules = [[fp4_kernel_schedule, fp4_epi_schedule]] CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) if isFp4: CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledSparseUniversal3x): # SM100 MMA with mixed F4/F6/F8 inputs + block scale + sparse instantiation_level = manifest.get_instantiation_level(pruned_level=591, default_level=591, exhaustive_level=9999) grouped = is_grouped(gemm_kind) if grouped: return # not support for grouped sparse kernels # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 128], [LayoutType.RowMajor, 0]], ] math_instructions_1sm, math_instructions_2sm = generate_sparse_mxf8f6f4_math_instructions_sm100(instantiation_level) acc_types = [ DataType.f32 ] def tile_schedulers(sfdtype): # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K. if sfdtype["type"] == DataType.void or grouped: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 8 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} }, ] for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for layout in layouts: for data_type in data_types: if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue epi_schedule = EpilogueScheduleType.TmaWarpSpecialized1SmMxf8f6f4 kernel_schedule = KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized1SmSm100 schedules = [[kernel_schedule, epi_schedule]] CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 8 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, # void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} }, # none void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, # void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} }, # none void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} }, ] for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for layout in layouts: for data_type in data_types: if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2SmMxf8f6f4 kernel_schedule = KernelScheduleType.SparseMxf8f6f4TmaWarpSpecialized2SmSm100 schedules = [[kernel_schedule, epi_schedule]] CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, schedules , tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind ) def GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.BlockScaledUniversal3x): # SM100 MMA with F4 + block scale if not CudaToolkitVersionSatisfies(cuda_version, 13, 0): return grouped = is_grouped(gemm_kind) # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.ColumnMajor, 0]], ] instruction_sizes_1sm = [ [128, 128, 96], ] instruction_sizes_2sm = [ [256, 128, 96], [256, 192, 96], [256, 256, 96] ] ab_types = [ DataType.f4, DataType.e2m1, ] sf_types = [ DataType.ue4m3, DataType.ue8m0 ] acc_types = [ DataType.f32 ] # Accumulator is always 32 bits for block scaled MMA instructions def tile_schedulers(sfdtype): # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K. if grouped: return [TileSchedulerType.Default] if sfdtype["type"] == DataType.void: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] min_cc = 103 max_cc = 103 epi_type = DataType.f32 math_instructions_1sm = [] is_runtime_datatype = lambda runtime_datatype: runtime_datatype in (DataType.f4, DataType.f6, DataType.f8) for instr_size, a_type, b_type, sf_type, acc_type in product(instruction_sizes_1sm, ab_types, ab_types, sf_types, acc_types): is_runtime_datatype_a = is_runtime_datatype(a_type) is_runtime_datatype_b = is_runtime_datatype(b_type) # A/B datatypes should be both static or dynamic if (is_runtime_datatype_a != is_runtime_datatype_b): continue math_instructions_1sm.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.BlockScaledTensorOp, MathOperation.multiply_add, sf_type) ) math_instructions_2sm = [] for instr_size, a_type, b_type, sf_type, acc_type in product(instruction_sizes_2sm, ab_types, ab_types, sf_types, acc_types): is_runtime_datatype_a = is_runtime_datatype(a_type) is_runtime_datatype_b = is_runtime_datatype(b_type) # A/B datatypes should be both static or dynamic if (is_runtime_datatype_a != is_runtime_datatype_b): continue math_instructions_2sm.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.BlockScaledTensorOp, MathOperation.multiply_add, sf_type) ) cluster_shapes_1sm = [ [1,1,1], # [1,2,1], [2,1,1], # [1,4,1], [4,4,1], DynamicClusterShape ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], 768], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: for data_type in data_types: # Set alignment d based on Destination format. if DataTypeSize[data_type["c_type"]] == 0 : layout[2][1] = 256 // DataTypeSize[data_type["d_type"]] else: layout[2][1] = min(256 // DataTypeSize[data_type["d_type"]], 256 // DataTypeSize[data_type["c_type"]]) if data_type["sfd_type"]["type"] != DataType.void and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue # E2M1 x E2M1, vector size 32, E8 isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 epilogue_1sm_schedule = to_grouped_schedule(EpilogueScheduleType.NoSmemWarpSpecialized1Sm, grouped) nvfp4_schedule = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103, grouped), epilogue_1sm_schedule] nvfp4_schedule_disable_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103DisablePrefetch, grouped), epilogue_1sm_schedule] nvfp4_schedule_tma_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs16Sm103TmaPrefetch, grouped), epilogue_1sm_schedule] fp4_schedule = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103, grouped), epilogue_1sm_schedule] fp4_schedule_disable_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103DisablePrefetch, grouped), epilogue_1sm_schedule] fp4_schedule_tma_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized1SmVs32Sm103TmaPrefetch, grouped), epilogue_1sm_schedule] nvfp4_schedules = [nvfp4_schedule, nvfp4_schedule_disable_prefetch, nvfp4_schedule_tma_prefetch] fp4_schedules = [fp4_schedule, fp4_schedule_disable_prefetch, fp4_schedule_tma_prefetch] CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) if isFp4: CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) cluster_shapes_2sm = [ [2,1,1], # [2,2,1], # [2,4,1], [4,1,1], # [4,2,1], [4,4,1], DynamicClusterShape ] for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 8 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: for data_type in data_types: # Set alignment d based on Destination format. if DataTypeSize[data_type["c_type"]] == 0 : layout[2][1] = 256 // DataTypeSize[data_type["d_type"]] else: layout[2][1] = min(256 // DataTypeSize[data_type["d_type"]], 256 // DataTypeSize[data_type["c_type"]]) if data_type["sfd_type"]["type"] != DataType.void and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.RowMajor): data_type["sfd_type"]["layout"] = layout[2][0] # For FP4 output , the scalefactor layout is same layout as D layout. if (data_type["sfd_type"]["type"] != DataType.void) and (data_type["d_type"] == DataType.e2m1) and (layout[2][0] == LayoutType.ColumnMajor): continue # E2M1 x E2M1, vector size 32, E8 isFp4 = math_inst.element_scale_factor == DataType.ue8m0 and math_inst.element_a == DataType.e2m1 and math_inst.element_b == DataType.e2m1 epilogue_2sm_schedule = to_grouped_schedule(EpilogueScheduleType.NoSmemWarpSpecialized2Sm, grouped) nvfp4_schedule = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103, grouped), epilogue_2sm_schedule] nvfp4_schedule_disable_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103DisablePrefetch, grouped), epilogue_2sm_schedule] nvfp4_schedule_tma_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs16Sm103TmaPrefetch, grouped), epilogue_2sm_schedule] fp4_schedule = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103, grouped), epilogue_2sm_schedule] fp4_schedule_disable_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103DisablePrefetch, grouped), epilogue_2sm_schedule] fp4_schedule_tma_prefetch = [to_grouped_schedule(KernelScheduleType.MxNvf4UltraTmaWarpSpecialized2SmVs32Sm103TmaPrefetch, grouped), epilogue_2sm_schedule] nvfp4_schedules = [nvfp4_schedule, nvfp4_schedule_disable_prefetch, nvfp4_schedule_tma_prefetch] fp4_schedules = [fp4_schedule, fp4_schedule_disable_prefetch, fp4_schedule_tma_prefetch] CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, nvfp4_schedules, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) if isFp4: CreateGemmUniversal3xOperator(manifest, [layout], tile_descriptions, data_type, fp4_schedules, tile_schedulers=tile_schedulers(data_type["sfd_type"]), gemm_kind=gemm_kind) def GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return # layouts for ABC and their alignments. layouts = [ [[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.ColumnMajor, 16], [LayoutType.RowMajor, 0]], [[LayoutType.RowMajor, 16], [LayoutType.RowMajor, 16], [LayoutType.RowMajor, 0]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) epi_type = DataType.f32 math_instructions_1sm = [ MathInstruction( [64, 128, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [128, 128, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add)] cluster_shapes_1sm = [[1,2,1], [2,1,1], [1,1,1], [1,4,1], [4,4,1] , DynamicClusterShape ] if thor_sm in manifest.compute_capabilities_baseline : cluster_shapes_1sm = [[1,2,1], [2,1,1], [1,1,1], [1,4,1] , DynamicClusterShape ] tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 4 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_accumulator, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, [[KernelScheduleType.TmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_types_mixed = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_a, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types_mixed[0]["d_type"]] CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed, [[KernelScheduleType.TmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels math_instructions_2sm = [ MathInstruction( [128, 128, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [256, 128, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 32], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.TensorOp, MathOperation.multiply_add), ] cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1] , DynamicClusterShape ] if thor_sm in manifest.compute_capabilities_baseline : cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1] , DynamicClusterShape ] for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 4 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_accumulator, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_accumulator, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] if math_inst.instruction_shape[0] == 128: epi_schedule = EpilogueScheduleType.TmaWarpSpecialized2Sm else: epi_schedule = EpilogueScheduleType.ScheduleAuto CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types, [[KernelScheduleType.TmaWarpSpecialized2SmSm100, epi_schedule]], tile_schedulers=tile_schedulers) # for mixed precision kernels, also generate kernels that write output matrix in the A/B format # Avoid emitting two kernels if the accumulator type does not differ from the input type (e.g. F16 accumulation) if math_inst.element_a != math_inst.element_accumulator: data_types_mixed = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : math_inst.element_a, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : math_inst.element_a, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, }, ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_types_mixed, [[KernelScheduleType.TmaWarpSpecialized2SmSm100, epi_schedule]], tile_schedulers=tile_schedulers) def GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return # layouts for ABC and their alignments. layouts = [ # Alignment requirement will be over-write below [[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] kernel_data_types = [ # void_c { "a_type" : DataType.f32, "b_type" : DataType.f32, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, # none void_c { "a_type" : DataType.f32, "b_type" : DataType.f32, "c_type" : DataType.f32, "d_type" : DataType.f32, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, ] math_instructions_1sm = [ MathInstruction( [128, 128, 16], DataType.tf32, DataType.tf32, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 16], DataType.tf32, DataType.tf32, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] math_instructions_2sm = [ MathInstruction( [256, 128, 16], DataType.tf32, DataType.tf32, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 16], DataType.tf32, DataType.tf32, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 2 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_2sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 2 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) def GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return # layouts for ABC and their alignments. layouts = [ # Alignment requirement will be over-write below [[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] kernel_data_types = [ # void_c { "a_type" : DataType.f16, "b_type" : DataType.f16, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, # none void_c { "a_type" : DataType.f16, "b_type" : DataType.f16, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, ] math_instructions_1sm = [ MathInstruction( [128, 128, 32], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 32], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] math_instructions_2sm = [ MathInstruction( [256, 128, 32], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 32], DataType.f16, DataType.f16, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 2 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_2sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 2 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) def GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return # layouts for ABC and their alignments. layouts = [ # Alignment requirement will be over-write below [[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] kernel_data_types = [ # void_c { "a_type" : DataType.s8, "b_type" : DataType.s8, "c_type" : DataType.void, "d_type" : DataType.s8, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, # none void_c { "a_type" : DataType.s8, "b_type" : DataType.s8, "c_type" : DataType.s8, "d_type" : DataType.s8, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, ] math_instructions_1sm = [ MathInstruction( [128, 128, 64], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 64], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add)] math_instructions_2sm = [ MathInstruction( [256, 128, 64], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 64], DataType.s8, DataType.s8, DataType.s32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 2 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_2sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 2 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) def GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return # layouts for ABC and their alignments. layouts = [ # Alignment requirement will be over-write below [[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] kernel_data_types = [ # NOTE: a/b type in kernel will be overwrite below. #* void_c # f8_f8_f32_void_f16 { "a_type" : DataType.e4m3, "b_type" : DataType.e4m3, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, #* non-void_c # f8_f8_f32_f16_f8 { "a_type" : DataType.e4m3, "b_type" : DataType.e4m3, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : DataType.f32, "epi_type" : DataType.f32, }, ] math_instructions_1sm = [ # Runtime DType MathInstruction( [128, 128, 64], DataType.f8, DataType.f8, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 64], DataType.f8, DataType.f8, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] math_instructions_2sm = [ # Runtime DType MathInstruction( [256, 128, 64], DataType.f8, DataType.f8, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 64], DataType.f8, DataType.f8, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 2 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update input AB type kernel_data_type["a_type"] = math_inst.element_a kernel_data_type["b_type"] = math_inst.element_b # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_2sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 2 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) for kernel_data_type in kernel_data_types: # Update input AB type kernel_data_type["a_type"] = math_inst.element_a kernel_data_type["b_type"] = math_inst.element_b # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_copy = copy.deepcopy(layouts) for layout in layouts_copy: # alignment for a, 2 for sparsity layout[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) CreateSparseGemmUniversal3xOperator(manifest, layouts_copy, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) def GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return # layouts for ABC and their alignments. layouts = [ # Alignment requirement will be over-write below [[LayoutType.RowMajor, -1], [LayoutType.ColumnMajor, -1], [LayoutType.RowMajor, -1]], ] thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) tile_schedulers = [ TileSchedulerType.Default, TileSchedulerType.StreamK ] math_instructions_1sm = [ # Runtime Dtype MathInstruction( [128, 128, 64], DataType.f4, DataType.f4, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 64], DataType.f4, DataType.f4, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 128, 64], DataType.f6, DataType.f6, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [128, 256, 64], DataType.f6, DataType.f6, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] math_instructions_2sm = [ # Runtime DType MathInstruction( [256, 128, 64], DataType.f4, DataType.f4, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 64], DataType.f4, DataType.f4, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 128, 64], DataType.f6, DataType.f6, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), MathInstruction( [256, 256, 64], DataType.f6, DataType.f6, DataType.f32, OpcodeClass.SparseTensorOp, MathOperation.multiply_add), ] # 1xSM MMA kernels for math_inst in math_instructions_1sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_1sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_1sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else cluster_shape tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_1sm[0], math_inst.instruction_shape[1] * multiplier_1sm[1], math_inst.instruction_shape[2] * 2 * multiplier_1sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) kernel_data_types = [ # void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, }, # none void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, }, ] for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_filtered = [] for layout in layouts: layout_filter = copy.deepcopy(layout) # * A_K : Logical TileShape_K % 256 == 0 # * A_M : TileShape_M % 128 == 0 # * B_N : TileSize_N % 128 == 0 # * B_K : TileSize_K % 128 == 0 if ((layout_filter[0][0] == LayoutType.RowMajor and (math_inst.instruction_shape[2] * 2) % 256 == 0) or \ (layout_filter[0][0] == LayoutType.ColumnMajor and math_inst.instruction_shape[0] % 128 == 0)) and \ ((layout_filter[1][0] == LayoutType.RowMajor and math_inst.instruction_shape[1] % 128 == 0) or \ (layout_filter[1][0] == LayoutType.ColumnMajor and (math_inst.instruction_shape[0] * 2) % 128 == 0)): # alignment for a, 2 for sparsity layout_filter[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout_filter[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout_filter[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) layouts_filtered.append(layout_filter) CreateSparseGemmUniversal3xOperator(manifest, layouts_filtered, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized1SmSm100, EpilogueScheduleType.TmaWarpSpecialized1Sm]], tile_schedulers=tile_schedulers) # 2xSM MMA kernels for math_inst in math_instructions_2sm: tile_descriptions = [] for cluster_shape in sm100_cluster_shape_2sm: if thor_sm in manifest.compute_capabilities_baseline : if cluster_shape == [4,4,1] : continue multiplier_2sm = (1, 1, 1) if cluster_shape == DynamicClusterShape else (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) tile_descriptions.append( TileDescription([ math_inst.instruction_shape[0] * multiplier_2sm[0], math_inst.instruction_shape[1] * multiplier_2sm[1], math_inst.instruction_shape[2] * 2 * multiplier_2sm[2]], 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) kernel_data_types = [ # void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, }, # none void_c { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32, }, ] for kernel_data_type in kernel_data_types: # Update layout alignment # alignment for d might be different for each kernel_data_type layouts_filtered = [] for layout in layouts: layout_filter = copy.deepcopy(layout) # * A_K : Logical TileShape_K % 256 == 0 # * A_M : TileShape_M % 128 == 0 # * B_N : TileSize_N % 256 == 0 # * B_K : TileSize_K % 128 == 0 if ((layout_filter[0][0] == LayoutType.RowMajor and (math_inst.instruction_shape[2] * 2) % 256 == 0) or \ (layout_filter[0][0] == LayoutType.ColumnMajor and math_inst.instruction_shape[0] % 128 == 0)) and \ ((layout_filter[1][0] == LayoutType.RowMajor and math_inst.instruction_shape[1] % 256 == 0) or \ (layout_filter[1][0] == LayoutType.ColumnMajor and (math_inst.instruction_shape[0] * 2) % 128 == 0)): # alignment for a, 2 for sparsity layout_filter[0][1] = get_tma_alignment_elt(kernel_data_type["a_type"]) * ( 2 if layout[0][0] == LayoutType.RowMajor else 1) # alignment for b layout_filter[1][1] = get_tma_alignment_elt(kernel_data_type["b_type"]) # alignment for d layout_filter[2][1] = get_tma_alignment_elt(kernel_data_type["d_type"]) layouts_filtered.append(layout_filter) CreateSparseGemmUniversal3xOperator(manifest, layouts_filtered, tile_descriptions, [kernel_data_type], [[KernelScheduleType.SparseTmaWarpSpecialized2SmSm100, EpilogueScheduleType.TmaWarpSpecialized2Sm]], tile_schedulers=tile_schedulers) # Conv Utility functions def make_dims_and_alignments_triple(dim: int, bit_per_element_A: int, bit_per_element_B: int, bit_per_element_C: int): bit_alignment_required_by_tma = 128 return ((dim, bit_alignment_required_by_tma // bit_per_element_A), # A (dim, bit_alignment_required_by_tma // bit_per_element_B), # B (dim, bit_alignment_required_by_tma // bit_per_element_C)) # C def make_math_instruction_w_output(data_types: Tuple[DataType, DataType, DataType, DataType], instruction_shape: Tuple[int, int, int]) -> (MathInstruction, DataType): default_opcode = OpcodeClass.TensorOp default_math_op = MathOperation.multiply_add [A_data_type, B_data_type, Acc_data_type, Out_data_type] = data_types return (MathInstruction( instruction_shape, A_data_type, B_data_type, Acc_data_type, default_opcode, default_math_op ), Out_data_type) """ Generate CUTLASS 3 convolution kernel(s) for SM100. This is meant to be called from GenerateSM100. """ def GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version, log_indent_level: int = 0): log_debug_line('GenerateSM100_TensorOp_16b_UMMA_conv3x', log_indent_level) log_indent_level = log_indent_level + 1 if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) spatial_dims = [2, 3] conv_kinds = [ ConvKind.Fprop, ConvKind.Dgrad, ConvKind.Wgrad ] stages = 0 # zero means "deduce the number of stages automatically" data_types_and_instruction_shapes_1sm = [ # ((A,B,Acc,C/D), (InstM,InstN,InstK)) ((DataType.f16, DataType.f16, DataType.f16, DataType.f16), (64, 128, 16)), ((DataType.f16, DataType.f16, DataType.f16, DataType.f16), (128, 128, 16)), ((DataType.f16, DataType.f16, DataType.f16, DataType.f16), (128, 256, 16)), ((DataType.f16, DataType.f16, DataType.f32, DataType.f16), (64, 128, 16)), ((DataType.f16, DataType.f16, DataType.f32, DataType.f16), (128, 128, 16)), ((DataType.f16, DataType.f16, DataType.f32, DataType.f16), (128, 256, 16)), ((DataType.bf16, DataType.bf16, DataType.f32, DataType.bf16), (64, 128, 16)), ((DataType.bf16, DataType.bf16, DataType.f32, DataType.bf16), (128, 128, 16)), ((DataType.bf16, DataType.bf16, DataType.f32, DataType.bf16), (128, 256, 16)), ] math_instructions_w_output_1sm = map(lambda x: make_math_instruction_w_output(*x), data_types_and_instruction_shapes_1sm) cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1],[4,4,1]] if thor_sm in manifest.compute_capabilities_baseline : cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1]] # tile_descriptions is a 2-level list. # Each inner list is for each cluster shape. for math_inst, output_type in math_instructions_w_output_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: cluster_multiplier = cluster_shape # Unlike SM90, SM100 tile shape calculation includes cluster shape. tile_shape = [ math_inst.instruction_shape[0] * cluster_multiplier[0], math_inst.instruction_shape[1] * cluster_multiplier[1], math_inst.instruction_shape[2] * 4 * cluster_multiplier[2] ] warp_count = [4, 1, 1] tile_description = TileDescription( tile_shape, stages, warp_count, math_inst, min_cc, max_cc, cluster_shape) tile_descriptions.append(tile_description) # It's typical to get the data types from the math instruction. data_type = { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : output_type, "d_type" : output_type, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator } dims_and_alignments = [make_dims_and_alignments_triple(dim, DataTypeSize[data_type["a_type"]], DataTypeSize[data_type["b_type"]], DataTypeSize[data_type["d_type"]]) for dim in spatial_dims] # Schedules mainloop_schedule = KernelScheduleType.ImplicitTmaWarpSpecialized1SmSm100 epilogue_schedule = EpilogueScheduleType.ScheduleAuto schedule_pairs = [ (mainloop_schedule, epilogue_schedule) ] for conv_kind in conv_kinds: CreateConvOperator3x(manifest, dims_and_alignments = dims_and_alignments, tile_descriptions = tile_descriptions, data_types = data_type, schedule_pairs = schedule_pairs, conv_kind = conv_kind, log_indent_level = log_indent_level) data_types_and_instruction_shapes_2sm = [ # ((A,B,Acc,C/D), (InstM,InstN,InstK)) ((DataType.f16, DataType.f16, DataType.f16, DataType.f16), (128, 128, 16)), ((DataType.f16, DataType.f16, DataType.f16, DataType.f16), (128, 256, 16)), ((DataType.f16, DataType.f16, DataType.f16, DataType.f16), (256, 256, 16)), ((DataType.f16, DataType.f16, DataType.f32, DataType.f16), (128, 128, 16)), ((DataType.f16, DataType.f16, DataType.f32, DataType.f16), (128, 256, 16)), ((DataType.f16, DataType.f16, DataType.f32, DataType.f16), (256, 256, 16)), ((DataType.bf16, DataType.bf16, DataType.f32, DataType.bf16), (128, 128, 16)), ((DataType.bf16, DataType.bf16, DataType.f32, DataType.bf16), (128, 256, 16)), ((DataType.bf16, DataType.bf16, DataType.f32, DataType.bf16), (256, 256, 16)), ] math_instructions_w_output_2sm = map(lambda x: make_math_instruction_w_output(*x), data_types_and_instruction_shapes_2sm) cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1]] if thor_sm in manifest.compute_capabilities_baseline : cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]] for math_inst, output_type in math_instructions_w_output_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: cluster_multiplier = (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) # Unlike SM90, SM100 tile shape calculation includes cluster shape. tile_shape = [ math_inst.instruction_shape[0] * cluster_multiplier[0], math_inst.instruction_shape[1] * cluster_multiplier[1], math_inst.instruction_shape[2] * 4 * cluster_multiplier[2] ] warp_count = [4, 1, 1] tile_description = TileDescription( tile_shape, stages, warp_count, math_inst, min_cc, max_cc, cluster_shape) tile_descriptions.append(tile_description) # It's typical to get the data types from the math instruction. data_type = { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : output_type, "d_type" : output_type, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator } dims_and_alignments = [make_dims_and_alignments_triple(dim, DataTypeSize[data_type["a_type"]], DataTypeSize[data_type["b_type"]], DataTypeSize[data_type["d_type"]]) for dim in spatial_dims] # Schedules mainloop_schedule = KernelScheduleType.ImplicitTmaWarpSpecialized2SmSm100 epilogue_schedule = EpilogueScheduleType.ScheduleAuto schedule_pairs = [ (mainloop_schedule, epilogue_schedule) ] for conv_kind in conv_kinds: CreateConvOperator3x(manifest, dims_and_alignments = dims_and_alignments, tile_descriptions = tile_descriptions, data_types = data_type, schedule_pairs = schedule_pairs, conv_kind = conv_kind, log_indent_level = log_indent_level) def GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version, log_indent_level: int = 0): # Instantiate Fp8 Fprop kernels with e4m3 A/B, f32 Acc, e4m3/bf16/f16/f32 C/D log_debug_line('GenerateSM100_TensorOp_fp8_UMMA_conv3x', log_indent_level) log_indent_level = log_indent_level + 1 if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return thor_sm = ThorSMRenumbering(cuda_version) min_cc = 100 max_cc = 100 max_cc = max(max_cc, thor_sm) spatial_dims = [2, 3] stages = 0 # zero means "deduce the number of stages automatically" data_types_and_instruction_shapes_1sm = [ # ((A,B,Acc,C/D), (InstM,InstN,InstK)) ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.e4m3), (64, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.e4m3), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.e4m3), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f16), (64, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f16), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f16), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.bf16), (64, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.bf16), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.bf16), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f32), (64, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f32), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f32), (128, 256, 32)), ] math_instructions_w_output_1sm = map(lambda x: make_math_instruction_w_output(*x), data_types_and_instruction_shapes_1sm) cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1],[4,4,1]] if thor_sm in manifest.compute_capabilities_baseline : cluster_shapes_1sm = [[1,1,1], [1,2,1], [1,4,1]] for math_inst, output_type in math_instructions_w_output_1sm: tile_descriptions = [] for cluster_shape in cluster_shapes_1sm: cluster_multiplier = cluster_shape # Unlike SM90, SM100 tile shape calculation includes cluster shape. tile_shape = [ math_inst.instruction_shape[0] * cluster_multiplier[0], math_inst.instruction_shape[1] * cluster_multiplier[1], math_inst.instruction_shape[2] * 4 * cluster_multiplier[2] ] warp_count = [4, 1, 1] tile_description = TileDescription( tile_shape, stages, warp_count, math_inst, min_cc, max_cc, cluster_shape) tile_descriptions.append(tile_description) data_type = { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : output_type, "d_type" : output_type, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator } dims_and_alignments = [make_dims_and_alignments_triple(dim, DataTypeSize[data_type["a_type"]], DataTypeSize[data_type["b_type"]], DataTypeSize[data_type["d_type"]]) for dim in spatial_dims] # Schedules mainloop_schedule = KernelScheduleType.ImplicitTmaWarpSpecialized1SmSm100 epilogue_schedule = EpilogueScheduleType.ScheduleAuto schedule_pairs = [ (mainloop_schedule, epilogue_schedule) ] CreateConvOperator3x(manifest, dims_and_alignments = dims_and_alignments, tile_descriptions = tile_descriptions, data_types = data_type, schedule_pairs = schedule_pairs, conv_kind = ConvKind.Fprop, log_indent_level = log_indent_level) data_types_and_instruction_shapes_2sm = [ # ((A,B,Acc,C/D), (InstM,InstN,InstK)) ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.e4m3), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.e4m3), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.e4m3), (256, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f16), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f16), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f16), (256, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.bf16), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.bf16), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.bf16), (256, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f32), (128, 128, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f32), (128, 256, 32)), ((DataType.e4m3, DataType.e4m3, DataType.f32, DataType.f32), (256, 256, 32)), ] math_instructions_w_output_2sm = map(lambda x: make_math_instruction_w_output(*x), data_types_and_instruction_shapes_2sm) cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1], [4,4,1]] if thor_sm in manifest.compute_capabilities_baseline : cluster_shapes_2sm = [[2,1,1], [2,2,1], [2,4,1], [4,1,1], [4,2,1]] for math_inst, output_type in math_instructions_w_output_2sm: tile_descriptions = [] for cluster_shape in cluster_shapes_2sm: cluster_multiplier = (cluster_shape[0] // 2, cluster_shape[1], cluster_shape[2]) # Unlike SM90, SM100 tile shape calculation includes cluster shape. tile_shape = [ math_inst.instruction_shape[0] * cluster_multiplier[0], math_inst.instruction_shape[1] * cluster_multiplier[1], math_inst.instruction_shape[2] * 4 * cluster_multiplier[2] ] warp_count = [4, 1, 1] tile_description = TileDescription( tile_shape, stages, warp_count, math_inst, min_cc, max_cc, cluster_shape) tile_descriptions.append(tile_description) data_type = { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : output_type, "d_type" : output_type, "acc_type" : math_inst.element_accumulator, "epi_type" : math_inst.element_accumulator } dims_and_alignments = [make_dims_and_alignments_triple(dim, DataTypeSize[data_type["a_type"]], DataTypeSize[data_type["b_type"]], DataTypeSize[data_type["d_type"]]) for dim in spatial_dims] # Schedules mainloop_schedule = KernelScheduleType.ImplicitTmaWarpSpecialized2SmSm100 epilogue_schedule = EpilogueScheduleType.ScheduleAuto schedule_pairs = [ (mainloop_schedule, epilogue_schedule) ] CreateConvOperator3x(manifest, dims_and_alignments = dims_and_alignments, tile_descriptions = tile_descriptions, data_types = data_type, schedule_pairs = schedule_pairs, conv_kind = ConvKind.Fprop, log_indent_level = log_indent_level) def GenerateSM120_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version): # SM120 MMA with mixed F4/F6/F8 inputs + block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return layouts = [ [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]] ] instruction_sizes = [ [16, 8, 32] ] tile_sizes = [ [128, 128, 128] ] cluster_shape = [1,1,1] ab_types = [ DataType.e2m1, DataType.e2m3, DataType.e3m2, DataType.e5m2, DataType.e4m3, ] acc_types = [ DataType.f32 ] def is_pingpong(kernel_schedule): if kernel_schedule == KernelScheduleType.Mxf8f6f4TmaWarpSpecializedPingpongSm120: return True else: return False def tile_schedulers(sfdtype, kernel_schedule): # Pingpong kernel schedule doesn't support stream-K. # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K if is_pingpong(kernel_schedule): return [TileSchedulerType.Default] elif sfdtype["type"] == DataType.void: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] min_cc = 120 max_cc = 121 epi_type = DataType.f32 math_instructions = [] kernel_schedules = [ KernelScheduleType.Mxf8f6f4TmaWarpSpecializedCooperativeSm120, KernelScheduleType.Mxf8f6f4TmaWarpSpecializedPingpongSm120 ] for instr_size, a_type, b_type, acc_type in product(instruction_sizes, ab_types, ab_types, acc_types): math_instructions.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.BlockScaledTensorOp, MathOperation.multiply_add, DataType.ue8m0) ) for math_inst in math_instructions: tile_descriptions = [] for tile_size in tile_sizes: tile_descriptions.append( TileDescription(tile_size, 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e3m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type, kernel_schedule in product(data_types, kernel_schedules): CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, EpilogueScheduleType.ScheduleAuto]], tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule), gemm_kind = GemmKind.BlockScaledUniversal3x ) def GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version): # SM120 MMA with with F4 + block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 32], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]] ] instruction_sizes = [ [16, 8, 64] ] tile_sizes_cooperative = [ [128, 128, 128], [128, 128, 256], [256, 128, 128] ] tile_sizes_pingpong = [ [128, 128, 128], [128, 128, 256] ] cluster_shape = [1,1,1] ab_types = [ DataType.e2m1 ] sf_types = [ DataType.ue4m3, DataType.ue8m0 ] acc_types = [ DataType.f32 ] def is_pingpong(kernel_schedule): if kernel_schedule == KernelScheduleType.Nvf4TmaWarpSpecializedPingpongSm120 or \ kernel_schedule == KernelScheduleType.Mxf4TmaWarpSpecializedPingpongSm120: return True else: return False def is_nvf4(kernel_schedule): if kernel_schedule == KernelScheduleType.Nvf4TmaWarpSpecializedCooperativeSm120 or \ kernel_schedule == KernelScheduleType.Nvf4TmaWarpSpecializedPingpongSm120: return True else: return False def tile_schedulers(sfdtype, kernel_schedule): # Pingpong kernel schedule doesn't support stream-K. # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K if is_pingpong(kernel_schedule): return [TileSchedulerType.Default] elif sfdtype["type"] == DataType.void: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] min_cc = 120 max_cc = 121 epi_type = DataType.f32 math_instructions = [] kernel_schedules = [ KernelScheduleType.Nvf4TmaWarpSpecializedCooperativeSm120, KernelScheduleType.Nvf4TmaWarpSpecializedPingpongSm120, KernelScheduleType.Mxf4TmaWarpSpecializedCooperativeSm120, KernelScheduleType.Mxf4TmaWarpSpecializedPingpongSm120 ] for instr_size, a_type, b_type, acc_type, sf_type in product(instruction_sizes, ab_types, ab_types, acc_types, sf_types): math_instructions.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.BlockScaledTensorOp, MathOperation.multiply_add, sf_type) ) for math_inst in math_instructions: for kernel_schedule in kernel_schedules: tile_descriptions = [] tile_sizes = tile_sizes_pingpong if is_pingpong(kernel_schedule) else tile_sizes_cooperative for tile_size in tile_sizes: # nvf4 kernel only supports ue4m3 SF # mxf4 kernel only supports ue8m0 SF if (math_inst.element_scale_factor == DataType.ue4m3 and is_nvf4(kernel_schedule)) or \ (math_inst.element_scale_factor == DataType.ue8m0 and not is_nvf4(kernel_schedule)): tile_descriptions.append( TileDescription(tile_size, 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 16, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type in data_types: CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, EpilogueScheduleType.ScheduleAuto]], tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule), gemm_kind = GemmKind.BlockScaledUniversal3x ) def GenerateSM120_Sparse_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version): # SM120 MMA with mixed F4/F6/F8 inputs + block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return layouts = [ [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]] ] instruction_sizes = [ [16, 8, 64] ] tile_sizes = [ [128, 64, 256], [128, 128, 256] ] cluster_shape = [1,1,1] ab_types = [ # DataType.e2m1, # DataType.e2m3, # DataType.e3m2, # DataType.e5m2, DataType.e4m3, ] acc_types = [ DataType.f32 ] def tile_schedulers(sfdtype, kernel_schedule): # Pingpong kernel schedule doesn't support stream-K. # Only use the stream-K scheduler for non-void SFD to limit kernel count. When SFD is void, # the epilogue is the traditional linear combination, for which we already have tests with stream-K if sfdtype["type"] == DataType.void: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] min_cc = 120 max_cc = 121 epi_type = DataType.f32 math_instructions = [] kernel_schedules = [ KernelScheduleType.SparseMxf8f6f4TmaWarpSpecializedAcc2x4Sm120 ] for instr_size, a_type, b_type, acc_type in product(instruction_sizes, ab_types, ab_types, acc_types): math_instructions.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.BlockScaledSparseTensorOp, MathOperation.multiply_add, DataType.ue8m0) ) for math_inst in math_instructions: tile_descriptions = [] for tile_size in tile_sizes: tile_descriptions.append( TileDescription(tile_size, 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 64, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type, kernel_schedule in product(data_types, kernel_schedules): CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, EpilogueScheduleType.SparseTmaWarpSpecializedCooperativeSm120]], tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule), gemm_kind = GemmKind.BlockScaledSparseUniversal3x ) def GenerateSM120_Sparse_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version): # SM120 MMA with with F4 + block scale if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return # layouts for ABC and their alignments. layouts = [ [[LayoutType.RowMajor, 64], [LayoutType.ColumnMajor, 32], [LayoutType.RowMajor, 0]] ] instruction_sizes = [ [16, 8, 64] ] tile_sizes = [ [128, 64, 256], [128, 128, 256] ] cluster_shape = [1,1,1] ab_types = [ DataType.e2m1 ] sf_types = [ DataType.ue4m3, # DataType.ue8m0 ] acc_types = [ DataType.f32 ] def is_nvf4(kernel_schedule): if kernel_schedule == KernelScheduleType.SparseNvf4TmaWarpSpecializedSm120: return True else: return False def tile_schedulers(sfdtype, kernel_schedule): # When SFD is void, the epilogue is the traditional linear combination, for which we already have tests with stream-K if sfdtype["type"] == DataType.void: return [TileSchedulerType.Default] else: return [TileSchedulerType.Default, TileSchedulerType.StreamK] min_cc = 120 max_cc = 121 epi_type = DataType.f32 math_instructions = [] kernel_schedules = [ KernelScheduleType.SparseNvf4TmaWarpSpecializedSm120, KernelScheduleType.SparseMxf4TmaWarpSpecializedSm120, ] for instr_size, a_type, b_type, acc_type, sf_type in product(instruction_sizes, ab_types, ab_types, acc_types, sf_types): math_instructions.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.BlockScaledSparseTensorOp, MathOperation.multiply_add, sf_type) ) for math_inst in math_instructions: for kernel_schedule in kernel_schedules: tile_descriptions = [] for tile_size in tile_sizes: # nvf4 kernel only supports ue4m3 SF # mxf4 kernel only supports ue8m0 SF if (math_inst.element_scale_factor == DataType.ue4m3 and is_nvf4(kernel_schedule)) or \ (math_inst.element_scale_factor == DataType.ue8m0 and not is_nvf4(kernel_schedule)): tile_descriptions.append( TileDescription(tile_size, 0, [4, 1, 1], math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.void, "vector_size": None, "layout" : None} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e2m1, "acc_type" : math_inst.element_accumulator, "epi_type" : epi_type, "sf_type" : math_inst.element_scale_factor, "sfd_type" : {"type": DataType.ue8m0, "vector_size": 32, "layout" : LayoutType.RowMajor} } ] # Set alignment d based on Destination format. for layout in layouts: layout[2][1] = 128 // DataTypeSize[data_types[0]["d_type"]] for data_type in data_types: CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, EpilogueScheduleType.SparseTmaWarpSpecializedCooperativeSm120]], tile_schedulers = tile_schedulers(data_type["sfd_type"], kernel_schedule), gemm_kind = GemmKind.BlockScaledSparseUniversal3x ) def GenerateSM120_Sparse_TensorOp_gemm(manifest, cuda_version): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return layouts = [ [[LayoutType.RowMajor, 256], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 0]] ] tile_sizes = [ [128, 128, 256] ] cluster_shape = [1,1,1] warp_count = [4, 2, 1] acc_types = [ DataType.f32 ] instruction_sizes_mxf8f6f4 = [ [16, 8, 64] ] ab_types_mxf8f6f4 = [ DataType.e2m1, #DataType.e2m3, DataType.e3m2, #DataType.e5m2, DataType.e4m3, ] def tile_schedulers(kernel_schedule): return [TileSchedulerType.Default] min_cc = 120 max_cc = 121 kernel_schedules = [ KernelScheduleType.F8f6f4SparseTmaWarpSpecializedCooperativeSm120, ] math_instructions_mxf8f6f4 = [] for instr_size, a_type, b_type, acc_type in product(instruction_sizes_mxf8f6f4, ab_types_mxf8f6f4, ab_types_mxf8f6f4, acc_types): math_instructions_mxf8f6f4.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.SparseTensorOp, MathOperation.multiply_add) ) # Create gemm operator for mxf8f6f4 for math_inst in math_instructions_mxf8f6f4: tile_descriptions_mxf8f6f4 = [] for tile_size in tile_sizes: tile_descriptions_mxf8f6f4.append( TileDescription(tile_size, 0, warp_count, math_inst, min_cc, max_cc, cluster_shape)) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f32, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.e5m2, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.e4m3, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 } ] for data_type, kernel_schedule in product(data_types, kernel_schedules): # Set alignment d based on Destination format for layout in layouts: layout[2][1] = int(128 // DataTypeSize[data_type["d_type"]]) # Create gemm operator CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions_mxf8f6f4, data_type, [[kernel_schedule, EpilogueScheduleType.ScheduleAuto]], tile_schedulers = tile_schedulers(kernel_schedule), gemm_kind = GemmKind.SparseUniversal3x) def GenerateSM120_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.BlockwiseUniversal3x): if not CudaToolkitVersionSatisfies(cuda_version, 12, 8): return layouts = [ [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.RowMajor, 16]], [[LayoutType.RowMajor, 128], [LayoutType.ColumnMajor, 128], [LayoutType.ColumnMajor, 16]] ] cooperative_tile_sizes = [ [128, 128, 128] ] pingpong_tile_sizes = [ [64, 128, 128] ] def get_tile_sizes(kernel_scheduler): if kernel_scheduler == KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120: return pingpong_tile_sizes return cooperative_tile_sizes def get_warp_count(kernel_scheduler): if kernel_scheduler == KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120: return [2, 2, 1] return [4, 2, 1] def get_sf_sizes(tile_size): sf_sizes = [] for vec_m in [1, 128]: if tile_size[0] % vec_m > 0: continue for vec_n in [1, 128]: if tile_size[1] % vec_m > 0: continue sf_sizes.append( [vec_m, vec_n, 128] ) return sf_sizes cluster_shape = [1,1,1] acc_types = [ DataType.f32 ] instruction_sizes = [ [16, 8, 32] ] def tile_schedulers(kernel_schedule): return [TileSchedulerType.Default] min_cc = 120 max_cc = 121 kernel_schedulers = [ KernelScheduleType.BlockwiseTmaWarpSpecializedCooperativeSm120, KernelScheduleType.BlockwiseTmaWarpSpecializedPingpongSm120 ] ab_types = [ [DataType.e4m3, DataType.e4m3], [DataType.e4m3, DataType.e5m2] ] math_instructions = [] for instr_size, ab_type, acc_type in product(instruction_sizes, ab_types, acc_types): a_type, b_type = ab_type math_instructions.append( MathInstruction( instr_size, a_type, b_type, acc_type, OpcodeClass.TensorOp, MathOperation.multiply_add) ) # Create gemm operator for mxf8f6f4 for kernel_schedule in kernel_schedulers: tile_sizes = get_tile_sizes(kernel_schedule) warp_count = get_warp_count(kernel_schedule) for math_inst in math_instructions: tile_descriptions = [] for tile_size in tile_sizes: sf_sizes = get_sf_sizes(tile_size) for sf_size in sf_sizes: tile_descriptions.append( TileDescription(tile_size, 0, warp_count, math_inst, min_cc, max_cc, cluster_shape, explicit_vector_sizes=sf_size) ) data_types = [ { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.f16, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.bf16, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.f16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 }, { "a_type" : math_inst.element_a, "b_type" : math_inst.element_b, "c_type" : DataType.void, "d_type" : DataType.bf16, "acc_type" : math_inst.element_accumulator, "epi_type" : DataType.f32 } ] for data_type in data_types: # Set alignment d based on Destination format for layout in layouts: layout[2][1] = int(128 // DataTypeSize[data_type["d_type"]]) # Create gemm operator CreateGemmUniversal3xOperator(manifest, layouts, tile_descriptions, data_type, [[kernel_schedule, EpilogueScheduleType.ScheduleAuto]], tile_schedulers = tile_schedulers(kernel_schedule), gemm_kind = gemm_kind) def GenerateSM100(manifest, cuda_version): arch_family_cc = ['100f', '101f', '103a'] if CudaToolkitVersionSatisfies(cuda_version, 13, 0): for old_cc, new_cc in [('101f', '110f')]: arch_family_cc = [cc.replace(old_cc, new_cc) for cc in arch_family_cc] # # Dense Gemm # GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version) GenerateSM100_TensorOp_32b_UMMA_gemm(manifest, cuda_version) if not bool(set(manifest.compute_capabilities_feature_set).intersection(arch_family_cc)): GenerateSM100_TensorOp_int8_UMMA_gemm(manifest, cuda_version) GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version) # grouped GEMM GenerateSM100_TensorOp_fp8_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x) GenerateSM100_TensorOp_16b_UMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x) # StreamK is included in regular generation GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version) # Blockwise kernels GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version) GenerateSM100_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockwiseUniversal3x) # # Sparse Gemm # GenerateSM100_SparseTensorOp_32b_UMMA_gemm(manifest, cuda_version) GenerateSM100_SparseTensorOp_16b_UMMA_gemm(manifest, cuda_version) if not bool(set(manifest.compute_capabilities_feature_set).intersection(arch_family_cc)): GenerateSM100_SparseTensorOp_int8_UMMA_gemm(manifest, cuda_version) GenerateSM100_SparseTensorOp_fp8_UMMA_gemm(manifest, cuda_version) GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm(manifest, cuda_version) # # Block Scaled Gemm # GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version) GenerateSM100_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockScaledUniversal3x) GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version) GenerateSM100_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockScaledUniversal3x) GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version) GenerateSM103_TensorOp_fp4_ultra_UMMA_gemm_with_block_scaled(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockScaledUniversal3x) # # Block Scaled Sparse Gemm # GenerateSM100_SparseTensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version) GenerateSM100_SparseTensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version) # # Conv # GenerateSM100_TensorOp_16b_UMMA_conv3x(manifest, cuda_version) GenerateSM100_TensorOp_fp8_UMMA_conv3x(manifest, cuda_version) def GenerateSM120(manifest, cuda_version): # StreamK is included in regular generation # # # Dense Block Scaled Gemm # GenerateSM120_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version) GenerateSM120_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version) # # Sparse Block Scaled Gemm # GenerateSM120_Sparse_TensorOp_mixed_8bits_UMMA_gemm_with_block_scaled(manifest, cuda_version) GenerateSM120_Sparse_TensorOp_fp4_UMMA_gemm_with_block_scaled(manifest, cuda_version) # # Sparse Gemm # GenerateSM120_Sparse_TensorOp_gemm(manifest, cuda_version) GenerateSM120_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version) GenerateSM120_TensorOp_fp8_UMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockwiseUniversal3x) ################################################################################################### def GenerateSM90_Conv3x(manifest, cuda_version, log_indent_level: int = 0): """ Generate CUTLASS 3 convolution kernel(s) for SM90. This is meant to be called from GenerateSM90. """ log_debug_line('GenerateSM90_Conv3x', log_indent_level) log_indent_level = log_indent_level + 1 if not CudaToolkitVersionSatisfies(cuda_version, 12, 0): return min_cc = 90 max_cc = 90 spatial_dims = (2, 3) # MMA shapes (MMA_M, MMA_N, MMA_K): # # Different hardware MMA instructions may have different MMA shapes. # This function may generate kernels with different MMA shapes for # different data types, either because the hardware only supports # certain shapes for certain types, or for performance reasons # (CUTLASS doesn't need to generate all valid kernels for the # profiler library, just the best-performing ones). # # The kernel names refer to tile shapes (TILE_M, TILE_N, TILE_K) # instead of MMA shapes. For SM >= 90 kernels, TILE_K = 4 * MMA_K, # where 4, the "number of MMA instructions per tile," is determined # through some combination of modeling and experiment. # # For performance on sm90, generally CUTLASS generates 64x128 # instead of 128x64. mma_64x64x16 = ( 64, 64, 16) mma_64x64x8 = ( 64, 64, 8) num_mma_per_tile = 4 # Cluster shapes (1, 1, 1) and (2, 2, 1) are valid, # but not included, because they tend not to perform as well. cluster_shapes = ( (2, 1, 1), (1, 2, 1), ) fp16 = DataType.f16 bf16 = DataType.bf16 fp32 = DataType.f32 s8 = DataType.s8 s32 = DataType.s32 # When generating kernels, the usual way is to specify 4 types, # (A, B, Acc, C/D). Tests instead have 5 types, # (ElementAct, ElementFlt, ElementOut, ElementAcc, ElementCompute), # where ElementCompute is also called 'epi_type', # and corresponds to the type of epilogue activations. # This script maps tests' 5 types to 4 types # by making ElementCompute the same as ElementOut. fp16_fp32_fp16_fp32 = { 'a_type': fp16, # ElementAct(ivation) 'b_type': fp16, # ElementF(i)lt(er) 'c_type': fp32, # ElementAcc 'd_type': fp32, # ElementOut (used only by CollectiveEpilogue) 'acc_type': fp16, # ElementAcc 'epi_type': fp32, # ElementCompute (used only by CollectiveEpilogue) 'alignment_A': 8, # tma alignment elements of A 'alignment_B': 8, # tma alignment elements of B 'alignment_C': 4, # tma alignment elements of C } fp16_fp32_fp32_fp32 = { 'a_type': fp16, 'b_type': fp16, 'c_type': fp32, 'd_type': fp32, 'acc_type': fp32, 'epi_type': fp32, 'alignment_A': 8, 'alignment_B': 8, 'alignment_C': 4, } fp32_fp32_fp32_fp32 = { 'a_type': fp32, 'b_type': fp32, 'c_type': fp32, 'd_type': fp32, 'acc_type': fp32, 'epi_type': fp32, 'alignment_A': 4, 'alignment_B': 4, 'alignment_C': 4, } s8_s32_s32_s32 = { 'a_type': s8, 'b_type': s8, 'c_type': s32, 'd_type': s32, 'acc_type': s32, 'epi_type': s32, 'alignment_A': 16, 'alignment_B': 16, 'alignment_C': 4, } # Other NVIDIA libraries may have the habit of specifying data types like this. bf16bf16_bf16f32_f32 = { 'a_type': bf16, 'b_type': bf16, 'c_type': fp32, 'd_type': fp32, 'acc_type': fp32, 'epi_type': fp32, 'alignment_A': 8, 'alignment_B': 8, 'alignment_C': 4, } f16f16_f16f16_f16 = { 'a_type': fp16, 'b_type': fp16, 'c_type': fp16, 'd_type': fp16, 'acc_type': fp16, 'epi_type': fp16, 'alignment_A': 8, 'alignment_B': 8, 'alignment_C': 8, } f16f16_f16f32_f32 = { 'a_type': fp16, 'b_type': fp16, 'c_type': fp16, 'd_type': fp16, 'acc_type': fp32, 'epi_type': fp32, 'alignment_A': 8, 'alignment_B': 8, 'alignment_C': 8, } f32f32_tf32f32_f32 = fp32_fp32_fp32_fp32 i8i8_i8i32_f32 = { 'a_type': s8, 'b_type': s8, 'c_type': s32, 'd_type': s32, 'acc_type': s32, 'epi_type': s32, 'alignment_A': 16, 'alignment_B': 16, 'alignment_C': 4, } # Each element in the outermost iterable is one combination of # # (ConvKind, spatial_dimension, data_types, byte_alignments, mma_sizes, cluster_sizes) # # for which to generate a kernel. spatial_dimension is the spatial # dimension of the convolution: either 1, 2, or 3. byte_alignments # is a triple of required minimum byte alignments for A, B, and C. # # Note that itertools functions produce a single-pass generator. # The code doesn't need a multipass iterable, but if one did, one # could call `tuple` or `list` on the generator. # # While this happens to use the same cluster sizes for each element, # the code doesn't require that. Different convolution kinds, data # types, or mma sizes might have different optimal cluster sizes. combinations_of_parameters = chain( # The following are all the kernels exercised in the unit tests. # Please try to keep in sync with the unit tests. product( ( ConvKind.Fprop, ), spatial_dims, ( fp16_fp32_fp16_fp32, fp16_fp32_fp32_fp32, s8_s32_s32_s32, ), ( mma_64x64x16, ), cluster_shapes ), product( ( ConvKind.Fprop, ), spatial_dims, ( fp32_fp32_fp32_fp32, ), ( mma_64x64x8, ), cluster_shapes ), product( ( ConvKind.Dgrad, ConvKind.Wgrad ), spatial_dims, ( fp16_fp32_fp16_fp32, fp16_fp32_fp32_fp32, ), ( mma_64x64x16, ), cluster_shapes ), # Kernels not necessarily in the unit tests, but used elsewhere # and thus useful to have generated for profiling. They may # duplicate kernels above. All of them are 2-D. In general, # CUTLASS prefers 64 x 128 to 128 x 64 on sm90, even if the # hardware permits 128 x 64. ( # Fprop # # bf16bf16_bf16f32_f32 # # cluster shape (2, 1, 1) # (ConvKind.Fprop, 2, bf16bf16_bf16f32_f32, (128, 256, 8), (2, 1, 1)), (ConvKind.Fprop, 2, bf16bf16_bf16f32_f32, (128, 256, 16), (2, 1, 1)), (ConvKind.Fprop, 2, bf16bf16_bf16f32_f32, (256, 128, 8), (2, 1, 1)), (ConvKind.Fprop, 2, bf16bf16_bf16f32_f32, (256, 128, 16), (2, 1, 1)), # # f16f16_f16f16_f16 # # cluster shape (1, 1, 1) # (ConvKind.Fprop, 2, f16f16_f16f16_f16, ( 64, 64, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, ( 64, 64, 16), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, ( 64, 128, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, ( 64, 128, 16), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, ( 64, 256, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, ( 64, 256, 16), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (128, 128, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (128, 128, 16), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (128, 256, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (128, 256, 16), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (256, 64, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (256, 64, 16), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (256, 128, 8), (1, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f16_f16, (256, 128, 16), (1, 1, 1)), # # f16f16_f16f32_f32 # # cluster shape (2, 1, 1) # (ConvKind.Fprop, 2, f16f16_f16f32_f32, (128, 192, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (128, 192, 16), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (128, 256, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (128, 256, 16), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (256, 96, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (256, 96, 16), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (256, 128, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f16f16_f16f32_f32, (256, 128, 16), (2, 1, 1)), # # f32f32_tf32f32_f32 # # cluster shape (2, 1, 1) # (ConvKind.Fprop, 2, f32f32_tf32f32_f32, (128, 192, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f32f32_tf32f32_f32, (128, 256, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f32f32_tf32f32_f32, (256, 128, 8), (2, 1, 1)), (ConvKind.Fprop, 2, f32f32_tf32f32_f32, (256, 96, 8), (2, 1, 1)), # # i8i8_i8i32_f32 # # cluster shape (2, 1, 1) # (ConvKind.Fprop, 2, i8i8_i8i32_f32, (128, 256, 16), (2, 1, 1)), (ConvKind.Fprop, 2, i8i8_i8i32_f32, (128, 256, 32), (2, 1, 1)), (ConvKind.Fprop, 2, i8i8_i8i32_f32, (256, 128, 16), (2, 1, 1)), (ConvKind.Fprop, 2, i8i8_i8i32_f32, (256, 128, 32), (2, 1, 1)), # # Dgrad # # bf16bf16_bf16f32_f32 # # cluster shape (2, 1, 1) # (ConvKind.Dgrad, 2, bf16bf16_bf16f32_f32, (128, 256, 8), (2, 1, 1)), (ConvKind.Dgrad, 2, bf16bf16_bf16f32_f32, (128, 256, 16), (2, 1, 1)), (ConvKind.Dgrad, 2, bf16bf16_bf16f32_f32, (256, 128, 8), (2, 1, 1)), (ConvKind.Dgrad, 2, bf16bf16_bf16f32_f32, (256, 128, 16), (2, 1, 1)), # # f16f16_f16f16_f16 # # cluster shape (1, 1, 1) # (ConvKind.Dgrad, 2, f16f16_f16f16_f16, ( 64, 64, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, ( 64, 64, 16), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, ( 64, 128, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, ( 64, 128, 16), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, ( 64, 256, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, ( 64, 256, 16), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (128, 128, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (128, 128, 16), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (128, 256, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (128, 256, 16), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (256, 64, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (256, 64, 16), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (256, 128, 8), (1, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f16_f16, (256, 128, 16), (1, 1, 1)), # # f16f16_f16f32_f32 # # cluster shape (2, 1, 1) # (ConvKind.Dgrad, 2, f16f16_f16f32_f32, (128, 256, 8), (2, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f32_f32, (128, 256, 16), (2, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f32_f32, (256, 128, 8), (2, 1, 1)), (ConvKind.Dgrad, 2, f16f16_f16f32_f32, (256, 128, 16), (2, 1, 1)), ), ) # SM >= 90 kernels don't actually use warp_count, but the # TileDescription class needs it. The 4 in the default # warp_count has nothing to do with num_mma_per_tile. warp_count = [4, 1, 1] stages = 0 # zero means "deduce the number of stages automatically" mainloop_schedule = KernelScheduleType.ImplicitTmaWarpSpecializedSm90 epilogue_schedule = EpilogueScheduleType.TmaWarpSpecialized schedule_pairs = ( (mainloop_schedule, epilogue_schedule), ) tile_schedulers = ( TileSchedulerType.Default, # -> void ) def make_math_instruction(data_types: Dict[str, DataType], mma_shape: Tuple[int, int, int]) -> MathInstruction: default_opcode = OpcodeClass.TensorOp default_math_op = MathOperation.multiply_add return MathInstruction( mma_shape, data_types['a_type'], data_types['b_type'], data_types['c_type'], default_opcode, default_math_op ) for (conv_kind, spatial_dim, data_types, mma_shape, cluster_shape) in combinations_of_parameters: math_inst = make_math_instruction(data_types, mma_shape) tile_shape = (mma_shape[0], mma_shape[1], num_mma_per_tile * mma_shape[2]) tile_description = TileDescription(tile_shape, stages, warp_count, math_inst, min_cc, max_cc, cluster_shape) assert(isinstance(spatial_dim, int)) dims_and_alignments = ( ( (spatial_dim, data_types['alignment_A']), (spatial_dim, data_types['alignment_B']), (spatial_dim, data_types['alignment_C']), ), ) CreateConvOperator3x(manifest, dims_and_alignments = dims_and_alignments, tile_descriptions = [tile_description], data_types = data_types, schedule_pairs = schedule_pairs, tile_schedulers = tile_schedulers, conv_kind = conv_kind, log_indent_level = log_indent_level) def GenerateSM90(manifest, cuda_version): GenerateSM90_TensorOp_16b_WGMMA_gemm(manifest, cuda_version) GenerateSM90_TensorOp_16b_WGMMA_alignx_gemm(manifest, cuda_version) GenerateSM90_TensorOp_tf32_WGMMA_gemm(manifest, cuda_version) GenerateSM90_TensorOp_tf32_WGMMA_alignx_gemm(manifest, cuda_version) GenerateSM90_TensorOp_int8_WGMMA_gemm(manifest, cuda_version) GenerateSM90_TensorOp_int8_WGMMA_alignx_gemm(manifest, cuda_version) GenerateSM90_TensorOp_fp8_WGMMA_gemm(manifest, cuda_version) GenerateSM90_TensorOp_fp8_WGMMA_alignx_gemm(manifest, cuda_version) GenerateSM90_TensorOp_mixed_dtype_WGMMA_gemm(manifest, cuda_version) GenerateSM90_TensorOp_1684(manifest, cuda_version) GenerateSM90_TensorOp_16b_WGMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x) GenerateSM90_TensorOp_fp8_WGMMA_gemm(manifest, cuda_version, gemm_kind=GemmKind.GroupedUniversal3x) GenerateSM90_TensorOp_1684_complex(manifest, cuda_version) GenerateSM90_TensorOp_1684_complex_gaussian(manifest, cuda_version) GenerateSM90_TensorOp_1684_rank_k(manifest, cuda_version) GenerateSM90_TensorOp_1684_rank_k_complex(manifest, cuda_version) GenerateSM90_TensorOp_1684_rank_k_complex_gaussian(manifest, cuda_version) GenerateSM90_TensorOp_1684_trmm(manifest, cuda_version) GenerateSM90_TensorOp_1684_trmm_complex(manifest, cuda_version) GenerateSM90_TensorOp_1684_trmm_complex_gaussian(manifest, cuda_version) GenerateSM90_TensorOp_1684_symm(manifest, cuda_version) GenerateSM90_TensorOp_1684_symm_complex(manifest, cuda_version) GenerateSM90_TensorOp_1684_symm_complex_gaussian(manifest, cuda_version) GenerateSM90_Conv3x(manifest, cuda_version) GenerateSM90_SparseTensorOp_16b_WGMMA_gemm(manifest, cuda_version) GenerateSM90_SparseTensorOp_tf32_WGMMA_gemm(manifest, cuda_version) GenerateSM90_SparseTensorOp_int8_WGMMA_gemm(manifest, cuda_version) GenerateSM90_SparseTensorOp_fp8_WGMMA_gemm(manifest, cuda_version) GenerateSM90_TensorOp_fp8_WGMMA_gemm_with_blockwise(manifest, cuda_version) GenerateSM90_TensorOp_fp8_WGMMA_gemm_with_blockwise(manifest, cuda_version, gemm_kind=GemmKind.GroupedBlockwiseUniversal3x) ################################################################################################### def numeric_log_level(log_level: str) -> int: """ Converts the string identifier of the log level into the numeric identifier used in setting the log level. :param x: string representation of log level (e.g., 'INFO', 'DEBUG') :type x: str :return: numeric representation of log level :rtype: int """ numeric_level = getattr(logging, log_level.upper(), None) if not isinstance(numeric_level, int): raise ValueError(f'Invalid log level: {log_level}') return numeric_level # This function for defining the ArgumentParser is used to make it easy for the CUTLASS Python interface # to leverage the functionality in this file without running this script via a shell prompt. def define_parser(): parser = argparse.ArgumentParser(description="Generates device kernel registration code for CUTLASS Kernels") parser.add_argument("--operations", default="all", help="Specifies the operation to generate (gemm, all)") parser.add_argument("--build-dir", default=".", required=False, help="CUTLASS top-level build directory") parser.add_argument("--curr-build-dir", default=".", help="CUTLASS current build directory. cmake files will be emitted in this directory") parser.add_argument("--generator-target", default='library', help="Target of CUTLASS Library Generator.") parser.add_argument("--architectures", default='53;60;61;70;75;80;90;100', help="Target compute architectures") parser.add_argument("--kernels", default='', help='Comma-delimited list to filter kernels by name. ' + 'Specifying this as \"all\" includes ALL the kernels, ' + 'while not specifying this includes only the default set of kernels.') parser.add_argument("--ignore-kernels", default='', help='Comma-delimited list of kernels ' + 'to exclude from build. For backwards compatibility reasons, ' + 'this option only takes effect if --kernels is set to a nonempty value.') parser.add_argument("--exclude-kernels", default='', help='Comma-delimited list of kernels ' + 'to exclude from build. In contrast to --ignore-kernels, ' + 'this option always takes effect, ' + 'whether or not --kernels is set to a nonempty value. ' + 'It also can exclude kernels from the filter file ' + '(see --kernel-filter-file option below).') parser.add_argument("--filter-by-cc", default='True', type=str, help='If enabled, kernels whose compute capability range is not satisfied by the build target are excluded.') parser.add_argument("--cuda-version", default="11.0.0", help="Semantic version string of CUDA Toolkit") parser.add_argument('--kernel-filter-file', type=str, default=None, required=False, help='Full path of filter file') parser.add_argument('--heuristics-problems-file', type=str, default=None, required=False, help='Full path of heuristics problem size description file, as a json list') parser.add_argument('--heuristics-testlist-file', type=str, default=None, required=False, help='Full path of heuristics testlist CSV file, to be passed to cutlass_profiler') parser.add_argument('--heuristics-gpu', type=str, default=None, required=False, help='GPU to use for evaluating heuristics offline. None or `auto` to autodetect using cuda', choices=['', 'auto', 'H100_SXM', 'H100_PCIE', 'H100_NVL', 'H200_SXM', 'H20_SXM', 'B200', 'GB200_NVL', 'RTX_5080', 'RTX_5090', 'RTX_PRO_6000']) parser.add_argument('--heuristics-configs-per-problem', type=int, default=10, required=False, help='Number of kernel configs to generate for each problem in the problem list') parser.add_argument('--heuristics-restrict-kernels', action='store_true', help='Restrict heuristics mode to use only the default set of kernels emitted by generator.py') parser.add_argument('--selected-kernel-list', type=str, default=None, required=False, help='Specify the output log file containing all enabled kernels in this build') parser.add_argument("--interface-dir", default=None, required=False, help="Interface header to kernels") parser.add_argument("--disable-full-archs-compilation", action="store_true", required=False, help="Disable compilation for every archs in --architectures") parser.add_argument("--log-level", default='info', type=numeric_log_level, required=False, help='Logging level to be used by the generator script') parser.add_argument('--instantiation-level', type=str, default="", required=False, help="Instantiation level for SM90 kernels. Set to `max` and make sure `--kernels` is not empty to generate all possible configurations.") _add_package_disablement_flag(parser) return parser if __name__ == "__main__": parser = define_parser() args = parser.parse_args() # Set the logging level based on the user-provided `--log-level` command-line option logging.basicConfig(level=args.log_level) manifest = Manifest(args) archs = args.architectures.split(';') if args.heuristics_problems_file: filter_manifest_and_write_heuristics_file(manifest, args) GenerateSM50(manifest, args.cuda_version) GenerateSM60(manifest, args.cuda_version) GenerateSM61(manifest, args.cuda_version) GenerateSM70(manifest, args.cuda_version) GenerateSM75(manifest, args.cuda_version) GenerateSM80(manifest, args.cuda_version) GenerateSM89(manifest, args.cuda_version) GenerateSM90(manifest, args.cuda_version) blackwell_arch_list = [ "100a", "100f", "101a", "101f", "103a", "103f", "110a", "110f", "120a", "120f", "121a", "121f", ] blackwell_enabled_arch = any(arch in blackwell_arch_list for arch in archs) if blackwell_enabled_arch: GenerateSM100(manifest, args.cuda_version) GenerateSM120(manifest, args.cuda_version) if 'library' in args.generator_target.split(','): manifest.emit(GeneratorTarget.Library) if 'kernel_testlist_l0' in args.generator_target.split(','): emit_gemm_kernel_testlist(manifest, args.curr_build_dir, args.architectures, "functional_L0") if 'kernel_testlist_l1' in args.generator_target.split(','): emit_gemm_kernel_testlist(manifest, args.curr_build_dir, args.architectures, "functional_L1") if args.selected_kernel_list is not None: if len(manifest.selected_kernels) > 0: with open(args.selected_kernel_list, 'w') as file_writer: for line in manifest.selected_kernels: file_writer.write("%s\n" % line) ###################################################################################################