# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name, dangerous-default-value # ruff: noqa: E501 """Conv2d kernel generator and profiler for CUTLASS.""" import os import pickle from functools import partial from .conv2d_operation import Conv2dOperation, EmitConv2dInstance from .conv2d_profiler import Conv2dProfilerEmitter from .gen_gemm import CutlassGemmProfiler from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine from .library import ( ConvKind, DataType, EpilogueFunctor, IteratorAlgorithm, LayoutType, StrideSupport, SwizzlingFunctor, TensorDescription, ) def create_conv2d_operator_with_epilogue( conv_kind, stride_support, op_type, tile_description, data_type, alignment, alignment_epilogue, swizzling_functor, split_k_slices, ): """ Instantiate a cutlass kernel from the given configuration, along with the epilouge functor """ if "residual" in op_type: activation_map = { "cutlass.conv2d_bias_hardswish": "cutlass::epilogue::thread::HardSwish", "cutlass.conv2d_bias_silu": "cutlass::epilogue::thread::SiLu", "cutlass.conv2d_bias_sigmoid": "cutlass::epilogue::thread::Sigmoid", "cutlass.conv2d_bias_relu": "cutlass::epilogue::thread::ReLu", "cutlass.conv2d_bias": "cutlass::epilogue::thread::Identity", } prefix = op_type[: op_type.find("_residual")] activation = activation_map[prefix] binary_op = "cutlass::multiplies" if "residual_multiply" in op_type else "cutlass::plus" unary_op = ( "cutlass::epilogue::thread::ReLu" if op_type.endswith("relu") else "cutlass::epilogue::thread::Identity" ) residual_block_info = { "activation": activation, "binary_op": binary_op, "unary_op": unary_op, } epilogue = EpilogueFunctor.LinearCombinationResidualBlock no_beta_scaling = False else: residual_block_info = None epilogue, no_beta_scaling = EPILOGUE_MAP[op_type] element_a, element_b, element_c, element_epilogue = data_type A = TensorDescription(element_a, LayoutType.TensorNHWC, alignment) B = TensorDescription(element_b, LayoutType.TensorNHWC, alignment) C = TensorDescription(element_c, LayoutType.TensorNHWC, alignment_epilogue) op = Conv2dOperation( conv_kind, IteratorAlgorithm.Optimized, tile_description.minimum_compute_capability, tile_description, A, B, C, element_epilogue, stride_support, epilogue, swizzling_functor, split_k_slices, ) name = op.procedural_name() opdef = EmitConv2dInstance().emit( op, no_beta_scaling=no_beta_scaling, residual_block_info=residual_block_info, emit_reduction=split_k_slices > 1, ) return name, opdef def enumerate_conv2d_operators( conv_kind, stride_support, split_k_slices, alignment_c, tile_descriptions, data_type, alignment_constraints, swizzling_functor=SwizzlingFunctor.Identity4, ): """Exhaustively instantiate all kernels from a given configuration.""" ret = [] kernel_emitter = EmitConv2dInstance() profiler_emitter = Conv2dProfilerEmitter() element_a, element_b, element_c, element_epilogue = data_type if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided: swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1 for split_k_slice in split_k_slices: for tile in tile_descriptions: for alignmentAB in alignment_constraints: for alignmentC in alignment_c: A = TensorDescription(element_a, LayoutType.TensorNHWC, alignmentAB) B = TensorDescription(element_b, LayoutType.TensorNHWC, alignmentAB) C = TensorDescription(element_c, LayoutType.TensorNHWC, alignmentC) if element_c == DataType.s32 and A.alignment == 1: tile.threadblock_shape[0] = min(tile.threadblock_shape[0], 128) tile.threadblock_shape[1] = min(tile.threadblock_shape[1], 128) op = Conv2dOperation( conv_kind, IteratorAlgorithm.Optimized, tile.minimum_compute_capability, tile, A, B, C, element_epilogue, stride_support, EpilogueFunctor.LinearCombination, swizzling_functor, split_k_slice, ) ret.append( { "src": profiler_emitter.emit( kernel_emitter.emit(op, emit_reduction=split_k_slice > 1), op.procedural_name(), element_output=element_c, split_k_slices=split_k_slice, ), "name": op.procedural_name(), "tile_description": tile, "alignment": alignmentAB, "alignment_epilogue": alignmentC, "data_type": data_type, "swizzle_functor": swizzling_functor, "split_k_slices": split_k_slice, } ) return ret class CutlassConv2DProfiler: """Profile all candidate kernels and select the best one.""" def __init__(self, sm, cutlass_path, binary_path): self.gemm_profiler = CutlassGemmProfiler(sm, cutlass_path, binary_path) self.sm = sm assert sm in GENERATOR_FUNC_TABLE, f"sm{sm} not supported yet." self.engine = ProfilerEngine(sm, cutlass_path, binary_path) self.cache_path = os.path.join(binary_path, "cutlass_conv2d_cache.pickle") if os.path.exists(self.cache_path): self.cache = pickle.load(open(self.cache_path, "rb")) else: self.cache = {} def get_default( self, op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32, conv_kind=ConvKind.Fprop, stride=(1, 1), ): """Return the default kernel for the requested architecture. For now, the default kernel was picked arbitrary. """ gemm_profile_result = self.gemm_profiler.get_default( op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32 ) tile_description = gemm_profile_result["tile_description"] alignment = gemm_profile_result["alignment"] data_type = gemm_profile_result["data_type"] stride_support = StrideSupport.Strided if stride[0] > 1 else StrideSupport.Unity if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided: swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1 else: swizzling_functor = SwizzlingFunctor.Identity4 name, opdef = create_conv2d_operator_with_epilogue( conv_kind, stride_support, op_type, tile_description, data_type, alignment, alignment, swizzling_functor, split_k_slices=1, ) return {"name": name, "opdef": opdef} def select_op( self, d_shape, w_shape, padding, stride, dilation, out_dtype, data_dtype, weight_dtype, use_3xtf32, conv_kind, stride_support, split_k_slices, profile_all_alignments=False, find_first_valid=False, use_multiprocessing=False, ): """ Profile and select the best kernel from candidate kernels. See the documentation for the profile method below. """ N, H, W, IC = d_shape OC, R, S, _ = w_shape workload = ( N, H, W, IC, OC, R, S, padding[0], padding[1], stride[0], stride[1], dilation[0], dilation[1], ) if workload in self.cache: return self.cache[workload] def alignments(dtype): if dtype in ["float16"]: alignments = [8, 4, 2, 1] elif dtype in ["float", "float32"]: alignments = [4, 2, 1] else: raise ValueError(f"Unsupported data type: {dtype}") return alignments alignments_c = [align for align in alignments(out_dtype) if OC % align == 0] if not profile_all_alignments: alignments_c = [alignments_c[0]] ops = GENERATOR_FUNC_TABLE[self.sm]( out_dtype, data_dtype, weight_dtype, partial( enumerate_conv2d_operators, conv_kind, stride_support, split_k_slices, alignments_c, ), lambda align: all([dim % align == 0 for dim in [IC]]), use_3xtf32, profile_all_alignments, # Use fp32 accumulation for wgrad to align with cuDNN accumlator_dtype="float32" if conv_kind == ConvKind.Wgrad else out_dtype, ) if not find_first_valid: self.engine.compile_all(ops, use_multiprocessing) args = "--n={} --h={} --w={} --c={} --k={} --r={} --s={} --pad_h={} --pad_w={} --stride_h={} --stride_w={} --dilation_h={} --dilation_w={}".format( *workload ) for op in ops: out = self.engine.evaluate(op, args.split(" ")) op["runtime"] = out if out < float("inf") and find_first_valid: self.cache[workload] = op return op op = min(ops, key=lambda i: i["runtime"]) self.cache[workload] = op with open(self.cache_path, "wb") as f: pickle.dump(self.cache, f) return op def profile( self, op_type, d_shape, w_shape, padding, stride, dilation, out_dtype, data_dtype, weight_dtype, use_3xtf32=True, conv_kind=ConvKind.Fprop, split_k_slices=[1], profile_all_alignments=False, find_first_valid=False, use_multiprocessing=False, ): """Profile and select the best kernel from candidate kernels. If find_first_valid is True, return immediately after the first applicable kernel is found. If use_multiprocessing is True, compile all profiler executables in parallel. """ # Dgrad requires Unity stride when stride == (1, 1) stride_support = ( StrideSupport.Unity if stride[0] == 1 and stride[1] == 1 and conv_kind == ConvKind.Dgrad else StrideSupport.Strided ) op = self.select_op( d_shape, w_shape, padding, stride, dilation, out_dtype, data_dtype, weight_dtype, use_3xtf32, conv_kind, stride_support, split_k_slices, profile_all_alignments, find_first_valid, use_multiprocessing, ) name, opdef = create_conv2d_operator_with_epilogue( conv_kind, stride_support, op_type, op["tile_description"], op["data_type"], op["alignment"], op["alignment_epilogue"], op["swizzle_functor"], op["split_k_slices"], ) return name, opdef, op["runtime"]