From aa2491b0b6ceff1fdb9c602da63829de8d73260a Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Thu, 20 Nov 2025 15:34:03 +0100 Subject: [PATCH 01/44] Added Triton support from https://github.com/super-nexus/kernel_tuner/tree/triton_support --- examples/triton/conv2d_tuning.py | 306 ++++++++++++++++++ examples/triton/vec_add.py | 46 +++ kernel_tuner/backends/triton.py | 174 ++++++++++ kernel_tuner/core.py | 247 ++++---------- kernel_tuner/interface.py | 7 +- kernel_tuner/kernel_sources/__init__.py | 0 kernel_tuner/kernel_sources/kernel_source.py | 77 +++++ .../kernel_sources/kernel_source_fn.py | 160 +++++++++ .../kernel_sources/kernel_source_str.py | 174 ++++++++++ kernel_tuner/kernel_sources/model/__init__.py | 0 .../model/prepared_kernel_source_data.py | 10 + kernel_tuner/kernelbuilder.py | 5 +- kernel_tuner/language.py | 10 + kernel_tuner/observers/triton.py | 32 ++ kernel_tuner/util.py | 7 +- 15 files changed, 1072 insertions(+), 183 deletions(-) create mode 100644 examples/triton/conv2d_tuning.py create mode 100644 examples/triton/vec_add.py create mode 100644 kernel_tuner/backends/triton.py create mode 100644 kernel_tuner/kernel_sources/__init__.py create mode 100644 kernel_tuner/kernel_sources/kernel_source.py create mode 100644 kernel_tuner/kernel_sources/kernel_source_fn.py create mode 100644 kernel_tuner/kernel_sources/kernel_source_str.py create mode 100644 kernel_tuner/kernel_sources/model/__init__.py create mode 100644 kernel_tuner/kernel_sources/model/prepared_kernel_source_data.py create mode 100644 kernel_tuner/language.py create mode 100644 kernel_tuner/observers/triton.py diff --git a/examples/triton/conv2d_tuning.py b/examples/triton/conv2d_tuning.py new file mode 100644 index 000000000..5389b4a28 --- /dev/null +++ b/examples/triton/conv2d_tuning.py @@ -0,0 +1,306 @@ +import torch +import triton.language as tl +import numpy as np +from kernel_tuner.interface import tune_kernel +import os +import json +from datetime import datetime + +# Check for required environment variable +cache_dir = os.getenv('KERNEL_TUNER_CACHE_DIR') +cache_file_name = os.getenv('KERNEL_TUNER_CACHE_FILE', 'conv2d_tuning_results.json') + +if cache_dir is None: + raise ValueError("Environment variable KERNEL_TUNER_CACHE_DIR must be set") + +cache_file = os.path.join(cache_dir, cache_file_name) + + +def conv2d_output_size( + in_size: int, + kernel_size: int, + stride: int, + padding: int, + dilation: int, +) -> int: + """ + Determines the output size of a 2D convolution operation. + + Args: + in_size: Input size. + kernel_size: Kernel size. + stride: Stride. + padding: Padding. + + Returns: + Output size of 2D convolution. + """ + return (in_size + 2 * padding - dilation * (kernel_size - 1) - 1) // stride + 1 + + +def conv2d_forward_kernel( + input_pointer, + weight_pointer, + output_pointer, + in_n, + input_height, + input_width, + out_c, + out_height, + out_width, + input_n_stride, + input_c_stride, + input_height_stride, + input_width_stride, + weight_n_stride, + weight_c_stride, + weight_height_stride, + weight_width_stride, + output_n_stride, + output_c_stride, + output_height_stride, + output_width_stride, + weight_c: tl.constexpr, + weight_height: tl.constexpr, + weight_width: tl.constexpr, + stride_height: tl.constexpr, + stride_width: tl.constexpr, + padding_height: tl.constexpr, + padding_width: tl.constexpr, + dilation_height: tl.constexpr, + dilation_width: tl.constexpr, + groups: tl.constexpr, + BLOCK_NI_HO_WO: tl.constexpr, + BLOCK_CI: tl.constexpr, + BLOCK_CO: tl.constexpr, +): + pid_ni_ho_wo = tl.program_id(0) + pid_co = tl.program_id(1) + pid_group = tl.program_id(2) + + # caculate in_n out_height out_weight value in kernel + ni_ho_wo_offset = pid_ni_ho_wo * BLOCK_NI_HO_WO + tl.arange(0, BLOCK_NI_HO_WO) + ni_ho_offset = ni_ho_wo_offset // out_width + in_n_point_value = ni_ho_offset // out_height + output_height_point_value = ni_ho_offset % out_height + output_width_point_value = ni_ho_wo_offset % out_width + + # Load the input and weight pointers. input and weight are of shape + # [in_n, groups, in_c, input_height, input_width] and [groups, out_c, in_c, weight_height, weight_width] + out_per_group_c = out_c // groups + output_c_offset = pid_co * BLOCK_CO + tl.arange(0, BLOCK_CO) + input_pointer += ( + input_n_stride * in_n_point_value + input_c_stride * pid_group * weight_c + )[:, None] + weight_pointer += ( + weight_n_stride * output_c_offset + + weight_n_stride * pid_group * out_per_group_c + )[None, :] + + accum = tl.zeros((BLOCK_NI_HO_WO, BLOCK_CO), dtype=tl.float32) + BLOCK_CI_COUNT = (weight_c + BLOCK_CI - 1) // BLOCK_CI + for hwc in range(weight_height * weight_width * BLOCK_CI_COUNT): + c = (hwc % BLOCK_CI_COUNT) * BLOCK_CI + hw = hwc // BLOCK_CI_COUNT + h = hw // weight_width + w = hw % weight_width + + input_c_offset = c + tl.arange(0, BLOCK_CI) + input_height_offset = ( + h * dilation_height + - padding_height + + stride_height * output_height_point_value + ) + input_width_offset = ( + w * dilation_width - padding_width + stride_width * output_width_point_value + ) + + curr_input_pointer = ( + input_pointer + + (input_c_stride * input_c_offset)[None, :] + + (input_height_stride * input_height_offset)[:, None] + + (input_width_stride * input_width_offset)[:, None] + ) + curr_weight_pointer = ( + weight_pointer + + (weight_c_stride * input_c_offset)[:, None] + + (weight_height_stride * h) + + (weight_width_stride * w) + ) + + input_mask = ( + (in_n_point_value < in_n)[:, None] + & (input_c_offset < weight_c)[None, :] + & (0 <= input_height_offset)[:, None] + & (input_height_offset < input_height)[:, None] + & (0 <= input_width_offset)[:, None] + & (input_width_offset < input_width)[:, None] + ) + weight_mask = (input_c_offset < weight_c)[:, None] & ( + output_c_offset < out_per_group_c + )[None, :] + + input_block = tl.load(curr_input_pointer, mask=input_mask) + weight_block = tl.load(curr_weight_pointer, mask=weight_mask) + + accum += tl.dot(input_block, weight_block, allow_tf32=False) + + output_pointer += ( + (output_n_stride * in_n_point_value)[:, None] + + (output_c_stride * (pid_group * out_per_group_c + output_c_offset))[None, :] + + (output_height_stride * output_height_point_value)[:, None] + + (output_width_stride * output_width_point_value)[:, None] + ) + output_mask = ( + (in_n_point_value < in_n)[:, None] + & (output_c_offset < out_per_group_c)[None, :] + & (output_height_point_value < out_height)[:, None] + & (output_width_point_value < out_width)[:, None] + ) + + tl.store(output_pointer, accum, mask=output_mask) + + +def tune_conv2d(batch_size=1, in_channels=64, height=32, width=32, + out_channels=128, kernel_size=3, stride=1, padding=1, + groups=1): + """ + Tune the conv2d kernel with different configurations. + """ + # Create sample inputs + input = torch.randn(batch_size, in_channels, height, width, + device='cuda', dtype=torch.float32) + weight = torch.randn(out_channels, in_channels//groups, kernel_size, kernel_size, + device='cuda', dtype=torch.float32) + + # Calculate output dimensions + out_height = conv2d_output_size(height, kernel_size, stride, padding, 1) + out_width = conv2d_output_size(width, kernel_size, stride, padding, 1) + output = torch.empty((batch_size, out_channels, out_height, out_width), + device='cuda', dtype=torch.float32) + + # Prepare all arguments for the kernel + arguments = [ + input, weight, output, + np.int32(batch_size), + np.int32(height), + np.int32(width), + np.int32(out_channels), + np.int32(out_height), + np.int32(out_width), + np.int32(input.stride(0)), + np.int32(input.stride(1)), + np.int32(input.stride(2)), + np.int32(input.stride(3)), + np.int32(weight.stride(0)), + np.int32(weight.stride(1)), + np.int32(weight.stride(2)), + np.int32(weight.stride(3)), + np.int32(output.stride(0)), + np.int32(output.stride(1)), + np.int32(output.stride(2)), + np.int32(output.stride(3)), + np.int32(in_channels//groups), # weight_c + np.int32(kernel_size), # weight_height + np.int32(kernel_size), # weight_width + np.int32(stride), # stride_height + np.int32(stride), # stride_width + np.int32(padding), # padding_height + np.int32(padding), # padding_width + np.int32(1), # dilation_height + np.int32(1), # dilation_width + np.int32(groups), # groups + ] + + # Define tuning parameters - only powers of 2 + tune_params = { + 'BLOCK_NI_HO_WO': [2 ** i for i in range(4, 10)], + 'BLOCK_CI': [2 ** i for i in range(4, 10)], + 'BLOCK_CO': [2 ** i for i in range(4, 10)], + 'num_stages': [1, 2, 3, 4], + 'num_warps': [1, 2, 4, 8], + } + + print(tune_params) + + # Define constraints + constraints = [ + "BLOCK_CI <= %d" % (in_channels//groups), + "BLOCK_CO <= %d" % out_channels, + ] + + # Problem size for the grid + problem_size = ( + batch_size * out_height * out_width, # Grid dimension 0 + out_channels, # Grid dimension 1 + groups, # Grid dimension 2 + ) + + # Grid divisor expressions + grid_div_x = ["BLOCK_NI_HO_WO"] + grid_div_y = ["BLOCK_CO"] + grid_div_z = ["1"] + + results, env = tune_kernel( + kernel_name='conv2d_forward_kernel', + kernel_source=conv2d_forward_kernel, + problem_size=problem_size, + arguments=arguments, + tune_params=tune_params, + restrictions=constraints, + lang='TRITON', + grid_div_x=grid_div_x, + grid_div_y=grid_div_y, + grid_div_z=grid_div_z, + block_size_names=['BLOCK_NI_HO_WO', 'BLOCK_CI', 'BLOCK_CO'], + strategy='genetic_algorithm', + strategy_options={ + 'maxiter': 1000, + 'popsize': 100, + }, + cache=cache_file, + ) + + return results + + +if __name__ == '__main__': + # Run tuning with moderately large input dimensions + results = tune_conv2d( + batch_size=16, + in_channels=128, + height=112, + width=112, + out_channels=256, + kernel_size=3, + stride=1, + padding=1, + groups=1 + ) + + + # Filter out failed compilations and find best config + valid_results = [result for result in results if isinstance(result['time'], (int, float))] + if valid_results: + best_config = min(valid_results, key=lambda x: x['time']) + print("\nBest configuration:") + print(json.dumps(best_config, indent=2)) + else: + print("\nNo valid configurations found - all compilations failed") + + # Create results dictionary with GPU info + all_results = { + "gpu_info": { + "gpu_name": torch.cuda.get_device_name() + }, + "results": valid_results + } + + # Add timestamp to filename + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + output_file = f'conv2d_results_{timestamp}.json' + + # Save results + import json + with open(output_file, 'w') as f: + json.dump(all_results, f, indent=2) \ No newline at end of file diff --git a/examples/triton/vec_add.py b/examples/triton/vec_add.py new file mode 100644 index 000000000..555cfba46 --- /dev/null +++ b/examples/triton/vec_add.py @@ -0,0 +1,46 @@ +import numpy +import triton.language as tl +import torch +from kernel_tuner import tune_kernel, run_kernel +from kernel_tuner.file_utils import store_output_file, store_metadata_file + + +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + # note: `constexpr` so it can be used as a shape value. + ): + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + + +size = 10000000 + +a = torch.randn(size, dtype=torch.float32) +b = torch.randn(size, dtype=torch.float32) +c = torch.zeros_like(b) +n = torch.tensor(size, dtype=torch.int32) + +args = [a, b, c, n] + +tune_params = dict() +tune_params["block_size_x"] = [2**i for i in range(10)] + +results, env = tune_kernel( + kernel_name="add_kernel", + kernel_source=add_kernel, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="triton" +) + +print("Hello") \ No newline at end of file diff --git a/kernel_tuner/backends/triton.py b/kernel_tuner/backends/triton.py new file mode 100644 index 000000000..cf6fda504 --- /dev/null +++ b/kernel_tuner/backends/triton.py @@ -0,0 +1,174 @@ +import logging +import numpy as np + +from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.observers.triton import TritonRuntimeObserver + +try: + import torch +except ImportError: + logging.error("Torch not available") + +try: + import triton + import triton.language as tl +except ImportError: + triton = None + tl = None + logging.error("Unable to load triton") + + +class TritonFunctions(GPUBackend): + + def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): + if not triton or not torch: + logging.error("Triton or torch not available") + raise ImportError("Triton or torch not available") + + self.device_id = torch.cuda.current_device() + + self.device_properties = torch.cuda.get_device_properties(self.device_id) + self.name = torch.cuda.get_device_name(self.device_id) + self.max_threads = self.device_properties.max_threads_per_multi_processor + + env = dict() + env["device_name"] = self.name + env["max_threads"] = self.max_threads + env["iterations"] = iterations + env["compiler_options"] = compiler_options + self.env = env + + self.stream = torch.cuda.default_stream() + self.start = torch.cuda.Event(enable_timing=True) + self.end = torch.cuda.Event(enable_timing=True) + + # setup observers + self.observers = observers or [] + self.observers.append(TritonRuntimeObserver(self)) + for obs in self.observers: + obs.register_device(self) + + self.units = {"time": "ms", "power": "s,mW", "energy": "J"} + + super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) + + def ready_argument_list(self, arguments): + # Allocate memory here + torch_args = [] + + for arg in arguments: + if isinstance(arg, torch.Tensor) and arg.dim() > 0: + torch_args.append(arg.cuda()) + elif isinstance(arg, torch.Tensor) and arg.dim() == 0: + scalar_value = arg.item() + torch_args.append(scalar_value) + elif isinstance(arg, np.ndarray): + torch_arg = torch.from_numpy(arg) + torch_arg_gpu = torch_arg.cuda() + torch_args.append(torch_arg_gpu) + elif isinstance(arg, np.generic): + scalar_value = arg.item() + torch_args.append(scalar_value) + else: + logging.warning("Unknown instance in triton functions") + + return torch_args + + def compile(self, kernel_instance, gpu_args=None): + logging.debug("Compiling triton kernel") + + if kernel_instance.kernel_fn is None: + raise ValueError("kernel_fn is None, currently Triton only supports callable kernel_source") + + if gpu_args is None: + raise ValueError("gpu_args is None, Triton needs gpu args to compile the kernel") + + grid = kernel_instance.grid + threads = kernel_instance.threads + jit_function = triton.jit(kernel_instance.kernel_fn) + params = kernel_instance.params + gpu_kwargs = self.build_gpu_kwargs(jit_function, threads, params) + + # Call the jit function in order to compile it + jit_function[grid](*gpu_args, **gpu_kwargs) + + return jit_function + + def start_event(self): + logging.debug("Start triton event") + self.start.record() + + def stop_event(self): + logging.debug("Stop triton event") + self.end.record() + + def kernel_finished(self): + logging.debug("Checking if kernel has finished") + return self.end.query() + + def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): + if params is None: + raise ValueError("params is None, Triton needs params in order to set num_warps, num_ctas, etc.") + + # Run the kernel + if stream is None: + stream = self.stream + + gpu_kwargs = self.build_gpu_kwargs(func, threads, params) + + with torch.cuda.stream(stream): + logging.debug("Running triton kernel") + func[grid](*gpu_args, **gpu_kwargs) + + def build_gpu_kwargs(self, jit_fn, threads, params=None): + gpu_kwargs = {} + + if 'BLOCK_SIZE' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE'] = threads[0] + + if 'BLOCK_SIZE_X' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE_X'] = threads[0] + + if 'BLOCK_SIZE_Y' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE_Y'] = threads[1] + + if 'BLOCK_SIZE_Z' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE_Z'] = threads[2] + + if params is None: + return gpu_kwargs + + for param in params: + if param in jit_fn.arg_names: + gpu_kwargs[param] = params[param] + + # Check for Triton specific parameters + if 'num_warps' in params: + gpu_kwargs['num_warps'] = params['num_warps'] + if 'num_ctas' in params: + gpu_kwargs['num_ctas'] = params['num_ctas'] + if 'num_stages' in params: + gpu_kwargs['num_stages'] = params['num_stages'] + + return gpu_kwargs + + def synchronize(self): + torch.cuda.synchronize() + + def memset(self, allocation, value, size): + pass + + def memcpy_dtoh(self, dest, src): + pass + + def memcpy_htod(self, dest, src): + pass + + def copy_constant_memory_args(self, cmem_args): + raise NotImplementedError("Triton does not support constant memory") + + def copy_shared_memory_args(self, smem_args): + raise NotImplementedError("Triton does not support shared memory") + + def copy_texture_memory_args(self, texmem_args): + raise NotImplementedError("Triton does not support texture memory") \ No newline at end of file diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 5352ced74..1a7e67b34 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -21,6 +21,8 @@ from kernel_tuner.backends.nvcuda import CudaFunctions from kernel_tuner.backends.opencl import OpenCLFunctions from kernel_tuner.backends.pycuda import PyCudaFunctions +from kernel_tuner.backends.triton import TritonFunctions +from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.observers.nvml import NVMLObserver from kernel_tuner.observers.observer import ContinuousObserver, OutputObserver, PrologueObserver from kernel_tuner.observers.tegra import TegraObserver @@ -41,6 +43,7 @@ "name", "kernel_source", "kernel_string", + "kernel_fn", "temp_files", "threads", "grid", @@ -53,13 +56,27 @@ class KernelInstance(_KernelInstance): """Class that represents the specific parameterized instance of a kernel.""" + def __new__(cls, *args, **kwargs): + # Detect old-style calls (without kernel_fn) + if len(args) == 8: # old version + name, kernel_source, kernel_string, temp_files, threads, grid, params, arguments = args + kernel_fn = None + args = (name, kernel_source, kernel_string, kernel_fn, temp_files, threads, grid, params, arguments) + elif "kernel_fn" not in kwargs and len(args) == 8: + kwargs["kernel_fn"] = None + return super().__new__(cls, *args, **kwargs) + def delete_temp_files(self): """Delete any generated temp files.""" - for v in self.temp_files.values(): - util.delete_temp_file(v) + tmp_files_list = self.temp_files.values() if isinstance(self.temp_files, dict) else self.temp_files + for tmp_file in tmp_files_list: + util.delete_temp_file(tmp_file) def prepare_temp_files_for_error_msg(self): """Prepare temp file with source code, and return list of temp file names.""" + if type(self.kernel_source).__name__ == "KernelSourceFn": + return [] # TODO what do we want to return here? + temp_filename = util.get_temp_filename(suffix=self.kernel_source.get_suffix()) util.write_file(temp_filename, self.kernel_string) ret = [temp_filename] @@ -67,160 +84,7 @@ def prepare_temp_files_for_error_msg(self): return ret -class KernelSource(object): - """Class that holds the kernel sources. - - There is a primary kernel source, which can be either a source string, - a filename (indicating a file containing the kernel source code), - or a callable (generating the kernel source code). - There can additionally be (one or multiple) secondary kernel sources, which - must be filenames. - """ - - def __init__(self, kernel_name, kernel_sources, lang, defines=None): - if not isinstance(kernel_sources, list): - kernel_sources = [kernel_sources] - self.kernel_sources = kernel_sources - self.kernel_name = kernel_name - self.defines = defines - if lang is None: - if callable(self.kernel_sources[0]): - raise TypeError("Please specify language when using a code generator function") - kernel_string = self.get_kernel_string(0) - lang = util.detect_language(kernel_string) - - # The validity of lang is checked later, when creating the DeviceInterface - self.lang = lang.upper() - - def get_kernel_string(self, index=0, params=None): - """Retrieve the kernel source with the given index and return as a string. - - See util.get_kernel_string() for details. - - :param index: Index of the kernel source in the list of sources. - :type index: int - - :param params: Dictionary containing the tunable parameters for this specific - kernel instance, only needed when kernel_source is a generator. - :type param: dict - - :returns: A string containing the kernel code. - :rtype: string - """ - logging.debug("get_kernel_string called") - - if hasattr(self, 'lang') and self.lang.upper() == "HYPERTUNER": - return "" - - kernel_source = self.kernel_sources[index] - return util.get_kernel_string(kernel_source, params) - - def prepare_list_of_files( - self, kernel_name, params, grid, threads, block_size_names - ): - """Prepare the kernel string along with any additional files. - - The first file in the list is allowed to include or read in the others - The files beyond the first are considered additional files that may also contain tunable parameters - - For each file beyond the first this function creates a temporary file with - preprocessors statements inserted. Occurrences of the original filenames in the - first file are replaced with their temporary counterparts. - - :param kernel_name: A string specifying the kernel name. - :type kernel_name: string - - :param params: A dictionary with the tunable parameters for this particular - instance. - :type params: dict() - - :param grid: The grid dimensions for this instance. The grid dimensions are - also inserted into the code as if they are tunable parameters for - convenience. - :type grid: tuple() - - :param threads: The thread block dimensions for this instance. The thread block are - also inserted into the code as if they are tunable parameters for - convenience. - :type threads: tuple() - - :param block_size_names: A list of strings that denote the names - for the thread block dimensions. - :type block_size_names: list(string) - - """ - temp_files = dict() - - if self.lang.upper() == "HYPERTUNER": - return tuple(["", "", temp_files]) - - for i, f in enumerate(self.kernel_sources): - if i > 0 and not util.looks_like_a_filename(f): - raise ValueError("When passing multiple kernel sources, the secondary entries must be filenames") - - ks = self.get_kernel_string(i, params) - # add preprocessor statements - n, ks = util.prepare_kernel_string( - kernel_name, - ks, - params, - grid, - threads, - block_size_names, - self.lang, - self.defines, - ) - - if i == 0: - # primary kernel source - name = n - kernel_string = ks - continue - - # save secondary kernel sources to temporary files - - # generate temp filename with the same extension - temp_file = util.get_temp_filename(suffix="." + f.split(".")[-1]) - temp_files[f] = temp_file - util.write_file(temp_file, ks) - # replace occurrences of the additional file's name in the first kernel_string with the name of the temp file - kernel_string = kernel_string.replace(f, temp_file) - - return name, kernel_string, temp_files - - def get_user_suffix(self, index=0): - """Get the suffix of the kernel filename, if the user specified one. Return None otherwise.""" - if util.looks_like_a_filename(self.kernel_sources[index]) and ("." in self.kernel_sources[index]): - return "." + self.kernel_sources[index].split(".")[-1] - return None - - def get_suffix(self, index=0): - """Return a suitable suffix for a kernel filename. - - This uses the user-specified suffix if available, or one based on the - lang/backend otherwise. - """ - # TODO: Consider delegating this to the backend - suffix = self.get_user_suffix(index) - if suffix is not None: - return suffix - - _suffixes = {"CUDA": ".cu", "OpenCL": ".cl", "C": ".c"} - try: - return _suffixes[self.lang] - except KeyError: - return ".c" - - def check_argument_lists(self, kernel_name, arguments): - """Check if the kernel arguments have the correct types. - - This is done by calling util.check_argument_list on each kernel string. - """ - for i, f in enumerate(self.kernel_sources): - if not callable(f): - util.check_argument_list(kernel_name, self.get_kernel_string(i), arguments) - else: - logging.debug("Checking of arguments list not supported yet for code generators.") +# KernelSource has been moved to kernel_sources/kernel_source_str.py class DeviceInterface(object): @@ -320,9 +184,16 @@ def __init__( compiler_options=compiler_options ) self.requires_warmup = False + elif lang.upper() == "TRITON": + dev = TritonFunctions( + device, + compiler_options=compiler_options, + iterations=iterations, + observers=observers + ) else: raise NotImplementedError( - "Sorry, support for languages other than CUDA, OpenCL, HIP, C, and Fortran is not implemented yet" + "Sorry, support for languages other than CUDA, OpenCL, HIP, C, Triton and Fortran is not implemented yet" ) self.dev = dev @@ -361,17 +232,23 @@ def __init__( if not quiet: print("Using: " + self.dev.name) - def benchmark_prologue(self, func, gpu_args, threads, grid, result): + def run_kernel_bench(self, func, gpu_args, threads, grid, stream=None, params=None): + if isinstance(self.dev, TritonFunctions): + self.dev.run_kernel(func, gpu_args, threads, grid, params=params) + else: + self.dev.run_kernel(func, gpu_args, threads, grid) + + def benchmark_prologue(self, func, gpu_args, threads, grid, result, params=None): """Benchmark prologue one kernel execution per PrologueObserver.""" for obs in self.prologue_observers: self.dev.synchronize() obs.before_start() - self.dev.run_kernel(func, gpu_args, threads, grid) + self.run_kernel_bench(func, gpu_args, threads, grid, stream=None, params=params) self.dev.synchronize() obs.after_finish() result.update(obs.get_results()) - def benchmark_default(self, func, gpu_args, threads, grid, result): + def benchmark_default(self, func, gpu_args, threads, grid, result, params=None): """Benchmark one kernel execution for 'iterations' at a time.""" self.dev.synchronize() for _ in range(self.iterations): @@ -379,7 +256,7 @@ def benchmark_default(self, func, gpu_args, threads, grid, result): obs.before_start() self.dev.synchronize() self.dev.start_event() - self.dev.run_kernel(func, gpu_args, threads, grid) + self.run_kernel_bench(func, gpu_args, threads, grid, stream=None, params=params) self.dev.stop_event() for obs in self.benchmark_observers: obs.after_start() @@ -395,7 +272,7 @@ def benchmark_default(self, func, gpu_args, threads, grid, result): result.update(obs.get_results()) - def benchmark_continuous(self, func, gpu_args, threads, grid, result, duration): + def benchmark_continuous(self, func, gpu_args, threads, grid, result, duration, params=None): """Benchmark continuously for at least 'duration' seconds.""" iterations = int(np.ceil(duration / (result["time"] / 1000))) self.dev.synchronize() @@ -403,7 +280,7 @@ def benchmark_continuous(self, func, gpu_args, threads, grid, result, duration): obs.before_start() self.dev.start_event() for _ in range(iterations): - self.dev.run_kernel(func, gpu_args, threads, grid) + self.run_kernel_bench(func, gpu_args, threads, grid, stream=None, params=params) self.dev.stop_event() for obs in self.continuous_observers: obs.after_start() @@ -453,8 +330,8 @@ def benchmark(self, func, gpu_args, instance, verbose, objective, skip_nvml_sett result = {} try: - self.benchmark_prologue(func, gpu_args, instance.threads, instance.grid, result) - self.benchmark_default(func, gpu_args, instance.threads, instance.grid, result) + self.benchmark_prologue(func, gpu_args, instance.threads, instance.grid, result, instance.params) + self.benchmark_default(func, gpu_args, instance.threads, instance.grid, result, instance.params) if self.continuous_observers: duration = 1 @@ -462,7 +339,7 @@ def benchmark(self, func, gpu_args, instance, verbose, objective, skip_nvml_sett obs.results = result duration = max(duration, obs.continuous_duration) - self.benchmark_continuous(func, gpu_args, instance.threads, instance.grid, result, duration) + self.benchmark_continuous(func, gpu_args, instance.threads, instance.grid, result, duration, instance.params) except Exception as e: # some launches may fail because too many registers are required @@ -507,7 +384,7 @@ def check_kernel_output( self.dev.refresh_memory(gpu_args, instance.arguments, should_sync) # run the kernel - check = self.run_kernel(func, gpu_args, instance) + check = self.run_kernel_check(func, gpu_args, instance) if not check: # runtime failure occurred that should be ignored, skip correctness check return @@ -575,7 +452,7 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, try: # compile the kernel start_compilation = time.perf_counter() - func = self.compile_kernel(instance, verbose) + func = self.compile_kernel(instance, verbose, gpu_args) if not func: result[to.objective] = util.CompilationFailedConfig() else: @@ -624,14 +501,17 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, return result - def compile_kernel(self, instance, verbose): + def compile_kernel(self, instance, verbose, gpu_args=None): """Compile the kernel for this specific instance.""" logging.debug("compile_kernel " + instance.name) # compile kernel_string into device func func = None try: - func = self.dev.compile(instance) + if isinstance(self.dev, TritonFunctions): + func = self.dev.compile(instance, gpu_args) + else: + func = self.dev.compile(instance) except Exception as e: # compiles may fail because certain kernel configurations use too # much shared memory for example, the desired behavior is to simply @@ -692,20 +572,21 @@ def create_kernel_instance(self, kernel_source, kernel_options, params, verbose) params, kernel_options.block_size_names, ) - if np.prod(threads) > self.dev.max_threads: + if kernel_source.lang != 'TRITON' and np.prod(threads) > self.dev.max_threads: if verbose: print(f"skipping config {util.get_instance_string(params)} reason: too many threads per block") return util.InvalidConfig() # obtain the kernel_string and prepare additional files, if any - name, kernel_string, temp_files = kernel_source.prepare_list_of_files( - kernel_options.kernel_name, + instance_data = kernel_source.prepare_kernel_instance( + kernel_options, params, grid, threads, - kernel_options.block_size_names, ) + name = instance_data.kernel_name + kernel_string=instance_data.kernel_str # check for templated kernel if kernel_source.lang in ["CUDA", "NVCUDA", "HIP"] and "<" in name and ">" in name: kernel_string, name = wrap_templated_kernel(kernel_string, name) @@ -714,7 +595,17 @@ def create_kernel_instance(self, kernel_source, kernel_options, params, verbose) arguments = _preprocess_gpu_arguments(kernel_options.arguments, params) # collect everything we know about this instance and return it - return KernelInstance(name, kernel_source, kernel_string, temp_files, threads, grid, params, arguments) + return KernelInstance( + name=name, + kernel_source=kernel_source, + kernel_string=kernel_string, + kernel_fn=instance_data.kernel_fn, + temp_files=instance_data.temp_files, + threads=threads, + grid=grid, + params=params, + arguments=arguments, + ) def get_environment(self): """Return dictionary with information about the environment.""" @@ -751,14 +642,15 @@ def ready_argument_list(self, arguments): return gpu_args - def run_kernel(self, func, gpu_args, instance): + + def run_kernel_check(self, func, gpu_args, instance): """Run a compiled kernel instance on a device.""" logging.debug("run_kernel %s", instance.name) logging.debug("thread block dims (%d, %d, %d)", *instance.threads) logging.debug("grid dims (%d, %d, %d)", *instance.grid) try: - self.dev.run_kernel(func, gpu_args, instance.threads, instance.grid) + self.run_kernel_bench(func, gpu_args, instance.threads, instance.grid, stream=None, params=instance.params) except Exception as e: if "too many resources requested for launch" in str(e) or "OUT_OF_RESOURCES" in str(e): logging.debug("ignoring runtime failure due to too many resources required") @@ -767,6 +659,7 @@ def run_kernel(self, func, gpu_args, instance): logging.debug("encountered unexpected runtime failure: " + str(e)) raise e return True + def _preprocess_gpu_arguments(old_arguments, params): diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index 32e91c86f..bad921d24 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -39,6 +39,7 @@ import kernel_tuner.util as util from kernel_tuner.file_utils import get_input_file, get_t4_metadata, get_t4_results, import_class_from_file from kernel_tuner.integration import get_objective_defaults +from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.runners.sequential import SequentialRunner from kernel_tuner.runners.simulation import SimulationRunner from kernel_tuner.searchspace import Searchspace @@ -592,7 +593,7 @@ def tune_kernel( if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = core.KernelSource(kernel_name, kernel_source, lang, defines) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -777,7 +778,7 @@ def run_kernel( if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = core.KernelSource(kernel_name, kernel_source, lang, defines) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -825,7 +826,7 @@ def run_kernel( instance.delete_temp_files() # run the kernel - if not dev.run_kernel(func, gpu_args, instance): + if not dev.run_kernel_check(func, gpu_args, instance): raise RuntimeError("runtime error occured, too many resources requested") # copy data in GPU memory back to the host diff --git a/kernel_tuner/kernel_sources/__init__.py b/kernel_tuner/kernel_sources/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/kernel_tuner/kernel_sources/kernel_source.py b/kernel_tuner/kernel_sources/kernel_source.py new file mode 100644 index 000000000..be8bbb265 --- /dev/null +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -0,0 +1,77 @@ +import inspect +import kernel_tuner.util as util + +from abc import abstractmethod + +from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData + + + + +class KernelSource: + def __new__(cls, kernel_name, kernel_sources, lang, defines=None): + """Factory behavior""" + if cls is KernelSource: + if inspect.isfunction(kernel_sources) and (lang and lang.upper() == "TRITON"): # TODO should this be isfunction? + from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn + print("CREATING KSFN") + return KernelSourceFn(kernel_name, kernel_sources, lang, defines) + else: + from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr + print("CREATING KSSTR") + return KernelSourceStr(kernel_name, kernel_sources, lang, defines) + # otherwise, normal subclass init + return super().__new__(cls) + + def __init__(self, kernel_name, kernel_sources, lang, defines=None): + if not isinstance(kernel_sources, list): + kernel_sources = [kernel_sources] + self.kernel_sources = kernel_sources + self.kernel_name = kernel_name + self.defines = defines + + if lang is None: + if callable(self.kernel_sources[0]): + raise TypeError("Please specify language when using a code generator function") + kernel_string = self.get_kernel_string(0) + self.lang = util.detect_language(kernel_string) + else: + self.lang = lang + + @abstractmethod + def prepare_kernel_instance(self, kernel_options, params, grid, threads) -> PreparedKernelSourceData: + raise NotImplementedError("create_kernel_instance not implemented") + + @abstractmethod + def check_argument_lists(self, kernel_name, arguments): + raise NotImplementedError("check_argument_lists not implemented") + + + +''' +class KernelSource: + def __init__(self, kernel_name, kernel_sources, lang, defines=None): + if not isinstance(kernel_sources, list): + kernel_sources = [kernel_sources] + + self.kernel_sources = kernel_sources + self.kernel_name = kernel_name + self.defines = defines + + if lang is None: + if callable(self.kernel_sources[0]): + raise TypeError("Please specify language when using a code generator function") + kernel_string = self.get_kernel_string(0) + self.lang = util.detect_language(kernel_string) + else: + self.lang = lang + + @abstractmethod + def prepare_kernel_instance(self, kernel_options, params, grid, threads) -> PreparedKernelSourceData: + raise NotImplementedError("create_kernel_instance not implemented") + + @abstractmethod + def check_argument_lists(self, kernel_name, arguments): + raise NotImplementedError("check_argument_lists not implemented") + +''' \ No newline at end of file diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py new file mode 100644 index 000000000..5df6fcb94 --- /dev/null +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -0,0 +1,160 @@ +import inspect +import ast +import copy +import uuid +import sys + +import astor +import tempfile +import importlib.util + +from typing import Any + +from kernel_tuner.kernel_sources.kernel_source import KernelSource +from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData + + +class KernelSourceFn(KernelSource): + + def __init__(self, kernel_name, kernel_source, lang, defines=None): + super().__init__(kernel_name, kernel_source, lang, defines) + if isinstance(kernel_source, list): + raise ValueError("KernelSourceFn only supports a single kernel source function") + + self.source_kernel_fn = kernel_source + self.kernel_fn = self.source_kernel_fn + self.source = inspect.getsource(kernel_source) + self.source_tree = ast.parse(self.source) + self.triton_import_nodes = [ + ast.Import(names=[ast.alias(name='triton', asname=None)]), + ast.ImportFrom( + module='triton', + names=[ast.alias(name='language', asname='tl')], + level=0 + ) + ] + + # Find the module where the kernel function is defined + self.module = inspect.getmodule(kernel_source) + # Get dependencies by analyzing the AST + self.dependencies = self._find_triton_dependencies() + + def _find_function_dependencies(self): + """Find all function calls in the kernel.""" + class FunctionCallVisitor(ast.NodeVisitor): + def __init__(self): + self.called_functions = set() + + def visit_Call(self, node): + if isinstance(node.func, ast.Name): + self.called_functions.add(node.func.id) + self.generic_visit(node) + + visitor = FunctionCallVisitor() + visitor.visit(self.source_tree) + return visitor.called_functions + + def _is_triton_jit_function(self, node): + """Check if a function has the @triton.jit decorator.""" + if not isinstance(node, ast.FunctionDef): + return False + + for decorator in node.decorator_list: + if (isinstance(decorator, ast.Name) and decorator.id == 'jit' or + isinstance(decorator, ast.Attribute) and decorator.attr == 'jit' or + isinstance(decorator, ast.Call) and + isinstance(decorator.func, ast.Attribute) and + decorator.func.attr == 'jit'): + return True + return False + + def _find_triton_dependencies(self): + """Find all Triton JIT functions that are called by the kernel.""" + dependencies = {} + called_functions = self._find_function_dependencies() + + # Parse the entire module to find Triton JIT functions + module_source = inspect.getsource(self.module) + module_tree = ast.parse(module_source) + + for node in module_tree.body: + if (self._is_triton_jit_function(node) and + node.name in called_functions and + node.name != self.kernel_name): # Skip the main kernel itself + dependencies[node.name] = node + + return dependencies + + def prepare_kernel_instance(self, kernel_options, params, grid, threads): + new_kernel_fn, temp_file_path = self.apply_params_to_source_fn(params) + self.kernel_fn = new_kernel_fn + + return PreparedKernelSourceData( + temp_files=[temp_file_path], + kernel_name=self.kernel_name, + kernel_fn=new_kernel_fn, + kernel_str=None + ) + + def check_argument_lists(self, kernel_name, arguments): + return True + + def apply_params_to_source_fn(self, params): + transformer = ReplaceVars(params) + + # Create a new module with all necessary functions + new_module = ast.Module(body=[], type_ignores=[]) + + # Add both triton imports + new_module.body.extend(self.triton_import_nodes) + + # Add all Triton JIT dependencies first + for dep_node in self.dependencies.values(): + new_module.body.append(copy.deepcopy(dep_node)) + + # Add transformed main kernel + source_tree_copy = copy.deepcopy(self.source_tree) + transformed_tree = transformer.visit(source_tree_copy) + new_module.body.extend(transformed_tree.body) + + # Fix locations and generate source + ast.fix_missing_locations(new_module) + new_source = astor.to_source(new_module) + + # Create a unique module name + module_name = f'temp_kernel_module_{uuid.uuid4().hex}' + + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file: + temp_file.write(new_source) + temp_file_path = temp_file.name + + spec = importlib.util.spec_from_file_location(module_name, temp_file_path) + temp_module = importlib.util.module_from_spec(spec) + + # Register the module in sys.modules before executing it + sys.modules[module_name] = temp_module + spec.loader.exec_module(temp_module) + new_fn = getattr(temp_module, self.kernel_name) + + return new_fn, temp_file_path + + def __del__(self): + # Clean up temporary modules when the instance is destroyed + for key in list(sys.modules.keys()): + if key.startswith('temp_kernel_module_'): + del sys.modules[key] + + +class ReplaceVars(ast.NodeTransformer): + + def __init__(self, params: dict): + self.params = params + + def visit_Name(self, node: ast.Name) -> Any: + if isinstance(node.ctx, ast.Load) and node.id in self.params.keys(): + return ast.copy_location( + ast.Constant(value=self.params[node.id]), + node + ) + + return node \ No newline at end of file diff --git a/kernel_tuner/kernel_sources/kernel_source_str.py b/kernel_tuner/kernel_sources/kernel_source_str.py new file mode 100644 index 000000000..f470b939f --- /dev/null +++ b/kernel_tuner/kernel_sources/kernel_source_str.py @@ -0,0 +1,174 @@ +import kernel_tuner.util as util +import logging + +from kernel_tuner.kernel_sources.kernel_source import KernelSource +from kernel_tuner.core import wrap_templated_kernel +from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData +from kernel_tuner.language import Language + + +class KernelSourceStr(KernelSource): + """Class that holds the kernel sources. + + There is a primary kernel source for string-based kernels., which can be either a source string, + a filename (indicating a file containing the kernel source code), + or a callable (generating the kernel source code). + There can additionally be (one or multiple) secondary kernel sources, which + must be filenames. + """ + + def __init__(self, kernel_name, kernel_sources, lang, defines=None): + super().__init__(kernel_name, kernel_sources, lang, defines) + + def prepare_kernel_instance(self, kernel_options, params, grid, threads): + name, kernel_string, temp_files = self.prepare_list_of_files( + kernel_name=kernel_options.kernel_name, + params=params, + grid=grid, + threads=threads, + block_size_names=kernel_options.block_size_names, + ) + + lang_is_cuda = self.lang in [Language.CUDA, Language.NVCUDA] + + if lang_is_cuda and "<" in name and ">" in name: + kernel_string, name = wrap_templated_kernel(kernel_string, name) + + return PreparedKernelSourceData( + temp_files=temp_files, + kernel_name=name, + kernel_str=kernel_string, + kernel_fn=None, + ) + + + + def get_kernel_string(self, index=0, params=None): + """Retrieve the kernel source with the given index and return as a string. + + See util.get_kernel_string() for details. + + :param index: Index of the kernel source in the list of sources. + :type index: int + + :param params: Dictionary containing the tunable parameters for this specific + kernel instance, only needed when kernel_source is a generator. + :type param: dict + + :returns: A string containing the kernel code. + :rtype: string + """ + logging.debug("get_kernel_string called") + + if hasattr(self, 'lang') and self.lang.upper() == "HYPERTUNER": + return "" + + kernel_source = self.kernel_sources[index] + return util.get_kernel_string(kernel_source, params) + + def prepare_list_of_files( + self, kernel_name, params, grid, threads, block_size_names + ): + """Prepare the kernel string along with any additional files. + + The first file in the list is allowed to include or read in the others + The files beyond the first are considered additional files that may also contain tunable parameters + + For each file beyond the first this function creates a temporary file with + preprocessors statements inserted. Occurrences of the original filenames in the + first file are replaced with their temporary counterparts. + + :param kernel_name: A string specifying the kernel name. + :type kernel_name: string + + :param params: A dictionary with the tunable parameters for this particular + instance. + :type params: dict() + + :param grid: The grid dimensions for this instance. The grid dimensions are + also inserted into the code as if they are tunable parameters for + convenience. + :type grid: tuple() + + :param threads: The thread block dimensions for this instance. The thread block are + also inserted into the code as if they are tunable parameters for + convenience. + :type threads: tuple() + + :param block_size_names: A list of strings that denote the names + for the thread block dimensions. + :type block_size_names: list(string) + + """ + temp_files = dict() + + if self.lang.upper() == "HYPERTUNER": + return tuple(["", "", temp_files]) + + for i, f in enumerate(self.kernel_sources): + if i > 0 and not util.looks_like_a_filename(f): + raise ValueError("When passing multiple kernel sources, the secondary entries must be filenames") + + ks = self.get_kernel_string(i, params) + # add preprocessor statements + n, ks = util.prepare_kernel_string( + kernel_name, + ks, + params, + grid, + threads, + block_size_names, + self.lang, + self.defines, + ) + + if i == 0: + # primary kernel source + name = n + kernel_string = ks + continue + + # save secondary kernel sources to temporary files + + # generate temp filename with the same extension + temp_file = util.get_temp_filename(suffix="." + f.split(".")[-1]) + temp_files[f] = temp_file + util.write_file(temp_file, ks) + # replace occurrences of the additional file's name in the first kernel_string with the name of the temp file + kernel_string = kernel_string.replace(f, temp_file) + + return name, kernel_string, temp_files + + def get_user_suffix(self, index=0): + """Get the suffix of the kernel filename, if the user specified one. Return None otherwise.""" + if util.looks_like_a_filename(self.kernel_sources[index]) and ("." in self.kernel_sources[index]): + return "." + self.kernel_sources[index].split(".")[-1] + return None + + def get_suffix(self, index=0): + """Return a suitable suffix for a kernel filename. + + This uses the user-specified suffix if available, or one based on the + lang/backend otherwise. + """ + # TODO: Consider delegating this to the backend + suffix = self.get_user_suffix(index) + if suffix is not None: + return suffix + + _suffixes = {"CUDA": ".cu", "OpenCL": ".cl", "C": ".c"} + try: + return _suffixes[self.lang] + except KeyError: + return ".c" + + def check_argument_lists(self, kernel_name, arguments): + """Check if the kernel arguments have the correct types. + + This is done by calling util.check_argument_list on each kernel string. + """ + for i, f in enumerate(self.kernel_sources): + if not callable(f): + util.check_argument_list(kernel_name, self.get_kernel_string(i), arguments) + else: + logging.debug("Checking of arguments list not supported yet for code generators.") \ No newline at end of file diff --git a/kernel_tuner/kernel_sources/model/__init__.py b/kernel_tuner/kernel_sources/model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/kernel_tuner/kernel_sources/model/prepared_kernel_source_data.py b/kernel_tuner/kernel_sources/model/prepared_kernel_source_data.py new file mode 100644 index 000000000..b9c904b38 --- /dev/null +++ b/kernel_tuner/kernel_sources/model/prepared_kernel_source_data.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass +from typing import Any + + +@dataclass +class PreparedKernelSourceData: + temp_files: Any + kernel_name: str + kernel_str: Any + kernel_fn: Any \ No newline at end of file diff --git a/kernel_tuner/kernelbuilder.py b/kernel_tuner/kernelbuilder.py index 0f3f6154f..e2039e66d 100644 --- a/kernel_tuner/kernelbuilder.py +++ b/kernel_tuner/kernelbuilder.py @@ -4,6 +4,7 @@ from kernel_tuner.interface import Options, _kernel_options from kernel_tuner.integration import TuneResults +from kernel_tuner.kernel_sources.kernel_source import KernelSource class PythonKernel(object): @@ -30,7 +31,7 @@ def __init__(self, kernel_name, kernel_string, problem_size, arguments, params=N """ #construct device interface - kernel_source = core.KernelSource(kernel_name, kernel_string, lang) + kernel_source = KernelSource(kernel_name, kernel_string, lang) self.dev = core.DeviceInterface(kernel_source, device=device, quiet=True) if not params: params = {} @@ -92,7 +93,7 @@ def run_kernel(self, args): :type args: list(np.ndarray or np.generic) """ self.update_gpu_args(args) - self.dev.run_kernel(self.func, self.gpu_args, self.kernel_instance) + self.dev.run_kernel_check(self.func, self.gpu_args, self.kernel_instance) return self.get_gpu_result(args) def __call__(self, *args): diff --git a/kernel_tuner/language.py b/kernel_tuner/language.py new file mode 100644 index 000000000..ad510b55b --- /dev/null +++ b/kernel_tuner/language.py @@ -0,0 +1,10 @@ +from enum import Enum + +class Language(Enum): + CUDA = "CUDA" + OpenCL = "OPENCL" + C = "C" + HIP = "HIP" + TRITON = "TRITON" + FORTRAN = "FORTRAN" + NVCUDA = "NVCUDA" \ No newline at end of file diff --git a/kernel_tuner/observers/triton.py b/kernel_tuner/observers/triton.py new file mode 100644 index 000000000..878e8f27c --- /dev/null +++ b/kernel_tuner/observers/triton.py @@ -0,0 +1,32 @@ +import numpy as np + +from kernel_tuner.observers.observer import BenchmarkObserver + +try: + import torch +except (ImportError, RuntimeError): + torch = None + + +class TritonRuntimeObserver(BenchmarkObserver): + """Observer that measures time using CUDA events during benchmarking.""" + + def __init__(self, dev): + if torch is None: + raise ImportError("Unable to load torch") + + self.dev = dev + self.stream = dev.stream + self.start = dev.start + self.end = dev.end + self.times = [] + + def after_finish(self): + # Time is measured in milliseconds + event_elapsed_time = self.start.elapsed_time(self.end) + self.times.append(event_elapsed_time) + + def get_results(self): + results = {"time": np.average(self.times), "times": self.times.copy()} + self.times = [] + return results \ No newline at end of file diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index 2d9e3f1b3..fd4cb5697 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -136,6 +136,9 @@ def check_argument_type(dtype, kernel_argument): def check_argument_list(kernel_name, kernel_string, args): """Raise an exception if kernel arguments do not match host arguments.""" + if kernel_string is None: + return + kernel_arguments = list() collected_errors = list() @@ -162,7 +165,7 @@ def check_argument_list(kernel_name, kernel_string, args): if not isinstance(arg, (np.ndarray, np.generic, cp.ndarray, torch.Tensor, DeviceArray)): raise TypeError( f"Argument at position {i} of type: {type(arg)} should be of type " - "np.ndarray, numpy scalar, or HIP Python DeviceArray type" +- "np.ndarray, numpy scalar, or HIP Python DeviceArray type" ) correct = True @@ -415,6 +418,8 @@ def detect_language(kernel_string): lang = "CUDA" elif "__kernel" in kernel_string: lang = "OpenCL" + elif "@triton.jit" in kernel_string: + lang = "Triton" else: lang = "C" return lang From 68d43ec4540b24fa13358fe717a739fe5fa69182 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Sat, 20 Dec 2025 12:51:19 +0100 Subject: [PATCH 02/44] Started tilus support --- examples/tilus/vec_add.py | 67 ++++++ kernel_tuner/backends/tilus.py | 194 ++++++++++++++++++ kernel_tuner/core.py | 8 + kernel_tuner/kernel_sources/kernel_source.py | 13 +- .../kernel_sources/kernel_source_fn.py | 49 +++-- kernel_tuner/language.py | 10 +- kernel_tuner/observers/generic_python.py | 33 +++ 7 files changed, 356 insertions(+), 18 deletions(-) create mode 100644 examples/tilus/vec_add.py create mode 100644 kernel_tuner/backends/tilus.py create mode 100644 kernel_tuner/observers/generic_python.py diff --git a/examples/tilus/vec_add.py b/examples/tilus/vec_add.py new file mode 100644 index 000000000..60dcc6074 --- /dev/null +++ b/examples/tilus/vec_add.py @@ -0,0 +1,67 @@ +import tilus +from tilus import float32, int32 +from tilus.utils import cdiv, benchmark_func +import torch + + +class VecAddV(tilus.Script): + def __init__(self): + super().__init__() + self.block_size = 256 # number of threads per block + + def __call__( + self, + n_size: int32, # size of the vectors + a_ptr: ~float32, # input vector A + b_ptr: ~float32, # input vector B + c_ptr: ~float32 # output vector C + ): + # compute the number of blocks needed + self.attrs.blocks = [cdiv(n_size, self.block_size)] + self.attrs.warps = 1 # number of warps per block + + # calculate the offset for this block + offset: int32 = self.block_size * self.blockIdx.x + + # create global views for input/output vectors + ga = self.global_view(a_ptr, dtype=float32, shape=[n_size]) + gb = self.global_view(b_ptr, dtype=float32, shape=[n_size]) + gc = self.global_view(c_ptr, dtype=float32, shape=[n_size]) + + # load a block of A and B into registers + a = self.load_global(ga, offsets=[offset], shape=[self.block_size]) + b = self.load_global(gb, offsets=[offset], shape=[self.block_size]) + + # perform element-wise addition + c = a + b + + # store the result back to global memory + self.store_global(gc, c, offsets=[offset]) + + +def main(): + N = 1 << 20 # vector size + vecadd = VecAddV() + + a = torch.rand(N, dtype=torch.float32).cuda() + b = torch.rand(N, dtype=torch.float32).cuda() + c_actual = torch.empty_like(a) + c_expect = a + b + + torch.cuda.synchronize() + vecadd(N, a, b, c_actual) + torch.cuda.synchronize() + + # correctness check + torch.testing.assert_close(c_expect, c_actual, atol=1e-6, rtol=1e-6) + + # benchmark + for name, func in [ + ("torch", lambda: a + b), + ("tilus", lambda: vecadd(N, a, b, c_actual)), + ]: + latency = benchmark_func(func, warmup=5, repeat=20) + print(f"{name} latency: {latency:.3f} ms") + +if __name__ == "__main__": + main() diff --git a/kernel_tuner/backends/tilus.py b/kernel_tuner/backends/tilus.py new file mode 100644 index 000000000..d3782fb0d --- /dev/null +++ b/kernel_tuner/backends/tilus.py @@ -0,0 +1,194 @@ +import logging +import numpy as np +import inspect + +from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.observers.generic_python import GenericPythonRuntimeObserver + +try: + import torch +except ImportError: + logging.error("Torch not available") + +try: + import tilus +except ImportError: + tilus = None + logging.error("Unable to load Tilus") + + +class TilusFunctions(GPUBackend): + + def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): + ''' + In here, everyting is generic if the language uses Torch as backend + ''' + if not tilus or not torch: + logging.error("Tilus or Torch not available") + raise ImportError("Tilus or Torch not available") + + self.device_id = torch.cuda.current_device() + + self.device_properties = torch.cuda.get_device_properties(self.device_id) + self.name = torch.cuda.get_device_name(self.device_id) + self.max_threads = self.device_properties.max_threads_per_multi_processor + + env = dict() + env["device_name"] = self.name + env["max_threads"] = self.max_threads + env["iterations"] = iterations + env["compiler_options"] = compiler_options + self.env = env + + self.stream = torch.cuda.default_stream() + self.start = torch.cuda.Event(enable_timing=True) + self.end = torch.cuda.Event(enable_timing=True) + + # setup observers + self.observers = observers or [] + self.observers.append(GenericPythonRuntimeObserver(self)) + for obs in self.observers: + obs.register_device(self) + + self.units = {"time": "ms", "power": "s,mW", "energy": "J"} + + super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) + + def ready_argument_list(self, arguments): + ''' + This seems to work for Torch as backend. However, another idea for generic would be + to let the user provide two argument lists: one with numpy and one with the actual arguments + that will be passed to the kernel. TODO: why do we need the numpy list anyway? + ''' + torch_args = [] + + for arg in arguments: + if isinstance(arg, torch.Tensor) and arg.dim() > 0: + torch_args.append(arg.cuda()) + elif isinstance(arg, torch.Tensor) and arg.dim() == 0: + scalar_value = arg.item() + torch_args.append(scalar_value) + elif isinstance(arg, np.ndarray): + torch_arg = torch.from_numpy(arg) + torch_arg_gpu = torch_arg.cuda() + torch_args.append(torch_arg_gpu) + elif isinstance(arg, np.generic): + scalar_value = arg.item() + torch_args.append(scalar_value) + else: + logging.warning("Unknown instance in Tilus functions") + + return torch_args + + def compile(self, kernel_instance, gpu_args=None): + logging.debug("Compiling Tilus kernel") + + if kernel_instance.kernel_fn is None: + raise ValueError("kernel_fn is None, currently Tilus only supports callable kernel_source") + + ''' + if gpu_args is None: + raise ValueError("gpu_args is None, Triton needs gpu args to compile the kernel") + TODO: find out about these GPU args + ''' + + grid = kernel_instance.grid # TODO might need to use this? + threads = kernel_instance.threads + params = kernel_instance.params + kernel_function = kernel_instance.kernel_fn + gpu_kwargs = self.build_gpu_kwargs(kernel_function, threads, params) + + # Call the jit function in order to compile it + kernel_function(*gpu_args, **gpu_kwargs) + + return kernel_function + + def start_event(self): + logging.debug("Start triton event") + self.start.record() + + def stop_event(self): + logging.debug("Stop triton event") + self.end.record() + + def kernel_finished(self): + logging.debug("Checking if kernel has finished") + return self.end.query() + + def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): + ''' + if params is None: + raise ValueError("params is None, Triton needs params in order to set num_warps, num_ctas, etc.") + ''' + + # Run the kernel + if stream is None: + stream = self.stream + + gpu_kwargs = self.build_gpu_kwargs(func, threads, params) + + with torch.cuda.stream(stream): + logging.debug("Running Tilus kernel") + func(*gpu_args, **gpu_kwargs) + + + + def build_gpu_kwargs(self, kernel_function, threads, params=None): + gpu_kwargs = {} + + ''' + if 'BLOCK_SIZE' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE'] = threads[0] + + if 'BLOCK_SIZE_X' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE_X'] = threads[0] + + if 'BLOCK_SIZE_Y' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE_Y'] = threads[1] + + if 'BLOCK_SIZE_Z' in jit_fn.arg_names: + gpu_kwargs['BLOCK_SIZE_Z'] = threads[2] + ''' + + if params is None: + return gpu_kwargs + + + sig = inspect.signature(kernel_function) + for name, p in sig.parameters.items(): + if name in params: + gpu_kwargs[name] = params[name] + + ''' + # Check for Triton specific parameters + if 'num_warps' in params: + gpu_kwargs['num_warps'] = params['num_warps'] + if 'num_ctas' in params: + gpu_kwargs['num_ctas'] = params['num_ctas'] + if 'num_stages' in params: + gpu_kwargs['num_stages'] = params['num_stages'] + ''' + + return gpu_kwargs + + + def synchronize(self): + torch.cuda.synchronize() + + def memset(self, allocation, value, size): + pass + + def memcpy_dtoh(self, dest, src): + pass + + def memcpy_htod(self, dest, src): + pass + + def copy_constant_memory_args(self, cmem_args): + raise NotImplementedError("Tilus does not support constant memory") + + def copy_shared_memory_args(self, smem_args): + raise NotImplementedError("Tilus does not support shared memory") + + def copy_texture_memory_args(self, texmem_args): + raise NotImplementedError("Tilus does not support texture memory") \ No newline at end of file diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 1a7e67b34..3b545d459 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -22,6 +22,7 @@ from kernel_tuner.backends.opencl import OpenCLFunctions from kernel_tuner.backends.pycuda import PyCudaFunctions from kernel_tuner.backends.triton import TritonFunctions +from kernel_tuner.backends.tilus import TilusFunctions from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.observers.nvml import NVMLObserver from kernel_tuner.observers.observer import ContinuousObserver, OutputObserver, PrologueObserver @@ -191,6 +192,13 @@ def __init__( iterations=iterations, observers=observers ) + elif lang.upper() == "GENERIC_PYTHON": + dev = TilusFunctions( + device, + compiler_options=compiler_options, + iterations=iterations, + observers=observers + ) else: raise NotImplementedError( "Sorry, support for languages other than CUDA, OpenCL, HIP, C, Triton and Fortran is not implemented yet" diff --git a/kernel_tuner/kernel_sources/kernel_source.py b/kernel_tuner/kernel_sources/kernel_source.py index be8bbb265..70c94903d 100644 --- a/kernel_tuner/kernel_sources/kernel_source.py +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -3,6 +3,7 @@ from abc import abstractmethod +from kernel_tuner.language import Language from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData @@ -11,8 +12,16 @@ class KernelSource: def __new__(cls, kernel_name, kernel_sources, lang, defines=None): """Factory behavior""" + if lang == None: + language = None + else: + try: + language = Language(lang) + except ValueError: + raise TypeError(f"Supported languages are {[l.value for l in Language]}") + if cls is KernelSource: - if inspect.isfunction(kernel_sources) and (lang and lang.upper() == "TRITON"): # TODO should this be isfunction? + if inspect.isfunction(kernel_sources) and (language and (language == Language.TRITON or language == Language.GENERIC_PYTHON)): # TODO should this be isfunction? from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn print("CREATING KSFN") return KernelSourceFn(kernel_name, kernel_sources, lang, defines) @@ -20,9 +29,11 @@ def __new__(cls, kernel_name, kernel_sources, lang, defines=None): from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr print("CREATING KSSTR") return KernelSourceStr(kernel_name, kernel_sources, lang, defines) + # otherwise, normal subclass init return super().__new__(cls) + def __init__(self, kernel_name, kernel_sources, lang, defines=None): if not isinstance(kernel_sources, list): kernel_sources = [kernel_sources] diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 5df6fcb94..223e0e84f 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -10,10 +10,12 @@ from typing import Any +from kernel_tuner.language import Language from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData + class KernelSourceFn(KernelSource): def __init__(self, kernel_name, kernel_source, lang, defines=None): @@ -21,23 +23,35 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None): if isinstance(kernel_source, list): raise ValueError("KernelSourceFn only supports a single kernel source function") + + try: + self.lang = Language(lang) + except ValueError: + raise TypeError(f"Supported languages are {[l.value for l in Language]}") self.source_kernel_fn = kernel_source self.kernel_fn = self.source_kernel_fn self.source = inspect.getsource(kernel_source) self.source_tree = ast.parse(self.source) - self.triton_import_nodes = [ - ast.Import(names=[ast.alias(name='triton', asname=None)]), - ast.ImportFrom( - module='triton', - names=[ast.alias(name='language', asname='tl')], - level=0 - ) - ] + if self.lang == Language.TRITON: + self.import_nodes = [ + ast.Import(names=[ast.alias(name='triton', asname=None)]), + ast.ImportFrom( + module='triton', + names=[ast.alias(name='language', asname='tl')], + level=0 + ) + ] + else: + self.import_nodes = [n for n in self.source_tree.body if isinstance(n, (ast.Import, ast.ImportFrom))] + # Find the module where the kernel function is defined self.module = inspect.getmodule(kernel_source) # Get dependencies by analyzing the AST - self.dependencies = self._find_triton_dependencies() + if self.lang == Language.lang: + self.dependencies = self._find_triton_dependencies() + else: + self.dependencies = None # TODO def _find_function_dependencies(self): """Find all function calls in the kernel.""" @@ -55,7 +69,7 @@ def visit_Call(self, node): return visitor.called_functions def _is_triton_jit_function(self, node): - """Check if a function has the @triton.jit decorator.""" + """Check if a function has the @triton.jit decorator. Is triton specific""" if not isinstance(node, ast.FunctionDef): return False @@ -69,7 +83,7 @@ def _is_triton_jit_function(self, node): return False def _find_triton_dependencies(self): - """Find all Triton JIT functions that are called by the kernel.""" + """Find all Triton JIT functions that are called by the kernel. Is Triton specific""" dependencies = {} called_functions = self._find_function_dependencies() @@ -105,12 +119,15 @@ def apply_params_to_source_fn(self, params): # Create a new module with all necessary functions new_module = ast.Module(body=[], type_ignores=[]) - # Add both triton imports - new_module.body.extend(self.triton_import_nodes) + # Add imports + new_module.body.extend(self.import_nodes) - # Add all Triton JIT dependencies first - for dep_node in self.dependencies.values(): - new_module.body.append(copy.deepcopy(dep_node)) + if self.lang == Language.TRITON: + # Add all Triton JIT dependencies first + for dep_node in self.dependencies.values(): + new_module.body.append(copy.deepcopy(dep_node)) + + # TODO the kernel can not have dependencies yet. Obviously this needs fixing. # Add transformed main kernel source_tree_copy = copy.deepcopy(self.source_tree) diff --git a/kernel_tuner/language.py b/kernel_tuner/language.py index ad510b55b..d33933cca 100644 --- a/kernel_tuner/language.py +++ b/kernel_tuner/language.py @@ -7,4 +7,12 @@ class Language(Enum): HIP = "HIP" TRITON = "TRITON" FORTRAN = "FORTRAN" - NVCUDA = "NVCUDA" \ No newline at end of file + NVCUDA = "NVCUDA" + GENERIC_PYTHON = "GENERIC_PYTHON" + + @classmethod + def is_valid(cls, value: str) -> bool: + """ + Test if a language is valid in Kernel Tuner framework. + """ + return value in cls._value2member_map_ \ No newline at end of file diff --git a/kernel_tuner/observers/generic_python.py b/kernel_tuner/observers/generic_python.py new file mode 100644 index 000000000..896b4f9cc --- /dev/null +++ b/kernel_tuner/observers/generic_python.py @@ -0,0 +1,33 @@ +import numpy as np + +from kernel_tuner.observers.observer import BenchmarkObserver + +try: + import torch +except (ImportError, RuntimeError): + torch = None + + +class GenericPythonRuntimeObserver(BenchmarkObserver): + """Observer that measures time using CUDA events during benchmarking. + TODO think about: do we need torch? Does CUDA work for all languages?""" + + def __init__(self, dev): + if torch is None: + raise ImportError("Unable to load torch") + + self.dev = dev + self.stream = dev.stream + self.start = dev.start + self.end = dev.end + self.times = [] + + def after_finish(self): + # Time is measured in milliseconds + event_elapsed_time = self.start.elapsed_time(self.end) + self.times.append(event_elapsed_time) + + def get_results(self): + results = {"time": np.average(self.times), "times": self.times.copy()} + self.times = [] + return results \ No newline at end of file From 542f5bae3b2bf29df9926dc2a22d770d7b5e71a8 Mon Sep 17 00:00:00 2001 From: "I.C. van Ooijen" Date: Fri, 6 Feb 2026 15:20:16 +0100 Subject: [PATCH 03/44] Working implementation tested with basic kernels. Might need some refactoring. Definately needs comments --- .../generic_python/matmul/triton_matmul.py | 154 ++++++++++ examples/generic_python/numba_vec_add.py | 68 ++++ examples/generic_python/tilus_naive_matmul.py | 124 ++++++++ .../generic_python/tilus_splitk_matmul.py | 267 ++++++++++++++++ .../generic_python/tilus_tunable_precision.py | 120 ++++++++ .../tilus_vec_add.py} | 64 ++-- examples/generic_python/triton_vec_add.py | 76 +++++ examples/generic_python/warp_vec_add.py | 98 ++++++ examples/triton/vec_add.py | 46 +-- kernel_tuner/backends/generic_python.py | 290 ++++++++++++++++++ kernel_tuner/backends/pycuda.py | 2 + kernel_tuner/backends/tilus.py | 194 ------------ kernel_tuner/core.py | 149 +++++---- kernel_tuner/interface.py | 43 ++- kernel_tuner/kernel_sources/kernel_source.py | 86 +++--- .../kernel_sources/kernel_source_fn.py | 167 ++++++++-- .../kernel_sources/kernel_source_str.py | 6 +- kernel_tuner/language.py | 4 +- kernel_tuner/observers/generic_python.py | 3 +- kernel_tuner/util.py | 2 +- test/context.py | 9 + test/test_backend.py | 8 +- test/test_generic_python_functions.py | 116 +++++++ test/test_kernel_source_fn.py | 160 ++++++++++ 24 files changed, 1866 insertions(+), 390 deletions(-) create mode 100644 examples/generic_python/matmul/triton_matmul.py create mode 100644 examples/generic_python/numba_vec_add.py create mode 100644 examples/generic_python/tilus_naive_matmul.py create mode 100644 examples/generic_python/tilus_splitk_matmul.py create mode 100644 examples/generic_python/tilus_tunable_precision.py rename examples/{tilus/vec_add.py => generic_python/tilus_vec_add.py} (53%) create mode 100644 examples/generic_python/triton_vec_add.py create mode 100644 examples/generic_python/warp_vec_add.py create mode 100644 kernel_tuner/backends/generic_python.py delete mode 100644 kernel_tuner/backends/tilus.py create mode 100644 test/test_generic_python_functions.py create mode 100644 test/test_kernel_source_fn.py diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py new file mode 100644 index 000000000..7115ea5cc --- /dev/null +++ b/examples/generic_python/matmul/triton_matmul.py @@ -0,0 +1,154 @@ +import torch + +import triton +import triton.language as tl + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + +def get_cuda_autotune_config(): + return [ + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, + num_warps=8), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, + num_warps=2), + triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, + num_warps=2), + # Good config for fp8 inputs. + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, + num_warps=8), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, + num_warps=8), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4) + ] + +''' +@triton.autotune( + configs=get_cuda_autotune_config(), + key=['M', 'N', 'K'], +) +''' +@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, b_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, # + stride_bk, stride_bn, # + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # + GROUP_SIZE_M: tl.constexpr, # +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ----------------------------------------------------------- + # Add some integer bound assumptions. + # This helps to guide integer analysis in the backend to optimize + # load/store offset address calculation + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + # We accumulate along the K dimension. + accumulator = tl.dot(a, b, accumulator) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + c = accumulator.to(tl.float16) + + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def matmul(a, b): + # Check constraints. + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.is_contiguous(), "Matrix A must be contiguous" + M, K = a.shape + K, N = b.shape + # Allocates output. + c = torch.empty((M, N), device=a.device, dtype=torch.float16) + # 1D launch kernel where each block gets its own program. + grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), ) + matmul_kernel[grid]( + a, b, c, # + M, N, K, # + a.stride(0), a.stride(1), # + b.stride(0), b.stride(1), # + c.stride(0), c.stride(1) + ) + return c + + diff --git a/examples/generic_python/numba_vec_add.py b/examples/generic_python/numba_vec_add.py new file mode 100644 index 000000000..3732aa018 --- /dev/null +++ b/examples/generic_python/numba_vec_add.py @@ -0,0 +1,68 @@ +import numpy as np +from numba import cuda +from math import ceil +from kernel_tuner import tune_kernel, run_kernel + +#@cuda.jit +def f(a, b, c): + tid = cuda.grid(1) + size = len(c) + + if tid < size: + c[tid] = a[tid] + b[tid] + + +def call_numba(kernel_function, args, kwargs, grid, threads, params): + kernel_function[grid[0], threads[0]](*args, **kwargs) + + +def verify(answer, result_host, atol): + correct = True + for i, ans in enumerate(answer): + if ans is None: + continue + res = result_host[i].copy_to_host() + if not np.allclose(ans, res, atol=atol): + correct = False + + return correct + + +N = 100000000 +a = cuda.to_device(np.random.random(N)) +b = cuda.to_device(np.random.random(N)) +c = cuda.device_array_like(a) +c_expect = a.copy_to_host() + b.copy_to_host() + +args = [a, b, c] +tune_params = {"block_size_x": [2**i for i in range(10)]} + +results = run_kernel( + kernel_name="f", + kernel_source=f, + problem_size=N, + arguments=args, + params={"block_size_x": 32}, + lang="generic_python", + call_function=call_numba, + decorator="@cuda.jit" +) + + +print(np.allclose(results[2], c_expect)) + +''' +results, env = tune_kernel( + kernel_name="f", + kernel_source=f, + problem_size=N, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, c_expect], + verify=verify, + call_function=call_numba, + decorator="@cuda.jit" +) +''' + diff --git a/examples/generic_python/tilus_naive_matmul.py b/examples/generic_python/tilus_naive_matmul.py new file mode 100644 index 000000000..2f358ef48 --- /dev/null +++ b/examples/generic_python/tilus_naive_matmul.py @@ -0,0 +1,124 @@ +from tilus import float16, float32, int32 +from tilus.utils import cdiv +import tilus +from kernel_tuner import tune_kernel, run_kernel +import math +import torch + + +class MatmulV0(tilus.Script): + def __init__(self): + super().__init__() + # we define three hyperparameters: ``block_m``, ``block_n``, and ``block_k`` to determine the tile size on + # m, n, and k dimensions for each `thread block` of the kernel. + self.block_m = 64 + self.block_n = 64 + self.block_k = 16 + + def __call__( + self, + m_size: int32, # the size of the m dimension of the input matrix A and output matrix C + n_size: int, # the size of the n dimension of the input matrix B and output matrix C + k_size: int, # the size of the k dimension of the input matrix A and B + a_ptr: ~float16, # the pointer to the input matrix A, which is a 2D tensor of shape [m_size, k_size] + b_ptr: ~float16, # the pointer to the input matrix B, which is a 2D tensor of shape [k_size, n_size] + c_ptr: ~float16, # the pointer to the output matrix C, which is a 2D tensor of shape [m_size, n_size] + ): + self.attrs.blocks = [ + cdiv(m_size, self.block_m), # the x dimension size of the grid + cdiv(n_size, self.block_n), # the y dimension size of the grid + ] + self.attrs.warps = 1 # the number of warps per thread block, must be a compile-time known integer + + # define two int32 variables to store the offsets of the m and n dimensions for the current thread block. + offset_m: int32 = self.block_m * self.blockIdx.x + offset_n: int32 = self.block_n * self.blockIdx.y + + # create two global tensors `ga` and `gb` to represent the input matrices A and B, respectively. + ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) + gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) + + # create a register tensor `acc` to accumulate the results of the matrix multiplication. + acc = self.register_tensor( + dtype=float32, shape=[self.block_m, self.block_n], init=0.0 + ) + + # iterate over the k dimension in blocks of size `block_k`. + for k in range(cdiv(k_size, self.block_k)): + # calculate the offset for the current block in the k dimension + offset_k = k * self.block_k + + # load a block of matrix A and B into register tensors `a` and `b`. + a = self.load_global( + ga, offsets=[offset_m, offset_k], shape=[self.block_m, self.block_k] + ) + b = self.load_global( + gb, offsets=[offset_k, offset_n], shape=[self.block_k, self.block_n] + ) + + # perform the dot product: acc = a @ b + acc + self.dot(a, b, acc, out=acc) + + # after the loop, we cast the accumulated result `acc` to float16 type and store it back to the output matrix C. + acc_f16 = self.cast(acc, dtype=float16) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + self.store_global(gc, acc_f16, offsets=[offset_m, offset_n]) + + +def call_tilus(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + +def main(): + m, n, k = 4096, 4096, 4096 + + # create an instance of the kernel we have just defined + matmul = MatmulV0() + + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + ''' + torch.cuda.synchronize() + # launch the kernel by passing required arguments + matmul(m, n, k, a, b, c_actual) + torch.cuda.synchronize() + + # check correctness + torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) + ''' + + args = [m, n, k, a, b, c_actual] + tune_params = dict() + tune_params["block_m"] = [16, 32, 64, 128, 256] + tune_params["block_n"] = [16, 32, 64, 128, 256] + tune_params["block_k"] = [16, 32, 64, 128, 256] + + + restrictions = [ + "block_m * block_n <= 4096", + "block_m * block_k <= 2048", + "block_k * block_n <= 4096", + "block_k >= 16", + ] + + results, env = tune_kernel( + kernel_name="MatmulV0", + kernel_source=MatmulV0, + problem_size=[m, n], + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, None, None, None, c_expect.cpu()], + restrictions=restrictions, + block_size_names=["block_m", "block_n", "block_k"], + call_function=call_tilus, + strategy="simulated_annealing" + ) + + + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/generic_python/tilus_splitk_matmul.py b/examples/generic_python/tilus_splitk_matmul.py new file mode 100644 index 000000000..f0dff1ac3 --- /dev/null +++ b/examples/generic_python/tilus_splitk_matmul.py @@ -0,0 +1,267 @@ +import tilus +from tilus import float16, float32, int32 +from tilus.utils import cdiv, benchmark_func +import torch +import math +from kernel_tuner import tune_kernel, run_kernel + +''' +@tilus.autotune("num_warps", [4, 8]) +@tilus.autotune("block_m, block_n", [(128, 128), (128, 64), (64, 128), (32, 256)]) +@tilus.autotune("block_k", [16, 32]) +@tilus.autotune("num_stages", [3, 4, 5]) +@tilus.autotune("split_k_factor", [1, 4, 12, 16]) +''' +class MatmulV5(tilus.Script): + ''' + def __init__(self, block_m, block_n, block_k, num_warps, num_stages, split_k_factor): + super().__init__() + self.block_m = block_m + self.block_n = block_n + self.block_k = block_k + self.num_warps = num_warps + self.num_stages = num_stages + self.split_k_factor = split_k_factor + ''' + def __init__(self): + super().__init__() + self.block_m = 128 + self.block_n = 128 + self.block_k = 16 + self.num_warps = 4 + self.num_stages = 4 + self.split_k_factor = 4 + + def __call__( + self, + m_size: int32, + n_size: int, + k_size: int, + a_ptr: ~float16, + b_ptr: ~float16, + c_ptr: ~float16, + ): + self.attrs.blocks = [ + cdiv(m_size, self.block_m), + cdiv(n_size, self.block_n), + self.split_k_factor, + ] + self.attrs.warps = self.num_warps + + # the k_size for each thread block + block_k_size = ( + cdiv(cdiv(k_size, self.split_k_factor), self.block_k) * self.block_k + ) + start_offset_k = self.blockIdx.z * block_k_size + end_offset_k = min(start_offset_k + block_k_size, k_size) + + block_m, block_n, block_k = self.block_m, self.block_n, self.block_k + offset_m: int32 = block_m * self.blockIdx.x + offset_n: int32 = block_n * self.blockIdx.y + + ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) + gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) + sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) + sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) + acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) + + for stage in range(self.num_stages - 1): + offset_k = start_offset_k + stage * self.block_k + self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) + self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) + self.copy_async_commit_group() + + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + current_stage: int32 = 0 + preload_stage: int32 = self.num_stages - 1 + for offset_k in self.range( + start_offset_k, end_offset_k, block_k, unroll=self.num_stages + ): + # computation for current tile + a = self.load_shared(sa[current_stage]) + b = self.load_shared(sb[current_stage]) + self.dot(a, b, acc, out=acc) + + # preload the next tile of A and B into shared memory + preload_offset_k = offset_k + (self.num_stages - 1) * block_k + if preload_offset_k < end_offset_k: + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) + self.copy_async_commit_group() + + # update the stage + current_stage = (current_stage + 1) % self.num_stages + preload_stage = (preload_stage + 1) % self.num_stages + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + # free the shared memory tensors for A and B + self.free_shared(sa) + self.free_shared(sb) + + # cast the accumulator to float16 and change the register tensor's layout + sc = self.shared_tensor(dtype=float16, shape=[block_m, block_n]) + casted_acc = self.cast(acc, dtype=float16) + self.store_shared(sc, casted_acc) + self.sync() + rc = self.load_shared(sc) + self.free_shared(sc) + + m_blocks, n_blocks = cdiv(m_size, block_m), cdiv(n_size, block_n) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + if self.split_k_factor == 0: + self.store_global(gc, rc, offsets=[offset_m, offset_n]) + else: + semaphores = self.global_tensor( + dtype=int32, shape=[m_blocks, n_blocks], requires_clean=True + ) + semaphore: ~int32 = ~semaphores[self.blockIdx.x, self.blockIdx.y] + + # load and accumulate the partial result in global memory + if self.blockIdx.z > 0: + self.lock_semaphore(semaphore, value=self.blockIdx.z) + partial_rc = self.load_global( + gc, offsets=[offset_m, offset_n], shape=[block_m, block_n] + ) + self.add(rc, partial_rc, out=rc) + + # store the result to global memory and release the semaphore + self.store_global(gc, rc, offsets=[offset_m, offset_n]) + + # release the semaphore + self.sync() # we need to make sure the previous store_global is finished + self.release_semaphore( + semaphore, value=(self.blockIdx.z + 1) % self.split_k_factor + ) + + +def call_tilus(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + + +def without_kernel_tuner(): + tilus.option.clear_cache = True + + tilus.option.verbose_autotune = True + m, n, k = 4096, 4096, 4096 + + # create an instance of the kernel we have just defined + matmul = MatmulV5() + + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + + torch.cuda.synchronize() + # launch the kernel by passing required arguments + matmul(m, n, k, a, b, c_actual) + torch.cuda.synchronize() + + # check correctness + torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) + + + import pandas + rows = [] + headers = ["m", "n", "k", "name", "latency (ms)", "tflops"] + # benchmark + for name, func in [ + ("torch", lambda: torch.matmul(a, b, out=c_expect)), + ("tilus", lambda: matmul(m, n, k, a, b, c_actual)), + ]: + latency = benchmark_func(func, warmup=5, repeat=20) + tflops = 2 * m * n * k / latency * 1e-9 + rows.append([m, n, k, name, latency, tflops]) + + df = pandas.DataFrame(rows, columns=headers) + print(df) + + + +#best performing configuration: +#block_m=128, block_n=64, block_k=16, num_warps=4, num_stages=4, split_k_factor=1, time=2.027ms + +def main(): + m, n, k = 4096, 4096, 4096 + + # create an instance of the kernel we have just defined + #matmul = MatmulV5() + + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + ''' + torch.cuda.synchronize() + # launch the kernel by passing required arguments + matmul(m, n, k, a, b, c_actual) + torch.cuda.synchronize() + + # check correctness + torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) + ''' + + args = [m, n, k, a, b, c_actual] + tune_params = dict() + tune_params["block_m"] = [32, 64, 128] #[16, 32, 64, 128, 256] + tune_params["block_n"] = [32, 64, 128] #[16, 32, 64, 128, 256] + tune_params["block_k"] = [16, 32] #[16, 32, 64, 128, 256] + tune_params["num_warps"] = [4, 8] #[2, 4, 8, 16] + tune_params["num_stages"] = [3, 4, 5] #[2, 3, 4, 5, 6] + tune_params["split_k_factor"] = [1, 4, 12, 16] #[1, 4, 12, 16, 20] + +#@tilus.autotune("num_warps", [4, 8]) +#@tilus.autotune("block_m, block_n", [(128, 128), (128, 64), (64, 128), (32, 256)]) +#@tilus.autotune("block_k", [16, 32]) +#@tilus.autotune("num_stages", [3, 4, 5]) +#@tilus.autotune("split_k_factor", [1, 4, 12, 16]) + + ''' + restrictions = [ + "block_m * block_n <= 4096", + "block_m * block_k <= 2048", + "block_k * block_n <= 4096", + "block_k >= 16", + "2 * (num_stages * block_k * (block_m + block_n) + block_m * block_n) <= 65536", # shared mem + "num_warps * 32 <= 1024", + "block_k * split_k_factor <= 4096", + ] + ''' + restrictions = ["block_m * block_n >= 8192", "block_m * block_n <= 16384"] + + + + results, env = tune_kernel( + kernel_name="MatmulV5", # This has to be a string of the actual name. TODO is this always the case? + kernel_source=MatmulV5, + problem_size=[m, n], + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, None, None, None, c_expect.cpu()], + atol=1e-2, + restrictions=restrictions, + block_size_names=["block_m", "block_n", "block_k"], + call_function=call_tilus, + #strategy="random_sample", + ) + + + + +if __name__ == "__main__": + main() + #without_kernel_tuner() \ No newline at end of file diff --git a/examples/generic_python/tilus_tunable_precision.py b/examples/generic_python/tilus_tunable_precision.py new file mode 100644 index 000000000..1cc4a7499 --- /dev/null +++ b/examples/generic_python/tilus_tunable_precision.py @@ -0,0 +1,120 @@ +import tilus +import hidet +from hidet import float32, float16, int32 +#from tilus import float32, int32, float16 +from tilus.utils import cdiv, benchmark_func +import torch +from kernel_tuner import tune_kernel, run_kernel +from kernel_tuner.accuracy import Tunable, AccuracyObserver +import numpy as np + +INPUT_TYPE = float32 +OUTPUT_TYPE = float32 + +class VecAddV(tilus.Script): + def __init__(self): + super().__init__() + self.block_size_x = 32 # number of threads per block + + def __call__( + self, + n_size: int32, # size of the vectors + a_ptr: ~INPUT_TYPE, # input vector A + b_ptr: ~INPUT_TYPE, # input vector B + c_ptr: ~OUTPUT_TYPE # output vector C + ): + + # compute the number of blocks needed + self.attrs.blocks = [cdiv(n_size, self.block_size_x)] + self.attrs.warps = 1 # number of warps per block + + # calculate the offset for this block + offset: int32 = self.block_size_x * self.blockIdx.x + + # create global views for input/output vectors + ga = self.global_view(a_ptr, dtype=INPUT_TYPE, shape=[n_size]) + gb = self.global_view(b_ptr, dtype=INPUT_TYPE, shape=[n_size]) + gc = self.global_view(c_ptr, dtype=OUTPUT_TYPE, shape=[n_size]) + + a = self.load_global(ga, offsets=[offset], shape=[self.block_size_x]) + b = self.load_global(gb, offsets=[offset], shape=[self.block_size_x]) + c = a + b + self.store_global(gc, c, offsets=[offset]) + + +def call_tilus(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + + +def verify(answer, result_host, atol): + correct = True + for i, ans in enumerate(answer): + if ans is None: + continue + res = result_host[i].cpu() + if not torch.allcose(ans, res, atol=atol): + correct = False + + return correct + + + +def main(): + + size = 1024000 + + a_32 = torch.randn(size, dtype=torch.float32) + b_32 = torch.randn(size, dtype=torch.float32) + c_32 = torch.zeros_like(b_32) + c_expect = a_32 + b_32 + + a_16 = a_32.to(torch.float16) + b_16 = b_32.to(torch.float16) + c_16 = c_32.to(torch.float16) + + + tune_params = dict() + tune_params["block_size_x"] = [32, 64, 128, 256, 512, 1024] + tune_params["INPUT_TYPE"] = [tilus.float16, tilus.float32] + tune_params["OUTPUT_TYPE"] = [tilus.float16, tilus.float32] + + args = [ + size, + Tunable("INPUT_TYPE", { + tilus.float32: a_32, + tilus.float16: a_16, + }), + Tunable("INPUT_TYPE", { + tilus.float32: b_32, + tilus.float16: b_16, + }), + Tunable("OUTPUT_TYPE", { + tilus.float32: c_32, + tilus.float16: c_16, + }), + ] + + print(tune_params) + + observers = [AccuracyObserver("RMSE")] + + + + results, env = tune_kernel( + kernel_name="VecAddV", + kernel_source=VecAddV, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, None, c_expect.cpu()], + observers=observers, + call_function=call_tilus, + verify=verify, + verbose=True, + ) + + + +if __name__ == "__main__": + main() diff --git a/examples/tilus/vec_add.py b/examples/generic_python/tilus_vec_add.py similarity index 53% rename from examples/tilus/vec_add.py rename to examples/generic_python/tilus_vec_add.py index 60dcc6074..dce785909 100644 --- a/examples/tilus/vec_add.py +++ b/examples/generic_python/tilus_vec_add.py @@ -2,12 +2,14 @@ from tilus import float32, int32 from tilus.utils import cdiv, benchmark_func import torch +from kernel_tuner import tune_kernel, run_kernel + class VecAddV(tilus.Script): def __init__(self): super().__init__() - self.block_size = 256 # number of threads per block + self.block_size_x = 32 # number of threads per block def __call__( self, @@ -15,53 +17,59 @@ def __call__( a_ptr: ~float32, # input vector A b_ptr: ~float32, # input vector B c_ptr: ~float32 # output vector C - ): + ): + # compute the number of blocks needed - self.attrs.blocks = [cdiv(n_size, self.block_size)] + self.attrs.blocks = [cdiv(n_size, self.block_size_x)] self.attrs.warps = 1 # number of warps per block # calculate the offset for this block - offset: int32 = self.block_size * self.blockIdx.x + offset: int32 = self.block_size_x * self.blockIdx.x # create global views for input/output vectors ga = self.global_view(a_ptr, dtype=float32, shape=[n_size]) gb = self.global_view(b_ptr, dtype=float32, shape=[n_size]) gc = self.global_view(c_ptr, dtype=float32, shape=[n_size]) - # load a block of A and B into registers - a = self.load_global(ga, offsets=[offset], shape=[self.block_size]) - b = self.load_global(gb, offsets=[offset], shape=[self.block_size]) - - # perform element-wise addition + a = self.load_global(ga, offsets=[offset], shape=[self.block_size_x]) + b = self.load_global(gb, offsets=[offset], shape=[self.block_size_x]) c = a + b - - # store the result back to global memory self.store_global(gc, c, offsets=[offset]) +def call_tilus(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + + def main(): - N = 1 << 20 # vector size - vecadd = VecAddV() - a = torch.rand(N, dtype=torch.float32).cuda() - b = torch.rand(N, dtype=torch.float32).cuda() - c_actual = torch.empty_like(a) + size = 1024000 + + a = torch.randn(size, dtype=torch.float32) + b = torch.randn(size, dtype=torch.float32) + c = torch.zeros_like(b) c_expect = a + b + + + args = [size, a, b, c] + tune_params = dict() + tune_params["block_size_x"] = [32, 64, 128, 256, 512, 1024] - torch.cuda.synchronize() - vecadd(N, a, b, c_actual) - torch.cuda.synchronize() - # correctness check - torch.testing.assert_close(c_expect, c_actual, atol=1e-6, rtol=1e-6) + + results, env = tune_kernel( + kernel_name="VecAddV", + kernel_source=VecAddV, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, None, c_expect.cpu()], + call_function=call_tilus, + verbose=True, + ) - # benchmark - for name, func in [ - ("torch", lambda: a + b), - ("tilus", lambda: vecadd(N, a, b, c_actual)), - ]: - latency = benchmark_func(func, warmup=5, repeat=20) - print(f"{name} latency: {latency:.3f} ms") + if __name__ == "__main__": main() diff --git a/examples/generic_python/triton_vec_add.py b/examples/generic_python/triton_vec_add.py new file mode 100644 index 000000000..be28f2003 --- /dev/null +++ b/examples/generic_python/triton_vec_add.py @@ -0,0 +1,76 @@ +import numpy as np +import triton.language as tl +import torch +from kernel_tuner import tune_kernel, run_kernel +from kernel_tuner.file_utils import store_output_file, store_metadata_file +import triton +from math import ceil + + +@triton.jit +def add_op(x, y): + return x + y + + +#triton.jit +def add_kernel(x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + block_size_x: tl.constexpr, # Number of elements each program should process. + # note: `constexpr` so it can be used as a shape value. + ): + pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. + block_start = pid * block_size_x + offsets = block_start + tl.arange(0, block_size_x) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = add_op(x, y) + tl.store(output_ptr + offsets, output, mask=mask) + + +def call_triton(kernel_function, args, kwargs, grid, threads, params): + kernel_function[grid](*args, **kwargs) + +# NOTE: can the python file be changed in between? what happens? +# NOTE: tune params in the funcion signature are supported as key word arguments. Do not pass them as args, these +# will be inserted automatically. You can use them in the call function (kwargs). + +def tune_with_generic(): + size = 10000000 + + a = torch.randn(size, device='cuda', dtype=torch.float32) + b = torch.randn(size, device='cuda', dtype=torch.float32) + c = torch.empty_like(b) + n = torch.tensor(size, dtype=torch.int32) + c_expect = a + b + + args = [a, b, c, size] + tune_params = dict() + tune_params["block_size_x"] = [2**i for i in range(10)] + + ''' + result = run_kernel("add_kernel", add_kernel, size, args, {"block_size_x": 256}, + lang="generic_python", call_function=call_triton, decorator="@triton.jit") + print(np.allclose(c_expect.cpu(), result[2])) + ''' + + results, env = tune_kernel( + kernel_name="add_kernel", + kernel_source=add_kernel, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, c_expect.cpu(), None], + call_function=call_triton, + decorator="@triton.jit" + ) + + + + + + +tune_with_generic() \ No newline at end of file diff --git a/examples/generic_python/warp_vec_add.py b/examples/generic_python/warp_vec_add.py new file mode 100644 index 000000000..dc08dd6d0 --- /dev/null +++ b/examples/generic_python/warp_vec_add.py @@ -0,0 +1,98 @@ +import warp as wp +import numpy as np +from kernel_tuner import tune_kernel, run_kernel + + + +@wp.func +def add_op(x: float, y: float): + return x + y + + +#@wp.kernel +def vec_add(a: wp.array(dtype=float), + b: wp.array(dtype=float), + c: wp.array(dtype=float), + n: int, + work_per_thread: int): + + tid = wp.tid() + base = tid * work_per_thread + + for i in range(work_per_thread): + idx = base + i + if idx < n: + c[idx] = add_op(a[idx], b[idx]) + + +# TODO do we allways want the call function to have the same parameters +# or do we only require some of them? +def call_warp(kernel_function, args, kwargs, grid, threads, params): + final_args = list(args) + final_args.extend(kwargs.values()) + dim = args[3] + wp.launch(kernel=kernel_function, dim=dim, inputs=final_args) + + +# NOTE default verify function only works for numpy/cupy ndarray, torch Tensor or numpy scalar +# That is why we need a costum verify function for warp. +def verify(answer, result_host, atol): + correct = True + for i, ans in enumerate(answer): + if ans is None: + continue + res = result_host[i].numpy() + if not np.allclose(ans, res, atol=atol): + correct = False + + return correct + + + +def tune(): + n = 1024 + + # Create host arrays + a_np = np.arange(n, dtype=np.float32) + b_np = np.arange(n, 0, -1, dtype=np.float32) + c_np = np.zeros(n, dtype=np.float32) + c_expect = a_np + b_np + + # Create Warp arrays on GPU + a = wp.array(a_np, dtype=float) + b = wp.array(b_np, dtype=float) + c = wp.array(c_np, dtype=float) + + tune_params = dict() + tune_params["work_per_thread"] = [2**i for i in range(10)] + args = [a, b, c, n] + + + ''' + results = run_kernel( + kernel_name="vec_add", + kernel_source=vec_add, + problem_size=n, + arguments=args, + params={"work_per_thread": 16}, + lang="generic_python", + call_function=call_warp, + decorator="@wp.kernel" + ) + ''' + + results, env = tune_kernel( + kernel_name="vec_add", + kernel_source=vec_add, + problem_size=n, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, c_expect, None], + verify=verify, + call_function=call_warp, + decorator="@wp.kernel" + ) + + +tune() \ No newline at end of file diff --git a/examples/triton/vec_add.py b/examples/triton/vec_add.py index 555cfba46..42c3e52c7 100644 --- a/examples/triton/vec_add.py +++ b/examples/triton/vec_add.py @@ -3,18 +3,19 @@ import torch from kernel_tuner import tune_kernel, run_kernel from kernel_tuner.file_utils import store_output_file, store_metadata_file +import triton - +#@triton.jit def add_kernel(x_ptr, # *Pointer* to first input vector. y_ptr, # *Pointer* to second input vector. output_ptr, # *Pointer* to output vector. n_elements, # Size of the vector. - BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. + block_size_x: tl.constexpr, # Number of elements each program should process. # note: `constexpr` so it can be used as a shape value. ): pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) + block_start = pid * block_size_x + offsets = block_start + tl.arange(0, block_size_x) mask = offsets < n_elements x = tl.load(x_ptr + offsets, mask=mask) y = tl.load(y_ptr + offsets, mask=mask) @@ -22,25 +23,28 @@ def add_kernel(x_ptr, # *Pointer* to first input vector. tl.store(output_ptr + offsets, output, mask=mask) -size = 10000000 -a = torch.randn(size, dtype=torch.float32) -b = torch.randn(size, dtype=torch.float32) -c = torch.zeros_like(b) -n = torch.tensor(size, dtype=torch.int32) +def tune_with_triton(): + size = 10000000 + + a = torch.randn(size, device='cuda', dtype=torch.float32) + b = torch.randn(size, device='cuda', dtype=torch.float32) + c = torch.empty_like(b) + n = torch.tensor(size, dtype=torch.int32) + + args = [a, b, c, n] -args = [a, b, c, n] + tune_params = dict() + tune_params["block_size_x"] = [2**i for i in range(10)] -tune_params = dict() -tune_params["block_size_x"] = [2**i for i in range(10)] + results, env = tune_kernel( + kernel_name="add_kernel", + kernel_source=add_kernel, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="triton" + ) -results, env = tune_kernel( - kernel_name="add_kernel", - kernel_source=add_kernel, - problem_size=size, - arguments=args, - tune_params=tune_params, - lang="triton" -) -print("Hello") \ No newline at end of file +tune_with_triton() \ No newline at end of file diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py new file mode 100644 index 000000000..fbe0f0750 --- /dev/null +++ b/kernel_tuner/backends/generic_python.py @@ -0,0 +1,290 @@ +import logging +import inspect +import copy +import traceback # for compile error handling +import re + +from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.observers.generic_python import GenericPythonRuntimeObserver + +try: + import torch +except ImportError: + torch = None + + +# TODO delete temp file + + + +class GenericPythonFunctions(GPUBackend): + + def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): + ''' + In here, everyting is generic if the language uses CUDA as backend + ''' + if not torch: + logging.error("Unable to import Torch") + raise ImportError("Unable to import Torch") + + self.device_id = torch.cuda.current_device() + + self.device_properties = torch.cuda.get_device_properties(self.device_id) + self.name = torch.cuda.get_device_name(self.device_id) + self.max_threads = self.device_properties.max_threads_per_multi_processor + + env = dict() + env["device_name"] = self.name + env["max_threads"] = self.max_threads + env["iterations"] = iterations + env["compiler_options"] = compiler_options + self.env = env + + self.stream = torch.cuda.default_stream() + self.start = torch.cuda.Event(enable_timing=True) + self.end = torch.cuda.Event(enable_timing=True) + + # setup observers + self.observers = observers or [] + self.observers.append(GenericPythonRuntimeObserver(self)) + for obs in self.observers: + obs.register_device(self) + + self.units = {"time": "ms", "power": "s,mW", "energy": "J"} + + # Variables to be filled in at compile time: + self.call_function = None + self.signature = None + self.gpu_kwargs = None + + super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) + + def ready_argument_list(self, arguments): + ''' + The user already supplies the arguments in the correct format, because we are working with + a Python based language anyway. TODO probably only works with torch and numpy? + ''' + return copy.deepcopy(arguments) + + def compile(self, kernel_instance, gpu_args=None): + logging.debug("Compiling Generic Python kernel") + + if kernel_instance.kernel_fn is None: + raise ValueError("kernel_fn is None, currently Generic Python only supports callable kernel_source") + + if gpu_args is None: + raise ValueError("gpu_args is None, Generic Python needs gpu args to compile the kernel") + + # The first time we compile, we also set the call function and the signature + if self.call_function is None or self.signature is None: + self.call_function = kernel_instance.kernel_source.call_function + self.signature = kernel_instance.kernel_source.signature + + grid = kernel_instance.grid + threads = kernel_instance.threads + params = kernel_instance.params + if inspect.isclass(kernel_instance.kernel_fn): + kernel_function = kernel_instance.kernel_fn() + elif callable(kernel_instance.kernel_fn): + # Handles functions and decroators that return callable objects + kernel_function = kernel_instance.kernel_fn + else: + raise TypeError("kernel function is not a class or function") + + self.gpu_kwargs = self.build_gpu_kwargs(params) + + # Call the jit function in order to compile it + self.synchronize() + self.call_function(kernel_function, gpu_args, self.gpu_kwargs, grid, threads, params) + self.synchronize() + + + return kernel_function + + def start_event(self): + logging.debug("Start Generic Python event") + self.start.record() + + def stop_event(self): + logging.debug("Stop Generic Python event") + self.end.record() + + def kernel_finished(self): + logging.debug("Checking if kernel has finished") + return self.end.query() + + def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): + + # Run the kernel + if stream is None: + stream = self.stream + + with torch.cuda.stream(stream): + logging.debug("Running Generic Python kernel") + self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) + + + + def build_gpu_kwargs(self, params=None): + gpu_kwargs = {} + + if params is None: + return gpu_kwargs + + for name, p in self.signature.parameters.items(): + if name in params: + gpu_kwargs[name] = params[name] + + return gpu_kwargs + + + def synchronize(self): + torch.cuda.synchronize() + + def memset(self, allocation, value, size): + pass + + def memcpy_dtoh(self, dest, src): + pass + + def memcpy_htod(self, dest, src): + pass + + def copy_constant_memory_args(self, cmem_args): + raise NotImplementedError("Generic Python does not support constant memory") + + def copy_shared_memory_args(self, smem_args): + raise NotImplementedError("Generic Python does not support shared memory") + + def copy_texture_memory_args(self, texmem_args): + raise NotImplementedError("Generic Python does not support texture memory") + + def refresh_memory(self, gpu_memory, host_arguments, should_sync): + """Refresh the GPU memory with the untouched host arguments. We overwrite the standard function + because Python DSLs do usually do not manage memory explicitely""" + for i, arg in enumerate(host_arguments): + if should_sync[i]: + gpu_memory[i] = copy.deepcopy(arg) + + + def classify_compile_exception(self, e): + """Best effort to differentiate between a user error and a resource error. Input is Exception""" + + RESOURCE_KEYWORDS = ( + # Shared memory + "shared memory", + "smem", + "uses too much shared", + "exceeds shared memory", + "shared memory limit", + + # Registers / occupancy + "too many registers", + "register spill", + "uses too many registers", + "out of registers", + + # Launch configuration + "invalid launch configuration", + "invalid configuration argument", + "threads per block", + "block size", + "grid size", + "num_warps", + "num_ctas", + + # Generic resource exhaustion + "too many resources", + "out of resources", + "exceeds maximum", + "exceeds limit", + + # Compiler-level indicators + "ptxas error", + "ptxas fatal", + "nvcc error", + "cuda error", + "llvm error", + "mlir error", + "lowering failed", + ) + + USER_ERROR_KEYWORDS = ( + # Undefined / missing symbols + "not defined", + "undefined variable", + "without definition", + "unknown variable", + "unbound", + + # Type / shape errors + "type mismatch", + "invalid type", + "cannot convert", + "expected .* but got", + "incompatible types", + + # Indexing / bounds + "index out of bounds", + "out of bounds access", + "invalid index", + + # IR / AST construction + "failed to build", + "invalid expression", + "malformed", + "illegal operation", + ) + + RESOURCE_ORIGINS = ( + "ptxas", + "nvcc", + "cuda", + "llvm", + "mlir", + "cubin", + "fatbin", + ) + + USER_ORIGINS = ( + "transpiler", + "scheduler", + "hidet", + "frontend", + "ast", + ) + + + USER_ERROR_TYPES = ( + NameError, + UnboundLocalError, + AttributeError, + TypeError, + SyntaxError, + IndentationError, + ) + + if isinstance(e, USER_ERROR_TYPES): + return "user_error" + + + def match_any(patterns, text): + return any(re.search(p, text) for p in patterns) + + msg = str(e).lower() + tb = "".join(traceback.format_tb(e.__traceback__)).lower() + + + if match_any(RESOURCE_KEYWORDS, msg): + return "resource_error" + + if match_any(RESOURCE_ORIGINS, msg + tb): + return "resource_error" + + if match_any(USER_ERROR_KEYWORDS, msg): + return "user_error" + + if match_any(USER_ORIGINS, msg + tb): + return "user_error" + + return "unknown" diff --git a/kernel_tuner/backends/pycuda.py b/kernel_tuner/backends/pycuda.py index c8f3e689a..d5d9c49ad 100644 --- a/kernel_tuner/backends/pycuda.py +++ b/kernel_tuner/backends/pycuda.py @@ -375,6 +375,7 @@ def memset(self, allocation, value, size): drv.memset_d8(allocation, value, size) def memcpy_dtoh(self, dest, src): + print("dtoh called") """Perform a device to host memory copy. :param dest: A numpy array in host memory to store the data @@ -389,6 +390,7 @@ def memcpy_dtoh(self, dest, src): dest[:] = src def memcpy_htod(self, dest, src): + print("htod called") """Perform a host to device memory copy. :param dest: A GPU memory allocation unit diff --git a/kernel_tuner/backends/tilus.py b/kernel_tuner/backends/tilus.py deleted file mode 100644 index d3782fb0d..000000000 --- a/kernel_tuner/backends/tilus.py +++ /dev/null @@ -1,194 +0,0 @@ -import logging -import numpy as np -import inspect - -from kernel_tuner.backends.backend import GPUBackend -from kernel_tuner.observers.generic_python import GenericPythonRuntimeObserver - -try: - import torch -except ImportError: - logging.error("Torch not available") - -try: - import tilus -except ImportError: - tilus = None - logging.error("Unable to load Tilus") - - -class TilusFunctions(GPUBackend): - - def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): - ''' - In here, everyting is generic if the language uses Torch as backend - ''' - if not tilus or not torch: - logging.error("Tilus or Torch not available") - raise ImportError("Tilus or Torch not available") - - self.device_id = torch.cuda.current_device() - - self.device_properties = torch.cuda.get_device_properties(self.device_id) - self.name = torch.cuda.get_device_name(self.device_id) - self.max_threads = self.device_properties.max_threads_per_multi_processor - - env = dict() - env["device_name"] = self.name - env["max_threads"] = self.max_threads - env["iterations"] = iterations - env["compiler_options"] = compiler_options - self.env = env - - self.stream = torch.cuda.default_stream() - self.start = torch.cuda.Event(enable_timing=True) - self.end = torch.cuda.Event(enable_timing=True) - - # setup observers - self.observers = observers or [] - self.observers.append(GenericPythonRuntimeObserver(self)) - for obs in self.observers: - obs.register_device(self) - - self.units = {"time": "ms", "power": "s,mW", "energy": "J"} - - super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) - - def ready_argument_list(self, arguments): - ''' - This seems to work for Torch as backend. However, another idea for generic would be - to let the user provide two argument lists: one with numpy and one with the actual arguments - that will be passed to the kernel. TODO: why do we need the numpy list anyway? - ''' - torch_args = [] - - for arg in arguments: - if isinstance(arg, torch.Tensor) and arg.dim() > 0: - torch_args.append(arg.cuda()) - elif isinstance(arg, torch.Tensor) and arg.dim() == 0: - scalar_value = arg.item() - torch_args.append(scalar_value) - elif isinstance(arg, np.ndarray): - torch_arg = torch.from_numpy(arg) - torch_arg_gpu = torch_arg.cuda() - torch_args.append(torch_arg_gpu) - elif isinstance(arg, np.generic): - scalar_value = arg.item() - torch_args.append(scalar_value) - else: - logging.warning("Unknown instance in Tilus functions") - - return torch_args - - def compile(self, kernel_instance, gpu_args=None): - logging.debug("Compiling Tilus kernel") - - if kernel_instance.kernel_fn is None: - raise ValueError("kernel_fn is None, currently Tilus only supports callable kernel_source") - - ''' - if gpu_args is None: - raise ValueError("gpu_args is None, Triton needs gpu args to compile the kernel") - TODO: find out about these GPU args - ''' - - grid = kernel_instance.grid # TODO might need to use this? - threads = kernel_instance.threads - params = kernel_instance.params - kernel_function = kernel_instance.kernel_fn - gpu_kwargs = self.build_gpu_kwargs(kernel_function, threads, params) - - # Call the jit function in order to compile it - kernel_function(*gpu_args, **gpu_kwargs) - - return kernel_function - - def start_event(self): - logging.debug("Start triton event") - self.start.record() - - def stop_event(self): - logging.debug("Stop triton event") - self.end.record() - - def kernel_finished(self): - logging.debug("Checking if kernel has finished") - return self.end.query() - - def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): - ''' - if params is None: - raise ValueError("params is None, Triton needs params in order to set num_warps, num_ctas, etc.") - ''' - - # Run the kernel - if stream is None: - stream = self.stream - - gpu_kwargs = self.build_gpu_kwargs(func, threads, params) - - with torch.cuda.stream(stream): - logging.debug("Running Tilus kernel") - func(*gpu_args, **gpu_kwargs) - - - - def build_gpu_kwargs(self, kernel_function, threads, params=None): - gpu_kwargs = {} - - ''' - if 'BLOCK_SIZE' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE'] = threads[0] - - if 'BLOCK_SIZE_X' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE_X'] = threads[0] - - if 'BLOCK_SIZE_Y' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE_Y'] = threads[1] - - if 'BLOCK_SIZE_Z' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE_Z'] = threads[2] - ''' - - if params is None: - return gpu_kwargs - - - sig = inspect.signature(kernel_function) - for name, p in sig.parameters.items(): - if name in params: - gpu_kwargs[name] = params[name] - - ''' - # Check for Triton specific parameters - if 'num_warps' in params: - gpu_kwargs['num_warps'] = params['num_warps'] - if 'num_ctas' in params: - gpu_kwargs['num_ctas'] = params['num_ctas'] - if 'num_stages' in params: - gpu_kwargs['num_stages'] = params['num_stages'] - ''' - - return gpu_kwargs - - - def synchronize(self): - torch.cuda.synchronize() - - def memset(self, allocation, value, size): - pass - - def memcpy_dtoh(self, dest, src): - pass - - def memcpy_htod(self, dest, src): - pass - - def copy_constant_memory_args(self, cmem_args): - raise NotImplementedError("Tilus does not support constant memory") - - def copy_shared_memory_args(self, smem_args): - raise NotImplementedError("Tilus does not support shared memory") - - def copy_texture_memory_args(self, texmem_args): - raise NotImplementedError("Tilus does not support texture memory") \ No newline at end of file diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 3b545d459..7705d828c 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -22,11 +22,12 @@ from kernel_tuner.backends.opencl import OpenCLFunctions from kernel_tuner.backends.pycuda import PyCudaFunctions from kernel_tuner.backends.triton import TritonFunctions -from kernel_tuner.backends.tilus import TilusFunctions +from kernel_tuner.backends.generic_python import GenericPythonFunctions from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.observers.nvml import NVMLObserver from kernel_tuner.observers.observer import ContinuousObserver, OutputObserver, PrologueObserver from kernel_tuner.observers.tegra import TegraObserver +from kernel_tuner.language import Language try: import torch @@ -58,12 +59,12 @@ class KernelInstance(_KernelInstance): """Class that represents the specific parameterized instance of a kernel.""" def __new__(cls, *args, **kwargs): - # Detect old-style calls (without kernel_fn) + # Detect old-style calls (without kernel_fn and gpu_kwargs) for old tests if len(args) == 8: # old version name, kernel_source, kernel_string, temp_files, threads, grid, params, arguments = args kernel_fn = None args = (name, kernel_source, kernel_string, kernel_fn, temp_files, threads, grid, params, arguments) - elif "kernel_fn" not in kwargs and len(args) == 8: + elif "kernel_fn" not in kwargs and len(args) < 9: kwargs["kernel_fn"] = None return super().__new__(cls, *args, **kwargs) @@ -76,7 +77,7 @@ def delete_temp_files(self): def prepare_temp_files_for_error_msg(self): """Prepare temp file with source code, and return list of temp file names.""" if type(self.kernel_source).__name__ == "KernelSourceFn": - return [] # TODO what do we want to return here? + return [] # already done during compilation temp_filename = util.get_temp_filename(suffix=self.kernel_source.get_suffix()) util.write_file(temp_filename, self.kernel_string) @@ -136,28 +137,28 @@ def __init__( logging.debug("DeviceInterface instantiated, lang=%s", lang) - if lang.upper() == "CUDA": + if lang == Language.CUDA: dev = PyCudaFunctions( device, compiler_options=compiler_options, iterations=iterations, observers=observers, ) - elif lang.upper() == "CUPY": + elif lang == Language.CUPY: dev = CupyFunctions( device, compiler_options=compiler_options, iterations=iterations, observers=observers, ) - elif lang.upper() == "NVCUDA": + elif lang == Language.NVCUDA: dev = CudaFunctions( device, compiler_options=compiler_options, iterations=iterations, observers=observers, ) - elif lang.upper() == "OPENCL": + elif lang == Language.OPENCL: dev = OpenCLFunctions( device, platform, @@ -165,35 +166,35 @@ def __init__( iterations=iterations, observers=observers, ) - elif lang.upper() in ["C", "FORTRAN"]: + elif lang == Language.C or lang == Language.FORTRAN: dev = CompilerFunctions( compiler=compiler, compiler_options=compiler_options, iterations=iterations, observers=observers, ) - elif lang.upper() == "HIP": + elif lang == Language.HIP: dev = HipFunctions( device, compiler_options=compiler_options, iterations=iterations, observers=observers, ) - elif lang.upper() == "HYPERTUNER": + elif lang == Language.HYPERTUNER: dev = HypertunerFunctions( iterations=iterations, compiler_options=compiler_options ) self.requires_warmup = False - elif lang.upper() == "TRITON": + elif lang == Language.TRITON: dev = TritonFunctions( device, compiler_options=compiler_options, iterations=iterations, observers=observers ) - elif lang.upper() == "GENERIC_PYTHON": - dev = TilusFunctions( + elif lang == Language.GENERIC_PYTHON: + dev = GenericPythonFunctions( device, compiler_options=compiler_options, iterations=iterations, @@ -241,7 +242,7 @@ def __init__( print("Using: " + self.dev.name) def run_kernel_bench(self, func, gpu_args, threads, grid, stream=None, params=None): - if isinstance(self.dev, TritonFunctions): + if isinstance(self.dev, TritonFunctions) or isinstance(self.dev, GenericPythonFunctions): self.dev.run_kernel(func, gpu_args, threads, grid, params=params) else: self.dev.run_kernel(func, gpu_args, threads, grid) @@ -392,31 +393,44 @@ def check_kernel_output( self.dev.refresh_memory(gpu_args, instance.arguments, should_sync) # run the kernel + self.dev.synchronize() check = self.run_kernel_check(func, gpu_args, instance) + self.dev.synchronize() if not check: # runtime failure occurred that should be ignored, skip correctness check return # retrieve gpu results to host memory result_host = [] - for i, arg in enumerate(instance.arguments): - if should_sync[i]: - if isinstance(arg, (np.ndarray, cp.ndarray)): - result_host.append(np.zeros_like(arg)) - self.dev.memcpy_dtoh(result_host[-1], gpu_args[i]) - elif isinstance(arg, torch.Tensor) and isinstance(answer[i], torch.Tensor): - if not answer[i].is_cuda: - # if the answer is on the host, copy gpu output to host as well - result_host.append(torch.zeros_like(answer[i])) - self.dev.memcpy_dtoh(result_host[-1], gpu_args[i].tensor) + if instance.kernel_source is not None and instance.kernel_source.lang == Language.GENERIC_PYTHON: + # Python DSLs do not explicitly manage memory. Therefore, we can use gpu args direclty + for i, arg in enumerate(gpu_args): + if should_sync[i]: + if isinstance(arg, torch.Tensor): + result_host.append(arg.cpu()) else: - result_host.append(gpu_args[i].tensor) + result_host.append(arg) + else: + result_host.append(None) + else: + for i, arg in enumerate(instance.arguments): + if should_sync[i]: + if isinstance(arg, (np.ndarray, cp.ndarray)): + result_host.append(np.zeros_like(arg)) + self.dev.memcpy_dtoh(result_host[-1], gpu_args[i]) + elif isinstance(arg, torch.Tensor) and isinstance(answer[i], torch.Tensor): + if not answer[i].is_cuda: + # if the answer is on the host, copy gpu output to host as well + result_host.append(torch.zeros_like(answer[i])) + self.dev.memcpy_dtoh(result_host[-1], gpu_args[i].tensor) + else: + result_host.append(gpu_args[i].tensor) + else: + # We should sync this argument, but we do not know how to transfer this type of argument + # What do we do? Should we throw an error? + result_host.append(None) else: - # We should sync this argument, but we do not know how to transfer this type of argument - # What do we do? Should we throw an error? result_host.append(None) - else: - result_host.append(None) # Call the output observers for obs in self.output_observers: @@ -434,6 +448,7 @@ def check_kernel_output( correct = True if not correct: + print("expected: ", answer, "\ngot: ", result_host) raise RuntimeError("Kernel result verification failed for: " + util.get_config_string(instance.params)) def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, to): @@ -449,7 +464,6 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, instance_string = util.get_instance_string(params) logging.debug("compile_and_benchmark " + instance_string) - instance = self.create_kernel_instance(kernel_source, kernel_options, params, verbose) if isinstance(instance, util.ErrorConfig): result[to.objective] = util.InvalidConfig() @@ -515,33 +529,55 @@ def compile_kernel(self, instance, verbose, gpu_args=None): # compile kernel_string into device func func = None - try: - if isinstance(self.dev, TritonFunctions): + + if isinstance(self.dev, TritonFunctions) or isinstance(self.dev, GenericPythonFunctions): + try: func = self.dev.compile(instance, gpu_args) - else: + except Exception as e: + if isinstance(self.dev, GenericPythonFunctions): + exception_type = self.dev.classify_compile_exception(e) + else: + exception_type = "unkown" + + if exception_type == "user_error": + error_message = str(e.stderr) if hasattr(e, "stderr") else str(e) + print("compile_kernel failed due to error: " + error_message) + print("Error while compiling:", instance.name) + raise e + else: + logging.debug( + "compile_kernel failed due to kernel using too many resources" + ) + if verbose: + print( + f"skipping config {util.get_instance_string(instance.params)} reason: too many resources" + ) + + else: + try: func = self.dev.compile(instance) - except Exception as e: - # compiles may fail because certain kernel configurations use too - # much shared memory for example, the desired behavior is to simply - # skip over this configuration and try the next one - shared_mem_error_messages = [ - "uses too much shared data", - "local memory limit exceeded", - r"local memory \(\d+\) exceeds limit \(\d+\)", - ] - error_message = str(e.stderr) if hasattr(e, "stderr") else str(e) - if any(re.search(msg, error_message) for msg in shared_mem_error_messages): - logging.debug( - "compile_kernel failed due to kernel using too much shared memory" - ) - if verbose: - print( - f"skipping config {util.get_instance_string(instance.params)} reason: too much shared memory used" + except Exception as e: + # compiles may fail because certain kernel configurations use too + # much shared memory for example, the desired behavior is to simply + # skip over this configuration and try the next one + shared_mem_error_messages = [ + "uses too much shared data", + "local memory limit exceeded", + r"local memory \(\d+\) exceeds limit \(\d+\)", + ] + error_message = str(e.stderr) if hasattr(e, "stderr") else str(e) + if any(re.search(msg, error_message) for msg in shared_mem_error_messages): + logging.debug( + "compile_kernel failed due to kernel using too much shared memory" ) - else: - print("compile_kernel failed due to error: " + error_message) - print("Error while compiling:", instance.name) - raise e + if verbose: + print( + f"skipping config {util.get_instance_string(instance.params)} reason: too much shared memory used" + ) + else: + print("compile_kernel failed due to error: " + error_message) + print("Error while compiling:", instance.name) + raise e return func @staticmethod @@ -580,7 +616,8 @@ def create_kernel_instance(self, kernel_source, kernel_options, params, verbose) params, kernel_options.block_size_names, ) - if kernel_source.lang != 'TRITON' and np.prod(threads) > self.dev.max_threads: + + if kernel_source.lang not in [Language.TRITON, Language.GENERIC_PYTHON] and np.prod(threads) > self.dev.max_threads: if verbose: print(f"skipping config {util.get_instance_string(params)} reason: too many threads per block") return util.InvalidConfig() diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index bad921d24..ce9d74080 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -43,6 +43,7 @@ from kernel_tuner.runners.sequential import SequentialRunner from kernel_tuner.runners.simulation import SimulationRunner from kernel_tuner.searchspace import Searchspace +from kernel_tuner.language import Language try: import torch @@ -588,12 +589,15 @@ def tune_kernel( observers=None, objective=None, objective_higher_is_better=None, + call_function=None, + decorator=None ): + start_overhead_time = perf_counter() if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = KernelSource(kernel_name, kernel_source, lang, defines) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function, decorator) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -774,11 +778,13 @@ def run_kernel( block_size_names=None, quiet=False, log=None, + call_function=None, + decorator=None ): if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = KernelSource(kernel_name, kernel_source, lang, defines) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function, decorator) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -806,8 +812,11 @@ def run_kernel( # see if the kernel arguments have correct type util.check_argument_list(instance.name, instance.kernel_string, arguments) - # compile the kernel - func = dev.compile_kernel(instance, False) + # compile the kernel, Extra logic needed to check if we have a python function + if lang.upper() == "TRITON" or lang.upper() == "GENERIC_PYTHON": + func = dev.compile_kernel(instance, False, gpu_args) + else: + func = dev.compile_kernel(instance, False) if func is None: raise RuntimeError("cannot compile kernel, too much shared memory used") @@ -829,16 +838,26 @@ def run_kernel( if not dev.run_kernel_check(func, gpu_args, instance): raise RuntimeError("runtime error occured, too many resources requested") + + # copy data in GPU memory back to the host results = [] - for i, arg in enumerate(arguments): - if numpy.isscalar(arg): - results.append(arg) - elif isinstance(arg, torch.Tensor): - results.append(arg.cpu()) - else: - results.append(numpy.zeros_like(arg)) - dev.memcpy_dtoh(results[-1], gpu_args[i]) + if instance.kernel_source.lang == Language.GENERIC_PYTHON: + for arg in gpu_args: + if isinstance(arg, torch.Tensor): + results.append(arg.cpu()) + else: + results.append(arg) + else: + for i, arg in enumerate(arguments): + if numpy.isscalar(arg): + results.append(arg) + elif isinstance(arg, torch.Tensor): + results.append(arg.cpu()) + else: + results.append(numpy.zeros_like(arg)) + dev.memcpy_dtoh(results[-1], gpu_args[i]) + return results diff --git a/kernel_tuner/kernel_sources/kernel_source.py b/kernel_tuner/kernel_sources/kernel_source.py index 70c94903d..9a16d9b6c 100644 --- a/kernel_tuner/kernel_sources/kernel_source.py +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -6,65 +6,50 @@ from kernel_tuner.language import Language from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData - - - -class KernelSource: - def __new__(cls, kernel_name, kernel_sources, lang, defines=None): +# We use this pattern because we would otherwise get an init twice +# We create a kernelsouce by calling KernelSource(...). The call method for KernelSource is +# replaced by the call method in KernelSourceFactory, because KernelSource uses that class as +# metaclass. Inside the call method, a call to create either KernelSourceStr or KernelSourceFN +# is done. This triggers the __init__ call of those subclasses. In both subclasses, we call super().__init__ +# which initializes some sublclass specific variables and triggers the __init__ call of the KernelSource class. +class KernelSourceFactory(type): + def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function=None, decorator=None): """Factory behavior""" if lang == None: language = None else: try: - language = Language(lang) + language = Language(lang.upper()) except ValueError: raise TypeError(f"Supported languages are {[l.value for l in Language]}") + + + # Determine if we need to create a KernelSourceStr or a KernelSourceFn + if (language and (language == Language.TRITON or language == Language.GENERIC_PYTHON)): + ks_str = False + else: + ks_str = True + from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr + from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn if cls is KernelSource: - if inspect.isfunction(kernel_sources) and (language and (language == Language.TRITON or language == Language.GENERIC_PYTHON)): # TODO should this be isfunction? - from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn - print("CREATING KSFN") - return KernelSourceFn(kernel_name, kernel_sources, lang, defines) - else: - from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr - print("CREATING KSSTR") + if ks_str: return KernelSourceStr(kernel_name, kernel_sources, lang, defines) - - # otherwise, normal subclass init - return super().__new__(cls) - - - def __init__(self, kernel_name, kernel_sources, lang, defines=None): - if not isinstance(kernel_sources, list): - kernel_sources = [kernel_sources] - self.kernel_sources = kernel_sources - self.kernel_name = kernel_name - self.defines = defines + else: + return KernelSourceFn(kernel_name, kernel_sources, lang, defines, call_function, decorator) - if lang is None: - if callable(self.kernel_sources[0]): - raise TypeError("Please specify language when using a code generator function") - kernel_string = self.get_kernel_string(0) - self.lang = util.detect_language(kernel_string) + # Else, normal behaviour for subclasses + if ks_str: + return super().__call__(kernel_name, kernel_sources, lang, defines) else: - self.lang = lang - - @abstractmethod - def prepare_kernel_instance(self, kernel_options, params, grid, threads) -> PreparedKernelSourceData: - raise NotImplementedError("create_kernel_instance not implemented") - - @abstractmethod - def check_argument_lists(self, kernel_name, arguments): - raise NotImplementedError("check_argument_lists not implemented") - - + return super().__call__(kernel_name, kernel_sources, lang, defines, call_function, decorator) + +# TODO do we really want the Language enum? it's a lot of changes +class KernelSource(metaclass=KernelSourceFactory): -''' -class KernelSource: def __init__(self, kernel_name, kernel_sources, lang, defines=None): if not isinstance(kernel_sources, list): kernel_sources = [kernel_sources] - self.kernel_sources = kernel_sources self.kernel_name = kernel_name self.defines = defines @@ -73,9 +58,18 @@ def __init__(self, kernel_name, kernel_sources, lang, defines=None): if callable(self.kernel_sources[0]): raise TypeError("Please specify language when using a code generator function") kernel_string = self.get_kernel_string(0) - self.lang = util.detect_language(kernel_string) + language = util.detect_language(kernel_string) + try: + self.lang = Language(language.upper()) + except ValueError: + # TODO this should never happen + raise TypeError(f"Supported languages are {[l.value for l in Language]}, found {language}") else: - self.lang = lang + try: + self.lang = Language(lang.upper()) + except ValueError: + raise TypeError(f"Supported languages are {[l.value for l in Language]}") + @abstractmethod def prepare_kernel_instance(self, kernel_options, params, grid, threads) -> PreparedKernelSourceData: @@ -85,4 +79,4 @@ def prepare_kernel_instance(self, kernel_options, params, grid, threads) -> Prep def check_argument_lists(self, kernel_name, arguments): raise NotImplementedError("check_argument_lists not implemented") -''' \ No newline at end of file + diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 223e0e84f..edb75f95d 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -18,40 +18,62 @@ class KernelSourceFn(KernelSource): - def __init__(self, kernel_name, kernel_source, lang, defines=None): + def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function=None, decorator=None): super().__init__(kernel_name, kernel_source, lang, defines) if isinstance(kernel_source, list): raise ValueError("KernelSourceFn only supports a single kernel source function") - try: - self.lang = Language(lang) + self.lang = Language(lang.upper()) except ValueError: raise TypeError(f"Supported languages are {[l.value for l in Language]}") + + if call_function is None: + raise ValueError("call_function must be supplied for language Generic Python") + if not callable(call_function): + raise TypeError(f"call_function {call_function} is not a callable object.") + self.call_function = call_function # TODO ceck signature + + if decorator: + if not isinstance(decorator, str): + raise TypeError(f"{decorator} is not a decorator") + if decorator[0] != '@': + raise ValueError(f"{decorator} is not a valid decorator") + self.decorator = decorator + self.source_kernel_fn = kernel_source self.kernel_fn = self.source_kernel_fn - self.source = inspect.getsource(kernel_source) + self.signature = inspect.signature(kernel_source) + try: + self.source = inspect.getsource(kernel_source) + except TypeError as e: + raise TypeError( + f"{e}. Did you forget to remove a decorator before tuning?" + ) from e self.source_tree = ast.parse(self.source) - if self.lang == Language.TRITON: - self.import_nodes = [ - ast.Import(names=[ast.alias(name='triton', asname=None)]), - ast.ImportFrom( - module='triton', - names=[ast.alias(name='language', asname='tl')], - level=0 - ) - ] - else: - self.import_nodes = [n for n in self.source_tree.body if isinstance(n, (ast.Import, ast.ImportFrom))] - + self.import_nodes = self._find_import_nodes(inspect.getfile(kernel_source)) # Find the module where the kernel function is defined self.module = inspect.getmodule(kernel_source) # Get dependencies by analyzing the AST - if self.lang == Language.lang: + if self.lang == Language.TRITON: self.dependencies = self._find_triton_dependencies() else: - self.dependencies = None # TODO + self.dependencies = self._find_dependencies() + + + + + def _find_import_nodes(self, source_file): + with open(source_file, "r") as f: + tree = ast.parse(f.read(), filename=source_file) + + import_nodes = [] + for node in tree.body: + if isinstance(node, (ast.Import, ast.ImportFrom)): + import_nodes.append(node) + + return import_nodes def _find_function_dependencies(self): """Find all function calls in the kernel.""" @@ -99,6 +121,49 @@ def _find_triton_dependencies(self): return dependencies + def _find_function_dependecies2(self, tree, local_funcs): + # Non triton specific + class FunctionCallVisitor(ast.NodeVisitor): + def __init__(self): + self.called = set() + + def visit_Call(self, node): + if isinstance(node.func, ast.Name): + name = node.func.id + if name in local_funcs: + self.called.add(name) + self.generic_visit(node) + + visitor = FunctionCallVisitor() + visitor.visit(tree) + return visitor.called + + def _find_local_functions(self, tree): + local_funcs = set() + for node in tree.body: + if isinstance(node, ast.FunctionDef): + local_funcs.add(node.name) + return local_funcs + + def _find_dependencies(self): + source_file = inspect.getfile(self.source_kernel_fn) + with open(source_file, "r") as f: + source_code = f.read() + tree = ast.parse(source_code, filename=source_file) + local_funcs = self._find_local_functions(tree) + called_functions = self._find_function_dependecies2(self.source_tree, local_funcs) + + dependencies = {} + for node in tree.body: + if (isinstance(node, ast.FunctionDef) and node.name in called_functions and + node.name != self.kernel_name): # Skip the main kernel itself + dependencies[node.name] = node + + return dependencies + + def _add_decorator(self, function): + pass + def prepare_kernel_instance(self, kernel_options, params, grid, threads): new_kernel_fn, temp_file_path = self.apply_params_to_source_fn(params) self.kernel_fn = new_kernel_fn @@ -122,22 +187,34 @@ def apply_params_to_source_fn(self, params): # Add imports new_module.body.extend(self.import_nodes) - if self.lang == Language.TRITON: - # Add all Triton JIT dependencies first - for dep_node in self.dependencies.values(): - new_module.body.append(copy.deepcopy(dep_node)) - - # TODO the kernel can not have dependencies yet. Obviously this needs fixing. + # Add dependencies (functions) + for dep_node in self.dependencies.values(): + dep_node_copy = copy.deepcopy(dep_node) + transformed_dep_node = transformer.visit(dep_node_copy) + #new_module.body.append(copy.deepcopy(dep_node)) + new_module.body.append(transformed_dep_node) # Add transformed main kernel source_tree_copy = copy.deepcopy(self.source_tree) transformed_tree = transformer.visit(source_tree_copy) + + # Add decorator if needed + if self.decorator: + dummy = f"{self.decorator}\ndef _dummy():\n pass\n" + decorator_node = ast.parse(dummy).body[0].decorator_list[0] + for node in transformed_tree.body: + if isinstance(node, ast.FunctionDef): + node.decorator_list.insert(0, decorator_node) + break # only apply to the top level function + new_module.body.extend(transformed_tree.body) # Fix locations and generate source ast.fix_missing_locations(new_module) new_source = astor.to_source(new_module) + #print(new_source) + # Create a unique module name module_name = f'temp_kernel_module_{uuid.uuid4().hex}' @@ -173,5 +250,45 @@ def visit_Name(self, node: ast.Name) -> Any: ast.Constant(value=self.params[node.id]), node ) + return node + + def visit_Attribute(self, node: ast.Attribute): + self.generic_visit(node) + + # Replace self. with constant, but only if it's being read + if ( + isinstance(node.value, ast.Name) + and node.value.id == "self" + and node.attr in self.params + and isinstance(node.ctx, ast.Load) # <- check context + ): + return ast.copy_location(ast.Constant(self.params[node.attr]), node) + + return node + + + def visit_Assign(self, node: ast.Assign): + # Replace assignments like: mock_param = ... + if ( + len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id in self.params + ): + node.value = ast.Constant(value=self.params[node.targets[0].id]) + return node + + # Replace assignment like self. = ... + if ( + len(node.targets) == 1 + and isinstance(node.targets[0], ast.Attribute) + and isinstance(node.targets[0].value, ast.Name) + and node.targets[0].value.id == "self" + and node.targets[0].attr in self.params + ): + node.value = ast.Constant(value=self.params[node.targets[0].attr]) + return node + + return self.generic_visit(node) + + - return node \ No newline at end of file diff --git a/kernel_tuner/kernel_sources/kernel_source_str.py b/kernel_tuner/kernel_sources/kernel_source_str.py index f470b939f..4db0831a7 100644 --- a/kernel_tuner/kernel_sources/kernel_source_str.py +++ b/kernel_tuner/kernel_sources/kernel_source_str.py @@ -60,7 +60,7 @@ def get_kernel_string(self, index=0, params=None): """ logging.debug("get_kernel_string called") - if hasattr(self, 'lang') and self.lang.upper() == "HYPERTUNER": + if hasattr(self, 'lang') and self.lang == Language.HYPERTUNER: return "" kernel_source = self.kernel_sources[index] @@ -102,7 +102,7 @@ def prepare_list_of_files( """ temp_files = dict() - if self.lang.upper() == "HYPERTUNER": + if self.lang == Language.HYPERTUNER: return tuple(["", "", temp_files]) for i, f in enumerate(self.kernel_sources): @@ -156,7 +156,7 @@ def get_suffix(self, index=0): if suffix is not None: return suffix - _suffixes = {"CUDA": ".cu", "OpenCL": ".cl", "C": ".c"} + _suffixes = {Language.CUDA: ".cu", Language.OPENCL: ".cl", Language.C: ".c"} try: return _suffixes[self.lang] except KeyError: diff --git a/kernel_tuner/language.py b/kernel_tuner/language.py index d33933cca..b129e04d5 100644 --- a/kernel_tuner/language.py +++ b/kernel_tuner/language.py @@ -2,13 +2,15 @@ class Language(Enum): CUDA = "CUDA" - OpenCL = "OPENCL" + OPENCL = "OPENCL" C = "C" HIP = "HIP" TRITON = "TRITON" FORTRAN = "FORTRAN" NVCUDA = "NVCUDA" GENERIC_PYTHON = "GENERIC_PYTHON" + CUPY = "CUPY" + HYPERTUNER = "HYPERTUNER" @classmethod def is_valid(cls, value: str) -> bool: diff --git a/kernel_tuner/observers/generic_python.py b/kernel_tuner/observers/generic_python.py index 896b4f9cc..62dfa20c7 100644 --- a/kernel_tuner/observers/generic_python.py +++ b/kernel_tuner/observers/generic_python.py @@ -9,8 +9,7 @@ class GenericPythonRuntimeObserver(BenchmarkObserver): - """Observer that measures time using CUDA events during benchmarking. - TODO think about: do we need torch? Does CUDA work for all languages?""" + """Observer that measures time using CUDA events during benchmarking.""" def __init__(self, dev): if torch is None: diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index fd4cb5697..c28bafa0b 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -165,7 +165,7 @@ def check_argument_list(kernel_name, kernel_string, args): if not isinstance(arg, (np.ndarray, np.generic, cp.ndarray, torch.Tensor, DeviceArray)): raise TypeError( f"Argument at position {i} of type: {type(arg)} should be of type " -- "np.ndarray, numpy scalar, or HIP Python DeviceArray type" + "np.ndarray, numpy scalar, or HIP Python DeviceArray type" ) correct = True diff --git a/test/context.py b/test/context.py index bad152986..81637f24b 100644 --- a/test/context.py +++ b/test/context.py @@ -73,6 +73,12 @@ except ImportError: bayes_opt_gpytorch_present = False +try: + import torch + gen_python_torch_present = True +except ImportError: + gen_python_torch_present = False + try: import pyatf pyatf_present = True @@ -109,6 +115,7 @@ skip_if_no_hip = pytest.mark.skipif(not hip_present, reason="No HIP Python found") skip_if_no_pyatf = pytest.mark.skipif(not pyatf_present, reason="PyATF not installed") skip_if_no_methodology = pytest.mark.skipif(not methodology_present, reason="Autotuning Methodology not found") +skip_if_no_torch = pytest.mark.skipif(not gen_python_torch_present, reason="Torch not installed") def skip_backend(backend: str): @@ -128,3 +135,5 @@ def skip_backend(backend: str): pytest.skip("No nvc++ on PATH") elif backend.upper() == "HIP" and not hip_present: pytest.skip("HIP Python not installed") + elif backend.upper() == "GENERIC_PYTHON" and not gen_python_torch_present: + pytest.skip("Torch not installed") diff --git a/test/test_backend.py b/test/test_backend.py index e694649c1..af9971a48 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -4,8 +4,9 @@ skip_if_no_cuda, skip_if_no_opencl, skip_if_no_pycuda, + skip_if_no_torch, ) -from kernel_tuner.backends import backend, compiler, cupy, nvcuda, opencl, pycuda +from kernel_tuner.backends import backend, compiler, cupy, nvcuda, opencl, pycuda, generic_python class WrongBackend(backend.Backend): @@ -45,3 +46,8 @@ def test_opencl_backend(): @skip_if_no_pycuda def test_pycuda_backend(): dev = pycuda.PyCudaFunctions() + + +@skip_if_no_torch +def test_generic_python_backend(): + dev = generic_python.GenericPythonFunctions() \ No newline at end of file diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py new file mode 100644 index 000000000..a3d250b60 --- /dev/null +++ b/test/test_generic_python_functions.py @@ -0,0 +1,116 @@ +from .context import skip_if_no_torch +from .test_kernel_source_fn import mock_kernel, kernel_with_kwarg, call_mock +from kernel_tuner.core import DeviceInterface, KernelInstance +from kernel_tuner.kernel_sources.kernel_source import KernelSource +import numpy as np + +try: + import torch + torch_present = True +except ImportError: + pass + + +# Helper functions ------------------------------ + +def value_equal(a, b): + # Torch tensors + if isinstance(a, torch.Tensor): + return torch.equal(a, b) + + # NumPy arrays + if isinstance(a, np.ndarray): + return np.array_equal(a, b) + + # Fallback (ints, floats, strings, lists, tuples, dicts) + return a == b + +def get_context(): + params = {'mock_param': 64} + a = 42 + b = torch.randn(12, device='cuda', dtype=torch.float32) + args = [a, b] + ks = KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + return ks, args, params + + +# Tests ---------------------------------------------- + +@skip_if_no_torch +def test_ready_argument_list(): + ks, args, params = get_context() + dev = DeviceInterface(ks) + gpu_args = dev.ready_argument_list(args) + + assert len(args) == len(gpu_args) + + for i, _ in enumerate(gpu_args): + assert value_equal(args[i], gpu_args[i]) + if type(gpu_args[i]) in (list, dict, torch.Tensor, np.ndarray): + assert gpu_args[i] is not args[i] # Assure deep copy + + +@skip_if_no_torch +def test_compile(): + ks, args, params = get_context() + + instance_data = ks.prepare_kernel_instance( + kernel_options=None, + params=params, + grid=None, + threads=None, + ) + dev = DeviceInterface(ks) + instance = KernelInstance( + name="mock_kernel", + kernel_source=ks, + kernel_string=None, + kernel_fn=instance_data.kernel_fn, + temp_files=instance_data.temp_files, + threads=None, + grid=None, + params=params, + arguments=None, + ) + gpu_args = dev.ready_argument_list(args) + + callable_fn = dev.compile_kernel(instance, verbose=False, gpu_args=gpu_args) + + # The mock function return the mock_param, which should be set to 64. + assert callable_fn(*args) == 64 + + +@skip_if_no_torch +def test_gpu_kwargs(): + params = {'mock_param': 64} + a = torch.randn(12, device='cuda', dtype=torch.float32) + args = [a] # we do not have to specify the kwarg here + ks = KernelSource("kernel_with_kwarg", kernel_with_kwarg, "generic_python", call_function=call_mock) + dev = DeviceInterface(ks) + + instance_data = ks.prepare_kernel_instance( + kernel_options=None, + params=params, + grid=None, + threads=None, + ) + dev = DeviceInterface(ks) + instance = KernelInstance( + name="kernel_with_kwarg", + kernel_source=ks, + kernel_string=None, + kernel_fn=instance_data.kernel_fn, + temp_files=instance_data.temp_files, + threads=None, + grid=None, + params=params, + arguments=None, + ) + gpu_args = dev.ready_argument_list(args) + callable_fn = dev.compile_kernel(instance, verbose=False, gpu_args=gpu_args) + + kwargs = dev.dev.build_gpu_kwargs(params) + + assert kwargs["mock_param"] == 64 + assert callable_fn(*args, **kwargs) == 64 + diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py new file mode 100644 index 000000000..e23a41a15 --- /dev/null +++ b/test/test_kernel_source_fn.py @@ -0,0 +1,160 @@ +import pytest +import inspect +import ast +import functools # for an example decorator +from kernel_tuner.kernel_sources.kernel_source import KernelSource +# Note: need subclass imports to test for types and instantiate subclasses direclty. +# Normally, just importing KernelSource is enough. +from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn +from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr + + +# Helper functions -------------------------------------------- + +def normalize_ast(src: str): + return ast.dump(ast.parse(src), include_attributes=False) + +def mock_kernel(a, b): + mock_param = 256 + if a < mock_param: + b = 42 + return mock_param + +def kernel_with_kwarg(a, mock_param): + return mock_param + +# NOTE should kernel_with_kwarg now return the tuning param value for mock param or 5? +def kernel_with_dependency(): + mock_param = 42 + foo = kernel_with_kwarg(mock_param, 5) + return mock_kernel(42, 42) + +def call_mock(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + +class TilusLike(): + def __init__(self): + self.mock_param = 32 + + def __call__(self): + return self.mock_param + + def another_function(self): + return self.mock_param + + +# Tests ------------------------------------------------------ + +def test_factory_behaviour(): + # KernelSourceFn should only be created when language generic_python is supplied + ks_fn = KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + ks_str = KernelSource("vector_add", 'extern "C" __global__ void vector_add(float *c, float *a, float *b, int n) {', lang=None) + + assert isinstance(ks_fn, KernelSourceFn) + assert isinstance(ks_str, KernelSourceStr) + + +def test_initiation(): + ''' + Test invalid KernelSourceFn initations + ''' + with pytest.raises(ValueError, match=r"call_function must be supplied for language .*"): + KernelSource("mock_kernel", mock_kernel, "generic_python") + + with pytest.raises(TypeError, match=r".* is not a callable object"): + KernelSource("mock_kernel", "This is a string Kernel", "generic_python", call_function=call_mock) + + with pytest.raises(ValueError, match=r"KernelSourceFn only supports a single kernel source function"): + KernelSource("mock_kernel", [mock_kernel, kernel_with_dependency], "generic_python", call_function=call_mock) + + with pytest.raises(TypeError, match=r".* is not a callable object"): + KernelSource("mock_kernel", mock_kernel, "generic_python", call_function="not a function") + + with pytest.raises(ValueError, match=r".* is not a valid decorator"): + KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator="not a decorator") + + with pytest.raises(TypeError, match=r".* is not a decorator"): + KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator=mock_kernel) + + +def test_param_subsitution(): + params = {"mock_param": 128} + ks = KernelSourceFn("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + new_kernel_fn, _ = ks.apply_params_to_source_fn(params) + + actual_src = inspect.getsource(new_kernel_fn) + expected_src = """ +def mock_kernel(a, b): + mock_param = 128 + if a < 128: + b = 42 + return 128 +""" + + assert normalize_ast(actual_src) == normalize_ast(expected_src) + + +def test_imports(): + ''' + Import statements that are present in the file where the function lives + should also be present in the file where the new function lives. + ''' + + params = {"mock_param": 128} + ks = KernelSourceFn("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + _, temp_path = ks.apply_params_to_source_fn(params) + + # Check if imports are present + with open(temp_path) as f: + full_src = f.read() + + assert "import pytest" in full_src + assert "import inspect" in full_src + + +def test_param_substitution_class(): + ''' + Tilus uses the __call__ function of a class to define its kernel. Therefore, + param substitution should also work classes. Param substiution in other class + functions is also supported. + ''' + + params = {"mock_param": 128} + ks = KernelSourceFn("TilusLike", TilusLike, "generic_python", call_function=call_mock) + + new_kernel_fn, _ = ks.apply_params_to_source_fn(params) + + actual_src = inspect.getsource(new_kernel_fn) + expected_src = """ +class TilusLike(): + def __init__(self): + self.mock_param = 128 + + def __call__(self): + return 128 + + def another_function(self): + return 128 +""" + assert normalize_ast(actual_src) == normalize_ast(expected_src) + + + +def test_decorator(): + params = {"mock_param": 128} + ks = KernelSourceFn("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator="@functools.lru_cache()") + new_kernel_fn, _ = ks.apply_params_to_source_fn(params) + + assert hasattr(new_kernel_fn, "__wrapped__") + + +def test_dependencies(): + params = {"mock_param": 128} + ks = KernelSourceFn("kernel_with_dependency", kernel_with_dependency, "generic_python", call_function=call_mock) + new_kernel_fn, _ = ks.apply_params_to_source_fn(params) + res = new_kernel_fn() # This should not throw an error if the dependency exists in the module. + + assert res == 128 + + + From d21377e16758bd3544503b58076ca1130ccb19b9 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Sun, 1 Mar 2026 10:40:14 +0100 Subject: [PATCH 04/44] improved features, removed explicit Triton support (now available through generic Python), cleaned up code --- .../generic_python/matmul/tilelang_matmul.py | 212 ++++++++++++ .../generic_python/matmul/tilus_matmul.py | 267 +++++++++++++++ .../generic_python/matmul/triton_matmul.py | 263 ++++++++++++++- examples/generic_python/numba_vec_add.py | 38 ++- examples/generic_python/tilelang_vec_add.py | 57 ++++ examples/generic_python/tilus_vec_add.py | 58 +++- examples/generic_python/triton_vec_add.py | 8 +- examples/generic_python/warp_vec_add.py | 36 ++- examples/triton/conv2d_tuning.py | 306 ------------------ examples/triton/vec_add.py | 50 --- kernel_tuner/backends/generic_python.py | 189 ++++++++--- kernel_tuner/backends/triton.py | 174 ---------- kernel_tuner/core.py | 32 +- kernel_tuner/interface.py | 35 +- kernel_tuner/kernel_sources/kernel_source.py | 32 +- .../kernel_sources/kernel_source_fn.py | 300 +++++++++-------- .../kernel_sources/kernel_source_str.py | 2 +- kernel_tuner/language.py | 1 - kernel_tuner/observers/triton.py | 32 -- test/test_generic_python_functions.py | 2 +- test/test_kernel_source_fn.py | 6 +- 21 files changed, 1254 insertions(+), 846 deletions(-) create mode 100644 examples/generic_python/matmul/tilelang_matmul.py create mode 100644 examples/generic_python/matmul/tilus_matmul.py create mode 100644 examples/generic_python/tilelang_vec_add.py delete mode 100644 examples/triton/conv2d_tuning.py delete mode 100644 examples/triton/vec_add.py delete mode 100644 kernel_tuner/backends/triton.py delete mode 100644 kernel_tuner/observers/triton.py diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py new file mode 100644 index 000000000..2a7b5154d --- /dev/null +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -0,0 +1,212 @@ +import tilelang +import tilelang.language as T +import torch +from kernel_tuner import tune_kernel + + +#@tilelang.jit +def matmul(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): + @T.prim_func + def main( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), dtype), + ): + # Define a grid with enough blocks to cover M×N + num_threads=128 + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): + + # Allocate shared memory for the current tile of A and B + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + + # Allocate a local (register) fragment for partial accumulations + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Enable swizzle-based rasterization for better L2 locality + panel_size = 4 + T.use_swizzle(panel_size=panel_size, enable=True) + + # Initialize the local accumulation buffer to zero + T.clear(C_local) + + num_stages=3 + + # Loop over the K dimension in block_K chunks, using a pipeline + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + # Copy from global memory to shared memory + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[k * block_K, bx * block_N], B_shared) + + # Perform a matrix multiply-accumulate on the tile + T.gemm(A_shared, B_shared, C_local) + + # Copy the accumulated result from local memory (C_local) to global memory (C) + T.copy(C_local, C[by * block_M, bx * block_N]) + + return main + +@tilelang.jit +def matmul_with_decorator(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): + @T.prim_func + def main( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), dtype), + ): + # Define a grid with enough blocks to cover M×N + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + + # Allocate shared memory for the current tile of A and B + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + + # Allocate a local (register) fragment for partial accumulations + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Initialize the local accumulation buffer to zero + T.clear(C_local) + + # Loop over the K dimension in block_K chunks, using a 3-stage pipeline + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): + # Copy from global memory to shared memory + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[k * block_K, bx * block_N], B_shared) + + # Perform a matrix multiply-accumulate on the tile + T.gemm(A_shared, B_shared, C_local) + + # Copy the accumulated result from local memory (C_local) to global memory (C) + T.copy(C_local, C[by * block_M, bx * block_N]) + + return main + + + +def run(m, n, k): + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(k, n, device="cuda", dtype=torch.float16) + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + kernel = matmul(m, n, k, 128, 128, 32) + kernel(a, b, c) + ref_c = a @ b + tol = m * 2**(-11) + # Validate results + torch.testing.assert_close(c, ref_c, rtol=tol, atol=tol) + + +def call_tilelang(kernel_function, args, kwargs, grid, threads, params): + compiled_kernel = kernel_function(**kwargs) + compiled_kernel(*args) + + +def time(m, n, k): + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(k, n, device="cuda", dtype=torch.float16) + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + c_ans = a @ b + + args = [a, b, c] + tune_params = dict() + tune_params["M"] = [m] + tune_params["K"] = [k] + tune_params["N"] = [n] + tune_params["block_M"] = [64, 128] + tune_params["block_N"] = [64, 128] + tune_params["block_K"] = [32, 64] + + results_kt, env = tune_kernel("matmul", matmul, m * n, args, tune_params, lang="generic_python", + call_function=call_tilelang, decorator="@tilelang.jit", verbose=False, iterations=100) + + import time + num_repeats = 100 + times_direct = [] + for config in results_kt: + bs_m = config["block_M"] + bs_n = config["block_N"] + bs_k = config["block_K"] + + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + + kernel = matmul_with_decorator(m, n, k, bs_m, bs_n, bs_k) + kernel(a, b, c) + + torch.allclose(c.cpu(), c_ans.cpu(), atol=m * 2**(-11)) + + for i in range(num_repeats): + times = [] + + torch.cuda.synchronize() + start = time.time() + kernel(a, b, c) + torch.cuda.synchronize() + times.append(time.time() - start) + + avg_time_ms = round((1000 * sum(times) / len(times)), 3) + times_direct.append(avg_time_ms) + print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, time={avg_time_ms}ms") + + + import matplotlib.pyplot as plt + + # Extract times + times_kt = [cfg['time'] for cfg in results_kt] + + # x-axis labels + configs = [f"config{i}" for i in range(len(times_kt))] + x = range(len(configs)) + + plt.figure(figsize=(10,6)) + plt.plot(configs, times_kt, marker='s', label='KernelTuner') + plt.plot(configs, times_direct, marker='x', label='Direct') + plt.ylabel('Time (ms)') + plt.xlabel('Configuration') + plt.title('Kernel execution time per configuration') + plt.xticks(rotation=45) + plt.legend() + plt.grid(True) + plt.tight_layout() + plt.savefig("tilelang.png") + print("saved fig") + + +def tune(m, n, k): + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(k, n, device="cuda", dtype=torch.float16) + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + c_actual = a @ b + + args = [a, b, c] + tune_params = dict() + tune_params["M"] = [m] + tune_params["K"] = [k] + tune_params["N"] = [n] + tune_params["block_M"] = [64, 128, 256] + tune_params["block_N"] = [64, 128, 256] + tune_params["block_K"] = [32, 64, 128] + tune_params["num_stages"] = [2, 3, 4] + tune_params["panel_size"] = [4, 8] # equivalent to group size m in Triton + tune_params["num_threads"] = [64, 128, 256] + + restrictions = [ + # tile size budget + "block_M * block_N <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "block_M <= 4 * block_N", + "block_N <= 4 * block_M", + + # large K only with reasonably large M/N + "not (block_K == 128 and block_M < 64 and block_N < 64)", + ] + + tol = m * 2**(-11) + answer = [None, None, c_actual.cpu()] + + results, env = tune_kernel("matmul", matmul, m * n, args, tune_params, atol=tol, lang="generic_python", + call_function=call_tilelang, restrictions=restrictions, answer=answer, decorator="@tilelang.jit", verbose=False) + +if __name__ == "__main__": + #m, n, k = 1024, 1024, 1024 + m, n, k = 8192, 8192, 8192 + time(m, n, k) \ No newline at end of file diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py new file mode 100644 index 000000000..07532013f --- /dev/null +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -0,0 +1,267 @@ +import math + +import pandas +import tilus +import torch +from tilus import float16, float32, int32 +from tilus.utils import benchmark_func +from kernel_tuner import tune_kernel, run_kernel + + + +class MatmulV4(tilus.Script): + def __init__(self): + super().__init__() + self.block_m = 128 + self.block_n = 128 + self.block_k = 16 + self.num_warps = 4 + self.num_stages = 4 + + def __call__( + self, + m_size: int32, + n_size: int, + k_size: int, + a_ptr: ~float16, + b_ptr: ~float16, + c_ptr: ~float16, + ): + self.attrs.blocks = [ + self.utils.ceil_div(m_size, self.block_m), + self.utils.ceil_div(n_size, self.block_n), + ] + self.attrs.warps = self.num_warps + + block_m, block_n, block_k = self.block_m, self.block_n, self.block_k + offset_m: int32 = block_m * self.blockIdx.x + offset_n: int32 = block_n * self.blockIdx.y + + ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) + gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) + sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) + sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) + acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) + + for stage in range(self.num_stages - 1): + offset_k = stage * self.block_k + self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) + self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) + self.copy_async_commit_group() + + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + current_stage: int32 = 0 + preload_stage: int32 = self.num_stages - 1 + for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): + # computation for current tile + a = self.load_shared(sa[current_stage]) + b = self.load_shared(sb[current_stage]) + self.dot(a, b, acc, out=acc) + + # preload the next tile of A and B into shared memory + preload_offset_k = offset_k + (self.num_stages - 1) * block_k + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) + self.copy_async_commit_group() + + # update the stage + current_stage = (current_stage + 1) % self.num_stages + preload_stage = (preload_stage + 1) % self.num_stages + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + self.free_shared(sa) + self.free_shared(sb) + + casted_acc = self.cast(acc, dtype=float16) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + + + +class MatmulGroupedOrdering(tilus.Script): + def __init__(self): + super().__init__() + self.block_m = 128 + self.block_n = 128 + self.block_k = 16 + self.num_warps = 4 + self.num_stages = 4 + self.group_size_m = 8 + + def __call__( + self, + m_size: int32, + n_size: int, + k_size: int, + a_ptr: ~float16, + b_ptr: ~float16, + c_ptr: ~float16, + ): + block_m, block_n, block_k = self.block_m, self.block_n, self.block_k + + num_pid_m = self.utils.ceil_div(m_size, block_m) + num_pid_n = self.utils.ceil_div(n_size, block_n) + self.attrs.blocks = [num_pid_m * num_pid_n] + + pid = self.blockIdx.x + num_pid_in_group = self.group_size_m * num_pid_n + group_id = pid // num_pid_in_group + + first_pid_m = group_id * self.group_size_m + group_size_m = min(num_pid_m - first_pid_m, self.group_size_m) + + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + self.attrs.warps = self.num_warps + offset_m: int32 = pid_m * block_m + offset_n: int32 = pid_n * block_n + + ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) + gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) + sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) + sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) + acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) + + for stage in range(self.num_stages - 1): + offset_k = stage * self.block_k + self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) + self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) + self.copy_async_commit_group() + + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + current_stage: int32 = 0 + preload_stage: int32 = self.num_stages - 1 + for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): + # computation for current tile + a = self.load_shared(sa[current_stage]) + b = self.load_shared(sb[current_stage]) + self.dot(a, b, acc, out=acc) + + # preload the next tile of A and B into shared memory + preload_offset_k = offset_k + (self.num_stages - 1) * block_k + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) + self.copy_async_commit_group() + + # update the stage + current_stage = (current_stage + 1) % self.num_stages + preload_stage = (preload_stage + 1) % self.num_stages + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + self.free_shared(sa) + self.free_shared(sb) + + casted_acc = self.cast(acc, dtype=float16) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + + + +def main(): + headers = ["m", "n", "k", "name", "latency (ms)", "tflops"] + workloads = [ + [4096, 4096, 4096], + [1024, 1024, 14336], + ] + + rows = [] + for m, n, k in workloads: + matmul = MatmulGroupedOrdering() #MatmulV4() + + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + matmul(m, n, k, a, b, c_actual) + + # check correctness + torch.testing.assert_close(c_expect, c_actual) + + # benchmark + for name, func in [ + ("torch", lambda: torch.matmul(a, b, out=c_expect)), + ("tilus", lambda: matmul(m, n, k, a, b, c_actual)), + ]: + latency = benchmark_func(func, warmup=5, repeat=20) + tflops = 2 * m * n * k / latency * 1e-9 + rows.append([m, n, k, name, latency, tflops]) + + df = pandas.DataFrame(rows, columns=headers) + print(df) + +def call_tilus(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + + +def tune_matmul(m, n, k): + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n #(m, n) + args = [m, n, k, a, b, c_actual] + tune_params = dict() + tune_params["block_m"] = [32, 64, 128, 256] + tune_params["block_n"] = [32, 64, 128, 256] + tune_params["block_k"] = [32, 64, 128] + tune_params["group_size_m"] = [4, 8] + tune_params["num_stages"] = [2, 3, 4] + tune_params["num_warps"] = [4, 8] + + + restrictions = [ + # tile size budget + "block_m * block_n <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "block_m <= 4 * block_n", + "block_n <= 4 * block_m", + + # large K only with reasonably large M/N + "not (block_k == 128 and block_m < 64 and block_n < 64)", + + # 32x32 requires 8 warps + "not (block_m == 32 and block_n == 32 and num_warps < 8)", + ] + + + answer = [None] * 6 + answer[-1] = c_expect.cpu() + atol = 1e-2 #m * 2**(-11) + + results, env = tune_kernel("MatmulGroupedOrdering", MatmulGroupedOrdering, size, args, tune_params, grid_div_x = ["block_m", "block_n"], + answer = answer, atol=atol, restrictions=restrictions, + lang="generic_python", call_function=call_tilus, + block_size_names=["block_m", "block_n", "block_k"], strategy="simulated_annealing") + + +if __name__ == "__main__": + #m, n, k = 4096, 4096, 4096 + m, n, k = 8192, 8192, 8192 + tune_matmul(m, n, k) + + #main() \ No newline at end of file diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py index 7115ea5cc..238a61c9f 100644 --- a/examples/generic_python/matmul/triton_matmul.py +++ b/examples/generic_python/matmul/triton_matmul.py @@ -3,7 +3,8 @@ import triton import triton.language as tl -DEVICE = triton.runtime.driver.active.get_active_torch_device() +from kernel_tuner import tune_kernel +from kernel_tuner import run_kernel def get_cuda_autotune_config(): return [ @@ -48,7 +49,7 @@ def get_cuda_autotune_config(): key=['M', 'N', 'K'], ) ''' -@triton.jit +#@triton.jit def matmul_kernel( # Pointers to matrices a_ptr, b_ptr, c_ptr, @@ -62,7 +63,7 @@ def matmul_kernel( stride_cm, stride_cn, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # - GROUP_SIZE_M: tl.constexpr, # + GROUP_SIZE_M: tl.constexpr, num_stages, num_warps# ): """Kernel for computing the matmul C = A x B. A has shape (M, K), B has shape (K, N) and C has shape (M, N) @@ -132,23 +133,251 @@ def matmul_kernel( tl.store(c_ptrs, c, mask=c_mask) -def matmul(a, b): - # Check constraints. - assert a.shape[1] == b.shape[0], "Incompatible dimensions" - assert a.is_contiguous(), "Matrix A must be contiguous" - M, K = a.shape - K, N = b.shape - # Allocates output. - c = torch.empty((M, N), device=a.device, dtype=torch.float16) - # 1D launch kernel where each block gets its own program. - grid = lambda META: (triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']), ) +def run_matmul(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + grid = lambda META: (triton.cdiv(m, META['BLOCK_SIZE_M']) * triton.cdiv(n, META['BLOCK_SIZE_N']), ) + matmul_kernel[grid]( - a, b, c, # - M, N, K, # + a, b, c_actual, # + m, n, k, # a.stride(0), a.stride(1), # b.stride(0), b.stride(1), # - c.stride(0), c.stride(1) + c_actual.stride(0), c_actual.stride(1), + 128, 256, 64, 8 ) - return c + + torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) + + +def run_matmul_kt(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n + + args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1)] + + params = {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M":4, "num_stages":3, "num_warps":4} + + result = run_kernel("matmul_kernel", matmul_kernel, size, args, params=params, grid_div_x=["BLOCK_SIZE_N", "BLOCK_SIZE_M"], + lang="generic_python", decorator="@triton.jit", call_function=call_triton, + block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) + c_res = result[2] + + assert torch.allclose(c_res, c_expect.cpu(), atol=1e-2, rtol=1e-1) + + + + + + +def call_triton(kernel_function, args, kwargs, grid, threads, params): + #print("using grid: ", grid) + #print("args: ", args) + #print("kwargs: ", kwargs) + kernel_function[grid](*args, **kwargs) + + + + +def check_time(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n + args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1)] + tune_params = dict() + tune_params["BLOCK_SIZE_M"] = [64, 128] + tune_params["BLOCK_SIZE_N"] = [64, 128] + tune_params["BLOCK_SIZE_K"] = [32, 64] + tune_params["GROUP_SIZE_M"] = [4, 8] + tune_params["num_stages"] = [3] + tune_params["num_warps"] = [4] + + restrictions = [ + # tile size budget + "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", + "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", + + # large K only with reasonably large M/N + "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", + + # 32x32 requires 8 warps + "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", + ] + + grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] + + answer = [None] * 12 + answer[2] = c_expect.cpu() + atol = 1e-2 #m * 2**(-11) + + + + results_ours, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, + restrictions=restrictions, iterations=100, answer=answer, atol=atol, + lang="generic_python", decorator="@triton.jit", call_function=call_triton, + block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) + + + results_prev, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, + restrictions=restrictions, iterations=100, answer=answer, atol=atol, + lang="triton", block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) + + + + import time + num_repeats = 100 + times_direct = [] + for config in results_prev: + + bs_m = config["BLOCK_SIZE_M"] + bs_n = config["BLOCK_SIZE_N"] + bs_k = config["BLOCK_SIZE_K"] + gs_m = config["GROUP_SIZE_M"] + num_stages = config["num_stages"] + num_warps = config["num_warps"] + + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + + grid = (triton.cdiv(m, bs_m) * triton.cdiv(n, bs_n, ), ) + jit_function = triton.jit(matmul_kernel) + + + jit_function[grid]( + a, b, c_actual, + m, n, k, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1), + bs_m, bs_n, bs_k, gs_m, num_stages, num_warps + ) + + + torch.allclose(c_expect.cpu(), c_actual.cpu(), atol=1e-2) + + + + for i in range(num_repeats): + times = [] + + #c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + + torch.cuda.synchronize() + start = time.time() + jit_function[grid]( + a, b, c_actual, + m, n, k, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1), + bs_m, bs_n, bs_k, gs_m, num_stages, num_warps + ) + + torch.cuda.synchronize() + times.append(time.time() - start) + + avg_time_ms = round((1000 * sum(times) / len(times)), 3) + times_direct.append(avg_time_ms) + print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, GROUP_SIZE_M={gs_m}, num_stages={num_stages}, num_warps={num_warps}, time={avg_time_ms}ms") + + + import matplotlib.pyplot as plt + + # Extract times + times_prev = [cfg['time'] for cfg in results_prev] + times_ours = [cfg['time'] for cfg in results_ours] + + # x-axis labels + configs = [f"config{i}" for i in range(len(times_prev))] + x = range(len(configs)) + + plt.figure(figsize=(10,6)) + plt.plot(configs, times_prev, marker='o', label='Triton tuned') + plt.plot(configs, times_ours, marker='s', label='Generic tuned') + plt.plot(configs, times_direct, marker='x', label='Direct') + plt.ylabel('Time (ms)') + plt.xlabel('Configuration') + plt.title('Kernel execution time per configuration') + plt.xticks(rotation=45) + plt.legend() + plt.grid(True) + plt.tight_layout() + plt.savefig("ouptut.png") + + + + + + +def tune_matmul(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n + args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1)] + tune_params = dict() + tune_params["BLOCK_SIZE_M"] = [64, 128, 256] + tune_params["BLOCK_SIZE_N"] = [64, 128, 256] + tune_params["BLOCK_SIZE_K"] = [32, 64, 128] + tune_params["GROUP_SIZE_M"] = [4, 8] + tune_params["num_stages"] = [2, 3, 4] + tune_params["num_warps"] = [4, 8] + + restrictions = [ + # tile size budget + "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", + "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", + + # large K only with reasonably large M/N + "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", + + # 32x32 requires 8 warps + "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", + ] + + grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] + + answer = [None] * 12 + answer[2] = c_expect.cpu() + + results, env = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, + answer = answer, atol=4.0, restrictions=restrictions, + lang="generic_python", decorator="@triton.jit", call_function=call_triton, + block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], strategy="simulated_annealing") + + + + + +if __name__ == "__main__": + m, n, k = 8192, 8192, 8192 + #m, n, k = 4096, 4096, 4096 + #tune_matmul(m, n, k) + #check_time(m, n, k) + #run_matmul_kt(m, n, k) + + + + diff --git a/examples/generic_python/numba_vec_add.py b/examples/generic_python/numba_vec_add.py index 3732aa018..e49f9e372 100644 --- a/examples/generic_python/numba_vec_add.py +++ b/examples/generic_python/numba_vec_add.py @@ -1,7 +1,8 @@ -import numpy as np from numba import cuda from math import ceil from kernel_tuner import tune_kernel, run_kernel +import torch + #@cuda.jit def f(a, b, c): @@ -13,7 +14,13 @@ def f(a, b, c): def call_numba(kernel_function, args, kwargs, grid, threads, params): - kernel_function[grid[0], threads[0]](*args, **kwargs) + numba_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + numba_args.append(cuda.as_cuda_array(arg)) + else: + numba_args.append(arg) + kernel_function[grid, threads](*args, **kwargs) def verify(answer, result_host, atol): @@ -28,15 +35,25 @@ def verify(answer, result_host, atol): return correct +import numpy as np N = 100000000 -a = cuda.to_device(np.random.random(N)) -b = cuda.to_device(np.random.random(N)) -c = cuda.device_array_like(a) -c_expect = a.copy_to_host() + b.copy_to_host() + +a = np.random.random(N) +b = np.random.random(N) +c = np.zeros(N) +c_expect = a + b +#c_expect = a.copy_to_host() + b.copy_to_host() + + +#a_torch = torch.rand(N, dtype=torch.float32) +#b_torch = torch.rand(N, dtype=torch.float32) +#c_torch = torch.zeros(N, dtype=torch.float32) +#c_expect = (a_torch + b_torch).cpu() args = [a, b, c] tune_params = {"block_size_x": [2**i for i in range(10)]} +''' results = run_kernel( kernel_name="f", kernel_source=f, @@ -47,11 +64,12 @@ def verify(answer, result_host, atol): call_function=call_numba, decorator="@cuda.jit" ) +''' -print(np.allclose(results[2], c_expect)) +#print(np.allclose(results[2], c_expect)) + -''' results, env = tune_kernel( kernel_name="f", kernel_source=f, @@ -60,9 +78,9 @@ def verify(answer, result_host, atol): tune_params=tune_params, lang="generic_python", answer=[None, None, c_expect], - verify=verify, + #verify=verify, call_function=call_numba, decorator="@cuda.jit" ) -''' + diff --git a/examples/generic_python/tilelang_vec_add.py b/examples/generic_python/tilelang_vec_add.py new file mode 100644 index 000000000..3ff3c9376 --- /dev/null +++ b/examples/generic_python/tilelang_vec_add.py @@ -0,0 +1,57 @@ +import tilelang +import tilelang.language as T +import torch +from kernel_tuner import tune_kernel + +#@tilelang.jit # infers target from tensors at first call +def add(N: int, dtype: str = 'float32', block: int = 256,): + + @T.prim_func + def add_kernel( + A: T.Tensor((N,), dtype), + B: T.Tensor((N,), dtype), + C: T.Tensor((N,), dtype), + ): + with T.Kernel(T.ceildiv(N, block), threads=block) as bx: + for i in T.Parallel(block): + gi = bx * block + i + # Optional — LegalizeSafeMemoryAccess inserts a guard when an access may be OOB + C[gi] = A[gi] + B[gi] + + return add_kernel + + +def run_normal(): + # Host side (PyTorch shown; NumPy/DLPack also supported) + N = 1 << 20 + A = torch.randn(N, device='cuda', dtype=torch.float32) + B = torch.randn(N, device='cuda', dtype=torch.float32) + C = torch.empty(N, device='cuda', dtype=torch.float32) + + kernel = add(N) + kernel(A, B, C) # runs on GPU + torch.testing.assert_close(C, A + B) + print("done") + + +def call_tilelang(kernel_function, args, kwargs, grid, threads, params): + compiled_kernel = kernel_function(**kwargs) # cached, so second time only cache lookup is performed + compiled_kernel(*args) + +def tune(): + N = 1 << 20 + A = torch.randn(N, device='cuda', dtype=torch.float32) + B = torch.randn(N, device='cuda', dtype=torch.float32) + C = torch.empty(N, device='cuda', dtype=torch.float32) + + args = [A, B, C] + tune_params = dict() + tune_params["block"] = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] + tune_params["N"] = [N] # Not a tune param, but enables an easier call function + + answer = [None, None, (A + B).cpu()] + + res, env = tune_kernel("add", add, N, args, tune_params, lang="generic_python", + call_function=call_tilelang, decorator="@tilelang.jit", answer=answer) + +tune() \ No newline at end of file diff --git a/examples/generic_python/tilus_vec_add.py b/examples/generic_python/tilus_vec_add.py index dce785909..161575181 100644 --- a/examples/generic_python/tilus_vec_add.py +++ b/examples/generic_python/tilus_vec_add.py @@ -41,13 +41,10 @@ def call_tilus(kernel_function, args, kwargs, grid, threads, params): kernel_function(*args, **kwargs) -def main(): - - size = 1024000 - - a = torch.randn(size, dtype=torch.float32) - b = torch.randn(size, dtype=torch.float32) - c = torch.zeros_like(b) +def tune_vecadd(size): + a = torch.randn(size, dtype=torch.float32).cuda() + b = torch.randn(size, dtype=torch.float32).cuda() + c = torch.empty(size, dtype=torch.float32).cuda() c_expect = a + b @@ -55,8 +52,6 @@ def main(): tune_params = dict() tune_params["block_size_x"] = [32, 64, 128, 256, 512, 1024] - - results, env = tune_kernel( kernel_name="VecAddV", kernel_source=VecAddV, @@ -69,7 +64,50 @@ def main(): verbose=True, ) + +# TODO run kernel error handling same as tune_kernel +def run_vecadd(size): + a = torch.randn(size, dtype=torch.float32).cuda() + b = torch.randn(size, dtype=torch.float32).cuda() + c = torch.empty(size, dtype=torch.float32).cuda() + c_expect = a + b + args = [size, a, b, c] + + results = run_kernel( + kernel_name="VecAddV", + kernel_source=VecAddV, + problem_size=size, + arguments=args, + params={"block_size_x": 32}, + lang="generic_python", + call_function=call_tilus, + #verbose=True, + ) + + c_expect = c_expect.cpu() + + assert torch.allclose(results[-1], c_expect) + + + +def run_normal(size): + a = torch.randn(size, dtype=torch.float32).cuda() + b = torch.randn(size, dtype=torch.float32).cuda() + c = torch.empty(size, dtype=torch.float32).cuda() + c_expect = a + b + + vecadd = VecAddV() + vecadd(size, a, b, c) + print(c) + print(c_expect) + assert torch.allclose(c, c_expect) + + + if __name__ == "__main__": - main() + size = 1024 + tune_vecadd(size) + #run_vecadd(size) + #run_normal(size) diff --git a/examples/generic_python/triton_vec_add.py b/examples/generic_python/triton_vec_add.py index be28f2003..fcc019615 100644 --- a/examples/generic_python/triton_vec_add.py +++ b/examples/generic_python/triton_vec_add.py @@ -12,7 +12,7 @@ def add_op(x, y): return x + y -#triton.jit +#@triton.jit def add_kernel(x_ptr, # *Pointer* to first input vector. y_ptr, # *Pointer* to second input vector. output_ptr, # *Pointer* to output vector. @@ -50,11 +50,12 @@ def tune_with_generic(): tune_params = dict() tune_params["block_size_x"] = [2**i for i in range(10)] - ''' + result = run_kernel("add_kernel", add_kernel, size, args, {"block_size_x": 256}, lang="generic_python", call_function=call_triton, decorator="@triton.jit") print(np.allclose(c_expect.cpu(), result[2])) - ''' + + results, env = tune_kernel( kernel_name="add_kernel", @@ -69,6 +70,7 @@ def tune_with_generic(): ) + diff --git a/examples/generic_python/warp_vec_add.py b/examples/generic_python/warp_vec_add.py index dc08dd6d0..4a2c8d350 100644 --- a/examples/generic_python/warp_vec_add.py +++ b/examples/generic_python/warp_vec_add.py @@ -1,6 +1,9 @@ import warp as wp import numpy as np from kernel_tuner import tune_kernel, run_kernel +import torch + +wp.init() @@ -13,9 +16,9 @@ def add_op(x: float, y: float): def vec_add(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=float), - n: int, - work_per_thread: int): + n: int): + work_per_thread = 8 tid = wp.tid() base = tid * work_per_thread @@ -28,10 +31,14 @@ def vec_add(a: wp.array(dtype=float), # TODO do we allways want the call function to have the same parameters # or do we only require some of them? def call_warp(kernel_function, args, kwargs, grid, threads, params): - final_args = list(args) - final_args.extend(kwargs.values()) + warp_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + warp_args.append(wp.from_torch(arg)) + else: + warp_args.append(arg) dim = args[3] - wp.launch(kernel=kernel_function, dim=dim, inputs=final_args) + wp.launch(kernel=kernel_function, dim=dim, inputs=warp_args) # NOTE default verify function only works for numpy/cupy ndarray, torch Tensor or numpy scalar @@ -41,6 +48,8 @@ def verify(answer, result_host, atol): for i, ans in enumerate(answer): if ans is None: continue + print("res: ", type(result_host[i])) + print("expect: ", type(ans)) res = result_host[i].numpy() if not np.allclose(ans, res, atol=atol): correct = False @@ -53,19 +62,15 @@ def tune(): n = 1024 # Create host arrays - a_np = np.arange(n, dtype=np.float32) - b_np = np.arange(n, 0, -1, dtype=np.float32) - c_np = np.zeros(n, dtype=np.float32) - c_expect = a_np + b_np + a_torch = torch.arange(n, dtype=torch.float32, device="cuda") + b_torch = torch.arange(n, 0, -1, dtype=torch.float32, device="cuda") + c_torch = torch.zeros(n, dtype=torch.float32, device="cuda") + c_expect = a_torch + b_torch - # Create Warp arrays on GPU - a = wp.array(a_np, dtype=float) - b = wp.array(b_np, dtype=float) - c = wp.array(c_np, dtype=float) tune_params = dict() tune_params["work_per_thread"] = [2**i for i in range(10)] - args = [a, b, c, n] + args = [a_torch, b_torch, c_torch, n] ''' @@ -88,8 +93,7 @@ def tune(): arguments=args, tune_params=tune_params, lang="generic_python", - answer=[None, None, c_expect, None], - verify=verify, + answer=[None, None, c_expect.cpu(), None], call_function=call_warp, decorator="@wp.kernel" ) diff --git a/examples/triton/conv2d_tuning.py b/examples/triton/conv2d_tuning.py deleted file mode 100644 index 5389b4a28..000000000 --- a/examples/triton/conv2d_tuning.py +++ /dev/null @@ -1,306 +0,0 @@ -import torch -import triton.language as tl -import numpy as np -from kernel_tuner.interface import tune_kernel -import os -import json -from datetime import datetime - -# Check for required environment variable -cache_dir = os.getenv('KERNEL_TUNER_CACHE_DIR') -cache_file_name = os.getenv('KERNEL_TUNER_CACHE_FILE', 'conv2d_tuning_results.json') - -if cache_dir is None: - raise ValueError("Environment variable KERNEL_TUNER_CACHE_DIR must be set") - -cache_file = os.path.join(cache_dir, cache_file_name) - - -def conv2d_output_size( - in_size: int, - kernel_size: int, - stride: int, - padding: int, - dilation: int, -) -> int: - """ - Determines the output size of a 2D convolution operation. - - Args: - in_size: Input size. - kernel_size: Kernel size. - stride: Stride. - padding: Padding. - - Returns: - Output size of 2D convolution. - """ - return (in_size + 2 * padding - dilation * (kernel_size - 1) - 1) // stride + 1 - - -def conv2d_forward_kernel( - input_pointer, - weight_pointer, - output_pointer, - in_n, - input_height, - input_width, - out_c, - out_height, - out_width, - input_n_stride, - input_c_stride, - input_height_stride, - input_width_stride, - weight_n_stride, - weight_c_stride, - weight_height_stride, - weight_width_stride, - output_n_stride, - output_c_stride, - output_height_stride, - output_width_stride, - weight_c: tl.constexpr, - weight_height: tl.constexpr, - weight_width: tl.constexpr, - stride_height: tl.constexpr, - stride_width: tl.constexpr, - padding_height: tl.constexpr, - padding_width: tl.constexpr, - dilation_height: tl.constexpr, - dilation_width: tl.constexpr, - groups: tl.constexpr, - BLOCK_NI_HO_WO: tl.constexpr, - BLOCK_CI: tl.constexpr, - BLOCK_CO: tl.constexpr, -): - pid_ni_ho_wo = tl.program_id(0) - pid_co = tl.program_id(1) - pid_group = tl.program_id(2) - - # caculate in_n out_height out_weight value in kernel - ni_ho_wo_offset = pid_ni_ho_wo * BLOCK_NI_HO_WO + tl.arange(0, BLOCK_NI_HO_WO) - ni_ho_offset = ni_ho_wo_offset // out_width - in_n_point_value = ni_ho_offset // out_height - output_height_point_value = ni_ho_offset % out_height - output_width_point_value = ni_ho_wo_offset % out_width - - # Load the input and weight pointers. input and weight are of shape - # [in_n, groups, in_c, input_height, input_width] and [groups, out_c, in_c, weight_height, weight_width] - out_per_group_c = out_c // groups - output_c_offset = pid_co * BLOCK_CO + tl.arange(0, BLOCK_CO) - input_pointer += ( - input_n_stride * in_n_point_value + input_c_stride * pid_group * weight_c - )[:, None] - weight_pointer += ( - weight_n_stride * output_c_offset - + weight_n_stride * pid_group * out_per_group_c - )[None, :] - - accum = tl.zeros((BLOCK_NI_HO_WO, BLOCK_CO), dtype=tl.float32) - BLOCK_CI_COUNT = (weight_c + BLOCK_CI - 1) // BLOCK_CI - for hwc in range(weight_height * weight_width * BLOCK_CI_COUNT): - c = (hwc % BLOCK_CI_COUNT) * BLOCK_CI - hw = hwc // BLOCK_CI_COUNT - h = hw // weight_width - w = hw % weight_width - - input_c_offset = c + tl.arange(0, BLOCK_CI) - input_height_offset = ( - h * dilation_height - - padding_height - + stride_height * output_height_point_value - ) - input_width_offset = ( - w * dilation_width - padding_width + stride_width * output_width_point_value - ) - - curr_input_pointer = ( - input_pointer - + (input_c_stride * input_c_offset)[None, :] - + (input_height_stride * input_height_offset)[:, None] - + (input_width_stride * input_width_offset)[:, None] - ) - curr_weight_pointer = ( - weight_pointer - + (weight_c_stride * input_c_offset)[:, None] - + (weight_height_stride * h) - + (weight_width_stride * w) - ) - - input_mask = ( - (in_n_point_value < in_n)[:, None] - & (input_c_offset < weight_c)[None, :] - & (0 <= input_height_offset)[:, None] - & (input_height_offset < input_height)[:, None] - & (0 <= input_width_offset)[:, None] - & (input_width_offset < input_width)[:, None] - ) - weight_mask = (input_c_offset < weight_c)[:, None] & ( - output_c_offset < out_per_group_c - )[None, :] - - input_block = tl.load(curr_input_pointer, mask=input_mask) - weight_block = tl.load(curr_weight_pointer, mask=weight_mask) - - accum += tl.dot(input_block, weight_block, allow_tf32=False) - - output_pointer += ( - (output_n_stride * in_n_point_value)[:, None] - + (output_c_stride * (pid_group * out_per_group_c + output_c_offset))[None, :] - + (output_height_stride * output_height_point_value)[:, None] - + (output_width_stride * output_width_point_value)[:, None] - ) - output_mask = ( - (in_n_point_value < in_n)[:, None] - & (output_c_offset < out_per_group_c)[None, :] - & (output_height_point_value < out_height)[:, None] - & (output_width_point_value < out_width)[:, None] - ) - - tl.store(output_pointer, accum, mask=output_mask) - - -def tune_conv2d(batch_size=1, in_channels=64, height=32, width=32, - out_channels=128, kernel_size=3, stride=1, padding=1, - groups=1): - """ - Tune the conv2d kernel with different configurations. - """ - # Create sample inputs - input = torch.randn(batch_size, in_channels, height, width, - device='cuda', dtype=torch.float32) - weight = torch.randn(out_channels, in_channels//groups, kernel_size, kernel_size, - device='cuda', dtype=torch.float32) - - # Calculate output dimensions - out_height = conv2d_output_size(height, kernel_size, stride, padding, 1) - out_width = conv2d_output_size(width, kernel_size, stride, padding, 1) - output = torch.empty((batch_size, out_channels, out_height, out_width), - device='cuda', dtype=torch.float32) - - # Prepare all arguments for the kernel - arguments = [ - input, weight, output, - np.int32(batch_size), - np.int32(height), - np.int32(width), - np.int32(out_channels), - np.int32(out_height), - np.int32(out_width), - np.int32(input.stride(0)), - np.int32(input.stride(1)), - np.int32(input.stride(2)), - np.int32(input.stride(3)), - np.int32(weight.stride(0)), - np.int32(weight.stride(1)), - np.int32(weight.stride(2)), - np.int32(weight.stride(3)), - np.int32(output.stride(0)), - np.int32(output.stride(1)), - np.int32(output.stride(2)), - np.int32(output.stride(3)), - np.int32(in_channels//groups), # weight_c - np.int32(kernel_size), # weight_height - np.int32(kernel_size), # weight_width - np.int32(stride), # stride_height - np.int32(stride), # stride_width - np.int32(padding), # padding_height - np.int32(padding), # padding_width - np.int32(1), # dilation_height - np.int32(1), # dilation_width - np.int32(groups), # groups - ] - - # Define tuning parameters - only powers of 2 - tune_params = { - 'BLOCK_NI_HO_WO': [2 ** i for i in range(4, 10)], - 'BLOCK_CI': [2 ** i for i in range(4, 10)], - 'BLOCK_CO': [2 ** i for i in range(4, 10)], - 'num_stages': [1, 2, 3, 4], - 'num_warps': [1, 2, 4, 8], - } - - print(tune_params) - - # Define constraints - constraints = [ - "BLOCK_CI <= %d" % (in_channels//groups), - "BLOCK_CO <= %d" % out_channels, - ] - - # Problem size for the grid - problem_size = ( - batch_size * out_height * out_width, # Grid dimension 0 - out_channels, # Grid dimension 1 - groups, # Grid dimension 2 - ) - - # Grid divisor expressions - grid_div_x = ["BLOCK_NI_HO_WO"] - grid_div_y = ["BLOCK_CO"] - grid_div_z = ["1"] - - results, env = tune_kernel( - kernel_name='conv2d_forward_kernel', - kernel_source=conv2d_forward_kernel, - problem_size=problem_size, - arguments=arguments, - tune_params=tune_params, - restrictions=constraints, - lang='TRITON', - grid_div_x=grid_div_x, - grid_div_y=grid_div_y, - grid_div_z=grid_div_z, - block_size_names=['BLOCK_NI_HO_WO', 'BLOCK_CI', 'BLOCK_CO'], - strategy='genetic_algorithm', - strategy_options={ - 'maxiter': 1000, - 'popsize': 100, - }, - cache=cache_file, - ) - - return results - - -if __name__ == '__main__': - # Run tuning with moderately large input dimensions - results = tune_conv2d( - batch_size=16, - in_channels=128, - height=112, - width=112, - out_channels=256, - kernel_size=3, - stride=1, - padding=1, - groups=1 - ) - - - # Filter out failed compilations and find best config - valid_results = [result for result in results if isinstance(result['time'], (int, float))] - if valid_results: - best_config = min(valid_results, key=lambda x: x['time']) - print("\nBest configuration:") - print(json.dumps(best_config, indent=2)) - else: - print("\nNo valid configurations found - all compilations failed") - - # Create results dictionary with GPU info - all_results = { - "gpu_info": { - "gpu_name": torch.cuda.get_device_name() - }, - "results": valid_results - } - - # Add timestamp to filename - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - output_file = f'conv2d_results_{timestamp}.json' - - # Save results - import json - with open(output_file, 'w') as f: - json.dump(all_results, f, indent=2) \ No newline at end of file diff --git a/examples/triton/vec_add.py b/examples/triton/vec_add.py deleted file mode 100644 index 42c3e52c7..000000000 --- a/examples/triton/vec_add.py +++ /dev/null @@ -1,50 +0,0 @@ -import numpy -import triton.language as tl -import torch -from kernel_tuner import tune_kernel, run_kernel -from kernel_tuner.file_utils import store_output_file, store_metadata_file -import triton - -#@triton.jit -def add_kernel(x_ptr, # *Pointer* to first input vector. - y_ptr, # *Pointer* to second input vector. - output_ptr, # *Pointer* to output vector. - n_elements, # Size of the vector. - block_size_x: tl.constexpr, # Number of elements each program should process. - # note: `constexpr` so it can be used as a shape value. - ): - pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. - block_start = pid * block_size_x - offsets = block_start + tl.arange(0, block_size_x) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - tl.store(output_ptr + offsets, output, mask=mask) - - - -def tune_with_triton(): - size = 10000000 - - a = torch.randn(size, device='cuda', dtype=torch.float32) - b = torch.randn(size, device='cuda', dtype=torch.float32) - c = torch.empty_like(b) - n = torch.tensor(size, dtype=torch.int32) - - args = [a, b, c, n] - - tune_params = dict() - tune_params["block_size_x"] = [2**i for i in range(10)] - - results, env = tune_kernel( - kernel_name="add_kernel", - kernel_source=add_kernel, - problem_size=size, - arguments=args, - tune_params=tune_params, - lang="triton" - ) - - -tune_with_triton() \ No newline at end of file diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index fbe0f0750..5d867fdac 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -3,6 +3,9 @@ import copy import traceback # for compile error handling import re +import warnings +import builtins +import numpy as np from kernel_tuner.backends.backend import GPUBackend from kernel_tuner.observers.generic_python import GenericPythonRuntimeObserver @@ -13,22 +16,29 @@ torch = None -# TODO delete temp file +class GenericPythonFunctions(GPUBackend): + """Class that groups the Python DSL functions on maintains state about the device.""" + def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): + """Instantiate GenericPythonFunctions object used for interacting with the device. + Currently, only CUDA devices are supported. + Instantiating this object will inspect and store certain device properties at + runtime, which are used during compilation and/or execution of kernels by the + kernel tuner. Compiler options are not supported for GenericPython. -class GenericPythonFunctions(GPUBackend): + :param device: Number of CUDA device to use for this context + :type device: int + + :param iterations: Number of iterations used while benchmarking a kernel, 7 by default. + :type iterations: int + """ - def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): - ''' - In here, everyting is generic if the language uses CUDA as backend - ''' if not torch: logging.error("Unable to import Torch") raise ImportError("Unable to import Torch") self.device_id = torch.cuda.current_device() - self.device_properties = torch.cuda.get_device_properties(self.device_id) self.name = torch.cuda.get_device_name(self.device_id) self.max_threads = self.device_properties.max_threads_per_multi_processor @@ -52,7 +62,7 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None self.units = {"time": "ms", "power": "s,mW", "energy": "J"} - # Variables to be filled in at compile time: + # Variables to be filled in at compile time, needed for running the kernel after compilation. self.call_function = None self.signature = None self.gpu_kwargs = None @@ -60,15 +70,58 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) def ready_argument_list(self, arguments): - ''' - The user already supplies the arguments in the correct format, because we are working with - a Python based language anyway. TODO probably only works with torch and numpy? - ''' - return copy.deepcopy(arguments) + """Ready argument list to be passed to the kernel. Converts arguments to Torch GPU Tensors or Python + Scalars. Arguments of built-in Python types are left untouched. + + :param arguments: List of arguments to be passed to the kernel. + The order should match the argument list on the kernel. + Allowed values are: + - numpy.ndarray, and/or numpy.int32, numpy.float32, and so on. + - Torch Tensors + - Built-in Python types. + :type arguments: list + + :returns: A list of arguments that can be passed to the kernel. + :rtype: list( Torch.Tensor on GPU, int, float, bool, etc. ) + """ + torch_args = [] + + for arg in arguments: + if isinstance(arg, torch.Tensor): + if arg.dim() == 0: # Scalar tensor, convert to Python scalar + torch_args.append(arg.item()) + else: + if arg.is_cuda: # already on GPU, need deep copy to not overwrite + torch_args.append(arg.clone()) + else: # Copy from CPU to GPU + torch_args.append(arg.contiguous().to("cuda")) + elif isinstance(arg, np.ndarray): # Convert Numpy CPU array to Torch GPU Tensor + torch_args.append(torch.from_numpy(arg).to("cuda")) + elif isinstance(arg, np.generic): # Numpy scalar, convert to Python scalar + torch_args.append(arg.item()) + elif isinstance(arg, (int, float, bool, str)): # Is already Python + torch_args.append(arg) + elif type(arg) in vars(builtins).values(): + torch_args.append(arg) + else: + raise TypeError("Unknown argument type: ", type(arg)) + + return torch_args + def compile(self, kernel_instance, gpu_args=None): - logging.debug("Compiling Generic Python kernel") + """Compile the kernel by executing it once. This enforces that the kernel is cached. + + :param kernel_instance: The kernel instance containing information such as the kernel_source, + grid, threads, params, etc. + :type kernel_instance: KernelInstance + + :param gpu_args: arguments to be passed to the kernel. + :type gpu_args: list(any) + :returns: A kernel that can be called with the user-defined call function. + :rtype: callable + """ if kernel_instance.kernel_fn is None: raise ValueError("kernel_fn is None, currently Generic Python only supports callable kernel_source") @@ -76,6 +129,7 @@ def compile(self, kernel_instance, gpu_args=None): raise ValueError("gpu_args is None, Generic Python needs gpu args to compile the kernel") # The first time we compile, we also set the call function and the signature + # We need this later to run the kernel. if self.call_function is None or self.signature is None: self.call_function = kernel_instance.kernel_source.call_function self.signature = kernel_instance.kernel_source.signature @@ -83,92 +137,145 @@ def compile(self, kernel_instance, gpu_args=None): grid = kernel_instance.grid threads = kernel_instance.threads params = kernel_instance.params + + # If the kernel source is a class, we use the __call__ function as the kernel_function. if inspect.isclass(kernel_instance.kernel_fn): kernel_function = kernel_instance.kernel_fn() elif callable(kernel_instance.kernel_fn): - # Handles functions and decroators that return callable objects kernel_function = kernel_instance.kernel_fn else: raise TypeError("kernel function is not a class or function") - self.gpu_kwargs = self.build_gpu_kwargs(params) - - # Call the jit function in order to compile it + # Tuning params can contain kernel arguments. In such cases, create keyword arguments with + # the values of the tuning params. + self.gpu_kwargs = {} + if params is not None: + for name, p in self.signature.parameters.items(): + if name in params: + self.gpu_kwargs[name] = params[name] + + # Call the user-defined call function in order to compile the kernel. self.synchronize() self.call_function(kernel_function, gpu_args, self.gpu_kwargs, grid, threads, params) self.synchronize() - return kernel_function def start_event(self): - logging.debug("Start Generic Python event") + """Records the event that marks the start of a measurement.""" self.start.record() def stop_event(self): - logging.debug("Stop Generic Python event") + """Records the event that marks the end of a measurement.""" self.end.record() def kernel_finished(self): - logging.debug("Checking if kernel has finished") + """Returns True if the kernel has finished, False otherwise.""" return self.end.query() def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): + """Runs the Python kernel passed as 'func'. + + :param func: A cached Python kernel for this specific kernel configuration + :type func: callable + + :param gpu_args: A list of arguments to the kernel, order should match the + order in the code. + :type gpu_args: list(any) - # Run the kernel + :param threads: A tuple listing the number of threads in each dimension of + the thread block + :type threads: tuple(int, int, int) + + :param grid: A tuple listing the number of thread blocks in each dimension + of the grid + :type grid: tuple(int, int, int) + + :param params: A dictionary with the tuning params for this specific kernel + configuration + :type params: dict + """ if stream is None: stream = self.stream with torch.cuda.stream(stream): logging.debug("Running Generic Python kernel") self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) - - - - def build_gpu_kwargs(self, params=None): - gpu_kwargs = {} - - if params is None: - return gpu_kwargs - - for name, p in self.signature.parameters.items(): - if name in params: - gpu_kwargs[name] = params[name] - - return gpu_kwargs def synchronize(self): + """Halts execution until device has finished its tasks.""" torch.cuda.synchronize() + def memset(self, allocation, value, size): + """This method must implement setting the memory to a value on the device. + Not implemented: Python DSLs usually do not perform explicit memory handling.""" pass + def memcpy_dtoh(self, dest, src): + """This method must implement a device to host copy. + Not implemented: Python DSLs usually do not perform explicit memory handling.""" pass + def memcpy_htod(self, dest, src): + """This method must implement a host to device copy. + Not implemented: Python DSLs usually do not perform explicit memory handling.""" pass + def copy_constant_memory_args(self, cmem_args): raise NotImplementedError("Generic Python does not support constant memory") + def copy_shared_memory_args(self, smem_args): raise NotImplementedError("Generic Python does not support shared memory") + def copy_texture_memory_args(self, texmem_args): raise NotImplementedError("Generic Python does not support texture memory") + def refresh_memory(self, gpu_memory, host_arguments, should_sync): """Refresh the GPU memory with the untouched host arguments. We overwrite the standard function because Python DSLs do usually do not manage memory explicitely""" - for i, arg in enumerate(host_arguments): - if should_sync[i]: - gpu_memory[i] = copy.deepcopy(arg) + for i, host_arg in enumerate(host_arguments): + if should_sync[i]: + gpu_arg = gpu_memory[i] + + # Scalar Python type + if isinstance(gpu_arg, (int, float, bool)): + gpu_memory[i] = host_arg + elif type(gpu_arg) in vars(builtins).values(): + gpu_memory[i] = host_arg + + # GPU tensor + elif isinstance(gpu_arg, torch.Tensor): + if isinstance(host_arg, np.ndarray): + gpu_arg.copy_(torch.as_tensor(host_arg)) + elif isinstance(host_arg, torch.Tensor): + if host_arg.is_cuda and host_arg.device != gpu_arg.device: + gpu_arg.copy_(host_arg.to(gpu_arg.device)) # different gpu's, no direct copy allowed + else: + gpu_arg.copy_(host_arg) + else: + # host_arg is scalar, fill into tensor + gpu_arg.fill_(host_arg) + + def classify_compile_exception(self, e): - """Best effort to differentiate between a user error and a resource error. Input is Exception""" + """Best effort to differentiate between a user error and a resource error. + + :param e: the caught exception + :type e: exception + + :returns: "resource_error" , "user_error" or "unknown" + :rtype: string + """ RESOURCE_KEYWORDS = ( # Shared memory diff --git a/kernel_tuner/backends/triton.py b/kernel_tuner/backends/triton.py deleted file mode 100644 index cf6fda504..000000000 --- a/kernel_tuner/backends/triton.py +++ /dev/null @@ -1,174 +0,0 @@ -import logging -import numpy as np - -from kernel_tuner.backends.backend import GPUBackend -from kernel_tuner.observers.triton import TritonRuntimeObserver - -try: - import torch -except ImportError: - logging.error("Torch not available") - -try: - import triton - import triton.language as tl -except ImportError: - triton = None - tl = None - logging.error("Unable to load triton") - - -class TritonFunctions(GPUBackend): - - def __init__(self, device=0, iterations=7, compiler_options=None, observers=None): - if not triton or not torch: - logging.error("Triton or torch not available") - raise ImportError("Triton or torch not available") - - self.device_id = torch.cuda.current_device() - - self.device_properties = torch.cuda.get_device_properties(self.device_id) - self.name = torch.cuda.get_device_name(self.device_id) - self.max_threads = self.device_properties.max_threads_per_multi_processor - - env = dict() - env["device_name"] = self.name - env["max_threads"] = self.max_threads - env["iterations"] = iterations - env["compiler_options"] = compiler_options - self.env = env - - self.stream = torch.cuda.default_stream() - self.start = torch.cuda.Event(enable_timing=True) - self.end = torch.cuda.Event(enable_timing=True) - - # setup observers - self.observers = observers or [] - self.observers.append(TritonRuntimeObserver(self)) - for obs in self.observers: - obs.register_device(self) - - self.units = {"time": "ms", "power": "s,mW", "energy": "J"} - - super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) - - def ready_argument_list(self, arguments): - # Allocate memory here - torch_args = [] - - for arg in arguments: - if isinstance(arg, torch.Tensor) and arg.dim() > 0: - torch_args.append(arg.cuda()) - elif isinstance(arg, torch.Tensor) and arg.dim() == 0: - scalar_value = arg.item() - torch_args.append(scalar_value) - elif isinstance(arg, np.ndarray): - torch_arg = torch.from_numpy(arg) - torch_arg_gpu = torch_arg.cuda() - torch_args.append(torch_arg_gpu) - elif isinstance(arg, np.generic): - scalar_value = arg.item() - torch_args.append(scalar_value) - else: - logging.warning("Unknown instance in triton functions") - - return torch_args - - def compile(self, kernel_instance, gpu_args=None): - logging.debug("Compiling triton kernel") - - if kernel_instance.kernel_fn is None: - raise ValueError("kernel_fn is None, currently Triton only supports callable kernel_source") - - if gpu_args is None: - raise ValueError("gpu_args is None, Triton needs gpu args to compile the kernel") - - grid = kernel_instance.grid - threads = kernel_instance.threads - jit_function = triton.jit(kernel_instance.kernel_fn) - params = kernel_instance.params - gpu_kwargs = self.build_gpu_kwargs(jit_function, threads, params) - - # Call the jit function in order to compile it - jit_function[grid](*gpu_args, **gpu_kwargs) - - return jit_function - - def start_event(self): - logging.debug("Start triton event") - self.start.record() - - def stop_event(self): - logging.debug("Stop triton event") - self.end.record() - - def kernel_finished(self): - logging.debug("Checking if kernel has finished") - return self.end.query() - - def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): - if params is None: - raise ValueError("params is None, Triton needs params in order to set num_warps, num_ctas, etc.") - - # Run the kernel - if stream is None: - stream = self.stream - - gpu_kwargs = self.build_gpu_kwargs(func, threads, params) - - with torch.cuda.stream(stream): - logging.debug("Running triton kernel") - func[grid](*gpu_args, **gpu_kwargs) - - def build_gpu_kwargs(self, jit_fn, threads, params=None): - gpu_kwargs = {} - - if 'BLOCK_SIZE' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE'] = threads[0] - - if 'BLOCK_SIZE_X' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE_X'] = threads[0] - - if 'BLOCK_SIZE_Y' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE_Y'] = threads[1] - - if 'BLOCK_SIZE_Z' in jit_fn.arg_names: - gpu_kwargs['BLOCK_SIZE_Z'] = threads[2] - - if params is None: - return gpu_kwargs - - for param in params: - if param in jit_fn.arg_names: - gpu_kwargs[param] = params[param] - - # Check for Triton specific parameters - if 'num_warps' in params: - gpu_kwargs['num_warps'] = params['num_warps'] - if 'num_ctas' in params: - gpu_kwargs['num_ctas'] = params['num_ctas'] - if 'num_stages' in params: - gpu_kwargs['num_stages'] = params['num_stages'] - - return gpu_kwargs - - def synchronize(self): - torch.cuda.synchronize() - - def memset(self, allocation, value, size): - pass - - def memcpy_dtoh(self, dest, src): - pass - - def memcpy_htod(self, dest, src): - pass - - def copy_constant_memory_args(self, cmem_args): - raise NotImplementedError("Triton does not support constant memory") - - def copy_shared_memory_args(self, smem_args): - raise NotImplementedError("Triton does not support shared memory") - - def copy_texture_memory_args(self, texmem_args): - raise NotImplementedError("Triton does not support texture memory") \ No newline at end of file diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 7705d828c..1f130afa3 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -21,7 +21,6 @@ from kernel_tuner.backends.nvcuda import CudaFunctions from kernel_tuner.backends.opencl import OpenCLFunctions from kernel_tuner.backends.pycuda import PyCudaFunctions -from kernel_tuner.backends.triton import TritonFunctions from kernel_tuner.backends.generic_python import GenericPythonFunctions from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.observers.nvml import NVMLObserver @@ -86,7 +85,7 @@ def prepare_temp_files_for_error_msg(self): return ret -# KernelSource has been moved to kernel_sources/kernel_source_str.py +# KernelSource has been moved to kernel_sources/kernel_source_str.py to support Generic Python class DeviceInterface(object): @@ -186,13 +185,6 @@ def __init__( compiler_options=compiler_options ) self.requires_warmup = False - elif lang == Language.TRITON: - dev = TritonFunctions( - device, - compiler_options=compiler_options, - iterations=iterations, - observers=observers - ) elif lang == Language.GENERIC_PYTHON: dev = GenericPythonFunctions( device, @@ -202,7 +194,7 @@ def __init__( ) else: raise NotImplementedError( - "Sorry, support for languages other than CUDA, OpenCL, HIP, C, Triton and Fortran is not implemented yet" + "Sorry, support for languages other than CUDA, OpenCL, HIP, C, Fortran and Generic Python is not implemented yet" ) self.dev = dev @@ -242,7 +234,7 @@ def __init__( print("Using: " + self.dev.name) def run_kernel_bench(self, func, gpu_args, threads, grid, stream=None, params=None): - if isinstance(self.dev, TritonFunctions) or isinstance(self.dev, GenericPythonFunctions): + if isinstance(self.dev, GenericPythonFunctions): # Generic Python needs params to compile. self.dev.run_kernel(func, gpu_args, threads, grid, params=params) else: self.dev.run_kernel(func, gpu_args, threads, grid) @@ -403,7 +395,7 @@ def check_kernel_output( # retrieve gpu results to host memory result_host = [] if instance.kernel_source is not None and instance.kernel_source.lang == Language.GENERIC_PYTHON: - # Python DSLs do not explicitly manage memory. Therefore, we can use gpu args direclty + # arguments are either Torch tensors or built-in Python types for i, arg in enumerate(gpu_args): if should_sync[i]: if isinstance(arg, torch.Tensor): @@ -530,15 +522,12 @@ def compile_kernel(self, instance, verbose, gpu_args=None): # compile kernel_string into device func func = None - if isinstance(self.dev, TritonFunctions) or isinstance(self.dev, GenericPythonFunctions): + # Python DSLs have different error handling, so we make a case distinction. + if isinstance(self.dev, GenericPythonFunctions): try: func = self.dev.compile(instance, gpu_args) except Exception as e: - if isinstance(self.dev, GenericPythonFunctions): - exception_type = self.dev.classify_compile_exception(e) - else: - exception_type = "unkown" - + exception_type = self.dev.classify_compile_exception(e) if exception_type == "user_error": error_message = str(e.stderr) if hasattr(e, "stderr") else str(e) print("compile_kernel failed due to error: " + error_message) @@ -550,10 +539,9 @@ def compile_kernel(self, instance, verbose, gpu_args=None): ) if verbose: print( - f"skipping config {util.get_instance_string(instance.params)} reason: too many resources" + f"skipping config {util.get_instance_string(instance.params)} reason: \n{e}" ) - - else: + else: # All other languages try: func = self.dev.compile(instance) except Exception as e: @@ -617,7 +605,7 @@ def create_kernel_instance(self, kernel_source, kernel_options, params, verbose) kernel_options.block_size_names, ) - if kernel_source.lang not in [Language.TRITON, Language.GENERIC_PYTHON] and np.prod(threads) > self.dev.max_threads: + if kernel_source.lang is not Language.GENERIC_PYTHON and np.prod(threads) > self.dev.max_threads: if verbose: print(f"skipping config {util.get_instance_string(params)} reason: too many threads per block") return util.InvalidConfig() diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index ce9d74080..e0421f847 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -108,9 +108,11 @@ def __deepcopy__(self, _): ( "kernel_source", ( - """The CUDA, OpenCL, HIP, or C kernel code. + """The CUDA, OpenCL, HIP, C or Python DSL kernel code. It is allowed for the code to be passed as a string, a filename, a function that returns a string of code, or a list when the code needs auxilliary files. + In the case of a kernel in a Python DSL such as Triton, the reference to the + Python callable should be passed. To support combined host and device code tuning, a list of filenames can be passed. The first file in the list should be the @@ -134,7 +136,7 @@ def __deepcopy__(self, _): """Specifies the language used for GPU kernels. The kernel_tuner automatically detects the language, but if it fails, you may specify the language using this argument, currently supported: "CUDA", "CuPy", - "nvcuda", "OpenCL", "HIP", or "C".""", + "nvcuda", "OpenCL", "HIP", "C", or "Generic_Python.""", "string", ), ), @@ -181,7 +183,9 @@ def __deepcopy__(self, _): "arguments", ( """A list of kernel arguments, use numpy arrays for - arrays, use numpy.int32 or numpy.float32 for scalars.""", + arrays, use numpy.int32 or numpy.float32 for scalars. When the language + Generic Python is used, Torch Tensors and built-in Python data types are also + excepted""", "list", ), ), @@ -285,6 +289,31 @@ def __deepcopy__(self, _): "dict", ), ), + ( + "call_function", + ( + """When the language Generic Python is used, a call function that calls the kernel in the Python + DSL must be specified. The function must take the following arguments: + :kernel_function: the callable function with the tuning parameters inserted. If provided, the + kernel_function is decorated with the de decorator. + :args: list of kernel arguments, as provided by the user in the argument. + :kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, + the tuning parameter will be added as a keyword argument. + :grid: the launch grid (tuple with 3 values), as computed by KernelTuner + :threads: the thread block size (tuple with 3 values), as computed by KernelTuner + :params: dictionary with the values of the tuning params for the specific configuration.""", + "function", + ), + ), + ( + "decorator", + ( + """When the language Generic Python is used, a decorator can be provided in which the kernel source + will be wrapped internally by KernelTuner. Note that when passing the kernel to KernelTuner with the + ``kernel_source`` argument, the decorator should be removed from the kernel.""", + "string", + ), + ), ] ) diff --git a/kernel_tuner/kernel_sources/kernel_source.py b/kernel_tuner/kernel_sources/kernel_source.py index 9a16d9b6c..d18ecb284 100644 --- a/kernel_tuner/kernel_sources/kernel_source.py +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -1,4 +1,3 @@ -import inspect import kernel_tuner.util as util from abc import abstractmethod @@ -6,15 +5,20 @@ from kernel_tuner.language import Language from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData -# We use this pattern because we would otherwise get an init twice -# We create a kernelsouce by calling KernelSource(...). The call method for KernelSource is -# replaced by the call method in KernelSourceFactory, because KernelSource uses that class as -# metaclass. Inside the call method, a call to create either KernelSourceStr or KernelSourceFN -# is done. This triggers the __init__ call of those subclasses. In both subclasses, we call super().__init__ -# which initializes some sublclass specific variables and triggers the __init__ call of the KernelSource class. + class KernelSourceFactory(type): + ''' + Factory to dynamically determine if a KernelSource should be a KernelSourceStr or a KernelSourceFn. + if lang=generic_python, KernelSourceFn will be used. In all other cases, an instance of KernelSourceStr + is created. + + We create a kernelsouce by calling KernelSource(...). The call method for KernelSource is replaced by + the call method in KernelSourceFactory, because KernelSource uses that class as metaclass. Inside the + call method, a call to create either KernelSourceStr or KernelSourceFn is performed. This triggers the + __init__ call of the corresponding subclass. In both subclasses, we call super().__init__, which triggers + the __init__ call of the KernelSource class to initalize some common variables. + ''' def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function=None, decorator=None): - """Factory behavior""" if lang == None: language = None else: @@ -23,9 +27,8 @@ def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function except ValueError: raise TypeError(f"Supported languages are {[l.value for l in Language]}") - # Determine if we need to create a KernelSourceStr or a KernelSourceFn - if (language and (language == Language.TRITON or language == Language.GENERIC_PYTHON)): + if (language and language == Language.GENERIC_PYTHON): ks_str = False else: ks_str = True @@ -44,7 +47,8 @@ def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function else: return super().__call__(kernel_name, kernel_sources, lang, defines, call_function, decorator) -# TODO do we really want the Language enum? it's a lot of changes + + class KernelSource(metaclass=KernelSourceFactory): def __init__(self, kernel_name, kernel_sources, lang, defines=None): @@ -59,11 +63,7 @@ def __init__(self, kernel_name, kernel_sources, lang, defines=None): raise TypeError("Please specify language when using a code generator function") kernel_string = self.get_kernel_string(0) language = util.detect_language(kernel_string) - try: - self.lang = Language(language.upper()) - except ValueError: - # TODO this should never happen - raise TypeError(f"Supported languages are {[l.value for l in Language]}, found {language}") + self.lang = Language(language.upper()) else: try: self.lang = Language(lang.upper()) diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index edb75f95d..386685377 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -3,6 +3,7 @@ import copy import uuid import sys +import logging import astor import tempfile @@ -17,154 +18,66 @@ class KernelSourceFn(KernelSource): + """ + Class that holds the Python-function-based kernel sources. + + There is a primary kernel source for function-based kernels in Python. The source must be + a callable. This can be a function of a class. The function should not be decorated, but + a decorator can be supplied to wrap the function. + + A call function to specify how the kernel should be launched must be supplied. The call + function must take the following arguments: + - kernel_function: the callable function with the tuning parameters inserted. If provided, the + kernel_function is decorated with the de decorator. + - args: list of kernel arguments, as provided by the user in the argument. + - kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, + the tuning parameter will be added as a keyword argument. + - grid: the launch grid (tuple with 3 values), as computed by KernelTuner + - threads: the thread block size (tuple with 3 values), as computed by KernelTuner + - params: dictionary with the values of the tuning params for this configuration. + """ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function=None, decorator=None): super().__init__(kernel_name, kernel_source, lang, defines) if isinstance(kernel_source, list): raise ValueError("KernelSourceFn only supports a single kernel source function") - - try: - self.lang = Language(lang.upper()) - except ValueError: - raise TypeError(f"Supported languages are {[l.value for l in Language]}") - - if call_function is None: - raise ValueError("call_function must be supplied for language Generic Python") - if not callable(call_function): - raise TypeError(f"call_function {call_function} is not a callable object.") + + if self.lang == Language.GENERIC_PYTHON: + if call_function is None: + raise ValueError("call_function must be supplied for language Generic Python") + if not callable(call_function): + raise TypeError(f"call_function of type {type(call_function)} is not a callable object.") self.call_function = call_function # TODO ceck signature if decorator: if not isinstance(decorator, str): - raise TypeError(f"{decorator} is not a decorator") + raise TypeError(f"The decorator should be a string, got {type(decorator)} instead.") if decorator[0] != '@': - raise ValueError(f"{decorator} is not a valid decorator") + raise ValueError(f"The decorator should start with a '@', got {decorator} instead.") self.decorator = decorator - self.source_kernel_fn = kernel_source - self.kernel_fn = self.source_kernel_fn - self.signature = inspect.signature(kernel_source) + self.source_kernel_fn = kernel_source # This kernel source remains the original object + self.kernel_fn = self.source_kernel_fn # This is the kernel source that we will modify + try: - self.source = inspect.getsource(kernel_source) + self.source = inspect.getsource(self.source_kernel_fn) except TypeError as e: raise TypeError( f"{e}. Did you forget to remove a decorator before tuning?" ) from e - self.source_tree = ast.parse(self.source) - self.import_nodes = self._find_import_nodes(inspect.getfile(kernel_source)) - - # Find the module where the kernel function is defined - self.module = inspect.getmodule(kernel_source) - # Get dependencies by analyzing the AST - if self.lang == Language.TRITON: - self.dependencies = self._find_triton_dependencies() - else: - self.dependencies = self._find_dependencies() - - - - - def _find_import_nodes(self, source_file): - with open(source_file, "r") as f: - tree = ast.parse(f.read(), filename=source_file) - - import_nodes = [] - for node in tree.body: - if isinstance(node, (ast.Import, ast.ImportFrom)): - import_nodes.append(node) - - return import_nodes - - def _find_function_dependencies(self): - """Find all function calls in the kernel.""" - class FunctionCallVisitor(ast.NodeVisitor): - def __init__(self): - self.called_functions = set() - - def visit_Call(self, node): - if isinstance(node.func, ast.Name): - self.called_functions.add(node.func.id) - self.generic_visit(node) - - visitor = FunctionCallVisitor() - visitor.visit(self.source_tree) - return visitor.called_functions - - def _is_triton_jit_function(self, node): - """Check if a function has the @triton.jit decorator. Is triton specific""" - if not isinstance(node, ast.FunctionDef): - return False - - for decorator in node.decorator_list: - if (isinstance(decorator, ast.Name) and decorator.id == 'jit' or - isinstance(decorator, ast.Attribute) and decorator.attr == 'jit' or - isinstance(decorator, ast.Call) and - isinstance(decorator.func, ast.Attribute) and - decorator.func.attr == 'jit'): - return True - return False - - def _find_triton_dependencies(self): - """Find all Triton JIT functions that are called by the kernel. Is Triton specific""" - dependencies = {} - called_functions = self._find_function_dependencies() - - # Parse the entire module to find Triton JIT functions - module_source = inspect.getsource(self.module) - module_tree = ast.parse(module_source) - for node in module_tree.body: - if (self._is_triton_jit_function(node) and - node.name in called_functions and - node.name != self.kernel_name): # Skip the main kernel itself - dependencies[node.name] = node - - return dependencies - - def _find_function_dependecies2(self, tree, local_funcs): - # Non triton specific - class FunctionCallVisitor(ast.NodeVisitor): - def __init__(self): - self.called = set() - - def visit_Call(self, node): - if isinstance(node.func, ast.Name): - name = node.func.id - if name in local_funcs: - self.called.add(name) - self.generic_visit(node) - - visitor = FunctionCallVisitor() - visitor.visit(tree) - return visitor.called - - def _find_local_functions(self, tree): - local_funcs = set() - for node in tree.body: - if isinstance(node, ast.FunctionDef): - local_funcs.add(node.name) - return local_funcs - - def _find_dependencies(self): - source_file = inspect.getfile(self.source_kernel_fn) - with open(source_file, "r") as f: - source_code = f.read() - tree = ast.parse(source_code, filename=source_file) - local_funcs = self._find_local_functions(tree) - called_functions = self._find_function_dependecies2(self.source_tree, local_funcs) - - dependencies = {} - for node in tree.body: - if (isinstance(node, ast.FunctionDef) and node.name in called_functions and - node.name != self.kernel_name): # Skip the main kernel itself - dependencies[node.name] = node + self.signature = inspect.signature(self.source_kernel_fn) + self.source_tree = ast.parse(self.source) + self.import_nodes = self._find_import_nodes(inspect.getfile(self.source_kernel_fn)) + self.dependencies = self._find_dependencies() - return dependencies - def _add_decorator(self, function): - pass - def prepare_kernel_instance(self, kernel_options, params, grid, threads): + ''' + Given the dictionary of tuning parameter values for this configuration, + generate a kernel instance with these tuning parameters inserted. kernel_options, + grid and threads are not needed for Python kernels. + ''' new_kernel_fn, temp_file_path = self.apply_params_to_source_fn(params) self.kernel_fn = new_kernel_fn @@ -175,13 +88,26 @@ def prepare_kernel_instance(self, kernel_options, params, grid, threads): kernel_str=None ) + def check_argument_lists(self, kernel_name, arguments): + ''' + Check if the kernel arguments have the correct types. + Not implemented for Python, because type hinting is not always supplied. + ''' + logging.debug("Checking of arguments list not supported yet for Python kernels") return True + def apply_params_to_source_fn(self, params): + ''' + Create a module with the kernel imports and local dependencies from the kernel source file. + Find instances of the tuning parameters in the kernel and local dependencies and replace these + instances by the values of this configuration. + Return the new kernel and the path to the new module, where the new kernel lives. + ''' transformer = ReplaceVars(params) - # Create a new module with all necessary functions + # Create a new module to store the transformed AST in. new_module = ast.Module(body=[], type_ignores=[]) # Add imports @@ -190,11 +116,10 @@ def apply_params_to_source_fn(self, params): # Add dependencies (functions) for dep_node in self.dependencies.values(): dep_node_copy = copy.deepcopy(dep_node) - transformed_dep_node = transformer.visit(dep_node_copy) - #new_module.body.append(copy.deepcopy(dep_node)) + transformed_dep_node = transformer.visit(dep_node_copy) # apply tuning params to dependencies new_module.body.append(transformed_dep_node) - # Add transformed main kernel + # Transform main kernel source_tree_copy = copy.deepcopy(self.source_tree) transformed_tree = transformer.visit(source_tree_copy) @@ -207,39 +132,133 @@ def apply_params_to_source_fn(self, params): node.decorator_list.insert(0, decorator_node) break # only apply to the top level function + # Add transformed main kernel to new module new_module.body.extend(transformed_tree.body) # Fix locations and generate source ast.fix_missing_locations(new_module) new_source = astor.to_source(new_module) - #print(new_source) - - # Create a unique module name + # Create a unique module name and write new source to it. module_name = f'temp_kernel_module_{uuid.uuid4().hex}' - with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file: temp_file.write(new_source) temp_file_path = temp_file.name - spec = importlib.util.spec_from_file_location(module_name, temp_file_path) - temp_module = importlib.util.module_from_spec(spec) - # Register the module in sys.modules before executing it + spec = importlib.util.spec_from_file_location(module_name, temp_file_path) + temp_module = importlib.util.module_from_spec(spec) sys.modules[module_name] = temp_module spec.loader.exec_module(temp_module) new_fn = getattr(temp_module, self.kernel_name) return new_fn, temp_file_path + + def _find_import_nodes(self, source_file): + ''' + Parse kernel source file to find import statements. Return those + statements as AST import nodes. + ''' + with open(source_file, "r") as f: + tree = ast.parse(f.read(), filename=source_file) + + import_nodes = [] + for node in tree.body: + if isinstance(node, (ast.Import, ast.ImportFrom)): + import_nodes.append(node) + + return import_nodes + + + def _find_function_dependecies(self, tree, local_funcs): + ''' + Given an AST and a list of local function names, return the funtions that are invoked + somewhere in the file, and are in the local_funcs list. + ''' + class FunctionCallVisitor(ast.NodeVisitor): + def __init__(self): + self.called = set() + + def visit_Call(self, node): + if isinstance(node.func, ast.Name): + name = node.func.id + if name in local_funcs: + self.called.add(name) + self.generic_visit(node) + + visitor = FunctionCallVisitor() + visitor.visit(tree) + return visitor.called + + + def _find_dependencies(self): + ''' + Find all local function dependencies in the file where the kernel source is + defined. Return a dicionary indexed by the function node name with the function + body as value. + ''' + source_file = inspect.getfile(self.source_kernel_fn) + with open(source_file, "r") as f: + source_code = f.read() + tree = ast.parse(source_code, filename=source_file) + + # Find the locally defined functions in the source file + local_funcs = set() + for node in tree.body: + if isinstance(node, ast.FunctionDef): + local_funcs.add(node.name) + + # Given source file AST and the set of locally defined functions, find the names of + # the local functions that are invoked somewhere in the AST. + called_functions = self._find_function_dependecies(self.source_tree, local_funcs) + + # Traverse the tree one more time to save the dependency functions in a dictionary. + # Skip the main kernel. + dependencies = {} + for node in tree.body: + if (isinstance(node, ast.FunctionDef) and node.name in called_functions and + node.name != self.kernel_name): # Skip the main kernel itself + dependencies[node.name] = node + + return dependencies + + def __del__(self): - # Clean up temporary modules when the instance is destroyed + ''' + Clean up temporary modules when the instance is destroyed + ''' for key in list(sys.modules.keys()): if key.startswith('temp_kernel_module_'): del sys.modules[key] class ReplaceVars(ast.NodeTransformer): + ''' + AST transformer that replaces occurrences of tuning parameters with constant values. + The tuning parameters (params) should be provided as a dictionary. + The transformation supports the following cases: + + - Variable reads: + occurrence of a variable name that matches a key in `params` are replaced if the + value is being read. + + - Object attribute reads: + expressions of the form `self.` are replaced with constants if + `` exists in `params` and is being read (not assigned to). + + - Assignments to parameters: + assignments like `param = ...` or `self.param = ...` are overridden, + replacing the right-hand side with the constant value from `params`. + + Examples: + Given `params = {"block_size": 32}`: + x = 2 * block_size -> x = 2 * 32 + x = 2 * self.block_size -> x = 2 * 32 + self.block_size = 64 -> self.block_size = 32 + block_size = 64 -> block_size = 32 + + ''' def __init__(self, params: dict): self.params = params @@ -252,6 +271,7 @@ def visit_Name(self, node: ast.Name) -> Any: ) return node + def visit_Attribute(self, node: ast.Attribute): self.generic_visit(node) @@ -260,7 +280,7 @@ def visit_Attribute(self, node: ast.Attribute): isinstance(node.value, ast.Name) and node.value.id == "self" and node.attr in self.params - and isinstance(node.ctx, ast.Load) # <- check context + and isinstance(node.ctx, ast.Load) ): return ast.copy_location(ast.Constant(self.params[node.attr]), node) diff --git a/kernel_tuner/kernel_sources/kernel_source_str.py b/kernel_tuner/kernel_sources/kernel_source_str.py index 4db0831a7..095417392 100644 --- a/kernel_tuner/kernel_sources/kernel_source_str.py +++ b/kernel_tuner/kernel_sources/kernel_source_str.py @@ -8,7 +8,7 @@ class KernelSourceStr(KernelSource): - """Class that holds the kernel sources. + """Class that holds the string-based kernel sources. There is a primary kernel source for string-based kernels., which can be either a source string, a filename (indicating a file containing the kernel source code), diff --git a/kernel_tuner/language.py b/kernel_tuner/language.py index b129e04d5..02948c12b 100644 --- a/kernel_tuner/language.py +++ b/kernel_tuner/language.py @@ -5,7 +5,6 @@ class Language(Enum): OPENCL = "OPENCL" C = "C" HIP = "HIP" - TRITON = "TRITON" FORTRAN = "FORTRAN" NVCUDA = "NVCUDA" GENERIC_PYTHON = "GENERIC_PYTHON" diff --git a/kernel_tuner/observers/triton.py b/kernel_tuner/observers/triton.py deleted file mode 100644 index 878e8f27c..000000000 --- a/kernel_tuner/observers/triton.py +++ /dev/null @@ -1,32 +0,0 @@ -import numpy as np - -from kernel_tuner.observers.observer import BenchmarkObserver - -try: - import torch -except (ImportError, RuntimeError): - torch = None - - -class TritonRuntimeObserver(BenchmarkObserver): - """Observer that measures time using CUDA events during benchmarking.""" - - def __init__(self, dev): - if torch is None: - raise ImportError("Unable to load torch") - - self.dev = dev - self.stream = dev.stream - self.start = dev.start - self.end = dev.end - self.times = [] - - def after_finish(self): - # Time is measured in milliseconds - event_elapsed_time = self.start.elapsed_time(self.end) - self.times.append(event_elapsed_time) - - def get_results(self): - results = {"time": np.average(self.times), "times": self.times.copy()} - self.times = [] - return results \ No newline at end of file diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py index a3d250b60..eee0d3dd1 100644 --- a/test/test_generic_python_functions.py +++ b/test/test_generic_python_functions.py @@ -109,7 +109,7 @@ def test_gpu_kwargs(): gpu_args = dev.ready_argument_list(args) callable_fn = dev.compile_kernel(instance, verbose=False, gpu_args=gpu_args) - kwargs = dev.dev.build_gpu_kwargs(params) + kwargs = dev.dev.gpu_kwargs assert kwargs["mock_param"] == 64 assert callable_fn(*args, **kwargs) == 64 diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py index e23a41a15..85aa4e3eb 100644 --- a/test/test_kernel_source_fn.py +++ b/test/test_kernel_source_fn.py @@ -61,7 +61,7 @@ def test_initiation(): with pytest.raises(ValueError, match=r"call_function must be supplied for language .*"): KernelSource("mock_kernel", mock_kernel, "generic_python") - with pytest.raises(TypeError, match=r".* is not a callable object"): + with pytest.raises(TypeError, match=r".* Did you forget to remove a decorator before tuning\?"): KernelSource("mock_kernel", "This is a string Kernel", "generic_python", call_function=call_mock) with pytest.raises(ValueError, match=r"KernelSourceFn only supports a single kernel source function"): @@ -70,10 +70,10 @@ def test_initiation(): with pytest.raises(TypeError, match=r".* is not a callable object"): KernelSource("mock_kernel", mock_kernel, "generic_python", call_function="not a function") - with pytest.raises(ValueError, match=r".* is not a valid decorator"): + with pytest.raises(ValueError, match=r"The decorator should start with a '@', got .* instead."): KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator="not a decorator") - with pytest.raises(TypeError, match=r".* is not a decorator"): + with pytest.raises(TypeError, match=r"The decorator should be a string, got .* instead."): KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator=mock_kernel) From 81340389cb4c41321f24894a00d81db2585c99aa Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Thu, 5 Mar 2026 12:54:35 +0100 Subject: [PATCH 05/44] support for decorated kernels without having to remove the decorator --- .../generic_python/matmul/tilus_matmul.py | 1 - examples/generic_python/numba_vec_add.py | 88 +++++-------- examples/generic_python/tilelang_vec_add.py | 14 +- examples/generic_python/tilus_naive_matmul.py | 5 +- .../generic_python/tilus_splitk_matmul.py | 26 +--- .../generic_python/tilus_tunable_precision.py | 120 ------------------ examples/generic_python/tilus_vec_add.py | 36 ++---- examples/generic_python/triton_vec_add.py | 37 +++--- examples/generic_python/warp_vec_add.py | 63 +++------ kernel_tuner/backends/generic_python.py | 15 +-- kernel_tuner/interface.py | 21 +-- kernel_tuner/kernel_sources/kernel_source.py | 6 +- .../kernel_sources/kernel_source_fn.py | 74 +++++------ kernel_tuner/util.py | 78 +++++++++++- test/test_generic_python_functions.py | 27 +++- test/test_kernel_source_fn.py | 50 +++++--- 16 files changed, 258 insertions(+), 403 deletions(-) delete mode 100644 examples/generic_python/tilus_tunable_precision.py diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py index 07532013f..e8996c201 100644 --- a/examples/generic_python/matmul/tilus_matmul.py +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -88,7 +88,6 @@ def __call__( self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) - class MatmulGroupedOrdering(tilus.Script): def __init__(self): super().__init__() diff --git a/examples/generic_python/numba_vec_add.py b/examples/generic_python/numba_vec_add.py index e49f9e372..9f8dfe095 100644 --- a/examples/generic_python/numba_vec_add.py +++ b/examples/generic_python/numba_vec_add.py @@ -1,10 +1,13 @@ from numba import cuda -from math import ceil -from kernel_tuner import tune_kernel, run_kernel import torch +from kernel_tuner import tune_kernel, run_kernel +import numpy as np +from pathlib import Path + +FULL_PATH = Path(__file__).resolve() -#@cuda.jit +@cuda.jit def f(a, b, c): tid = cuda.grid(1) size = len(c) @@ -23,64 +26,31 @@ def call_numba(kernel_function, args, kwargs, grid, threads, params): kernel_function[grid, threads](*args, **kwargs) -def verify(answer, result_host, atol): - correct = True - for i, ans in enumerate(answer): - if ans is None: - continue - res = result_host[i].copy_to_host() - if not np.allclose(ans, res, atol=atol): - correct = False - return correct +def tune(): + + N = 100000 + + a = np.random.random(N) + b = np.random.random(N) + c = np.zeros(N) + c_expect = a + b + + args = [a, b, c] + tune_params = {"block_size_x": [2**i for i in range(10)]} -import numpy as np -N = 100000000 - -a = np.random.random(N) -b = np.random.random(N) -c = np.zeros(N) -c_expect = a + b -#c_expect = a.copy_to_host() + b.copy_to_host() - - -#a_torch = torch.rand(N, dtype=torch.float32) -#b_torch = torch.rand(N, dtype=torch.float32) -#c_torch = torch.zeros(N, dtype=torch.float32) -#c_expect = (a_torch + b_torch).cpu() - -args = [a, b, c] -tune_params = {"block_size_x": [2**i for i in range(10)]} - -''' -results = run_kernel( - kernel_name="f", - kernel_source=f, - problem_size=N, - arguments=args, - params={"block_size_x": 32}, - lang="generic_python", - call_function=call_numba, - decorator="@cuda.jit" -) -''' - - -#print(np.allclose(results[2], c_expect)) - - -results, env = tune_kernel( - kernel_name="f", - kernel_source=f, - problem_size=N, - arguments=args, - tune_params=tune_params, - lang="generic_python", - answer=[None, None, c_expect], - #verify=verify, - call_function=call_numba, - decorator="@cuda.jit" -) + results, env = tune_kernel( + kernel_name="f", + kernel_source=FULL_PATH, + problem_size=N, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, c_expect], + call_function=call_numba, + ) +if __name__ == "__main__": + tune() \ No newline at end of file diff --git a/examples/generic_python/tilelang_vec_add.py b/examples/generic_python/tilelang_vec_add.py index 3ff3c9376..601b6a7b7 100644 --- a/examples/generic_python/tilelang_vec_add.py +++ b/examples/generic_python/tilelang_vec_add.py @@ -2,8 +2,11 @@ import tilelang.language as T import torch from kernel_tuner import tune_kernel +from pathlib import Path -#@tilelang.jit # infers target from tensors at first call +FULL_PATH = Path(__file__).resolve() + +@tilelang.jit # infers target from tensors at first call def add(N: int, dtype: str = 'float32', block: int = 256,): @T.prim_func @@ -51,7 +54,8 @@ def tune(): answer = [None, None, (A + B).cpu()] - res, env = tune_kernel("add", add, N, args, tune_params, lang="generic_python", - call_function=call_tilelang, decorator="@tilelang.jit", answer=answer) - -tune() \ No newline at end of file + res, env = tune_kernel("add", FULL_PATH, N, args, tune_params, lang="generic_python", + call_function=call_tilelang, answer=answer) + +if __name__ == "__main__": + tune() \ No newline at end of file diff --git a/examples/generic_python/tilus_naive_matmul.py b/examples/generic_python/tilus_naive_matmul.py index 2f358ef48..42cdff86d 100644 --- a/examples/generic_python/tilus_naive_matmul.py +++ b/examples/generic_python/tilus_naive_matmul.py @@ -4,6 +4,9 @@ from kernel_tuner import tune_kernel, run_kernel import math import torch +from pathlib import Path + +FULL_PATH = Path(__file__).resolve() class MatmulV0(tilus.Script): @@ -105,7 +108,7 @@ def main(): results, env = tune_kernel( kernel_name="MatmulV0", - kernel_source=MatmulV0, + kernel_source=FULL_PATH, problem_size=[m, n], arguments=args, tune_params=tune_params, diff --git a/examples/generic_python/tilus_splitk_matmul.py b/examples/generic_python/tilus_splitk_matmul.py index f0dff1ac3..be797299e 100644 --- a/examples/generic_python/tilus_splitk_matmul.py +++ b/examples/generic_python/tilus_splitk_matmul.py @@ -4,17 +4,13 @@ import torch import math from kernel_tuner import tune_kernel, run_kernel +from pathlib import Path + +FULL_PATH = Path(__file__).resolve() -''' -@tilus.autotune("num_warps", [4, 8]) -@tilus.autotune("block_m, block_n", [(128, 128), (128, 64), (64, 128), (32, 256)]) -@tilus.autotune("block_k", [16, 32]) -@tilus.autotune("num_stages", [3, 4, 5]) -@tilus.autotune("split_k_factor", [1, 4, 12, 16]) -''' class MatmulV5(tilus.Script): - ''' - def __init__(self, block_m, block_n, block_k, num_warps, num_stages, split_k_factor): + def __init__(self, block_m=None, block_n=None, block_k=None, + num_warps=None, num_stages=None, split_k_factor=None): super().__init__() self.block_m = block_m self.block_n = block_n @@ -22,15 +18,7 @@ def __init__(self, block_m, block_n, block_k, num_warps, num_stages, split_k_fac self.num_warps = num_warps self.num_stages = num_stages self.split_k_factor = split_k_factor - ''' - def __init__(self): - super().__init__() - self.block_m = 128 - self.block_n = 128 - self.block_k = 16 - self.num_warps = 4 - self.num_stages = 4 - self.split_k_factor = 4 + def __call__( self, @@ -246,7 +234,7 @@ def main(): results, env = tune_kernel( kernel_name="MatmulV5", # This has to be a string of the actual name. TODO is this always the case? - kernel_source=MatmulV5, + kernel_source=FULL_PATH, problem_size=[m, n], arguments=args, tune_params=tune_params, diff --git a/examples/generic_python/tilus_tunable_precision.py b/examples/generic_python/tilus_tunable_precision.py deleted file mode 100644 index 1cc4a7499..000000000 --- a/examples/generic_python/tilus_tunable_precision.py +++ /dev/null @@ -1,120 +0,0 @@ -import tilus -import hidet -from hidet import float32, float16, int32 -#from tilus import float32, int32, float16 -from tilus.utils import cdiv, benchmark_func -import torch -from kernel_tuner import tune_kernel, run_kernel -from kernel_tuner.accuracy import Tunable, AccuracyObserver -import numpy as np - -INPUT_TYPE = float32 -OUTPUT_TYPE = float32 - -class VecAddV(tilus.Script): - def __init__(self): - super().__init__() - self.block_size_x = 32 # number of threads per block - - def __call__( - self, - n_size: int32, # size of the vectors - a_ptr: ~INPUT_TYPE, # input vector A - b_ptr: ~INPUT_TYPE, # input vector B - c_ptr: ~OUTPUT_TYPE # output vector C - ): - - # compute the number of blocks needed - self.attrs.blocks = [cdiv(n_size, self.block_size_x)] - self.attrs.warps = 1 # number of warps per block - - # calculate the offset for this block - offset: int32 = self.block_size_x * self.blockIdx.x - - # create global views for input/output vectors - ga = self.global_view(a_ptr, dtype=INPUT_TYPE, shape=[n_size]) - gb = self.global_view(b_ptr, dtype=INPUT_TYPE, shape=[n_size]) - gc = self.global_view(c_ptr, dtype=OUTPUT_TYPE, shape=[n_size]) - - a = self.load_global(ga, offsets=[offset], shape=[self.block_size_x]) - b = self.load_global(gb, offsets=[offset], shape=[self.block_size_x]) - c = a + b - self.store_global(gc, c, offsets=[offset]) - - -def call_tilus(kernel_function, args, kwargs, grid, threads, params): - kernel_function(*args, **kwargs) - - -def verify(answer, result_host, atol): - correct = True - for i, ans in enumerate(answer): - if ans is None: - continue - res = result_host[i].cpu() - if not torch.allcose(ans, res, atol=atol): - correct = False - - return correct - - - -def main(): - - size = 1024000 - - a_32 = torch.randn(size, dtype=torch.float32) - b_32 = torch.randn(size, dtype=torch.float32) - c_32 = torch.zeros_like(b_32) - c_expect = a_32 + b_32 - - a_16 = a_32.to(torch.float16) - b_16 = b_32.to(torch.float16) - c_16 = c_32.to(torch.float16) - - - tune_params = dict() - tune_params["block_size_x"] = [32, 64, 128, 256, 512, 1024] - tune_params["INPUT_TYPE"] = [tilus.float16, tilus.float32] - tune_params["OUTPUT_TYPE"] = [tilus.float16, tilus.float32] - - args = [ - size, - Tunable("INPUT_TYPE", { - tilus.float32: a_32, - tilus.float16: a_16, - }), - Tunable("INPUT_TYPE", { - tilus.float32: b_32, - tilus.float16: b_16, - }), - Tunable("OUTPUT_TYPE", { - tilus.float32: c_32, - tilus.float16: c_16, - }), - ] - - print(tune_params) - - observers = [AccuracyObserver("RMSE")] - - - - results, env = tune_kernel( - kernel_name="VecAddV", - kernel_source=VecAddV, - problem_size=size, - arguments=args, - tune_params=tune_params, - lang="generic_python", - answer=[None, None, None, c_expect.cpu()], - observers=observers, - call_function=call_tilus, - verify=verify, - verbose=True, - ) - - - -if __name__ == "__main__": - main() diff --git a/examples/generic_python/tilus_vec_add.py b/examples/generic_python/tilus_vec_add.py index 161575181..938d9026c 100644 --- a/examples/generic_python/tilus_vec_add.py +++ b/examples/generic_python/tilus_vec_add.py @@ -3,8 +3,9 @@ from tilus.utils import cdiv, benchmark_func import torch from kernel_tuner import tune_kernel, run_kernel +from pathlib import Path - +FULL_PATH = Path(__file__).resolve() class VecAddV(tilus.Script): def __init__(self): @@ -41,7 +42,7 @@ def call_tilus(kernel_function, args, kwargs, grid, threads, params): kernel_function(*args, **kwargs) -def tune_vecadd(size): +def tune(size): a = torch.randn(size, dtype=torch.float32).cuda() b = torch.randn(size, dtype=torch.float32).cuda() c = torch.empty(size, dtype=torch.float32).cuda() @@ -52,9 +53,10 @@ def tune_vecadd(size): tune_params = dict() tune_params["block_size_x"] = [32, 64, 128, 256, 512, 1024] + results, env = tune_kernel( kernel_name="VecAddV", - kernel_source=VecAddV, + kernel_source=FULL_PATH, problem_size=size, arguments=args, tune_params=tune_params, @@ -65,8 +67,7 @@ def tune_vecadd(size): ) -# TODO run kernel error handling same as tune_kernel -def run_vecadd(size): +def run(size): a = torch.randn(size, dtype=torch.float32).cuda() b = torch.randn(size, dtype=torch.float32).cuda() c = torch.empty(size, dtype=torch.float32).cuda() @@ -77,37 +78,20 @@ def run_vecadd(size): results = run_kernel( kernel_name="VecAddV", - kernel_source=VecAddV, + kernel_source=FULL_PATH, problem_size=size, arguments=args, params={"block_size_x": 32}, lang="generic_python", call_function=call_tilus, - #verbose=True, ) c_expect = c_expect.cpu() assert torch.allclose(results[-1], c_expect) - - - -def run_normal(size): - a = torch.randn(size, dtype=torch.float32).cuda() - b = torch.randn(size, dtype=torch.float32).cuda() - c = torch.empty(size, dtype=torch.float32).cuda() - c_expect = a + b - - vecadd = VecAddV() - vecadd(size, a, b, c) - print(c) - print(c_expect) - assert torch.allclose(c, c_expect) - if __name__ == "__main__": - size = 1024 - tune_vecadd(size) - #run_vecadd(size) - #run_normal(size) + size = 100000000 + tune(size) + run(size) diff --git a/examples/generic_python/triton_vec_add.py b/examples/generic_python/triton_vec_add.py index fcc019615..2e982abfa 100644 --- a/examples/generic_python/triton_vec_add.py +++ b/examples/generic_python/triton_vec_add.py @@ -1,18 +1,18 @@ import numpy as np -import triton.language as tl import torch -from kernel_tuner import tune_kernel, run_kernel -from kernel_tuner.file_utils import store_output_file, store_metadata_file import triton -from math import ceil +import triton.language as tl +from pathlib import Path +from kernel_tuner import tune_kernel, run_kernel + +FULL_PATH = Path(__file__).resolve() @triton.jit def add_op(x, y): return x + y - -#@triton.jit +@triton.jit def add_kernel(x_ptr, # *Pointer* to first input vector. y_ptr, # *Pointer* to second input vector. output_ptr, # *Pointer* to output vector. @@ -33,11 +33,8 @@ def add_kernel(x_ptr, # *Pointer* to first input vector. def call_triton(kernel_function, args, kwargs, grid, threads, params): kernel_function[grid](*args, **kwargs) -# NOTE: can the python file be changed in between? what happens? -# NOTE: tune params in the funcion signature are supported as key word arguments. Do not pass them as args, these -# will be inserted automatically. You can use them in the call function (kwargs). -def tune_with_generic(): +def tune(): size = 10000000 a = torch.randn(size, device='cuda', dtype=torch.float32) @@ -48,31 +45,27 @@ def tune_with_generic(): args = [a, b, c, size] tune_params = dict() - tune_params["block_size_x"] = [2**i for i in range(10)] + tune_params["block_size_x"] = [2**i for i in range(11)] - result = run_kernel("add_kernel", add_kernel, size, args, {"block_size_x": 256}, - lang="generic_python", call_function=call_triton, decorator="@triton.jit") - print(np.allclose(c_expect.cpu(), result[2])) - + result = run_kernel("add_kernel", FULL_PATH, size, args, {"block_size_x": 256}, + lang="generic_python", call_function=call_triton) + assert np.allclose(c_expect.cpu(), result[2]) + results, env = tune_kernel( kernel_name="add_kernel", - kernel_source=add_kernel, + kernel_source=FULL_PATH, problem_size=size, arguments=args, tune_params=tune_params, lang="generic_python", answer=[None, None, c_expect.cpu(), None], call_function=call_triton, - decorator="@triton.jit" ) - - - - -tune_with_generic() \ No newline at end of file +if __name__ == "__main__": + tune() \ No newline at end of file diff --git a/examples/generic_python/warp_vec_add.py b/examples/generic_python/warp_vec_add.py index 4a2c8d350..4f94c91d2 100644 --- a/examples/generic_python/warp_vec_add.py +++ b/examples/generic_python/warp_vec_add.py @@ -2,17 +2,16 @@ import numpy as np from kernel_tuner import tune_kernel, run_kernel import torch +from pathlib import Path wp.init() - - +FULL_PATH = Path(__file__).resolve() @wp.func def add_op(x: float, y: float): return x + y - -#@wp.kernel +@wp.kernel def vec_add(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=float), @@ -28,8 +27,6 @@ def vec_add(a: wp.array(dtype=float), c[idx] = add_op(a[idx], b[idx]) -# TODO do we allways want the call function to have the same parameters -# or do we only require some of them? def call_warp(kernel_function, args, kwargs, grid, threads, params): warp_args = [] for arg in args: @@ -37,66 +34,36 @@ def call_warp(kernel_function, args, kwargs, grid, threads, params): warp_args.append(wp.from_torch(arg)) else: warp_args.append(arg) - dim = args[3] + dim = params['size'] wp.launch(kernel=kernel_function, dim=dim, inputs=warp_args) -# NOTE default verify function only works for numpy/cupy ndarray, torch Tensor or numpy scalar -# That is why we need a costum verify function for warp. -def verify(answer, result_host, atol): - correct = True - for i, ans in enumerate(answer): - if ans is None: - continue - print("res: ", type(result_host[i])) - print("expect: ", type(ans)) - res = result_host[i].numpy() - if not np.allclose(ans, res, atol=atol): - correct = False - - return correct - - - def tune(): n = 1024 # Create host arrays - a_torch = torch.arange(n, dtype=torch.float32, device="cuda") - b_torch = torch.arange(n, 0, -1, dtype=torch.float32, device="cuda") - c_torch = torch.zeros(n, dtype=torch.float32, device="cuda") - c_expect = a_torch + b_torch - + a = np.arange(n, dtype=np.float32) + b = np.arange(n, 0, -1, dtype=np.float32) + c = np.zeros(n, dtype=np.float32) + c_expect = a + b + tune_params = dict() tune_params["work_per_thread"] = [2**i for i in range(10)] - args = [a_torch, b_torch, c_torch, n] - - - ''' - results = run_kernel( - kernel_name="vec_add", - kernel_source=vec_add, - problem_size=n, - arguments=args, - params={"work_per_thread": 16}, - lang="generic_python", - call_function=call_warp, - decorator="@wp.kernel" - ) - ''' + tune_params["size"] = [n] + args = [a, b, c, n] results, env = tune_kernel( kernel_name="vec_add", - kernel_source=vec_add, + kernel_source=FULL_PATH, problem_size=n, arguments=args, tune_params=tune_params, lang="generic_python", - answer=[None, None, c_expect.cpu(), None], + answer=[None, None, c_expect, None], call_function=call_warp, - decorator="@wp.kernel" ) -tune() \ No newline at end of file +if __name__ == "__main__": + tune() \ No newline at end of file diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 5d867fdac..39be59a23 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -104,7 +104,7 @@ def ready_argument_list(self, arguments): elif type(arg) in vars(builtins).values(): torch_args.append(arg) else: - raise TypeError("Unknown argument type: ", type(arg)) + raise TypeError("Unknown argument type: ", type(arg), ". Accepted types are Torch tenors, NumPy arrays and scalars and built-in Python types.") return torch_args @@ -150,9 +150,9 @@ def compile(self, kernel_instance, gpu_args=None): # the values of the tuning params. self.gpu_kwargs = {} if params is not None: - for name, p in self.signature.parameters.items(): - if name in params: - self.gpu_kwargs[name] = params[name] + for arg_name in self.signature: + if arg_name in params: + self.gpu_kwargs[arg_name] = params[arg_name] # Call the user-defined call function in order to compile the kernel. self.synchronize() @@ -173,7 +173,7 @@ def kernel_finished(self): """Returns True if the kernel has finished, False otherwise.""" return self.end.query() - def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): + def run_kernel(self, func, gpu_args, threads, grid, params=None): """Runs the Python kernel passed as 'func'. :param func: A cached Python kernel for this specific kernel configuration @@ -195,10 +195,7 @@ def run_kernel(self, func, gpu_args, threads, grid, stream=None, params=None): configuration :type params: dict """ - if stream is None: - stream = self.stream - - with torch.cuda.stream(stream): + with torch.cuda.stream(self.stream): logging.debug("Running Generic Python kernel") self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index e0421f847..d69ce3de0 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -111,8 +111,7 @@ def __deepcopy__(self, _): """The CUDA, OpenCL, HIP, C or Python DSL kernel code. It is allowed for the code to be passed as a string, a filename, a function that returns a string of code, or a list when the code needs auxilliary files. - In the case of a kernel in a Python DSL such as Triton, the reference to the - Python callable should be passed. + In the case of a kernel in a Python DSL such as Triton, only a filename is accepted. To support combined host and device code tuning, a list of filenames can be passed. The first file in the list should be the @@ -294,8 +293,7 @@ def __deepcopy__(self, _): ( """When the language Generic Python is used, a call function that calls the kernel in the Python DSL must be specified. The function must take the following arguments: - :kernel_function: the callable function with the tuning parameters inserted. If provided, the - kernel_function is decorated with the de decorator. + :kernel_function: the callable function with the tuning parameters inserted. :args: list of kernel arguments, as provided by the user in the argument. :kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, the tuning parameter will be added as a keyword argument. @@ -305,15 +303,6 @@ def __deepcopy__(self, _): "function", ), ), - ( - "decorator", - ( - """When the language Generic Python is used, a decorator can be provided in which the kernel source - will be wrapped internally by KernelTuner. Note that when passing the kernel to KernelTuner with the - ``kernel_source`` argument, the decorator should be removed from the kernel.""", - "string", - ), - ), ] ) @@ -619,14 +608,13 @@ def tune_kernel( objective=None, objective_higher_is_better=None, call_function=None, - decorator=None ): start_overhead_time = perf_counter() if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function, decorator) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -808,12 +796,11 @@ def run_kernel( quiet=False, log=None, call_function=None, - decorator=None ): if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function, decorator) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) diff --git a/kernel_tuner/kernel_sources/kernel_source.py b/kernel_tuner/kernel_sources/kernel_source.py index d18ecb284..51fb75d9b 100644 --- a/kernel_tuner/kernel_sources/kernel_source.py +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -18,7 +18,7 @@ class KernelSourceFactory(type): __init__ call of the corresponding subclass. In both subclasses, we call super().__init__, which triggers the __init__ call of the KernelSource class to initalize some common variables. ''' - def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function=None, decorator=None): + def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function=None): if lang == None: language = None else: @@ -39,13 +39,13 @@ def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function if ks_str: return KernelSourceStr(kernel_name, kernel_sources, lang, defines) else: - return KernelSourceFn(kernel_name, kernel_sources, lang, defines, call_function, decorator) + return KernelSourceFn(kernel_name, kernel_sources, lang, defines, call_function) # Else, normal behaviour for subclasses if ks_str: return super().__call__(kernel_name, kernel_sources, lang, defines) else: - return super().__call__(kernel_name, kernel_sources, lang, defines, call_function, decorator) + return super().__call__(kernel_name, kernel_sources, lang, defines, call_function) diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 386685377..22509ae37 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -14,21 +14,20 @@ from kernel_tuner.language import Language from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData - +from kernel_tuner.util import get_kernel_ast, get_arg_names class KernelSourceFn(KernelSource): """ Class that holds the Python-function-based kernel sources. - There is a primary kernel source for function-based kernels in Python. The source must be - a callable. This can be a function of a class. The function should not be decorated, but - a decorator can be supplied to wrap the function. + There is a primary kernel source for function-based kernels in Python. The kernel_source + must be a path to the file where the kernel with kernel_name lives. The kernel can be + decorated by a JIT decorator. A call function to specify how the kernel should be launched must be supplied. The call function must take the following arguments: - - kernel_function: the callable function with the tuning parameters inserted. If provided, the - kernel_function is decorated with the de decorator. + - kernel_function: the callable function with the tuning parameters inserted. - args: list of kernel arguments, as provided by the user in the argument. - kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, the tuning parameter will be added as a keyword argument. @@ -37,10 +36,10 @@ class KernelSourceFn(KernelSource): - params: dictionary with the values of the tuning params for this configuration. """ - def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function=None, decorator=None): + def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function=None): super().__init__(kernel_name, kernel_source, lang, defines) if isinstance(kernel_source, list): - raise ValueError("KernelSourceFn only supports a single kernel source function") + raise ValueError("KernelSourceFn only supports a single kernel source") if self.lang == Language.GENERIC_PYTHON: if call_function is None: @@ -49,27 +48,20 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function raise TypeError(f"call_function of type {type(call_function)} is not a callable object.") self.call_function = call_function # TODO ceck signature - if decorator: - if not isinstance(decorator, str): - raise TypeError(f"The decorator should be a string, got {type(decorator)} instead.") - if decorator[0] != '@': - raise ValueError(f"The decorator should start with a '@', got {decorator} instead.") - self.decorator = decorator - - self.source_kernel_fn = kernel_source # This kernel source remains the original object - self.kernel_fn = self.source_kernel_fn # This is the kernel source that we will modify - - try: - self.source = inspect.getsource(self.source_kernel_fn) - except TypeError as e: - raise TypeError( - f"{e}. Did you forget to remove a decorator before tuning?" - ) from e + if not isinstance(kernel_name, str): + raise TypeError("kernel_name should be a string, got ", type(kernel_name)) - self.signature = inspect.signature(self.source_kernel_fn) - self.source_tree = ast.parse(self.source) - self.import_nodes = self._find_import_nodes(inspect.getfile(self.source_kernel_fn)) - self.dependencies = self._find_dependencies() + source_ast = get_kernel_ast(kernel_name, kernel_source) + if isinstance(source_ast, tuple): # Class based kernel + self.source_tree = source_ast[0] + self.signature = get_arg_names(source_ast[1]) + else: + self.source_tree = source_ast + self.signature = get_arg_names(source_ast) + + self.kernel_fn = self.source_tree # This is where we will store the transformed source. + self.import_nodes = self._find_import_nodes(kernel_source) + self.dependencies = self._find_dependencies(kernel_source) def prepare_kernel_instance(self, kernel_options, params, grid, threads): @@ -123,22 +115,15 @@ def apply_params_to_source_fn(self, params): source_tree_copy = copy.deepcopy(self.source_tree) transformed_tree = transformer.visit(source_tree_copy) - # Add decorator if needed - if self.decorator: - dummy = f"{self.decorator}\ndef _dummy():\n pass\n" - decorator_node = ast.parse(dummy).body[0].decorator_list[0] - for node in transformed_tree.body: - if isinstance(node, ast.FunctionDef): - node.decorator_list.insert(0, decorator_node) - break # only apply to the top level function - # Add transformed main kernel to new module - new_module.body.extend(transformed_tree.body) + new_module.body.append(transformed_tree) # Fix locations and generate source ast.fix_missing_locations(new_module) new_source = astor.to_source(new_module) - + + #print(new_source) + # Create a unique module name and write new source to it. module_name = f'temp_kernel_module_{uuid.uuid4().hex}' with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file: @@ -151,10 +136,10 @@ def apply_params_to_source_fn(self, params): sys.modules[module_name] = temp_module spec.loader.exec_module(temp_module) new_fn = getattr(temp_module, self.kernel_name) - + return new_fn, temp_file_path - + def _find_import_nodes(self, source_file): ''' Parse kernel source file to find import statements. Return those @@ -192,16 +177,15 @@ def visit_Call(self, node): return visitor.called - def _find_dependencies(self): + def _find_dependencies(self, filepath): ''' Find all local function dependencies in the file where the kernel source is defined. Return a dicionary indexed by the function node name with the function body as value. ''' - source_file = inspect.getfile(self.source_kernel_fn) - with open(source_file, "r") as f: + with open(filepath, "r") as f: source_code = f.read() - tree = ast.parse(source_code, filename=source_file) + tree = ast.parse(source_code, filename=filepath) # Find the locally defined functions in the source file local_funcs = set() diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index c28bafa0b..24c622546 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -418,8 +418,6 @@ def detect_language(kernel_string): lang = "CUDA" elif "__kernel" in kernel_string: lang = "OpenCL" - elif "@triton.jit" in kernel_string: - lang = "Triton" else: lang = "C" return lang @@ -559,6 +557,82 @@ def get_kernel_string(kernel_source, params=None): return kernel_string +def get_kernel_ast(kernel_name, filepath): + ''' + Util function for Generic Python Backend that returns the kernel function as AST. + + :param kernel_name: name of the kernel (as passed by the user) + :type kernel_name: string + + :param filepath: the path to the file where the kernel lives. (passed by the user as kenrel_source) + :type filepath: string or Path containing a filename that points to the kernel source + + :returns: ast.FunctionDef node in case the kernel is a function or a tuple + (ast.ClassDef node, ast.FunctionDef node) in case the kernel is represented as the __call__ function + of a class (for Tilus support). + ''' + if isinstance(filepath, Path): + source = read_file(filepath) + elif isinstance(filepath, str): + with open(filepath, "r") as f: + source = f.read() + else: + raise TypeError("Error kernel_source does not specify a path to a file") + + tree = ast.parse(source) + + # Function based kernels + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == kernel_name: + return node + + # Class based kernels + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == kernel_name: + class_node = node + break + + if not class_node: + raise ValueError(f"Kernel {kernel_name} not found in {filepath}") + + # Search for __call__ function within class + for node in class_node.body: + if isinstance(node, ast.FunctionDef) and node.name == "__call__": + call_node = node + return (class_node, call_node) + + if not call_node: + raise ValueError(f"No __call__ function found inside Class {kernel_name}") + + + +def get_arg_names(func_node: ast.FunctionDef): + """ + Util to get the argument names from a Python AST function definition. This is needed to check + if there are any tunable parameters in the function signature. + """ + args = func_node.args + names = [] + + names += [arg.arg for arg in args.args] + names += [arg.arg for arg in args.posonlyargs] + names += [arg.arg for arg in args.kwonlyargs] + + # *args + if args.vararg: + names.append(args.vararg.arg) + + # **kwargs + if args.kwarg: + names.append(args.kwarg.arg) + + # Remove self for classes. + if 'self' in names: + names.remove('self') + + return names + + def get_problem_size(problem_size, params): """Compute current problem size.""" if callable(problem_size): diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py index eee0d3dd1..b8cddea88 100644 --- a/test/test_generic_python_functions.py +++ b/test/test_generic_python_functions.py @@ -1,8 +1,10 @@ from .context import skip_if_no_torch -from .test_kernel_source_fn import mock_kernel, kernel_with_kwarg, call_mock +from .test_kernel_source_fn import call_mock from kernel_tuner.core import DeviceInterface, KernelInstance from kernel_tuner.kernel_sources.kernel_source import KernelSource import numpy as np +from pathlib import Path +import os try: import torch @@ -10,6 +12,7 @@ except ImportError: pass +KS_FILE = os.path.join(Path(__file__).resolve().parent, "test_kernel_source_fn.py") # Helper functions ------------------------------ @@ -30,7 +33,7 @@ def get_context(): a = 42 b = torch.randn(12, device='cuda', dtype=torch.float32) args = [a, b] - ks = KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + ks = KernelSource("mock_kernel", KS_FILE, "generic_python", call_function=call_mock) return ks, args, params @@ -40,14 +43,24 @@ def get_context(): def test_ready_argument_list(): ks, args, params = get_context() dev = DeviceInterface(ks) + gpu_args = dev.ready_argument_list(args) assert len(args) == len(gpu_args) - for i, _ in enumerate(gpu_args): - assert value_equal(args[i], gpu_args[i]) - if type(gpu_args[i]) in (list, dict, torch.Tensor, np.ndarray): - assert gpu_args[i] is not args[i] # Assure deep copy + # arg 0: python scalar + assert isinstance(gpu_args[0], int) + assert gpu_args[0] == args[0] + + # arg 1: torch cuda tensor + assert isinstance(gpu_args[1], torch.Tensor) + assert gpu_args[1].is_cuda + + # values equal + assert torch.allclose(gpu_args[1], args[1]) + + # ensure deep copy + assert gpu_args[1] is not args[1] @skip_if_no_torch @@ -85,7 +98,7 @@ def test_gpu_kwargs(): params = {'mock_param': 64} a = torch.randn(12, device='cuda', dtype=torch.float32) args = [a] # we do not have to specify the kwarg here - ks = KernelSource("kernel_with_kwarg", kernel_with_kwarg, "generic_python", call_function=call_mock) + ks = KernelSource("kernel_with_kwarg", KS_FILE, "generic_python", call_function=call_mock) dev = DeviceInterface(ks) instance_data = ks.prepare_kernel_instance( diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py index 85aa4e3eb..d02be9e4d 100644 --- a/test/test_kernel_source_fn.py +++ b/test/test_kernel_source_fn.py @@ -1,3 +1,5 @@ +import os +from pathlib import Path import pytest import inspect import ast @@ -8,7 +10,7 @@ from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr - +KS_FILE = Path(__file__).resolve() # Helper functions -------------------------------------------- def normalize_ast(src: str): @@ -41,13 +43,18 @@ def __call__(self): def another_function(self): return self.mock_param + +@functools.lru_cache() +def kernel_with_decorator(): + mock_param = 42 + return mock_param # Tests ------------------------------------------------------ def test_factory_behaviour(): # KernelSourceFn should only be created when language generic_python is supplied - ks_fn = KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + ks_fn = KernelSource("mock_kernel", KS_FILE, "generic_python", call_function=call_mock) ks_str = KernelSource("vector_add", 'extern "C" __global__ void vector_add(float *c, float *a, float *b, int n) {', lang=None) assert isinstance(ks_fn, KernelSourceFn) @@ -59,27 +66,25 @@ def test_initiation(): Test invalid KernelSourceFn initations ''' with pytest.raises(ValueError, match=r"call_function must be supplied for language .*"): - KernelSource("mock_kernel", mock_kernel, "generic_python") + KernelSource("mock_kernel", KS_FILE, "generic_python") - with pytest.raises(TypeError, match=r".* Did you forget to remove a decorator before tuning\?"): + with pytest.raises(FileNotFoundError, match=r".* No such file or directory: .*"): KernelSource("mock_kernel", "This is a string Kernel", "generic_python", call_function=call_mock) - with pytest.raises(ValueError, match=r"KernelSourceFn only supports a single kernel source function"): - KernelSource("mock_kernel", [mock_kernel, kernel_with_dependency], "generic_python", call_function=call_mock) + with pytest.raises(TypeError, match="Error kernel_source does not specify a path to a file"): + KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + + with pytest.raises(ValueError, match=r"KernelSourceFn only supports a single kernel source"): + KernelSource("mock_kernel", [KS_FILE, "another file"], "generic_python", call_function=call_mock) with pytest.raises(TypeError, match=r".* is not a callable object"): - KernelSource("mock_kernel", mock_kernel, "generic_python", call_function="not a function") - - with pytest.raises(ValueError, match=r"The decorator should start with a '@', got .* instead."): - KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator="not a decorator") - - with pytest.raises(TypeError, match=r"The decorator should be a string, got .* instead."): - KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator=mock_kernel) + KernelSource("mock_kernel", KS_FILE, "generic_python", call_function="not a function") + def test_param_subsitution(): params = {"mock_param": 128} - ks = KernelSourceFn("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + ks = KernelSourceFn("mock_kernel", KS_FILE, "generic_python", call_function=call_mock) new_kernel_fn, _ = ks.apply_params_to_source_fn(params) actual_src = inspect.getsource(new_kernel_fn) @@ -101,7 +106,7 @@ def test_imports(): ''' params = {"mock_param": 128} - ks = KernelSourceFn("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) + ks = KernelSourceFn("mock_kernel", KS_FILE, "generic_python", call_function=call_mock) _, temp_path = ks.apply_params_to_source_fn(params) # Check if imports are present @@ -120,7 +125,7 @@ def test_param_substitution_class(): ''' params = {"mock_param": 128} - ks = KernelSourceFn("TilusLike", TilusLike, "generic_python", call_function=call_mock) + ks = KernelSourceFn("TilusLike", KS_FILE, "generic_python", call_function=call_mock) new_kernel_fn, _ = ks.apply_params_to_source_fn(params) @@ -142,15 +147,22 @@ def another_function(self): def test_decorator(): params = {"mock_param": 128} - ks = KernelSourceFn("mock_kernel", mock_kernel, "generic_python", call_function=call_mock, decorator="@functools.lru_cache()") + ks = KernelSourceFn("kernel_with_decorator", KS_FILE, "generic_python", call_function=call_mock) new_kernel_fn, _ = ks.apply_params_to_source_fn(params) - + actual_src = inspect.getsource(new_kernel_fn) + expected_src = """ +@functools.lru_cache() +def kernel_with_decorator(): + mock_param = 128 + return 128 +""" assert hasattr(new_kernel_fn, "__wrapped__") + assert normalize_ast(actual_src) == normalize_ast(expected_src) def test_dependencies(): params = {"mock_param": 128} - ks = KernelSourceFn("kernel_with_dependency", kernel_with_dependency, "generic_python", call_function=call_mock) + ks = KernelSourceFn("kernel_with_dependency", KS_FILE, "generic_python", call_function=call_mock) new_kernel_fn, _ = ks.apply_params_to_source_fn(params) res = new_kernel_fn() # This should not throw an error if the dependency exists in the module. From f9977f1344e8a9a73707acd2548b9649ab1654cc Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Sat, 7 Mar 2026 12:29:20 +0100 Subject: [PATCH 06/44] added normalize utility for the call function --- examples/generic_python/matmul/test_tilus.py | 26 ++ .../generic_python/matmul/tilelang_matmul.py | 221 +++------- .../generic_python/matmul/tilus_matmul.py | 319 ++++++--------- .../generic_python/matmul/triton_matmul.py | 371 ++++------------- .../matmul_old/tilelang_matmul.py | 212 ++++++++++ .../generic_python/matmul_old/tilus_matmul.py | 266 ++++++++++++ .../matmul_old/triton_matmul.py | 383 ++++++++++++++++++ examples/generic_python/numba_vec_add.py | 4 +- examples/generic_python/tilelang_vec_add.py | 2 +- examples/generic_python/tilus_naive_matmul.py | 2 +- .../generic_python/tilus_splitk_matmul.py | 2 +- examples/generic_python/tilus_vec_add.py | 28 +- examples/generic_python/triton_vec_add.py | 2 +- kernel_tuner/interface.py | 15 +- .../kernel_sources/kernel_source_fn.py | 2 +- kernel_tuner/util.py | 27 ++ test/test_util_functions.py | 22 + 17 files changed, 1229 insertions(+), 675 deletions(-) create mode 100644 examples/generic_python/matmul/test_tilus.py create mode 100644 examples/generic_python/matmul_old/tilelang_matmul.py create mode 100644 examples/generic_python/matmul_old/tilus_matmul.py create mode 100644 examples/generic_python/matmul_old/triton_matmul.py diff --git a/examples/generic_python/matmul/test_tilus.py b/examples/generic_python/matmul/test_tilus.py new file mode 100644 index 000000000..6d4b2430d --- /dev/null +++ b/examples/generic_python/matmul/test_tilus.py @@ -0,0 +1,26 @@ +import torch +from tilus_matmul import MatmulBasic + +sizes = [ + (65, 65, 17), + (67, 71, 19), + (1, 1, 1), + (63, 63, 15), + (129, 130, 33), +] + +matmul = MatmulBasic() +for m, n, k in sizes: + print(m, n, k) + + a = torch.randn(m, k, dtype=torch.float16, device="cuda") + b = torch.randn(k, n, dtype=torch.float16, device="cuda") + + c_actual = torch.empty(m, n, dtype=torch.float16, device="cuda") + c_expect = a @ b + + matmul(m, n, k, a, b, c_actual) + + torch.cuda.synchronize() + + torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) \ No newline at end of file diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py index 2a7b5154d..895732629 100644 --- a/examples/generic_python/matmul/tilelang_matmul.py +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -1,212 +1,87 @@ +import torch import tilelang import tilelang.language as T -import torch -from kernel_tuner import tune_kernel -#@tilelang.jit -def matmul(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): +@tilelang.jit +def matmul_basic(M:int, N:int, K:int, block_M:int, block_N:int, block_K:int, dtype:str="float16", accum_dtype:str="float32"): @T.prim_func - def main( + def gemm( A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): - # Define a grid with enough blocks to cover M×N - num_threads=128 - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): - - # Allocate shared memory for the current tile of A and B + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + # We do use shared memory, even though this is a basic kernel. However, you don't + # really get around this because T.gemm can not handle global memory directly. A_shared = T.alloc_shared((block_M, block_K), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - - # Allocate a local (register) fragment for partial accumulations C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - - # Enable swizzle-based rasterization for better L2 locality - panel_size = 4 - T.use_swizzle(panel_size=panel_size, enable=True) - - # Initialize the local accumulation buffer to zero T.clear(C_local) - - num_stages=3 - - # Loop over the K dimension in block_K chunks, using a pipeline - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - # Copy from global memory to shared memory + + # We do use a pipelining optimization here, because this is 'the basic way' + # of writing for loops in TileLang. + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[k * block_K, bx * block_N], B_shared) - - # Perform a matrix multiply-accumulate on the tile T.gemm(A_shared, B_shared, C_local) - # Copy the accumulated result from local memory (C_local) to global memory (C) T.copy(C_local, C[by * block_M, bx * block_N]) + + return gemm - return main +# https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py @tilelang.jit -def matmul_with_decorator(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): +def matmul_opt(M, N, K, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): @T.prim_func - def main( + def gemm_autotune( A: T.Tensor((M, K), dtype), - B: T.Tensor((K, N), dtype), + B: T.Tensor((N, K), dtype), C: T.Tensor((M, N), dtype), ): - # Define a grid with enough blocks to cover M×N - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): - - # Allocate shared memory for the current tile of A and B + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - - # Allocate a local (register) fragment for partial accumulations + B_shared = T.alloc_shared((block_N, block_K), dtype) C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - - # Initialize the local accumulation buffer to zero + C_shared = T.alloc_shared((block_M, block_N), dtype) + T.use_swizzle(panel_size=10, enable=enable_rasteration) T.clear(C_local) - - # Loop over the K dimension in block_K chunks, using a 3-stage pipeline - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): - # Copy from global memory to shared memory + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[k * block_K, bx * block_N], B_shared) + T.copy(B[bx * block_N, k * block_K], B_shared) + T.gemm( + A_shared, + B_shared, + C_local, + transpose_B=True, + ) + T.copy(C_local, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return gemm_autotune - # Perform a matrix multiply-accumulate on the tile - T.gemm(A_shared, B_shared, C_local) - # Copy the accumulated result from local memory (C_local) to global memory (C) - T.copy(C_local, C[by * block_M, bx * block_N]) - return main +def main(): + kernel = matmul_basic(1024, 1024, 1024, 128, 128, 32) + import torch + + a = torch.randn(1024, 1024).cuda().half() + b = torch.randn(1024, 1024).cuda().half() + c = torch.empty((1024, 1024), device='cuda', dtype=torch.float16) -def run(m, n, k): - a = torch.randn(m, k, device="cuda", dtype=torch.float16) - b = torch.randn(k, n, device="cuda", dtype=torch.float16) - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - kernel = matmul(m, n, k, 128, 128, 32) kernel(a, b, c) + ref_c = a @ b - tol = m * 2**(-11) - # Validate results - torch.testing.assert_close(c, ref_c, rtol=tol, atol=tol) - - -def call_tilelang(kernel_function, args, kwargs, grid, threads, params): - compiled_kernel = kernel_function(**kwargs) - compiled_kernel(*args) - - -def time(m, n, k): - a = torch.randn(m, k, device="cuda", dtype=torch.float16) - b = torch.randn(k, n, device="cuda", dtype=torch.float16) - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - c_ans = a @ b - - args = [a, b, c] - tune_params = dict() - tune_params["M"] = [m] - tune_params["K"] = [k] - tune_params["N"] = [n] - tune_params["block_M"] = [64, 128] - tune_params["block_N"] = [64, 128] - tune_params["block_K"] = [32, 64] - - results_kt, env = tune_kernel("matmul", matmul, m * n, args, tune_params, lang="generic_python", - call_function=call_tilelang, decorator="@tilelang.jit", verbose=False, iterations=100) - - import time - num_repeats = 100 - times_direct = [] - for config in results_kt: - bs_m = config["block_M"] - bs_n = config["block_N"] - bs_k = config["block_K"] - - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - - kernel = matmul_with_decorator(m, n, k, bs_m, bs_n, bs_k) - kernel(a, b, c) - - torch.allclose(c.cpu(), c_ans.cpu(), atol=m * 2**(-11)) - - for i in range(num_repeats): - times = [] - - torch.cuda.synchronize() - start = time.time() - kernel(a, b, c) - torch.cuda.synchronize() - times.append(time.time() - start) - - avg_time_ms = round((1000 * sum(times) / len(times)), 3) - times_direct.append(avg_time_ms) - print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, time={avg_time_ms}ms") - - - import matplotlib.pyplot as plt - - # Extract times - times_kt = [cfg['time'] for cfg in results_kt] - - # x-axis labels - configs = [f"config{i}" for i in range(len(times_kt))] - x = range(len(configs)) - - plt.figure(figsize=(10,6)) - plt.plot(configs, times_kt, marker='s', label='KernelTuner') - plt.plot(configs, times_direct, marker='x', label='Direct') - plt.ylabel('Time (ms)') - plt.xlabel('Configuration') - plt.title('Kernel execution time per configuration') - plt.xticks(rotation=45) - plt.legend() - plt.grid(True) - plt.tight_layout() - plt.savefig("tilelang.png") - print("saved fig") - - -def tune(m, n, k): - a = torch.randn(m, k, device="cuda", dtype=torch.float16) - b = torch.randn(k, n, device="cuda", dtype=torch.float16) - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - c_actual = a @ b - - args = [a, b, c] - tune_params = dict() - tune_params["M"] = [m] - tune_params["K"] = [k] - tune_params["N"] = [n] - tune_params["block_M"] = [64, 128, 256] - tune_params["block_N"] = [64, 128, 256] - tune_params["block_K"] = [32, 64, 128] - tune_params["num_stages"] = [2, 3, 4] - tune_params["panel_size"] = [4, 8] # equivalent to group size m in Triton - tune_params["num_threads"] = [64, 128, 256] - - restrictions = [ - # tile size budget - "block_M * block_N <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "block_M <= 4 * block_N", - "block_N <= 4 * block_M", - - # large K only with reasonably large M/N - "not (block_K == 128 and block_M < 64 and block_N < 64)", - ] - - tol = m * 2**(-11) - answer = [None, None, c_actual.cpu()] - - results, env = tune_kernel("matmul", matmul, m * n, args, tune_params, atol=tol, lang="generic_python", - call_function=call_tilelang, restrictions=restrictions, answer=answer, decorator="@tilelang.jit", verbose=False) + + torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2) + print("All check passed.") + + + if __name__ == "__main__": - #m, n, k = 1024, 1024, 1024 - m, n, k = 8192, 8192, 8192 - time(m, n, k) \ No newline at end of file + main() diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py index e8996c201..dfe294761 100644 --- a/examples/generic_python/matmul/tilus_matmul.py +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -1,102 +1,98 @@ -import math - -import pandas import tilus -import torch from tilus import float16, float32, int32 -from tilus.utils import benchmark_func -from kernel_tuner import tune_kernel, run_kernel - +from tilus.utils import cdiv -class MatmulV4(tilus.Script): +# This kernel is copied from the Tilus project: +# https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_v0.py +# +# Original example: matmul_v0.py +# Copyright (c) the Tilus authors +class MatmulBasic(tilus.Script): def __init__(self): super().__init__() - self.block_m = 128 - self.block_n = 128 + # we define three hyperparameters: ``block_m``, ``block_n``, and ``block_k`` to determine the tile size on + # m, n, and k dimensions for each `thread block` of the kernel. + self.block_m = 64 + self.block_n = 64 self.block_k = 16 - self.num_warps = 4 - self.num_stages = 4 def __call__( self, - m_size: int32, - n_size: int, - k_size: int, - a_ptr: ~float16, - b_ptr: ~float16, - c_ptr: ~float16, + m_size: int32, # the size of the m dimension of the input matrix A and output matrix C + n_size: int, # the size of the n dimension of the input matrix B and output matrix C + k_size: int, # the size of the k dimension of the input matrix A and B + a_ptr: ~float16, # the pointer to the input matrix A, which is a 2D tensor of shape [m_size, k_size] + b_ptr: ~float16, # the pointer to the input matrix B, which is a 2D tensor of shape [k_size, n_size] + c_ptr: ~float16, # the pointer to the output matrix C, which is a 2D tensor of shape [m_size, n_size] ): self.attrs.blocks = [ - self.utils.ceil_div(m_size, self.block_m), - self.utils.ceil_div(n_size, self.block_n), + cdiv(m_size, self.block_m), # the x dimension size of the grid + cdiv(n_size, self.block_n), # the y dimension size of the grid ] - self.attrs.warps = self.num_warps + self.attrs.warps = 1 # the number of warps per thread block, must be a compile-time known integer - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k - offset_m: int32 = block_m * self.blockIdx.x - offset_n: int32 = block_n * self.blockIdx.y + # define two int32 variables to store the offsets of the m and n dimensions for the current thread block. + offset_m: int32 = self.block_m * self.blockIdx.x + offset_n: int32 = self.block_n * self.blockIdx.y + # create two global tensors `ga` and `gb` to represent the input matrices A and B, respectively. ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) - sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) - sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) - acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) - - for stage in range(self.num_stages - 1): - offset_k = stage * self.block_k - self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) - self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) - self.copy_async_commit_group() - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() + # create a register tensor `acc` to accumulate the results of the matrix multiplication. + acc = self.register_tensor( + dtype=float32, shape=[self.block_m, self.block_n], init=0.0 + ) - current_stage: int32 = 0 - preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): - # computation for current tile - a = self.load_shared(sa[current_stage]) - b = self.load_shared(sb[current_stage]) - self.dot(a, b, acc, out=acc) + # iterate over the k dimension in blocks of size `block_k`. + for k in range(cdiv(k_size, self.block_k)): + # calculate the offset for the current block in the k dimension + offset_k = k * self.block_k - # preload the next tile of A and B into shared memory - preload_offset_k = offset_k + (self.num_stages - 1) * block_k - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], + # load a block of matrix A and B into register tensors `a` and `b`. + a = self.load_global( + ga, offsets=[offset_m, offset_k], shape=[self.block_m, self.block_k] ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], + b = self.load_global( + gb, offsets=[offset_k, offset_n], shape=[self.block_k, self.block_n] ) - self.copy_async_commit_group() - # update the stage - current_stage = (current_stage + 1) % self.num_stages - preload_stage = (preload_stage + 1) % self.num_stages - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - self.free_shared(sa) - self.free_shared(sb) + # perform the dot product: acc = a @ b + acc + self.dot(a, b, acc, out=acc) - casted_acc = self.cast(acc, dtype=float16) + # after the loop, we cast the accumulated result `acc` to float16 type and store it back to the output matrix C. + acc_f16 = self.cast(acc, dtype=float16) gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + self.store_global(gc, acc_f16, offsets=[offset_m, offset_n]) -class MatmulGroupedOrdering(tilus.Script): - def __init__(self): + +# This kernel is copied from the Tilus project: +# https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_v5.py +# +# Original example: matmul_v5.py +# Copyright (c) the Tilus authors +# +# Modifications in file: +# - Removed auto-tuning decorators +class MatmulOpt(tilus.Script): + def __init__( + self, + num_warps=None, + block_m=None, + block_n=None, + block_k=None, + num_stages=None, + split_k_factor=None, + ): super().__init__() - self.block_m = 128 - self.block_n = 128 - self.block_k = 16 - self.num_warps = 4 - self.num_stages = 4 - self.group_size_m = 8 + self.block_m = block_m + self.block_n = block_n + self.block_k = block_k + self.num_warps = num_warps + self.num_stages = num_stages + self.split_k_factor = split_k_factor def __call__( self, @@ -107,25 +103,23 @@ def __call__( b_ptr: ~float16, c_ptr: ~float16, ): - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k - - num_pid_m = self.utils.ceil_div(m_size, block_m) - num_pid_n = self.utils.ceil_div(n_size, block_n) - self.attrs.blocks = [num_pid_m * num_pid_n] - - pid = self.blockIdx.x - num_pid_in_group = self.group_size_m * num_pid_n - group_id = pid // num_pid_in_group - - first_pid_m = group_id * self.group_size_m - group_size_m = min(num_pid_m - first_pid_m, self.group_size_m) + self.attrs.blocks = [ + cdiv(m_size, self.block_m), + cdiv(n_size, self.block_n), + self.split_k_factor, + ] + self.attrs.warps = self.num_warps - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m + # the k_size for each thread block + block_k_size = ( + cdiv(cdiv(k_size, self.split_k_factor), self.block_k) * self.block_k + ) + start_offset_k = self.blockIdx.z * block_k_size + end_offset_k = min(start_offset_k + block_k_size, k_size) - self.attrs.warps = self.num_warps - offset_m: int32 = pid_m * block_m - offset_n: int32 = pid_n * block_n + block_m, block_n, block_k = self.block_m, self.block_n, self.block_k + offset_m: int32 = block_m * self.blockIdx.x + offset_n: int32 = block_n * self.blockIdx.y ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) @@ -134,7 +128,7 @@ def __call__( acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) for stage in range(self.num_stages - 1): - offset_k = stage * self.block_k + offset_k = start_offset_k + stage * self.block_k self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) self.copy_async_commit_group() @@ -144,7 +138,9 @@ def __call__( current_stage: int32 = 0 preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): + for offset_k in self.range( + start_offset_k, end_offset_k, block_k, unroll=self.num_stages + ): # computation for current tile a = self.load_shared(sa[current_stage]) b = self.load_shared(sb[current_stage]) @@ -152,16 +148,17 @@ def __call__( # preload the next tile of A and B into shared memory preload_offset_k = offset_k + (self.num_stages - 1) * block_k - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], - ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], - ) + if preload_offset_k < end_offset_k: + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) self.copy_async_commit_group() # update the stage @@ -170,97 +167,41 @@ def __call__( self.copy_async_wait_group(n=self.num_stages - 2) self.sync() + # free the shared memory tensors for A and B self.free_shared(sa) self.free_shared(sb) + # cast the accumulator to float16 and change the register tensor's layout + sc = self.shared_tensor(dtype=float16, shape=[block_m, block_n]) casted_acc = self.cast(acc, dtype=float16) - gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) - - - -def main(): - headers = ["m", "n", "k", "name", "latency (ms)", "tflops"] - workloads = [ - [4096, 4096, 4096], - [1024, 1024, 14336], - ] - - rows = [] - for m, n, k in workloads: - matmul = MatmulGroupedOrdering() #MatmulV4() - - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - matmul(m, n, k, a, b, c_actual) - - # check correctness - torch.testing.assert_close(c_expect, c_actual) - - # benchmark - for name, func in [ - ("torch", lambda: torch.matmul(a, b, out=c_expect)), - ("tilus", lambda: matmul(m, n, k, a, b, c_actual)), - ]: - latency = benchmark_func(func, warmup=5, repeat=20) - tflops = 2 * m * n * k / latency * 1e-9 - rows.append([m, n, k, name, latency, tflops]) - - df = pandas.DataFrame(rows, columns=headers) - print(df) - -def call_tilus(kernel_function, args, kwargs, grid, threads, params): - kernel_function(*args, **kwargs) - - -def tune_matmul(m, n, k): - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n #(m, n) - args = [m, n, k, a, b, c_actual] - tune_params = dict() - tune_params["block_m"] = [32, 64, 128, 256] - tune_params["block_n"] = [32, 64, 128, 256] - tune_params["block_k"] = [32, 64, 128] - tune_params["group_size_m"] = [4, 8] - tune_params["num_stages"] = [2, 3, 4] - tune_params["num_warps"] = [4, 8] - - - restrictions = [ - # tile size budget - "block_m * block_n <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "block_m <= 4 * block_n", - "block_n <= 4 * block_m", - - # large K only with reasonably large M/N - "not (block_k == 128 and block_m < 64 and block_n < 64)", - - # 32x32 requires 8 warps - "not (block_m == 32 and block_n == 32 and num_warps < 8)", - ] - - - answer = [None] * 6 - answer[-1] = c_expect.cpu() - atol = 1e-2 #m * 2**(-11) - - results, env = tune_kernel("MatmulGroupedOrdering", MatmulGroupedOrdering, size, args, tune_params, grid_div_x = ["block_m", "block_n"], - answer = answer, atol=atol, restrictions=restrictions, - lang="generic_python", call_function=call_tilus, - block_size_names=["block_m", "block_n", "block_k"], strategy="simulated_annealing") - - -if __name__ == "__main__": - #m, n, k = 4096, 4096, 4096 - m, n, k = 8192, 8192, 8192 - tune_matmul(m, n, k) + self.store_shared(sc, casted_acc) + self.sync() + rc = self.load_shared(sc) + self.free_shared(sc) - #main() \ No newline at end of file + m_blocks, n_blocks = cdiv(m_size, block_m), cdiv(n_size, block_n) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + if self.split_k_factor == 0: + self.store_global(gc, rc, offsets=[offset_m, offset_n]) + else: + semaphores = self.global_tensor( + dtype=int32, shape=[m_blocks, n_blocks], requires_clean=True + ) + semaphore: ~int32 = ~semaphores[self.blockIdx.x, self.blockIdx.y] + + # load and accumulate the partial result in global memory + if self.blockIdx.z > 0: + self.lock_semaphore(semaphore, value=self.blockIdx.z) + partial_rc = self.load_global( + gc, offsets=[offset_m, offset_n], shape=[block_m, block_n] + ) + self.add(rc, partial_rc, out=rc) + + # store the result to global memory and release the semaphore + self.store_global(gc, rc, offsets=[offset_m, offset_n]) + + # release the semaphore + self.sync() # we need to make sure the previous store_global is finished + self.release_semaphore( + semaphore, value=(self.blockIdx.z + 1) % self.split_k_factor + ) \ No newline at end of file diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py index 238a61c9f..b127b0331 100644 --- a/examples/generic_python/matmul/triton_matmul.py +++ b/examples/generic_python/matmul/triton_matmul.py @@ -1,56 +1,81 @@ import torch - import triton import triton.language as tl -from kernel_tuner import tune_kernel -from kernel_tuner import run_kernel -def get_cuda_autotune_config(): - return [ - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, - num_warps=8), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, - num_warps=2), - triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, - num_warps=2), - # Good config for fp8 inputs. - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, - num_warps=8), - triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, - num_warps=8), - triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4) - ] +@triton.jit +def matmul_basic( + # Pointers + a_ptr, b_ptr, c_ptr, + # Matrix sizes + M, N, K, + # Strides + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + # Tile sizes + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + # Each program computes one BLOCK_SIZE_M x BLOCK_SIZE_N tile of C + pid_m = tl.program_id(axis=0) + pid_n = tl.program_id(axis=1) + + # Compute row/column indices for the tile + offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_k = tl.arange(0, BLOCK_SIZE_K) + + # Create pointers to A and B tiles + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # Accumulator + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + # Loop over K dimension + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load( + a_ptrs, + mask=(offs_am[:, None] < M) & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + + b = tl.load( + b_ptrs, + mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) & (offs_bn[None, :] < N), + other=0.0, + ) + + accumulator = tl.dot(a, b, accumulator) -''' -@triton.autotune( - configs=get_cuda_autotune_config(), - key=['M', 'N', 'K'], -) -''' -#@triton.jit -def matmul_kernel( + # advance K + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Store result + c = accumulator.to(tl.float16) + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=mask) + + + + +# This kernel is copied from the Triton project: +# https://github.com/triton-lang/triton/blob/main/python/tutorials/03-matrix-multiplication.py +# +# Original example: 03-matrix-multiplication.py +# Copyright (c) the Triton authors +# +# Modifications in file: +# - Removed auto-tuning decorators +# - Removed activation function +@triton.jit +def matmul_opt( # Pointers to matrices a_ptr, b_ptr, c_ptr, # Matrix dimensions @@ -63,7 +88,7 @@ def matmul_kernel( stride_cm, stride_cn, # Meta-parameters BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # - GROUP_SIZE_M: tl.constexpr, num_stages, num_warps# + GROUP_SIZE_M: tl.constexpr, # ): """Kernel for computing the matmul C = A x B. A has shape (M, K), B has shape (K, N) and C has shape (M, N) @@ -71,6 +96,7 @@ def matmul_kernel( # ----------------------------------------------------------- # Map program ids `pid` to the block of C it should compute. # This is done in a grouped ordering to promote L2 data reuse. + # See above `L2 Cache Optimizations` section for details. pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) @@ -100,6 +126,7 @@ def matmul_kernel( # and accumulate # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + # See above `Pointer Arithmetic` section for details offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N offs_k = tl.arange(0, BLOCK_SIZE_K) @@ -133,251 +160,3 @@ def matmul_kernel( tl.store(c_ptrs, c, mask=c_mask) -def run_matmul(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - grid = lambda META: (triton.cdiv(m, META['BLOCK_SIZE_M']) * triton.cdiv(n, META['BLOCK_SIZE_N']), ) - - matmul_kernel[grid]( - a, b, c_actual, # - m, n, k, # - a.stride(0), a.stride(1), # - b.stride(0), b.stride(1), # - c_actual.stride(0), c_actual.stride(1), - 128, 256, 64, 8 - ) - - torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) - - -def run_matmul_kt(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n - - args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1)] - - params = {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M":4, "num_stages":3, "num_warps":4} - - result = run_kernel("matmul_kernel", matmul_kernel, size, args, params=params, grid_div_x=["BLOCK_SIZE_N", "BLOCK_SIZE_M"], - lang="generic_python", decorator="@triton.jit", call_function=call_triton, - block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) - c_res = result[2] - - assert torch.allclose(c_res, c_expect.cpu(), atol=1e-2, rtol=1e-1) - - - - - - -def call_triton(kernel_function, args, kwargs, grid, threads, params): - #print("using grid: ", grid) - #print("args: ", args) - #print("kwargs: ", kwargs) - kernel_function[grid](*args, **kwargs) - - - - -def check_time(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n - args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1)] - tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [64, 128] - tune_params["BLOCK_SIZE_N"] = [64, 128] - tune_params["BLOCK_SIZE_K"] = [32, 64] - tune_params["GROUP_SIZE_M"] = [4, 8] - tune_params["num_stages"] = [3] - tune_params["num_warps"] = [4] - - restrictions = [ - # tile size budget - "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", - "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", - - # large K only with reasonably large M/N - "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", - - # 32x32 requires 8 warps - "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", - ] - - grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] - - answer = [None] * 12 - answer[2] = c_expect.cpu() - atol = 1e-2 #m * 2**(-11) - - - - results_ours, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, - restrictions=restrictions, iterations=100, answer=answer, atol=atol, - lang="generic_python", decorator="@triton.jit", call_function=call_triton, - block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) - - - results_prev, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, - restrictions=restrictions, iterations=100, answer=answer, atol=atol, - lang="triton", block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) - - - - import time - num_repeats = 100 - times_direct = [] - for config in results_prev: - - bs_m = config["BLOCK_SIZE_M"] - bs_n = config["BLOCK_SIZE_N"] - bs_k = config["BLOCK_SIZE_K"] - gs_m = config["GROUP_SIZE_M"] - num_stages = config["num_stages"] - num_warps = config["num_warps"] - - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - - grid = (triton.cdiv(m, bs_m) * triton.cdiv(n, bs_n, ), ) - jit_function = triton.jit(matmul_kernel) - - - jit_function[grid]( - a, b, c_actual, - m, n, k, - a.stride(0), a.stride(1), - b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1), - bs_m, bs_n, bs_k, gs_m, num_stages, num_warps - ) - - - torch.allclose(c_expect.cpu(), c_actual.cpu(), atol=1e-2) - - - - for i in range(num_repeats): - times = [] - - #c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - - torch.cuda.synchronize() - start = time.time() - jit_function[grid]( - a, b, c_actual, - m, n, k, - a.stride(0), a.stride(1), - b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1), - bs_m, bs_n, bs_k, gs_m, num_stages, num_warps - ) - - torch.cuda.synchronize() - times.append(time.time() - start) - - avg_time_ms = round((1000 * sum(times) / len(times)), 3) - times_direct.append(avg_time_ms) - print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, GROUP_SIZE_M={gs_m}, num_stages={num_stages}, num_warps={num_warps}, time={avg_time_ms}ms") - - - import matplotlib.pyplot as plt - - # Extract times - times_prev = [cfg['time'] for cfg in results_prev] - times_ours = [cfg['time'] for cfg in results_ours] - - # x-axis labels - configs = [f"config{i}" for i in range(len(times_prev))] - x = range(len(configs)) - - plt.figure(figsize=(10,6)) - plt.plot(configs, times_prev, marker='o', label='Triton tuned') - plt.plot(configs, times_ours, marker='s', label='Generic tuned') - plt.plot(configs, times_direct, marker='x', label='Direct') - plt.ylabel('Time (ms)') - plt.xlabel('Configuration') - plt.title('Kernel execution time per configuration') - plt.xticks(rotation=45) - plt.legend() - plt.grid(True) - plt.tight_layout() - plt.savefig("ouptut.png") - - - - - - -def tune_matmul(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n - args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1)] - tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [64, 128, 256] - tune_params["BLOCK_SIZE_N"] = [64, 128, 256] - tune_params["BLOCK_SIZE_K"] = [32, 64, 128] - tune_params["GROUP_SIZE_M"] = [4, 8] - tune_params["num_stages"] = [2, 3, 4] - tune_params["num_warps"] = [4, 8] - - restrictions = [ - # tile size budget - "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", - "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", - - # large K only with reasonably large M/N - "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", - - # 32x32 requires 8 warps - "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", - ] - - grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] - - answer = [None] * 12 - answer[2] = c_expect.cpu() - - results, env = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, - answer = answer, atol=4.0, restrictions=restrictions, - lang="generic_python", decorator="@triton.jit", call_function=call_triton, - block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], strategy="simulated_annealing") - - - - - -if __name__ == "__main__": - m, n, k = 8192, 8192, 8192 - #m, n, k = 4096, 4096, 4096 - #tune_matmul(m, n, k) - #check_time(m, n, k) - #run_matmul_kt(m, n, k) - - - - - - diff --git a/examples/generic_python/matmul_old/tilelang_matmul.py b/examples/generic_python/matmul_old/tilelang_matmul.py new file mode 100644 index 000000000..2a7b5154d --- /dev/null +++ b/examples/generic_python/matmul_old/tilelang_matmul.py @@ -0,0 +1,212 @@ +import tilelang +import tilelang.language as T +import torch +from kernel_tuner import tune_kernel + + +#@tilelang.jit +def matmul(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): + @T.prim_func + def main( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), dtype), + ): + # Define a grid with enough blocks to cover M×N + num_threads=128 + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): + + # Allocate shared memory for the current tile of A and B + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + + # Allocate a local (register) fragment for partial accumulations + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Enable swizzle-based rasterization for better L2 locality + panel_size = 4 + T.use_swizzle(panel_size=panel_size, enable=True) + + # Initialize the local accumulation buffer to zero + T.clear(C_local) + + num_stages=3 + + # Loop over the K dimension in block_K chunks, using a pipeline + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + # Copy from global memory to shared memory + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[k * block_K, bx * block_N], B_shared) + + # Perform a matrix multiply-accumulate on the tile + T.gemm(A_shared, B_shared, C_local) + + # Copy the accumulated result from local memory (C_local) to global memory (C) + T.copy(C_local, C[by * block_M, bx * block_N]) + + return main + +@tilelang.jit +def matmul_with_decorator(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): + @T.prim_func + def main( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), dtype), + ): + # Define a grid with enough blocks to cover M×N + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + + # Allocate shared memory for the current tile of A and B + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + + # Allocate a local (register) fragment for partial accumulations + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Initialize the local accumulation buffer to zero + T.clear(C_local) + + # Loop over the K dimension in block_K chunks, using a 3-stage pipeline + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): + # Copy from global memory to shared memory + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[k * block_K, bx * block_N], B_shared) + + # Perform a matrix multiply-accumulate on the tile + T.gemm(A_shared, B_shared, C_local) + + # Copy the accumulated result from local memory (C_local) to global memory (C) + T.copy(C_local, C[by * block_M, bx * block_N]) + + return main + + + +def run(m, n, k): + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(k, n, device="cuda", dtype=torch.float16) + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + kernel = matmul(m, n, k, 128, 128, 32) + kernel(a, b, c) + ref_c = a @ b + tol = m * 2**(-11) + # Validate results + torch.testing.assert_close(c, ref_c, rtol=tol, atol=tol) + + +def call_tilelang(kernel_function, args, kwargs, grid, threads, params): + compiled_kernel = kernel_function(**kwargs) + compiled_kernel(*args) + + +def time(m, n, k): + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(k, n, device="cuda", dtype=torch.float16) + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + c_ans = a @ b + + args = [a, b, c] + tune_params = dict() + tune_params["M"] = [m] + tune_params["K"] = [k] + tune_params["N"] = [n] + tune_params["block_M"] = [64, 128] + tune_params["block_N"] = [64, 128] + tune_params["block_K"] = [32, 64] + + results_kt, env = tune_kernel("matmul", matmul, m * n, args, tune_params, lang="generic_python", + call_function=call_tilelang, decorator="@tilelang.jit", verbose=False, iterations=100) + + import time + num_repeats = 100 + times_direct = [] + for config in results_kt: + bs_m = config["block_M"] + bs_n = config["block_N"] + bs_k = config["block_K"] + + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + + kernel = matmul_with_decorator(m, n, k, bs_m, bs_n, bs_k) + kernel(a, b, c) + + torch.allclose(c.cpu(), c_ans.cpu(), atol=m * 2**(-11)) + + for i in range(num_repeats): + times = [] + + torch.cuda.synchronize() + start = time.time() + kernel(a, b, c) + torch.cuda.synchronize() + times.append(time.time() - start) + + avg_time_ms = round((1000 * sum(times) / len(times)), 3) + times_direct.append(avg_time_ms) + print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, time={avg_time_ms}ms") + + + import matplotlib.pyplot as plt + + # Extract times + times_kt = [cfg['time'] for cfg in results_kt] + + # x-axis labels + configs = [f"config{i}" for i in range(len(times_kt))] + x = range(len(configs)) + + plt.figure(figsize=(10,6)) + plt.plot(configs, times_kt, marker='s', label='KernelTuner') + plt.plot(configs, times_direct, marker='x', label='Direct') + plt.ylabel('Time (ms)') + plt.xlabel('Configuration') + plt.title('Kernel execution time per configuration') + plt.xticks(rotation=45) + plt.legend() + plt.grid(True) + plt.tight_layout() + plt.savefig("tilelang.png") + print("saved fig") + + +def tune(m, n, k): + a = torch.randn(m, k, device="cuda", dtype=torch.float16) + b = torch.randn(k, n, device="cuda", dtype=torch.float16) + c = torch.empty(m, n, device="cuda", dtype=torch.float16) + c_actual = a @ b + + args = [a, b, c] + tune_params = dict() + tune_params["M"] = [m] + tune_params["K"] = [k] + tune_params["N"] = [n] + tune_params["block_M"] = [64, 128, 256] + tune_params["block_N"] = [64, 128, 256] + tune_params["block_K"] = [32, 64, 128] + tune_params["num_stages"] = [2, 3, 4] + tune_params["panel_size"] = [4, 8] # equivalent to group size m in Triton + tune_params["num_threads"] = [64, 128, 256] + + restrictions = [ + # tile size budget + "block_M * block_N <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "block_M <= 4 * block_N", + "block_N <= 4 * block_M", + + # large K only with reasonably large M/N + "not (block_K == 128 and block_M < 64 and block_N < 64)", + ] + + tol = m * 2**(-11) + answer = [None, None, c_actual.cpu()] + + results, env = tune_kernel("matmul", matmul, m * n, args, tune_params, atol=tol, lang="generic_python", + call_function=call_tilelang, restrictions=restrictions, answer=answer, decorator="@tilelang.jit", verbose=False) + +if __name__ == "__main__": + #m, n, k = 1024, 1024, 1024 + m, n, k = 8192, 8192, 8192 + time(m, n, k) \ No newline at end of file diff --git a/examples/generic_python/matmul_old/tilus_matmul.py b/examples/generic_python/matmul_old/tilus_matmul.py new file mode 100644 index 000000000..e8996c201 --- /dev/null +++ b/examples/generic_python/matmul_old/tilus_matmul.py @@ -0,0 +1,266 @@ +import math + +import pandas +import tilus +import torch +from tilus import float16, float32, int32 +from tilus.utils import benchmark_func +from kernel_tuner import tune_kernel, run_kernel + + + +class MatmulV4(tilus.Script): + def __init__(self): + super().__init__() + self.block_m = 128 + self.block_n = 128 + self.block_k = 16 + self.num_warps = 4 + self.num_stages = 4 + + def __call__( + self, + m_size: int32, + n_size: int, + k_size: int, + a_ptr: ~float16, + b_ptr: ~float16, + c_ptr: ~float16, + ): + self.attrs.blocks = [ + self.utils.ceil_div(m_size, self.block_m), + self.utils.ceil_div(n_size, self.block_n), + ] + self.attrs.warps = self.num_warps + + block_m, block_n, block_k = self.block_m, self.block_n, self.block_k + offset_m: int32 = block_m * self.blockIdx.x + offset_n: int32 = block_n * self.blockIdx.y + + ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) + gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) + sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) + sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) + acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) + + for stage in range(self.num_stages - 1): + offset_k = stage * self.block_k + self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) + self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) + self.copy_async_commit_group() + + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + current_stage: int32 = 0 + preload_stage: int32 = self.num_stages - 1 + for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): + # computation for current tile + a = self.load_shared(sa[current_stage]) + b = self.load_shared(sb[current_stage]) + self.dot(a, b, acc, out=acc) + + # preload the next tile of A and B into shared memory + preload_offset_k = offset_k + (self.num_stages - 1) * block_k + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) + self.copy_async_commit_group() + + # update the stage + current_stage = (current_stage + 1) % self.num_stages + preload_stage = (preload_stage + 1) % self.num_stages + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + self.free_shared(sa) + self.free_shared(sb) + + casted_acc = self.cast(acc, dtype=float16) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + + +class MatmulGroupedOrdering(tilus.Script): + def __init__(self): + super().__init__() + self.block_m = 128 + self.block_n = 128 + self.block_k = 16 + self.num_warps = 4 + self.num_stages = 4 + self.group_size_m = 8 + + def __call__( + self, + m_size: int32, + n_size: int, + k_size: int, + a_ptr: ~float16, + b_ptr: ~float16, + c_ptr: ~float16, + ): + block_m, block_n, block_k = self.block_m, self.block_n, self.block_k + + num_pid_m = self.utils.ceil_div(m_size, block_m) + num_pid_n = self.utils.ceil_div(n_size, block_n) + self.attrs.blocks = [num_pid_m * num_pid_n] + + pid = self.blockIdx.x + num_pid_in_group = self.group_size_m * num_pid_n + group_id = pid // num_pid_in_group + + first_pid_m = group_id * self.group_size_m + group_size_m = min(num_pid_m - first_pid_m, self.group_size_m) + + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + self.attrs.warps = self.num_warps + offset_m: int32 = pid_m * block_m + offset_n: int32 = pid_n * block_n + + ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) + gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) + sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) + sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) + acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) + + for stage in range(self.num_stages - 1): + offset_k = stage * self.block_k + self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) + self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) + self.copy_async_commit_group() + + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + current_stage: int32 = 0 + preload_stage: int32 = self.num_stages - 1 + for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): + # computation for current tile + a = self.load_shared(sa[current_stage]) + b = self.load_shared(sb[current_stage]) + self.dot(a, b, acc, out=acc) + + # preload the next tile of A and B into shared memory + preload_offset_k = offset_k + (self.num_stages - 1) * block_k + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) + self.copy_async_commit_group() + + # update the stage + current_stage = (current_stage + 1) % self.num_stages + preload_stage = (preload_stage + 1) % self.num_stages + self.copy_async_wait_group(n=self.num_stages - 2) + self.sync() + + self.free_shared(sa) + self.free_shared(sb) + + casted_acc = self.cast(acc, dtype=float16) + gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) + self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + + + +def main(): + headers = ["m", "n", "k", "name", "latency (ms)", "tflops"] + workloads = [ + [4096, 4096, 4096], + [1024, 1024, 14336], + ] + + rows = [] + for m, n, k in workloads: + matmul = MatmulGroupedOrdering() #MatmulV4() + + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + matmul(m, n, k, a, b, c_actual) + + # check correctness + torch.testing.assert_close(c_expect, c_actual) + + # benchmark + for name, func in [ + ("torch", lambda: torch.matmul(a, b, out=c_expect)), + ("tilus", lambda: matmul(m, n, k, a, b, c_actual)), + ]: + latency = benchmark_func(func, warmup=5, repeat=20) + tflops = 2 * m * n * k / latency * 1e-9 + rows.append([m, n, k, name, latency, tflops]) + + df = pandas.DataFrame(rows, columns=headers) + print(df) + +def call_tilus(kernel_function, args, kwargs, grid, threads, params): + kernel_function(*args, **kwargs) + + +def tune_matmul(m, n, k): + a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n #(m, n) + args = [m, n, k, a, b, c_actual] + tune_params = dict() + tune_params["block_m"] = [32, 64, 128, 256] + tune_params["block_n"] = [32, 64, 128, 256] + tune_params["block_k"] = [32, 64, 128] + tune_params["group_size_m"] = [4, 8] + tune_params["num_stages"] = [2, 3, 4] + tune_params["num_warps"] = [4, 8] + + + restrictions = [ + # tile size budget + "block_m * block_n <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "block_m <= 4 * block_n", + "block_n <= 4 * block_m", + + # large K only with reasonably large M/N + "not (block_k == 128 and block_m < 64 and block_n < 64)", + + # 32x32 requires 8 warps + "not (block_m == 32 and block_n == 32 and num_warps < 8)", + ] + + + answer = [None] * 6 + answer[-1] = c_expect.cpu() + atol = 1e-2 #m * 2**(-11) + + results, env = tune_kernel("MatmulGroupedOrdering", MatmulGroupedOrdering, size, args, tune_params, grid_div_x = ["block_m", "block_n"], + answer = answer, atol=atol, restrictions=restrictions, + lang="generic_python", call_function=call_tilus, + block_size_names=["block_m", "block_n", "block_k"], strategy="simulated_annealing") + + +if __name__ == "__main__": + #m, n, k = 4096, 4096, 4096 + m, n, k = 8192, 8192, 8192 + tune_matmul(m, n, k) + + #main() \ No newline at end of file diff --git a/examples/generic_python/matmul_old/triton_matmul.py b/examples/generic_python/matmul_old/triton_matmul.py new file mode 100644 index 000000000..238a61c9f --- /dev/null +++ b/examples/generic_python/matmul_old/triton_matmul.py @@ -0,0 +1,383 @@ +import torch + +import triton +import triton.language as tl + +from kernel_tuner import tune_kernel +from kernel_tuner import run_kernel + +def get_cuda_autotune_config(): + return [ + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, + num_warps=8), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, + num_warps=2), + triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, + num_warps=2), + # Good config for fp8 inputs. + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, + num_warps=8), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, + num_warps=8), + triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4), + triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, + num_warps=4) + ] + +''' +@triton.autotune( + configs=get_cuda_autotune_config(), + key=['M', 'N', 'K'], +) +''' +#@triton.jit +def matmul_kernel( + # Pointers to matrices + a_ptr, b_ptr, c_ptr, + # Matrix dimensions + M, N, K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` + # by to get the element one row down (A has M rows). + stride_am, stride_ak, # + stride_bk, stride_bn, # + stride_cm, stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # + GROUP_SIZE_M: tl.constexpr, num_stages, num_warps# +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ----------------------------------------------------------- + # Add some integer bound assumptions. + # This helps to guide integer analysis in the backend to optimize + # load/store offset address calculation + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + # ---------------------------------------------------------- + # Create pointers for the first blocks of A and B. + # We will advance this pointer as we move in the K direction + # and accumulate + # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers + # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # ----------------------------------------------------------- + # Iterate to compute a block of the C matrix. + # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block + # of fp32 values for higher accuracy. + # `accumulator` will be converted back to fp16 after the loop. + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + # We accumulate along the K dimension. + accumulator = tl.dot(a, b, accumulator) + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + c = accumulator.to(tl.float16) + + # ----------------------------------------------------------- + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def run_matmul(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + grid = lambda META: (triton.cdiv(m, META['BLOCK_SIZE_M']) * triton.cdiv(n, META['BLOCK_SIZE_N']), ) + + matmul_kernel[grid]( + a, b, c_actual, # + m, n, k, # + a.stride(0), a.stride(1), # + b.stride(0), b.stride(1), # + c_actual.stride(0), c_actual.stride(1), + 128, 256, 64, 8 + ) + + torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) + + +def run_matmul_kt(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n + + args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1)] + + params = {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M":4, "num_stages":3, "num_warps":4} + + result = run_kernel("matmul_kernel", matmul_kernel, size, args, params=params, grid_div_x=["BLOCK_SIZE_N", "BLOCK_SIZE_M"], + lang="generic_python", decorator="@triton.jit", call_function=call_triton, + block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) + c_res = result[2] + + assert torch.allclose(c_res, c_expect.cpu(), atol=1e-2, rtol=1e-1) + + + + + + +def call_triton(kernel_function, args, kwargs, grid, threads, params): + #print("using grid: ", grid) + #print("args: ", args) + #print("kwargs: ", kwargs) + kernel_function[grid](*args, **kwargs) + + + + +def check_time(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n + args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1)] + tune_params = dict() + tune_params["BLOCK_SIZE_M"] = [64, 128] + tune_params["BLOCK_SIZE_N"] = [64, 128] + tune_params["BLOCK_SIZE_K"] = [32, 64] + tune_params["GROUP_SIZE_M"] = [4, 8] + tune_params["num_stages"] = [3] + tune_params["num_warps"] = [4] + + restrictions = [ + # tile size budget + "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", + "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", + + # large K only with reasonably large M/N + "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", + + # 32x32 requires 8 warps + "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", + ] + + grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] + + answer = [None] * 12 + answer[2] = c_expect.cpu() + atol = 1e-2 #m * 2**(-11) + + + + results_ours, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, + restrictions=restrictions, iterations=100, answer=answer, atol=atol, + lang="generic_python", decorator="@triton.jit", call_function=call_triton, + block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) + + + results_prev, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, + restrictions=restrictions, iterations=100, answer=answer, atol=atol, + lang="triton", block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) + + + + import time + num_repeats = 100 + times_direct = [] + for config in results_prev: + + bs_m = config["BLOCK_SIZE_M"] + bs_n = config["BLOCK_SIZE_N"] + bs_k = config["BLOCK_SIZE_K"] + gs_m = config["GROUP_SIZE_M"] + num_stages = config["num_stages"] + num_warps = config["num_warps"] + + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + + grid = (triton.cdiv(m, bs_m) * triton.cdiv(n, bs_n, ), ) + jit_function = triton.jit(matmul_kernel) + + + jit_function[grid]( + a, b, c_actual, + m, n, k, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1), + bs_m, bs_n, bs_k, gs_m, num_stages, num_warps + ) + + + torch.allclose(c_expect.cpu(), c_actual.cpu(), atol=1e-2) + + + + for i in range(num_repeats): + times = [] + + #c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + + torch.cuda.synchronize() + start = time.time() + jit_function[grid]( + a, b, c_actual, + m, n, k, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1), + bs_m, bs_n, bs_k, gs_m, num_stages, num_warps + ) + + torch.cuda.synchronize() + times.append(time.time() - start) + + avg_time_ms = round((1000 * sum(times) / len(times)), 3) + times_direct.append(avg_time_ms) + print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, GROUP_SIZE_M={gs_m}, num_stages={num_stages}, num_warps={num_warps}, time={avg_time_ms}ms") + + + import matplotlib.pyplot as plt + + # Extract times + times_prev = [cfg['time'] for cfg in results_prev] + times_ours = [cfg['time'] for cfg in results_ours] + + # x-axis labels + configs = [f"config{i}" for i in range(len(times_prev))] + x = range(len(configs)) + + plt.figure(figsize=(10,6)) + plt.plot(configs, times_prev, marker='o', label='Triton tuned') + plt.plot(configs, times_ours, marker='s', label='Generic tuned') + plt.plot(configs, times_direct, marker='x', label='Direct') + plt.ylabel('Time (ms)') + plt.xlabel('Configuration') + plt.title('Kernel execution time per configuration') + plt.xticks(rotation=45) + plt.legend() + plt.grid(True) + plt.tight_layout() + plt.savefig("ouptut.png") + + + + + + +def tune_matmul(m, n, k): + a = torch.rand(m, k, dtype=torch.float16).cuda() + b = torch.rand(k, n, dtype=torch.float16).cuda() + c_actual = torch.empty(m, n, dtype=torch.float16).cuda() + c_expect = a @ b + + size = m * n + args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), + c_actual.stride(0), c_actual.stride(1)] + tune_params = dict() + tune_params["BLOCK_SIZE_M"] = [64, 128, 256] + tune_params["BLOCK_SIZE_N"] = [64, 128, 256] + tune_params["BLOCK_SIZE_K"] = [32, 64, 128] + tune_params["GROUP_SIZE_M"] = [4, 8] + tune_params["num_stages"] = [2, 3, 4] + tune_params["num_warps"] = [4, 8] + + restrictions = [ + # tile size budget + "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", + + # aspect ratio <= 4 (no max/min allowed, so expand manually) + "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", + "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", + + # large K only with reasonably large M/N + "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", + + # 32x32 requires 8 warps + "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", + ] + + grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] + + answer = [None] * 12 + answer[2] = c_expect.cpu() + + results, env = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, + answer = answer, atol=4.0, restrictions=restrictions, + lang="generic_python", decorator="@triton.jit", call_function=call_triton, + block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], strategy="simulated_annealing") + + + + + +if __name__ == "__main__": + m, n, k = 8192, 8192, 8192 + #m, n, k = 4096, 4096, 4096 + #tune_matmul(m, n, k) + #check_time(m, n, k) + #run_matmul_kt(m, n, k) + + + + + + diff --git a/examples/generic_python/numba_vec_add.py b/examples/generic_python/numba_vec_add.py index 9f8dfe095..b5a2c56f3 100644 --- a/examples/generic_python/numba_vec_add.py +++ b/examples/generic_python/numba_vec_add.py @@ -16,7 +16,7 @@ def f(a, b, c): c[tid] = a[tid] + b[tid] -def call_numba(kernel_function, args, kwargs, grid, threads, params): +def call_numba(kernel_function, args, kwargs, grid, threads): numba_args = [] for arg in args: if isinstance(arg, torch.Tensor): @@ -26,8 +26,6 @@ def call_numba(kernel_function, args, kwargs, grid, threads, params): kernel_function[grid, threads](*args, **kwargs) - - def tune(): N = 100000 diff --git a/examples/generic_python/tilelang_vec_add.py b/examples/generic_python/tilelang_vec_add.py index 601b6a7b7..99463d112 100644 --- a/examples/generic_python/tilelang_vec_add.py +++ b/examples/generic_python/tilelang_vec_add.py @@ -37,7 +37,7 @@ def run_normal(): print("done") -def call_tilelang(kernel_function, args, kwargs, grid, threads, params): +def call_tilelang(kernel_function, args, kwargs): compiled_kernel = kernel_function(**kwargs) # cached, so second time only cache lookup is performed compiled_kernel(*args) diff --git a/examples/generic_python/tilus_naive_matmul.py b/examples/generic_python/tilus_naive_matmul.py index 42cdff86d..c53c22f73 100644 --- a/examples/generic_python/tilus_naive_matmul.py +++ b/examples/generic_python/tilus_naive_matmul.py @@ -68,7 +68,7 @@ def __call__( self.store_global(gc, acc_f16, offsets=[offset_m, offset_n]) -def call_tilus(kernel_function, args, kwargs, grid, threads, params): +def call_tilus(kernel_function, args, kwargs): kernel_function(*args, **kwargs) def main(): diff --git a/examples/generic_python/tilus_splitk_matmul.py b/examples/generic_python/tilus_splitk_matmul.py index be797299e..0695dcceb 100644 --- a/examples/generic_python/tilus_splitk_matmul.py +++ b/examples/generic_python/tilus_splitk_matmul.py @@ -133,7 +133,7 @@ def __call__( ) -def call_tilus(kernel_function, args, kwargs, grid, threads, params): +def call_tilus(kernel_function, args, kwargs): kernel_function(*args, **kwargs) diff --git a/examples/generic_python/tilus_vec_add.py b/examples/generic_python/tilus_vec_add.py index 938d9026c..9f842637c 100644 --- a/examples/generic_python/tilus_vec_add.py +++ b/examples/generic_python/tilus_vec_add.py @@ -8,9 +8,9 @@ FULL_PATH = Path(__file__).resolve() class VecAddV(tilus.Script): - def __init__(self): + def __init__(self, block_size_x=None): super().__init__() - self.block_size_x = 32 # number of threads per block + self.block_size_x = block_size_x # number of threads per block def __call__( self, @@ -38,7 +38,7 @@ def __call__( self.store_global(gc, c, offsets=[offset]) -def call_tilus(kernel_function, args, kwargs, grid, threads, params): +def call_tilus(kernel_function, args, kwargs): kernel_function(*args, **kwargs) @@ -91,7 +91,23 @@ def run(size): assert torch.allclose(results[-1], c_expect) +def tune_with_builtin(size): + TunedVecAdd = tilus.autotune("block_size_x", [32, 64, 128, 256, 512, 1024])(VecAddV) + vecadd = TunedVecAdd() + + a = torch.randn(size, dtype=torch.float32).cuda() + b = torch.randn(size, dtype=torch.float32).cuda() + c = torch.empty(size, dtype=torch.float32).cuda() + c_expect = a + b + + vecadd(size, a, b, c) + torch.cuda.synchronize() + + torch.testing.assert_close(c_expect, c) + + if __name__ == "__main__": - size = 100000000 - tune(size) - run(size) + size = 10000000 + #tune(size) + #run(size) + tune_with_builtin(size) diff --git a/examples/generic_python/triton_vec_add.py b/examples/generic_python/triton_vec_add.py index 2e982abfa..f69ae96d7 100644 --- a/examples/generic_python/triton_vec_add.py +++ b/examples/generic_python/triton_vec_add.py @@ -30,7 +30,7 @@ def add_kernel(x_ptr, # *Pointer* to first input vector. tl.store(output_ptr + offsets, output, mask=mask) -def call_triton(kernel_function, args, kwargs, grid, threads, params): +def call_triton(kernel_function, args, kwargs, grid): kernel_function[grid](*args, **kwargs) diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index d69ce3de0..3417fb4d7 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -292,11 +292,12 @@ def __deepcopy__(self, _): "call_function", ( """When the language Generic Python is used, a call function that calls the kernel in the Python - DSL must be specified. The function must take the following arguments: + DSL must be specified. The function must take the following positional arguments: :kernel_function: the callable function with the tuning parameters inserted. :args: list of kernel arguments, as provided by the user in the argument. :kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, the tuning parameter will be added as a keyword argument. + Optionally, the following arguments can be used. The order and name of the arguments must match. :grid: the launch grid (tuple with 3 values), as computed by KernelTuner :threads: the thread block size (tuple with 3 values), as computed by KernelTuner :params: dictionary with the values of the tuning params for the specific configuration.""", @@ -614,7 +615,7 @@ def tune_kernel( if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, util.normalize_call_function(call_function)) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -653,6 +654,10 @@ def tune_kernel( logging.debug("tuning_options: %s", util.get_config_string(tuning_options)) logging.debug("device_options: %s", util.get_config_string(device_options)) + # the user-specific call function may of may not have optional grid, threads and params + # arguments. We normalize it so that it always accepts all arguments. + kernel_options.call_function = util.normalize_call_function(kernel_options.call_function) + # check whether the selected strategy and options are valid strategy_string = strategy if strategy: @@ -800,7 +805,7 @@ def run_kernel( if log: logging.basicConfig(filename=kernel_name + datetime.now().strftime("%Y%m%d-%H:%M:%S") + ".log", level=log) - kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, call_function) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, util.normalize_call_function(call_function)) _check_user_input(kernel_name, kernelsource, arguments, block_size_names) @@ -809,6 +814,10 @@ def run_kernel( kernel_options = Options([(k, opts[k]) for k in _kernel_options.keys()]) device_options = Options([(k, opts[k]) for k in _device_options.keys()]) + # the user-specific call function may of may not have optional grid, threads and params + # arguments. We normalize it so that it always accepts all arguments. + kernel_options.call_function = util.normalize_call_function(kernel_options.call_function) + # detect language and create the right device function interface dev = core.DeviceInterface(kernelsource, iterations=1, **device_options) diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 22509ae37..3eb27d86b 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -46,7 +46,7 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function raise ValueError("call_function must be supplied for language Generic Python") if not callable(call_function): raise TypeError(f"call_function of type {type(call_function)} is not a callable object.") - self.call_function = call_function # TODO ceck signature + self.call_function = call_function if not isinstance(kernel_name, str): raise TypeError("kernel_name should be a string, got ", type(kernel_name)) diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index 24c622546..beae6b5d5 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -978,6 +978,33 @@ def has_kw_argument(func, name): return lambda answer, result_host, atol: v(answer, result_host) +def normalize_call_function(v): + """Normalize a user-specified call function for language Generic Python. + + The user-specified function has three required positional arguments (kernel_function, + args, kwargs), and three optional keyword arguments: grid, threads and params. The + optional keyword arguments should appear in that order. We normalize the function + so that it always accepts grid, threads and params. + + Undefined behaviour if the passed function does not match the required signatures. + """ + def has_kw_argument(func, name): + sig = signature(func) + return name in sig.parameters + + if v is None: + return None + + if has_kw_argument(v, "grid"): + if has_kw_argument(v, "threads"): + if has_kw_argument(v, "params"): + return v + return lambda kernel_function, args, kwargs, grid, threads, params: v(kernel_function, args, kwargs, grid, threads) + return lambda kernel_function, args, kwargs, grid, threads, params: v(kernel_function, args, kwargs, grid) + return lambda kernel_function, args, kwargs, grid, threads, params: v(kernel_function, args, kwargs) + + + def parse_restrictions( restrictions: list[str], tune_params: dict, monolithic=False, format=None ) -> list[tuple[Union[Constraint, str], list[str]]]: diff --git a/test/test_util_functions.py b/test/test_util_functions.py index 4a1858f37..389da8b81 100644 --- a/test/test_util_functions.py +++ b/test/test_util_functions.py @@ -595,6 +595,28 @@ def verify2(answer, result_host, atol): assert v(1, 2, atol=3) + +def test_normalize_call_function_none(): + assert normalize_call_function(None) is None + +@pytest.mark.parametrize( + "func", + [ + lambda f, a, k: f(*a, **k), + lambda f, a, k, grid: f(*a, **k), + lambda f, a, k, grid, threads: f(*a, **k), + lambda f, a, k, grid, threads, params: f(*a, **k), + ], +) +def test_normalize_call_function(func): + v = normalize_call_function(func) + + def kernel(x): + return x + 1 + + assert v(kernel, (1,), {}, grid=1, threads=2, params=3) == 2 + + def test_process_cache(): def assert_open_cachefile_is_correctly_parsed(cache): with open(cache, "r") as cachefile: From 976da390af4fd1b5b821c4bc11510d5478ec8f7b Mon Sep 17 00:00:00 2001 From: "I.C. van Ooijen" Date: Sun, 8 Mar 2026 11:34:15 +0100 Subject: [PATCH 07/44] taichi matmul exmple --- .../generic_python/matmul/taichi_matmul.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 examples/generic_python/matmul/taichi_matmul.py diff --git a/examples/generic_python/matmul/taichi_matmul.py b/examples/generic_python/matmul/taichi_matmul.py new file mode 100644 index 000000000..6eca9d21e --- /dev/null +++ b/examples/generic_python/matmul/taichi_matmul.py @@ -0,0 +1,56 @@ +import torch +import taichi as ti + +from kernel_tuner import tune_kernel +from pathlib import Path + +FULL_PATH = Path(__file__).resolve() + +ti.init(arch=ti.gpu) + +# NOTE tuning works on vibranium. Taichi does not work on DAS6 +# TODO make sure this is zero-copy +@ti.kernel +def matmul(A: ti.types.ndarray(), B: ti.types.ndarray(), C: ti.types.ndarray()): + BLOCK_DIM = 16 + ti.loop_config(block_dim=BLOCK_DIM) + + K_dim = A.shape[1] + + for i, j in C: + sum = 0.0 + for k in range(K_dim): + sum += A[i, k] * B[k, j] + C[i, j] = sum + + +def call_taichi(kernel_function, args, kwargs): + kernel_function(*args, **kwargs) + +def tune(M, N, K): + torch_A = torch.rand((N, K), device='cuda', dtype=torch.float32) + torch_B = torch.rand((K, M), device='cuda', dtype=torch.float32) + torch_C = torch.empty((N, M), device='cuda', dtype=torch.float32) + + size = M * N + args = [torch_A, torch_B, torch_C] + tune_params = {"BLOCK_DIM": {4, 8, 16, 32, 64, 128, 256, 512, 1024}} + + answer = [None, None, (torch_A @ torch_B).cpu()] + + results, env = tune_kernel( + kernel_name="matmul", + kernel_source=FULL_PATH, + problem_size=size, + arguments=args, + tune_params=tune_params, + answer=answer, + lang="generic_python", + call_function=call_taichi, + ) + +if __name__ == "__main__": + tune(128, 128, 128) + + + From 51f5a6a4bed48b3b9e8c2d55492bc5c1c46afeb4 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Tue, 7 Apr 2026 10:56:26 +0200 Subject: [PATCH 08/44] working on experiments --- examples/__init__.py | 0 examples/generic_python/call_functions.py | 29 + examples/generic_python/cupy_copy.py | 19 + examples/generic_python/cutile_vec_add.py | 62 ++ .../flash_attention/tilelang_attention.py | 152 +++++ .../flash_attention/tilus_attention.py | 611 ++++++++++++++++++ .../flash_attention/triton_attention.py | 276 ++++++++ examples/generic_python/loopy_example.py | 38 ++ examples/generic_python/matmul/cupy_matmul.py | 105 +++ examples/generic_python/matmul/cute_matmul.py | 135 ++++ .../generic_python/matmul/helion_matmul.py | 20 + .../generic_python/matmul/numba_matmul.py | 148 +++++ examples/generic_python/matmul/test.py | 278 ++++++++ .../generic_python/matmul/tilelang_matmul.py | 18 +- .../generic_python/matmul/tilus_matmul.py | 3 +- examples/generic_python/matmul/warp_matmul.py | 131 ++++ .../normalization/tilelang_norm.py | 78 +++ .../normalization/tilus_norm.py | 132 ++++ .../normalization/triton_norm.py | 66 ++ examples/generic_python/pallas_vec_add.py | 23 + examples/generic_python/tilus_vec_add.py | 7 +- examples/generic_python/triton_vec_add.py | 2 + kernel_tuner/backends/generic_python.py | 5 +- kernel_tuner/core.py | 6 + .../kernel_sources/kernel_source_fn.py | 10 +- 25 files changed, 2336 insertions(+), 18 deletions(-) create mode 100644 examples/__init__.py create mode 100644 examples/generic_python/call_functions.py create mode 100644 examples/generic_python/cupy_copy.py create mode 100644 examples/generic_python/cutile_vec_add.py create mode 100644 examples/generic_python/flash_attention/tilelang_attention.py create mode 100644 examples/generic_python/flash_attention/tilus_attention.py create mode 100644 examples/generic_python/flash_attention/triton_attention.py create mode 100644 examples/generic_python/loopy_example.py create mode 100644 examples/generic_python/matmul/cupy_matmul.py create mode 100644 examples/generic_python/matmul/cute_matmul.py create mode 100644 examples/generic_python/matmul/helion_matmul.py create mode 100644 examples/generic_python/matmul/numba_matmul.py create mode 100644 examples/generic_python/matmul/test.py create mode 100644 examples/generic_python/matmul/warp_matmul.py create mode 100644 examples/generic_python/normalization/tilelang_norm.py create mode 100644 examples/generic_python/normalization/tilus_norm.py create mode 100644 examples/generic_python/normalization/triton_norm.py create mode 100644 examples/generic_python/pallas_vec_add.py diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/generic_python/call_functions.py b/examples/generic_python/call_functions.py new file mode 100644 index 000000000..6b321c4a2 --- /dev/null +++ b/examples/generic_python/call_functions.py @@ -0,0 +1,29 @@ +import torch + +def call_tilus(kernel_function, args, kwargs): + kernel_function(*args, **kwargs) + +def call_triton(kernel_function, args, kwargs, grid, threads, params): + if "num_warps" in params.keys(): + kwargs["num_warps"] = params["num_warps"] + if "num_stages" in params.keys(): + kwargs["num_stages"] = params["num_stages"] + + torch.cuda.nvtx.range_push("kt call") + kernel_function[grid](*args, **kwargs) + torch.cuda.nvtx.range_pop() + +def call_tilelang(kernel_function, args, kwargs): + compiled_kernel = kernel_function(**kwargs) + compiled_kernel(*args) + +def call_numba(kernel_function, args, kwargs, grid, threads): + from numba import cuda + + numba_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + numba_args.append(cuda.as_cuda_array(arg)) + else: + numba_args.append(arg) + kernel_function[grid, threads](*args, **kwargs) \ No newline at end of file diff --git a/examples/generic_python/cupy_copy.py b/examples/generic_python/cupy_copy.py new file mode 100644 index 000000000..8d9110209 --- /dev/null +++ b/examples/generic_python/cupy_copy.py @@ -0,0 +1,19 @@ +import cupy +from cupyx import jit + +@jit.rawkernel() +def elementwise_copy(x, y, size): + tid = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x + ntid = jit.gridDim.x * jit.blockDim.x + for i in range(tid, size, ntid): + y[i] = x[i] + +size = cupy.uint32(2 ** 22) +x = cupy.random.normal(size=(size,), dtype=cupy.float32) +y = cupy.empty((size,), dtype=cupy.float32) + +elementwise_copy((128,), (1024,), (x, y, size)) # RawKernel style +assert (x == y).all() + +elementwise_copy[128, 1024](x, y, size) # Numba style +assert (x == y).all() \ No newline at end of file diff --git a/examples/generic_python/cutile_vec_add.py b/examples/generic_python/cutile_vec_add.py new file mode 100644 index 000000000..e95302045 --- /dev/null +++ b/examples/generic_python/cutile_vec_add.py @@ -0,0 +1,62 @@ +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +""" +Example demonstrating simple vector addition. +Shows how to perform elementwise operations on vectors. +Does not work on das vu, because we need cuda 13.1 +""" + +import cupy as cp +import numpy as np +import cuda.tile as ct + + +@ct.kernel +def vector_add(a, b, c, tile_size: ct.Constant[int]): + # Get the 1D pid + pid = ct.bid(0) + + # Load input tiles + a_tile = ct.load(a, index=(pid,), shape=(tile_size,)) + b_tile = ct.load(b, index=(pid,), shape=(tile_size,)) + + # Perform elementwise addition + result = a_tile + b_tile + + # Store result + ct.store(c, index=(pid, ), tile=result) + + +def test(): + # Create input data + vector_size = 2**12 + tile_size = 2**4 + grid = (ct.cdiv(vector_size, tile_size), 1, 1) + + rng = cp.random.default_rng() + a = rng.random(vector_size) + b = rng.random(vector_size) + c = cp.zeros_like(a) + + # Launch kernel + ct.launch(cp.cuda.get_current_stream(), + grid, # 1D grid of processors + vector_add, + (a, b, c, tile_size)) + + # Copy to host only to compare + a_np = cp.asnumpy(a) + b_np = cp.asnumpy(b) + c_np = cp.asnumpy(c) + + # Verify results + expected = a_np + b_np + np.testing.assert_array_almost_equal(c_np, expected) + + print("✓ vector_add_example passed!") + + +if __name__ == "__main__": + test() \ No newline at end of file diff --git a/examples/generic_python/flash_attention/tilelang_attention.py b/examples/generic_python/flash_attention/tilelang_attention.py new file mode 100644 index 000000000..5b2be204a --- /dev/null +++ b/examples/generic_python/flash_attention/tilelang_attention.py @@ -0,0 +1,152 @@ +# example taken from https://github.com/tile-ai/tilelang/blob/main/examples/flash_attention/example_mha_fwd_bshd.py +import torch +import torch.nn.functional as F +import tilelang +from tilelang.autotuner import * +import tilelang.language as T +import itertools +import argparse +from functools import partial + + +def get_configs(): + iter_params = dict(block_M=[64], block_N=[64], num_stages=[1], threads=[128]) + return [dict(zip(iter_params, values)) for values in itertools.product(*iter_params.values())] + + +@autotune(configs=get_configs(), warmup=10, rep=10) +@tilelang.jit( + out_idx=[3], + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + }, +) +def flashattn(batch, heads, seq_len, dim, is_causal, block_M=64, block_N=64, num_stages=1, threads=128): + scale = (1.0 / dim) ** 0.5 * 1.44269504 # log2(e) + shape = [batch, seq_len, heads, dim] + dtype = T.float16 + accum_dtype = T.float32 + + @T.prim_func + def main( + Q: T.Tensor(shape, dtype), + K: T.Tensor(shape, dtype), + V: T.Tensor(shape, dtype), + Output: T.Tensor(shape, dtype), + ): + with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=threads) as (bx, by, bz): + Q_shared = T.alloc_shared([block_M, dim], dtype) + K_shared = T.alloc_shared([block_N, dim], dtype) + V_shared = T.alloc_shared([block_N, dim], dtype) + O_shared = T.alloc_shared([block_M, dim], dtype) + acc_s = T.alloc_fragment([block_M, block_N], accum_dtype) + acc_s_cast = T.alloc_fragment([block_M, block_N], dtype) + acc_o = T.alloc_fragment([block_M, dim], accum_dtype) + scores_max = T.alloc_fragment([block_M], accum_dtype) + scores_max_prev = T.alloc_fragment([block_M], accum_dtype) + scores_scale = T.alloc_fragment([block_M], accum_dtype) + scores_sum = T.alloc_fragment([block_M], accum_dtype) + logsum = T.alloc_fragment([block_M], accum_dtype) + + T.copy(Q[bz, bx * block_M : (bx + 1) * block_M, by, :], Q_shared) + T.fill(acc_o, 0) + T.fill(logsum, 0) + T.fill(scores_max, -T.infinity(accum_dtype)) + + loop_range = ( + T.min(T.ceildiv(seq_len, block_N), T.ceildiv((bx + 1) * block_M, block_N)) if is_causal else T.ceildiv(seq_len, block_N) + ) + + for k in T.Pipelined(loop_range, num_stages=num_stages): + T.copy(K[bz, k * block_N : (k + 1) * block_N, by, :], K_shared) + if is_causal: + for i, j in T.Parallel(block_M, block_N): + acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0, -T.infinity(acc_s.dtype)) + else: + for i, j in T.Parallel(block_M, block_N): + acc_s[i, j] = T.if_then_else(k * block_N + j >= seq_len, -T.infinity(acc_s.dtype), 0) + T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + T.copy(scores_max, scores_max_prev) + T.fill(scores_max, -T.infinity(accum_dtype)) + T.reduce_max(acc_s, scores_max, dim=1, clear=False) + for i in T.Parallel(block_M): + scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) + for i in T.Parallel(block_M): + scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) + for i, j in T.Parallel(block_M, block_N): + acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) + T.reduce_sum(acc_s, scores_sum, dim=1) + for i in T.Parallel(block_M): + logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] + T.copy(acc_s, acc_s_cast) + + for i, j in T.Parallel(block_M, dim): + acc_o[i, j] *= scores_scale[i] + + T.copy(V[bz, k * block_N : (k + 1) * block_N, by, :], V_shared) + T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Parallel(block_M, dim): + acc_o[i, j] /= logsum[i] + T.copy(acc_o, O_shared) + T.copy(O_shared, Output[bz, bx * block_M : (bx + 1) * block_M, by, :]) + + return main + + +def ref_program(Q, K, V, is_causal): + dim = Q.size(-1) + scores = torch.einsum("bqhd,bkhd->bhqk", Q, K) + scores = scores / torch.sqrt(torch.tensor(dim, dtype=scores.dtype)) + if is_causal: + seq_len = Q.size(1) + mask = torch.tril(torch.ones(seq_len, seq_len, device=scores.device)) + mask = mask.unsqueeze(0).unsqueeze(0) + scores = scores.masked_fill(mask == 0, float("-inf")) + attention_weights = F.softmax(scores, dim=-1) + output = torch.einsum("bhqk,bkhd->bqhd", attention_weights, V) + return output + + +def main( + batch: int = 8, + heads: int = 32, + seq_len: int = 4096, + dim: int = 128, + is_causal: bool = False, + tune: bool = False, +): + flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim + total_flops = 2 * flops_per_matmul + if is_causal: + total_flops *= 0.5 + + if not tune: + #kernel = flashattn(batch, heads, seq_len, dim, is_causal, block_M=128, block_N=128, num_stages=1, threads=128) + # Changed block size so fits in shared mem + kernel = flashattn(batch, heads, seq_len, dim, is_causal, block_M=64, block_N=64, num_stages=1, threads=128) + ref_program_processed = partial(ref_program, is_causal=is_causal) + profiler = kernel.get_profiler() + profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01) + print("All checks pass.") + latency = profiler.do_bench(ref_program_processed, warmup=500) + print("Ref: {:.2f} ms".format(latency)) + print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) + latency = profiler.do_bench(warmup=500) + print("Tile-lang: {:.2f} ms".format(latency)) + print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) + else: + best_result = flashattn(batch, heads, seq_len, dim, is_causal) + best_latency = best_result.latency + best_config = best_result.config + ref_latency = best_result.ref_latency + print(f"Best latency: {best_latency}") + print(f"Best TFlops: {total_flops / best_latency * 1e-9}") + print(f"Best config: {best_config}") + print(f"Ref latency: {ref_latency}") + + +if __name__ == "__main__": + + main() \ No newline at end of file diff --git a/examples/generic_python/flash_attention/tilus_attention.py b/examples/generic_python/flash_attention/tilus_attention.py new file mode 100644 index 000000000..5169d9316 --- /dev/null +++ b/examples/generic_python/flash_attention/tilus_attention.py @@ -0,0 +1,611 @@ +# Code taken from https://github.com/NVIDIA/tilus/blob/main/examples/attention/flash_attention_v3.py +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import numpy as np +import pandas as pd +import tilus +import torch +from tilus import boolean, f32, int32, void_p +from hidet.ir import DataType +from tilus.ir import RegisterTensor, SharedTensor +from tilus.ir.tensor import GlobalTensor +from tilus.utils import benchmark_func, cdiv + +pd.options.display.max_columns = None +pd.options.display.width = 1000 + + +@tilus.autotune("num_warps", [4, 8]) +@tilus.autotune("block_q", [32, 64, 128]) +@tilus.autotune("block_kv", [32, 64, 128]) +@tilus.autotune("split_kv", [-1, 512, 1024, 4096]) +@tilus.autotune("keep_q_in_regs", [False, True]) +class FlashAttention(tilus.Script): + LOG2_E = 1.4426950408889634 # log2(e) + + debug_schedule = dict( + num_warps=4, + block_q=64, + block_kv=64, + split_kv=-1, + keep_q_in_regs=False, + ) + + def __init__( + self, + dtype: DataType, + num_heads: int, + num_heads_kv: int, + head_size: int, + num_warps: int, + block_q: int, + block_kv: int, + split_kv: int, + keep_q_in_regs: bool, + ): + super().__init__() + self.dtype: DataType = dtype + self.num_heads = num_heads + self.num_heads_kv = num_heads_kv + self.head_size = head_size + self.num_warps = num_warps + self.block_q = block_q + self.block_kv = block_kv + self.split_kv = split_kv + self.keep_q_in_regs = keep_q_in_regs + self.score_scale = float(1.0 / np.sqrt(head_size)) + self.group_heads = num_heads // num_heads_kv + + assert self.split_kv % self.block_kv == 0 or self.split_kv == -1, ( + "split_kv must be a multiple of block_kv or -1" + ) + + # determine layout + self.sv_config = self.cuda.resolve_dot_config( + dtype, f32, m=block_q, n=head_size, k=block_kv, warp_m=num_warps, warp_n=1 + ) + + def apply_mask(self, score: RegisterTensor, q_offset: int32, kv_offset: int32): + mask = self.register_tensor( + dtype=boolean, + shape=[self.block_q, self.block_kv], + init=lambda i, j: i + q_offset >= j + kv_offset, + ) + self.assign(score, score + self.where(mask, x=0.0, y=-1e6)) + + def softmax_rescale( + self, + score: RegisterTensor, + m: RegisterTensor, + l: RegisterTensor, + o: RegisterTensor, + ) -> RegisterTensor: + scale = self.score_scale * self.LOG2_E # log2(e) * score_scale + cur_m = self.max(score, dim=1, keepdim=True) * scale # [block_q, 1] + new_m = self.maximum(m, cur_m) # [block_q, 1] + rp = self.exp2(score * scale - new_m) # [block_q, block_kv] + m_scale = self.exp2(m - new_m) + self.assign(o, o * m_scale) + self.assign(l, l * m_scale + self.sum(rp, dim=1, keepdim=True)) + self.assign(m, new_m) + return rp.to(self.dtype) + + def attention_iteration( + self, + bs: int32, + kv_offset: int32, + q_offset: int32, + head: int32, + gk: GlobalTensor, + gv: GlobalTensor, + sq: SharedTensor, # f16[block_q, head_size] + rq: RegisterTensor, # f16[block_q, head_size] + sk: SharedTensor, # f16[block_kv, head_size], + sv: SharedTensor, # f16[block_kv, head_size], + o: RegisterTensor, # f32[block_q, head_size] + m: RegisterTensor, # f32[block_q, 1] + l: RegisterTensor, # f32[block_q, 1] + check_bounds: bool, + ): + if not self.keep_q_in_regs: + self.load_shared(sq, out=rq) + # wait for the async copy of k to finish + self.copy_async_wait_group(0) + self.sync() + self.copy_async( + gv, + sv, + offsets=[bs, kv_offset, head // self.group_heads, 0], + dims=[1, 3], + check_bounds=check_bounds, + ) + self.copy_async_commit_group() + + # issue the async copy for v and perform dot(q, k) + rk = self.load_shared(sk) # [block_kv, head_size] + score = self.dot(rq, rk.transpose(), acc_dtype=f32) # [block_q, block_kv] + + if check_bounds: + self.apply_mask(score, q_offset, kv_offset) # apply causal mask + + # wait for the async copy of v to finish + self.copy_async_wait_group(0) + self.sync() + self.copy_async( + gk, + sk, + offsets=[bs, kv_offset + self.block_kv, head // self.group_heads, 0], + dims=[1, 3], + check_bounds=check_bounds, + ) + self.copy_async_commit_group() + + # load v to register + rv = self.load_shared(sv) # [block_kv, head_size] + + # online softmax + rp = self.softmax_rescale(score, m=m, l=l, o=o) + + # pv + cur_o = self.dot(rp, rv, acc_dtype=f32) # [block_q, head_size] + self.annotate_layout(cur_o, self.sv_config.lc) + self.assign(o, o + cur_o) + + def main_loop( + self, + gq: GlobalTensor, + gk: GlobalTensor, + gv: GlobalTensor, + o: RegisterTensor, + m: RegisterTensor, + l: RegisterTensor, + ): + # calculate offsets + q_offset = self.blockIdx.x * self.block_q + kv_start_offset = 0 if self.split_kv == -1 else self.blockIdx.y * self.split_kv + + if q_offset + self.block_q <= kv_start_offset: + return + + head = self.blockIdx.z % self.num_heads + bs = self.blockIdx.z // self.num_heads + + sq = self.shared_tensor(dtype=self.dtype, shape=[self.block_q, self.head_size]) + sk = self.shared_tensor(dtype=self.dtype, shape=[self.block_kv, self.head_size]) + sv = self.shared_tensor(dtype=self.dtype, shape=[self.block_kv, self.head_size]) + + rq = self.register_tensor(dtype=self.dtype, shape=[self.block_q, self.head_size]) + + # copy q to shared memory + self.copy_async( + gq, sq, offsets=[bs, q_offset, head, 0], dims=[1, 3], check_bounds=True + ) + self.copy_async_wait_all() + self.sync() + + # copy q to registers if not keeping in shared memory + if self.keep_q_in_regs: + self.load_shared(sq, out=rq) # [block_q, head_size] + self.free_shared(sq) + + # issue a copy of gk + self.copy_async(gk, sk, offsets=[bs, 0, head // self.group_heads, 0], dims=[1, 3]) + self.copy_async_commit_group() + + kv_offset_inner_end = (q_offset + 1) // self.block_kv * self.block_kv + if self.split_kv != -1: + kv_offset_inner_end = min( + kv_offset_inner_end, kv_start_offset + self.split_kv + ) + for kv_offset in range(kv_start_offset, kv_offset_inner_end, self.block_kv): + self.attention_iteration( + bs, + kv_offset, + q_offset, + head, + gk, + gv, + sq, + rq, + sk, + sv, + o, + m, + l, + check_bounds=False, + ) + + kv_offset_end = q_offset + self.block_q + if self.split_kv != -1: + kv_offset_end = min(kv_offset_end, kv_start_offset + self.split_kv) + for kv_offset in range(kv_offset_inner_end, kv_offset_end, self.block_kv): + self.attention_iteration( + bs, + kv_offset, + q_offset, + head, + gk, + gv, + sq, + rq, + sk, + sv, + o, + m, + l, + check_bounds=True, + ) + + self.copy_async_wait_group(0) + self.sync() + self.free_shared(sk) + self.free_shared(sv) + if not self.keep_q_in_regs: + self.free_shared(sq) + + def store_back( + self, + o: RegisterTensor, + l: RegisterTensor, + m: RegisterTensor, + o_ptr: void_p, + batch_size: int, + q_len: int32, + ): + # o: [block_q, head_size] + # m: [block_q, 1] + # l: [block_q, 1] + go = self.global_view( + o_ptr, + dtype=self.dtype, + shape=[batch_size, q_len, self.num_heads, self.head_size], + ) + o = o / l + o_f16 = self.cast(o, dtype=self.dtype) # [block_q, head_size] + so = self.shared_tensor(dtype=self.dtype, shape=[self.block_q, self.head_size]) + + head = self.blockIdx.z % self.num_heads + q_offset = self.blockIdx.x * self.block_q + bs = self.blockIdx.z // self.num_heads + + if self.split_kv == -1: + self.store_shared(so, o_f16) + self.sync() + self.store_global( + go, + self.load_shared(so), + offsets=[bs, q_offset, head, 0], + dims=[1, 3], + ) + else: + num_q_blocks = cdiv(q_len, self.block_q) + semaphores = self.global_tensor( + dtype=int32, + shape=[num_q_blocks, batch_size, self.num_heads], + requires_clean=True, + ) + gm = self.global_tensor( + dtype=f32, + shape=[num_q_blocks, batch_size, self.num_heads, self.block_q], + requires_clean=False, + ) + gl = self.global_tensor( + dtype=f32, + shape=[num_q_blocks, batch_size, self.num_heads, self.block_q], + requires_clean=False, + ) + semaphore = semaphores[self.blockIdx.x, bs, head].item_ptr() + + sm = self.shared_tensor(dtype=f32, shape=[self.block_q]) + sl = self.shared_tensor(dtype=f32, shape=[self.block_q]) + + self.lock_semaphore(semaphore, value=self.blockIdx.y) + + # load previous o, m and l and merge with the current results + if self.blockIdx.y > 0: + self.copy_async(gm, sm, offsets=[self.blockIdx.x, bs, head, 0], dims=[3]) + self.copy_async(gl, sl, offsets=[self.blockIdx.x, bs, head, 0], dims=[3]) + self.copy_async(go, so, offsets=[bs, q_offset, head, 0], dims=[1, 3]) + self.copy_async_wait_all() + self.sync() + lhs_o = self.load_shared(so) + lhs_m = self.load_shared(sm).unsqueeze(1) + lhs_l = self.load_shared(sl).unsqueeze(1) + rhs_o = o_f16 + rhs_m = m + rhs_l = l + m = self.maximum(lhs_m, rhs_m) + lhs_ll = lhs_l * self.exp(lhs_m - m) + rhs_ll = rhs_l * self.exp(rhs_m - m) + l = lhs_ll + rhs_ll + o_f16 = lhs_o * self.cast( + lhs_ll / l, dtype=self.dtype + ) + rhs_o * self.cast(rhs_ll / l, dtype=self.dtype) + self.sync() + + # store the results to so and load it + self.store_shared(so, o_f16) + self.store_shared(sm, m.squeeze(dim=1)) + self.store_shared(sl, l.squeeze(dim=1)) + self.sync() + + # store the results to global memory and release the semaphore + self.store_global( + go, + self.load_shared(so), + offsets=[bs, q_offset, head, 0], + dims=[1, 3], + ) + self.store_global( + gm, + self.load_shared(sm), + offsets=[self.blockIdx.x, bs, head, 0], + dims=[3], + ) + self.store_global( + gl, + self.load_shared(sl), + offsets=[self.blockIdx.x, bs, head, 0], + dims=[3], + ) + self.sync() + + self.free_shared(sm) + self.free_shared(sl) + + # release the semaphore + self.release_semaphore( + semaphore, + value=self.blockIdx.y + 1 + if (self.blockIdx.y + 1) * self.split_kv < q_offset + self.block_q + else 0, + ) + self.free_shared(so) + + def __call__( + self, + batch_size: int, + q_len: int32, + kv_len: int32, + q_ptr: void_p, + k_ptr: void_p, + v_ptr: void_p, + o_ptr: void_p, + ): + self.attrs.warps = self.num_warps + self.attrs.blocks = ( + cdiv(q_len, self.block_q), + cdiv(kv_len, self.split_kv) if self.split_kv != -1 else 1, + self.num_heads * batch_size, + ) + + gq = self.global_view( + q_ptr, + dtype=self.dtype, + shape=[batch_size, q_len, self.num_heads, self.head_size], + ) + gk = self.global_view( + k_ptr, + dtype=self.dtype, + shape=[batch_size, kv_len, self.num_heads_kv, self.head_size], + ) + gv = self.global_view( + v_ptr, + dtype=self.dtype, + shape=[batch_size, kv_len, self.num_heads_kv, self.head_size], + ) + + o = self.register_tensor( + dtype=f32, shape=[self.block_q, self.head_size], init=0.0 + ) + m = self.register_tensor( + dtype=f32, shape=[self.block_q, 1], init=-1e6 + ) # rowmax(score) + l = self.register_tensor( + dtype=f32, shape=[self.block_q, 1], init=0.0 + ) # rowsum(exp(score - m)) + + self.main_loop(gq, gk, gv, o, m, l) + + self.store_back(o, l, m, o_ptr=o_ptr, batch_size=batch_size, q_len=q_len) + + +def flash_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +): + """ + Flash attention function for variable length sequences. + + Parameters + ---------- + q: torch.Tensor + The query tensor of shape (bs, seqlen, num_heads, head_size). + + k: torch.Tensor + The key tensor of shape (bs, seqlen, num_heads_kv, head_size). + + v: torch.Tensor + The value tensor of shape (bs, seqlen, num_heads_kv, head_size). + + Returns + ------- + o: torch.Tensor + The output tensor of shape (bs, seqlen, num_heads, head_size). + """ + out = torch.empty_like(q) + FlashAttention( + dtype=tilus.float16, + num_heads=q.size(2), + num_heads_kv=k.size(2), + head_size=q.size(3), + )(q.size(0), q.size(1), k.size(1), q, k, v, out) + return out + + +def flash_attention_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +): + bs, seqlen, num_heads, head_size = q.size() + _, _, num_heads_kv, _ = k.size() + assert q.size(0) == k.size(0) == v.size(0), "Batch size must match for q, k, and v." + assert q.size(1) == k.size(1) == v.size(1), ( + "Sequence length must match for q, k, and v." + ) + assert q.size(3) == k.size(3) == v.size(3), "Head size must match for q, k, and v." + assert k.size(2) == v.size(2), "Number of heads in k and v must match." + assert num_heads % num_heads_kv == 0, ( + "Number of heads must be divisible by number of kv heads." + ) + + k = torch.repeat_interleave(k, num_heads // num_heads_kv, dim=2) + v = torch.repeat_interleave(v, num_heads // num_heads_kv, dim=2) + + q = torch.transpose(q, 1, 2).reshape(bs * num_heads, seqlen, head_size) + k = torch.transpose(k, 1, 2).reshape(bs * num_heads, seqlen, head_size) + v = torch.transpose(v, 1, 2).reshape(bs * num_heads, seqlen, head_size) + + score = torch.bmm(q, k.mT) / np.sqrt(head_size) # [bs * num_heads, seqlen, seqlen] + causal_mask = torch.tril(torch.ones(seqlen, seqlen, dtype=torch.bool), diagonal=0).to( + q.device + ) + causal_mask = causal_mask.unsqueeze(0) # [1, seqlen, seqlen] + causal_mask = causal_mask.expand( + bs * num_heads, seqlen, seqlen + ).contiguous() # [bs * num_heads, seqlen, seqlen] + score = score.masked_fill(causal_mask == 0, float("-inf")) + + o = torch.bmm( + torch.softmax(score.float(), dim=-1).to(q.dtype), v + ) # [bs * num_heads, seqlen, head_size] + o = o.reshape(bs, num_heads, seqlen, head_size).transpose(1, 2).contiguous() + return o + + +def flash_attention_flash_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, +): + try: + from flash_attn.cute.interface import flash_attn_func + + return flash_attn_func(q, k, v, causal=True) + except ImportError: + return flash_attention_reference(q, k, v) + + +def demo_flash_attention(): + for bs, seqlen, num_heads, head_size, num_heads_kv in [ + # [1, 8, 1, 64, 1], + [1, 4096, 32, 128, 8] + ]: + q = torch.rand(bs, seqlen, num_heads, head_size, dtype=torch.float16).cuda() + k = torch.rand(bs, seqlen, num_heads_kv, head_size, dtype=torch.float16).cuda() + v = torch.rand(bs, seqlen, num_heads_kv, head_size, dtype=torch.float16).cuda() + flash_attention(q, k, v) + torch.cuda.synchronize() + + +def main(bench=True): + headers = [ + "batch_size", + "seqlen", + "num_heads", + "head_size", + "num_heads_kv", + "name", + "latency (ms)", + "tflops", + ] + data = [] + for batch_size, seqlen, num_heads, head_size, num_heads_kv in [ + [1, 512, 32, 128, 8], + [1, 1024, 32, 128, 8], + [1, 2048, 32, 128, 8], + [1, 4096, 32, 128, 8], + #[1, 8192, 32, 128, 8], + [1, 512, 64, 128, 8], + [1, 1024, 64, 128, 8], + [1, 2048, 64, 128, 8], + [1, 4096, 64, 128, 8], + #[1, 8192, 64, 128, 8], + ]: + q = torch.rand( + batch_size, seqlen, num_heads, head_size, dtype=torch.float16 + ).cuda() + k = torch.rand( + batch_size, seqlen, num_heads_kv, head_size, dtype=torch.float16 + ).cuda() + v = torch.rand( + batch_size, seqlen, num_heads_kv, head_size, dtype=torch.float16 + ).cuda() + for name, runner in [ + ("flash-attn", flash_attention_flash_attn), + ("tilus", flash_attention), + ]: + print( + f"Running {name} with batch_size={batch_size}, seqlen={seqlen}, num_heads={num_heads}, head_size={head_size}, num_heads_kv={num_heads_kv}" + ) + try: + actual = runner(q, k, v) + except torch.OutOfMemoryError: + print("Out of memory, skipping this configuration.") + continue + + try: + expected = flash_attention_reference(q, k, v) + torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) + except torch.OutOfMemoryError: + pass + + latency = ( + benchmark_func( + lambda: runner(q, k, v), + warmup=20, + repeat=50, + ) + if bench + else float("nan") + ) + tflops = ( + 2 * batch_size * num_heads * seqlen * head_size * seqlen / latency * 1e-9 + ) + data.append( + [ + batch_size, + seqlen, + num_heads, + head_size, + num_heads_kv, + name, + latency, + tflops, + ] + ) + df = pd.DataFrame(data, columns=headers) + df_pivot = df.pivot( + index=[ + "batch_size", + "seqlen", + "num_heads", + "head_size", + "num_heads_kv", + ], + columns="name", + values=["latency (ms)", "tflops"], + ).reset_index() + # sort by (batch_size, num_heads, head_size, seqlen) + df_pivot = df_pivot.sort_values( + by=["batch_size", "num_heads", "head_size", "seqlen"], + ascending=[True, True, True, True], + ) + print(df_pivot) + + +if __name__ == "__main__": + main() + # ncu_run(main, bench=False, kernel_regex="flash_fwd|flash_attention") \ No newline at end of file diff --git a/examples/generic_python/flash_attention/triton_attention.py b/examples/generic_python/flash_attention/triton_attention.py new file mode 100644 index 000000000..e4f36a133 --- /dev/null +++ b/examples/generic_python/flash_attention/triton_attention.py @@ -0,0 +1,276 @@ +# example taken from https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html +# removed optional support for FP8 +# Removed backward pass +# removed checks for e.g. hopper/hip etc +import pytest +import torch +import os +import numpy as np + +import triton +import triton.language as tl +from triton.tools.tensor_descriptor import TensorDescriptor + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + +configs = [ + triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \ + for BM in [64, 128]\ + for BN in [32, 64, 128]\ + for s in [2, 3, 4] \ + for w in [4, 8]\ +] + + +def prune_invalid_configs(configs, named_args, **kwargs): + N_CTX = kwargs["N_CTX"] + STAGE = kwargs["STAGE"] + + # Filter out configs where BLOCK_M > N_CTX + # Filter out configs where BLOCK_M < BLOCK_N when causal is True + return [ + conf for conf in configs if conf.kwargs.get("BLOCK_M", 0) <= N_CTX and ( + conf.kwargs.get("BLOCK_M", 0) >= conf.kwargs.get("BLOCK_N", 0) or STAGE == 1) + ] + + +@triton.jit +def _attn_fwd_inner(acc, l_i, m_i, q, # + desc_k, desc_v, # + offset_y, dtype: tl.constexpr, start_m, qk_scale, # + BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # + STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, # + N_CTX: tl.constexpr, warp_specialize: tl.constexpr): + # range of values handled by this stage + if STAGE == 1: + lo, hi = 0, start_m * BLOCK_M + elif STAGE == 2: + lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M + lo = tl.multiple_of(lo, BLOCK_M) + # causal = False + else: + lo, hi = 0, N_CTX + offsetk_y = offset_y + lo + if dtype == tl.float8e5: + offsetv_y = offset_y * HEAD_DIM + lo + else: + offsetv_y = offset_y + lo + # loop over k, v and update accumulator + for start_n in tl.range(lo, hi, BLOCK_N, warp_specialize=warp_specialize): + start_n = tl.multiple_of(start_n, BLOCK_N) + # -- compute qk ---- + k = desc_k.load([offsetk_y, 0]).T + qk = tl.dot(q, k) + if STAGE == 2: + mask = offs_m[:, None] >= (start_n + offs_n[None, :]) + qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) + m_ij = tl.maximum(m_i, tl.max(qk, 1)) + qk -= m_ij[:, None] + else: + m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) + qk = qk * qk_scale - m_ij[:, None] + p = tl.math.exp2(qk) + # -- compute correction factor + alpha = tl.math.exp2(m_i - m_ij) + l_ij = tl.sum(p, 1) + # -- update output accumulator -- + if warp_specialize and BLOCK_M == 128 and HEAD_DIM == 128: + BM: tl.constexpr = acc.shape[0] + BN: tl.constexpr = acc.shape[1] + acc0, acc1 = acc.reshape([BM, 2, BN // 2]).permute(0, 2, 1).split() + acc0 = acc0 * alpha[:, None] + acc1 = acc1 * alpha[:, None] + acc = tl.join(acc0, acc1).permute(0, 2, 1).reshape([BM, BN]) + else: + acc = acc * alpha[:, None] + # prepare p and v for the dot + if dtype == tl.float8e5: + v = desc_v.load([0, offsetv_y]).T + else: + v = desc_v.load([offsetv_y, 0]) + p = p.to(dtype) + # note that this non transposed v for FP8 is only supported on Blackwell + acc = tl.dot(p, v, acc) + # update m_i and l_i + # place this at the end of the loop to reduce register pressure + l_i = l_i * alpha + l_ij + m_i = m_ij + offsetk_y += BLOCK_N + offsetv_y += BLOCK_N + return acc, l_i, m_i + + +@triton.jit +def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape): + if isinstance(desc_or_ptr, tl.tensor_descriptor): + return desc_or_ptr + else: + return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape) + + +@triton.autotune(configs=configs, key=["N_CTX", "HEAD_DIM", "warp_specialize"], + prune_configs_by={'early_config_prune': prune_invalid_configs}) +@triton.jit +def _attn_fwd(sm_scale, M, # + Z, H, desc_q, desc_k, desc_v, desc_o, N_CTX, # + HEAD_DIM: tl.constexpr, # + BLOCK_M: tl.constexpr, # + BLOCK_N: tl.constexpr, # + STAGE: tl.constexpr, # + warp_specialize: tl.constexpr, # + ): + dtype = tl.float16 + tl.static_assert(BLOCK_N <= HEAD_DIM) + start_m = tl.program_id(0) + off_hz = tl.program_id(1) + off_z = off_hz // H + off_h = off_hz % H + + y_dim = Z * H * N_CTX + desc_q = _maybe_make_tensor_desc(desc_q, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], + block_shape=[BLOCK_M, HEAD_DIM]) + desc_v = _maybe_make_tensor_desc(desc_v, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], + block_shape=[BLOCK_N, HEAD_DIM]) + desc_k = _maybe_make_tensor_desc(desc_k, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], + block_shape=[BLOCK_N, HEAD_DIM]) + desc_o = _maybe_make_tensor_desc(desc_o, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], + block_shape=[BLOCK_M, HEAD_DIM]) + + offset_y = off_z * (N_CTX * H) + off_h * N_CTX + qo_offset_y = offset_y + start_m * BLOCK_M + # initialize offsets + offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + # initialize pointer to m and l + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + # load scales + qk_scale = sm_scale + qk_scale *= 1.44269504 # 1/log(2) + # load q: it will stay in SRAM throughout + q = desc_q.load([qo_offset_y, 0]) + # stage 1: off-band + # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE + # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE + if STAGE & 1: + acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, # + desc_k, desc_v, # + offset_y, dtype, start_m, qk_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 4 - STAGE, offs_m, offs_n, N_CTX, # + warp_specialize) + # stage 2: on-band + if STAGE & 2: + acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, # + desc_k, desc_v, # + offset_y, dtype, start_m, qk_scale, # + BLOCK_M, HEAD_DIM, BLOCK_N, # + 2, offs_m, offs_n, N_CTX, # + warp_specialize) + # epilogue + m_i += tl.math.log2(l_i) + acc = acc / l_i[:, None] + m_ptrs = M + off_hz * N_CTX + offs_m + tl.store(m_ptrs, m_i) + desc_o.store([qo_offset_y, 0], acc.to(dtype)) + + + +def forward(q, k, v, causal, sm_scale, warp_specialize=True): + # shape constraints + HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] + # when v is in float8_e5m2 it is transposed. + HEAD_DIM_V = v.shape[-1] + assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V + assert HEAD_DIM_K in {16, 32, 64, 128, 256} + o = torch.empty_like(q) + stage = 3 if causal else 1 + extra_kern_args = {} + M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) + desc_q = q + desc_v = v + desc_k = k + desc_o = o + + def grid(META): + return (triton.cdiv(q.shape[2], META["BLOCK_M"]), q.shape[0] * q.shape[1], 1) + + _attn_fwd[grid]( + sm_scale, M, # + q.shape[0], q.shape[1], # + desc_q, desc_k, desc_v, desc_o, # + N_CTX=q.shape[2], # + HEAD_DIM=HEAD_DIM_K, # + STAGE=stage, # + warp_specialize=warp_specialize, # + **extra_kern_args) + + return o + + + + + +if __name__ == "__main__": + import torch + +# Triton kernel (your forward function) +# from your previous code: forward(ctx, Q, K, V, causal, sm_scale, warp_specialize=True) + +def flash_attention_reference(q, k, v, causal=False, sm_scale=1.0): + """ + Reference attention using batched matmul (like the pytest does). + Args: + q, k, v: [batch, heads, seq_len, head_dim] + causal: whether to apply causal mask + sm_scale: scaling factor for attention scores + Returns: + output: [batch, heads, seq_len, head_dim] + """ + # Batched matrix multiply + # scores: [batch, heads, seq_len, seq_len] + scores = torch.matmul(q, k.transpose(-2, -1)) * sm_scale + + if causal: + seq_len = q.shape[2] + mask = torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool)) + scores = scores.masked_fill(~mask[None, None, :, :], float("-inf")) + + # Softmax along last dimension + attention = torch.softmax(scores.float(), dim=-1) + + # Weighted sum with V + output = torch.matmul(attention, v) + return output + + + +# Example test +if __name__ == "__main__": + forward_simple() + exit(0) + + # Parameters + B, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 + causal = True + sm_scale = 1.0 + + # Random inputs + Q = torch.randn(B, H, N_CTX, HEAD_DIM, dtype=torch.float16, device="cuda") + K = torch.randn_like(Q) + V = torch.randn_like(Q) + + # Triton output + output_triton = forward(Q, K, V, causal=causal, sm_scale=sm_scale, warp_specialize=True) + + # Reference output + output_ref = flash_attention_reference(Q.float(), K.float(), V.float(), causal=causal, sm_scale=sm_scale) + + # Compare + max_diff = (output_triton.float() - output_ref).abs().max() + mean_diff = (output_triton.float() - output_ref).abs().mean() + print(f"Max difference: {max_diff.item():.6f}") + print(f"Mean difference: {mean_diff.item():.6f}") + + diff --git a/examples/generic_python/loopy_example.py b/examples/generic_python/loopy_example.py new file mode 100644 index 000000000..49ad131f6 --- /dev/null +++ b/examples/generic_python/loopy_example.py @@ -0,0 +1,38 @@ +import numpy as np + +import pyopencl as cl +import pyopencl.array + +import loopy as lp +from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_2 # noqa: F401 + + +# setup +# ----- +ctx = cl.create_some_context() +queue = cl.CommandQueue(ctx) + +n = 15 * 10**6 +a = cl.array.arange(queue, n, dtype=np.float32) + +# create +# ------ +knl = lp.make_kernel( + "{ [i]: 0<=i torch.Tensor: + m, k = x.size() + k, n = y.size() + out = torch.empty([m, n], dtype=x.dtype, device=x.device) + + for tile_m, tile_n in hl.tile([m, n]): + acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) + for tile_k in hl.tile(k): + acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n]) + out[tile_m, tile_n] = acc + + return out + + +out = matmul(torch.randn([2048, 2048], device="cuda"), + torch.randn([2048, 2048], device="cuda")) \ No newline at end of file diff --git a/examples/generic_python/matmul/numba_matmul.py b/examples/generic_python/matmul/numba_matmul.py new file mode 100644 index 000000000..00102a382 --- /dev/null +++ b/examples/generic_python/matmul/numba_matmul.py @@ -0,0 +1,148 @@ +import torch +import numpy as np +from numba import cuda +from kernel_tuner import tune_kernel +from pathlib import Path + +FULL_PATH = Path(__file__).resolve() + +# Example taken from https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication +@cuda.jit(cache=True) +def matmul(A, B, C): + """Perform square matrix multiplication of C = A * B + """ + i, j = cuda.grid(2) + if i < C.shape[0] and j < C.shape[1]: + tmp = 0. + for k in range(A.shape[1]): + tmp += A[i, k] * B[k, j] + C[i, j] = tmp + + +# Example taken from https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication +# Changed data type from float32 to float16 +@cuda.jit(cache=True) +def fast_matmul(A, B, C): + # Define an array in the shared memory + # The size and type of the arrays must be known at compile time + TPB = 16 # TEMP voor overhead testing + sA = cuda.shared.array(shape=(TPB, TPB), dtype=np.float16) + sB = cuda.shared.array(shape=(TPB, TPB), dtype=np.float16) + + x, y = cuda.grid(2) + + tx = cuda.threadIdx.x + ty = cuda.threadIdx.y + bpg = cuda.gridDim.x # blocks per grid + + if x >= C.shape[0] and y >= C.shape[1]: + # Quit if (x, y) is outside of valid C boundary + return + + # Each thread computes one element in the result matrix. + # The dot product is chunked into dot products of TPB-long vectors. + tmp = 0. + for i in range(bpg): + # Preload data into shared memory + sA[tx, ty] = A[x, ty + i * TPB] + sB[tx, ty] = B[tx + i * TPB, y] + + # Wait until all threads finish preloading + cuda.syncthreads() + + # Computes partial product on the shared memory + for j in range(TPB): + tmp += sA[tx, j] * sB[j, ty] + + # Wait until all threads finish computing + cuda.syncthreads() + + C[x, y] = tmp + + +def run_matmul(M, N, K): + + # create numpy arrays + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) + + # copy to GPU + A_d = cuda.to_device(A) + B_d = cuda.to_device(B) + C_d = cuda.to_device(C) + + # threads per block + threads = (16, 16) + + # compute grid size (ceil division) + blocks = ( + (M + threads[0] - 1) // threads[0], + (N + threads[1] - 1) // threads[1], + ) + + # launch kernel + fast_matmul[blocks, threads](A_d, B_d, C_d) + + # copy result back + C_result = C_d.copy_to_host() + + # check + np.testing.assert_allclose(C_result, A @ B, rtol=1e-2) + + print("Correct!") + + +def call_numba(kernel_function, args, kwargs, grid, threads): + numba_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + numba_args.append(cuda.as_cuda_array(arg)) + else: + numba_args.append(arg) + kernel_function[grid, threads](*args, **kwargs) + + +def tune(M, N, K): + # create numpy arrays + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) + + size = (M, N) + args = [A, B, C] + tune_params = dict() + tune_params["block_size_x"] = [4, 8, 16, 32, 64, 128, 256] + tune_params["block_size_y"] = [4, 8, 16, 32, 64, 128, 256] + tune_params["TPB"] = [4, 8, 16, 32, 64, 128, 256] + restrictions = ["block_size_x == block_size_y", "block_size_x == TPB"] + + answer = [None, None, A @ B] + atol = M * 2**(-11) + + results, env = tune_kernel( + kernel_name="fast_matmul", + kernel_source=FULL_PATH, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=answer, + atol=atol, + call_function=call_numba, + restrictions=restrictions + ) + + +if __name__ == "__main__": + run_matmul(128, 96, 64) + #tune(1024, 1024, 1024) + + ''' + 128: 16, 8 + 256: 8, 16 + 512: 8, 16 + 1024: 16, 8 + 4096: 16, 8 + 8192: 16, 8 + ''' \ No newline at end of file diff --git a/examples/generic_python/matmul/test.py b/examples/generic_python/matmul/test.py new file mode 100644 index 000000000..991ef2cc4 --- /dev/null +++ b/examples/generic_python/matmul/test.py @@ -0,0 +1,278 @@ +import argparse +import itertools +import tilelang as tl +import tilelang.language as T +from tilelang.autotuner import AutoTuner +from tilelang.carver.template import MatmulTemplate +from tilelang.carver.arch import CUDA +from tilelang.carver.arch import CDNA +from tilelang.carver.roller.rasterization import NoRasterization +import torch + + +def ref_program(A, B): + """ + Compute the matrix product of A and the transpose of B. + + A and B are expected to be 2-D tensors where A has shape (M, K) and B has shape (N, K). The result is a tensor with shape (M, N) equal to A @ B.T, using the inputs' dtypes. + """ + return A @ B.T + + +def get_configs(M, N, K, with_roller=False, topk=20): + """ + Generate a list of kernel tuning configuration dictionaries for a tiled matrix-multiply. + + When with_roller is True this queries the MatmulTemplate roller to produce up to `topk` recommended + configurations (device-specific TensorCore-friendly tilings). Each returned dict contains: + - block_M, block_N, block_K: tile sizes + - num_stages: pipeline staging (0 means no explicit staging) + - thread_num: total threads used for the block + - enable_rasteration: whether a rasterization/swizzle layout was recommended (note spelling) + + When with_roller is False this returns the Cartesian product of a fixed set of candidate + parameters; the returned dicts use the backward-compatible key name "enable_rasteration" for that flag. + + Parameters: + M, N, K (int): GEMM dimensions used to generate valid tile sizes. + with_roller (bool): If True, use MatmulTemplate's roller to generate device-aware hints; + otherwise use a predefined candidate grid. + topk (int): Maximum number of roller hints to request when with_roller is True. + + Returns: + List[dict]: A list of configuration dictionaries as described above. + + Raises: + ValueError: if with_roller is True but the roller returns no hints. + """ + if with_roller: + arch = CUDA("cuda") if torch.version.hip is None else CDNA("hip") + carve_template = MatmulTemplate( + M=M, + N=N, + K=K, + in_dtype=T.float16, + out_dtype=T.float16, + accum_dtype=T.float32, + ).with_arch(arch) + + func = carve_template.equivalent_function() + assert func is not None, "Function is None" + roller_hints = carve_template.recommend_hints(topk=topk) + if roller_hints is None: + raise ValueError("No Roller Hints Found for TensorCore Scheduling") + configs = [] + for hint in roller_hints: + config = {} + block_m, block_n = hint.block + warp_m, warp_n = hint.warp + # block_rows, block_cols represents warp partitioning + block_rows, block_cols = block_m // warp_m, block_n // warp_n + config["block_M"] = block_m + config["block_N"] = block_n + config["block_K"] = hint.rstep[0] + config["num_stages"] = hint.pipeline_stage if hint.pipeline_stage > 1 else 0 + config["thread_num"] = block_rows * block_cols * 32 + config["enable_rasteration"] = hint.rasterization_plan is not NoRasterization + configs.append(config) + else: + block_M = [64, 128, 256] + block_N = [64, 128, 256] + block_K = [32, 64] + num_stages = [0, 1, 2, 3] + thread_num = [128, 256] + enable_rasterization = [True, False] + _configs = list( + itertools.product( + block_M, + block_N, + block_K, + num_stages, + thread_num, + enable_rasterization, + ) + ) + + configs = [ + { + "block_M": c[0], + "block_N": c[1], + "block_K": c[2], + "num_stages": c[3], + "thread_num": c[4], + "enable_rasteration": c[5], # keep param name for backward-compat + } + for c in _configs + ] + return configs + + +def get_best_config( + M, + N, + K, + with_roller: bool = False, + profile_backend: str = "event", +): + def kernel( + block_M=None, + block_N=None, + block_K=None, + num_stages=None, + thread_num=None, + enable_rasteration=None, + ): + dtype = T.bfloat16 + accum_dtype = T.float32 + + @T.prim_func + def main( + A: T.Tensor((M, K), dtype), + B: T.Tensor((N, K), dtype), + C: T.Tensor((M, N), dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_N, block_K), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_shared = T.alloc_shared((block_M, block_N), dtype) + T.use_swizzle(panel_size=10, enable=enable_rasteration) + T.clear(C_local) + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[bx * block_N, k * block_K], B_shared) + T.gemm( + A_shared, + B_shared, + C_local, + transpose_B=True, + ) + T.copy(C_local, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return main + + autotuner = ( + AutoTuner.from_kernel(kernel=kernel, configs=get_configs(M, N, K, with_roller)) + .set_compile_args( + out_idx=[-1], + target="auto", + ) + .set_profile_args( + supply_type=tl.TensorSupplyType.Integer, + ref_prog=ref_program, + skip_check=False, + backend=profile_backend, + ) + ) + return autotuner.run(warmup=3, rep=20) + + +def get_heuristic_config() -> dict: + # Get CUDA device properties + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + sm_major, sm_minor = torch.cuda.get_device_capability(device) + sm_version = sm_major * 10 + sm_minor + print(f"CUDA device capability: {sm_version}") + if sm_version in {80}: + return {"block_M": 128, "block_N": 256, "block_K": 32, "num_stages": 2, "thread_num": 128, "enable_rasteration": True} + elif sm_version in {90}: + return {"block_M": 128, "block_N": 256, "block_K": 64, "num_stages": 3, "thread_num": 256, "enable_rasteration": True} + else: + return {"block_M": 128, "block_N": 256, "block_K": 32, "num_stages": 0, "thread_num": 128, "enable_rasteration": True} + + +@tl.jit(out_idx=[-1]) +def matmul(M, N, K, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): + @T.prim_func + def gemm_autotune( + A: T.Tensor((M, K), dtype), + B: T.Tensor((N, K), dtype), + C: T.Tensor((M, N), dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_N, block_K), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_shared = T.alloc_shared((block_M, block_N), dtype) + T.use_swizzle(panel_size=10, enable=enable_rasteration) + T.clear(C_local) + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[bx * block_N, k * block_K], B_shared) + T.gemm( + A_shared, + B_shared, + C_local, + transpose_B=True, + ) + T.copy(C_local, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return gemm_autotune + + +def main( + M: int = 4096, + N: int = 4096, + K: int = 4096, + use_autotune: bool = False, + with_roller: bool = False, + profile_backend: str = "event", +): + if use_autotune: + result = get_best_config( + M, + N, + K, + with_roller=with_roller, + profile_backend=profile_backend, + ) + print(result.config) + kernel = result.kernel + else: + config = get_heuristic_config() + kernel = matmul(M, N, K, **config) + + # benchmark + profiler = kernel.get_profiler(tensor_supply_type=tl.TensorSupplyType.Auto) + tilelang_latency = profiler.do_bench( + backend=profile_backend, + ) + ref_latency = profiler.do_bench( + ref_program, + backend=profile_backend, + ) + profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2) + print(f"TileLang latency: {tilelang_latency}") + print(f"Ref latency: {ref_latency}") + print(f"TileLang TFlops: {2 * M * N * K / tilelang_latency * 1e-9}") + print(f"Ref TFlops: {2 * M * N * K / ref_latency * 1e-9}") + + +def run_regression_perf(M: int = 4096, N: int = 4096, K: int = 4096): + config = get_heuristic_config() + kernel = matmul(M, N, K, **config) + profiler = kernel.get_profiler(tensor_supply_type=tl.TensorSupplyType.Auto) + return profiler.do_bench(backend="cupti") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Autotuned MatMul Benchmark") + parser.add_argument("--m", type=int, default=4096, help="Matrix dimension M") + parser.add_argument("--n", type=int, default=4096, help="Matrix dimension N") + parser.add_argument("--k", type=int, default=4096, help="Matrix dimension K") + parser.add_argument("--use_autotune", action="store_true", default=False, help="Whether to use autotune for matmul configs") + parser.add_argument("--with_roller", action="store_true", default=False, help="Whether to enable BitBLAS roller for search space") + parser.add_argument("--profile_backend", type=str, default="event", help="Profiler backend") + args = parser.parse_args() + main( + args.m, + args.n, + args.k, + args.use_autotune, + args.with_roller, + args.profile_backend, + ) \ No newline at end of file diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py index 895732629..7ced8cc83 100644 --- a/examples/generic_python/matmul/tilelang_matmul.py +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -32,40 +32,42 @@ def gemm( # https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py -@tilelang.jit +# changed gemm_autotune to gemm +# originally, B was transposed. I changed this so the kernel input is the same as in other languages. +#@tilelang.jit def matmul_opt(M, N, K, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): @T.prim_func - def gemm_autotune( + def gemm( A: T.Tensor((M, K), dtype), - B: T.Tensor((N, K), dtype), + B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_N, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) C_local = T.alloc_fragment((block_M, block_N), accum_dtype) C_shared = T.alloc_shared((block_M, block_N), dtype) T.use_swizzle(panel_size=10, enable=enable_rasteration) T.clear(C_local) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[bx * block_N, k * block_K], B_shared) + T.copy(B[k * block_K, bx * block_N], B_shared) T.gemm( A_shared, B_shared, C_local, - transpose_B=True, + transpose_B=False, ) T.copy(C_local, C_shared) T.copy(C_shared, C[by * block_M, bx * block_N]) - return gemm_autotune + return gemm def main(): - kernel = matmul_basic(1024, 1024, 1024, 128, 128, 32) + kernel = matmul_opt(1024, 1024, 1024, 128, 128, 32, 3, 128, False) import torch diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py index dfe294761..d893b9490 100644 --- a/examples/generic_python/matmul/tilus_matmul.py +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -68,8 +68,9 @@ def __call__( + # This kernel is copied from the Tilus project: -# https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_v5.py +# https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_+v5.py # # Original example: matmul_v5.py # Copyright (c) the Tilus authors diff --git a/examples/generic_python/matmul/warp_matmul.py b/examples/generic_python/matmul/warp_matmul.py new file mode 100644 index 000000000..9e6b7dcf3 --- /dev/null +++ b/examples/generic_python/matmul/warp_matmul.py @@ -0,0 +1,131 @@ +import torch +import numpy as np +import warp as wp + +from kernel_tuner import tune_kernel +from pathlib import Path + +wp.init() + +FULL_PATH = Path(__file__).resolve() + +# tile size +TILE_M = wp.constant(8) +TILE_N = wp.constant(4) +TILE_K = wp.constant(8) + +# num threads per-tile +TILE_THREADS = 64 + +# GEMM example from https://nvidia.github.io/warp/user_guide/tiles.html +@wp.kernel +def tile_gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)): + + # output tile index + i, j = wp.tid() + + sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32) + + M = A.shape[0] + N = B.shape[1] + K = A.shape[1] + + count = (K + TILE_K - 1) // TILE_K + + for k in range(0, count): + a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i*TILE_M, k*TILE_K), bounds_check=True) + b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k*TILE_K, j*TILE_N), bounds_check=True) + + # sum += a*b + wp.tile_matmul(a, b, sum) + + wp.tile_store(C, sum, offset=(i*TILE_M, j*TILE_N), bounds_check=True) + + +def run_kernel_direct(M, N, K): + rng = np.random.default_rng(42) + A = rng.random((M, K), dtype=np.float32) + B = rng.random((K, N), dtype=np.float32) + C = np.zeros((M, N), dtype=np.float32) + + A_wp = wp.array(A) + B_wp = wp.array(B) + C_wp = wp.array(C) + + with wp.Tape() as tape: + wp.launch_tiled( + tile_gemm, + dim=((M + TILE_M - 1) // TILE_M, (N + TILE_N - 1) // TILE_N), + inputs=[A_wp, B_wp, C_wp], + block_dim=TILE_THREADS) + + np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-3) + + print("Example matrix multiplication passed") + + + +def call_warp(kernel_function, args, kwargs, grid, threads, params): + # Convert Torch tensors to Warp args + warp_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + warp_args.append(wp.from_torch(arg)) + else: + warp_args.append(arg) + + # launch kernel + with wp.Tape() as tape: + wp.launch_tiled( + kernel_function, + dim=grid, + inputs=warp_args, + block_dim=params["TILE_THREADS"], # We could directly take threads, but in the given example this is a constant + ) + + +def tune(M, K, N): + rng = np.random.default_rng(42) + A = rng.random((M, K), dtype=np.float32) + B = rng.random((K, N), dtype=np.float32) + C = np.zeros((M, N), dtype=np.float32) + + size = (M, N) + block_size_names = ["TILE_M", "TILE_N"] + + tune_params = dict() + tune_params["TILE_M"] = [4, 8, 16] + tune_params["TILE_N"] = [2, 4, 8] + tune_params["TILE_K"] = [4, 8, 16] + tune_params["TILE_THREADS"] = [32, 64, 128] + + args = [A, B, C] + answer = [None, None, A @ B] + + results, env = tune_kernel( + kernel_name="tile_gemm", + kernel_source=FULL_PATH, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=answer, + call_function=call_warp, + block_size_names=block_size_names, + ) + + +if __name__ == "__main__": + #tune(128, 128, 128) + + sizes = [ + (65, 65, 17), + (67, 71, 19), + (1, 1, 1), + (63, 63, 15), + (129, 130, 33), + ] + + for size in sizes: + print(size) + run_kernel_direct(*size) diff --git a/examples/generic_python/normalization/tilelang_norm.py b/examples/generic_python/normalization/tilelang_norm.py new file mode 100644 index 000000000..7c26d8f5f --- /dev/null +++ b/examples/generic_python/normalization/tilelang_norm.py @@ -0,0 +1,78 @@ +# taken from https://github.com/tile-ai/tilelang/blob/main/examples/norm/rms_norm.py +# TODO not comparable +import torch +import tilelang +import tilelang.language as T + + +def rms_norm_splitk(M, N, blk_m, blk_k): + dtype = T.float + + @T.prim_func + def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)): + with T.Kernel(T.ceildiv(M, blk_m), threads=128) as bx: + A_shared = T.alloc_shared((blk_m, blk_k), dtype) + A_local = T.alloc_fragment((blk_m, blk_k), dtype) + A_powsum = T.alloc_fragment((blk_m,), dtype) + + num_k_step = T.ceildiv(N, blk_k) + T.clear(A_local) + for k in range(num_k_step): + T.copy(A[bx * blk_m, k * blk_k], A_shared) + for i, j in T.Parallel(blk_m, blk_k): + A_local[i, j] += A_shared[i, j] * A_shared[i, j] + T.reduce_sum(A_local, A_powsum, dim=1) + for i in T.Parallel(blk_m): + A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12) + + for k in range(num_k_step): + # reverse, better cache hit rate + T.copy(A[bx * blk_m, (num_k_step - 1 - k) * blk_k], A_shared) + for i, j in T.Parallel(blk_m, blk_k): + A_shared[i, j] *= A_powsum[i] + T.copy(A_shared, B[bx * blk_m, (num_k_step - 1 - k) * blk_k]) + + return main + + +@tilelang.jit(out_idx=[-1], pass_configs={"tl.disable_tma_lower": True}) +def rms_norm(M, N, blk_m): + dtype = T.float + + @T.prim_func + def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)): + with T.Kernel(T.ceildiv(M, blk_m), threads=128) as bx: + A_shared = T.alloc_shared((blk_m, N), dtype) + A_pow_local = T.alloc_fragment((blk_m, N), dtype) + A_local = T.alloc_fragment((blk_m, N), dtype) + A_powsum = T.alloc_fragment((blk_m,), dtype) + + T.copy(A[bx * blk_m : (bx + 1) * blk_m, :], A_shared) + T.copy(A_shared, A_local) + for i, j in T.Parallel(blk_m, N): + A_pow_local[i, j] = A_local[i, j] * A_local[i, j] + T.reduce_sum(A_pow_local, A_powsum, dim=1) + for i in T.Parallel(blk_m): + A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12) + for i, j in T.Parallel(blk_m, N): + A_local[i, j] *= A_powsum[i] + T.copy(A_local, B[bx * blk_m : (bx + 1) * blk_m, :]) + + return main + + +def ref_program(x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-12) + + +if __name__ == "__main__": + M, N, blk_m, blk_k = 8192, 8192, 1, 512 + kernel = rms_norm(M, N, blk_m) + profiler = kernel.get_profiler() + profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) + print("All checks pass.") + + latency = profiler.do_bench(ref_program, warmup=500) + print("Ref: {:.2f} ms".format(latency)) + latency = profiler.do_bench(warmup=500) + print("Tile-lang: {:.2f} ms".format(latency)) \ No newline at end of file diff --git a/examples/generic_python/normalization/tilus_norm.py b/examples/generic_python/normalization/tilus_norm.py new file mode 100644 index 000000000..cfe569e80 --- /dev/null +++ b/examples/generic_python/normalization/tilus_norm.py @@ -0,0 +1,132 @@ +# taken from https://github.com/NVIDIA/tilus/blob/main/examples/norm/layer_norm.py +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +import pandas +import tilus +import torch +from tilus import float16, float32, int32 +from tilus.utils import benchmark_func, cdiv + + +@tilus.autotune("block_m", [1, 8]) +@tilus.autotune("block_n", [128, 256, 512, 1024]) +@tilus.autotune("warps", [2, 4, 8]) +class LayerNorm(tilus.Script): + """Forward-only layer normalization tilus kernel. + + This implements the per-row LayerNorm used in many transformer blocks. It + computes: y = (x - mean) / sqrt(var + eps) * gamma + beta + + Only the forward is provided. + """ + + def __init__(self, block_m: int, block_n: int, warps: int): + super().__init__() + self.block_m: int = block_m + self.block_n: int = block_n + self.warps: int = warps + + def __call__( + self, + m_size: int, + n_size: int32, + x_ptr: ~float16, + gamma_ptr: ~float16, + beta_ptr: ~float16, + y_ptr: ~float16, + eps: float, + ): + self.attrs.blocks = (cdiv(m_size, self.block_m),) + self.attrs.warps = self.warps + + offset_m = self.blockIdx.x * self.block_m + + g_x = self.global_view(x_ptr, dtype=float16, shape=[m_size, n_size]) + g_y = self.global_view(y_ptr, dtype=float16, shape=[m_size, n_size]) + g_gamma = self.global_view(gamma_ptr, dtype=float16, shape=[n_size]) + g_beta = self.global_view(beta_ptr, dtype=float16, shape=[n_size]) + + # Register accumulators for mean and variance (computed in float32) + r_sum = self.register_tensor( + dtype=float32, shape=[self.block_m, self.block_n], init=0.0 + ) + r_square = self.register_tensor( + dtype=float32, shape=[self.block_m, self.block_n], init=0.0 + ) + + # first pass: compute mean and variance + for offset_n in range(0, n_size, self.block_n): + r_x = self.load_global( + g_x, offsets=[offset_m, offset_n], shape=[self.block_m, self.block_n] + ).to(float32) # [block_m, block_n] + r_sum = r_sum + r_x + r_square = r_square + self.square(r_x) + + # finalize mean and variance + r_mean = self.sum(r_sum, dim=1, keepdim=True) / n_size # [block_m, 1] + r_var = ( + self.sum(r_square, dim=1, keepdim=True) / n_size - r_mean * r_mean + ) # [block_m, 1], var = E[x^2] - (E[x])^2 + r_rstd = self.rsqrt(r_var + eps) + + # second pass: y = (x - mean) * rstd * gamma + beta + for offset_n in range(0, n_size, self.block_n): + r_x = self.load_global( + g_x, offsets=[offset_m, offset_n], shape=[self.block_m, self.block_n] + ).to(float32) # [block_m, block_n] + r_gamma = self.load_global( + g_gamma, offsets=[offset_n], shape=[self.block_n] + ).to(float32) # [block_n] + r_beta = self.load_global( + g_beta, offsets=[offset_n], shape=[self.block_n] + ).to(float32) # [block_n] + r_x_hat = (r_x - r_mean) * r_rstd + r_y = r_x_hat * r_gamma + r_beta + self.store_global(g_y, r_y.to(float16), offsets=[offset_m, offset_n]) + + +def main(): + headers = ["m_size", "n_size", "dtype", "torch (ms)", "tilus (ms)"] + rows = [] + for i in [1, 2, 4, 8]: + m_size = n_size = 1024 * i + + tilus_layer_norm = LayerNorm() + + x = (torch.rand(m_size, n_size, dtype=torch.float16).cuda() - 0.5) * 2.0 + gamma = torch.rand(n_size, dtype=torch.float16).cuda() + beta = torch.rand(n_size, dtype=torch.float16).cuda() + y_actual = torch.empty_like(x) + + tilus_layer_norm(m_size, n_size, x, gamma, beta, y_actual, 1e-5) + y_expected = torch.nn.functional.layer_norm( + x, normalized_shape=[n_size], weight=gamma, bias=beta, eps=1e-5 + ) + + torch.testing.assert_close(y_actual, y_expected, atol=1e-2, rtol=1e-2) + + rows.append( + [ + m_size, + n_size, + "float16", + benchmark_func( + lambda: torch.nn.functional.layer_norm( + x, normalized_shape=[n_size], weight=gamma, bias=beta, eps=1e-5 + ) + ), + benchmark_func( + lambda: tilus_layer_norm( + m_size, n_size, x, gamma, beta, y_actual, 1e-5 + ) + ), + ] + ) + print(f"LayerNorm forward matches reference for size ({m_size}, {n_size})") + + df = pandas.DataFrame(rows, columns=headers) + print(df) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/generic_python/normalization/triton_norm.py b/examples/generic_python/normalization/triton_norm.py new file mode 100644 index 000000000..365f541c4 --- /dev/null +++ b/examples/generic_python/normalization/triton_norm.py @@ -0,0 +1,66 @@ + +#Example taken from https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +import torch + +import triton +import triton.language as tl + +try: + # This is https://github.com/NVIDIA/apex, NOT the apex on PyPi, so it + # should not be added to extras_require in setup.py. + import apex + HAS_APEX = True +except ModuleNotFoundError: + HAS_APEX = False + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def _layer_norm_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + stride, # how much to increase the pointer when moving by 1 row + N, # number of columns in X + eps, # epsilon to avoid division by zero + BLOCK_SIZE: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + row = tl.program_id(0) + Y += row * stride + X += row * stride + # Compute mean + mean = 0 + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=0) / N + # Compute variance + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) + x = tl.where(cols < N, x - mean, 0.) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + # Write mean / rstd + tl.store(Mean + row, mean) + tl.store(Rstd + row, rstd) + # Normalize and apply linear transformation + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask) + b = tl.load(B + cols, mask=mask) + x = tl.load(X + cols, mask=mask, other=0.).to(tl.float32) + x_hat = (x - mean) * rstd + y = x_hat * w + b + # Write output + tl.store(Y + cols, y, mask=mask) \ No newline at end of file diff --git a/examples/generic_python/pallas_vec_add.py b/examples/generic_python/pallas_vec_add.py new file mode 100644 index 000000000..874392dd9 --- /dev/null +++ b/examples/generic_python/pallas_vec_add.py @@ -0,0 +1,23 @@ +# Could not get working + +from functools import partial + +import jax +from jax.experimental import pallas as pl +import jax.numpy as jnp +import numpy as np + +def add_vectors_kernel(x_ref, y_ref, o_ref): + x, y = x_ref[...], y_ref[...] + o_ref[...] = x + y + + +@jax.jit +def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array: + return pl.pallas_call( + add_vectors_kernel, + out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype) + )(x, y) + + +add_vectors(jnp.arange(8), jnp.arange(8)) \ No newline at end of file diff --git a/examples/generic_python/tilus_vec_add.py b/examples/generic_python/tilus_vec_add.py index 9f842637c..d00c77972 100644 --- a/examples/generic_python/tilus_vec_add.py +++ b/examples/generic_python/tilus_vec_add.py @@ -8,7 +8,7 @@ FULL_PATH = Path(__file__).resolve() class VecAddV(tilus.Script): - def __init__(self, block_size_x=None): + def __init__(self, block_size_x=None, num_warps=None): super().__init__() self.block_size_x = block_size_x # number of threads per block @@ -22,7 +22,7 @@ def __call__( # compute the number of blocks needed self.attrs.blocks = [cdiv(n_size, self.block_size_x)] - self.attrs.warps = 1 # number of warps per block + self.attrs.warps = 4 # number of warps per block # calculate the offset for this block offset: int32 = self.block_size_x * self.blockIdx.x @@ -52,6 +52,7 @@ def tune(size): args = [size, a, b, c] tune_params = dict() tune_params["block_size_x"] = [32, 64, 128, 256, 512, 1024] + tune_params["num_warps"] = [4, 8] results, env = tune_kernel( @@ -100,7 +101,7 @@ def tune_with_builtin(size): c = torch.empty(size, dtype=torch.float32).cuda() c_expect = a + b - vecadd(size, a, b, c) + vecadd(size, a, b, c) # This is where the actual tuning takes place torch.cuda.synchronize() torch.testing.assert_close(c_expect, c) diff --git a/examples/generic_python/triton_vec_add.py b/examples/generic_python/triton_vec_add.py index f69ae96d7..00bd9dfbc 100644 --- a/examples/generic_python/triton_vec_add.py +++ b/examples/generic_python/triton_vec_add.py @@ -65,6 +65,8 @@ def tune(): call_function=call_triton, ) + print(results) + if __name__ == "__main__": diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 39be59a23..4845fc4e2 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -195,9 +195,7 @@ def run_kernel(self, func, gpu_args, threads, grid, params=None): configuration :type params: dict """ - with torch.cuda.stream(self.stream): - logging.debug("Running Generic Python kernel") - self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) + self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) def synchronize(self): @@ -239,7 +237,6 @@ def refresh_memory(self, gpu_memory, host_arguments, should_sync): """Refresh the GPU memory with the untouched host arguments. We overwrite the standard function because Python DSLs do usually do not manage memory explicitely""" - for i, host_arg in enumerate(host_arguments): if should_sync[i]: gpu_arg = gpu_memory[i] diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 1f130afa3..942ff329a 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -269,6 +269,8 @@ def benchmark_default(self, func, gpu_args, threads, grid, result, params=None): for obs in self.benchmark_observers: obs.after_finish() + #time.sleep(0.1) # prevent termal throttling + for obs in self.benchmark_observers: result.update(obs.get_results()) @@ -509,6 +511,10 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, # clean up any temporary files, if no error occurred instance.delete_temp_files() + # For Python DSLs, the compilation time also includes one kernel run, so we subtract the runtime + if self.lang == Language.GENERIC_PYTHON: + last_compilation_time -= result["time"] or 0 + result["compile_time"] = last_compilation_time or 0 result["verification_time"] = last_verification_time or 0 result["benchmark_time"] = last_benchmark_time or 0 diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 3eb27d86b..f806d5b87 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -70,7 +70,13 @@ def prepare_kernel_instance(self, kernel_options, params, grid, threads): generate a kernel instance with these tuning parameters inserted. kernel_options, grid and threads are not needed for Python kernels. ''' - new_kernel_fn, temp_file_path = self.apply_params_to_source_fn(params) + # Test: see if removing signature params helps in Triton overhead + filtered_params = {} + for k, v in params.items(): + if k not in self.signature: + filtered_params[k] = v + + new_kernel_fn, temp_file_path = self.apply_params_to_source_fn(filtered_params) self.kernel_fn = new_kernel_fn return PreparedKernelSourceData( @@ -121,7 +127,7 @@ def apply_params_to_source_fn(self, params): # Fix locations and generate source ast.fix_missing_locations(new_module) new_source = astor.to_source(new_module) - + #print(new_source) # Create a unique module name and write new source to it. From 775e8f565dbd35666db200312b450ea0b1852780 Mon Sep 17 00:00:00 2001 From: "I.C. van Ooijen" Date: Sun, 12 Apr 2026 12:01:32 +0200 Subject: [PATCH 09/44] matmul examples --- examples/generic_python/call_functions.py | 78 +- examples/generic_python/cute_vec_add.py | 92 ++ examples/generic_python/matmul/cute_matmul.py | 1019 +++++++++++++++-- .../generic_python/matmul/numba_matmul.py | 133 ++- .../generic_python/matmul/taichi_matmul.py | 18 +- .../generic_python/matmul/tilelang_matmul.py | 107 +- examples/generic_python/matmul/warp_matmul.py | 29 +- 7 files changed, 1370 insertions(+), 106 deletions(-) create mode 100644 examples/generic_python/cute_vec_add.py diff --git a/examples/generic_python/call_functions.py b/examples/generic_python/call_functions.py index 6b321c4a2..734568efd 100644 --- a/examples/generic_python/call_functions.py +++ b/examples/generic_python/call_functions.py @@ -3,20 +3,21 @@ def call_tilus(kernel_function, args, kwargs): kernel_function(*args, **kwargs) + def call_triton(kernel_function, args, kwargs, grid, threads, params): if "num_warps" in params.keys(): kwargs["num_warps"] = params["num_warps"] if "num_stages" in params.keys(): kwargs["num_stages"] = params["num_stages"] - torch.cuda.nvtx.range_push("kt call") kernel_function[grid](*args, **kwargs) - torch.cuda.nvtx.range_pop() + def call_tilelang(kernel_function, args, kwargs): compiled_kernel = kernel_function(**kwargs) compiled_kernel(*args) + def call_numba(kernel_function, args, kwargs, grid, threads): from numba import cuda @@ -26,4 +27,75 @@ def call_numba(kernel_function, args, kwargs, grid, threads): numba_args.append(cuda.as_cuda_array(arg)) else: numba_args.append(arg) - kernel_function[grid, threads](*args, **kwargs) \ No newline at end of file + + kernel_function[grid, threads](*args, **kwargs) + + +def call_cupyx(kernel_function, args, kwargs, grid, threads): + import cupy as cp + + cupy_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + cupy_args.append(cp.from_dlpack(arg)) + else: + cupy_args.append(arg) + kernel_function(grid, threads, tuple(cupy_args)) + + +def call_cute(kernel_function, args, kwargs, grid, threads, params): + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + # Initialize cache if it does not exist + if not hasattr(call_cute, "custom_cache"): + call_cute.custom_cache = {} + + # Convert Torch tensors to CuTe tensors with correct layout + cute_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + arg_ = from_dlpack(arg) + cute_args.append(arg_) + else: + cute_args.append(arg) + + # Form cache key from tuning parameters + param_keys = sorted(params.keys()) + cache_str = type(kernel_function).__name__ + for k in param_keys: + cache_str += "_" + str(params[k]) + + # Check if kernel exists in cache. Otherwise, compile and save + if cache_str in call_cute.custom_cache: + compiled_kernel = call_cute.custom_cache[cache_str] + else: + compiled_kernel = cute.compile(kernel_function, *cute_args) + call_cute.custom_cache[cache_str] = compiled_kernel + + compiled_kernel(*cute_args, **kwargs) + + +def call_taichi(kernel_function, args, kwargs): + kernel_function(*args, **kwargs) + + +def call_warp(kernel_function, args, kwargs, grid, threads, params): + import warp as wp + + # Convert Torch tensors to Warp args + warp_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + warp_args.append(wp.from_torch(arg)) + else: + warp_args.append(arg) + + # launch kernel + with wp.Tape() as tape: + wp.launch_tiled( + kernel_function, + dim=grid, + inputs=warp_args, + block_dim=params["TILE_THREADS"], # We could directly take threads, but in the given example this is a constant + ) diff --git a/examples/generic_python/cute_vec_add.py b/examples/generic_python/cute_vec_add.py new file mode 100644 index 000000000..a83694ab9 --- /dev/null +++ b/examples/generic_python/cute_vec_add.py @@ -0,0 +1,92 @@ +import torch +from functools import partial +from typing import List +import time + +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from kernel_tuner import tune_kernel + +@cute.kernel +def vec_add_kernel( + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + size: cute.Int32, +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + bdim, _, _ = cute.arch.block_dim() + thread_id = bdim * bidx + tidx + + if thread_id < size: + gC[thread_id] = gA[thread_id] + gB[thread_id] + + + +@cute.jit +def vec_add( + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + size: cute.Int32, +): + num_threads_per_block = 256 + + kernel = vec_add_kernel(mA, mB, mC, size) + + kernel.launch( + grid=(cute.ceil_div(size, num_threads_per_block), 1, 1), + block = (num_threads_per_block, 1, 1), + ) + + + + +def call_cute(kernel_function, args, kwargs, grid, threads, params): + cute_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + arg_ = from_dlpack(arg) + cute_args.append(arg_) + else: + cute_args.append(arg) + + kernel_function(*cute_args, **kwargs) + + + + +def main(): + size = 16384 + a = torch.randn(size, device="cuda", dtype=torch.float16) + b = torch.randn(size, device="cuda", dtype=torch.float16) + c = torch.zeros(size, device="cuda", dtype=torch.float16) + + ''' + a_ = from_dlpack(a, assumed_align=16) + b_ = from_dlpack(b, assumed_align=16) + c_ = from_dlpack(c, assumed_align=16) + ''' + + args = [a, b, c, size] + tune_params = {"num_threads_per_block": [1, 2, 4, 8, 16, 32, 64, 128, 265, 512, 1024]} + answer = [None, None, (a+b).cpu(), None] + + from pathlib import Path + FULL_PATH = Path(__file__).resolve() + tune_kernel("vec_add", FULL_PATH, size, args, tune_params, answer=answer, + lang="generic_python", call_function=call_cute, verbose=True) + #naive_elementwise_add_ = cute.compile(naive_elementwise_add, a_, b_, c_) + + #naive_elementwise_add_(a_, b_, c_) + #vec_add(a_, b_, c_, size) + + + + #torch.testing.assert_close(c, a+b) + +main() + diff --git a/examples/generic_python/matmul/cute_matmul.py b/examples/generic_python/matmul/cute_matmul.py index 90e6e67e9..8654829fa 100644 --- a/examples/generic_python/matmul/cute_matmul.py +++ b/examples/generic_python/matmul/cute_matmul.py @@ -1,32 +1,38 @@ +import math import torch +from typing import Tuple + import cutlass import cutlass.cute as cute - +import cutlass.utils as utils from cutlass.cute.runtime import from_dlpack + from kernel_tuner import tune_kernel from pathlib import Path +from examples.generic_python.call_functions import call_cute FULL_PATH = Path(__file__).resolve() # need export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH to work -# Maybe use this as optimized kernel: https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/ampere/tensorop_gemm.py - +## Basic Matmul ================================================================ @cute.kernel -def naive_matmul_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor): - - tidx, _, _ = cute.arch.thread_idx() - bidx, _, _ = cute.arch.block_idx() - bdim, _, _ = cute.arch.block_dim() - - thread_idx = bidx * bdim + tidx - - M, K = gA.shape - _, N = gB.shape +def naive_matmul_kernel( + gA: cute.Tensor, + gB: cute.Tensor, + gC: cute.Tensor, + M: int, + K: int, + N: int, +): + tx, ty, _ = cute.arch.thread_idx() + bx, by, _ = cute.arch.block_idx() + bdx, bdy, _ = cute.arch.block_dim() - n = thread_idx % N - m = thread_idx // N + # Global indices + n = bx * bdx + tx + m = by * bdy + ty if m < M and n < N: acc = cutlass.Float32(0.0) @@ -39,97 +45,964 @@ def naive_matmul_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor): gC[m, n] = acc.to(gC.element_type) - @cute.jit def matmul( mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor, ): + block_size_x = 16 + block_size_y = 16 + block = (block_size_x, block_size_y, 1) - threads_per_block = 256 - - M, _ = mA.shape + M, K = mA.shape _, N = mB.shape - total_outputs = M * N - - kernel = naive_matmul_kernel(mA, mB, mC) - - kernel.launch( - grid=((total_outputs + threads_per_block - 1) // threads_per_block, 1, 1), - block=(threads_per_block, 1, 1), + grid = ( + (N + block[0] - 1) // block[0], + (M + block[1] - 1) // block[1], + 1, ) + kernel = naive_matmul_kernel(mA, mB, mC, M, K, N) -# Here we formulate the jit wrapper a bit differently, becuase the grid and block sizes -# are computed by kernel tuner. But we could also use threads per block as a tuning param and -# keep the same jit wrapper as above -@cute.jit -def matmul_kt( - mA: cute.Tensor, - mB: cute.Tensor, - mC: cute.Tensor, - grid: cutlass.Constexpr, - block: cutlass.Constexpr, -): - - kernel = naive_matmul_kernel(mA, mB, mC) - - kernel.launch( - grid=grid, - block=block, - ) + kernel.launch(grid=grid, block=block) -def run_example(M, N, K): +def tune_naive_matmul(M, N, K): a = torch.randn(M, K, device="cuda", dtype=torch.float16) b = torch.randn(K, N, device="cuda", dtype=torch.float16) c = torch.zeros(M, N, device="cuda", dtype=torch.float16) - a_ = from_dlpack(a, assumed_align=16) - b_ = from_dlpack(b, assumed_align=16) - c_ = from_dlpack(c, assumed_align=16) - - matmul_ = cute.compile(matmul, a_, b_, c_) - matmul_(a_, b_, c_) - - torch.testing.assert_close(c, a @ b, atol=1e-2, rtol=1e-2) - + args = [a, b, c] + size = M * N + answer = [None, None, (a @ b).cpu()] + tune_params = dict() + tune_params["block_size_x"] = [8, 16, 32, 64, 128] + tune_params["block_size_y"] = [8, 16, 32, 64, 128] -def call_cute(kernel_function, args, kwargs, grid, threads, params): + results, env = tune_kernel("matmul", FULL_PATH, size, args, tune_params, lang="generic_python", + call_function=call_cute, answer=answer, atol=4, verbose=False) + + + +## Optimized matmul ========================================================== +# Copyright (c) 2025 - 2026 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. + + + +""" +A dense GEMM (C = A * B) example for the NVIDIA Ampere architecture using CUTE DSL. +- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") +- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") +- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") + +This GEMM kernel supports the following features: + - Utilizes Ampere's tensor cores for matrix multiply-accumulate (MMA) operations + - Threadblock rasterization to improve data re-use + - Supports multi-stage pipeline to overlap computation and memory access + - Implements shared memory buffering for epilogue to increase coalesced global memory access + +This GEMM works as follows: +1. Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using asynchronous copies. +2. Perform matrix multiply-accumulate (MMA) operations. +3. Store results from registers (RMEM) to shared memory (SMEM), then to global memory (GMEM). + +The Ampere tensor core instruction used operates as follows: +- Read matrix A from SMEM +- Read matrix B from SMEM +- Perform MMA operation and store the result in Accumulator(register) + +To run this example: + +.. code-block:: bash + + python examples/ampere/tensorop_gemm.py \ + --mnkl 8192,8192,8192,1 --atom_layout_mnk 2,2,1 \ + --ab_dtype Float16 \ + --c_dtype Float16 --acc_dtype Float32 \ + --a_major m --b_major n --c_major n + +The above example command computes with M=8192, N=8192, K=8192, +batch_count=1. The atom layout's shape is 2x2x1 and the input, mma +accumulator, and output data type are set as fp16, fp32 and fp16, +respectively. + +To collect performance with NCU profiler: + +.. code-block:: bash + + ncu python examples/ampere/tensorop_gemm.py \ + --mnkl 8192,8192,8192,1 --atom_layout_mnk 2,2,1 \ + --ab_dtype Float16 \ + --c_dtype Float16 --acc_dtype Float32 \ + --a_major m --b_major n --c_major n \ + --skip_ref_check --iterations 2 + +Constraints: +* Supported input and output data types: fp16 +* Support accumulator data types: f32 +* Default tile shape is set to be 128x128x32 +* Atom layout's MNK shape is set so that tile shape can be divided by MMA + instruction shape +* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, + i.e, number of elements is a multiple of 8 +""" + + +class TensorOpGemm: + def __init__(self): + self.ab_dtype = cutlass.Float16 + self.c_dtype = cutlass.Float16 + self.acc_dtype = cutlass.Float32 + tile_m = 128 # Added for KT support + tile_n = 128 + tile_k = 32 + self.cta_tiler = (tile_m, tile_n, tile_k) + self.num_stages = 3 + self.atom_layout_mnk = (2, 2, 1) # moved from paramter to here for KT support + atom_lay_M, atom_lay_N, atom_lay_K = self.atom_layout_mnk + self.num_threads = atom_lay_M * atom_lay_N * atom_lay_K * 32 + + self.bM, self.bN, self.bK = self.cta_tiler + self.mma_inst_shape = (16, 8, 16) + mmaM, mmaN, mmaK = self.mma_inst_shape + + assert self.bM % (atom_lay_M * mmaM) == 0, ( + "bM must be divisible by MMA instruction" + ) + assert self.bN % (atom_lay_N * mmaN) == 0, ( + "bN must be divisible by MMA instruction" + ) + assert atom_lay_K == 1, "this example does not support atom layout K > 1" + assert self.bK % mmaK == 0, "bK must be divisible by MMA instruction" + assert self.num_stages >= 3, "num_stages must be greater than or equal to 3" + + @cute.jit + def __call__( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # The grid divides the problems's M, N, and L dimensions by the + # respective modes of the tile shape (bM, bN, 1). The K dimension is + # handled within a block via a multistage process. + + self.a_major_mode = utils.LayoutEnum.from_tensor(mA) + self.b_major_mode = utils.LayoutEnum.from_tensor(mB) + self.c_major_mode = utils.LayoutEnum.from_tensor(mC) + + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: + # /////////////////////////////////////////////////////////////////////////////// + + # Creates a layout with the size required for the provided tile + # size and num stages (stages are used for K dimension) that is also + # sectioned into 64x8 or 8x32 layout atoms. The swizzle is set so that + # the atom for the shared memory -> register copy does not encounter + # bank conflicts + + # assume the input is 16B align + ab_copy_bits = 128 + sA_layout = self._make_smem_layout_AB( + mA.element_type, + self.a_major_mode, + ab_copy_bits, + (self.cta_tiler[0], self.cta_tiler[2], self.num_stages), + ) + sB_layout = self._make_smem_layout_AB( + mB.element_type, + self.b_major_mode, + ab_copy_bits, + (self.cta_tiler[1], self.cta_tiler[2], self.num_stages), + ) + + # Creates a similar layout but without num_stages or layout atoms + sC_layout = self._make_smem_layout_C( + mC.element_type, + self.c_major_mode, + ab_copy_bits, + (self.cta_tiler[0], self.cta_tiler[1]), + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Tiled copy: + # The majorness of tA/tB/tC follows the majorness of gA/gB/gC, + # enabling merged accesses to global memory for faster data + # transfer between global and shared memory. + # /////////////////////////////////////////////////////////////////////////////// + + # Create a copy atom for a global to shared memory asynchronous copy + atom_async_copy = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp( + cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL + ), + mA.element_type, + num_bits_per_copy=ab_copy_bits, + ) + + # Create thread layouts for tiled copy from the copy atom where the + # thread layout simply follows the leading dimension of the tensor + tiled_copy_A = self._make_gmem_tiled_copy_AB( + atom_async_copy, mA.element_type, self.a_major_mode, ab_copy_bits + ) + tiled_copy_B = self._make_gmem_tiled_copy_AB( + atom_async_copy, mB.element_type, self.b_major_mode, ab_copy_bits + ) + + # Creates a synchronous copy atom and thread layouts for the epilogue + c_copy_bits = 128 + atom_sync_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mC.element_type, + num_bits_per_copy=c_copy_bits, + ) + tiled_copy_C = self._make_gmem_tiled_copy_C( + atom_sync_copy, mC.element_type, self.c_major_mode, c_copy_bits + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Tiled MMA + # /////////////////////////////////////////////////////////////////////////////// + + # Creates a mma atom with 16x8x16 shape for MNK + op = cute.nvgpu.warp.MmaF16BF16Op( + self.ab_dtype, self.acc_dtype, self.mma_inst_shape + ) + + permutation_mnk = ( + self.atom_layout_mnk[0] * self.mma_inst_shape[0], + # if atom layout's N-mode is 1, to leverage the largest coalesced + # shared memory -> register copy, set the tiled mma's N mode to 16 + self.atom_layout_mnk[1] * self.mma_inst_shape[1] * 2, + self.atom_layout_mnk[2] * self.mma_inst_shape[2], + ) + + # Created a tiled mma that tiles the atom according to specified layout. + # For a 2x2x1 atom layout, the mma atom is duplicated 4 times, twice + # across M and twice across N + tC = cute.make_layout(self.atom_layout_mnk) + tiled_mma = cute.make_tiled_mma( + op, + tC, + permutation_mnk=permutation_mnk, + ) + + # grid_dim: ((m + BLK_M - 1) // BLK_M, (n + BLK_N - 1) // BLK_N, l) + grid_dim = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) + + # Add threadblock rasterization to improve re-use of data + raster_factor = 1 + grid_dim_n = cute.size(grid_dim[1]) + # Thresholds picked so that it doesn't cause too many no-op CTAs + if grid_dim_n > 5: + raster_factor = 8 + elif grid_dim_n > 2: + raster_factor = 4 + elif grid_dim_n > 1: + raster_factor = 2 + rasterization_remap_grid_dim = ( + cute.size(grid_dim[0]) * raster_factor, + (cute.size(grid_dim[1]) + raster_factor - 1) // raster_factor, + cute.size(grid_dim[2]), + ) + + self.kernel( + mA, + mB, + mC, + sA_layout, + sB_layout, + sC_layout, + tiled_copy_A, + tiled_copy_B, + tiled_copy_C, + tiled_mma, + raster_factor, + epilogue_op, + ).launch( + grid=rasterization_remap_grid_dim, + block=[self.num_threads, 1, 1], + ) + + @cute.kernel + def kernel( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC: cute.Tensor, + sA_layout: cute.ComposedLayout, + sB_layout: cute.ComposedLayout, + sC_layout: cute.ComposedLayout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_copy_C: cute.TiledCopy, + tiled_mma: cute.TiledMma, + rasterization_factor: cutlass.Int32, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + grid_dim = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) + offset_tile_x, offset_tile_y = self.raster_tile( + bidx, bidy, rasterization_factor + ) + # Early exit if CTA is out of range + if grid_dim[0] <= offset_tile_x or grid_dim[1] <= offset_tile_y: + pass + else: + tiler_coord = (offset_tile_x, offset_tile_y, None) + + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # gA: (BLK_M, BLK_N, k), gB: (BLK_N, BLK_K, k), gC: (BLK_M, BLK_N) + # /////////////////////////////////////////////////////////////////////////////// + gA = cute.local_tile( + mA[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, None, 1), + ) + gB = cute.local_tile( + mB[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(None, 1, 1), + ) + gC = cute.local_tile( + mC[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + + # By default, if the tensor k mode does not divide into the tile k + # size, then last tiles in the k dimension are irregular. + # Instead, make the first tiles irregular when k is irregular. + # This allows us to handle the irregular tile first to avoid + # checking for this condition within the mainloop. + + # residual_k is a negative number indicating the amount needed to + # shift the pointer by in dimension k + residual_k = cute.size(mA, mode=[1]) - cutlass.Int32(self.bK) * cute.size( + gA, mode=[2] + ) + + # move the pointer of gA/gB in the `-k` direction + gA = cute.domain_offset((0, residual_k, 0), gA) + gB = cute.domain_offset((0, residual_k, 0), gB) + # input is 16B aligned + gA = cute.make_tensor(gA.iterator.align(16), gA.layout) + gB = cute.make_tensor(gB.iterator.align(16), gB.layout) + + # Construct identity layout for sA and sB (mirrors global tensors, + # used for predication only) + mcA = cute.make_identity_tensor(mA.layout.shape) + mcB = cute.make_identity_tensor(mB.layout.shape) + cA = cute.local_tile( + mcA[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, None, 1), + ) + cB = cute.local_tile( + mcB[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(None, 1, 1), + ) + + cA = cute.domain_offset((0, residual_k, 0), cA) + cB = cute.domain_offset((0, residual_k, 0), cB) + + # /////////////////////////////////////////////////////////////////////////////// + # Create shared memory buffers and get the appropriate fragments for this thread. + # sA: (BLK_M, BLK_K, PIPE) , sB: (BLK_N, BLK_K, PIPE) + # tAgA: (CPY, CPY_M, CPY_K, k) , tBgB: (CPY, CPY_N, CPY_K, k) + # tAsA: (CPY, CPY_M, CPY_K, PIPE) , tBsB: (CPY, CPY_N, CPY_K, PIPE) + # /////////////////////////////////////////////////////////////////////////////// + @cute.struct + class SharedStorageAB: + a: cute.struct.Align[ + cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], + 16, + ] + b: cute.struct.Align[ + cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], + 16, + ] + + @cute.struct + class SharedStorageC: + c: cute.struct.Align[ + cute.struct.MemRange[mC.element_type, cute.cosize(sC_layout)], + 16, + ] + + # Shared memory buffer + smem = cutlass.utils.SmemAllocator() + # Shared memory allocated for operations with A, B will be + # overwritten for operations on C. This is to improve performance + # by reducing the size of shared memory requested by each block + storage = smem.allocate( + max(SharedStorageAB.size_in_bytes(), SharedStorageC.size_in_bytes()), + byte_alignment=16, + ) + sA = SharedStorageAB(storage).a.get_tensor(sA_layout) + sB = SharedStorageAB(storage).b.get_tensor(sB_layout) + sC = SharedStorageC(storage).c.get_tensor(sC_layout) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + thr_copy_C = tiled_copy_C.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB = thr_copy_B.partition_S(gB) + tBsB = thr_copy_B.partition_D(sB) + tCsC_epilogue = thr_copy_C.partition_S(sC) + tCgC_epilogue = thr_copy_C.partition_D(gC) + + # Repeat the partitioning with identity layouts + tAcA = thr_copy_A.partition_S(cA) + tBcB = thr_copy_B.partition_S(cB) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + + # For predication over the tensors A (M/K), B (N/K), and (in the + # epilogue) C (M/N), we will compute it in a fashion similar to an + # outer product. The predication along one of the dimensions is + # evaluated and stored in a predication tensor. Then, the + # predication for the remaining dimension is handled later via an + # if/else branch at the copy. + # For A and B, predication booleans along M/N are stored in a + # predication tensor and along K is handled via a if/else branch. + + # Allocate predicate tensors for M and N. Predication is checked + # at the granularity of a copy atom, so the predicate tensor does not + # need separate booleans for individual elements within a copy + # atom (for example, the elements of tAgA.shape[0][0].) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAgA.shape[0][1], + cute.size(tAgA, mode=[1]), + cute.size(tAgA, mode=[2]), + ), + stride=(cute.size(tAgA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBsB.shape[0][1], + cute.size(tBsB, mode=[1]), + cute.size(tBsB, mode=[2]), + ), + stride=(cute.size(tBsB, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + # Set predicates for M/N bounds + for rest_v in range(tApA.shape[0]): + for m in range(tApA.shape[1]): + tApA[rest_v, m, 0] = cute.elem_less( + tAcA[(0, rest_v), m, 0, 0][0], mA.shape[0] + ) + for rest_v in range(tBpB.shape[0]): + for n in range(tBpB.shape[1]): + tBpB[rest_v, n, 0] = cute.elem_less( + tBcB[(0, rest_v), n, 0, 0][0], mB.shape[0] + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Prefetch Prologue + # /////////////////////////////////////////////////////////////////////////////// + # Clear the smem tiles to account for predicated off loads + tAsA.fill(0) + tBsB.fill(0) + cute.arch.sync_threads() + # Start async loads for the first k-tile. Here we take care of the k residue + # via if/else check along the k dimension. Because we shifted the identity tensor + # by the residue_k and because the identity tensor is a coord tensor, the + # values of any identity tensor element that is poison is less than -1 + num_smem_stages = cute.size(tAsA, mode=[3]) + k_tile_count = cute.size(tAgA, mode=[3]) + k_tile_index = cutlass.Int32(0) + + for k in range(tApA.shape[2]): + if cute.elem_less(cutlass.Int32(-1), tAcA[0, 0, k, 0][1]): + cute.copy( + tiled_copy_A, + tAgA[None, None, k, k_tile_index], + tAsA[None, None, k, 0], + pred=tApA[None, None, k], + ) + for k in range(tBpB.shape[2]): + if cute.elem_less(cutlass.Int32(-1), tBcB[0, 0, k, 0][1]): + cute.copy( + tiled_copy_B, + tBgB[None, None, k, k_tile_index], + tBsB[None, None, k, 0], + pred=tBpB[None, None, k], + ) + k_tile_index = k_tile_index + 1 + cute.arch.cp_async_commit_group() + + # Start async loads for rest of the k-tiles + for k_tile in range(1, num_smem_stages - 1): + if k_tile == k_tile_count: + tApA.fill(0) + tBpB.fill(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, k_tile], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, k_tile], + pred=tBpB, + ) + k_tile_index = k_tile_index + 1 + cute.arch.cp_async_commit_group() + + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + thr_mma = tiled_mma.get_slice(tidx) + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCsC = thr_mma.partition_C(sC) + tCgC = thr_mma.partition_C(gC) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrC = tiled_mma.make_fragment_C(tCgC) + # Clear the accumulator + tCrC.fill(0.0) + + # /////////////////////////////////////////////////////////////////////////////// + # Copy Atom A/B retiling + # /////////////////////////////////////////////////////////////////////////////// + + # Create the copy atoms for the copy from shared memory to register + atom_copy_s2r_A = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.a_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 + ), + mA.element_type, + ) + atom_copy_s2r_B = cute.make_copy_atom( + cute.nvgpu.warp.LdMatrix8x8x16bOp( + self.b_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 + ), + mB.element_type, + ) + + # Creates the tiled copy so that it matches the thread-value layout + # expected by the tiled mma + tiled_copy_s2r_A = cute.make_tiled_copy_A(atom_copy_s2r_A, tiled_mma) + tiled_copy_s2r_B = cute.make_tiled_copy_B(atom_copy_s2r_B, tiled_mma) + + thr_copy_ldmatrix_A = tiled_copy_s2r_A.get_slice(tidx) + thr_copy_ldmatrix_B = tiled_copy_s2r_B.get_slice(tidx) + tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA) + tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA) + tCsB_copy_view = thr_copy_ldmatrix_B.partition_S(sB) + tCrB_copy_view = thr_copy_ldmatrix_B.retile(tCrB) + + # Current pipe index in smem to read from / write to + smem_pipe_read = 0 + smem_pipe_write = num_smem_stages - 1 + + tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read] + tCsB_p = tCsB_copy_view[None, None, None, smem_pipe_read] + + # /////////////////////////////////////////////////////////////////////////////// + # PREFETCH register pipeline + # /////////////////////////////////////////////////////////////////////////////// + num_k_block = cute.size(tCrA, mode=[2]) + if num_k_block > 1: + # Wait until our first prefetched tile is loaded in + cute.arch.cp_async_wait_group(num_smem_stages - 2) + cute.arch.sync_threads() + # Prefetch the first k-block rmem from the first k-tile + cute.copy( + tiled_copy_s2r_A, + tCsA_p[None, None, 0], + tCrA_copy_view[None, None, 0], + ) + cute.copy( + tiled_copy_s2r_B, + tCsB_p[None, None, 0], + tCrB_copy_view[None, None, 0], + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Mainloop + # 1. Shared memory pipeline (gmem -> smem): + # The default smem pipeline depth is 3, meaning that for shared + # memory buffers, we allocate three times the size described by the + # CTA tiler. We prefetch 2 of these buffers before entering the main + # loop. Considering only the transfer from global memory to shared + # memory, the general structure of the mainloop is: + # (1) copy k-tile from gmem to smem; + # (2) perform gemm computation on k-tile; + # (3) wait for the next copy to finish. + # The `cute.arch.cp_async_wait_group(num_smem_stages - 2)` command + # waits for the number of unfinished 'copy' to be <= 1. The advantage + # of this approach is that it allows for simultaneous production + # (i.e., step (1)) and consumption (i.e., step (2)) of smem. + # A common misconception is to prefetch N buffers and rewrite + # the pipeline logic to wait on N-1 pending copies. The disadvantage + # of this approach is that it requires fully consuming a buffer in + # order to open an empty buffer for the next copy. + # 2. Register pipeline (smem -> register): + # Similarly, the register pipeline produces i+1, consumes i, and + # produces i+2... Notably, i and i+1 do not use the same register, + # eliminating dependencies on the same register for better parallelism. + # 3. Combining the smem and register pipelines results in the mainloop. + # /////////////////////////////////////////////////////////////////////////////// + for k_tile in range(k_tile_count): + for k_block in cutlass.range(num_k_block, unroll_full=True): + if k_block == num_k_block - 1: + tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read] + tCsB_p = tCsB_copy_view[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(num_smem_stages - 2) + cute.arch.sync_threads() + + # Load A, B from shared memory to registers for k_block + 1 + k_block_next = (k_block + 1) % num_k_block # static + cute.copy( + tiled_copy_s2r_A, + tCsA_p[None, None, k_block_next], + tCrA_copy_view[None, None, k_block_next], + ) + cute.copy( + tiled_copy_s2r_B, + tCsB_p[None, None, k_block_next], + tCrB_copy_view[None, None, k_block_next], + ) + + # Fetch next A: To better interleave global memory access and compute + # instructions, we intentionally use the sequence: copy A, perform GEMM, + # then copy B. + if k_block == 0: + if k_tile + num_smem_stages - 1 < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, k_tile_index], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + + # Thread-level register gemm for k_block + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + + # Fetch next B and update smem pipeline read/write + if k_block == 0: + if k_tile + num_smem_stages - 1 < k_tile_count: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, k_tile_index], + tBsB[None, None, None, smem_pipe_write], + pred=tBpB, + ) + k_tile_index = k_tile_index + 1 + cute.arch.cp_async_commit_group() + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == num_smem_stages: + smem_pipe_read = 0 + + # Sync before epilogue + cute.arch.cp_async_wait_group(0) + cute.arch.sync_threads() + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue with fusion + # /////////////////////////////////////////////////////////////////////////////// + tCrD = cute.make_fragment_like(tCrC, self.c_dtype) + tCrD[None] = epilogue_op(tCrC.load()).to(self.c_dtype) + + # Copy results of D back to shared memory + cute.autovec_copy(tCrD, tCsC) + + # Create coord tensor for C + ceilM, ceilN, _ = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) + mcC = cute.make_identity_tensor( + ( + cute.size(ceilM) * self.cta_tiler[0], + cute.size(ceilN) * self.cta_tiler[1], + 1, + ) + ) + cC = cute.local_tile( + mcC[None, None, bidz], + tiler=self.cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + tCcC = thr_copy_C.partition_S(cC) + + tCrC_epilogue = cute.make_fragment_like(tCsC_epilogue) + # Wait for all writes to shared memory to finish before starting copies + # using the new layouts + cute.arch.sync_threads() + cute.autovec_copy(tCsC_epilogue, tCrC_epilogue) + + # Create predication tensor for m + tCpC = cute.make_rmem_tensor( + cute.make_layout( + ( + tCgC_epilogue.shape[0][1], + cute.size(tCgC_epilogue, mode=[1]), + cute.size(tCgC_epilogue, mode=[2]), + ), + stride=(cute.size(tCgC_epilogue, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tCpC.shape[0]): + for m in range(tCpC.shape[1]): + tCpC[rest_v, m, 0] = cute.elem_less( + tCcC[(0, rest_v), m, 0][0], mC.shape[0] + ) + + # Copy to global memory using better vectorization + for rest_v in range(tCpC.shape[0]): + for n in range(tCpC.shape[2]): + if cute.elem_less(tCcC[(0, rest_v), 0, n][1], mC.shape[1]): + cute.copy( + tiled_copy_C, + tCrC_epilogue[None, None, n], + tCgC_epilogue[None, None, n], + pred=tCpC[None, None, n], + ) + return + + def _make_smem_layout_AB(self, dtype, major_mode, copy_bits, smem_tiler): + major_mode_size = ( + smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] + ) + major_mode_size = 64 if major_mode_size >= 64 else major_mode_size + + swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) + swizzle_bits = min(swizzle_bits, 3) + + layout_atom_outer = ( + cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) + ) + layout_atom = cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 3, 3), + 0, + layout_atom_outer, + ) + layout = cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2)) + return layout + + def _make_smem_layout_C(self, dtype, major_mode, copy_bits, smem_tiler): + major_mode_size = ( + smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] + ) + + swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) + swizzle_bits = min(swizzle_bits, 3) + + layout_atom_outer = ( + cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) + ) + layout_atom = cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, 3, 4), + 0, + layout_atom_outer, + ) + + # Due to the thread layout of the mma, remove swizzle in C to + # prevent shared memory fragments owned by an single thread from + # holding swizzles + if major_mode == utils.LayoutEnum.COL_MAJOR: + layout_atom = cute.make_composed_layout( + cute.make_swizzle(0, 3, 4), 0, layout_atom_outer + ) + layout = cute.tile_to_shape( + layout_atom, + smem_tiler, + (0, 1), + ) + return layout + + def _make_gmem_tiled_copy_AB(self, atom_copy, dtype, major_mode, copy_bits): + copy_elems = copy_bits // dtype.width + shape_dim_1 = cute.size(self.bK) // copy_elems + # thread layout for copy + thread_layout = cute.make_layout( + (self.num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) + ) + if major_mode != utils.LayoutEnum.ROW_MAJOR: + shape_dim_0 = cute.size(self.bM) // copy_elems + thread_layout = cute.make_layout( + (shape_dim_0, self.num_threads // shape_dim_0), stride=(1, shape_dim_0) + ) + # Value layout for copy + value_layout = ( + cute.make_layout((1, copy_elems)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((copy_elems, 1)) + ) + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) + + def _make_gmem_tiled_copy_C(self, atom_copy, dtype, major_mode, copy_bits): + copy_elems = copy_bits // dtype.width + shape_dim_1 = cute.size(self.bN) // copy_elems + # thread layout for copy + thread_layout = cute.make_layout( + (self.num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) + ) + if major_mode != utils.LayoutEnum.ROW_MAJOR: + shape_dim_0 = cute.size(self.bM) // copy_elems + thread_layout = cute.make_layout( + (shape_dim_0, self.num_threads // shape_dim_0), stride=(1, shape_dim_0) + ) + value_layout = ( + cute.make_layout((1, copy_elems)) + if major_mode == utils.LayoutEnum.ROW_MAJOR + else cute.make_layout((copy_elems, 1)) + ) + return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) + + def raster_tile(self, i, j, f): + new_i = i // f + new_j = (i % f) + (j * f) + return (new_i, new_j) + + +# Need a custum call function for the optimized kernel because of the tensor layout +def call_cute_custom(kernel_function, args, kwargs, grid, threads, params): + # Initialize cache if it does not exist + if not hasattr(call_cute_custom, "custom_cache"): + call_cute_custom.custom_cache = {} + + # Convert Torch tensors to CuTe tensors with correct layout cute_args = [] for arg in args: if isinstance(arg, torch.Tensor): - arg_ = from_dlpack(arg) - cute_args.append(arg_) + cute_tensor = ( + from_dlpack(arg, assumed_align=16) + .mark_layout_dynamic(leading_dim=1) + .mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=8, + ) + ) + cute_args.append(cute_tensor) else: cute_args.append(arg) - kernel_function(*cute_args, grid, threads) + # Form cache key from tuning parameters + param_keys = sorted(params.keys()) + cache_str = type(kernel_function).__name__ + for k in param_keys: + cache_str += "_" + str(params[k]) + + # Check if kernel exists in cache. Otherwise, compile and save + if cache_str in call_cute_custom.custom_cache: + compiled_kernel = call_cute_custom.custom_cache[cache_str] + else: + # Wrap in try-except because CuTe gives TypeError for certain block sizes. This + # stops the tuning process and we do not want this + try: + compiled_kernel = cute.compile(kernel_function, *cute_args) + call_cute_custom.custom_cache[cache_str] = compiled_kernel + except: + raise RuntimeError("Invalid configuration") + + + + compiled_kernel(*cute_args) + + + +def tune_optimized(mnkl: Tuple[int, int, int, int]): + + M, N, K, L = mnkl + + a = torch.randn((L, M, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) + b = torch.randn((L, N, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) # b is transposed + c = torch.empty((L, M, N), device="cuda", dtype=torch.float16).permute(1, 2, 0) + + c_ref = torch.matmul( + a.permute(2, 0, 1), + b.permute(2, 1, 0) + ).permute(1, 2, 0) + + args = [a, b, c] + tune_params = { + "tile_m": [64, 128, 256], + "tile_n": [64, 128, 256], + "tile_k": [16, 32, 64], + } -def tune(M, N, K): - a = torch.randn(M, K, device="cuda", dtype=torch.float16) - b = torch.randn(K, N, device="cuda", dtype=torch.float16) - c = torch.zeros(M, N, device="cuda", dtype=torch.float16) + constraints = ["3 * (tile_m * tile_k + tile_n * tile_k) * 2 <= 100 * 1024"] # SMEM constraint - args = [a, b, c] - size = M * N - answer = [None, None, (a @ b).cpu()] - tune_params = dict() - tune_params["block_size_x"] = [16, 32, 64, 128] - tune_params["block_size_y"] = [16, 32, 64, 128] - results, env = tune_kernel("matmul_kt", FULL_PATH, size, args, tune_params, lang="generic_python", - call_function=call_cute, answer=answer, atol=1e-2, verbose=True) + results, env = tune_kernel("TensorOpGemm", FULL_PATH, M * N, args, tune_params, verbose=True, restrictions=constraints, + lang="generic_python", call_function=call_cute_custom, answer=[None, None, c_ref.cpu()], atol=4) +if __name__ == "__main__": + #m, n, k = 4096, 4096, 4096 + #tune_naive_matmul(m, n, k) + mnkl = (4096, 4096, 4096, 1) + tune_optimized(mnkl) -if __name__ == "__main__": - run_example(128, 128, 128) - tune(128, 128, 128) diff --git a/examples/generic_python/matmul/numba_matmul.py b/examples/generic_python/matmul/numba_matmul.py index 00102a382..af6cce347 100644 --- a/examples/generic_python/matmul/numba_matmul.py +++ b/examples/generic_python/matmul/numba_matmul.py @@ -1,6 +1,6 @@ import torch import numpy as np -from numba import cuda +from numba import cuda, float32 from kernel_tuner import tune_kernel from pathlib import Path @@ -21,7 +21,7 @@ def matmul(A, B, C): # Example taken from https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication # Changed data type from float32 to float16 -@cuda.jit(cache=True) +@cuda.jit(cache=True, fastmath=True) def fast_matmul(A, B, C): # Define an array in the shared memory # The size and type of the arrays must be known at compile time @@ -60,6 +60,90 @@ def fast_matmul(A, B, C): C[x, y] = tmp +# Translated to numba-cuda from https://github.com/cupy/cupy/blob/main/examples/gemm/sgemm.cu +@cuda.jit(cache=True) +def optimized_matmul(M, N, K, A, B, C): + DIM_X = 16 + DIM_Y = 16 + BLK_M = 64 + BLK_N = 64 + BLK_K = 16 + THR_M = 4 + THR_N = 4 + + # thread indices + idx = cuda.threadIdx.x + idy = cuda.threadIdx.y + + blx = cuda.blockIdx.x + bly = cuda.blockIdx.y + + # shared memory + sA = cuda.shared.array((BLK_K, BLK_M), np.float16) + sB = cuda.shared.array((BLK_N, BLK_K), np.float16) + + # registers + rC = cuda.local.array((THR_N, THR_M), float32) + rA = cuda.local.array(THR_M, float32) + rB = cuda.local.array(THR_N, float32) + + # init accumulator + for n in range(THR_N): + for m in range(THR_M): + rC[n][m] = 0.0 + + # global indices + base_row = blx * BLK_M + base_col = bly * BLK_N + + # loop over K tiles + for kk in range(0, K, BLK_K): + + # load A tile into shared memory + for i in range(idy, BLK_K, DIM_Y): + for j in range(idx, BLK_M, DIM_X): + row = base_row + j + col = kk + i + if row < M and col < K: + sA[i, j] = A[row, col] + else: + sA[i, j] = 0.0 + + # load B tile + for i in range(idy, BLK_N, DIM_Y): + for j in range(idx, BLK_K, DIM_X): + row = kk + j + col = base_col + i + if row < K and col < N: + sB[i, j] = B[row, col] + else: + sB[i, j] = 0.0 + + cuda.syncthreads() + + # compute + for k in range(BLK_K): + for m in range(THR_M): + rA[m] = float32(sA[k, m * DIM_X + idx]) + + for n in range(THR_N): + rB[n] = float32(sB[n * DIM_Y + idy, k]) + + for n in range(THR_N): + for m in range(THR_M): + rC[n][m] += rA[m] * rB[n] + + cuda.syncthreads() + + # write back + for n in range(THR_N): + col = base_col + n * DIM_Y + idy + for m in range(THR_M): + row = base_row + m * DIM_X + idx + if row < M and col < N: + C[row, col] = np.float16(rC[n][m]) + + def run_matmul(M, N, K): # create numpy arrays @@ -135,9 +219,52 @@ def tune(M, N, K): if __name__ == "__main__": - run_matmul(128, 96, 64) + + #run_matmul(128, 96, 64) #tune(1024, 1024, 1024) + + M, N, K = 1024, 1024, 1024 + + # random FP16 inputs + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + + # output + C = np.zeros((M, N), dtype=np.float16) + + # copy to device + dA = cuda.to_device(A) + dB = cuda.to_device(B) + dC = cuda.to_device(C) + + # launch config + threads_per_block = (16, 16) + + blocks_per_grid_x = (M + 64 - 1) // 64 + blocks_per_grid_y = (N + 64 - 1) // 64 + blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y) + + # run kernel + optimized_matmul[blocks_per_grid, threads_per_block](M, N, K, dA, dB, dC) + + # copy result back + C_result = dC.copy_to_host() + + # reference (FP32 accumulate like your kernel) + C_ref = (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) + + # check error + max_error = np.max(np.abs(C_result - C_ref)) + print("Max error:", max_error) + + # tolerance (FP16 is noisy) + if max_error < 4: + print("✅ Looks correct") + else: + print("❌ Something is off") + print("Expected: ", C_ref, "\nGot: ", C_result) + ''' 128: 16, 8 256: 8, 16 diff --git a/examples/generic_python/matmul/taichi_matmul.py b/examples/generic_python/matmul/taichi_matmul.py index 6eca9d21e..b8a1c9388 100644 --- a/examples/generic_python/matmul/taichi_matmul.py +++ b/examples/generic_python/matmul/taichi_matmul.py @@ -2,16 +2,15 @@ import taichi as ti from kernel_tuner import tune_kernel -from pathlib import Path -FULL_PATH = Path(__file__).resolve() + ti.init(arch=ti.gpu) -# NOTE tuning works on vibranium. Taichi does not work on DAS6 + # TODO make sure this is zero-copy @ti.kernel -def matmul(A: ti.types.ndarray(), B: ti.types.ndarray(), C: ti.types.ndarray()): +def matmul(A: ti.types.ndarray(dtype=ti.f16, ndim=2), B: ti.types.ndarray(dtype=ti.f16, ndim=2), C: ti.types.ndarray(dtype=ti.f16, ndim=2)): BLOCK_DIM = 16 ti.loop_config(block_dim=BLOCK_DIM) @@ -21,16 +20,16 @@ def matmul(A: ti.types.ndarray(), B: ti.types.ndarray(), C: ti.types.ndarray()): sum = 0.0 for k in range(K_dim): sum += A[i, k] * B[k, j] - C[i, j] = sum + C[i, j] = ti.cast(sum, ti.f16) def call_taichi(kernel_function, args, kwargs): kernel_function(*args, **kwargs) def tune(M, N, K): - torch_A = torch.rand((N, K), device='cuda', dtype=torch.float32) - torch_B = torch.rand((K, M), device='cuda', dtype=torch.float32) - torch_C = torch.empty((N, M), device='cuda', dtype=torch.float32) + torch_A = torch.rand((N, K), device='cuda', dtype=torch.float16) + torch_B = torch.rand((K, M), device='cuda', dtype=torch.float16) + torch_C = torch.empty((N, M), device='cuda', dtype=torch.float16) size = M * N args = [torch_A, torch_B, torch_C] @@ -40,13 +39,14 @@ def tune(M, N, K): results, env = tune_kernel( kernel_name="matmul", - kernel_source=FULL_PATH, + kernel_source=__file__, problem_size=size, arguments=args, tune_params=tune_params, answer=answer, lang="generic_python", call_function=call_taichi, + atol=1e-1, ) if __name__ == "__main__": diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py index 7ced8cc83..2ea510d91 100644 --- a/examples/generic_python/matmul/tilelang_matmul.py +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -2,6 +2,8 @@ import tilelang import tilelang.language as T +import itertools + @tilelang.jit def matmul_basic(M:int, N:int, K:int, block_M:int, block_N:int, block_K:int, dtype:str="float16", accum_dtype:str="float32"): @@ -31,11 +33,13 @@ def gemm( return gemm + + # https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py # changed gemm_autotune to gemm # originally, B was transposed. I changed this so the kernel input is the same as in other languages. -#@tilelang.jit -def matmul_opt(M, N, K, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): +@tilelang.jit +def matmul_opt(M, N, K, dummy, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): @T.prim_func def gemm( A: T.Tensor((M, K), dtype), @@ -65,10 +69,105 @@ def gemm( -def main(): - kernel = matmul_opt(1024, 1024, 1024, 128, 128, 32, 3, 128, False) +def get_configs(kt=False): + """ + Generate a list of kernel tuning configuration dictionaries for a tiled matrix-multiply. + This function is used for tuning experiments with the built-in tuner + + Returns: + List[dict]: A list of configuration dictionaries + """ + + ''' + block_M = [64, 128, 256] + block_N = [64, 128, 256] + block_K = [32, 64] + num_stages = [0, 1, 2, 3] + thread_num = [128, 256] + enable_rasterization = [True, False] + ''' + + block_M = [64] + block_N = [64] + block_K = [32, 64] + num_stages = [0, 3] + thread_num = [128] + enable_rasterization = [True] + + _configs = list( + itertools.product( + block_M, + block_N, + block_K, + num_stages, + thread_num, + enable_rasterization, + ) + ) + + if kt: + configs = { + "block_M": block_M, + "block_N": block_N, + "block_K": block_K, + "num_stages": num_stages, + "thread_num": thread_num, + "enable_rasteration": enable_rasterization, + } + else: + configs = [ + { + "block_M": c[0], + "block_N": c[1], + "block_K": c[2], + "num_stages": c[3], + "thread_num": c[4], + "enable_rasteration": c[5], # keep param name for backward-compat + } + for c in _configs + ] + return configs + +# https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py +# changed gemm_autotune to gemm +# originally, B was transposed. I changed this so the kernel input is the same as in other languages. +@tilelang.autotune(configs=get_configs()) +@tilelang.jit +def matmul_opt_autotune(M, N, K, dummy, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): + @T.prim_func + def gemm_autotune( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_shared = T.alloc_shared((block_M, block_N), dtype) + T.use_swizzle(panel_size=10, enable=enable_rasteration) + T.clear(C_local) + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + T.copy(A[by * block_M, k * block_K], A_shared) + T.copy(B[k * block_K, bx * block_N], B_shared) + T.gemm( + A_shared, + B_shared, + C_local, + transpose_B=False, + ) + T.copy(C_local, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return gemm_autotune + + + + +def main(): + kernel = matmul_opt_autotune(1024, 1024, 1024) import torch a = torch.randn(1024, 1024).cuda().half() diff --git a/examples/generic_python/matmul/warp_matmul.py b/examples/generic_python/matmul/warp_matmul.py index 9e6b7dcf3..78a90ce65 100644 --- a/examples/generic_python/matmul/warp_matmul.py +++ b/examples/generic_python/matmul/warp_matmul.py @@ -19,7 +19,7 @@ # GEMM example from https://nvidia.github.io/warp/user_guide/tiles.html @wp.kernel -def tile_gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)): +def tile_gemm(A: wp.array2d(dtype=wp.float16), B: wp.array2d(dtype=wp.float16), C: wp.array2d(dtype=wp.float16)): # output tile index i, j = wp.tid() @@ -36,17 +36,16 @@ def tile_gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.arra a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i*TILE_M, k*TILE_K), bounds_check=True) b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k*TILE_K, j*TILE_N), bounds_check=True) - # sum += a*b wp.tile_matmul(a, b, sum) - wp.tile_store(C, sum, offset=(i*TILE_M, j*TILE_N), bounds_check=True) + wp.tile_store(C, wp.tile_astype(sum, wp.float16), offset=(i*TILE_M, j*TILE_N), bounds_check=True) def run_kernel_direct(M, N, K): rng = np.random.default_rng(42) - A = rng.random((M, K), dtype=np.float32) - B = rng.random((K, N), dtype=np.float32) - C = np.zeros((M, N), dtype=np.float32) + A = rng.random((M, K)).astype(np.float16) + B = rng.random((K, N)).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) A_wp = wp.array(A) B_wp = wp.array(B) @@ -59,7 +58,7 @@ def run_kernel_direct(M, N, K): inputs=[A_wp, B_wp, C_wp], block_dim=TILE_THREADS) - np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-3) + np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-1) print("Example matrix multiplication passed") @@ -80,15 +79,15 @@ def call_warp(kernel_function, args, kwargs, grid, threads, params): kernel_function, dim=grid, inputs=warp_args, - block_dim=params["TILE_THREADS"], # We could directly take threads, but in the given example this is a constant + block_dim=params["TILE_THREADS"], ) def tune(M, K, N): rng = np.random.default_rng(42) - A = rng.random((M, K), dtype=np.float32) - B = rng.random((K, N), dtype=np.float32) - C = np.zeros((M, N), dtype=np.float32) + A = rng.random((M, K)).astype(np.float16) + B = rng.random((K, N)).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) size = (M, N) block_size_names = ["TILE_M", "TILE_N"] @@ -126,6 +125,8 @@ def tune(M, K, N): (129, 130, 33), ] - for size in sizes: - print(size) - run_kernel_direct(*size) + #for size in sizes: + # print(size) + # run_kernel_direct(*size) + + run_kernel_direct(1024, 1024, 1024) From 68a48468f0b1453f63bd3eb5e530aee5021c5382 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Fri, 1 May 2026 08:47:50 +0200 Subject: [PATCH 10/44] examples --- examples/generic_python/call_functions.py | 27 ++- examples/generic_python/matmul/cupy_matmul.py | 98 +++----- examples/generic_python/matmul/cute_matmul.py | 115 ++++++--- .../generic_python/matmul/numba_matmul.py | 218 ++++++++---------- .../generic_python/matmul/taichi_matmul.py | 43 ++-- .../generic_python/matmul/tilelang_matmul.py | 212 ++++++++++++++--- .../generic_python/matmul/tilus_matmul.py | 213 +++++++++++------ .../generic_python/matmul/triton_matmul.py | 166 ++++++++++++- examples/generic_python/matmul/warp_matmul.py | 129 ++++++----- kernel_tuner/backends/generic_python.py | 4 +- 10 files changed, 816 insertions(+), 409 deletions(-) diff --git a/examples/generic_python/call_functions.py b/examples/generic_python/call_functions.py index 734568efd..8c74790e1 100644 --- a/examples/generic_python/call_functions.py +++ b/examples/generic_python/call_functions.py @@ -91,11 +91,24 @@ def call_warp(kernel_function, args, kwargs, grid, threads, params): else: warp_args.append(arg) + # Check if block_dim is in the tuning parameters. Otherwise, use + # the computed thread block dimensions. + if 'block_dim' in params.keys(): + threads_per_block = params['block_dim'] + else: + threads_per_block = threads[0] * threads[1] * threads[2] + + # Check if dim is in the tuning parameters. Otherwise, compute from + # grid and threads. + if 'dim' in params.keys(): + dimensions = params['dim'] + else: + dimensions = [grid[i] * threads[i] for i in range(len(grid))] + # launch kernel - with wp.Tape() as tape: - wp.launch_tiled( - kernel_function, - dim=grid, - inputs=warp_args, - block_dim=params["TILE_THREADS"], # We could directly take threads, but in the given example this is a constant - ) + wp.launch( + kernel_function, + dim=dimensions, + inputs=warp_args, + block_dim=threads_per_block + ) \ No newline at end of file diff --git a/examples/generic_python/matmul/cupy_matmul.py b/examples/generic_python/matmul/cupy_matmul.py index acd3703fb..1101b5fae 100644 --- a/examples/generic_python/matmul/cupy_matmul.py +++ b/examples/generic_python/matmul/cupy_matmul.py @@ -1,105 +1,71 @@ import cupy as cp from cupyx import jit import numpy as np -import torch from kernel_tuner import tune_kernel -from pathlib import Path - - -@jit.rawkernel() -def gemm_raw_strided(a, b, c, M, N, K): # with outer loop so that we can have more work per thread - # global thread indices - row = jit.blockIdx.y * jit.blockDim.y + jit.threadIdx.y - col = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x - - # grid stride increments - stride_y = jit.gridDim.y * jit.blockDim.y - stride_x = jit.gridDim.x * jit.blockDim.x - - i = row - while i < M: - j = col - while j < N: - acc = 0.0 - for kk in range(K): - acc += a[i, kk] * b[kk, j] - c[i, j] = acc - j += stride_x - i += stride_y +from examples.generic_python.call_functions import call_cupyx @jit.rawkernel() def gemm(a, b, c, M, N, K): row = jit.blockIdx.y * jit.blockDim.y + jit.threadIdx.y col = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x + + if row < M and col < N: + acc = 0.0 + for kk in range(K): + acc += a[row, kk] * b[kk, col] + c[row, col] = acc - acc = 0.0 - for kk in range(K): - acc += a[row, kk] * b[kk, col] - c[row, col] = acc - - -def run(): - M, K, N = 128, 256, 64 - # random test data - A = cp.random.random((M, K), dtype=cp.float32) - B = cp.random.random((K, N), dtype=cp.float32) - C = cp.zeros((M, N), dtype=cp.float32) +def run(M, N, K): + # float16 matrices on GPU + a = cp.random.random((M, K)).astype(cp.float16) + b = cp.random.random((K, N)).astype(cp.float16) + c = cp.zeros((M, N), dtype=cp.float16) - # launch parameters + # block / grid configuration block = (16, 16) - grid = ((N + block[0] - 1) // block[0], - (M + block[1] - 1) // block[1]) + grid = ((N + block[0] - 1) // block[0], (M + block[1] - 1) // block[1]) # launch kernel - gemm_raw_strided(grid, block, (A, B, C, M, N, K)) + gemm[grid, block](a, b, c, M, N, K) + cp.cuda.Device().synchronize() - # validate - C_ref = A.dot(B) - print("max error:", float(cp.max(cp.abs(C - C_ref)))) + # Correctness verification + c_ref = cp.matmul(a, b) + assert cp.allclose(c, c_ref, rtol=1e-2, atol=1e-1) + print("Succes") -def call_cupyx(kernel_function, args, kwargs, grid, threads): - cupy_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - cupy_args.append(cp.from_dlpack(arg)) - else: - cupy_args.append(arg) - kernel_function(grid, threads, tuple(cupy_args)) -def tune(): - M, K, N = 128, 256, 64 - - # random test data. Here we had to change cupy to numpy arrays. +def tune(M, N, K): + # random test data. Here we had to use numpy arrays instead of cupy. A = np.random.rand(M, K).astype(np.float16) B = np.random.rand(K, N).astype(np.float16) C = np.zeros((M, N), dtype=np.float16) - args = [A, B, C, M, N, K] size = (N, M) - tune_params = {"block_size_x": [2**i for i in range(10)], "block_size_y": [2**i for i in range(10)]} - restrictions = ["block_size_x == block_size_y"] - source = Path(__file__).resolve() + tune_params = {"block_size_x": [2**i for i in range(11)], "block_size_y": [2**i for i in range(11)]} + restrictions = ["block_size_x * block_size_y <= 1024"] results, env = tune_kernel( kernel_name="gemm", - kernel_source=source, + kernel_source=__file__, problem_size=size, arguments=args, tune_params=tune_params, - answer=[None, None, A.dot(B), None, None], - atol=1e-2, + answer=[None, None, A.dot(B), None, None, None], + atol=1e-1, call_function=call_cupyx, lang="generic_python", - verbose=True, - restrictions=restrictions, + restrictions=restrictions, + verbose=True, ) if __name__ == "__main__": - #tune() - run() \ No newline at end of file + M, N, K = 1024, 1024, 1024 + tune(M, N, K) + \ No newline at end of file diff --git a/examples/generic_python/matmul/cute_matmul.py b/examples/generic_python/matmul/cute_matmul.py index 8654829fa..4b5637c82 100644 --- a/examples/generic_python/matmul/cute_matmul.py +++ b/examples/generic_python/matmul/cute_matmul.py @@ -8,12 +8,9 @@ from cutlass.cute.runtime import from_dlpack from kernel_tuner import tune_kernel -from pathlib import Path from examples.generic_python.call_functions import call_cute -FULL_PATH = Path(__file__).resolve() - -# need export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH to work +# might need export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH to work ## Basic Matmul ================================================================ @@ -69,6 +66,23 @@ def matmul( kernel.launch(grid=grid, block=block) +def run_naive_matmul(M, N, K): + a = torch.randn(M, K, device="cuda", dtype=torch.float16) + b = torch.randn(K, N, device="cuda", dtype=torch.float16) + c = torch.zeros(M, N, device="cuda", dtype=torch.float16) + c_ref = a @ b + + # Convert to CuTe tensors + a_ = from_dlpack(a, assumed_align=16) + b_ = from_dlpack(b, assumed_align=16) + c_ = from_dlpack(c, assumed_align=16) + + compiled_kernel = cute.compile(matmul, a_, b_, c_) + compiled_kernel(a_, b_, c_) + + assert torch.allclose(c, c_ref, atol=M * 2 **(-11), rtol=1e-2) + print("Succes") + def tune_naive_matmul(M, N, K): a = torch.randn(M, K, device="cuda", dtype=torch.float16) @@ -79,11 +93,12 @@ def tune_naive_matmul(M, N, K): size = M * N answer = [None, None, (a @ b).cpu()] tune_params = dict() - tune_params["block_size_x"] = [8, 16, 32, 64, 128] - tune_params["block_size_y"] = [8, 16, 32, 64, 128] + tune_params["block_size_x"] = [2**i for i in range(1, 10)] + tune_params["block_size_y"] = [2**i for i in range(1, 10)] + restrictions = ["block_size_x * block_size_y <= 1024"] - results, env = tune_kernel("matmul", FULL_PATH, size, args, tune_params, lang="generic_python", - call_function=call_cute, answer=answer, atol=4, verbose=False) + results, env = tune_kernel("matmul", __file__, size, args, tune_params, lang="generic_python", + call_function=call_cute, answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) @@ -182,16 +197,17 @@ def __init__(self): self.ab_dtype = cutlass.Float16 self.c_dtype = cutlass.Float16 self.acc_dtype = cutlass.Float32 - tile_m = 128 # Added for KT support - tile_n = 128 - tile_k = 32 - self.cta_tiler = (tile_m, tile_n, tile_k) + self.bM = 128 # extracted from cta_tiler for KT support + self.bN = 128 # extracted from cta_tiler for KT support + self.bK = 32 # extracted from cta_tiler for KT support + self.cta_tiler = (self.bM, self.bN, self.bK) self.num_stages = 3 - self.atom_layout_mnk = (2, 2, 1) # moved from paramter to here for KT support - atom_lay_M, atom_lay_N, atom_lay_K = self.atom_layout_mnk + atom_lay_M = 2 # extracted from atom_layout_mnk for KT support + atom_lay_N = 2 # extracted from atom_layout_mnk for KT support + atom_lay_K = 1 # extracted from atom_layout_mnk for KT support + self.atom_layout_mnk = (atom_lay_M, atom_lay_N, atom_lay_K) # moved from init parameter to here for KT support self.num_threads = atom_lay_M * atom_lay_N * atom_lay_K * 32 - self.bM, self.bN, self.bK = self.cta_tiler self.mma_inst_shape = (16, 8, 16) mmaM, mmaN, mmaK = self.mma_inst_shape @@ -969,10 +985,39 @@ def call_cute_custom(kernel_function, args, kwargs, grid, threads, params): -def tune_optimized(mnkl: Tuple[int, int, int, int]): - - M, N, K, L = mnkl - +def run_optimized(M, N, K, L=1): + a = torch.randn((L, M, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) + b = torch.randn((L, N, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) # b is transposed + c = torch.empty((L, M, N), device="cuda", dtype=torch.float16).permute(1, 2, 0) + + c_ref = torch.matmul( + a.permute(2, 0, 1), + b.permute(2, 1, 0) + ).permute(1, 2, 0) + + def create_cute_tensor(torch_tensor): + cute_tensor = ( + from_dlpack(torch_tensor, assumed_align=16) + .mark_layout_dynamic(leading_dim=1) + .mark_compact_shape_dynamic( + mode=1, + stride_order=(2, 0, 1), + divisibility=8, + ) + ) + return cute_tensor + + a_, b_, c_ = create_cute_tensor(a), create_cute_tensor(b), create_cute_tensor(c) + + gemm = TensorOpGemm() + compiled_kernel = cute.compile(gemm, a_, b_, c_) + compiled_kernel(a_, b_, c_) + + assert torch.allclose(c_ref, c, atol= M * 2**(-11), rtol=1e-2) + print("Sucess") + + +def tune_optimized(M, N, K, L=1): a = torch.randn((L, M, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) b = torch.randn((L, N, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) # b is transposed c = torch.empty((L, M, N), device="cuda", dtype=torch.float16).permute(1, 2, 0) @@ -985,24 +1030,34 @@ def tune_optimized(mnkl: Tuple[int, int, int, int]): args = [a, b, c] tune_params = { - "tile_m": [64, 128, 256], - "tile_n": [64, 128, 256], - "tile_k": [16, 32, 64], + "bM":[16, 32, 64, 128, 256], + "bN":[32, 64, 128, 256], + "bK": [16, 32, 64, 128], # must be divisable by 16 + "num_stages": [3, 4, 5], # restricted to >= 3 by the kernel + "atom_lay_M": [1, 2, 4], + "atom_lay_N": [1, 2], } - constraints = ["3 * (tile_m * tile_k + tile_n * tile_k) * 2 <= 100 * 1024"] # SMEM constraint + restrictions = [ + "bM % (atom_lay_M * 16) == 0", "bN % (atom_lay_N * 8) == 0", # layout + "atom_lay_M * atom_lay_N * 32 <= 1024", # number of threads + "num_stages * (bK * (bM + bN)) * 2 <= 49152", # SMEM + "bM >= atom_lay_M * 32", # ensure each atom sees at least 2 MMA tiles per dimension + "bN >= atom_lay_N * 16", # ensure each atom sees at least 2 MMA tiles per dimension + ] - results, env = tune_kernel("TensorOpGemm", FULL_PATH, M * N, args, tune_params, verbose=True, restrictions=constraints, - lang="generic_python", call_function=call_cute_custom, answer=[None, None, c_ref.cpu()], atol=4) + + results, env = tune_kernel("TensorOpGemm", __file__, M * N, args, tune_params, verbose=True, restrictions=restrictions, #strategy="bayes_opt", + lang="generic_python", call_function=call_cute_custom, answer=[None, None, c_ref.cpu()], atol=M * 2**(-10)) if __name__ == "__main__": - #m, n, k = 4096, 4096, 4096 + m, n, k = 4096, 4096, 4096 + #run_naive_matmul(m, n, k) #tune_naive_matmul(m, n, k) - mnkl = (4096, 4096, 4096, 1) - tune_optimized(mnkl) - - + # mnkl = (4096, 4096, 4096, 1) + tune_optimized(m, n, k) + #run_optimized(m, n, k) diff --git a/examples/generic_python/matmul/numba_matmul.py b/examples/generic_python/matmul/numba_matmul.py index af6cce347..d8e23d6fe 100644 --- a/examples/generic_python/matmul/numba_matmul.py +++ b/examples/generic_python/matmul/numba_matmul.py @@ -1,16 +1,13 @@ -import torch import numpy as np from numba import cuda, float32 + from kernel_tuner import tune_kernel -from pathlib import Path +from examples.generic_python.call_functions import call_numba -FULL_PATH = Path(__file__).resolve() -# Example taken from https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication +# Source: https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication @cuda.jit(cache=True) def matmul(A, B, C): - """Perform square matrix multiplication of C = A * B - """ i, j = cuda.grid(2) if i < C.shape[0] and j < C.shape[1]: tmp = 0. @@ -19,48 +16,7 @@ def matmul(A, B, C): C[i, j] = tmp -# Example taken from https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication -# Changed data type from float32 to float16 -@cuda.jit(cache=True, fastmath=True) -def fast_matmul(A, B, C): - # Define an array in the shared memory - # The size and type of the arrays must be known at compile time - TPB = 16 # TEMP voor overhead testing - sA = cuda.shared.array(shape=(TPB, TPB), dtype=np.float16) - sB = cuda.shared.array(shape=(TPB, TPB), dtype=np.float16) - - x, y = cuda.grid(2) - - tx = cuda.threadIdx.x - ty = cuda.threadIdx.y - bpg = cuda.gridDim.x # blocks per grid - - if x >= C.shape[0] and y >= C.shape[1]: - # Quit if (x, y) is outside of valid C boundary - return - - # Each thread computes one element in the result matrix. - # The dot product is chunked into dot products of TPB-long vectors. - tmp = 0. - for i in range(bpg): - # Preload data into shared memory - sA[tx, ty] = A[x, ty + i * TPB] - sB[tx, ty] = B[tx + i * TPB, y] - - # Wait until all threads finish preloading - cuda.syncthreads() - - # Computes partial product on the shared memory - for j in range(TPB): - tmp += sA[tx, j] * sB[j, ty] - - # Wait until all threads finish computing - cuda.syncthreads() - - C[x, y] = tmp - - -# Translated to numba-cuda from https://github.com/cupy/cupy/blob/main/examples/gemm/sgemm.cu +# Translated to Numba-CUDA from https://github.com/cupy/cupy/blob/main/examples/gemm/sgemm.cu @cuda.jit(cache=True) def optimized_matmul(M, N, K, A, B, C): DIM_X = 16 @@ -68,8 +24,8 @@ def optimized_matmul(M, N, K, A, B, C): BLK_M = 64 BLK_N = 64 BLK_K = 16 - THR_M = 4 - THR_N = 4 + THR_M = 4 # Should be equal to BLK_M / DIM_X + THR_N = 4 # Should be equal to BLK_N / DIM_Y # thread indices idx = cuda.threadIdx.x @@ -144,12 +100,12 @@ def optimized_matmul(M, N, K, A, B, C): C[row, col] = np.float16(rC[n][m]) -def run_matmul(M, N, K): - +def run_basic(M, N, K): # create numpy arrays A = np.random.rand(M, K).astype(np.float16) B = np.random.rand(K, N).astype(np.float16) C = np.zeros((M, N), dtype=np.float16) + C_ref = A @ B # copy to GPU A_d = cuda.to_device(A) @@ -166,29 +122,52 @@ def run_matmul(M, N, K): ) # launch kernel - fast_matmul[blocks, threads](A_d, B_d, C_d) + matmul[blocks, threads](A_d, B_d, C_d) + cuda.synchronize() # copy result back C_result = C_d.copy_to_host() # check - np.testing.assert_allclose(C_result, A @ B, rtol=1e-2) + np.testing.assert_allclose(C_result, C_ref, rtol=1e-2, atol=M * 2**(-11)) + print("Succes") - print("Correct!") +def run_optimized(M, N, K): + # inputs + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) + C_ref = A @ B + + # move to GPU + dA = cuda.to_device(A) + dB = cuda.to_device(B) + dC = cuda.to_device(C) -def call_numba(kernel_function, args, kwargs, grid, threads): - numba_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - numba_args.append(cuda.as_cuda_array(arg)) - else: - numba_args.append(arg) - kernel_function[grid, threads](*args, **kwargs) + threads_per_block = (16, 16) + blocks_per_grid = ( + (M + 63) // 64, # BLK_M = 64 + (N + 63) // 64, # BLK_N = 64 + ) + # launch + for i in range(10000): + optimized_matmul[blocks_per_grid, threads_per_block](M, N, K, dA, dB, dC) + cuda.synchronize() -def tune(M, N, K): - # create numpy arrays + # copy result back + C_result = dC.copy_to_host() + + # check + np.testing.assert_allclose(C_result, C_ref, rtol=1e-2, atol=M * 2**(-11)) + print("Succes") + + + + +def tune_basic(M, N, K): + # create inputs as normal, but do not copy to device A = np.random.rand(M, K).astype(np.float16) B = np.random.rand(K, N).astype(np.float16) C = np.zeros((M, N), dtype=np.float16) @@ -196,80 +175,87 @@ def tune(M, N, K): size = (M, N) args = [A, B, C] tune_params = dict() - tune_params["block_size_x"] = [4, 8, 16, 32, 64, 128, 256] - tune_params["block_size_y"] = [4, 8, 16, 32, 64, 128, 256] - tune_params["TPB"] = [4, 8, 16, 32, 64, 128, 256] - restrictions = ["block_size_x == block_size_y", "block_size_x == TPB"] + tune_params["block_size_x"] = [2**i for i in range(1, 10)] + tune_params["block_size_y"] = [2**i for i in range(1, 10)] + + restrictions = ["block_size_x * block_size_y <= 1024"] answer = [None, None, A @ B] atol = M * 2**(-11) results, env = tune_kernel( - kernel_name="fast_matmul", - kernel_source=FULL_PATH, + kernel_name="matmul", + kernel_source=__file__, problem_size=size, arguments=args, tune_params=tune_params, lang="generic_python", answer=answer, atol=atol, + restrictions=restrictions, call_function=call_numba, - restrictions=restrictions ) -if __name__ == "__main__": +def tune_optimized(M, N, K): + # create inputs as normal, but do not copy to device + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) + + size = (M, N) + args = [M, N, K, A, B, C] + tune_params = { + "DIM_X": [2**i for i in range(1, 10)], + "DIM_Y": [2**i for i in range(1, 10)], + "BLK_M": [2**i for i in range(1, 10)], + "BLK_N": [2**i for i in range(1, 10)], + "BLK_K": [2**i for i in range(1, 10)], + "THR_M": [2**i for i in range(1, 9)], # Restricted to BLK_M / DIM_X + "THR_N": [2**i for i in range(1, 9)], # Restricted to BLK_N / DIM_Y + } - #run_matmul(128, 96, 64) - #tune(1024, 1024, 1024) + answer = [None, None, None, None, None, A @ B] + atol = M * 2**(-11) + restrictions = [ + "BLK_M % DIM_X == 0", + "BLK_N % DIM_Y == 0", + "THR_M == BLK_M / DIM_X", + "THR_N == BLK_N / DIM_Y", + "DIM_X * DIM_Y <= 1024", + "DIM_x * DIM_Y >= 32", + ] - M, N, K = 1024, 1024, 1024 + results, env = tune_kernel( + kernel_name="optimized_matmul", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=answer, + atol=atol, + restrictions = restrictions, + block_size_names = ["DIM_X", "DIM_Y"], + grid_div_x = ["BLK_M"], + grid_div_y = ["BLK_N"], + strategy = "bayes_opt", + strategy_options = {"max_fevals": 100}, + call_function=call_numba, + ) - # random FP16 inputs - A = np.random.rand(M, K).astype(np.float16) - B = np.random.rand(K, N).astype(np.float16) - # output - C = np.zeros((M, N), dtype=np.float16) - # copy to device - dA = cuda.to_device(A) - dB = cuda.to_device(B) - dC = cuda.to_device(C) +if __name__ == "__main__": + M, N, K = 1024, 1024, 1024 - # launch config - threads_per_block = (16, 16) + run_basic(M, N, K) + #tune_basic(M, N, K) - blocks_per_grid_x = (M + 64 - 1) // 64 - blocks_per_grid_y = (N + 64 - 1) // 64 - blocks_per_grid = (blocks_per_grid_x, blocks_per_grid_y) + #run_optimized(M, N, K) + #tune_optimized(M, N, K) - # run kernel - optimized_matmul[blocks_per_grid, threads_per_block](M, N, K, dA, dB, dC) - # copy result back - C_result = dC.copy_to_host() - # reference (FP32 accumulate like your kernel) - C_ref = (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) - - # check error - max_error = np.max(np.abs(C_result - C_ref)) - print("Max error:", max_error) - - # tolerance (FP16 is noisy) - if max_error < 4: - print("✅ Looks correct") - else: - print("❌ Something is off") - print("Expected: ", C_ref, "\nGot: ", C_result) - - ''' - 128: 16, 8 - 256: 8, 16 - 512: 8, 16 - 1024: 16, 8 - 4096: 16, 8 - 8192: 16, 8 - ''' \ No newline at end of file + \ No newline at end of file diff --git a/examples/generic_python/matmul/taichi_matmul.py b/examples/generic_python/matmul/taichi_matmul.py index b8a1c9388..bb3c53a81 100644 --- a/examples/generic_python/matmul/taichi_matmul.py +++ b/examples/generic_python/matmul/taichi_matmul.py @@ -1,21 +1,17 @@ -import torch +import numpy as np import taichi as ti from kernel_tuner import tune_kernel - - +from examples.generic_python.call_functions import call_taichi ti.init(arch=ti.gpu) - -# TODO make sure this is zero-copy +BLOCK_DIM = 512 @ti.kernel def matmul(A: ti.types.ndarray(dtype=ti.f16, ndim=2), B: ti.types.ndarray(dtype=ti.f16, ndim=2), C: ti.types.ndarray(dtype=ti.f16, ndim=2)): - BLOCK_DIM = 16 - ti.loop_config(block_dim=BLOCK_DIM) - K_dim = A.shape[1] + ti.loop_config(block_dim=BLOCK_DIM) for i, j in C: sum = 0.0 for k in range(K_dim): @@ -23,19 +19,28 @@ def matmul(A: ti.types.ndarray(dtype=ti.f16, ndim=2), B: ti.types.ndarray(dtype= C[i, j] = ti.cast(sum, ti.f16) -def call_taichi(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) +def run(M, N, K): + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) + C_ref = A @ B + + matmul(A, B, C) + + np.testing.assert_allclose(C, C_ref, rtol=1e-2, atol=M * 2**(-11)) + print("Succes") + def tune(M, N, K): - torch_A = torch.rand((N, K), device='cuda', dtype=torch.float16) - torch_B = torch.rand((K, M), device='cuda', dtype=torch.float16) - torch_C = torch.empty((N, M), device='cuda', dtype=torch.float16) + A = np.random.rand(M, K).astype(np.float16) + B = np.random.rand(K, N).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) size = M * N - args = [torch_A, torch_B, torch_C] - tune_params = {"BLOCK_DIM": {4, 8, 16, 32, 64, 128, 256, 512, 1024}} + args = [A, B, C] + tune_params = {"BLOCK_DIM": [2**i for i in range(5, 11)]} - answer = [None, None, (torch_A @ torch_B).cpu()] + answer = [None, None, A @ B] results, env = tune_kernel( kernel_name="matmul", @@ -46,11 +51,13 @@ def tune(M, N, K): answer=answer, lang="generic_python", call_function=call_taichi, - atol=1e-1, + atol=M * 2**(-11), + block_size_names = ["BLOCK_DIM"], ) if __name__ == "__main__": - tune(128, 128, 128) + run(1024, 1024, 1024) + tune(1024, 1024, 1024) diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py index 2ea510d91..1067d9ff5 100644 --- a/examples/generic_python/matmul/tilelang_matmul.py +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -4,7 +4,11 @@ import itertools +from kernel_tuner import tune_kernel +from examples.generic_python.call_functions import call_tilelang +# https://github.com/tile-ai/tilelang/tree/main/examples/gemm +# num_threads and num_stages added as variables to enable tuning. @tilelang.jit def matmul_basic(M:int, N:int, K:int, block_M:int, block_N:int, block_K:int, dtype:str="float16", accum_dtype:str="float32"): @T.prim_func @@ -13,7 +17,9 @@ def gemm( B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + num_threads = 128 + num_stages = 3 + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): # We do use shared memory, even though this is a basic kernel. However, you don't # really get around this because T.gemm can not handle global memory directly. A_shared = T.alloc_shared((block_M, block_K), dtype) @@ -23,7 +29,7 @@ def gemm( # We do use a pipelining optimization here, because this is 'the basic way' # of writing for loops in TileLang. - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): + for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[k * block_K, bx * block_N], B_shared) T.gemm(A_shared, B_shared, C_local) @@ -35,41 +41,50 @@ def gemm( -# https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py -# changed gemm_autotune to gemm -# originally, B was transposed. I changed this so the kernel input is the same as in other languages. +# https://github.com/tile-ai/tilelang/tree/main/examples/gemm +# num_threads and num_stages were added as metaparameters to make it possible to compare the built-in tuner +# with Kernel Tuner. +# Removed annotated memory layout for tiles A and B @tilelang.jit -def matmul_opt(M, N, K, dummy, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): +def matmul_opt(M, N, K, block_M, block_N, block_K, num_threads=128, num_stages=3, dummy=0, dtype=T.float16, accum_dtype=T.float): @T.prim_func - def gemm( + def main( A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): + # Allocate shared and local fragments A_shared = T.alloc_shared((block_M, block_K), dtype) B_shared = T.alloc_shared((block_K, block_N), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - C_shared = T.alloc_shared((block_M, block_N), dtype) - T.use_swizzle(panel_size=10, enable=enable_rasteration) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + + # Enable swizzle-based rasterization for better L2 locality + T.use_swizzle(panel_size=10, enable=True) + + # Clear the local accumulation buffer T.clear(C_local) - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[k * block_K, bx * block_N], B_shared) - T.gemm( - A_shared, - B_shared, - C_local, - transpose_B=False, - ) - T.copy(C_local, C_shared) - T.copy(C_shared, C[by * block_M, bx * block_N]) - return gemm + # Pipelined iteration over K dimension + for idx in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + # Copy tile of A + T.copy(A[by * block_M, idx * block_K], A_shared) + # Parallel copy tile of B + for ko, j in T.Parallel(block_K, block_N): + B_shared[ko, j] = B[idx * block_K + ko, bx * block_N + j] + # Perform local GEMM on the shared-memory tiles + T.gemm(A_shared, B_shared, C_local) + # Copy the result tile back + T.copy(C_local, C[by * block_M, bx * block_N]) + return main + + + +# TODO def get_configs(kt=False): """ Generate a list of kernel tuning configuration dictionaries for a tiled matrix-multiply. @@ -130,6 +145,8 @@ def get_configs(kt=False): return configs + +# TODO # https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py # changed gemm_autotune to gemm # originally, B was transposed. I changed this so the kernel input is the same as in other languages. @@ -137,35 +154,35 @@ def get_configs(kt=False): @tilelang.jit def matmul_opt_autotune(M, N, K, dummy, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): @T.prim_func - def gemm_autotune( + def gemm( A: T.Tensor((M, K), dtype), - B: T.Tensor((K, N), dtype), + B: T.Tensor((N, K), dtype), C: T.Tensor((M, N), dtype), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) + B_shared = T.alloc_shared((block_N, block_K), dtype) C_local = T.alloc_fragment((block_M, block_N), accum_dtype) C_shared = T.alloc_shared((block_M, block_N), dtype) T.use_swizzle(panel_size=10, enable=enable_rasteration) T.clear(C_local) for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[k * block_K, bx * block_N], B_shared) + T.copy(B[bx * block_N, k * block_K], B_shared) T.gemm( A_shared, B_shared, C_local, - transpose_B=False, + transpose_B=True, ) T.copy(C_local, C_shared) T.copy(C_shared, C[by * block_M, bx * block_N]) - return gemm_autotune - + return gemm +# DEZE mag denk ik weg def main(): kernel = matmul_opt_autotune(1024, 1024, 1024) import torch @@ -184,5 +201,138 @@ def main(): +def run_basic(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + C_ref = A @ B + + kernel = matmul_basic( + M, N, K, + block_M=64, + block_N=64, + block_K=32, + ) + + kernel(A, B, C) + + assert torch.allclose(C, C_ref, atol=M * 2 **(-11), rtol=1e-2) + print("Succes") + + +def run_opt(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + C_ref = A @ B + + kernel = matmul_opt( + M, N, K, + dummy=0, + block_M=64, + block_N=64, + block_K=32, + num_stages=3, + num_threads=128, + ) + + kernel(A, B, C) + + assert torch.allclose(C, C_ref, atol=M * 2 **(-11), rtol=1e-2) + print("Succes") + + +def tune_basic(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + size = (M, N) + + + args = [A, B, C] + + tune_params = dict() + tune_params["block_M"] = [2**i for i in range(4, 10)] + tune_params["block_N"] = [2**i for i in range(4, 10)] + tune_params["block_K"] = [2**i for i in range(4, 10)] + tune_params["num_threads"] = [2**i for i in range(5, 11)] + tune_params["num_stages"] = [1, 2, 3, 4, 5] + tune_params["M"] = [M] + tune_params["N"] = [N] + tune_params["K"] = [K] + + restrictions = [ + "2 * num_stages * block_K * (block_M + block_N) <= 49152" + ] + + results, env = tune_kernel( + kernel_name="matmul_basic", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, C_ref.cpu()], + atol=M * 2**(-11), + strategy = "bayes_opt", + strategy_options = {"max_fevals": 100}, + call_function=call_tilelang, + verbose=True, + restrictions=restrictions, + ) + + +def tune_opt(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + size = (M, N) + + args = [A, B, C] + + tune_params = dict() + tune_params["block_M"] = [2**i for i in range(4, 10)] + tune_params["block_N"] = [2**i for i in range(4, 10)] + tune_params["block_K"] = [2**i for i in range(4, 10)] + tune_params["num_threads"] = [2**i for i in range(5, 11)] + tune_params["num_stages"] = [1, 2, 3, 4, 5] + tune_params["M"] = [M] + tune_params["N"] = [N] + tune_params["K"] = [K] + + restrictions = [ + "2 * num_stages * block_K * (block_M + block_N) <= 49152" + ] + + results, env = tune_kernel( + kernel_name="matmul_opt", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, C_ref.cpu()], + atol=M * 2**(-11), + strategy = "bayes_opt", + strategy_options = {"max_fevals": 100}, + call_function=call_tilelang, + verbose=True, + restrictions=restrictions, + ) + + + + if __name__ == "__main__": - main() + M, N, K = 4096, 4096, 4096 + #run_basic(M, N, K) + #run_opt(M, N, K) + + #tune_basic(M, N, K) + tune_opt(M, N, K) + + diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py index d893b9490..38f6523f9 100644 --- a/examples/generic_python/matmul/tilus_matmul.py +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -1,7 +1,12 @@ +import torch + import tilus from tilus import float16, float32, int32 from tilus.utils import cdiv +from kernel_tuner import tune_kernel +from examples.generic_python.call_functions import call_tilus + # This kernel is copied from the Tilus project: # https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_v0.py @@ -30,7 +35,8 @@ def __call__( cdiv(m_size, self.block_m), # the x dimension size of the grid cdiv(n_size, self.block_n), # the y dimension size of the grid ] - self.attrs.warps = 1 # the number of warps per thread block, must be a compile-time known integer + num_warps = 1 # added for tuning + self.attrs.warps = num_warps # the number of warps per thread block, must be a compile-time known integer # define two int32 variables to store the offsets of the m and n dimensions for the current thread block. offset_m: int32 = self.block_m * self.blockIdx.x @@ -70,30 +76,22 @@ def __call__( # This kernel is copied from the Tilus project: -# https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_+v5.py +# https://nvidia.github.io/tilus/stable/tutorials/matmul-ampere/matmul/matmul_v4.html # -# Original example: matmul_v5.py +# Original example: matmul_v4.py # Copyright (c) the Tilus authors # # Modifications in file: # - Removed auto-tuning decorators +# - Added default values (None) for the __init__ parameters class MatmulOpt(tilus.Script): - def __init__( - self, - num_warps=None, - block_m=None, - block_n=None, - block_k=None, - num_stages=None, - split_k_factor=None, - ): + def __init__(self, num_warps=None, block_m=None, block_n=None, block_k=None, num_stages=None): super().__init__() self.block_m = block_m self.block_n = block_n self.block_k = block_k self.num_warps = num_warps self.num_stages = num_stages - self.split_k_factor = split_k_factor def __call__( self, @@ -104,20 +102,9 @@ def __call__( b_ptr: ~float16, c_ptr: ~float16, ): - self.attrs.blocks = [ - cdiv(m_size, self.block_m), - cdiv(n_size, self.block_n), - self.split_k_factor, - ] + self.attrs.blocks = [cdiv(m_size, self.block_m), cdiv(n_size, self.block_n)] self.attrs.warps = self.num_warps - # the k_size for each thread block - block_k_size = ( - cdiv(cdiv(k_size, self.split_k_factor), self.block_k) * self.block_k - ) - start_offset_k = self.blockIdx.z * block_k_size - end_offset_k = min(start_offset_k + block_k_size, k_size) - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k offset_m: int32 = block_m * self.blockIdx.x offset_n: int32 = block_n * self.blockIdx.y @@ -129,7 +116,7 @@ def __call__( acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) for stage in range(self.num_stages - 1): - offset_k = start_offset_k + stage * self.block_k + offset_k = stage * self.block_k self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) self.copy_async_commit_group() @@ -139,9 +126,7 @@ def __call__( current_stage: int32 = 0 preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range( - start_offset_k, end_offset_k, block_k, unroll=self.num_stages - ): + for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): # computation for current tile a = self.load_shared(sa[current_stage]) b = self.load_shared(sb[current_stage]) @@ -149,17 +134,16 @@ def __call__( # preload the next tile of A and B into shared memory preload_offset_k = offset_k + (self.num_stages - 1) * block_k - if preload_offset_k < end_offset_k: - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], - ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], - ) + self.copy_async( + src=ga, + dst=sa[preload_stage], + offsets=[offset_m, preload_offset_k], + ) + self.copy_async( + src=gb, + dst=sb[preload_stage], + offsets=[preload_offset_k, offset_n], + ) self.copy_async_commit_group() # update the stage @@ -168,41 +152,124 @@ def __call__( self.copy_async_wait_group(n=self.num_stages - 2) self.sync() - # free the shared memory tensors for A and B self.free_shared(sa) self.free_shared(sb) - # cast the accumulator to float16 and change the register tensor's layout - sc = self.shared_tensor(dtype=float16, shape=[block_m, block_n]) casted_acc = self.cast(acc, dtype=float16) - self.store_shared(sc, casted_acc) - self.sync() - rc = self.load_shared(sc) - self.free_shared(sc) - - m_blocks, n_blocks = cdiv(m_size, block_m), cdiv(n_size, block_n) gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - if self.split_k_factor == 0: - self.store_global(gc, rc, offsets=[offset_m, offset_n]) - else: - semaphores = self.global_tensor( - dtype=int32, shape=[m_blocks, n_blocks], requires_clean=True - ) - semaphore: ~int32 = ~semaphores[self.blockIdx.x, self.blockIdx.y] - - # load and accumulate the partial result in global memory - if self.blockIdx.z > 0: - self.lock_semaphore(semaphore, value=self.blockIdx.z) - partial_rc = self.load_global( - gc, offsets=[offset_m, offset_n], shape=[block_m, block_n] - ) - self.add(rc, partial_rc, out=rc) - - # store the result to global memory and release the semaphore - self.store_global(gc, rc, offsets=[offset_m, offset_n]) - - # release the semaphore - self.sync() # we need to make sure the previous store_global is finished - self.release_semaphore( - semaphore, value=(self.blockIdx.z + 1) % self.split_k_factor - ) \ No newline at end of file + self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) + + + +def run_basic(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + matmul = MatmulBasic() + torch.cuda.synchronize() + matmul(M, N, K, A, B, C) + torch.cuda.synchronize() + + torch.testing.assert_close(C_ref, C, atol=M * 2**(-11), rtol=1e-2) + print("Succes") + + +def run_optmized(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + matmul = MatmulOpt(num_warps=4, block_m=64, block_n=64, block_k=16, num_stages=3) + torch.cuda.synchronize() + + matmul(M, N, K, A, B, C) + torch.cuda.synchronize() + + torch.testing.assert_close(C_ref, C, atol=M * 2**(-11), rtol=1e-2) + print("Succes") + + +def tune_basic(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + size = (M, N) + + args = [M, N, K, A, B, C] + + tune_params = dict() + tune_params["block_m"] = [16, 32, 64, 128] + tune_params["block_n"] = [16, 32, 64, 128] + tune_params["block_k"] = [16, 32, 64, 128] + tune_params["num_warps"] = [2, 4, 8, 16] + + + results, env = tune_kernel( + kernel_name="MatmulBasic", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, None, None, None, C_ref.cpu()], + atol=M * 2**(-11), + strategy = "bayes_opt", + strategy_options = {"max_fevals": 200}, + call_function=call_tilus, + ) + + +def tune_opt(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + size = (M, N) + + args = [M, N, K, A, B, C] + + tune_params = dict() + tune_params["block_m"] = [16, 32, 64, 128, 256] + tune_params["block_n"] = [16, 32, 64, 128, 256] + tune_params["block_k"] = [16, 32, 64, 128] + tune_params["num_warps"] = [2, 4, 8, 16] + tune_params["num_stages"] = [2, 3, 4, 5, 6] + + # Shared memory restriction + restrictions = ["2 * num_stages * block_k * (block_m + block_n) <= 49152"] + + results, env = tune_kernel( + kernel_name="MatmulOpt", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, None, None, None, C_ref.cpu()], + atol=M * 2**(-11), + restrictions=restrictions, + strategy = "bayes_opt", + strategy_options = {"max_fevals": 200}, + call_function=call_tilus, + ) + + + + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + #M, N, K = 8192, 8192, 8192 + #run_basic(M, N, K) + + run_optmized(M, N, K) + + #tune_basic(M, N, K) + + #tune_opt(M, N, K) \ No newline at end of file diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py index b127b0331..dc81c5bd0 100644 --- a/examples/generic_python/matmul/triton_matmul.py +++ b/examples/generic_python/matmul/triton_matmul.py @@ -2,19 +2,17 @@ import triton import triton.language as tl +from kernel_tuner import tune_kernel +from examples.generic_python.call_functions import call_triton @triton.jit def matmul_basic( - # Pointers - a_ptr, b_ptr, c_ptr, - # Matrix sizes - M, N, K, - # Strides - stride_am, stride_ak, + a_ptr, b_ptr, c_ptr, # Pointers + M, N, K, # Matrix sizes + stride_am, stride_ak, # Strides stride_bk, stride_bn, stride_cm, stride_cn, - # Tile sizes - BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, # Tile sizes BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, ): @@ -41,7 +39,7 @@ def matmul_basic( mask=(offs_am[:, None] < M) & (offs_k[None, :] < K - k * BLOCK_SIZE_K), other=0.0, ) - + b = tl.load( b_ptrs, mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) & (offs_bn[None, :] < N), @@ -160,3 +158,153 @@ def matmul_opt( tl.store(c_ptrs, c, mask=c_mask) + +def run_basic(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + BLOCK_SIZE_M = 64 + BLOCK_SIZE_N = 64 + BLOCK_SIZE_K = 32 + + grid = ( + triton.cdiv(M, BLOCK_SIZE_M), + triton.cdiv(N, BLOCK_SIZE_N), + ) + + matmul_basic[grid]( + A, B, C, + M, N, K, + A.stride(0), A.stride(1), + B.stride(0), B.stride(1), + C.stride(0), C.stride(1), + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + ) + + assert torch.allclose(C, C_ref, rtol=1e-2, atol= M * 2**(-11)) + + print("Passed") + + +def run_opt(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + BLOCK_SIZE_M = 64 + BLOCK_SIZE_N = 64 + BLOCK_SIZE_K = 32 + GROUP_SIZE_M = 32 + + grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),) + + matmul_opt[grid]( + A, B, C, + M, N, K, + A.stride(0), A.stride(1), + B.stride(0), B.stride(1), + C.stride(0), C.stride(1), + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=GROUP_SIZE_M, + ) + + assert torch.allclose(C, C_ref, rtol=1e-2, atol= M * 2**(-11)) + + print("Passed") + + +def tune_basic(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + size = (M, N) + + args = [A, B, C, + M, N, K, + A.stride(0), A.stride(1), + B.stride(0), B.stride(1), + C.stride(0), C.stride(1), + ] + + tune_params = dict() + tune_params["BLOCK_SIZE_M"] = [2**i for i in range(1, 10)] + tune_params["BLOCK_SIZE_N"] = [2**i for i in range(1, 10)] + tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 10)] # tl.dot requires K >= 16 + tune_params["num_warps"] = [1, 2, 4, 8, 16, 32] + tune_params["num_stages"] = [1, 2, 3, 4, 5] + + results, env = tune_kernel( + kernel_name="matmul_basic", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, C_ref.cpu(), None, None, None, None, None, None, None, None, None], + atol=M * 2**(-11), + block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_M"], + strategy = "bayes_opt", + strategy_options = {"max_fevals": 100}, + call_function=call_triton, + ) + + +def tune_opt(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.empty((M, N), device='cuda', dtype=torch.float16) + C_ref = A @ B + + size = M * N + + args = [A, B, C, + M, N, K, + A.stride(0), A.stride(1), + B.stride(0), B.stride(1), + C.stride(0), C.stride(1), + ] + + tune_params = dict() + tune_params["BLOCK_SIZE_M"] = [2**i for i in range(1, 10)] + tune_params["BLOCK_SIZE_N"] = [2**i for i in range(1, 10)] + tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 10)] # tl.dot requires K >= 16 + tune_params["GROUP_SIZE_M"] = [1, 2, 4, 8, 16, 32, 64] + tune_params["num_warps"] = [1, 2, 4, 8, 16, 32] + tune_params["num_stages"] = [1, 2, 3, 4, 5] + + results, env = tune_kernel( + kernel_name="matmul_opt", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + lang="generic_python", + answer=[None, None, C_ref.cpu(), None, None, None, None, None, None, None, None, None], + atol=M * 2**(-11), + block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_M"], + grid_div_x = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], + strategy = "bayes_opt", + strategy_options = {"max_fevals": 100}, + call_function=call_triton, + ) + + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + #M, N, K = 8192, 8192, 8192 + #run_basic(M, N, K) + #run_opt(M, N, K) + + + #tune_basic(M, N, K) + #tune_opt(M, N, K) diff --git a/examples/generic_python/matmul/warp_matmul.py b/examples/generic_python/matmul/warp_matmul.py index 78a90ce65..388ae3a1b 100644 --- a/examples/generic_python/matmul/warp_matmul.py +++ b/examples/generic_python/matmul/warp_matmul.py @@ -3,33 +3,54 @@ import warp as wp from kernel_tuner import tune_kernel -from pathlib import Path +from examples.generic_python.call_functions import call_warp wp.init() +wp.config.enable_backward = False + + +@wp.kernel() +def gemm( + A: wp.array2d(dtype=wp.float16), B: wp.array2d(dtype=wp.float16), C: wp.array2d(dtype=wp.float16) +): + i, j = wp.tid() + M = A.shape[0] + N = B.shape[1] + K = A.shape[1] + + if i >= M or j >= N: + return + + # compute dot product + sum = wp.float32(0.0) + for k in range(K): + sum += wp.float32(A[i, k]) * wp.float32(B[k, j]) + + # write result + C[i, j] = wp.float16(sum) -FULL_PATH = Path(__file__).resolve() # tile size -TILE_M = wp.constant(8) -TILE_N = wp.constant(4) -TILE_K = wp.constant(8) +TILE_M = 32 +TILE_N = 32 +TILE_K = 32 # num threads per-tile -TILE_THREADS = 64 +TILE_THREADS = 1024 # GEMM example from https://nvidia.github.io/warp/user_guide/tiles.html -@wp.kernel -def tile_gemm(A: wp.array2d(dtype=wp.float16), B: wp.array2d(dtype=wp.float16), C: wp.array2d(dtype=wp.float16)): - +@wp.kernel() +def tile_gemm( + A: wp.array2d(dtype=wp.float16), + B: wp.array2d(dtype=wp.float16), + C: wp.array2d(dtype=wp.float16) +): # output tile index i, j = wp.tid() sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32) - M = A.shape[0] - N = B.shape[1] K = A.shape[1] - count = (K + TILE_K - 1) // TILE_K for k in range(0, count): @@ -41,7 +62,7 @@ def tile_gemm(A: wp.array2d(dtype=wp.float16), B: wp.array2d(dtype=wp.float16), wp.tile_store(C, wp.tile_astype(sum, wp.float16), offset=(i*TILE_M, j*TILE_N), bounds_check=True) -def run_kernel_direct(M, N, K): +def run_gemm(M, N, K): rng = np.random.default_rng(42) A = rng.random((M, K)).astype(np.float16) B = rng.random((K, N)).astype(np.float16) @@ -51,36 +72,33 @@ def run_kernel_direct(M, N, K): B_wp = wp.array(B) C_wp = wp.array(C) - with wp.Tape() as tape: - wp.launch_tiled( - tile_gemm, - dim=((M + TILE_M - 1) // TILE_M, (N + TILE_N - 1) // TILE_N), - inputs=[A_wp, B_wp, C_wp], - block_dim=TILE_THREADS) + wp.launch(gemm, dim=(M, N), inputs=[A_wp, B_wp, C_wp]) - np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-1) + np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-2, atol=M * 2**(-11)) - print("Example matrix multiplication passed") + print("Succes") +def run_gemm_tiled(M, N, K): + rng = np.random.default_rng(42) + A = rng.random((M, K)).astype(np.float16) + B = rng.random((K, N)).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) -def call_warp(kernel_function, args, kwargs, grid, threads, params): - # Convert Torch tensors to Warp args - warp_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - warp_args.append(wp.from_torch(arg)) - else: - warp_args.append(arg) + A_wp = wp.array(A) + B_wp = wp.array(B) + C_wp = wp.array(C) - # launch kernel - with wp.Tape() as tape: - wp.launch_tiled( - kernel_function, - dim=grid, - inputs=warp_args, - block_dim=params["TILE_THREADS"], - ) + + wp.launch_tiled( + tile_gemm, + dim=((M + TILE_M - 1) // TILE_M, (N + TILE_N - 1) // TILE_N), + inputs=[A_wp, B_wp, C_wp], + block_dim=TILE_THREADS) + + np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-2, atol=M * 2**(-11)) + + print("Succes") def tune(M, K, N): @@ -90,43 +108,40 @@ def tune(M, K, N): C = np.zeros((M, N), dtype=np.float16) size = (M, N) - block_size_names = ["TILE_M", "TILE_N"] tune_params = dict() - tune_params["TILE_M"] = [4, 8, 16] - tune_params["TILE_N"] = [2, 4, 8] - tune_params["TILE_K"] = [4, 8, 16] - tune_params["TILE_THREADS"] = [32, 64, 128] + tune_params["block_dim"] = [2**i for i in range(5, 11)] + tune_params["dim"] = [size] args = [A, B, C] answer = [None, None, A @ B] results, env = tune_kernel( - kernel_name="tile_gemm", - kernel_source=FULL_PATH, + kernel_name="gemm", + kernel_source=__file__, problem_size=size, arguments=args, tune_params=tune_params, lang="generic_python", answer=answer, call_function=call_warp, - block_size_names=block_size_names, ) + if __name__ == "__main__": - #tune(128, 128, 128) + M, N, K = 1024, 1024, 1024 + + #run_gemm(M, N, K) + #run_gemm_tiled(M, N, K) + + tune(M, N, K) + + + - sizes = [ - (65, 65, 17), - (67, 71, 19), - (1, 1, 1), - (63, 63, 15), - (129, 130, 33), - ] + - #for size in sizes: - # print(size) - # run_kernel_direct(*size) + - run_kernel_direct(1024, 1024, 1024) + diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 4845fc4e2..2247d5e87 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -41,11 +41,11 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None self.device_id = torch.cuda.current_device() self.device_properties = torch.cuda.get_device_properties(self.device_id) self.name = torch.cuda.get_device_name(self.device_id) - self.max_threads = self.device_properties.max_threads_per_multi_processor + self.max_threads = 10**18 # 'inf' to support tile based programming models, which can use less threads then the size of a tile. env = dict() env["device_name"] = self.name - env["max_threads"] = self.max_threads + env["max_threads"] = self.max_threads env["iterations"] = iterations env["compiler_options"] = compiler_options self.env = env From 3cccf50215a178d7b58d239339d53b9563a452c5 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Sat, 6 Jun 2026 10:56:49 +0200 Subject: [PATCH 11/44] modified examples --- examples/generic_python/matmul/cute_matmul.py | 20 +-- .../generic_python/matmul/numba_matmul.py | 5 +- .../generic_python/matmul/tilelang_matmul.py | 144 ++++++++---------- .../generic_python/matmul/tilus_matmul.py | 36 ++--- .../generic_python/matmul/triton_matmul.py | 22 +-- 5 files changed, 105 insertions(+), 122 deletions(-) diff --git a/examples/generic_python/matmul/cute_matmul.py b/examples/generic_python/matmul/cute_matmul.py index 4b5637c82..09ced866e 100644 --- a/examples/generic_python/matmul/cute_matmul.py +++ b/examples/generic_python/matmul/cute_matmul.py @@ -15,14 +15,7 @@ ## Basic Matmul ================================================================ @cute.kernel -def naive_matmul_kernel( - gA: cute.Tensor, - gB: cute.Tensor, - gC: cute.Tensor, - M: int, - K: int, - N: int, -): +def naive_matmul_kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor, M: int, K: int, N: int): tx, ty, _ = cute.arch.thread_idx() bx, by, _ = cute.arch.block_idx() bdx, bdy, _ = cute.arch.block_dim() @@ -95,13 +88,22 @@ def tune_naive_matmul(M, N, K): tune_params = dict() tune_params["block_size_x"] = [2**i for i in range(1, 10)] tune_params["block_size_y"] = [2**i for i in range(1, 10)] - restrictions = ["block_size_x * block_size_y <= 1024"] + restrictions = ["block_size_x * block_size_y >= 32, block_size_x * block_size_y <= 1024"] results, env = tune_kernel("matmul", __file__, size, args, tune_params, lang="generic_python", call_function=call_cute, answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) +# This kernel is copied from the CuTe DSL project: +# https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/ampere/tensorop_gemm.py +# +# Original example: tensorop_gemm.py +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Modifications in file: +# - Tuning paramters made more explicit + ## Optimized matmul ========================================================== # Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause diff --git a/examples/generic_python/matmul/numba_matmul.py b/examples/generic_python/matmul/numba_matmul.py index d8e23d6fe..6486ec540 100644 --- a/examples/generic_python/matmul/numba_matmul.py +++ b/examples/generic_python/matmul/numba_matmul.py @@ -152,9 +152,8 @@ def run_optimized(M, N, K): ) # launch - for i in range(10000): - optimized_matmul[blocks_per_grid, threads_per_block](M, N, K, dA, dB, dC) - cuda.synchronize() + optimized_matmul[blocks_per_grid, threads_per_block](M, N, K, dA, dB, dC) + cuda.synchronize() # copy result back C_result = dC.copy_to_host() diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py index 1067d9ff5..0c476c786 100644 --- a/examples/generic_python/matmul/tilelang_matmul.py +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -40,11 +40,16 @@ def gemm( - -# https://github.com/tile-ai/tilelang/tree/main/examples/gemm -# num_threads and num_stages were added as metaparameters to make it possible to compare the built-in tuner -# with Kernel Tuner. -# Removed annotated memory layout for tiles A and B +# This kernel is copied from the TileLang project: +# https://github.com/tile-ai/tilelang/tree/main/examples/gemm/example_gemm.py +# +# Original example: example_gemm.py +# Copyright (c) the TileLang authors +# +# Modifications in file: +# - num_threads and num_stages added as metaparameters +# - Removed annotated memory layout for tiles A and B +# - dummy parameter to trigger fresh compilation when timing repeated tuning @tilelang.jit def matmul_opt(M, N, K, block_M, block_N, block_K, num_threads=128, num_stages=3, dummy=0, dtype=T.float16, accum_dtype=T.float): @T.prim_func @@ -60,7 +65,8 @@ def main( C_local = T.alloc_fragment((block_M, block_N), accum_dtype) # Enable swizzle-based rasterization for better L2 locality - T.use_swizzle(panel_size=10, enable=True) + panel_size = 10 + T.use_swizzle(panel_size=panel_size, enable=True) # Clear the local accumulation buffer T.clear(C_local) @@ -84,7 +90,7 @@ def main( -# TODO + def get_configs(kt=False): """ Generate a list of kernel tuning configuration dictionaries for a tiled matrix-multiply. @@ -94,30 +100,22 @@ def get_configs(kt=False): List[dict]: A list of configuration dictionaries """ - ''' - block_M = [64, 128, 256] - block_N = [64, 128, 256] - block_K = [32, 64] - num_stages = [0, 1, 2, 3] - thread_num = [128, 256] - enable_rasterization = [True, False] - ''' - - block_M = [64] - block_N = [64] - block_K = [32, 64] - num_stages = [0, 3] - thread_num = [128] - enable_rasterization = [True] + block_M = [32, 64, 128, 256] + block_N = [32, 64, 128, 256] + block_K = [16, 32, 64, 128] + num_threads = [256]#[64, 128, 256, 512] + num_stages = [3]#[0, 1, 2, 3, 4] + panel_size = [10]#[4, 6, 8, 10] + _configs = list( itertools.product( block_M, block_N, block_K, + num_threads, num_stages, - thread_num, - enable_rasterization, + panel_size, ) ) @@ -127,8 +125,8 @@ def get_configs(kt=False): "block_N": block_N, "block_K": block_K, "num_stages": num_stages, - "thread_num": thread_num, - "enable_rasteration": enable_rasterization, + "num_threads": num_threads, + "panel_size": panel_size, } else: configs = [ @@ -136,9 +134,9 @@ def get_configs(kt=False): "block_M": c[0], "block_N": c[1], "block_K": c[2], - "num_stages": c[3], - "thread_num": c[4], - "enable_rasteration": c[5], # keep param name for backward-compat + "num_threads": c[3], + "num_stages": c[4], + "panel_size": c[5], } for c in _configs ] @@ -146,58 +144,45 @@ def get_configs(kt=False): -# TODO +# For autotuning experiment # https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py -# changed gemm_autotune to gemm -# originally, B was transposed. I changed this so the kernel input is the same as in other languages. -@tilelang.autotune(configs=get_configs()) +@tilelang.autotune(configs=get_configs(), warmup=1, rep=32) @tilelang.jit -def matmul_opt_autotune(M, N, K, dummy, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): +def matmul_opt_autotune(M, N, K, block_M, block_N, block_K, num_threads=128, num_stages=3, panel_size=10, dummy=0, dtype=T.float16, accum_dtype=T.float): @T.prim_func - def gemm( + def main( A: T.Tensor((M, K), dtype), - B: T.Tensor((N, K), dtype), + B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): + # Allocate shared and local fragments A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_N, block_K), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - C_shared = T.alloc_shared((block_M, block_N), dtype) - T.use_swizzle(panel_size=10, enable=enable_rasteration) - T.clear(C_local) - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[bx * block_N, k * block_K], B_shared) - T.gemm( - A_shared, - B_shared, - C_local, - transpose_B=True, - ) - T.copy(C_local, C_shared) - T.copy(C_shared, C[by * block_M, bx * block_N]) - - return gemm - + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + # Enable swizzle-based rasterization for better L2 locality + T.use_swizzle(panel_size=panel_size, enable=True) -# DEZE mag denk ik weg -def main(): - kernel = matmul_opt_autotune(1024, 1024, 1024) - import torch + # Clear the local accumulation buffer + T.clear(C_local) - a = torch.randn(1024, 1024).cuda().half() - b = torch.randn(1024, 1024).cuda().half() - c = torch.empty((1024, 1024), device='cuda', dtype=torch.float16) + # Pipelined iteration over K dimension + for idx in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): + # Copy tile of A + T.copy(A[by * block_M, idx * block_K], A_shared) - kernel(a, b, c) + # Parallel copy tile of B + for ko, j in T.Parallel(block_K, block_N): + B_shared[ko, j] = B[idx * block_K + ko, bx * block_N + j] - ref_c = a @ b + # Perform local GEMM on the shared-memory tiles + T.gemm(A_shared, B_shared, C_local) - torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2) - print("All check passed.") + # Copy the result tile back + T.copy(C_local, C[by * block_M, bx * block_N]) + return main @@ -254,11 +239,11 @@ def tune_basic(M, N, K): args = [A, B, C] tune_params = dict() - tune_params["block_M"] = [2**i for i in range(4, 10)] - tune_params["block_N"] = [2**i for i in range(4, 10)] - tune_params["block_K"] = [2**i for i in range(4, 10)] - tune_params["num_threads"] = [2**i for i in range(5, 11)] - tune_params["num_stages"] = [1, 2, 3, 4, 5] + tune_params["block_M"] = [2**i for i in range(5, 9)] + tune_params["block_N"] = [2**i for i in range(5, 9)] + tune_params["block_K"] = [2**i for i in range(4, 8)] + tune_params["num_threads"] = [64, 128, 256, 512] + tune_params["num_stages"] = [0, 1, 2, 3, 4] tune_params["M"] = [M] tune_params["N"] = [N] tune_params["K"] = [K] @@ -295,11 +280,12 @@ def tune_opt(M, N, K): args = [A, B, C] tune_params = dict() - tune_params["block_M"] = [2**i for i in range(4, 10)] - tune_params["block_N"] = [2**i for i in range(4, 10)] - tune_params["block_K"] = [2**i for i in range(4, 10)] - tune_params["num_threads"] = [2**i for i in range(5, 11)] - tune_params["num_stages"] = [1, 2, 3, 4, 5] + tune_params["block_M"] = [2**i for i in range(5, 9)] + tune_params["block_N"] = [2**i for i in range(5, 9)] + tune_params["block_K"] = [2**i for i in range(4, 8)] + tune_params["num_threads"] = [64, 128, 256, 512] + tune_params["num_stages"] = [0, 1, 2, 3, 4] + tune_params["panel_size"] = [4, 6, 8, 10] tune_params["M"] = [M] tune_params["N"] = [N] tune_params["K"] = [K] @@ -332,7 +318,7 @@ def tune_opt(M, N, K): #run_basic(M, N, K) #run_opt(M, N, K) - #tune_basic(M, N, K) - tune_opt(M, N, K) + tune_basic(M, N, K) + #tune_opt(M, N, K) diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py index 38f6523f9..811d5c2d7 100644 --- a/examples/generic_python/matmul/tilus_matmul.py +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -24,12 +24,8 @@ def __init__(self): def __call__( self, - m_size: int32, # the size of the m dimension of the input matrix A and output matrix C - n_size: int, # the size of the n dimension of the input matrix B and output matrix C - k_size: int, # the size of the k dimension of the input matrix A and B - a_ptr: ~float16, # the pointer to the input matrix A, which is a 2D tensor of shape [m_size, k_size] - b_ptr: ~float16, # the pointer to the input matrix B, which is a 2D tensor of shape [k_size, n_size] - c_ptr: ~float16, # the pointer to the output matrix C, which is a 2D tensor of shape [m_size, n_size] + m_size: int32, n_size: int, k_size: int, # Matrix dimensions + a_ptr: ~float16, b_ptr: ~float16, c_ptr: ~float16, # Matrix pointers ): self.attrs.blocks = [ cdiv(m_size, self.block_m), # the x dimension size of the grid @@ -95,12 +91,8 @@ def __init__(self, num_warps=None, block_m=None, block_n=None, block_k=None, num def __call__( self, - m_size: int32, - n_size: int, - k_size: int, - a_ptr: ~float16, - b_ptr: ~float16, - c_ptr: ~float16, + m_size: int32, n_size: int, k_size: int, + a_ptr: ~float16, b_ptr: ~float16, c_ptr: ~float16, ): self.attrs.blocks = [cdiv(m_size, self.block_m), cdiv(n_size, self.block_n)] self.attrs.warps = self.num_warps @@ -176,7 +168,7 @@ def run_basic(M, N, K): print("Succes") -def run_optmized(M, N, K): +def run_optimized(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) C = torch.empty((M, N), device='cuda', dtype=torch.float16) @@ -203,9 +195,9 @@ def tune_basic(M, N, K): args = [M, N, K, A, B, C] tune_params = dict() - tune_params["block_m"] = [16, 32, 64, 128] - tune_params["block_n"] = [16, 32, 64, 128] - tune_params["block_k"] = [16, 32, 64, 128] + tune_params["block_m"] = [2**i for i in range(5, 9)] + tune_params["block_n"] = [2**i for i in range(5, 9)] + tune_params["block_k"] = [2**i for i in range(4, 8)] tune_params["num_warps"] = [2, 4, 8, 16] @@ -235,11 +227,11 @@ def tune_opt(M, N, K): args = [M, N, K, A, B, C] tune_params = dict() - tune_params["block_m"] = [16, 32, 64, 128, 256] - tune_params["block_n"] = [16, 32, 64, 128, 256] - tune_params["block_k"] = [16, 32, 64, 128] + tune_params["block_m"] = [2**i for i in range(5, 9)] + tune_params["block_n"] = [2**i for i in range(5, 9)] + tune_params["block_k"] = [2**i for i in range(4, 8)] tune_params["num_warps"] = [2, 4, 8, 16] - tune_params["num_stages"] = [2, 3, 4, 5, 6] + tune_params["num_stages"] = [2, 3, 4, 5] # Shared memory restriction restrictions = ["2 * num_stages * block_k * (block_m + block_n) <= 49152"] @@ -268,8 +260,8 @@ def tune_opt(M, N, K): #M, N, K = 8192, 8192, 8192 #run_basic(M, N, K) - run_optmized(M, N, K) + #run_optmized(M, N, K) - #tune_basic(M, N, K) + tune_basic(M, N, K) #tune_opt(M, N, K) \ No newline at end of file diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py index dc81c5bd0..cac3cab73 100644 --- a/examples/generic_python/matmul/triton_matmul.py +++ b/examples/generic_python/matmul/triton_matmul.py @@ -236,10 +236,10 @@ def tune_basic(M, N, K): ] tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [2**i for i in range(1, 10)] - tune_params["BLOCK_SIZE_N"] = [2**i for i in range(1, 10)] - tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 10)] # tl.dot requires K >= 16 - tune_params["num_warps"] = [1, 2, 4, 8, 16, 32] + tune_params["BLOCK_SIZE_M"] = [2**i for i in range(5, 9)] + tune_params["BLOCK_SIZE_N"] = [2**i for i in range(5, 9)] + tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 8)] + tune_params["num_warps"] = [2, 4, 8, 16] tune_params["num_stages"] = [1, 2, 3, 4, 5] results, env = tune_kernel( @@ -274,12 +274,15 @@ def tune_opt(M, N, K): ] tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [2**i for i in range(1, 10)] - tune_params["BLOCK_SIZE_N"] = [2**i for i in range(1, 10)] - tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 10)] # tl.dot requires K >= 16 - tune_params["GROUP_SIZE_M"] = [1, 2, 4, 8, 16, 32, 64] - tune_params["num_warps"] = [1, 2, 4, 8, 16, 32] + tune_params["BLOCK_SIZE_M"] = [2**i for i in range(5, 9)] + tune_params["BLOCK_SIZE_N"] = [2**i for i in range(5, 9)] + tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 8)] + tune_params["GROUP_SIZE_M"] = [4, 6, 8, 10] + tune_params["num_warps"] = [2, 4, 8, 16] tune_params["num_stages"] = [1, 2, 3, 4, 5] + tune_params["M_dim"] = [M] + + restrictions = ["GROUP_SIZE_M * BLOCK_SIZE_M < M_dim"] # Otherwise grouped ordering has no effect results, env = tune_kernel( kernel_name="matmul_opt", @@ -292,6 +295,7 @@ def tune_opt(M, N, K): atol=M * 2**(-11), block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_M"], grid_div_x = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], + restrictions = restrictions, strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, call_function=call_triton, From 68afbc1c22dba73010ddbb6a40094e34ce2fe0a6 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Sat, 6 Jun 2026 11:09:54 +0200 Subject: [PATCH 12/44] Restore workflow file --- .github/workflows/test-python-package.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-python-package.yml b/.github/workflows/test-python-package.yml index f86240c7f..dec26fdba 100644 --- a/.github/workflows/test-python-package.yml +++ b/.github/workflows/test-python-package.yml @@ -21,7 +21,7 @@ jobs: strategy: matrix: - os: [ubuntu-latest, macos-latest] + os: [ubuntu-latest, macos-13] steps: - uses: actions/checkout@v4 From 9d11e5e927ab038193a6c578dc333ffc12c41343 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Sat, 6 Jun 2026 11:33:31 +0200 Subject: [PATCH 13/44] fixed failed tests --- kernel_tuner/core.py | 2 +- test/test_observers.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index aa9cada75..1fe5ffbe6 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -400,6 +400,7 @@ def check_kernel_output( ): """Runs the kernel once and checks the result against answer.""" logging.debug("check_kernel_output") + cp = _get_cupy() # if not using custom verify function, check if the length is the same if answer: @@ -408,7 +409,6 @@ def check_kernel_output( should_sync = [answer[i] is not None for i, arg in enumerate(instance.arguments)] else: - cp = _get_cupy() cupy_ndarray = (cp.ndarray,) if cp is not None else () should_sync = [ isinstance(arg, (np.ndarray, torch.Tensor, DeviceArray) + cupy_ndarray) diff --git a/test/test_observers.py b/test/test_observers.py index 69a94a9bb..d1a27179c 100644 --- a/test/test_observers.py +++ b/test/test_observers.py @@ -4,6 +4,7 @@ from kernel_tuner.observers.nvml import NVMLObserver from kernel_tuner.observers.observer import BenchmarkObserver from kernel_tuner.observers.register import RegisterObserver +from kernel_tuner.language import Language from .context import ( skip_if_no_cuda, @@ -59,7 +60,7 @@ def get_results(self): result, _ = kernel_tuner.tune_kernel(*env_compiler, observers=[lambda args: MyObserver(args)], compiler_options=["-fopenmp"]) # Check if the observer has correctly received the lang option - assert result[0]["observer_args"]["lang"] == "C" + assert result[0]["observer_args"]["lang"] == Language.C @skip_if_no_pycuda def test_register_observer_pycuda(env): From 221c912a06bb177608f5949ef98b00f600c3d6d2 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Thu, 18 Jun 2026 10:40:53 +0200 Subject: [PATCH 14/44] finalized examples --- examples/generic_python/cupy_copy.py | 19 - examples/generic_python/cute_vec_add.py | 39 +- examples/generic_python/cutile_vec_add.py | 62 -- .../flash_attention/tilelang_attention.py | 152 ----- .../flash_attention/tilus_attention.py | 611 ------------------ .../flash_attention/triton_attention.py | 276 -------- examples/generic_python/loopy_example.py | 38 -- examples/generic_python/matmul/cute_matmul.py | 8 +- .../generic_python/matmul/helion_matmul.py | 20 - .../generic_python/matmul/numba_matmul.py | 6 +- examples/generic_python/matmul/test.py | 278 -------- examples/generic_python/matmul/test_tilus.py | 26 - .../generic_python/matmul/tilelang_matmul.py | 6 +- .../generic_python/matmul/tilus_matmul.py | 10 +- .../generic_python/matmul/triton_matmul.py | 10 +- examples/generic_python/matmul/warp_matmul.py | 59 +- .../matmul_old/tilelang_matmul.py | 212 ------ .../generic_python/matmul_old/tilus_matmul.py | 266 -------- .../matmul_old/triton_matmul.py | 383 ----------- .../normalization/tilelang_norm.py | 78 --- .../normalization/tilus_norm.py | 132 ---- .../normalization/triton_norm.py | 66 -- examples/generic_python/numba_vec_add.py | 18 +- examples/generic_python/pallas_vec_add.py | 23 - examples/generic_python/tilelang_vec_add.py | 23 +- examples/generic_python/tilus_naive_matmul.py | 127 ---- .../generic_python/tilus_splitk_matmul.py | 255 -------- examples/generic_python/tilus_vec_add.py | 35 +- examples/generic_python/triton_vec_add.py | 13 +- examples/generic_python/warp_vec_add.py | 10 +- 30 files changed, 46 insertions(+), 3215 deletions(-) delete mode 100644 examples/generic_python/cupy_copy.py delete mode 100644 examples/generic_python/cutile_vec_add.py delete mode 100644 examples/generic_python/flash_attention/tilelang_attention.py delete mode 100644 examples/generic_python/flash_attention/tilus_attention.py delete mode 100644 examples/generic_python/flash_attention/triton_attention.py delete mode 100644 examples/generic_python/loopy_example.py delete mode 100644 examples/generic_python/matmul/helion_matmul.py delete mode 100644 examples/generic_python/matmul/test.py delete mode 100644 examples/generic_python/matmul/test_tilus.py delete mode 100644 examples/generic_python/matmul_old/tilelang_matmul.py delete mode 100644 examples/generic_python/matmul_old/tilus_matmul.py delete mode 100644 examples/generic_python/matmul_old/triton_matmul.py delete mode 100644 examples/generic_python/normalization/tilelang_norm.py delete mode 100644 examples/generic_python/normalization/tilus_norm.py delete mode 100644 examples/generic_python/normalization/triton_norm.py delete mode 100644 examples/generic_python/pallas_vec_add.py delete mode 100644 examples/generic_python/tilus_naive_matmul.py delete mode 100644 examples/generic_python/tilus_splitk_matmul.py diff --git a/examples/generic_python/cupy_copy.py b/examples/generic_python/cupy_copy.py deleted file mode 100644 index 8d9110209..000000000 --- a/examples/generic_python/cupy_copy.py +++ /dev/null @@ -1,19 +0,0 @@ -import cupy -from cupyx import jit - -@jit.rawkernel() -def elementwise_copy(x, y, size): - tid = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x - ntid = jit.gridDim.x * jit.blockDim.x - for i in range(tid, size, ntid): - y[i] = x[i] - -size = cupy.uint32(2 ** 22) -x = cupy.random.normal(size=(size,), dtype=cupy.float32) -y = cupy.empty((size,), dtype=cupy.float32) - -elementwise_copy((128,), (1024,), (x, y, size)) # RawKernel style -assert (x == y).all() - -elementwise_copy[128, 1024](x, y, size) # Numba style -assert (x == y).all() \ No newline at end of file diff --git a/examples/generic_python/cute_vec_add.py b/examples/generic_python/cute_vec_add.py index a83694ab9..2744abb77 100644 --- a/examples/generic_python/cute_vec_add.py +++ b/examples/generic_python/cute_vec_add.py @@ -1,13 +1,10 @@ import torch -from functools import partial -from typing import List -import time import cutlass import cutlass.cute as cute -from cutlass.cute.runtime import from_dlpack from kernel_tuner import tune_kernel +from call_functions import call_cute @cute.kernel def vec_add_kernel( @@ -25,7 +22,6 @@ def vec_add_kernel( gC[thread_id] = gA[thread_id] + gB[thread_id] - @cute.jit def vec_add( mA: cute.Tensor, @@ -43,50 +39,19 @@ def vec_add( ) - - -def call_cute(kernel_function, args, kwargs, grid, threads, params): - cute_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - arg_ = from_dlpack(arg) - cute_args.append(arg_) - else: - cute_args.append(arg) - - kernel_function(*cute_args, **kwargs) - - - - def main(): size = 16384 a = torch.randn(size, device="cuda", dtype=torch.float16) b = torch.randn(size, device="cuda", dtype=torch.float16) c = torch.zeros(size, device="cuda", dtype=torch.float16) - ''' - a_ = from_dlpack(a, assumed_align=16) - b_ = from_dlpack(b, assumed_align=16) - c_ = from_dlpack(c, assumed_align=16) - ''' - args = [a, b, c, size] tune_params = {"num_threads_per_block": [1, 2, 4, 8, 16, 32, 64, 128, 265, 512, 1024]} answer = [None, None, (a+b).cpu(), None] - from pathlib import Path - FULL_PATH = Path(__file__).resolve() - tune_kernel("vec_add", FULL_PATH, size, args, tune_params, answer=answer, + tune_kernel("vec_add", __file__, size, args, tune_params, answer=answer, lang="generic_python", call_function=call_cute, verbose=True) - #naive_elementwise_add_ = cute.compile(naive_elementwise_add, a_, b_, c_) - - #naive_elementwise_add_(a_, b_, c_) - #vec_add(a_, b_, c_, size) - - - #torch.testing.assert_close(c, a+b) main() diff --git a/examples/generic_python/cutile_vec_add.py b/examples/generic_python/cutile_vec_add.py deleted file mode 100644 index e95302045..000000000 --- a/examples/generic_python/cutile_vec_add.py +++ /dev/null @@ -1,62 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# SPDX-License-Identifier: Apache-2.0 - -""" -Example demonstrating simple vector addition. -Shows how to perform elementwise operations on vectors. -Does not work on das vu, because we need cuda 13.1 -""" - -import cupy as cp -import numpy as np -import cuda.tile as ct - - -@ct.kernel -def vector_add(a, b, c, tile_size: ct.Constant[int]): - # Get the 1D pid - pid = ct.bid(0) - - # Load input tiles - a_tile = ct.load(a, index=(pid,), shape=(tile_size,)) - b_tile = ct.load(b, index=(pid,), shape=(tile_size,)) - - # Perform elementwise addition - result = a_tile + b_tile - - # Store result - ct.store(c, index=(pid, ), tile=result) - - -def test(): - # Create input data - vector_size = 2**12 - tile_size = 2**4 - grid = (ct.cdiv(vector_size, tile_size), 1, 1) - - rng = cp.random.default_rng() - a = rng.random(vector_size) - b = rng.random(vector_size) - c = cp.zeros_like(a) - - # Launch kernel - ct.launch(cp.cuda.get_current_stream(), - grid, # 1D grid of processors - vector_add, - (a, b, c, tile_size)) - - # Copy to host only to compare - a_np = cp.asnumpy(a) - b_np = cp.asnumpy(b) - c_np = cp.asnumpy(c) - - # Verify results - expected = a_np + b_np - np.testing.assert_array_almost_equal(c_np, expected) - - print("✓ vector_add_example passed!") - - -if __name__ == "__main__": - test() \ No newline at end of file diff --git a/examples/generic_python/flash_attention/tilelang_attention.py b/examples/generic_python/flash_attention/tilelang_attention.py deleted file mode 100644 index 5b2be204a..000000000 --- a/examples/generic_python/flash_attention/tilelang_attention.py +++ /dev/null @@ -1,152 +0,0 @@ -# example taken from https://github.com/tile-ai/tilelang/blob/main/examples/flash_attention/example_mha_fwd_bshd.py -import torch -import torch.nn.functional as F -import tilelang -from tilelang.autotuner import * -import tilelang.language as T -import itertools -import argparse -from functools import partial - - -def get_configs(): - iter_params = dict(block_M=[64], block_N=[64], num_stages=[1], threads=[128]) - return [dict(zip(iter_params, values)) for values in itertools.product(*iter_params.values())] - - -@autotune(configs=get_configs(), warmup=10, rep=10) -@tilelang.jit( - out_idx=[3], - pass_configs={ - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, - }, -) -def flashattn(batch, heads, seq_len, dim, is_causal, block_M=64, block_N=64, num_stages=1, threads=128): - scale = (1.0 / dim) ** 0.5 * 1.44269504 # log2(e) - shape = [batch, seq_len, heads, dim] - dtype = T.float16 - accum_dtype = T.float32 - - @T.prim_func - def main( - Q: T.Tensor(shape, dtype), - K: T.Tensor(shape, dtype), - V: T.Tensor(shape, dtype), - Output: T.Tensor(shape, dtype), - ): - with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=threads) as (bx, by, bz): - Q_shared = T.alloc_shared([block_M, dim], dtype) - K_shared = T.alloc_shared([block_N, dim], dtype) - V_shared = T.alloc_shared([block_N, dim], dtype) - O_shared = T.alloc_shared([block_M, dim], dtype) - acc_s = T.alloc_fragment([block_M, block_N], accum_dtype) - acc_s_cast = T.alloc_fragment([block_M, block_N], dtype) - acc_o = T.alloc_fragment([block_M, dim], accum_dtype) - scores_max = T.alloc_fragment([block_M], accum_dtype) - scores_max_prev = T.alloc_fragment([block_M], accum_dtype) - scores_scale = T.alloc_fragment([block_M], accum_dtype) - scores_sum = T.alloc_fragment([block_M], accum_dtype) - logsum = T.alloc_fragment([block_M], accum_dtype) - - T.copy(Q[bz, bx * block_M : (bx + 1) * block_M, by, :], Q_shared) - T.fill(acc_o, 0) - T.fill(logsum, 0) - T.fill(scores_max, -T.infinity(accum_dtype)) - - loop_range = ( - T.min(T.ceildiv(seq_len, block_N), T.ceildiv((bx + 1) * block_M, block_N)) if is_causal else T.ceildiv(seq_len, block_N) - ) - - for k in T.Pipelined(loop_range, num_stages=num_stages): - T.copy(K[bz, k * block_N : (k + 1) * block_N, by, :], K_shared) - if is_causal: - for i, j in T.Parallel(block_M, block_N): - acc_s[i, j] = T.if_then_else(bx * block_M + i >= k * block_N + j, 0, -T.infinity(acc_s.dtype)) - else: - for i, j in T.Parallel(block_M, block_N): - acc_s[i, j] = T.if_then_else(k * block_N + j >= seq_len, -T.infinity(acc_s.dtype), 0) - T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - T.copy(scores_max, scores_max_prev) - T.fill(scores_max, -T.infinity(accum_dtype)) - T.reduce_max(acc_s, scores_max, dim=1, clear=False) - for i in T.Parallel(block_M): - scores_max[i] = T.max(scores_max[i], scores_max_prev[i]) - for i in T.Parallel(block_M): - scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale) - for i, j in T.Parallel(block_M, block_N): - acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale) - T.reduce_sum(acc_s, scores_sum, dim=1) - for i in T.Parallel(block_M): - logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i] - T.copy(acc_s, acc_s_cast) - - for i, j in T.Parallel(block_M, dim): - acc_o[i, j] *= scores_scale[i] - - T.copy(V[bz, k * block_N : (k + 1) * block_N, by, :], V_shared) - T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) - - for i, j in T.Parallel(block_M, dim): - acc_o[i, j] /= logsum[i] - T.copy(acc_o, O_shared) - T.copy(O_shared, Output[bz, bx * block_M : (bx + 1) * block_M, by, :]) - - return main - - -def ref_program(Q, K, V, is_causal): - dim = Q.size(-1) - scores = torch.einsum("bqhd,bkhd->bhqk", Q, K) - scores = scores / torch.sqrt(torch.tensor(dim, dtype=scores.dtype)) - if is_causal: - seq_len = Q.size(1) - mask = torch.tril(torch.ones(seq_len, seq_len, device=scores.device)) - mask = mask.unsqueeze(0).unsqueeze(0) - scores = scores.masked_fill(mask == 0, float("-inf")) - attention_weights = F.softmax(scores, dim=-1) - output = torch.einsum("bhqk,bkhd->bqhd", attention_weights, V) - return output - - -def main( - batch: int = 8, - heads: int = 32, - seq_len: int = 4096, - dim: int = 128, - is_causal: bool = False, - tune: bool = False, -): - flops_per_matmul = 2.0 * batch * heads * seq_len * seq_len * dim - total_flops = 2 * flops_per_matmul - if is_causal: - total_flops *= 0.5 - - if not tune: - #kernel = flashattn(batch, heads, seq_len, dim, is_causal, block_M=128, block_N=128, num_stages=1, threads=128) - # Changed block size so fits in shared mem - kernel = flashattn(batch, heads, seq_len, dim, is_causal, block_M=64, block_N=64, num_stages=1, threads=128) - ref_program_processed = partial(ref_program, is_causal=is_causal) - profiler = kernel.get_profiler() - profiler.assert_allclose(ref_program_processed, rtol=0.01, atol=0.01) - print("All checks pass.") - latency = profiler.do_bench(ref_program_processed, warmup=500) - print("Ref: {:.2f} ms".format(latency)) - print("Ref: {:.2f} TFlops".format(total_flops / latency * 1e-9)) - latency = profiler.do_bench(warmup=500) - print("Tile-lang: {:.2f} ms".format(latency)) - print("Tile-lang: {:.2f} TFlops".format(total_flops / latency * 1e-9)) - else: - best_result = flashattn(batch, heads, seq_len, dim, is_causal) - best_latency = best_result.latency - best_config = best_result.config - ref_latency = best_result.ref_latency - print(f"Best latency: {best_latency}") - print(f"Best TFlops: {total_flops / best_latency * 1e-9}") - print(f"Best config: {best_config}") - print(f"Ref latency: {ref_latency}") - - -if __name__ == "__main__": - - main() \ No newline at end of file diff --git a/examples/generic_python/flash_attention/tilus_attention.py b/examples/generic_python/flash_attention/tilus_attention.py deleted file mode 100644 index 5169d9316..000000000 --- a/examples/generic_python/flash_attention/tilus_attention.py +++ /dev/null @@ -1,611 +0,0 @@ -# Code taken from https://github.com/NVIDIA/tilus/blob/main/examples/attention/flash_attention_v3.py -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -import numpy as np -import pandas as pd -import tilus -import torch -from tilus import boolean, f32, int32, void_p -from hidet.ir import DataType -from tilus.ir import RegisterTensor, SharedTensor -from tilus.ir.tensor import GlobalTensor -from tilus.utils import benchmark_func, cdiv - -pd.options.display.max_columns = None -pd.options.display.width = 1000 - - -@tilus.autotune("num_warps", [4, 8]) -@tilus.autotune("block_q", [32, 64, 128]) -@tilus.autotune("block_kv", [32, 64, 128]) -@tilus.autotune("split_kv", [-1, 512, 1024, 4096]) -@tilus.autotune("keep_q_in_regs", [False, True]) -class FlashAttention(tilus.Script): - LOG2_E = 1.4426950408889634 # log2(e) - - debug_schedule = dict( - num_warps=4, - block_q=64, - block_kv=64, - split_kv=-1, - keep_q_in_regs=False, - ) - - def __init__( - self, - dtype: DataType, - num_heads: int, - num_heads_kv: int, - head_size: int, - num_warps: int, - block_q: int, - block_kv: int, - split_kv: int, - keep_q_in_regs: bool, - ): - super().__init__() - self.dtype: DataType = dtype - self.num_heads = num_heads - self.num_heads_kv = num_heads_kv - self.head_size = head_size - self.num_warps = num_warps - self.block_q = block_q - self.block_kv = block_kv - self.split_kv = split_kv - self.keep_q_in_regs = keep_q_in_regs - self.score_scale = float(1.0 / np.sqrt(head_size)) - self.group_heads = num_heads // num_heads_kv - - assert self.split_kv % self.block_kv == 0 or self.split_kv == -1, ( - "split_kv must be a multiple of block_kv or -1" - ) - - # determine layout - self.sv_config = self.cuda.resolve_dot_config( - dtype, f32, m=block_q, n=head_size, k=block_kv, warp_m=num_warps, warp_n=1 - ) - - def apply_mask(self, score: RegisterTensor, q_offset: int32, kv_offset: int32): - mask = self.register_tensor( - dtype=boolean, - shape=[self.block_q, self.block_kv], - init=lambda i, j: i + q_offset >= j + kv_offset, - ) - self.assign(score, score + self.where(mask, x=0.0, y=-1e6)) - - def softmax_rescale( - self, - score: RegisterTensor, - m: RegisterTensor, - l: RegisterTensor, - o: RegisterTensor, - ) -> RegisterTensor: - scale = self.score_scale * self.LOG2_E # log2(e) * score_scale - cur_m = self.max(score, dim=1, keepdim=True) * scale # [block_q, 1] - new_m = self.maximum(m, cur_m) # [block_q, 1] - rp = self.exp2(score * scale - new_m) # [block_q, block_kv] - m_scale = self.exp2(m - new_m) - self.assign(o, o * m_scale) - self.assign(l, l * m_scale + self.sum(rp, dim=1, keepdim=True)) - self.assign(m, new_m) - return rp.to(self.dtype) - - def attention_iteration( - self, - bs: int32, - kv_offset: int32, - q_offset: int32, - head: int32, - gk: GlobalTensor, - gv: GlobalTensor, - sq: SharedTensor, # f16[block_q, head_size] - rq: RegisterTensor, # f16[block_q, head_size] - sk: SharedTensor, # f16[block_kv, head_size], - sv: SharedTensor, # f16[block_kv, head_size], - o: RegisterTensor, # f32[block_q, head_size] - m: RegisterTensor, # f32[block_q, 1] - l: RegisterTensor, # f32[block_q, 1] - check_bounds: bool, - ): - if not self.keep_q_in_regs: - self.load_shared(sq, out=rq) - # wait for the async copy of k to finish - self.copy_async_wait_group(0) - self.sync() - self.copy_async( - gv, - sv, - offsets=[bs, kv_offset, head // self.group_heads, 0], - dims=[1, 3], - check_bounds=check_bounds, - ) - self.copy_async_commit_group() - - # issue the async copy for v and perform dot(q, k) - rk = self.load_shared(sk) # [block_kv, head_size] - score = self.dot(rq, rk.transpose(), acc_dtype=f32) # [block_q, block_kv] - - if check_bounds: - self.apply_mask(score, q_offset, kv_offset) # apply causal mask - - # wait for the async copy of v to finish - self.copy_async_wait_group(0) - self.sync() - self.copy_async( - gk, - sk, - offsets=[bs, kv_offset + self.block_kv, head // self.group_heads, 0], - dims=[1, 3], - check_bounds=check_bounds, - ) - self.copy_async_commit_group() - - # load v to register - rv = self.load_shared(sv) # [block_kv, head_size] - - # online softmax - rp = self.softmax_rescale(score, m=m, l=l, o=o) - - # pv - cur_o = self.dot(rp, rv, acc_dtype=f32) # [block_q, head_size] - self.annotate_layout(cur_o, self.sv_config.lc) - self.assign(o, o + cur_o) - - def main_loop( - self, - gq: GlobalTensor, - gk: GlobalTensor, - gv: GlobalTensor, - o: RegisterTensor, - m: RegisterTensor, - l: RegisterTensor, - ): - # calculate offsets - q_offset = self.blockIdx.x * self.block_q - kv_start_offset = 0 if self.split_kv == -1 else self.blockIdx.y * self.split_kv - - if q_offset + self.block_q <= kv_start_offset: - return - - head = self.blockIdx.z % self.num_heads - bs = self.blockIdx.z // self.num_heads - - sq = self.shared_tensor(dtype=self.dtype, shape=[self.block_q, self.head_size]) - sk = self.shared_tensor(dtype=self.dtype, shape=[self.block_kv, self.head_size]) - sv = self.shared_tensor(dtype=self.dtype, shape=[self.block_kv, self.head_size]) - - rq = self.register_tensor(dtype=self.dtype, shape=[self.block_q, self.head_size]) - - # copy q to shared memory - self.copy_async( - gq, sq, offsets=[bs, q_offset, head, 0], dims=[1, 3], check_bounds=True - ) - self.copy_async_wait_all() - self.sync() - - # copy q to registers if not keeping in shared memory - if self.keep_q_in_regs: - self.load_shared(sq, out=rq) # [block_q, head_size] - self.free_shared(sq) - - # issue a copy of gk - self.copy_async(gk, sk, offsets=[bs, 0, head // self.group_heads, 0], dims=[1, 3]) - self.copy_async_commit_group() - - kv_offset_inner_end = (q_offset + 1) // self.block_kv * self.block_kv - if self.split_kv != -1: - kv_offset_inner_end = min( - kv_offset_inner_end, kv_start_offset + self.split_kv - ) - for kv_offset in range(kv_start_offset, kv_offset_inner_end, self.block_kv): - self.attention_iteration( - bs, - kv_offset, - q_offset, - head, - gk, - gv, - sq, - rq, - sk, - sv, - o, - m, - l, - check_bounds=False, - ) - - kv_offset_end = q_offset + self.block_q - if self.split_kv != -1: - kv_offset_end = min(kv_offset_end, kv_start_offset + self.split_kv) - for kv_offset in range(kv_offset_inner_end, kv_offset_end, self.block_kv): - self.attention_iteration( - bs, - kv_offset, - q_offset, - head, - gk, - gv, - sq, - rq, - sk, - sv, - o, - m, - l, - check_bounds=True, - ) - - self.copy_async_wait_group(0) - self.sync() - self.free_shared(sk) - self.free_shared(sv) - if not self.keep_q_in_regs: - self.free_shared(sq) - - def store_back( - self, - o: RegisterTensor, - l: RegisterTensor, - m: RegisterTensor, - o_ptr: void_p, - batch_size: int, - q_len: int32, - ): - # o: [block_q, head_size] - # m: [block_q, 1] - # l: [block_q, 1] - go = self.global_view( - o_ptr, - dtype=self.dtype, - shape=[batch_size, q_len, self.num_heads, self.head_size], - ) - o = o / l - o_f16 = self.cast(o, dtype=self.dtype) # [block_q, head_size] - so = self.shared_tensor(dtype=self.dtype, shape=[self.block_q, self.head_size]) - - head = self.blockIdx.z % self.num_heads - q_offset = self.blockIdx.x * self.block_q - bs = self.blockIdx.z // self.num_heads - - if self.split_kv == -1: - self.store_shared(so, o_f16) - self.sync() - self.store_global( - go, - self.load_shared(so), - offsets=[bs, q_offset, head, 0], - dims=[1, 3], - ) - else: - num_q_blocks = cdiv(q_len, self.block_q) - semaphores = self.global_tensor( - dtype=int32, - shape=[num_q_blocks, batch_size, self.num_heads], - requires_clean=True, - ) - gm = self.global_tensor( - dtype=f32, - shape=[num_q_blocks, batch_size, self.num_heads, self.block_q], - requires_clean=False, - ) - gl = self.global_tensor( - dtype=f32, - shape=[num_q_blocks, batch_size, self.num_heads, self.block_q], - requires_clean=False, - ) - semaphore = semaphores[self.blockIdx.x, bs, head].item_ptr() - - sm = self.shared_tensor(dtype=f32, shape=[self.block_q]) - sl = self.shared_tensor(dtype=f32, shape=[self.block_q]) - - self.lock_semaphore(semaphore, value=self.blockIdx.y) - - # load previous o, m and l and merge with the current results - if self.blockIdx.y > 0: - self.copy_async(gm, sm, offsets=[self.blockIdx.x, bs, head, 0], dims=[3]) - self.copy_async(gl, sl, offsets=[self.blockIdx.x, bs, head, 0], dims=[3]) - self.copy_async(go, so, offsets=[bs, q_offset, head, 0], dims=[1, 3]) - self.copy_async_wait_all() - self.sync() - lhs_o = self.load_shared(so) - lhs_m = self.load_shared(sm).unsqueeze(1) - lhs_l = self.load_shared(sl).unsqueeze(1) - rhs_o = o_f16 - rhs_m = m - rhs_l = l - m = self.maximum(lhs_m, rhs_m) - lhs_ll = lhs_l * self.exp(lhs_m - m) - rhs_ll = rhs_l * self.exp(rhs_m - m) - l = lhs_ll + rhs_ll - o_f16 = lhs_o * self.cast( - lhs_ll / l, dtype=self.dtype - ) + rhs_o * self.cast(rhs_ll / l, dtype=self.dtype) - self.sync() - - # store the results to so and load it - self.store_shared(so, o_f16) - self.store_shared(sm, m.squeeze(dim=1)) - self.store_shared(sl, l.squeeze(dim=1)) - self.sync() - - # store the results to global memory and release the semaphore - self.store_global( - go, - self.load_shared(so), - offsets=[bs, q_offset, head, 0], - dims=[1, 3], - ) - self.store_global( - gm, - self.load_shared(sm), - offsets=[self.blockIdx.x, bs, head, 0], - dims=[3], - ) - self.store_global( - gl, - self.load_shared(sl), - offsets=[self.blockIdx.x, bs, head, 0], - dims=[3], - ) - self.sync() - - self.free_shared(sm) - self.free_shared(sl) - - # release the semaphore - self.release_semaphore( - semaphore, - value=self.blockIdx.y + 1 - if (self.blockIdx.y + 1) * self.split_kv < q_offset + self.block_q - else 0, - ) - self.free_shared(so) - - def __call__( - self, - batch_size: int, - q_len: int32, - kv_len: int32, - q_ptr: void_p, - k_ptr: void_p, - v_ptr: void_p, - o_ptr: void_p, - ): - self.attrs.warps = self.num_warps - self.attrs.blocks = ( - cdiv(q_len, self.block_q), - cdiv(kv_len, self.split_kv) if self.split_kv != -1 else 1, - self.num_heads * batch_size, - ) - - gq = self.global_view( - q_ptr, - dtype=self.dtype, - shape=[batch_size, q_len, self.num_heads, self.head_size], - ) - gk = self.global_view( - k_ptr, - dtype=self.dtype, - shape=[batch_size, kv_len, self.num_heads_kv, self.head_size], - ) - gv = self.global_view( - v_ptr, - dtype=self.dtype, - shape=[batch_size, kv_len, self.num_heads_kv, self.head_size], - ) - - o = self.register_tensor( - dtype=f32, shape=[self.block_q, self.head_size], init=0.0 - ) - m = self.register_tensor( - dtype=f32, shape=[self.block_q, 1], init=-1e6 - ) # rowmax(score) - l = self.register_tensor( - dtype=f32, shape=[self.block_q, 1], init=0.0 - ) # rowsum(exp(score - m)) - - self.main_loop(gq, gk, gv, o, m, l) - - self.store_back(o, l, m, o_ptr=o_ptr, batch_size=batch_size, q_len=q_len) - - -def flash_attention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, -): - """ - Flash attention function for variable length sequences. - - Parameters - ---------- - q: torch.Tensor - The query tensor of shape (bs, seqlen, num_heads, head_size). - - k: torch.Tensor - The key tensor of shape (bs, seqlen, num_heads_kv, head_size). - - v: torch.Tensor - The value tensor of shape (bs, seqlen, num_heads_kv, head_size). - - Returns - ------- - o: torch.Tensor - The output tensor of shape (bs, seqlen, num_heads, head_size). - """ - out = torch.empty_like(q) - FlashAttention( - dtype=tilus.float16, - num_heads=q.size(2), - num_heads_kv=k.size(2), - head_size=q.size(3), - )(q.size(0), q.size(1), k.size(1), q, k, v, out) - return out - - -def flash_attention_reference( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, -): - bs, seqlen, num_heads, head_size = q.size() - _, _, num_heads_kv, _ = k.size() - assert q.size(0) == k.size(0) == v.size(0), "Batch size must match for q, k, and v." - assert q.size(1) == k.size(1) == v.size(1), ( - "Sequence length must match for q, k, and v." - ) - assert q.size(3) == k.size(3) == v.size(3), "Head size must match for q, k, and v." - assert k.size(2) == v.size(2), "Number of heads in k and v must match." - assert num_heads % num_heads_kv == 0, ( - "Number of heads must be divisible by number of kv heads." - ) - - k = torch.repeat_interleave(k, num_heads // num_heads_kv, dim=2) - v = torch.repeat_interleave(v, num_heads // num_heads_kv, dim=2) - - q = torch.transpose(q, 1, 2).reshape(bs * num_heads, seqlen, head_size) - k = torch.transpose(k, 1, 2).reshape(bs * num_heads, seqlen, head_size) - v = torch.transpose(v, 1, 2).reshape(bs * num_heads, seqlen, head_size) - - score = torch.bmm(q, k.mT) / np.sqrt(head_size) # [bs * num_heads, seqlen, seqlen] - causal_mask = torch.tril(torch.ones(seqlen, seqlen, dtype=torch.bool), diagonal=0).to( - q.device - ) - causal_mask = causal_mask.unsqueeze(0) # [1, seqlen, seqlen] - causal_mask = causal_mask.expand( - bs * num_heads, seqlen, seqlen - ).contiguous() # [bs * num_heads, seqlen, seqlen] - score = score.masked_fill(causal_mask == 0, float("-inf")) - - o = torch.bmm( - torch.softmax(score.float(), dim=-1).to(q.dtype), v - ) # [bs * num_heads, seqlen, head_size] - o = o.reshape(bs, num_heads, seqlen, head_size).transpose(1, 2).contiguous() - return o - - -def flash_attention_flash_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, -): - try: - from flash_attn.cute.interface import flash_attn_func - - return flash_attn_func(q, k, v, causal=True) - except ImportError: - return flash_attention_reference(q, k, v) - - -def demo_flash_attention(): - for bs, seqlen, num_heads, head_size, num_heads_kv in [ - # [1, 8, 1, 64, 1], - [1, 4096, 32, 128, 8] - ]: - q = torch.rand(bs, seqlen, num_heads, head_size, dtype=torch.float16).cuda() - k = torch.rand(bs, seqlen, num_heads_kv, head_size, dtype=torch.float16).cuda() - v = torch.rand(bs, seqlen, num_heads_kv, head_size, dtype=torch.float16).cuda() - flash_attention(q, k, v) - torch.cuda.synchronize() - - -def main(bench=True): - headers = [ - "batch_size", - "seqlen", - "num_heads", - "head_size", - "num_heads_kv", - "name", - "latency (ms)", - "tflops", - ] - data = [] - for batch_size, seqlen, num_heads, head_size, num_heads_kv in [ - [1, 512, 32, 128, 8], - [1, 1024, 32, 128, 8], - [1, 2048, 32, 128, 8], - [1, 4096, 32, 128, 8], - #[1, 8192, 32, 128, 8], - [1, 512, 64, 128, 8], - [1, 1024, 64, 128, 8], - [1, 2048, 64, 128, 8], - [1, 4096, 64, 128, 8], - #[1, 8192, 64, 128, 8], - ]: - q = torch.rand( - batch_size, seqlen, num_heads, head_size, dtype=torch.float16 - ).cuda() - k = torch.rand( - batch_size, seqlen, num_heads_kv, head_size, dtype=torch.float16 - ).cuda() - v = torch.rand( - batch_size, seqlen, num_heads_kv, head_size, dtype=torch.float16 - ).cuda() - for name, runner in [ - ("flash-attn", flash_attention_flash_attn), - ("tilus", flash_attention), - ]: - print( - f"Running {name} with batch_size={batch_size}, seqlen={seqlen}, num_heads={num_heads}, head_size={head_size}, num_heads_kv={num_heads_kv}" - ) - try: - actual = runner(q, k, v) - except torch.OutOfMemoryError: - print("Out of memory, skipping this configuration.") - continue - - try: - expected = flash_attention_reference(q, k, v) - torch.testing.assert_close(actual, expected, atol=1e-2, rtol=1e-2) - except torch.OutOfMemoryError: - pass - - latency = ( - benchmark_func( - lambda: runner(q, k, v), - warmup=20, - repeat=50, - ) - if bench - else float("nan") - ) - tflops = ( - 2 * batch_size * num_heads * seqlen * head_size * seqlen / latency * 1e-9 - ) - data.append( - [ - batch_size, - seqlen, - num_heads, - head_size, - num_heads_kv, - name, - latency, - tflops, - ] - ) - df = pd.DataFrame(data, columns=headers) - df_pivot = df.pivot( - index=[ - "batch_size", - "seqlen", - "num_heads", - "head_size", - "num_heads_kv", - ], - columns="name", - values=["latency (ms)", "tflops"], - ).reset_index() - # sort by (batch_size, num_heads, head_size, seqlen) - df_pivot = df_pivot.sort_values( - by=["batch_size", "num_heads", "head_size", "seqlen"], - ascending=[True, True, True, True], - ) - print(df_pivot) - - -if __name__ == "__main__": - main() - # ncu_run(main, bench=False, kernel_regex="flash_fwd|flash_attention") \ No newline at end of file diff --git a/examples/generic_python/flash_attention/triton_attention.py b/examples/generic_python/flash_attention/triton_attention.py deleted file mode 100644 index e4f36a133..000000000 --- a/examples/generic_python/flash_attention/triton_attention.py +++ /dev/null @@ -1,276 +0,0 @@ -# example taken from https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html -# removed optional support for FP8 -# Removed backward pass -# removed checks for e.g. hopper/hip etc -import pytest -import torch -import os -import numpy as np - -import triton -import triton.language as tl -from triton.tools.tensor_descriptor import TensorDescriptor - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - -configs = [ - triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w) \ - for BM in [64, 128]\ - for BN in [32, 64, 128]\ - for s in [2, 3, 4] \ - for w in [4, 8]\ -] - - -def prune_invalid_configs(configs, named_args, **kwargs): - N_CTX = kwargs["N_CTX"] - STAGE = kwargs["STAGE"] - - # Filter out configs where BLOCK_M > N_CTX - # Filter out configs where BLOCK_M < BLOCK_N when causal is True - return [ - conf for conf in configs if conf.kwargs.get("BLOCK_M", 0) <= N_CTX and ( - conf.kwargs.get("BLOCK_M", 0) >= conf.kwargs.get("BLOCK_N", 0) or STAGE == 1) - ] - - -@triton.jit -def _attn_fwd_inner(acc, l_i, m_i, q, # - desc_k, desc_v, # - offset_y, dtype: tl.constexpr, start_m, qk_scale, # - BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # - STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, # - N_CTX: tl.constexpr, warp_specialize: tl.constexpr): - # range of values handled by this stage - if STAGE == 1: - lo, hi = 0, start_m * BLOCK_M - elif STAGE == 2: - lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M - lo = tl.multiple_of(lo, BLOCK_M) - # causal = False - else: - lo, hi = 0, N_CTX - offsetk_y = offset_y + lo - if dtype == tl.float8e5: - offsetv_y = offset_y * HEAD_DIM + lo - else: - offsetv_y = offset_y + lo - # loop over k, v and update accumulator - for start_n in tl.range(lo, hi, BLOCK_N, warp_specialize=warp_specialize): - start_n = tl.multiple_of(start_n, BLOCK_N) - # -- compute qk ---- - k = desc_k.load([offsetk_y, 0]).T - qk = tl.dot(q, k) - if STAGE == 2: - mask = offs_m[:, None] >= (start_n + offs_n[None, :]) - qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) - m_ij = tl.maximum(m_i, tl.max(qk, 1)) - qk -= m_ij[:, None] - else: - m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) - qk = qk * qk_scale - m_ij[:, None] - p = tl.math.exp2(qk) - # -- compute correction factor - alpha = tl.math.exp2(m_i - m_ij) - l_ij = tl.sum(p, 1) - # -- update output accumulator -- - if warp_specialize and BLOCK_M == 128 and HEAD_DIM == 128: - BM: tl.constexpr = acc.shape[0] - BN: tl.constexpr = acc.shape[1] - acc0, acc1 = acc.reshape([BM, 2, BN // 2]).permute(0, 2, 1).split() - acc0 = acc0 * alpha[:, None] - acc1 = acc1 * alpha[:, None] - acc = tl.join(acc0, acc1).permute(0, 2, 1).reshape([BM, BN]) - else: - acc = acc * alpha[:, None] - # prepare p and v for the dot - if dtype == tl.float8e5: - v = desc_v.load([0, offsetv_y]).T - else: - v = desc_v.load([offsetv_y, 0]) - p = p.to(dtype) - # note that this non transposed v for FP8 is only supported on Blackwell - acc = tl.dot(p, v, acc) - # update m_i and l_i - # place this at the end of the loop to reduce register pressure - l_i = l_i * alpha + l_ij - m_i = m_ij - offsetk_y += BLOCK_N - offsetv_y += BLOCK_N - return acc, l_i, m_i - - -@triton.jit -def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape): - if isinstance(desc_or_ptr, tl.tensor_descriptor): - return desc_or_ptr - else: - return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape) - - -@triton.autotune(configs=configs, key=["N_CTX", "HEAD_DIM", "warp_specialize"], - prune_configs_by={'early_config_prune': prune_invalid_configs}) -@triton.jit -def _attn_fwd(sm_scale, M, # - Z, H, desc_q, desc_k, desc_v, desc_o, N_CTX, # - HEAD_DIM: tl.constexpr, # - BLOCK_M: tl.constexpr, # - BLOCK_N: tl.constexpr, # - STAGE: tl.constexpr, # - warp_specialize: tl.constexpr, # - ): - dtype = tl.float16 - tl.static_assert(BLOCK_N <= HEAD_DIM) - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - - y_dim = Z * H * N_CTX - desc_q = _maybe_make_tensor_desc(desc_q, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_M, HEAD_DIM]) - desc_v = _maybe_make_tensor_desc(desc_v, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_N, HEAD_DIM]) - desc_k = _maybe_make_tensor_desc(desc_k, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_N, HEAD_DIM]) - desc_o = _maybe_make_tensor_desc(desc_o, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_M, HEAD_DIM]) - - offset_y = off_z * (N_CTX * H) + off_h * N_CTX - qo_offset_y = offset_y + start_m * BLOCK_M - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - # initialize pointer to m and l - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 - acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - # load scales - qk_scale = sm_scale - qk_scale *= 1.44269504 # 1/log(2) - # load q: it will stay in SRAM throughout - q = desc_q.load([qo_offset_y, 0]) - # stage 1: off-band - # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE - # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE - if STAGE & 1: - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, # - desc_k, desc_v, # - offset_y, dtype, start_m, qk_scale, # - BLOCK_M, HEAD_DIM, BLOCK_N, # - 4 - STAGE, offs_m, offs_n, N_CTX, # - warp_specialize) - # stage 2: on-band - if STAGE & 2: - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, # - desc_k, desc_v, # - offset_y, dtype, start_m, qk_scale, # - BLOCK_M, HEAD_DIM, BLOCK_N, # - 2, offs_m, offs_n, N_CTX, # - warp_specialize) - # epilogue - m_i += tl.math.log2(l_i) - acc = acc / l_i[:, None] - m_ptrs = M + off_hz * N_CTX + offs_m - tl.store(m_ptrs, m_i) - desc_o.store([qo_offset_y, 0], acc.to(dtype)) - - - -def forward(q, k, v, causal, sm_scale, warp_specialize=True): - # shape constraints - HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] - # when v is in float8_e5m2 it is transposed. - HEAD_DIM_V = v.shape[-1] - assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V - assert HEAD_DIM_K in {16, 32, 64, 128, 256} - o = torch.empty_like(q) - stage = 3 if causal else 1 - extra_kern_args = {} - M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) - desc_q = q - desc_v = v - desc_k = k - desc_o = o - - def grid(META): - return (triton.cdiv(q.shape[2], META["BLOCK_M"]), q.shape[0] * q.shape[1], 1) - - _attn_fwd[grid]( - sm_scale, M, # - q.shape[0], q.shape[1], # - desc_q, desc_k, desc_v, desc_o, # - N_CTX=q.shape[2], # - HEAD_DIM=HEAD_DIM_K, # - STAGE=stage, # - warp_specialize=warp_specialize, # - **extra_kern_args) - - return o - - - - - -if __name__ == "__main__": - import torch - -# Triton kernel (your forward function) -# from your previous code: forward(ctx, Q, K, V, causal, sm_scale, warp_specialize=True) - -def flash_attention_reference(q, k, v, causal=False, sm_scale=1.0): - """ - Reference attention using batched matmul (like the pytest does). - Args: - q, k, v: [batch, heads, seq_len, head_dim] - causal: whether to apply causal mask - sm_scale: scaling factor for attention scores - Returns: - output: [batch, heads, seq_len, head_dim] - """ - # Batched matrix multiply - # scores: [batch, heads, seq_len, seq_len] - scores = torch.matmul(q, k.transpose(-2, -1)) * sm_scale - - if causal: - seq_len = q.shape[2] - mask = torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool)) - scores = scores.masked_fill(~mask[None, None, :, :], float("-inf")) - - # Softmax along last dimension - attention = torch.softmax(scores.float(), dim=-1) - - # Weighted sum with V - output = torch.matmul(attention, v) - return output - - - -# Example test -if __name__ == "__main__": - forward_simple() - exit(0) - - # Parameters - B, H, N_CTX, HEAD_DIM = 2, 4, 128, 64 - causal = True - sm_scale = 1.0 - - # Random inputs - Q = torch.randn(B, H, N_CTX, HEAD_DIM, dtype=torch.float16, device="cuda") - K = torch.randn_like(Q) - V = torch.randn_like(Q) - - # Triton output - output_triton = forward(Q, K, V, causal=causal, sm_scale=sm_scale, warp_specialize=True) - - # Reference output - output_ref = flash_attention_reference(Q.float(), K.float(), V.float(), causal=causal, sm_scale=sm_scale) - - # Compare - max_diff = (output_triton.float() - output_ref).abs().max() - mean_diff = (output_triton.float() - output_ref).abs().mean() - print(f"Max difference: {max_diff.item():.6f}") - print(f"Mean difference: {mean_diff.item():.6f}") - - diff --git a/examples/generic_python/loopy_example.py b/examples/generic_python/loopy_example.py deleted file mode 100644 index 49ad131f6..000000000 --- a/examples/generic_python/loopy_example.py +++ /dev/null @@ -1,38 +0,0 @@ -import numpy as np - -import pyopencl as cl -import pyopencl.array - -import loopy as lp -from loopy.version import LOOPY_USE_LANGUAGE_VERSION_2018_2 # noqa: F401 - - -# setup -# ----- -ctx = cl.create_some_context() -queue = cl.CommandQueue(ctx) - -n = 15 * 10**6 -a = cl.array.arange(queue, n, dtype=np.float32) - -# create -# ------ -knl = lp.make_kernel( - "{ [i]: 0<=i torch.Tensor: - m, k = x.size() - k, n = y.size() - out = torch.empty([m, n], dtype=x.dtype, device=x.device) - - for tile_m, tile_n in hl.tile([m, n]): - acc = hl.zeros([tile_m, tile_n], dtype=torch.float32) - for tile_k in hl.tile(k): - acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n]) - out[tile_m, tile_n] = acc - - return out - - -out = matmul(torch.randn([2048, 2048], device="cuda"), - torch.randn([2048, 2048], device="cuda")) \ No newline at end of file diff --git a/examples/generic_python/matmul/numba_matmul.py b/examples/generic_python/matmul/numba_matmul.py index 6486ec540..71a6b1b3a 100644 --- a/examples/generic_python/matmul/numba_matmul.py +++ b/examples/generic_python/matmul/numba_matmul.py @@ -250,10 +250,10 @@ def tune_optimized(M, N, K): M, N, K = 1024, 1024, 1024 run_basic(M, N, K) - #tune_basic(M, N, K) + tune_basic(M, N, K) - #run_optimized(M, N, K) - #tune_optimized(M, N, K) + run_optimized(M, N, K) + tune_optimized(M, N, K) diff --git a/examples/generic_python/matmul/test.py b/examples/generic_python/matmul/test.py deleted file mode 100644 index 991ef2cc4..000000000 --- a/examples/generic_python/matmul/test.py +++ /dev/null @@ -1,278 +0,0 @@ -import argparse -import itertools -import tilelang as tl -import tilelang.language as T -from tilelang.autotuner import AutoTuner -from tilelang.carver.template import MatmulTemplate -from tilelang.carver.arch import CUDA -from tilelang.carver.arch import CDNA -from tilelang.carver.roller.rasterization import NoRasterization -import torch - - -def ref_program(A, B): - """ - Compute the matrix product of A and the transpose of B. - - A and B are expected to be 2-D tensors where A has shape (M, K) and B has shape (N, K). The result is a tensor with shape (M, N) equal to A @ B.T, using the inputs' dtypes. - """ - return A @ B.T - - -def get_configs(M, N, K, with_roller=False, topk=20): - """ - Generate a list of kernel tuning configuration dictionaries for a tiled matrix-multiply. - - When with_roller is True this queries the MatmulTemplate roller to produce up to `topk` recommended - configurations (device-specific TensorCore-friendly tilings). Each returned dict contains: - - block_M, block_N, block_K: tile sizes - - num_stages: pipeline staging (0 means no explicit staging) - - thread_num: total threads used for the block - - enable_rasteration: whether a rasterization/swizzle layout was recommended (note spelling) - - When with_roller is False this returns the Cartesian product of a fixed set of candidate - parameters; the returned dicts use the backward-compatible key name "enable_rasteration" for that flag. - - Parameters: - M, N, K (int): GEMM dimensions used to generate valid tile sizes. - with_roller (bool): If True, use MatmulTemplate's roller to generate device-aware hints; - otherwise use a predefined candidate grid. - topk (int): Maximum number of roller hints to request when with_roller is True. - - Returns: - List[dict]: A list of configuration dictionaries as described above. - - Raises: - ValueError: if with_roller is True but the roller returns no hints. - """ - if with_roller: - arch = CUDA("cuda") if torch.version.hip is None else CDNA("hip") - carve_template = MatmulTemplate( - M=M, - N=N, - K=K, - in_dtype=T.float16, - out_dtype=T.float16, - accum_dtype=T.float32, - ).with_arch(arch) - - func = carve_template.equivalent_function() - assert func is not None, "Function is None" - roller_hints = carve_template.recommend_hints(topk=topk) - if roller_hints is None: - raise ValueError("No Roller Hints Found for TensorCore Scheduling") - configs = [] - for hint in roller_hints: - config = {} - block_m, block_n = hint.block - warp_m, warp_n = hint.warp - # block_rows, block_cols represents warp partitioning - block_rows, block_cols = block_m // warp_m, block_n // warp_n - config["block_M"] = block_m - config["block_N"] = block_n - config["block_K"] = hint.rstep[0] - config["num_stages"] = hint.pipeline_stage if hint.pipeline_stage > 1 else 0 - config["thread_num"] = block_rows * block_cols * 32 - config["enable_rasteration"] = hint.rasterization_plan is not NoRasterization - configs.append(config) - else: - block_M = [64, 128, 256] - block_N = [64, 128, 256] - block_K = [32, 64] - num_stages = [0, 1, 2, 3] - thread_num = [128, 256] - enable_rasterization = [True, False] - _configs = list( - itertools.product( - block_M, - block_N, - block_K, - num_stages, - thread_num, - enable_rasterization, - ) - ) - - configs = [ - { - "block_M": c[0], - "block_N": c[1], - "block_K": c[2], - "num_stages": c[3], - "thread_num": c[4], - "enable_rasteration": c[5], # keep param name for backward-compat - } - for c in _configs - ] - return configs - - -def get_best_config( - M, - N, - K, - with_roller: bool = False, - profile_backend: str = "event", -): - def kernel( - block_M=None, - block_N=None, - block_K=None, - num_stages=None, - thread_num=None, - enable_rasteration=None, - ): - dtype = T.bfloat16 - accum_dtype = T.float32 - - @T.prim_func - def main( - A: T.Tensor((M, K), dtype), - B: T.Tensor((N, K), dtype), - C: T.Tensor((M, N), dtype), - ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_N, block_K), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - C_shared = T.alloc_shared((block_M, block_N), dtype) - T.use_swizzle(panel_size=10, enable=enable_rasteration) - T.clear(C_local) - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[bx * block_N, k * block_K], B_shared) - T.gemm( - A_shared, - B_shared, - C_local, - transpose_B=True, - ) - T.copy(C_local, C_shared) - T.copy(C_shared, C[by * block_M, bx * block_N]) - - return main - - autotuner = ( - AutoTuner.from_kernel(kernel=kernel, configs=get_configs(M, N, K, with_roller)) - .set_compile_args( - out_idx=[-1], - target="auto", - ) - .set_profile_args( - supply_type=tl.TensorSupplyType.Integer, - ref_prog=ref_program, - skip_check=False, - backend=profile_backend, - ) - ) - return autotuner.run(warmup=3, rep=20) - - -def get_heuristic_config() -> dict: - # Get CUDA device properties - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - sm_major, sm_minor = torch.cuda.get_device_capability(device) - sm_version = sm_major * 10 + sm_minor - print(f"CUDA device capability: {sm_version}") - if sm_version in {80}: - return {"block_M": 128, "block_N": 256, "block_K": 32, "num_stages": 2, "thread_num": 128, "enable_rasteration": True} - elif sm_version in {90}: - return {"block_M": 128, "block_N": 256, "block_K": 64, "num_stages": 3, "thread_num": 256, "enable_rasteration": True} - else: - return {"block_M": 128, "block_N": 256, "block_K": 32, "num_stages": 0, "thread_num": 128, "enable_rasteration": True} - - -@tl.jit(out_idx=[-1]) -def matmul(M, N, K, block_M, block_N, block_K, num_stages, thread_num, enable_rasteration, dtype=T.float16, accum_dtype=T.float32): - @T.prim_func - def gemm_autotune( - A: T.Tensor((M, K), dtype), - B: T.Tensor((N, K), dtype), - C: T.Tensor((M, N), dtype), - ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by): - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_N, block_K), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - C_shared = T.alloc_shared((block_M, block_N), dtype) - T.use_swizzle(panel_size=10, enable=enable_rasteration) - T.clear(C_local) - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[bx * block_N, k * block_K], B_shared) - T.gemm( - A_shared, - B_shared, - C_local, - transpose_B=True, - ) - T.copy(C_local, C_shared) - T.copy(C_shared, C[by * block_M, bx * block_N]) - - return gemm_autotune - - -def main( - M: int = 4096, - N: int = 4096, - K: int = 4096, - use_autotune: bool = False, - with_roller: bool = False, - profile_backend: str = "event", -): - if use_autotune: - result = get_best_config( - M, - N, - K, - with_roller=with_roller, - profile_backend=profile_backend, - ) - print(result.config) - kernel = result.kernel - else: - config = get_heuristic_config() - kernel = matmul(M, N, K, **config) - - # benchmark - profiler = kernel.get_profiler(tensor_supply_type=tl.TensorSupplyType.Auto) - tilelang_latency = profiler.do_bench( - backend=profile_backend, - ) - ref_latency = profiler.do_bench( - ref_program, - backend=profile_backend, - ) - profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2) - print(f"TileLang latency: {tilelang_latency}") - print(f"Ref latency: {ref_latency}") - print(f"TileLang TFlops: {2 * M * N * K / tilelang_latency * 1e-9}") - print(f"Ref TFlops: {2 * M * N * K / ref_latency * 1e-9}") - - -def run_regression_perf(M: int = 4096, N: int = 4096, K: int = 4096): - config = get_heuristic_config() - kernel = matmul(M, N, K, **config) - profiler = kernel.get_profiler(tensor_supply_type=tl.TensorSupplyType.Auto) - return profiler.do_bench(backend="cupti") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Autotuned MatMul Benchmark") - parser.add_argument("--m", type=int, default=4096, help="Matrix dimension M") - parser.add_argument("--n", type=int, default=4096, help="Matrix dimension N") - parser.add_argument("--k", type=int, default=4096, help="Matrix dimension K") - parser.add_argument("--use_autotune", action="store_true", default=False, help="Whether to use autotune for matmul configs") - parser.add_argument("--with_roller", action="store_true", default=False, help="Whether to enable BitBLAS roller for search space") - parser.add_argument("--profile_backend", type=str, default="event", help="Profiler backend") - args = parser.parse_args() - main( - args.m, - args.n, - args.k, - args.use_autotune, - args.with_roller, - args.profile_backend, - ) \ No newline at end of file diff --git a/examples/generic_python/matmul/test_tilus.py b/examples/generic_python/matmul/test_tilus.py deleted file mode 100644 index 6d4b2430d..000000000 --- a/examples/generic_python/matmul/test_tilus.py +++ /dev/null @@ -1,26 +0,0 @@ -import torch -from tilus_matmul import MatmulBasic - -sizes = [ - (65, 65, 17), - (67, 71, 19), - (1, 1, 1), - (63, 63, 15), - (129, 130, 33), -] - -matmul = MatmulBasic() -for m, n, k in sizes: - print(m, n, k) - - a = torch.randn(m, k, dtype=torch.float16, device="cuda") - b = torch.randn(k, n, dtype=torch.float16, device="cuda") - - c_actual = torch.empty(m, n, dtype=torch.float16, device="cuda") - c_expect = a @ b - - matmul(m, n, k, a, b, c_actual) - - torch.cuda.synchronize() - - torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) \ No newline at end of file diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/generic_python/matmul/tilelang_matmul.py index 0c476c786..751c0add4 100644 --- a/examples/generic_python/matmul/tilelang_matmul.py +++ b/examples/generic_python/matmul/tilelang_matmul.py @@ -315,10 +315,10 @@ def tune_opt(M, N, K): if __name__ == "__main__": M, N, K = 4096, 4096, 4096 - #run_basic(M, N, K) - #run_opt(M, N, K) + run_basic(M, N, K) + run_opt(M, N, K) tune_basic(M, N, K) - #tune_opt(M, N, K) + tune_opt(M, N, K) diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/generic_python/matmul/tilus_matmul.py index 811d5c2d7..c65cdeaa7 100644 --- a/examples/generic_python/matmul/tilus_matmul.py +++ b/examples/generic_python/matmul/tilus_matmul.py @@ -257,11 +257,7 @@ def tune_opt(M, N, K): if __name__ == "__main__": M, N, K = 4096, 4096, 4096 - #M, N, K = 8192, 8192, 8192 - #run_basic(M, N, K) - - #run_optmized(M, N, K) - + run_basic(M, N, K) + run_optimized(M, N, K) tune_basic(M, N, K) - - #tune_opt(M, N, K) \ No newline at end of file + tune_opt(M, N, K) \ No newline at end of file diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py index cac3cab73..d5522c627 100644 --- a/examples/generic_python/matmul/triton_matmul.py +++ b/examples/generic_python/matmul/triton_matmul.py @@ -305,10 +305,8 @@ def tune_opt(M, N, K): if __name__ == "__main__": M, N, K = 4096, 4096, 4096 - #M, N, K = 8192, 8192, 8192 - #run_basic(M, N, K) - #run_opt(M, N, K) + run_basic(M, N, K) + run_opt(M, N, K) - - #tune_basic(M, N, K) - #tune_opt(M, N, K) + tune_basic(M, N, K) + tune_opt(M, N, K) diff --git a/examples/generic_python/matmul/warp_matmul.py b/examples/generic_python/matmul/warp_matmul.py index 388ae3a1b..c4a218b21 100644 --- a/examples/generic_python/matmul/warp_matmul.py +++ b/examples/generic_python/matmul/warp_matmul.py @@ -1,4 +1,3 @@ -import torch import numpy as np import warp as wp @@ -30,38 +29,6 @@ def gemm( C[i, j] = wp.float16(sum) -# tile size -TILE_M = 32 -TILE_N = 32 -TILE_K = 32 - -# num threads per-tile -TILE_THREADS = 1024 - -# GEMM example from https://nvidia.github.io/warp/user_guide/tiles.html -@wp.kernel() -def tile_gemm( - A: wp.array2d(dtype=wp.float16), - B: wp.array2d(dtype=wp.float16), - C: wp.array2d(dtype=wp.float16) -): - # output tile index - i, j = wp.tid() - - sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32) - - K = A.shape[1] - count = (K + TILE_K - 1) // TILE_K - - for k in range(0, count): - a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i*TILE_M, k*TILE_K), bounds_check=True) - b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k*TILE_K, j*TILE_N), bounds_check=True) - - wp.tile_matmul(a, b, sum) - - wp.tile_store(C, wp.tile_astype(sum, wp.float16), offset=(i*TILE_M, j*TILE_N), bounds_check=True) - - def run_gemm(M, N, K): rng = np.random.default_rng(42) A = rng.random((M, K)).astype(np.float16) @@ -79,28 +46,6 @@ def run_gemm(M, N, K): print("Succes") -def run_gemm_tiled(M, N, K): - rng = np.random.default_rng(42) - A = rng.random((M, K)).astype(np.float16) - B = rng.random((K, N)).astype(np.float16) - C = np.zeros((M, N), dtype=np.float16) - - A_wp = wp.array(A) - B_wp = wp.array(B) - C_wp = wp.array(C) - - - wp.launch_tiled( - tile_gemm, - dim=((M + TILE_M - 1) // TILE_M, (N + TILE_N - 1) // TILE_N), - inputs=[A_wp, B_wp, C_wp], - block_dim=TILE_THREADS) - - np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-2, atol=M * 2**(-11)) - - print("Succes") - - def tune(M, K, N): rng = np.random.default_rng(42) A = rng.random((M, K)).astype(np.float16) @@ -132,9 +77,7 @@ def tune(M, K, N): if __name__ == "__main__": M, N, K = 1024, 1024, 1024 - #run_gemm(M, N, K) - #run_gemm_tiled(M, N, K) - + run_gemm(M, N, K) tune(M, N, K) diff --git a/examples/generic_python/matmul_old/tilelang_matmul.py b/examples/generic_python/matmul_old/tilelang_matmul.py deleted file mode 100644 index 2a7b5154d..000000000 --- a/examples/generic_python/matmul_old/tilelang_matmul.py +++ /dev/null @@ -1,212 +0,0 @@ -import tilelang -import tilelang.language as T -import torch -from kernel_tuner import tune_kernel - - -#@tilelang.jit -def matmul(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): - @T.prim_func - def main( - A: T.Tensor((M, K), dtype), - B: T.Tensor((K, N), dtype), - C: T.Tensor((M, N), dtype), - ): - # Define a grid with enough blocks to cover M×N - num_threads=128 - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): - - # Allocate shared memory for the current tile of A and B - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - - # Allocate a local (register) fragment for partial accumulations - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - - # Enable swizzle-based rasterization for better L2 locality - panel_size = 4 - T.use_swizzle(panel_size=panel_size, enable=True) - - # Initialize the local accumulation buffer to zero - T.clear(C_local) - - num_stages=3 - - # Loop over the K dimension in block_K chunks, using a pipeline - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - # Copy from global memory to shared memory - T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[k * block_K, bx * block_N], B_shared) - - # Perform a matrix multiply-accumulate on the tile - T.gemm(A_shared, B_shared, C_local) - - # Copy the accumulated result from local memory (C_local) to global memory (C) - T.copy(C_local, C[by * block_M, bx * block_N]) - - return main - -@tilelang.jit -def matmul_with_decorator(M, N, K, block_M, block_N, block_K, dtype: str = 'float16', accum_dtype: str = 'float32'): - @T.prim_func - def main( - A: T.Tensor((M, K), dtype), - B: T.Tensor((K, N), dtype), - C: T.Tensor((M, N), dtype), - ): - # Define a grid with enough blocks to cover M×N - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): - - # Allocate shared memory for the current tile of A and B - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - - # Allocate a local (register) fragment for partial accumulations - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - - # Initialize the local accumulation buffer to zero - T.clear(C_local) - - # Loop over the K dimension in block_K chunks, using a 3-stage pipeline - for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): - # Copy from global memory to shared memory - T.copy(A[by * block_M, k * block_K], A_shared) - T.copy(B[k * block_K, bx * block_N], B_shared) - - # Perform a matrix multiply-accumulate on the tile - T.gemm(A_shared, B_shared, C_local) - - # Copy the accumulated result from local memory (C_local) to global memory (C) - T.copy(C_local, C[by * block_M, bx * block_N]) - - return main - - - -def run(m, n, k): - a = torch.randn(m, k, device="cuda", dtype=torch.float16) - b = torch.randn(k, n, device="cuda", dtype=torch.float16) - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - kernel = matmul(m, n, k, 128, 128, 32) - kernel(a, b, c) - ref_c = a @ b - tol = m * 2**(-11) - # Validate results - torch.testing.assert_close(c, ref_c, rtol=tol, atol=tol) - - -def call_tilelang(kernel_function, args, kwargs, grid, threads, params): - compiled_kernel = kernel_function(**kwargs) - compiled_kernel(*args) - - -def time(m, n, k): - a = torch.randn(m, k, device="cuda", dtype=torch.float16) - b = torch.randn(k, n, device="cuda", dtype=torch.float16) - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - c_ans = a @ b - - args = [a, b, c] - tune_params = dict() - tune_params["M"] = [m] - tune_params["K"] = [k] - tune_params["N"] = [n] - tune_params["block_M"] = [64, 128] - tune_params["block_N"] = [64, 128] - tune_params["block_K"] = [32, 64] - - results_kt, env = tune_kernel("matmul", matmul, m * n, args, tune_params, lang="generic_python", - call_function=call_tilelang, decorator="@tilelang.jit", verbose=False, iterations=100) - - import time - num_repeats = 100 - times_direct = [] - for config in results_kt: - bs_m = config["block_M"] - bs_n = config["block_N"] - bs_k = config["block_K"] - - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - - kernel = matmul_with_decorator(m, n, k, bs_m, bs_n, bs_k) - kernel(a, b, c) - - torch.allclose(c.cpu(), c_ans.cpu(), atol=m * 2**(-11)) - - for i in range(num_repeats): - times = [] - - torch.cuda.synchronize() - start = time.time() - kernel(a, b, c) - torch.cuda.synchronize() - times.append(time.time() - start) - - avg_time_ms = round((1000 * sum(times) / len(times)), 3) - times_direct.append(avg_time_ms) - print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, time={avg_time_ms}ms") - - - import matplotlib.pyplot as plt - - # Extract times - times_kt = [cfg['time'] for cfg in results_kt] - - # x-axis labels - configs = [f"config{i}" for i in range(len(times_kt))] - x = range(len(configs)) - - plt.figure(figsize=(10,6)) - plt.plot(configs, times_kt, marker='s', label='KernelTuner') - plt.plot(configs, times_direct, marker='x', label='Direct') - plt.ylabel('Time (ms)') - plt.xlabel('Configuration') - plt.title('Kernel execution time per configuration') - plt.xticks(rotation=45) - plt.legend() - plt.grid(True) - plt.tight_layout() - plt.savefig("tilelang.png") - print("saved fig") - - -def tune(m, n, k): - a = torch.randn(m, k, device="cuda", dtype=torch.float16) - b = torch.randn(k, n, device="cuda", dtype=torch.float16) - c = torch.empty(m, n, device="cuda", dtype=torch.float16) - c_actual = a @ b - - args = [a, b, c] - tune_params = dict() - tune_params["M"] = [m] - tune_params["K"] = [k] - tune_params["N"] = [n] - tune_params["block_M"] = [64, 128, 256] - tune_params["block_N"] = [64, 128, 256] - tune_params["block_K"] = [32, 64, 128] - tune_params["num_stages"] = [2, 3, 4] - tune_params["panel_size"] = [4, 8] # equivalent to group size m in Triton - tune_params["num_threads"] = [64, 128, 256] - - restrictions = [ - # tile size budget - "block_M * block_N <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "block_M <= 4 * block_N", - "block_N <= 4 * block_M", - - # large K only with reasonably large M/N - "not (block_K == 128 and block_M < 64 and block_N < 64)", - ] - - tol = m * 2**(-11) - answer = [None, None, c_actual.cpu()] - - results, env = tune_kernel("matmul", matmul, m * n, args, tune_params, atol=tol, lang="generic_python", - call_function=call_tilelang, restrictions=restrictions, answer=answer, decorator="@tilelang.jit", verbose=False) - -if __name__ == "__main__": - #m, n, k = 1024, 1024, 1024 - m, n, k = 8192, 8192, 8192 - time(m, n, k) \ No newline at end of file diff --git a/examples/generic_python/matmul_old/tilus_matmul.py b/examples/generic_python/matmul_old/tilus_matmul.py deleted file mode 100644 index e8996c201..000000000 --- a/examples/generic_python/matmul_old/tilus_matmul.py +++ /dev/null @@ -1,266 +0,0 @@ -import math - -import pandas -import tilus -import torch -from tilus import float16, float32, int32 -from tilus.utils import benchmark_func -from kernel_tuner import tune_kernel, run_kernel - - - -class MatmulV4(tilus.Script): - def __init__(self): - super().__init__() - self.block_m = 128 - self.block_n = 128 - self.block_k = 16 - self.num_warps = 4 - self.num_stages = 4 - - def __call__( - self, - m_size: int32, - n_size: int, - k_size: int, - a_ptr: ~float16, - b_ptr: ~float16, - c_ptr: ~float16, - ): - self.attrs.blocks = [ - self.utils.ceil_div(m_size, self.block_m), - self.utils.ceil_div(n_size, self.block_n), - ] - self.attrs.warps = self.num_warps - - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k - offset_m: int32 = block_m * self.blockIdx.x - offset_n: int32 = block_n * self.blockIdx.y - - ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) - gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) - sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) - sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) - acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) - - for stage in range(self.num_stages - 1): - offset_k = stage * self.block_k - self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) - self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) - self.copy_async_commit_group() - - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - current_stage: int32 = 0 - preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): - # computation for current tile - a = self.load_shared(sa[current_stage]) - b = self.load_shared(sb[current_stage]) - self.dot(a, b, acc, out=acc) - - # preload the next tile of A and B into shared memory - preload_offset_k = offset_k + (self.num_stages - 1) * block_k - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], - ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], - ) - self.copy_async_commit_group() - - # update the stage - current_stage = (current_stage + 1) % self.num_stages - preload_stage = (preload_stage + 1) % self.num_stages - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - self.free_shared(sa) - self.free_shared(sb) - - casted_acc = self.cast(acc, dtype=float16) - gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) - - -class MatmulGroupedOrdering(tilus.Script): - def __init__(self): - super().__init__() - self.block_m = 128 - self.block_n = 128 - self.block_k = 16 - self.num_warps = 4 - self.num_stages = 4 - self.group_size_m = 8 - - def __call__( - self, - m_size: int32, - n_size: int, - k_size: int, - a_ptr: ~float16, - b_ptr: ~float16, - c_ptr: ~float16, - ): - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k - - num_pid_m = self.utils.ceil_div(m_size, block_m) - num_pid_n = self.utils.ceil_div(n_size, block_n) - self.attrs.blocks = [num_pid_m * num_pid_n] - - pid = self.blockIdx.x - num_pid_in_group = self.group_size_m * num_pid_n - group_id = pid // num_pid_in_group - - first_pid_m = group_id * self.group_size_m - group_size_m = min(num_pid_m - first_pid_m, self.group_size_m) - - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - self.attrs.warps = self.num_warps - offset_m: int32 = pid_m * block_m - offset_n: int32 = pid_n * block_n - - ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) - gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) - sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) - sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) - acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) - - for stage in range(self.num_stages - 1): - offset_k = stage * self.block_k - self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) - self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) - self.copy_async_commit_group() - - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - current_stage: int32 = 0 - preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): - # computation for current tile - a = self.load_shared(sa[current_stage]) - b = self.load_shared(sb[current_stage]) - self.dot(a, b, acc, out=acc) - - # preload the next tile of A and B into shared memory - preload_offset_k = offset_k + (self.num_stages - 1) * block_k - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], - ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], - ) - self.copy_async_commit_group() - - # update the stage - current_stage = (current_stage + 1) % self.num_stages - preload_stage = (preload_stage + 1) % self.num_stages - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - self.free_shared(sa) - self.free_shared(sb) - - casted_acc = self.cast(acc, dtype=float16) - gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) - - - -def main(): - headers = ["m", "n", "k", "name", "latency (ms)", "tflops"] - workloads = [ - [4096, 4096, 4096], - [1024, 1024, 14336], - ] - - rows = [] - for m, n, k in workloads: - matmul = MatmulGroupedOrdering() #MatmulV4() - - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - matmul(m, n, k, a, b, c_actual) - - # check correctness - torch.testing.assert_close(c_expect, c_actual) - - # benchmark - for name, func in [ - ("torch", lambda: torch.matmul(a, b, out=c_expect)), - ("tilus", lambda: matmul(m, n, k, a, b, c_actual)), - ]: - latency = benchmark_func(func, warmup=5, repeat=20) - tflops = 2 * m * n * k / latency * 1e-9 - rows.append([m, n, k, name, latency, tflops]) - - df = pandas.DataFrame(rows, columns=headers) - print(df) - -def call_tilus(kernel_function, args, kwargs, grid, threads, params): - kernel_function(*args, **kwargs) - - -def tune_matmul(m, n, k): - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n #(m, n) - args = [m, n, k, a, b, c_actual] - tune_params = dict() - tune_params["block_m"] = [32, 64, 128, 256] - tune_params["block_n"] = [32, 64, 128, 256] - tune_params["block_k"] = [32, 64, 128] - tune_params["group_size_m"] = [4, 8] - tune_params["num_stages"] = [2, 3, 4] - tune_params["num_warps"] = [4, 8] - - - restrictions = [ - # tile size budget - "block_m * block_n <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "block_m <= 4 * block_n", - "block_n <= 4 * block_m", - - # large K only with reasonably large M/N - "not (block_k == 128 and block_m < 64 and block_n < 64)", - - # 32x32 requires 8 warps - "not (block_m == 32 and block_n == 32 and num_warps < 8)", - ] - - - answer = [None] * 6 - answer[-1] = c_expect.cpu() - atol = 1e-2 #m * 2**(-11) - - results, env = tune_kernel("MatmulGroupedOrdering", MatmulGroupedOrdering, size, args, tune_params, grid_div_x = ["block_m", "block_n"], - answer = answer, atol=atol, restrictions=restrictions, - lang="generic_python", call_function=call_tilus, - block_size_names=["block_m", "block_n", "block_k"], strategy="simulated_annealing") - - -if __name__ == "__main__": - #m, n, k = 4096, 4096, 4096 - m, n, k = 8192, 8192, 8192 - tune_matmul(m, n, k) - - #main() \ No newline at end of file diff --git a/examples/generic_python/matmul_old/triton_matmul.py b/examples/generic_python/matmul_old/triton_matmul.py deleted file mode 100644 index 238a61c9f..000000000 --- a/examples/generic_python/matmul_old/triton_matmul.py +++ /dev/null @@ -1,383 +0,0 @@ -import torch - -import triton -import triton.language as tl - -from kernel_tuner import tune_kernel -from kernel_tuner import run_kernel - -def get_cuda_autotune_config(): - return [ - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=3, - num_warps=8), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, - num_warps=2), - triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 8}, num_stages=5, - num_warps=2), - # Good config for fp8 inputs. - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, - num_warps=8), - triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=3, - num_warps=8), - triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 128, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4), - triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 32, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8}, num_stages=4, - num_warps=4) - ] - -''' -@triton.autotune( - configs=get_cuda_autotune_config(), - key=['M', 'N', 'K'], -) -''' -#@triton.jit -def matmul_kernel( - # Pointers to matrices - a_ptr, b_ptr, c_ptr, - # Matrix dimensions - M, N, K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` - # by to get the element one row down (A has M rows). - stride_am, stride_ak, # - stride_bk, stride_bn, # - stride_cm, stride_cn, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # - GROUP_SIZE_M: tl.constexpr, num_stages, num_warps# -): - """Kernel for computing the matmul C = A x B. - A has shape (M, K), B has shape (K, N) and C has shape (M, N) - """ - # ----------------------------------------------------------- - # Map program ids `pid` to the block of C it should compute. - # This is done in a grouped ordering to promote L2 data reuse. - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - # ----------------------------------------------------------- - # Add some integer bound assumptions. - # This helps to guide integer analysis in the backend to optimize - # load/store offset address calculation - tl.assume(pid_m >= 0) - tl.assume(pid_n >= 0) - tl.assume(stride_am > 0) - tl.assume(stride_ak > 0) - tl.assume(stride_bn > 0) - tl.assume(stride_bk > 0) - tl.assume(stride_cm > 0) - tl.assume(stride_cn > 0) - - # ---------------------------------------------------------- - # Create pointers for the first blocks of A and B. - # We will advance this pointer as we move in the K direction - # and accumulate - # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers - # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - # ----------------------------------------------------------- - # Iterate to compute a block of the C matrix. - # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block - # of fp32 values for higher accuracy. - # `accumulator` will be converted back to fp16 after the loop. - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - # Load the next block of A and B, generate a mask by checking the K dimension. - # If it is out of bounds, set it to 0. - a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) - # We accumulate along the K dimension. - accumulator = tl.dot(a, b, accumulator) - # Advance the ptrs to the next K block. - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - c = accumulator.to(tl.float16) - - # ----------------------------------------------------------- - # Write back the block of the output matrix C with masks. - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=c_mask) - - -def run_matmul(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - grid = lambda META: (triton.cdiv(m, META['BLOCK_SIZE_M']) * triton.cdiv(n, META['BLOCK_SIZE_N']), ) - - matmul_kernel[grid]( - a, b, c_actual, # - m, n, k, # - a.stride(0), a.stride(1), # - b.stride(0), b.stride(1), # - c_actual.stride(0), c_actual.stride(1), - 128, 256, 64, 8 - ) - - torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) - - -def run_matmul_kt(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n - - args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1)] - - params = {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M":4, "num_stages":3, "num_warps":4} - - result = run_kernel("matmul_kernel", matmul_kernel, size, args, params=params, grid_div_x=["BLOCK_SIZE_N", "BLOCK_SIZE_M"], - lang="generic_python", decorator="@triton.jit", call_function=call_triton, - block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) - c_res = result[2] - - assert torch.allclose(c_res, c_expect.cpu(), atol=1e-2, rtol=1e-1) - - - - - - -def call_triton(kernel_function, args, kwargs, grid, threads, params): - #print("using grid: ", grid) - #print("args: ", args) - #print("kwargs: ", kwargs) - kernel_function[grid](*args, **kwargs) - - - - -def check_time(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n - args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1)] - tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [64, 128] - tune_params["BLOCK_SIZE_N"] = [64, 128] - tune_params["BLOCK_SIZE_K"] = [32, 64] - tune_params["GROUP_SIZE_M"] = [4, 8] - tune_params["num_stages"] = [3] - tune_params["num_warps"] = [4] - - restrictions = [ - # tile size budget - "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", - "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", - - # large K only with reasonably large M/N - "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", - - # 32x32 requires 8 warps - "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", - ] - - grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] - - answer = [None] * 12 - answer[2] = c_expect.cpu() - atol = 1e-2 #m * 2**(-11) - - - - results_ours, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, - restrictions=restrictions, iterations=100, answer=answer, atol=atol, - lang="generic_python", decorator="@triton.jit", call_function=call_triton, - block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) - - - results_prev, _ = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, - restrictions=restrictions, iterations=100, answer=answer, atol=atol, - lang="triton", block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"]) - - - - import time - num_repeats = 100 - times_direct = [] - for config in results_prev: - - bs_m = config["BLOCK_SIZE_M"] - bs_n = config["BLOCK_SIZE_N"] - bs_k = config["BLOCK_SIZE_K"] - gs_m = config["GROUP_SIZE_M"] - num_stages = config["num_stages"] - num_warps = config["num_warps"] - - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - - grid = (triton.cdiv(m, bs_m) * triton.cdiv(n, bs_n, ), ) - jit_function = triton.jit(matmul_kernel) - - - jit_function[grid]( - a, b, c_actual, - m, n, k, - a.stride(0), a.stride(1), - b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1), - bs_m, bs_n, bs_k, gs_m, num_stages, num_warps - ) - - - torch.allclose(c_expect.cpu(), c_actual.cpu(), atol=1e-2) - - - - for i in range(num_repeats): - times = [] - - #c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - - torch.cuda.synchronize() - start = time.time() - jit_function[grid]( - a, b, c_actual, - m, n, k, - a.stride(0), a.stride(1), - b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1), - bs_m, bs_n, bs_k, gs_m, num_stages, num_warps - ) - - torch.cuda.synchronize() - times.append(time.time() - start) - - avg_time_ms = round((1000 * sum(times) / len(times)), 3) - times_direct.append(avg_time_ms) - print(f"BLOCK_SIZE_M={bs_m}, BLOCK_SIZE_N={bs_n}, BLOCK_SIZE_K={bs_k}, GROUP_SIZE_M={gs_m}, num_stages={num_stages}, num_warps={num_warps}, time={avg_time_ms}ms") - - - import matplotlib.pyplot as plt - - # Extract times - times_prev = [cfg['time'] for cfg in results_prev] - times_ours = [cfg['time'] for cfg in results_ours] - - # x-axis labels - configs = [f"config{i}" for i in range(len(times_prev))] - x = range(len(configs)) - - plt.figure(figsize=(10,6)) - plt.plot(configs, times_prev, marker='o', label='Triton tuned') - plt.plot(configs, times_ours, marker='s', label='Generic tuned') - plt.plot(configs, times_direct, marker='x', label='Direct') - plt.ylabel('Time (ms)') - plt.xlabel('Configuration') - plt.title('Kernel execution time per configuration') - plt.xticks(rotation=45) - plt.legend() - plt.grid(True) - plt.tight_layout() - plt.savefig("ouptut.png") - - - - - - -def tune_matmul(m, n, k): - a = torch.rand(m, k, dtype=torch.float16).cuda() - b = torch.rand(k, n, dtype=torch.float16).cuda() - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - size = m * n - args = [a, b, c_actual, m, n, k, a.stride(0), a.stride(1), b.stride(0), b.stride(1), - c_actual.stride(0), c_actual.stride(1)] - tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [64, 128, 256] - tune_params["BLOCK_SIZE_N"] = [64, 128, 256] - tune_params["BLOCK_SIZE_K"] = [32, 64, 128] - tune_params["GROUP_SIZE_M"] = [4, 8] - tune_params["num_stages"] = [2, 3, 4] - tune_params["num_warps"] = [4, 8] - - restrictions = [ - # tile size budget - "BLOCK_SIZE_M * BLOCK_SIZE_N <= 16384", - - # aspect ratio <= 4 (no max/min allowed, so expand manually) - "BLOCK_SIZE_M <= 4 * BLOCK_SIZE_N", - "BLOCK_SIZE_N <= 4 * BLOCK_SIZE_M", - - # large K only with reasonably large M/N - "not (BLOCK_SIZE_K == 128 and BLOCK_SIZE_M < 64 and BLOCK_SIZE_N < 64)", - - # 32x32 requires 8 warps - "not (BLOCK_SIZE_M == 32 and BLOCK_SIZE_N == 32 and num_warps < 8)", - ] - - grid_div = ["BLOCK_SIZE_N", "BLOCK_SIZE_M"] - - answer = [None] * 12 - answer[2] = c_expect.cpu() - - results, env = tune_kernel("matmul_kernel", matmul_kernel, size, args, tune_params, grid_div_x = grid_div, - answer = answer, atol=4.0, restrictions=restrictions, - lang="generic_python", decorator="@triton.jit", call_function=call_triton, - block_size_names=["BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K"], strategy="simulated_annealing") - - - - - -if __name__ == "__main__": - m, n, k = 8192, 8192, 8192 - #m, n, k = 4096, 4096, 4096 - #tune_matmul(m, n, k) - #check_time(m, n, k) - #run_matmul_kt(m, n, k) - - - - - - diff --git a/examples/generic_python/normalization/tilelang_norm.py b/examples/generic_python/normalization/tilelang_norm.py deleted file mode 100644 index 7c26d8f5f..000000000 --- a/examples/generic_python/normalization/tilelang_norm.py +++ /dev/null @@ -1,78 +0,0 @@ -# taken from https://github.com/tile-ai/tilelang/blob/main/examples/norm/rms_norm.py -# TODO not comparable -import torch -import tilelang -import tilelang.language as T - - -def rms_norm_splitk(M, N, blk_m, blk_k): - dtype = T.float - - @T.prim_func - def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)): - with T.Kernel(T.ceildiv(M, blk_m), threads=128) as bx: - A_shared = T.alloc_shared((blk_m, blk_k), dtype) - A_local = T.alloc_fragment((blk_m, blk_k), dtype) - A_powsum = T.alloc_fragment((blk_m,), dtype) - - num_k_step = T.ceildiv(N, blk_k) - T.clear(A_local) - for k in range(num_k_step): - T.copy(A[bx * blk_m, k * blk_k], A_shared) - for i, j in T.Parallel(blk_m, blk_k): - A_local[i, j] += A_shared[i, j] * A_shared[i, j] - T.reduce_sum(A_local, A_powsum, dim=1) - for i in T.Parallel(blk_m): - A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12) - - for k in range(num_k_step): - # reverse, better cache hit rate - T.copy(A[bx * blk_m, (num_k_step - 1 - k) * blk_k], A_shared) - for i, j in T.Parallel(blk_m, blk_k): - A_shared[i, j] *= A_powsum[i] - T.copy(A_shared, B[bx * blk_m, (num_k_step - 1 - k) * blk_k]) - - return main - - -@tilelang.jit(out_idx=[-1], pass_configs={"tl.disable_tma_lower": True}) -def rms_norm(M, N, blk_m): - dtype = T.float - - @T.prim_func - def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)): - with T.Kernel(T.ceildiv(M, blk_m), threads=128) as bx: - A_shared = T.alloc_shared((blk_m, N), dtype) - A_pow_local = T.alloc_fragment((blk_m, N), dtype) - A_local = T.alloc_fragment((blk_m, N), dtype) - A_powsum = T.alloc_fragment((blk_m,), dtype) - - T.copy(A[bx * blk_m : (bx + 1) * blk_m, :], A_shared) - T.copy(A_shared, A_local) - for i, j in T.Parallel(blk_m, N): - A_pow_local[i, j] = A_local[i, j] * A_local[i, j] - T.reduce_sum(A_pow_local, A_powsum, dim=1) - for i in T.Parallel(blk_m): - A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12) - for i, j in T.Parallel(blk_m, N): - A_local[i, j] *= A_powsum[i] - T.copy(A_local, B[bx * blk_m : (bx + 1) * blk_m, :]) - - return main - - -def ref_program(x): - return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-12) - - -if __name__ == "__main__": - M, N, blk_m, blk_k = 8192, 8192, 1, 512 - kernel = rms_norm(M, N, blk_m) - profiler = kernel.get_profiler() - profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01) - print("All checks pass.") - - latency = profiler.do_bench(ref_program, warmup=500) - print("Ref: {:.2f} ms".format(latency)) - latency = profiler.do_bench(warmup=500) - print("Tile-lang: {:.2f} ms".format(latency)) \ No newline at end of file diff --git a/examples/generic_python/normalization/tilus_norm.py b/examples/generic_python/normalization/tilus_norm.py deleted file mode 100644 index cfe569e80..000000000 --- a/examples/generic_python/normalization/tilus_norm.py +++ /dev/null @@ -1,132 +0,0 @@ -# taken from https://github.com/NVIDIA/tilus/blob/main/examples/norm/layer_norm.py -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -import pandas -import tilus -import torch -from tilus import float16, float32, int32 -from tilus.utils import benchmark_func, cdiv - - -@tilus.autotune("block_m", [1, 8]) -@tilus.autotune("block_n", [128, 256, 512, 1024]) -@tilus.autotune("warps", [2, 4, 8]) -class LayerNorm(tilus.Script): - """Forward-only layer normalization tilus kernel. - - This implements the per-row LayerNorm used in many transformer blocks. It - computes: y = (x - mean) / sqrt(var + eps) * gamma + beta - - Only the forward is provided. - """ - - def __init__(self, block_m: int, block_n: int, warps: int): - super().__init__() - self.block_m: int = block_m - self.block_n: int = block_n - self.warps: int = warps - - def __call__( - self, - m_size: int, - n_size: int32, - x_ptr: ~float16, - gamma_ptr: ~float16, - beta_ptr: ~float16, - y_ptr: ~float16, - eps: float, - ): - self.attrs.blocks = (cdiv(m_size, self.block_m),) - self.attrs.warps = self.warps - - offset_m = self.blockIdx.x * self.block_m - - g_x = self.global_view(x_ptr, dtype=float16, shape=[m_size, n_size]) - g_y = self.global_view(y_ptr, dtype=float16, shape=[m_size, n_size]) - g_gamma = self.global_view(gamma_ptr, dtype=float16, shape=[n_size]) - g_beta = self.global_view(beta_ptr, dtype=float16, shape=[n_size]) - - # Register accumulators for mean and variance (computed in float32) - r_sum = self.register_tensor( - dtype=float32, shape=[self.block_m, self.block_n], init=0.0 - ) - r_square = self.register_tensor( - dtype=float32, shape=[self.block_m, self.block_n], init=0.0 - ) - - # first pass: compute mean and variance - for offset_n in range(0, n_size, self.block_n): - r_x = self.load_global( - g_x, offsets=[offset_m, offset_n], shape=[self.block_m, self.block_n] - ).to(float32) # [block_m, block_n] - r_sum = r_sum + r_x - r_square = r_square + self.square(r_x) - - # finalize mean and variance - r_mean = self.sum(r_sum, dim=1, keepdim=True) / n_size # [block_m, 1] - r_var = ( - self.sum(r_square, dim=1, keepdim=True) / n_size - r_mean * r_mean - ) # [block_m, 1], var = E[x^2] - (E[x])^2 - r_rstd = self.rsqrt(r_var + eps) - - # second pass: y = (x - mean) * rstd * gamma + beta - for offset_n in range(0, n_size, self.block_n): - r_x = self.load_global( - g_x, offsets=[offset_m, offset_n], shape=[self.block_m, self.block_n] - ).to(float32) # [block_m, block_n] - r_gamma = self.load_global( - g_gamma, offsets=[offset_n], shape=[self.block_n] - ).to(float32) # [block_n] - r_beta = self.load_global( - g_beta, offsets=[offset_n], shape=[self.block_n] - ).to(float32) # [block_n] - r_x_hat = (r_x - r_mean) * r_rstd - r_y = r_x_hat * r_gamma + r_beta - self.store_global(g_y, r_y.to(float16), offsets=[offset_m, offset_n]) - - -def main(): - headers = ["m_size", "n_size", "dtype", "torch (ms)", "tilus (ms)"] - rows = [] - for i in [1, 2, 4, 8]: - m_size = n_size = 1024 * i - - tilus_layer_norm = LayerNorm() - - x = (torch.rand(m_size, n_size, dtype=torch.float16).cuda() - 0.5) * 2.0 - gamma = torch.rand(n_size, dtype=torch.float16).cuda() - beta = torch.rand(n_size, dtype=torch.float16).cuda() - y_actual = torch.empty_like(x) - - tilus_layer_norm(m_size, n_size, x, gamma, beta, y_actual, 1e-5) - y_expected = torch.nn.functional.layer_norm( - x, normalized_shape=[n_size], weight=gamma, bias=beta, eps=1e-5 - ) - - torch.testing.assert_close(y_actual, y_expected, atol=1e-2, rtol=1e-2) - - rows.append( - [ - m_size, - n_size, - "float16", - benchmark_func( - lambda: torch.nn.functional.layer_norm( - x, normalized_shape=[n_size], weight=gamma, bias=beta, eps=1e-5 - ) - ), - benchmark_func( - lambda: tilus_layer_norm( - m_size, n_size, x, gamma, beta, y_actual, 1e-5 - ) - ), - ] - ) - print(f"LayerNorm forward matches reference for size ({m_size}, {n_size})") - - df = pandas.DataFrame(rows, columns=headers) - print(df) - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/examples/generic_python/normalization/triton_norm.py b/examples/generic_python/normalization/triton_norm.py deleted file mode 100644 index 365f541c4..000000000 --- a/examples/generic_python/normalization/triton_norm.py +++ /dev/null @@ -1,66 +0,0 @@ - -#Example taken from https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -import torch - -import triton -import triton.language as tl - -try: - # This is https://github.com/NVIDIA/apex, NOT the apex on PyPi, so it - # should not be added to extras_require in setup.py. - import apex - HAS_APEX = True -except ModuleNotFoundError: - HAS_APEX = False - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - - -@triton.jit -def _layer_norm_fwd_fused( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - Mean, # pointer to the mean - Rstd, # pointer to the 1/std - stride, # how much to increase the pointer when moving by 1 row - N, # number of columns in X - eps, # epsilon to avoid division by zero - BLOCK_SIZE: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - Y += row * stride - X += row * stride - # Compute mean - mean = 0 - _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) - _mean += a - mean = tl.sum(_mean, axis=0) / N - # Compute variance - _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) - x = tl.where(cols < N, x - mean, 0.) - _var += x * x - var = tl.sum(_var, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - # Write mean / rstd - tl.store(Mean + row, mean) - tl.store(Rstd + row, rstd) - # Normalize and apply linear transformation - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - mask = cols < N - w = tl.load(W + cols, mask=mask) - b = tl.load(B + cols, mask=mask) - x = tl.load(X + cols, mask=mask, other=0.).to(tl.float32) - x_hat = (x - mean) * rstd - y = x_hat * w + b - # Write output - tl.store(Y + cols, y, mask=mask) \ No newline at end of file diff --git a/examples/generic_python/numba_vec_add.py b/examples/generic_python/numba_vec_add.py index b5a2c56f3..3bc90de86 100644 --- a/examples/generic_python/numba_vec_add.py +++ b/examples/generic_python/numba_vec_add.py @@ -1,10 +1,8 @@ from numba import cuda -import torch -from kernel_tuner import tune_kernel, run_kernel import numpy as np -from pathlib import Path -FULL_PATH = Path(__file__).resolve() +from kernel_tuner import tune_kernel +from call_functions import call_numba @cuda.jit @@ -16,16 +14,6 @@ def f(a, b, c): c[tid] = a[tid] + b[tid] -def call_numba(kernel_function, args, kwargs, grid, threads): - numba_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - numba_args.append(cuda.as_cuda_array(arg)) - else: - numba_args.append(arg) - kernel_function[grid, threads](*args, **kwargs) - - def tune(): N = 100000 @@ -41,7 +29,7 @@ def tune(): results, env = tune_kernel( kernel_name="f", - kernel_source=FULL_PATH, + kernel_source=__file__, problem_size=N, arguments=args, tune_params=tune_params, diff --git a/examples/generic_python/pallas_vec_add.py b/examples/generic_python/pallas_vec_add.py deleted file mode 100644 index 874392dd9..000000000 --- a/examples/generic_python/pallas_vec_add.py +++ /dev/null @@ -1,23 +0,0 @@ -# Could not get working - -from functools import partial - -import jax -from jax.experimental import pallas as pl -import jax.numpy as jnp -import numpy as np - -def add_vectors_kernel(x_ref, y_ref, o_ref): - x, y = x_ref[...], y_ref[...] - o_ref[...] = x + y - - -@jax.jit -def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array: - return pl.pallas_call( - add_vectors_kernel, - out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype) - )(x, y) - - -add_vectors(jnp.arange(8), jnp.arange(8)) \ No newline at end of file diff --git a/examples/generic_python/tilelang_vec_add.py b/examples/generic_python/tilelang_vec_add.py index 99463d112..38b4ab309 100644 --- a/examples/generic_python/tilelang_vec_add.py +++ b/examples/generic_python/tilelang_vec_add.py @@ -1,10 +1,10 @@ import tilelang import tilelang.language as T import torch + from kernel_tuner import tune_kernel -from pathlib import Path +from call_functions import call_tilelang -FULL_PATH = Path(__file__).resolve() @tilelang.jit # infers target from tensors at first call def add(N: int, dtype: str = 'float32', block: int = 256,): @@ -24,23 +24,6 @@ def add_kernel( return add_kernel -def run_normal(): - # Host side (PyTorch shown; NumPy/DLPack also supported) - N = 1 << 20 - A = torch.randn(N, device='cuda', dtype=torch.float32) - B = torch.randn(N, device='cuda', dtype=torch.float32) - C = torch.empty(N, device='cuda', dtype=torch.float32) - - kernel = add(N) - kernel(A, B, C) # runs on GPU - torch.testing.assert_close(C, A + B) - print("done") - - -def call_tilelang(kernel_function, args, kwargs): - compiled_kernel = kernel_function(**kwargs) # cached, so second time only cache lookup is performed - compiled_kernel(*args) - def tune(): N = 1 << 20 A = torch.randn(N, device='cuda', dtype=torch.float32) @@ -54,7 +37,7 @@ def tune(): answer = [None, None, (A + B).cpu()] - res, env = tune_kernel("add", FULL_PATH, N, args, tune_params, lang="generic_python", + res, env = tune_kernel("add", __file__, N, args, tune_params, lang="generic_python", call_function=call_tilelang, answer=answer) if __name__ == "__main__": diff --git a/examples/generic_python/tilus_naive_matmul.py b/examples/generic_python/tilus_naive_matmul.py deleted file mode 100644 index c53c22f73..000000000 --- a/examples/generic_python/tilus_naive_matmul.py +++ /dev/null @@ -1,127 +0,0 @@ -from tilus import float16, float32, int32 -from tilus.utils import cdiv -import tilus -from kernel_tuner import tune_kernel, run_kernel -import math -import torch -from pathlib import Path - -FULL_PATH = Path(__file__).resolve() - - -class MatmulV0(tilus.Script): - def __init__(self): - super().__init__() - # we define three hyperparameters: ``block_m``, ``block_n``, and ``block_k`` to determine the tile size on - # m, n, and k dimensions for each `thread block` of the kernel. - self.block_m = 64 - self.block_n = 64 - self.block_k = 16 - - def __call__( - self, - m_size: int32, # the size of the m dimension of the input matrix A and output matrix C - n_size: int, # the size of the n dimension of the input matrix B and output matrix C - k_size: int, # the size of the k dimension of the input matrix A and B - a_ptr: ~float16, # the pointer to the input matrix A, which is a 2D tensor of shape [m_size, k_size] - b_ptr: ~float16, # the pointer to the input matrix B, which is a 2D tensor of shape [k_size, n_size] - c_ptr: ~float16, # the pointer to the output matrix C, which is a 2D tensor of shape [m_size, n_size] - ): - self.attrs.blocks = [ - cdiv(m_size, self.block_m), # the x dimension size of the grid - cdiv(n_size, self.block_n), # the y dimension size of the grid - ] - self.attrs.warps = 1 # the number of warps per thread block, must be a compile-time known integer - - # define two int32 variables to store the offsets of the m and n dimensions for the current thread block. - offset_m: int32 = self.block_m * self.blockIdx.x - offset_n: int32 = self.block_n * self.blockIdx.y - - # create two global tensors `ga` and `gb` to represent the input matrices A and B, respectively. - ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) - gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) - - # create a register tensor `acc` to accumulate the results of the matrix multiplication. - acc = self.register_tensor( - dtype=float32, shape=[self.block_m, self.block_n], init=0.0 - ) - - # iterate over the k dimension in blocks of size `block_k`. - for k in range(cdiv(k_size, self.block_k)): - # calculate the offset for the current block in the k dimension - offset_k = k * self.block_k - - # load a block of matrix A and B into register tensors `a` and `b`. - a = self.load_global( - ga, offsets=[offset_m, offset_k], shape=[self.block_m, self.block_k] - ) - b = self.load_global( - gb, offsets=[offset_k, offset_n], shape=[self.block_k, self.block_n] - ) - - # perform the dot product: acc = a @ b + acc - self.dot(a, b, acc, out=acc) - - # after the loop, we cast the accumulated result `acc` to float16 type and store it back to the output matrix C. - acc_f16 = self.cast(acc, dtype=float16) - gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - self.store_global(gc, acc_f16, offsets=[offset_m, offset_n]) - - -def call_tilus(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) - -def main(): - m, n, k = 4096, 4096, 4096 - - # create an instance of the kernel we have just defined - matmul = MatmulV0() - - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - ''' - torch.cuda.synchronize() - # launch the kernel by passing required arguments - matmul(m, n, k, a, b, c_actual) - torch.cuda.synchronize() - - # check correctness - torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) - ''' - - args = [m, n, k, a, b, c_actual] - tune_params = dict() - tune_params["block_m"] = [16, 32, 64, 128, 256] - tune_params["block_n"] = [16, 32, 64, 128, 256] - tune_params["block_k"] = [16, 32, 64, 128, 256] - - - restrictions = [ - "block_m * block_n <= 4096", - "block_m * block_k <= 2048", - "block_k * block_n <= 4096", - "block_k >= 16", - ] - - results, env = tune_kernel( - kernel_name="MatmulV0", - kernel_source=FULL_PATH, - problem_size=[m, n], - arguments=args, - tune_params=tune_params, - lang="generic_python", - answer=[None, None, None, None, None, c_expect.cpu()], - restrictions=restrictions, - block_size_names=["block_m", "block_n", "block_k"], - call_function=call_tilus, - strategy="simulated_annealing" - ) - - - - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/examples/generic_python/tilus_splitk_matmul.py b/examples/generic_python/tilus_splitk_matmul.py deleted file mode 100644 index 0695dcceb..000000000 --- a/examples/generic_python/tilus_splitk_matmul.py +++ /dev/null @@ -1,255 +0,0 @@ -import tilus -from tilus import float16, float32, int32 -from tilus.utils import cdiv, benchmark_func -import torch -import math -from kernel_tuner import tune_kernel, run_kernel -from pathlib import Path - -FULL_PATH = Path(__file__).resolve() - -class MatmulV5(tilus.Script): - def __init__(self, block_m=None, block_n=None, block_k=None, - num_warps=None, num_stages=None, split_k_factor=None): - super().__init__() - self.block_m = block_m - self.block_n = block_n - self.block_k = block_k - self.num_warps = num_warps - self.num_stages = num_stages - self.split_k_factor = split_k_factor - - - def __call__( - self, - m_size: int32, - n_size: int, - k_size: int, - a_ptr: ~float16, - b_ptr: ~float16, - c_ptr: ~float16, - ): - self.attrs.blocks = [ - cdiv(m_size, self.block_m), - cdiv(n_size, self.block_n), - self.split_k_factor, - ] - self.attrs.warps = self.num_warps - - # the k_size for each thread block - block_k_size = ( - cdiv(cdiv(k_size, self.split_k_factor), self.block_k) * self.block_k - ) - start_offset_k = self.blockIdx.z * block_k_size - end_offset_k = min(start_offset_k + block_k_size, k_size) - - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k - offset_m: int32 = block_m * self.blockIdx.x - offset_n: int32 = block_n * self.blockIdx.y - - ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) - gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) - sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) - sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) - acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) - - for stage in range(self.num_stages - 1): - offset_k = start_offset_k + stage * self.block_k - self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) - self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) - self.copy_async_commit_group() - - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - current_stage: int32 = 0 - preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range( - start_offset_k, end_offset_k, block_k, unroll=self.num_stages - ): - # computation for current tile - a = self.load_shared(sa[current_stage]) - b = self.load_shared(sb[current_stage]) - self.dot(a, b, acc, out=acc) - - # preload the next tile of A and B into shared memory - preload_offset_k = offset_k + (self.num_stages - 1) * block_k - if preload_offset_k < end_offset_k: - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], - ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], - ) - self.copy_async_commit_group() - - # update the stage - current_stage = (current_stage + 1) % self.num_stages - preload_stage = (preload_stage + 1) % self.num_stages - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - # free the shared memory tensors for A and B - self.free_shared(sa) - self.free_shared(sb) - - # cast the accumulator to float16 and change the register tensor's layout - sc = self.shared_tensor(dtype=float16, shape=[block_m, block_n]) - casted_acc = self.cast(acc, dtype=float16) - self.store_shared(sc, casted_acc) - self.sync() - rc = self.load_shared(sc) - self.free_shared(sc) - - m_blocks, n_blocks = cdiv(m_size, block_m), cdiv(n_size, block_n) - gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - if self.split_k_factor == 0: - self.store_global(gc, rc, offsets=[offset_m, offset_n]) - else: - semaphores = self.global_tensor( - dtype=int32, shape=[m_blocks, n_blocks], requires_clean=True - ) - semaphore: ~int32 = ~semaphores[self.blockIdx.x, self.blockIdx.y] - - # load and accumulate the partial result in global memory - if self.blockIdx.z > 0: - self.lock_semaphore(semaphore, value=self.blockIdx.z) - partial_rc = self.load_global( - gc, offsets=[offset_m, offset_n], shape=[block_m, block_n] - ) - self.add(rc, partial_rc, out=rc) - - # store the result to global memory and release the semaphore - self.store_global(gc, rc, offsets=[offset_m, offset_n]) - - # release the semaphore - self.sync() # we need to make sure the previous store_global is finished - self.release_semaphore( - semaphore, value=(self.blockIdx.z + 1) % self.split_k_factor - ) - - -def call_tilus(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) - - -def without_kernel_tuner(): - tilus.option.clear_cache = True - - tilus.option.verbose_autotune = True - m, n, k = 4096, 4096, 4096 - - # create an instance of the kernel we have just defined - matmul = MatmulV5() - - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - - torch.cuda.synchronize() - # launch the kernel by passing required arguments - matmul(m, n, k, a, b, c_actual) - torch.cuda.synchronize() - - # check correctness - torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) - - - import pandas - rows = [] - headers = ["m", "n", "k", "name", "latency (ms)", "tflops"] - # benchmark - for name, func in [ - ("torch", lambda: torch.matmul(a, b, out=c_expect)), - ("tilus", lambda: matmul(m, n, k, a, b, c_actual)), - ]: - latency = benchmark_func(func, warmup=5, repeat=20) - tflops = 2 * m * n * k / latency * 1e-9 - rows.append([m, n, k, name, latency, tflops]) - - df = pandas.DataFrame(rows, columns=headers) - print(df) - - - -#best performing configuration: -#block_m=128, block_n=64, block_k=16, num_warps=4, num_stages=4, split_k_factor=1, time=2.027ms - -def main(): - m, n, k = 4096, 4096, 4096 - - # create an instance of the kernel we have just defined - #matmul = MatmulV5() - - a = (torch.rand(m, k, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - b = (torch.rand(k, n, dtype=torch.float16).cuda() - 0.5) / math.sqrt(k) - c_actual = torch.empty(m, n, dtype=torch.float16).cuda() - c_expect = a @ b - - ''' - torch.cuda.synchronize() - # launch the kernel by passing required arguments - matmul(m, n, k, a, b, c_actual) - torch.cuda.synchronize() - - # check correctness - torch.testing.assert_close(c_expect, c_actual, atol=1e-2, rtol=1e-2) - ''' - - args = [m, n, k, a, b, c_actual] - tune_params = dict() - tune_params["block_m"] = [32, 64, 128] #[16, 32, 64, 128, 256] - tune_params["block_n"] = [32, 64, 128] #[16, 32, 64, 128, 256] - tune_params["block_k"] = [16, 32] #[16, 32, 64, 128, 256] - tune_params["num_warps"] = [4, 8] #[2, 4, 8, 16] - tune_params["num_stages"] = [3, 4, 5] #[2, 3, 4, 5, 6] - tune_params["split_k_factor"] = [1, 4, 12, 16] #[1, 4, 12, 16, 20] - -#@tilus.autotune("num_warps", [4, 8]) -#@tilus.autotune("block_m, block_n", [(128, 128), (128, 64), (64, 128), (32, 256)]) -#@tilus.autotune("block_k", [16, 32]) -#@tilus.autotune("num_stages", [3, 4, 5]) -#@tilus.autotune("split_k_factor", [1, 4, 12, 16]) - - ''' - restrictions = [ - "block_m * block_n <= 4096", - "block_m * block_k <= 2048", - "block_k * block_n <= 4096", - "block_k >= 16", - "2 * (num_stages * block_k * (block_m + block_n) + block_m * block_n) <= 65536", # shared mem - "num_warps * 32 <= 1024", - "block_k * split_k_factor <= 4096", - ] - ''' - restrictions = ["block_m * block_n >= 8192", "block_m * block_n <= 16384"] - - - - results, env = tune_kernel( - kernel_name="MatmulV5", # This has to be a string of the actual name. TODO is this always the case? - kernel_source=FULL_PATH, - problem_size=[m, n], - arguments=args, - tune_params=tune_params, - lang="generic_python", - answer=[None, None, None, None, None, c_expect.cpu()], - atol=1e-2, - restrictions=restrictions, - block_size_names=["block_m", "block_n", "block_k"], - call_function=call_tilus, - #strategy="random_sample", - ) - - - - -if __name__ == "__main__": - main() - #without_kernel_tuner() \ No newline at end of file diff --git a/examples/generic_python/tilus_vec_add.py b/examples/generic_python/tilus_vec_add.py index d00c77972..8ee96aa07 100644 --- a/examples/generic_python/tilus_vec_add.py +++ b/examples/generic_python/tilus_vec_add.py @@ -1,11 +1,13 @@ import tilus from tilus import float32, int32 -from tilus.utils import cdiv, benchmark_func +from tilus.utils import cdiv import torch + from kernel_tuner import tune_kernel, run_kernel -from pathlib import Path +from call_functions import call_tilus + + -FULL_PATH = Path(__file__).resolve() class VecAddV(tilus.Script): def __init__(self, block_size_x=None, num_warps=None): @@ -38,8 +40,7 @@ def __call__( self.store_global(gc, c, offsets=[offset]) -def call_tilus(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) + def tune(size): @@ -57,7 +58,7 @@ def tune(size): results, env = tune_kernel( kernel_name="VecAddV", - kernel_source=FULL_PATH, + kernel_source=__file__, problem_size=size, arguments=args, tune_params=tune_params, @@ -79,7 +80,7 @@ def run(size): results = run_kernel( kernel_name="VecAddV", - kernel_source=FULL_PATH, + kernel_source=__file__, problem_size=size, arguments=args, params={"block_size_x": 32}, @@ -92,23 +93,9 @@ def run(size): assert torch.allclose(results[-1], c_expect) -def tune_with_builtin(size): - TunedVecAdd = tilus.autotune("block_size_x", [32, 64, 128, 256, 512, 1024])(VecAddV) - vecadd = TunedVecAdd() - - a = torch.randn(size, dtype=torch.float32).cuda() - b = torch.randn(size, dtype=torch.float32).cuda() - c = torch.empty(size, dtype=torch.float32).cuda() - c_expect = a + b - - vecadd(size, a, b, c) # This is where the actual tuning takes place - torch.cuda.synchronize() - - torch.testing.assert_close(c_expect, c) - if __name__ == "__main__": size = 10000000 - #tune(size) - #run(size) - tune_with_builtin(size) + tune(size) + run(size) + diff --git a/examples/generic_python/triton_vec_add.py b/examples/generic_python/triton_vec_add.py index 00bd9dfbc..cafba79a9 100644 --- a/examples/generic_python/triton_vec_add.py +++ b/examples/generic_python/triton_vec_add.py @@ -2,11 +2,9 @@ import torch import triton import triton.language as tl -from pathlib import Path from kernel_tuner import tune_kernel, run_kernel - -FULL_PATH = Path(__file__).resolve() +from call_functions import call_triton @triton.jit def add_op(x, y): @@ -30,9 +28,6 @@ def add_kernel(x_ptr, # *Pointer* to first input vector. tl.store(output_ptr + offsets, output, mask=mask) -def call_triton(kernel_function, args, kwargs, grid): - kernel_function[grid](*args, **kwargs) - def tune(): size = 10000000 @@ -48,7 +43,7 @@ def tune(): tune_params["block_size_x"] = [2**i for i in range(11)] - result = run_kernel("add_kernel", FULL_PATH, size, args, {"block_size_x": 256}, + result = run_kernel("add_kernel", __file__, size, args, {"block_size_x": 256}, lang="generic_python", call_function=call_triton) assert np.allclose(c_expect.cpu(), result[2]) @@ -56,7 +51,7 @@ def tune(): results, env = tune_kernel( kernel_name="add_kernel", - kernel_source=FULL_PATH, + kernel_source=__file__, problem_size=size, arguments=args, tune_params=tune_params, @@ -65,8 +60,6 @@ def tune(): call_function=call_triton, ) - print(results) - if __name__ == "__main__": diff --git a/examples/generic_python/warp_vec_add.py b/examples/generic_python/warp_vec_add.py index 4f94c91d2..81d7fb19e 100644 --- a/examples/generic_python/warp_vec_add.py +++ b/examples/generic_python/warp_vec_add.py @@ -1,11 +1,13 @@ import warp as wp import numpy as np -from kernel_tuner import tune_kernel, run_kernel import torch -from pathlib import Path + +from kernel_tuner import tune_kernel +from call_functions import call_warp + wp.init() -FULL_PATH = Path(__file__).resolve() + @wp.func def add_op(x: float, y: float): @@ -55,7 +57,7 @@ def tune(): results, env = tune_kernel( kernel_name="vec_add", - kernel_source=FULL_PATH, + kernel_source=__file__, problem_size=n, arguments=args, tune_params=tune_params, From 05f1332c00730c6cb846b8519784cefa4187eca6 Mon Sep 17 00:00:00 2001 From: Imke van Ooijen Date: Thu, 2 Jul 2026 17:58:09 +0200 Subject: [PATCH 15/44] more simple basic triton kernel --- .../generic_python/matmul/triton_matmul.py | 85 +++++++------------ 1 file changed, 32 insertions(+), 53 deletions(-) diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/generic_python/matmul/triton_matmul.py index d5522c627..9f34356ab 100644 --- a/examples/generic_python/matmul/triton_matmul.py +++ b/examples/generic_python/matmul/triton_matmul.py @@ -5,61 +5,46 @@ from kernel_tuner import tune_kernel from examples.generic_python.call_functions import call_triton + + @triton.jit def matmul_basic( - a_ptr, b_ptr, c_ptr, # Pointers - M, N, K, # Matrix sizes - stride_am, stride_ak, # Strides - stride_bk, stride_bn, - stride_cm, stride_cn, - BLOCK_SIZE_M: tl.constexpr, # Tile sizes - BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, + A_ptr, B_ptr, C_ptr, + M, N, K, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, ): - # Each program computes one BLOCK_SIZE_M x BLOCK_SIZE_N tile of C - pid_m = tl.program_id(axis=0) - pid_n = tl.program_id(axis=1) + m = tl.program_id(0) + n = tl.program_id(1) - # Compute row/column indices for the tile - offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - offs_k = tl.arange(0, BLOCK_SIZE_K) + # Base offsets for this tile + offs_m = m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - # Create pointers to A and B tiles - a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + # Accumulator for C tile + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - # Accumulator - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - - # Loop over K dimension + # Loop over K in chunks for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load( - a_ptrs, - mask=(offs_am[:, None] < M) & (offs_k[None, :] < K - k * BLOCK_SIZE_K), - other=0.0, - ) - - b = tl.load( - b_ptrs, - mask=(offs_k[:, None] < K - k * BLOCK_SIZE_K) & (offs_bn[None, :] < N), - other=0.0, - ) + offs_k = k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) - accumulator = tl.dot(a, b, accumulator) + # Load tiles of A and B + a_ptrs = A_ptr + offs_m[:, None] * K + offs_k[None, :] # [BLOCK_M, BLOCK_K] + b_ptrs = B_ptr + offs_k[:, None] * N + offs_n[None, :] # [BLOCK_K, BLOCK_N] - # advance K - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk + mask_a = (offs_m[:, None] < M) & (offs_k[None, :] < K) + mask_b = (offs_k[:, None] < K) & (offs_n[None, :] < N) - # Store result - c = accumulator.to(tl.float16) - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=mask) + a = tl.load(a_ptrs, mask=mask_a, other=0.0) + b = tl.load(b_ptrs, mask=mask_b, other=0.0) + acc += tl.dot(a, b) + + # Store result + c_ptrs = C_ptr + offs_m[:, None] * N + offs_n[None, :] + mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(c_ptrs, acc, mask=mask_c) @@ -177,9 +162,6 @@ def run_basic(M, N, K): matmul_basic[grid]( A, B, C, M, N, K, - A.stride(0), A.stride(1), - B.stride(0), B.stride(1), - C.stride(0), C.stride(1), BLOCK_SIZE_M=BLOCK_SIZE_M, BLOCK_SIZE_N=BLOCK_SIZE_N, BLOCK_SIZE_K=BLOCK_SIZE_K, @@ -230,9 +212,6 @@ def tune_basic(M, N, K): args = [A, B, C, M, N, K, - A.stride(0), A.stride(1), - B.stride(0), B.stride(1), - C.stride(0), C.stride(1), ] tune_params = dict() @@ -249,7 +228,7 @@ def tune_basic(M, N, K): arguments=args, tune_params=tune_params, lang="generic_python", - answer=[None, None, C_ref.cpu(), None, None, None, None, None, None, None, None, None], + answer=[None, None, C_ref.cpu(), None, None, None], atol=M * 2**(-11), block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_M"], strategy = "bayes_opt", @@ -306,7 +285,7 @@ def tune_opt(M, N, K): if __name__ == "__main__": M, N, K = 4096, 4096, 4096 run_basic(M, N, K) - run_opt(M, N, K) + #run_opt(M, N, K) tune_basic(M, N, K) - tune_opt(M, N, K) + #tune_opt(M, N, K) From 2ff0fdeb448103fb06350af0c63e251237049e6b Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Fri, 18 Sep 2026 19:58:41 +0200 Subject: [PATCH 16/44] fix output verification for torch answers Only convert Julia arrays in the expected answers to numpy arrays, other answers such as PyTorch tensors are kept as-is. Pass equal_nan to allclose as a bool, as torch.allclose does not accept a tensor. --- kernel_tuner/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index d9c288aeb..742d5548a 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -433,7 +433,7 @@ def check_kernel_output(self, func, gpu_args, instance, answer, atol, verify, ve # convert juliacall arrays to numpy arrays where necessary if answer is not None: - answer = [None if a is None else np.array(a) for a in util.possible_julia_vector_to_list(answer)] + answer = [np.array(a) if util.is_julia_array(a) else a for a in util.possible_julia_vector_to_list(answer)] for i, arg in enumerate(instance.arguments): if util.is_julia_array(arg) and isinstance(answer[i], np.ndarray): instance.arguments[i] = np.array(arg, dtype=answer[i].dtype) @@ -883,7 +883,7 @@ def _flatten(a): else np ) expected_nan = lib.isnan(expected) - output_test = lib.allclose(expected, result, atol=atol, equal_nan=expected_nan.any()) + output_test = lib.allclose(expected, result, atol=atol, equal_nan=bool(expected_nan.any())) if expected_nan.any(): warn( f"Answer contains {expected_nan.sum()} NaNs. NaN values will now be considered equal in comparison." From c032b0e1b7c83cc1eac05fc961223c0426808645 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Fri, 18 Sep 2026 19:58:43 +0200 Subject: [PATCH 17/44] support CUDA Array Interface arguments in nvcuda and pycuda backends Device arrays implementing __cuda_array_interface__ (e.g. PyTorch CUDA tensors, CuPy arrays) and CPU PyTorch tensors are now copied to a separate device allocation in both backends, like numpy arrays. This adds PyTorch support to the nvcuda backend and replaces the Holder wrapper in the pycuda backend, so outputs are correctly reset between kernel runs and the user's tensors are no longer modified. Neither backend imports torch anymore. Also removes leftover debug prints from the pycuda backend, adds regression tests, and updates the changelog. --- CHANGELOG.md | 5 ++ kernel_tuner/backends/backend.py | 29 ++++++++ kernel_tuner/backends/nvcuda.py | 28 +++++--- kernel_tuner/backends/pycuda.py | 60 ++++++----------- kernel_tuner/core.py | 14 ++-- kernel_tuner/interface.py | 3 +- test/context.py | 2 +- test/test_torch_arguments.py | 111 +++++++++++++++++++++++++++++++ 8 files changed, 195 insertions(+), 57 deletions(-) create mode 100644 test/test_torch_arguments.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f8e66ef3..1db199f1d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,11 @@ This project adheres to [Semantic Versioning](http://semver.org/). ## Unreleased - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers +- Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, and CuTe, using `lang="generic_python"` and a `call_function` +- Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend +- Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel +- Fix output verification failing for PyTorch tensors as expected answers +- Fix outputs not being reset between kernel runs for PyTorch tensors in the PyCUDA backend ## [1.5.0] - 2026-08-20 - Kernel Tuner can now be used from Julia, [see KernelTuner.jl](https://github.com/KernelTuner/KernelTuner.jl) diff --git a/kernel_tuner/backends/backend.py b/kernel_tuner/backends/backend.py index 6063dbb43..4b4e4aaba 100644 --- a/kernel_tuner/backends/backend.py +++ b/kernel_tuner/backends/backend.py @@ -3,6 +3,35 @@ from abc import ABC, abstractmethod +import numpy as np + + +def get_device_array(arg): + """Return (device pointer, size in bytes) if arg implements the CUDA Array Interface, None otherwise. + + This covers for example PyTorch CUDA tensors, CuPy arrays, and Numba device arrays, + without having to import any of these libraries. + """ + # CPU tensors raise AttributeError on this property, which getattr turns into None + cai = getattr(arg, "__cuda_array_interface__", None) + if cai is None: + return None + shape = tuple(cai["shape"]) + itemsize = np.dtype(cai["typestr"]).itemsize + strides = cai.get("strides") + if strides is not None: + contiguous = tuple(int(np.prod(shape[i + 1 :])) * itemsize for i in range(len(shape))) + if tuple(strides) != contiguous: + raise ValueError("Device arrays passed as kernel arguments must be C-contiguous") + return cai["data"][0], int(np.prod(shape)) * itemsize + + +def is_host_array(arg): + """Return True if arg is a numpy array or a host array that converts to one, such as a CPU PyTorch tensor.""" + if isinstance(arg, np.ndarray): + return True + return hasattr(arg, "__array__") and get_device_array(arg) is None and not np.isscalar(arg) + class Backend(ABC): """Base class for kernel_tuner backends.""" diff --git a/kernel_tuner/backends/nvcuda.py b/kernel_tuner/backends/nvcuda.py index c5f9b577d..cdfcbc1ef 100644 --- a/kernel_tuner/backends/nvcuda.py +++ b/kernel_tuner/backends/nvcuda.py @@ -4,7 +4,7 @@ import uuid import os -from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.backends.backend import GPUBackend, get_device_array, is_host_array from kernel_tuner.observers.nvcuda import CudaRuntimeObserver from kernel_tuner.util import SkippableFailure from kernel_tuner.utils.nvcuda import cuda_error_check, to_valid_nvrtc_gpu_arch_cc, find_cuda_home, _check @@ -232,16 +232,20 @@ def ready_argument_list(self, arguments): :param arguments: List of arguments to be passed to the kernel. The order should match the argument list on the CUDA kernel. Allowed values are numpy.ndarray, and/or numpy.int32, numpy.float32, and so on. + Device arrays that implement the CUDA Array Interface (e.g. PyTorch CUDA tensors, + CuPy arrays) and CPU PyTorch tensors are also supported. :type arguments: list(numpy objects) :returns: A list of arguments that can be passed to an CUDA kernel. - :rtype: list( pycuda.driver.DeviceAllocation, numpy.int32, ... ) + :rtype: list( cuda.CUdeviceptr, numpy.int32, ... ) """ gpu_args = [] for arg in arguments: - # if arg is a numpy array copy it to device - if isinstance(arg, np.ndarray): - err, device_memory = driver.cuMemAlloc(arg.nbytes) + # arrays are copied to a new device allocation, so the kernel never modifies the user's data + device_array = get_device_array(arg) + if device_array is not None or is_host_array(arg): + nbytes = device_array[1] if device_array is not None else np.asarray(arg).nbytes + err, device_memory = driver.cuMemAlloc(nbytes) cuda_error_check(err) self.allocations.append(device_memory) gpu_args.append(device_memory) @@ -454,12 +458,13 @@ def memset(allocation, value, size): def memcpy_dtoh(dest, src): """Perform a device to host memory copy. - :param dest: A numpy array in host memory to store the data + :param dest: A numpy array (or CPU PyTorch tensor) in host memory to store the data :type dest: numpy.ndarray :param src: A GPU memory allocation unit :type src: cuda.CUdeviceptr """ + dest = np.asarray(dest) err = driver.cuMemcpyDtoH(dest, src, dest.nbytes) cuda_error_check(err) @@ -470,10 +475,17 @@ def memcpy_htod(dest, src): :param dest: A GPU memory allocation unit :type dest: cuda.CUdeviceptr - :param src: A numpy array in host memory to store the data + :param src: A numpy array (or CPU PyTorch tensor) in host memory, or a device array + that implements the CUDA Array Interface :type src: numpy.ndarray """ - err = driver.cuMemcpyHtoD(dest, src, src.nbytes) + device_array = get_device_array(src) + if device_array is not None: + ptr, nbytes = device_array + err = driver.cuMemcpyDtoD(dest, ptr, nbytes) + else: + src = np.asarray(src) + err = driver.cuMemcpyHtoD(dest, src, src.nbytes) cuda_error_check(err) units = {"time": "ms"} diff --git a/kernel_tuner/backends/pycuda.py b/kernel_tuner/backends/pycuda.py index 8f7360e5e..dca9fa1bf 100644 --- a/kernel_tuner/backends/pycuda.py +++ b/kernel_tuner/backends/pycuda.py @@ -6,10 +6,10 @@ import numpy as np -from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.backends.backend import GPUBackend, get_device_array, is_host_array from kernel_tuner.observers.nvml import nvml # noqa F401 from kernel_tuner.observers.pycuda import PyCudaRuntimeObserver -from kernel_tuner.util import SkippableFailure, TorchPlaceHolder +from kernel_tuner.util import SkippableFailure # embedded in try block to be able to generate documentation # and run tests without pycuda installed @@ -20,8 +20,7 @@ except ImportError: class PyCudaPlaceHolder: - def __init__(self): - self.PointerHolderBase = object + pass drv = PyCudaPlaceHolder() pycuda_available = False @@ -35,23 +34,6 @@ def __init__(self): except ImportError: DynamicSourceModule = None -try: - import torch -except ImportError: - torch = TorchPlaceHolder() - - -class Holder(drv.PointerHolderBase): - """class to interoperate torch device memory allocations with PyCUDA.""" - - def __init__(self, tensor): - super(Holder, self).__init__() - self.tensor = tensor - self.gpudata = tensor.data_ptr() - - def get_pointer(self): - return self.t.data_ptr() - class PyCudaFunctions(GPUBackend): """Class that groups the CUDA functions on maintains state about the device.""" @@ -157,6 +139,8 @@ def ready_argument_list(self, arguments): :param arguments: List of arguments to be passed to the kernel. The order should match the argument list on the CUDA kernel. Allowed values are numpy.ndarray, and/or numpy.int32, numpy.float32, and so on. + Device arrays that implement the CUDA Array Interface (e.g. PyTorch CUDA tensors, + CuPy arrays) and CPU PyTorch tensors are also supported. :type arguments: list(numpy objects) :returns: A list of arguments that can be passed to an CUDA kernel. @@ -164,17 +148,14 @@ def ready_argument_list(self, arguments): """ gpu_args = [] for arg in arguments: - # if arg i is a numpy array copy to device - if isinstance(arg, np.ndarray): - alloc = drv.mem_alloc(arg.nbytes) + # arrays are copied to a new device allocation, so the kernel never modifies the user's data + device_array = get_device_array(arg) + if device_array is not None or is_host_array(arg): + nbytes = device_array[1] if device_array is not None else np.asarray(arg).nbytes + alloc = drv.mem_alloc(nbytes) self.allocations.append(alloc) gpu_args.append(alloc) - drv.memcpy_htod(gpu_args[-1], arg) - elif isinstance(arg, torch.Tensor): - if arg.is_cuda: - gpu_args.append(Holder(arg)) - else: - gpu_args.append(Holder(arg.cuda())) + self.memcpy_htod(alloc, arg) # pycuda does not support bool, convert to uint8 instead elif isinstance(arg, np.bool_): gpu_args.append(arg.astype(np.uint8)) @@ -366,32 +347,31 @@ def memset(self, allocation, value, size): drv.memset_d8(allocation, value, size) def memcpy_dtoh(self, dest, src): - print("dtoh called") """Perform a device to host memory copy. - :param dest: A numpy array in host memory to store the data + :param dest: A numpy array (or CPU PyTorch tensor) in host memory to store the data :type dest: numpy.ndarray :param src: A GPU memory allocation unit :type src: pycuda.driver.DeviceAllocation """ - if isinstance(src, drv.DeviceAllocation): - drv.memcpy_dtoh(dest, src) - elif isinstance(src, torch.Tensor): - dest[:] = src + drv.memcpy_dtoh(np.asarray(dest), src) def memcpy_htod(self, dest, src): - print("htod called") """Perform a host to device memory copy. :param dest: A GPU memory allocation unit :type dest: pycuda.driver.DeviceAllocation - :param src: A numpy array in host memory to store the data + :param src: A numpy array (or CPU PyTorch tensor) in host memory, or a device array + that implements the CUDA Array Interface :type src: numpy.ndarray """ - if isinstance(dest, drv.DeviceAllocation): - drv.memcpy_htod(dest, src) + device_array = get_device_array(src) + if device_array is not None: + drv.memcpy_dtod(dest, *device_array) + else: + drv.memcpy_htod(dest, np.asarray(src)) units = {"time": "ms", "power": "s,mW", "energy": "J"} diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 742d5548a..a78fa1ac6 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -754,13 +754,13 @@ def retrieve_results_to_host(self, arguments: list, should_sync: list[bool], gpu if isinstance(arg, (np.ndarray,) + cupy_ndarray) or util.is_julia_array(arg): result_host.append(np.zeros_like(arg)) self.dev.memcpy_dtoh(result_host[-1], gpu_args[i]) - elif isinstance(arg, torch.Tensor) and isinstance(answer[i], torch.Tensor): - if not answer[i].is_cuda: - # if the answer is on the host, copy gpu output to host as well - result_host.append(torch.zeros_like(answer[i])) - self.dev.memcpy_dtoh(result_host[-1], gpu_args[i].tensor) - else: - result_host.append(gpu_args[i].tensor) + elif isinstance(arg, torch.Tensor): + result = torch.empty(arg.shape, dtype=arg.dtype, device="cpu") + self.dev.memcpy_dtoh(result, gpu_args[i]) + # verification compares on the device of the answer + if answer is not None and isinstance(answer[i], torch.Tensor) and answer[i].is_cuda: + result = result.to(answer[i].device) + result_host.append(result) else: # We should sync this argument, but we do not know how to transfer this type of argument # What do we do? Should we throw an error? diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index b29dc270a..fc1b736ba 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -1061,7 +1061,8 @@ def run_kernel( if numpy.isscalar(arg): results.append(arg) elif isinstance(arg, torch.Tensor): - results.append(arg.cpu()) + results.append(torch.empty(arg.shape, dtype=arg.dtype, device="cpu")) + dev.memcpy_dtoh(results[-1], gpu_args[i]) else: results.append(numpy.zeros_like(arg)) dev.memcpy_dtoh(results[-1], gpu_args[i]) diff --git a/test/context.py b/test/context.py index 7c30a33c9..cd7d2b194 100644 --- a/test/context.py +++ b/test/context.py @@ -134,7 +134,7 @@ def skip_backend(backend: str): - if backend.upper() == "CUDA" and not pycuda_present: + if backend.upper() in ("CUDA", "PYCUDA") and not pycuda_present: pytest.skip("PyCuda not installed or no CUDA device detected") elif backend.upper() == "CUPY" and not cupy_present: pytest.skip("CuPy not installed or no CUDA device detected") diff --git a/test/test_torch_arguments.py b/test/test_torch_arguments.py new file mode 100644 index 000000000..73953c166 --- /dev/null +++ b/test/test_torch_arguments.py @@ -0,0 +1,111 @@ +"""Tests for passing PyTorch tensors and other CUDA Array Interface objects as kernel arguments.""" + +import numpy as np +import pytest + +from kernel_tuner import core, run_kernel, tune_kernel +from kernel_tuner.backends.backend import get_device_array, is_host_array + +from .context import skip_backend + +try: + import torch + + torch_cuda_present = torch.cuda.is_available() +except ImportError: + torch = None + torch_cuda_present = False + +skip_if_no_torch = pytest.mark.skipif(torch is None, reason="Torch not installed") +skip_if_no_torch_cuda = pytest.mark.skipif(not torch_cuda_present, reason="Torch not installed or no CUDA device") + +# accumulates into c, so verification only passes if outputs are reset between kernel runs +accumulate_kernel = """ +__global__ void vector_add(float *c, float *a, float *b, int n) { + int i = blockIdx.x * block_size_x + threadIdx.x; + if (i < n) { + c[i] += a[i] + b[i]; + } +} +""" + + +class FakeDeviceArray: + def __init__(self, shape, typestr, strides=None): + self.__cuda_array_interface__ = { + "shape": shape, + "typestr": typestr, + "data": (1234, False), + "strides": strides, + "version": 3, + } + + +def test_get_device_array(): + assert get_device_array(FakeDeviceArray((10, 4), " Date: Fri, 25 Sep 2026 15:46:13 +0200 Subject: [PATCH 18/44] rename examples/generic_python to examples/python --- examples/{generic_python => python}/call_functions.py | 0 examples/{generic_python => python}/cute_vec_add.py | 0 examples/{generic_python => python}/matmul/cupy_matmul.py | 0 examples/{generic_python => python}/matmul/cute_matmul.py | 0 examples/{generic_python => python}/matmul/numba_matmul.py | 0 examples/{generic_python => python}/matmul/taichi_matmul.py | 0 examples/{generic_python => python}/matmul/tilelang_matmul.py | 0 examples/{generic_python => python}/matmul/tilus_matmul.py | 0 examples/{generic_python => python}/matmul/triton_matmul.py | 0 examples/{generic_python => python}/matmul/warp_matmul.py | 0 examples/{generic_python => python}/numba_vec_add.py | 0 examples/{generic_python => python}/tilelang_vec_add.py | 0 examples/{generic_python => python}/tilus_vec_add.py | 0 examples/{generic_python => python}/triton_vec_add.py | 0 examples/{generic_python => python}/warp_vec_add.py | 0 15 files changed, 0 insertions(+), 0 deletions(-) rename examples/{generic_python => python}/call_functions.py (100%) rename examples/{generic_python => python}/cute_vec_add.py (100%) rename examples/{generic_python => python}/matmul/cupy_matmul.py (100%) rename examples/{generic_python => python}/matmul/cute_matmul.py (100%) rename examples/{generic_python => python}/matmul/numba_matmul.py (100%) rename examples/{generic_python => python}/matmul/taichi_matmul.py (100%) rename examples/{generic_python => python}/matmul/tilelang_matmul.py (100%) rename examples/{generic_python => python}/matmul/tilus_matmul.py (100%) rename examples/{generic_python => python}/matmul/triton_matmul.py (100%) rename examples/{generic_python => python}/matmul/warp_matmul.py (100%) rename examples/{generic_python => python}/numba_vec_add.py (100%) rename examples/{generic_python => python}/tilelang_vec_add.py (100%) rename examples/{generic_python => python}/tilus_vec_add.py (100%) rename examples/{generic_python => python}/triton_vec_add.py (100%) rename examples/{generic_python => python}/warp_vec_add.py (100%) diff --git a/examples/generic_python/call_functions.py b/examples/python/call_functions.py similarity index 100% rename from examples/generic_python/call_functions.py rename to examples/python/call_functions.py diff --git a/examples/generic_python/cute_vec_add.py b/examples/python/cute_vec_add.py similarity index 100% rename from examples/generic_python/cute_vec_add.py rename to examples/python/cute_vec_add.py diff --git a/examples/generic_python/matmul/cupy_matmul.py b/examples/python/matmul/cupy_matmul.py similarity index 100% rename from examples/generic_python/matmul/cupy_matmul.py rename to examples/python/matmul/cupy_matmul.py diff --git a/examples/generic_python/matmul/cute_matmul.py b/examples/python/matmul/cute_matmul.py similarity index 100% rename from examples/generic_python/matmul/cute_matmul.py rename to examples/python/matmul/cute_matmul.py diff --git a/examples/generic_python/matmul/numba_matmul.py b/examples/python/matmul/numba_matmul.py similarity index 100% rename from examples/generic_python/matmul/numba_matmul.py rename to examples/python/matmul/numba_matmul.py diff --git a/examples/generic_python/matmul/taichi_matmul.py b/examples/python/matmul/taichi_matmul.py similarity index 100% rename from examples/generic_python/matmul/taichi_matmul.py rename to examples/python/matmul/taichi_matmul.py diff --git a/examples/generic_python/matmul/tilelang_matmul.py b/examples/python/matmul/tilelang_matmul.py similarity index 100% rename from examples/generic_python/matmul/tilelang_matmul.py rename to examples/python/matmul/tilelang_matmul.py diff --git a/examples/generic_python/matmul/tilus_matmul.py b/examples/python/matmul/tilus_matmul.py similarity index 100% rename from examples/generic_python/matmul/tilus_matmul.py rename to examples/python/matmul/tilus_matmul.py diff --git a/examples/generic_python/matmul/triton_matmul.py b/examples/python/matmul/triton_matmul.py similarity index 100% rename from examples/generic_python/matmul/triton_matmul.py rename to examples/python/matmul/triton_matmul.py diff --git a/examples/generic_python/matmul/warp_matmul.py b/examples/python/matmul/warp_matmul.py similarity index 100% rename from examples/generic_python/matmul/warp_matmul.py rename to examples/python/matmul/warp_matmul.py diff --git a/examples/generic_python/numba_vec_add.py b/examples/python/numba_vec_add.py similarity index 100% rename from examples/generic_python/numba_vec_add.py rename to examples/python/numba_vec_add.py diff --git a/examples/generic_python/tilelang_vec_add.py b/examples/python/tilelang_vec_add.py similarity index 100% rename from examples/generic_python/tilelang_vec_add.py rename to examples/python/tilelang_vec_add.py diff --git a/examples/generic_python/tilus_vec_add.py b/examples/python/tilus_vec_add.py similarity index 100% rename from examples/generic_python/tilus_vec_add.py rename to examples/python/tilus_vec_add.py diff --git a/examples/generic_python/triton_vec_add.py b/examples/python/triton_vec_add.py similarity index 100% rename from examples/generic_python/triton_vec_add.py rename to examples/python/triton_vec_add.py diff --git a/examples/generic_python/warp_vec_add.py b/examples/python/warp_vec_add.py similarity index 100% rename from examples/generic_python/warp_vec_add.py rename to examples/python/warp_vec_add.py From c995d75d39bc55dead6e5a2f1f8d6d737fad4355 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Fri, 25 Sep 2026 15:46:16 +0200 Subject: [PATCH 19/44] fix Python DSL examples after the move to examples/python The matmul examples imported the call functions from the old examples.generic_python package. They now add examples/python to sys.path and import call_functions directly, like the vec_add examples. The tilus vec_add example hard-coded the number of warps, so the num_warps tunable parameter had no effect. It is now taken from the tuning parameters, and run() passes a value for it. --- examples/python/matmul/cupy_matmul.py | 6 +++++- examples/python/matmul/cute_matmul.py | 6 +++++- examples/python/matmul/numba_matmul.py | 6 +++++- examples/python/matmul/taichi_matmul.py | 6 +++++- examples/python/matmul/tilelang_matmul.py | 6 +++++- examples/python/matmul/tilus_matmul.py | 6 +++++- examples/python/matmul/triton_matmul.py | 6 +++++- examples/python/matmul/warp_matmul.py | 6 +++++- examples/python/tilus_vec_add.py | 7 ++++--- 9 files changed, 44 insertions(+), 11 deletions(-) diff --git a/examples/python/matmul/cupy_matmul.py b/examples/python/matmul/cupy_matmul.py index 1101b5fae..be95fad98 100644 --- a/examples/python/matmul/cupy_matmul.py +++ b/examples/python/matmul/cupy_matmul.py @@ -3,7 +3,11 @@ import numpy as np from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_cupyx +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_cupyx # noqa: E402 @jit.rawkernel() diff --git a/examples/python/matmul/cute_matmul.py b/examples/python/matmul/cute_matmul.py index d584a6aa8..1ef6d911f 100644 --- a/examples/python/matmul/cute_matmul.py +++ b/examples/python/matmul/cute_matmul.py @@ -8,7 +8,11 @@ from cutlass.cute.runtime import from_dlpack from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_cute +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_cute # noqa: E402 # might need export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH to work diff --git a/examples/python/matmul/numba_matmul.py b/examples/python/matmul/numba_matmul.py index 71a6b1b3a..7ed9b79ff 100644 --- a/examples/python/matmul/numba_matmul.py +++ b/examples/python/matmul/numba_matmul.py @@ -2,7 +2,11 @@ from numba import cuda, float32 from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_numba +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_numba # noqa: E402 # Source: https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication diff --git a/examples/python/matmul/taichi_matmul.py b/examples/python/matmul/taichi_matmul.py index bb3c53a81..4992ed086 100644 --- a/examples/python/matmul/taichi_matmul.py +++ b/examples/python/matmul/taichi_matmul.py @@ -2,7 +2,11 @@ import taichi as ti from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_taichi +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_taichi # noqa: E402 ti.init(arch=ti.gpu) diff --git a/examples/python/matmul/tilelang_matmul.py b/examples/python/matmul/tilelang_matmul.py index 751c0add4..0e7658639 100644 --- a/examples/python/matmul/tilelang_matmul.py +++ b/examples/python/matmul/tilelang_matmul.py @@ -5,7 +5,11 @@ import itertools from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_tilelang +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_tilelang # noqa: E402 # https://github.com/tile-ai/tilelang/tree/main/examples/gemm # num_threads and num_stages added as variables to enable tuning. diff --git a/examples/python/matmul/tilus_matmul.py b/examples/python/matmul/tilus_matmul.py index c65cdeaa7..7a75197e7 100644 --- a/examples/python/matmul/tilus_matmul.py +++ b/examples/python/matmul/tilus_matmul.py @@ -5,7 +5,11 @@ from tilus.utils import cdiv from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_tilus +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_tilus # noqa: E402 # This kernel is copied from the Tilus project: diff --git a/examples/python/matmul/triton_matmul.py b/examples/python/matmul/triton_matmul.py index 9f34356ab..f52b984c1 100644 --- a/examples/python/matmul/triton_matmul.py +++ b/examples/python/matmul/triton_matmul.py @@ -3,7 +3,11 @@ import triton.language as tl from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_triton +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_triton # noqa: E402 diff --git a/examples/python/matmul/warp_matmul.py b/examples/python/matmul/warp_matmul.py index c4a218b21..ac6734749 100644 --- a/examples/python/matmul/warp_matmul.py +++ b/examples/python/matmul/warp_matmul.py @@ -2,7 +2,11 @@ import warp as wp from kernel_tuner import tune_kernel -from examples.generic_python.call_functions import call_warp +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_warp # noqa: E402 wp.init() wp.config.enable_backward = False diff --git a/examples/python/tilus_vec_add.py b/examples/python/tilus_vec_add.py index 8ee96aa07..ca215503a 100644 --- a/examples/python/tilus_vec_add.py +++ b/examples/python/tilus_vec_add.py @@ -12,7 +12,8 @@ class VecAddV(tilus.Script): def __init__(self, block_size_x=None, num_warps=None): super().__init__() - self.block_size_x = block_size_x # number of threads per block + self.block_size_x = block_size_x # number of elements per block + self.num_warps = num_warps # number of warps per block def __call__( self, @@ -24,7 +25,7 @@ def __call__( # compute the number of blocks needed self.attrs.blocks = [cdiv(n_size, self.block_size_x)] - self.attrs.warps = 4 # number of warps per block + self.attrs.warps = self.num_warps # calculate the offset for this block offset: int32 = self.block_size_x * self.blockIdx.x @@ -83,7 +84,7 @@ def run(size): kernel_source=__file__, problem_size=size, arguments=args, - params={"block_size_x": 32}, + params={"block_size_x": 32, "num_warps": 4}, lang="generic_python", call_function=call_tilus, ) From 0df1bd3a8544096aedf491eef3173ddda838bacb Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Fri, 25 Sep 2026 15:49:38 +0200 Subject: [PATCH 20/44] remove trailing whitespace and unused imports in generic Python code Whitespace cleanup in the generic Python backend and the Python DSL examples, and removes the unused copy and warnings imports from the generic Python backend. --- examples/python/call_functions.py | 18 ++++---- examples/python/triton_vec_add.py | 18 ++++---- kernel_tuner/backends/generic_python.py | 58 ++++++++++++------------- 3 files changed, 45 insertions(+), 49 deletions(-) diff --git a/examples/python/call_functions.py b/examples/python/call_functions.py index 8c74790e1..5d805960b 100644 --- a/examples/python/call_functions.py +++ b/examples/python/call_functions.py @@ -9,7 +9,7 @@ def call_triton(kernel_function, args, kwargs, grid, threads, params): kwargs["num_warps"] = params["num_warps"] if "num_stages" in params.keys(): kwargs["num_stages"] = params["num_stages"] - + kernel_function[grid](*args, **kwargs) @@ -20,7 +20,7 @@ def call_tilelang(kernel_function, args, kwargs): def call_numba(kernel_function, args, kwargs, grid, threads): from numba import cuda - + numba_args = [] for arg in args: if isinstance(arg, torch.Tensor): @@ -49,7 +49,7 @@ def call_cute(kernel_function, args, kwargs, grid, threads, params): # Initialize cache if it does not exist if not hasattr(call_cute, "custom_cache"): - call_cute.custom_cache = {} + call_cute.custom_cache = {} # Convert Torch tensors to CuTe tensors with correct layout cute_args = [] @@ -64,12 +64,12 @@ def call_cute(kernel_function, args, kwargs, grid, threads, params): param_keys = sorted(params.keys()) cache_str = type(kernel_function).__name__ for k in param_keys: - cache_str += "_" + str(params[k]) - + cache_str += "_" + str(params[k]) + # Check if kernel exists in cache. Otherwise, compile and save if cache_str in call_cute.custom_cache: compiled_kernel = call_cute.custom_cache[cache_str] - else: + else: compiled_kernel = cute.compile(kernel_function, *cute_args) call_cute.custom_cache[cache_str] = compiled_kernel @@ -82,7 +82,7 @@ def call_taichi(kernel_function, args, kwargs): def call_warp(kernel_function, args, kwargs, grid, threads, params): import warp as wp - + # Convert Torch tensors to Warp args warp_args = [] for arg in args: @@ -104,11 +104,11 @@ def call_warp(kernel_function, args, kwargs, grid, threads, params): dimensions = params['dim'] else: dimensions = [grid[i] * threads[i] for i in range(len(grid))] - + # launch kernel wp.launch( kernel_function, dim=dimensions, inputs=warp_args, block_dim=threads_per_block - ) \ No newline at end of file + ) diff --git a/examples/python/triton_vec_add.py b/examples/python/triton_vec_add.py index cafba79a9..c4f6bb643 100644 --- a/examples/python/triton_vec_add.py +++ b/examples/python/triton_vec_add.py @@ -38,17 +38,16 @@ def tune(): n = torch.tensor(size, dtype=torch.int32) c_expect = a + b - args = [a, b, c, size] + args = [a, b, c, size] tune_params = dict() tune_params["block_size_x"] = [2**i for i in range(11)] - - result = run_kernel("add_kernel", __file__, size, args, {"block_size_x": 256}, - lang="generic_python", call_function=call_triton) + + result = run_kernel("add_kernel", __file__, size, args, {"block_size_x": 256}, + lang="generic_python", call_function=call_triton) assert np.allclose(c_expect.cpu(), result[2]) - - - + + results, env = tune_kernel( kernel_name="add_kernel", kernel_source=__file__, @@ -60,7 +59,6 @@ def tune(): call_function=call_triton, ) - - + if __name__ == "__main__": - tune() \ No newline at end of file + tune() diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 2247d5e87..5adf2d7c0 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -1,9 +1,7 @@ import logging import inspect -import copy import traceback # for compile error handling import re -import warnings import builtins import numpy as np @@ -34,7 +32,7 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None :type iterations: int """ - if not torch: + if not torch: logging.error("Unable to import Torch") raise ImportError("Unable to import Torch") @@ -43,9 +41,9 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None self.name = torch.cuda.get_device_name(self.device_id) self.max_threads = 10**18 # 'inf' to support tile based programming models, which can use less threads then the size of a tile. - env = dict() + env = {} env["device_name"] = self.name - env["max_threads"] = self.max_threads + env["max_threads"] = self.max_threads env["iterations"] = iterations env["compiler_options"] = compiler_options self.env = env @@ -64,13 +62,13 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None # Variables to be filled in at compile time, needed for running the kernel after compilation. self.call_function = None - self.signature = None - self.gpu_kwargs = None + self.signature = None + self.gpu_kwargs = None super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) def ready_argument_list(self, arguments): - """Ready argument list to be passed to the kernel. Converts arguments to Torch GPU Tensors or Python + """Ready argument list to be passed to the kernel. Converts arguments to Torch GPU Tensors or Python Scalars. Arguments of built-in Python types are left untouched. :param arguments: List of arguments to be passed to the kernel. @@ -107,11 +105,11 @@ def ready_argument_list(self, arguments): raise TypeError("Unknown argument type: ", type(arg), ". Accepted types are Torch tenors, NumPy arrays and scalars and built-in Python types.") return torch_args - + def compile(self, kernel_instance, gpu_args=None): """Compile the kernel by executing it once. This enforces that the kernel is cached. - + :param kernel_instance: The kernel instance containing information such as the kernel_source, grid, threads, params, etc. :type kernel_instance: KernelInstance @@ -127,25 +125,25 @@ def compile(self, kernel_instance, gpu_args=None): if gpu_args is None: raise ValueError("gpu_args is None, Generic Python needs gpu args to compile the kernel") - + # The first time we compile, we also set the call function and the signature # We need this later to run the kernel. if self.call_function is None or self.signature is None: self.call_function = kernel_instance.kernel_source.call_function self.signature = kernel_instance.kernel_source.signature - + grid = kernel_instance.grid threads = kernel_instance.threads params = kernel_instance.params # If the kernel source is a class, we use the __call__ function as the kernel_function. - if inspect.isclass(kernel_instance.kernel_fn): + if inspect.isclass(kernel_instance.kernel_fn): kernel_function = kernel_instance.kernel_fn() elif callable(kernel_instance.kernel_fn): kernel_function = kernel_instance.kernel_fn else: raise TypeError("kernel function is not a class or function") - + # Tuning params can contain kernel arguments. In such cases, create keyword arguments with # the values of the tuning params. self.gpu_kwargs = {} @@ -156,7 +154,7 @@ def compile(self, kernel_instance, gpu_args=None): # Call the user-defined call function in order to compile the kernel. self.synchronize() - self.call_function(kernel_function, gpu_args, self.gpu_kwargs, grid, threads, params) + self.call_function(kernel_function, gpu_args, self.gpu_kwargs, grid, threads, params) self.synchronize() return kernel_function @@ -191,13 +189,13 @@ def run_kernel(self, func, gpu_args, threads, grid, params=None): of the grid :type grid: tuple(int, int, int) - :param params: A dictionary with the tuning params for this specific kernel + :param params: A dictionary with the tuning params for this specific kernel configuration :type params: dict """ - self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) - - + self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) + + def synchronize(self): """Halts execution until device has finished its tasks.""" torch.cuda.synchronize() @@ -231,8 +229,8 @@ def copy_shared_memory_args(self, smem_args): def copy_texture_memory_args(self, texmem_args): raise NotImplementedError("Generic Python does not support texture memory") - - + + def refresh_memory(self, gpu_memory, host_arguments, should_sync): """Refresh the GPU memory with the untouched host arguments. We overwrite the standard function because Python DSLs do usually do not manage memory explicitely""" @@ -250,20 +248,20 @@ def refresh_memory(self, gpu_memory, host_arguments, should_sync): # GPU tensor elif isinstance(gpu_arg, torch.Tensor): if isinstance(host_arg, np.ndarray): - gpu_arg.copy_(torch.as_tensor(host_arg)) + gpu_arg.copy_(torch.as_tensor(host_arg)) elif isinstance(host_arg, torch.Tensor): if host_arg.is_cuda and host_arg.device != gpu_arg.device: gpu_arg.copy_(host_arg.to(gpu_arg.device)) # different gpu's, no direct copy allowed else: - gpu_arg.copy_(host_arg) + gpu_arg.copy_(host_arg) else: # host_arg is scalar, fill into tensor gpu_arg.fill_(host_arg) - + def classify_compile_exception(self, e): - """Best effort to differentiate between a user error and a resource error. - + """Best effort to differentiate between a user error and a resource error. + :param e: the caught exception :type e: exception @@ -367,15 +365,15 @@ def classify_compile_exception(self, e): if isinstance(e, USER_ERROR_TYPES): return "user_error" - - + + def match_any(patterns, text): return any(re.search(p, text) for p in patterns) msg = str(e).lower() tb = "".join(traceback.format_tb(e.__traceback__)).lower() - + if match_any(RESOURCE_KEYWORDS, msg): return "resource_error" @@ -384,7 +382,7 @@ def match_any(patterns, text): if match_any(USER_ERROR_KEYWORDS, msg): return "user_error" - + if match_any(USER_ORIGINS, msg + tb): return "user_error" From 3c13f697e1577585ab3b7ae2a388827f25bb2a31 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Tue, 29 Sep 2026 20:31:10 +0200 Subject: [PATCH 21/44] improve kernel lookup in get_kernel_ast Search functions and classes in a single breadth-first pass, return clear errors when the kernel file, kernel or __call__ method cannot be found, and report the kernel file name on syntax errors. Also accepts async function definitions. Updates the expected error messages in test_kernel_source_fn accordingly. --- kernel_tuner/util.py | 53 +++++++++++++---------------------- test/test_kernel_source_fn.py | 4 +-- 2 files changed, 22 insertions(+), 35 deletions(-) diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index 3937e9b8c..9320fba62 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -768,52 +768,39 @@ def get_kernel_string(kernel_source, params=None): def get_kernel_ast(kernel_name, filepath): - ''' - Util function for Generic Python Backend that returns the kernel function as AST. + """Util function for Generic Python Backend that returns the kernel function as AST. :param kernel_name: name of the kernel (as passed by the user) :type kernel_name: string - :param filepath: the path to the file where the kernel lives. (passed by the user as kenrel_source) + :param filepath: the path to the file where the kernel lives (passed by the user as kernel_source) :type filepath: string or Path containing a filename that points to the kernel source - :returns: ast.FunctionDef node in case the kernel is a function or a tuple + :returns: ast.FunctionDef node in case the kernel is a function or a tuple (ast.ClassDef node, ast.FunctionDef node) in case the kernel is represented as the __call__ function of a class (for Tilus support). - ''' - if isinstance(filepath, Path): - source = read_file(filepath) - elif isinstance(filepath, str): - with open(filepath, "r") as f: - source = f.read() - else: - raise TypeError("Error kernel_source does not specify a path to a file") + """ + if not isinstance(filepath, (str, Path)): + raise TypeError(f"kernel_source should be a path to a file, got {type(filepath)}") + if not Path(filepath).is_file(): + raise FileNotFoundError(f"Kernel source file {filepath} not found") - tree = ast.parse(source) + source = read_file(filepath) + tree = ast.parse(source, filename=str(filepath)) - # Function based kernels - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef) and node.name == kernel_name: - return node - - # Class based kernels + # ast.walk is breadth-first, so top-level definitions are found before nested ones for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == kernel_name: + # Function based kernels + return node if isinstance(node, ast.ClassDef) and node.name == kernel_name: - class_node = node - break + # Class based kernels, the kernel is the __call__ function of the class + for member in node.body: + if isinstance(member, (ast.FunctionDef, ast.AsyncFunctionDef)) and member.name == "__call__": + return (node, member) + raise ValueError(f"No __call__ function found inside class {kernel_name} in {filepath}") - if not class_node: - raise ValueError(f"Kernel {kernel_name} not found in {filepath}") - - # Search for __call__ function within class - for node in class_node.body: - if isinstance(node, ast.FunctionDef) and node.name == "__call__": - call_node = node - return (class_node, call_node) - - if not call_node: - raise ValueError(f"No __call__ function found inside Class {kernel_name}") - + raise ValueError(f"Kernel {kernel_name} not found in {filepath}") def get_arg_names(func_node: ast.FunctionDef): diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py index d02be9e4d..9749dc6e4 100644 --- a/test/test_kernel_source_fn.py +++ b/test/test_kernel_source_fn.py @@ -68,10 +68,10 @@ def test_initiation(): with pytest.raises(ValueError, match=r"call_function must be supplied for language .*"): KernelSource("mock_kernel", KS_FILE, "generic_python") - with pytest.raises(FileNotFoundError, match=r".* No such file or directory: .*"): + with pytest.raises(FileNotFoundError, match=r"Kernel source file .* not found"): KernelSource("mock_kernel", "This is a string Kernel", "generic_python", call_function=call_mock) - with pytest.raises(TypeError, match="Error kernel_source does not specify a path to a file"): + with pytest.raises(TypeError, match="kernel_source should be a path to a file"): KernelSource("mock_kernel", mock_kernel, "generic_python", call_function=call_mock) with pytest.raises(ValueError, match=r"KernelSourceFn only supports a single kernel source"): From 759470d533a629edd589827c3456c5b772d1e86d Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Tue, 29 Sep 2026 20:31:11 +0200 Subject: [PATCH 22/44] use BLOCK_SIZE_N for grid y-dimension in Triton matmul example With BLOCK_SIZE_N as the second block size name, the grid now covers all columns of C when BLOCK_SIZE_N differs from BLOCK_SIZE_M. --- examples/python/matmul/triton_matmul.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/python/matmul/triton_matmul.py b/examples/python/matmul/triton_matmul.py index f52b984c1..9ffd5c282 100644 --- a/examples/python/matmul/triton_matmul.py +++ b/examples/python/matmul/triton_matmul.py @@ -234,7 +234,7 @@ def tune_basic(M, N, K): lang="generic_python", answer=[None, None, C_ref.cpu(), None, None, None], atol=M * 2**(-11), - block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_M"], + block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, call_function=call_triton, @@ -276,7 +276,7 @@ def tune_opt(M, N, K): lang="generic_python", answer=[None, None, C_ref.cpu(), None, None, None, None, None, None, None, None, None], atol=M * 2**(-11), - block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_M"], + block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], grid_div_x = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], restrictions = restrictions, strategy = "bayes_opt", From 48637b87088f9d9e20e9df4f6a6052959e135f22 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 09:10:56 +0200 Subject: [PATCH 23/44] fix bugs in Python DSL examples - numba matmul: drop cache=True (fp16 kernels cannot be pickled by numba-cuda 0.30) and fix DIM_x typo in restrictions - cute matmul: split naive restrictions into separate strings, restrict optimized GEMM to tile sizes supported by the C smem layout, and report the underlying exception type on compile failures - call_numba: launch with the converted arguments - warp vec_add: launch ceil(size / work_per_thread) threads - cute vec_add: fix 265 -> 256 block size, add __main__ guard - warp matmul: fix swapped tune(M, K, N) parameter order --- examples/python/call_functions.py | 2 +- examples/python/cute_vec_add.py | 7 ++++--- examples/python/matmul/cute_matmul.py | 10 ++++++---- examples/python/matmul/numba_matmul.py | 6 +++--- examples/python/matmul/warp_matmul.py | 2 +- examples/python/warp_vec_add.py | 3 ++- 6 files changed, 17 insertions(+), 13 deletions(-) diff --git a/examples/python/call_functions.py b/examples/python/call_functions.py index 5d805960b..e3132771e 100644 --- a/examples/python/call_functions.py +++ b/examples/python/call_functions.py @@ -28,7 +28,7 @@ def call_numba(kernel_function, args, kwargs, grid, threads): else: numba_args.append(arg) - kernel_function[grid, threads](*args, **kwargs) + kernel_function[grid, threads](*numba_args, **kwargs) def call_cupyx(kernel_function, args, kwargs, grid, threads): diff --git a/examples/python/cute_vec_add.py b/examples/python/cute_vec_add.py index 2744abb77..ecd2386e4 100644 --- a/examples/python/cute_vec_add.py +++ b/examples/python/cute_vec_add.py @@ -1,4 +1,4 @@ -import torch +import torch import cutlass import cutlass.cute as cute @@ -46,12 +46,13 @@ def main(): c = torch.zeros(size, device="cuda", dtype=torch.float16) args = [a, b, c, size] - tune_params = {"num_threads_per_block": [1, 2, 4, 8, 16, 32, 64, 128, 265, 512, 1024]} + tune_params = {"num_threads_per_block": [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]} answer = [None, None, (a+b).cpu(), None] tune_kernel("vec_add", __file__, size, args, tune_params, answer=answer, lang="generic_python", call_function=call_cute, verbose=True) -main() +if __name__ == "__main__": + main() diff --git a/examples/python/matmul/cute_matmul.py b/examples/python/matmul/cute_matmul.py index 1ef6d911f..681218b40 100644 --- a/examples/python/matmul/cute_matmul.py +++ b/examples/python/matmul/cute_matmul.py @@ -92,7 +92,7 @@ def tune_naive_matmul(M, N, K): tune_params = dict() tune_params["block_size_x"] = [2**i for i in range(1, 10)] tune_params["block_size_y"] = [2**i for i in range(1, 10)] - restrictions = ["block_size_x * block_size_y >= 32, block_size_x * block_size_y <= 1024"] + restrictions = ["block_size_x * block_size_y >= 32", "block_size_x * block_size_y <= 1024"] results, env = tune_kernel("matmul", __file__, size, args, tune_params, lang="generic_python", call_function=call_cute, answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) @@ -979,11 +979,11 @@ def call_cute_custom(kernel_function, args, kwargs, grid, threads, params): else: # Wrap in try-except because CuTe gives TypeError for certain block sizes. This # stops the tuning process and we do not want this - try: + try: compiled_kernel = cute.compile(kernel_function, *cute_args) call_cute_custom.custom_cache[cache_str] = compiled_kernel - except: - raise RuntimeError("Invalid configuration") + except Exception as e: + raise RuntimeError(f"Invalid configuration, CuTe compile failed with {type(e).__name__}") from e @@ -1050,6 +1050,8 @@ def tune_optimized(M, N, K, L=1): "num_stages * (bK * (bM + bN)) * 2 <= 49152", # SMEM "bM >= atom_lay_M * 32", # ensure each atom sees at least 2 MMA tiles per dimension "bN >= atom_lay_N * 16", # ensure each atom sees at least 2 MMA tiles per dimension + # the swizzled smem layout of C only supports these tiles, others fail in autovec_copy + "bN >= 128", "bN <= 128 * atom_lay_N", ] diff --git a/examples/python/matmul/numba_matmul.py b/examples/python/matmul/numba_matmul.py index 7ed9b79ff..69d0585f9 100644 --- a/examples/python/matmul/numba_matmul.py +++ b/examples/python/matmul/numba_matmul.py @@ -10,7 +10,7 @@ # Source: https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication -@cuda.jit(cache=True) +@cuda.jit def matmul(A, B, C): i, j = cuda.grid(2) if i < C.shape[0] and j < C.shape[1]: @@ -21,7 +21,7 @@ def matmul(A, B, C): # Translated to Numba-CUDA from https://github.com/cupy/cupy/blob/main/examples/gemm/sgemm.cu -@cuda.jit(cache=True) +@cuda.jit def optimized_matmul(M, N, K, A, B, C): DIM_X = 16 DIM_Y = 16 @@ -227,7 +227,7 @@ def tune_optimized(M, N, K): "THR_M == BLK_M / DIM_X", "THR_N == BLK_N / DIM_Y", "DIM_X * DIM_Y <= 1024", - "DIM_x * DIM_Y >= 32", + "DIM_X * DIM_Y >= 32", ] results, env = tune_kernel( diff --git a/examples/python/matmul/warp_matmul.py b/examples/python/matmul/warp_matmul.py index ac6734749..051eb506f 100644 --- a/examples/python/matmul/warp_matmul.py +++ b/examples/python/matmul/warp_matmul.py @@ -50,7 +50,7 @@ def run_gemm(M, N, K): print("Succes") -def tune(M, K, N): +def tune(M, N, K): rng = np.random.default_rng(42) A = rng.random((M, K)).astype(np.float16) B = rng.random((K, N)).astype(np.float16) diff --git a/examples/python/warp_vec_add.py b/examples/python/warp_vec_add.py index 81d7fb19e..0ccb07a96 100644 --- a/examples/python/warp_vec_add.py +++ b/examples/python/warp_vec_add.py @@ -36,7 +36,8 @@ def call_warp(kernel_function, args, kwargs, grid, threads, params): warp_args.append(wp.from_torch(arg)) else: warp_args.append(arg) - dim = params['size'] + # each thread processes work_per_thread elements + dim = (params['size'] + params['work_per_thread'] - 1) // params['work_per_thread'] wp.launch(kernel=kernel_function, dim=dim, inputs=warp_args) From 288add13a85a99a948e86591d21bebb5a1badd12 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 10:38:18 +0200 Subject: [PATCH 24/44] classify generic Python compile errors by message only Matching the traceback made nearly every error look like a resource error, because DSL file paths contain "cuda", while words like "last" or "cast" matched the "ast" user error origin. Only the messages of the exception and the exceptions it was raised from are now inspected, with word boundaries for short tokens. Code errors (NameError, SyntaxError) are always user errors, and resource patterns are checked before other exception types. Configurations are now only skipped when classified as resource errors, so unrecognized errors are raised like for the other backends. --- kernel_tuner/backends/generic_python.py | 161 +++++++++++------------- kernel_tuner/core.py | 5 +- 2 files changed, 77 insertions(+), 89 deletions(-) diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 5adf2d7c0..5f26b8236 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -1,6 +1,5 @@ import logging import inspect -import traceback # for compile error handling import re import builtins import numpy as np @@ -262,6 +261,10 @@ def refresh_memory(self, gpu_memory, host_arguments, should_sync): def classify_compile_exception(self, e): """Best effort to differentiate between a user error and a resource error. + Only the messages of the exception and of the exceptions it was raised from are + inspected. The traceback is not used, because the file paths and source lines in it + contain words like "cuda" or "ast" for almost any DSL, regardless of the error. + :param e: the caught exception :type e: exception @@ -269,121 +272,105 @@ def classify_compile_exception(self, e): :rtype: string """ - RESOURCE_KEYWORDS = ( + RESOURCE_PATTERNS = ( # Shared memory - "shared memory", - "smem", - "uses too much shared", - "exceeds shared memory", - "shared memory limit", + r"shared memory", + r"\bsmem\b", + r"too much shared", # Registers / occupancy - "too many registers", - "register spill", - "uses too many registers", - "out of registers", + r"too many registers", + r"register spill", + r"out of registers", + r"register usage", # Launch configuration - "invalid launch configuration", - "invalid configuration argument", - "threads per block", - "block size", - "grid size", - "num_warps", - "num_ctas", + r"invalid launch configuration", + r"invalid configuration", + r"invalidconfiguration", # cudaErrorInvalidConfiguration + r"launch_invalid_config", + r"threads per block", + r"block size", + r"grid size", + r"\bnum_warps\b", + r"\bnum_ctas\b", + r"no valid warp partition", + r"no valid schedule", # Generic resource exhaustion - "too many resources", - "out of resources", - "exceeds maximum", - "exceeds limit", - - # Compiler-level indicators - "ptxas error", - "ptxas fatal", - "nvcc error", - "cuda error", - "llvm error", - "mlir error", - "lowering failed", + r"too many resources", + r"out of resources?\b", + r"out_of_resources", # CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES + r"exceeds maximum", + r"exceeds limit", + + # Errors reported by the compiler toolchain + r"\bptxas\b", + r"\bnvcc error", + r"\bllvm error", + r"\bmlir error", + r"lowering failed", ) - USER_ERROR_KEYWORDS = ( + USER_ERROR_PATTERNS = ( # Undefined / missing symbols - "not defined", - "undefined variable", - "without definition", - "unknown variable", - "unbound", + r"not defined", + r"undefined variable", + r"without definition", + r"unknown variable", + r"\bunbound\b", # Type / shape errors - "type mismatch", - "invalid type", - "cannot convert", - "expected .* but got", - "incompatible types", + r"type mismatch", + r"invalid type", + r"cannot convert", + r"expected .* but got", + r"incompatible types", # Indexing / bounds - "index out of bounds", - "out of bounds access", - "invalid index", + r"index out of bounds", + r"out of bounds access", + r"invalid index", # IR / AST construction - "failed to build", - "invalid expression", - "malformed", - "illegal operation", - ) - - RESOURCE_ORIGINS = ( - "ptxas", - "nvcc", - "cuda", - "llvm", - "mlir", - "cubin", - "fatbin", + r"failed to build", + r"invalid expression", + r"malformed", + r"illegal operation", + + # Frontend components + r"\btranspiler\b", + r"\bfrontend\b", + r"\bast\b", ) - USER_ORIGINS = ( - "transpiler", - "scheduler", - "hidet", - "frontend", - "ast", - ) - - - USER_ERROR_TYPES = ( - NameError, - UnboundLocalError, - AttributeError, - TypeError, - SyntaxError, - IndentationError, - ) + # Errors in the kernel code that no tuning parameter value can cause (IndentationError is a SyntaxError) + CODE_ERROR_TYPES = (NameError, UnboundLocalError, SyntaxError) - if isinstance(e, USER_ERROR_TYPES): - return "user_error" + # Usually user errors, but some DSLs also raise these for configurations they do not support + USER_ERROR_TYPES = (AttributeError, TypeError) + # Follow the chain of exceptions like Python does when printing a traceback, + # wrappers often hide the original error in __cause__ or __context__ + chain = [] + while e is not None and all(e is not seen for seen in chain): + chain.append(e) + e = e.__cause__ if e.__cause__ is not None else (None if e.__suppress_context__ else e.__context__) + msg = "\n".join(str(exc) for exc in chain).lower() def match_any(patterns, text): return any(re.search(p, text) for p in patterns) - msg = str(e).lower() - tb = "".join(traceback.format_tb(e.__traceback__)).lower() - - - if match_any(RESOURCE_KEYWORDS, msg): - return "resource_error" + if any(isinstance(exc, CODE_ERROR_TYPES) for exc in chain): + return "user_error" - if match_any(RESOURCE_ORIGINS, msg + tb): + if match_any(RESOURCE_PATTERNS, msg): return "resource_error" - if match_any(USER_ERROR_KEYWORDS, msg): + if any(isinstance(exc, USER_ERROR_TYPES) for exc in chain): return "user_error" - if match_any(USER_ORIGINS, msg + tb): + if match_any(USER_ERROR_PATTERNS, msg): return "user_error" return "unknown" diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index a78fa1ac6..2b82ab246 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -587,8 +587,9 @@ def compile_kernel(self, instance, verbose, gpu_args=None): # skip over this configuration and try the next one error_message = str(e.stderr) if hasattr(e, "stderr") else str(e) if self.lang.upper() == "GENERIC_PYTHON": - # Python DSLs raise many different errors, the backend classifies them - skippable = self.dev.classify_compile_exception(e) != "user_error" + # Python DSLs raise many different errors, the backend classifies them. + # Only skip configurations that are recognized as using too many resources. + skippable = self.dev.classify_compile_exception(e) == "resource_error" reason = f"\n{e}" else: shared_mem_error_messages = [ From fd934a1ac0cd415edbab0c6aa15971f17fe574ac Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 10:38:21 +0200 Subject: [PATCH 25/44] add tests for tuning kernels in each Python DSL Parametrized end-to-end tune_kernel test for Numba, CuPy (cupyx.jit), Warp, Taichi, CuTe, Triton, Tilus, and TileLang, each skipped when the DSL is not installed, plus unit tests for classify_compile_exception using error messages from real DSL runs. skip_if_no_torch now also requires a CUDA device. --- CHANGELOG.md | 2 +- test/context.py | 36 ++++- test/generic_python_kernels/__init__.py | 10 ++ .../cupyx_vector_add.py | 25 ++++ .../generic_python_kernels/cute_vector_add.py | 43 ++++++ .../numba_vector_add.py | 24 ++++ .../taichi_vector_add.py | 34 +++++ .../tilelang_vector_add.py | 32 +++++ .../tilus_vector_add.py | 37 +++++ .../triton_vector_add.py | 25 ++++ .../generic_python_kernels/warp_vector_add.py | 26 ++++ test/test_generic_python_dsls.py | 128 ++++++++++++++++++ 12 files changed, 418 insertions(+), 4 deletions(-) create mode 100644 test/generic_python_kernels/__init__.py create mode 100644 test/generic_python_kernels/cupyx_vector_add.py create mode 100644 test/generic_python_kernels/cute_vector_add.py create mode 100644 test/generic_python_kernels/numba_vector_add.py create mode 100644 test/generic_python_kernels/taichi_vector_add.py create mode 100644 test/generic_python_kernels/tilelang_vector_add.py create mode 100644 test/generic_python_kernels/tilus_vector_add.py create mode 100644 test/generic_python_kernels/triton_vector_add.py create mode 100644 test/generic_python_kernels/warp_vector_add.py create mode 100644 test/test_generic_python_dsls.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1db199f1d..fe3a75298 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). ## Unreleased - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers -- Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, and CuTe, using `lang="generic_python"` and a `call_function` +- Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` - Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend - Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel - Fix output verification failing for PyTorch tensors as expected answers diff --git a/test/context.py b/test/context.py index cd7d2b194..dd54b3462 100644 --- a/test/context.py +++ b/test/context.py @@ -2,6 +2,7 @@ import subprocess import sys import ctypes.util +import os from os import environ import importlib.util @@ -81,10 +82,32 @@ try: import torch - gen_python_torch_present = True + gen_python_torch_present = torch.cuda.is_available() except ImportError: gen_python_torch_present = False + +def python_dsl_present(module, subpackage=None): + """Check if a Python DSL is installed without importing it, some DSLs are slow to import.""" + try: + spec = importlib.util.find_spec(module) + except (ImportError, ValueError): + return False + if spec is None: + return False + if subpackage is None: + return True + return any(os.path.isdir(os.path.join(loc, subpackage)) for loc in spec.submodule_search_locations or []) + + +numba_cuda_present = python_dsl_present("numba_cuda") +warp_present = python_dsl_present("warp") +taichi_present = python_dsl_present("taichi") +cute_present = python_dsl_present("cutlass", "cute") +triton_present = python_dsl_present("triton") +tilus_present = python_dsl_present("tilus") +tilelang_present = python_dsl_present("tilelang") + try: import pyatf @@ -130,7 +153,14 @@ skip_if_no_pyatf = pytest.mark.skipif(not pyatf_present, reason="PyATF not installed") skip_if_no_methodology = pytest.mark.skipif(not methodology_present, reason="Autotuning Methodology not found") skip_if_no_pymoo = pytest.mark.skipif(not pymoo_present, reason="No PyMOO found") -skip_if_no_torch = pytest.mark.skipif(not gen_python_torch_present, reason="Torch not installed") +skip_if_no_torch = pytest.mark.skipif(not gen_python_torch_present, reason="Torch not installed or no CUDA device") +skip_if_no_numba_cuda = pytest.mark.skipif(not numba_cuda_present, reason="numba-cuda not installed") +skip_if_no_warp = pytest.mark.skipif(not warp_present, reason="Warp not installed") +skip_if_no_taichi = pytest.mark.skipif(not taichi_present, reason="Taichi not installed") +skip_if_no_cute = pytest.mark.skipif(not cute_present, reason="CuTe DSL not installed") +skip_if_no_triton = pytest.mark.skipif(not triton_present, reason="Triton not installed") +skip_if_no_tilus = pytest.mark.skipif(not tilus_present, reason="Tilus not installed") +skip_if_no_tilelang = pytest.mark.skipif(not tilelang_present, reason="TileLang not installed") def skip_backend(backend: str): @@ -153,4 +183,4 @@ def skip_backend(backend: str): elif backend.upper() == "JULIA" and not shutil.which("julia"): pytest.skip("No Julia on PATH") elif backend.upper() == "GENERIC_PYTHON" and not gen_python_torch_present: - pytest.skip("Torch not installed") + pytest.skip("Torch not installed or no CUDA device") diff --git a/test/generic_python_kernels/__init__.py b/test/generic_python_kernels/__init__.py new file mode 100644 index 000000000..57af35627 --- /dev/null +++ b/test/generic_python_kernels/__init__.py @@ -0,0 +1,10 @@ +"""Vector add kernels in Python DSLs, used by test_generic_python_dsls.py. + +Each DSL lives in its own module, because Kernel Tuner copies all imports of the kernel source file +into the module it generates for each configuration. Every module provides: + +- ``kernel_name``: name of the kernel function or class in the module +- ``arguments(c, a, b, n)``: the kernel arguments in the order the kernel expects them +- ``tune_params(n)``: a small set of tunable parameters +- ``call_function``: the call function that launches the kernel +""" diff --git a/test/generic_python_kernels/cupyx_vector_add.py b/test/generic_python_kernels/cupyx_vector_add.py new file mode 100644 index 000000000..d4d1f725e --- /dev/null +++ b/test/generic_python_kernels/cupyx_vector_add.py @@ -0,0 +1,25 @@ +import cupy as cp +import torch +from cupyx import jit + +kernel_name = "vector_add" + + +@jit.rawkernel() +def vector_add(c, a, b, n): + i = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x + if i < n: + c[i] = a[i] + b[i] + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256]} + + +def call_function(kernel_function, args, kwargs, grid, threads): + cupy_args = [cp.from_dlpack(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] + kernel_function(grid, threads, tuple(cupy_args)) diff --git a/test/generic_python_kernels/cute_vector_add.py b/test/generic_python_kernels/cute_vector_add.py new file mode 100644 index 000000000..0d5e60cf9 --- /dev/null +++ b/test/generic_python_kernels/cute_vector_add.py @@ -0,0 +1,43 @@ +import cutlass.cute as cute +import torch +from cutlass.cute.runtime import from_dlpack + +kernel_name = "vector_add" + + +@cute.kernel +def vector_add_kernel(gC: cute.Tensor, gA: cute.Tensor, gB: cute.Tensor, n: cute.Int32): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + bdim, _, _ = cute.arch.block_dim() + i = bdim * bidx + tidx + if i < n: + gC[i] = gA[i] + gB[i] + + +@cute.jit +def vector_add(mC: cute.Tensor, mA: cute.Tensor, mB: cute.Tensor, n: cute.Int32): + block_size_x = 128 + vector_add_kernel(mC, mA, mB, n).launch( + grid=(cute.ceil_div(n, block_size_x), 1, 1), + block=(block_size_x, 1, 1), + ) + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256]} + + +# kernels compiled with cute.compile, indexed by the kernel function of each configuration +compiled_kernels = {} + + +def call_function(kernel_function, args, kwargs): + cute_args = [from_dlpack(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] + if kernel_function not in compiled_kernels: + compiled_kernels[kernel_function] = cute.compile(kernel_function, *cute_args) + compiled_kernels[kernel_function](*cute_args, **kwargs) diff --git a/test/generic_python_kernels/numba_vector_add.py b/test/generic_python_kernels/numba_vector_add.py new file mode 100644 index 000000000..57299bada --- /dev/null +++ b/test/generic_python_kernels/numba_vector_add.py @@ -0,0 +1,24 @@ +import torch +from numba import cuda + +kernel_name = "vector_add" + + +@cuda.jit +def vector_add(c, a, b, n): + i = cuda.grid(1) + if i < n: + c[i] = a[i] + b[i] + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256]} + + +def call_function(kernel_function, args, kwargs, grid, threads): + numba_args = [cuda.as_cuda_array(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] + kernel_function[grid, threads](*numba_args, **kwargs) diff --git a/test/generic_python_kernels/taichi_vector_add.py b/test/generic_python_kernels/taichi_vector_add.py new file mode 100644 index 000000000..1169e20a8 --- /dev/null +++ b/test/generic_python_kernels/taichi_vector_add.py @@ -0,0 +1,34 @@ +import taichi as ti + +ti.init(arch=ti.cuda) + +# Taichi silently falls back to the CPU when it cannot use CUDA +requires_skip = None if ti.lang.impl.current_cfg().arch == ti.cuda else "Taichi cannot use CUDA" + +kernel_name = "vector_add" + +block_size_x = 128 # replaced by the value of the tunable parameter + + +@ti.kernel +def vector_add( + c: ti.types.ndarray(dtype=ti.f32, ndim=1), + a: ti.types.ndarray(dtype=ti.f32, ndim=1), + b: ti.types.ndarray(dtype=ti.f32, ndim=1), + n: ti.i32, +): + ti.loop_config(block_dim=block_size_x) + for i in range(n): + c[i] = a[i] + b[i] + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256]} + + +def call_function(kernel_function, args, kwargs): + kernel_function(*args, **kwargs) diff --git a/test/generic_python_kernels/tilelang_vector_add.py b/test/generic_python_kernels/tilelang_vector_add.py new file mode 100644 index 000000000..8471aafe0 --- /dev/null +++ b/test/generic_python_kernels/tilelang_vector_add.py @@ -0,0 +1,32 @@ +import tilelang +import tilelang.language as T + +kernel_name = "vector_add" + + +@tilelang.jit +def vector_add(n: int, block_size_x: int = 128): + @T.prim_func + def vector_add_kernel( + C: T.Tensor((n,), "float32"), + A: T.Tensor((n,), "float32"), + B: T.Tensor((n,), "float32"), + ): + with T.Kernel(T.ceildiv(n, block_size_x), threads=block_size_x) as bx: + for i in T.Parallel(block_size_x): + C[bx * block_size_x + i] = A[bx * block_size_x + i] + B[bx * block_size_x + i] + + return vector_add_kernel + + +def arguments(c, a, b, n): + # n is passed to the kernel factory through the tunable parameters + return [c, a, b] + + +def tune_params(n): + return {"block_size_x": [128, 256], "n": [n]} + + +def call_function(kernel_function, args, kwargs): + kernel_function(**kwargs)(*args) diff --git a/test/generic_python_kernels/tilus_vector_add.py b/test/generic_python_kernels/tilus_vector_add.py new file mode 100644 index 000000000..69fd9e3d1 --- /dev/null +++ b/test/generic_python_kernels/tilus_vector_add.py @@ -0,0 +1,37 @@ +import tilus +from tilus import float32, int32 +from tilus.utils import cdiv + +kernel_name = "VectorAdd" + + +class VectorAdd(tilus.Script): + def __init__(self): + super().__init__() + self.block_size_x = 128 + self.num_warps = 4 + + def __call__(self, c_ptr: ~float32, a_ptr: ~float32, b_ptr: ~float32, n: int32): + self.attrs.blocks = [cdiv(n, self.block_size_x)] + self.attrs.warps = self.num_warps + + offset: int32 = self.block_size_x * self.blockIdx.x + ga = self.global_view(a_ptr, dtype=float32, shape=[n]) + gb = self.global_view(b_ptr, dtype=float32, shape=[n]) + gc = self.global_view(c_ptr, dtype=float32, shape=[n]) + + a = self.load_global(ga, offsets=[offset], shape=[self.block_size_x]) + b = self.load_global(gb, offsets=[offset], shape=[self.block_size_x]) + self.store_global(gc, a + b, offsets=[offset]) + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256], "num_warps": [4]} + + +def call_function(kernel_function, args, kwargs): + kernel_function(*args, **kwargs) diff --git a/test/generic_python_kernels/triton_vector_add.py b/test/generic_python_kernels/triton_vector_add.py new file mode 100644 index 000000000..03a50ef29 --- /dev/null +++ b/test/generic_python_kernels/triton_vector_add.py @@ -0,0 +1,25 @@ +import triton +import triton.language as tl + +kernel_name = "vector_add" + + +@triton.jit +def vector_add(c_ptr, a_ptr, b_ptr, n, block_size_x: tl.constexpr): + offsets = tl.program_id(axis=0) * block_size_x + tl.arange(0, block_size_x) + mask = offsets < n + a = tl.load(a_ptr + offsets, mask=mask) + b = tl.load(b_ptr + offsets, mask=mask) + tl.store(c_ptr + offsets, a + b, mask=mask) + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256]} + + +def call_function(kernel_function, args, kwargs, grid): + kernel_function[grid](*args, **kwargs) diff --git a/test/generic_python_kernels/warp_vector_add.py b/test/generic_python_kernels/warp_vector_add.py new file mode 100644 index 000000000..e30629a11 --- /dev/null +++ b/test/generic_python_kernels/warp_vector_add.py @@ -0,0 +1,26 @@ +import torch +import warp as wp + +wp.init() + +kernel_name = "vector_add" + + +@wp.kernel +def vector_add(c: wp.array(dtype=float), a: wp.array(dtype=float), b: wp.array(dtype=float), n: int): + i = wp.tid() + if i < n: + c[i] = a[i] + b[i] + + +def arguments(c, a, b, n): + return [c, a, b, n] + + +def tune_params(n): + return {"block_size_x": [128, 256]} + + +def call_function(kernel_function, args, kwargs, grid, threads): + warp_args = [wp.from_torch(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] + wp.launch(kernel_function, dim=grid[0] * threads[0], inputs=warp_args, block_dim=threads[0]) diff --git a/test/test_generic_python_dsls.py b/test/test_generic_python_dsls.py new file mode 100644 index 000000000..6c88bb82e --- /dev/null +++ b/test/test_generic_python_dsls.py @@ -0,0 +1,128 @@ +"""End-to-end tests for tuning kernels written in Python DSLs with lang="generic_python".""" + +import importlib + +import numpy as np +import pytest + +from kernel_tuner import tune_kernel +from kernel_tuner.backends.generic_python import GenericPythonFunctions + +from .context import ( + skip_if_no_cupy, + skip_if_no_cute, + skip_if_no_numba_cuda, + skip_if_no_taichi, + skip_if_no_tilelang, + skip_if_no_tilus, + skip_if_no_torch, + skip_if_no_triton, + skip_if_no_warp, +) + +# every DSL needs Torch with a CUDA device, which the generic_python backend uses for memory and timing +dsls = [ + pytest.param("numba", marks=[skip_if_no_torch, skip_if_no_numba_cuda]), + pytest.param("cupyx", marks=[skip_if_no_torch, skip_if_no_cupy]), + pytest.param("warp", marks=[skip_if_no_torch, skip_if_no_warp]), + pytest.param("taichi", marks=[skip_if_no_torch, skip_if_no_taichi]), + pytest.param("cute", marks=[skip_if_no_torch, skip_if_no_cute]), + pytest.param("triton", marks=[skip_if_no_torch, skip_if_no_triton]), + pytest.param("tilus", marks=[skip_if_no_torch, skip_if_no_tilus]), + pytest.param("tilelang", marks=[skip_if_no_torch, skip_if_no_tilelang]), +] + + +@pytest.mark.parametrize("dsl", dsls) +def test_generic_python_dsl(dsl): + # the kernel modules import their DSL, so only import them when the DSL is installed + kernel = importlib.import_module(f".generic_python_kernels.{dsl}_vector_add", package=__package__) + if getattr(kernel, "requires_skip", None): + pytest.skip(kernel.requires_skip) + + n = 4096 + a = np.random.randn(n).astype(np.float32) + b = np.random.randn(n).astype(np.float32) + c = np.zeros_like(a) + args = kernel.arguments(c, a, b, n) + answer = [a + b if arg is c else None for arg in args] + tune_params = kernel.tune_params(n) + + results, _ = tune_kernel( + kernel.kernel_name, + kernel.__file__, + n, + args, + tune_params, + lang="generic_python", + call_function=kernel.call_function, + answer=answer, + ) + + num_configs = np.prod([len(v) for v in tune_params.values()]) + assert len(results) == num_configs + # verification raises an exception on wrong results, so every configuration must have a valid time + failed = [result for result in results if not isinstance(result.get("time"), float)] + assert not failed, f"configurations failed: {failed}" + + +def wrapped(error, wrapper=None): + """Raise error and wrap it the way DSLs and call functions do, returning the wrapper.""" + wrapper = wrapper or RuntimeError("wrapper without details") + try: + try: + raise error + except Exception as e: + raise wrapper from e + except Exception as e: + return e + + +# messages taken from the errors DSLs raise for configurations that use too many resources +resource_errors = [ + RuntimeError("out of resource: shared memory, Required: 163840, Hardware limit: 101376."), # Triton + ValueError("The register usage (256) of given config is too high."), # Tilus + RuntimeError("No valid schedule found during building programs:"), # Tilus + RuntimeError("Failed to set the allowed dynamic shared memory size to 131072"), # TileLang + RuntimeError("No valid warp partition for T.gemm: M=32, N=32 cannot be evenly covered by 16 warps"), # TileLang + RuntimeError("CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES: This indicates that a launch did not occur"), # Numba + RuntimeError("ERROR_PTX_COMPILE (4)\nnvJitLink error log: ptxas error : Entry function uses too much shared data"), + RuntimeError("error[CUDA_LAUNCH_INVALID_CONFIG]: CUDA launch failed: cudaErrorInvalidValue (1)"), # CuTe + RuntimeError("error: cudaErrorInvalidConfiguration (error code: 9)"), # CuTe + TypeError("Invalid configuration, CuTe compile failed with TypeError"), # resource message wins over type + wrapped(RuntimeError("out of resource: shared memory")), # message only in the chained exception +] + +user_errors = [ + NameError("name 'num_warps' is not defined"), # code errors win over resource keywords + SyntaxError("invalid syntax"), + AttributeError("module 'triton.language' has no attribute 'lod'"), + TypeError("vector_add() takes 4 positional arguments but 5 were given"), + RuntimeError("Type mismatch: float32 and int32"), + wrapped(NameError("name 'x' is not defined")), + wrapped(TypeError("unsupported operand")), +] + +unknown_errors = [ + # words in file paths or source lines, such as "cuda" and "ast", should not affect classification + RuntimeError("illegal memory access in /site-packages/numba_cuda/cudadrv/driver.py, last call failed"), + ValueError("expected a positive value"), + # a problem with the installation, not with the configuration + RuntimeError("ERROR_OUTDATED_LIBRARY (14)\nnvJitLink error log: ERROR NVVM_ERROR_INVALID_INPUT (4)"), + wrapped(ValueError("something went wrong"), RuntimeError("")), +] + + +@pytest.mark.parametrize("error", resource_errors, ids=repr) +def test_classify_compile_exception_resource(error): + assert GenericPythonFunctions.classify_compile_exception(None, error) == "resource_error" + + +@pytest.mark.parametrize("error", user_errors, ids=repr) +def test_classify_compile_exception_user(error): + assert GenericPythonFunctions.classify_compile_exception(None, error) == "user_error" + + +@pytest.mark.parametrize("error", unknown_errors, ids=repr) +def test_classify_compile_exception_unknown(error): + assert GenericPythonFunctions.classify_compile_exception(None, error) == "unknown" From 44ce47680820a2ab6c0c00d5f4c24f327a09023e Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 11:58:50 +0200 Subject: [PATCH 26/44] add cuTile support to generic Python backend - add call_cutile, which appends tunable kernel arguments to the positional arguments of ct.launch and limits the tileiras compile time - add cuTile vec_add and matmul (basic and swizzled) examples - classify compile timeouts and non power of two tile sizes as errors that skip the configuration - add cuTile to the parametrized DSL tests with a skip_if_no_cutile marker --- CHANGELOG.md | 2 +- examples/python/call_functions.py | 11 ++ examples/python/cutile_vec_add.py | 34 ++++++ examples/python/matmul/cutile_matmul.py | 107 ++++++++++++++++++ kernel_tuner/backends/generic_python.py | 4 + test/context.py | 2 + .../cutile_vector_add.py | 27 +++++ test/test_generic_python_dsls.py | 5 + 8 files changed, 191 insertions(+), 1 deletion(-) create mode 100644 examples/python/cutile_vec_add.py create mode 100644 examples/python/matmul/cutile_matmul.py create mode 100644 test/generic_python_kernels/cutile_vector_add.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fe3a75298..aafcc53c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). ## Unreleased - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers -- Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` +- Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, cuTile, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` - Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend - Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel - Fix output verification failing for PyTorch tensors as expected answers diff --git a/examples/python/call_functions.py b/examples/python/call_functions.py index e3132771e..63dc7b6e5 100644 --- a/examples/python/call_functions.py +++ b/examples/python/call_functions.py @@ -112,3 +112,14 @@ def call_warp(kernel_function, args, kwargs, grid, threads, params): inputs=warp_args, block_dim=threads_per_block ) + + +def call_cutile(kernel_function, args, kwargs, grid): + import cuda.tile as ct + + # ct.launch only accepts positional arguments, tunable parameters that are also + # kernel arguments are passed in kwargs and should be the last kernel arguments. + # Large tiles can take very long to compile, so limit the compile time, the + # resulting timeout error makes Kernel Tuner skip the configuration. + with ct.compiler_timeout(60): + ct.launch(torch.cuda.current_stream(), grid, kernel_function, (*args, *kwargs.values())) diff --git a/examples/python/cutile_vec_add.py b/examples/python/cutile_vec_add.py new file mode 100644 index 000000000..d81d344fe --- /dev/null +++ b/examples/python/cutile_vec_add.py @@ -0,0 +1,34 @@ +import cuda.tile as ct +import torch + +from kernel_tuner import tune_kernel +from call_functions import call_cutile + + +@ct.kernel +def vec_add(a, b, c, tile_size: ct.Constant[int]): + pid = ct.bid(0) + a_tile = ct.load(a, index=(pid,), shape=(tile_size,)) + b_tile = ct.load(b, index=(pid,), shape=(tile_size,)) + ct.store(c, index=(pid,), tile=a_tile + b_tile) + + +def tune(): + size = 1 << 24 + a = torch.randn(size, device="cuda", dtype=torch.float32) + b = torch.randn(size, device="cuda", dtype=torch.float32) + c = torch.zeros(size, device="cuda", dtype=torch.float32) + + args = [a, b, c] + # tile sizes must be powers of 2 + tune_params = {"tile_size": [2**i for i in range(5, 13)]} + answer = [None, None, (a + b).cpu()] + + # Each block processes one tile, so the tile size acts as the block size to compute the grid + results, env = tune_kernel("vec_add", __file__, size, args, tune_params, answer=answer, + block_size_names=["tile_size"], lang="generic_python", + call_function=call_cutile) + + +if __name__ == "__main__": + tune() diff --git a/examples/python/matmul/cutile_matmul.py b/examples/python/matmul/cutile_matmul.py new file mode 100644 index 000000000..cc6d34ea1 --- /dev/null +++ b/examples/python/matmul/cutile_matmul.py @@ -0,0 +1,107 @@ +import cuda.tile as ct +import torch + +from kernel_tuner import tune_kernel +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py +from call_functions import call_cutile # noqa: E402 + + +@ct.kernel +def matmul(A, B, C, tile_m: ct.Constant[int], tile_n: ct.Constant[int], tile_k: ct.Constant[int]): + bidx = ct.bid(0) + bidy = ct.bid(1) + + # accumulate in fp32 over all tiles in the K dimension + acc = ct.full((tile_m, tile_n), 0, dtype=ct.float32) + for k in range(ct.num_tiles(A, axis=1, shape=(tile_m, tile_k))): + a = ct.load(A, index=(bidx, k), shape=(tile_m, tile_k), padding_mode=ct.PaddingMode.ZERO) + b = ct.load(B, index=(k, bidy), shape=(tile_k, tile_n), padding_mode=ct.PaddingMode.ZERO) + acc = ct.mma(a, b, acc) + + ct.store(C, index=(bidx, bidy), tile=ct.astype(acc, C.dtype)) + + +@ct.kernel +def matmul_swizzled(A, B, C, tile_m: ct.Constant[int], tile_n: ct.Constant[int], tile_k: ct.Constant[int], + group_size_m: ct.Constant[int]): + # map the 1D block index to a tile of C in groups of group_size_m rows of tiles, + # so consecutive blocks reuse the same tiles of B, which improves L2 cache reuse + bid = ct.bid(0) + num_bid_m = ct.cdiv(A.shape[0], tile_m) + num_bid_n = ct.cdiv(B.shape[1], tile_n) + num_bid_in_group = group_size_m * num_bid_n + first_bid_m = (bid // num_bid_in_group) * group_size_m + group_size = ct.minimum(num_bid_m - first_bid_m, group_size_m) + bidx = first_bid_m + (bid % num_bid_in_group) % group_size + bidy = (bid % num_bid_in_group) // group_size + + acc = ct.full((tile_m, tile_n), 0, dtype=ct.float32) + for k in range(ct.num_tiles(A, axis=1, shape=(tile_m, tile_k))): + a = ct.load(A, index=(bidx, k), shape=(tile_m, tile_k), padding_mode=ct.PaddingMode.ZERO) + b = ct.load(B, index=(k, bidy), shape=(tile_k, tile_n), padding_mode=ct.PaddingMode.ZERO) + acc = ct.mma(a, b, acc) + + ct.store(C, index=(bidx, bidy), tile=ct.astype(acc, C.dtype)) + + +def run(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + + tile_m, tile_n, tile_k = 64, 64, 32 + grid = (ct.cdiv(M, tile_m), ct.cdiv(N, tile_n), 1) + ct.launch(torch.cuda.current_stream(), grid, matmul, (A, B, C, tile_m, tile_n, tile_k)) + + assert torch.allclose(C, A @ B, rtol=1e-2, atol=M * 2**(-11)) + print("Passed") + + +def tune(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + + args = [A, B, C] + answer = [None, None, (A @ B).cpu()] + + tune_params = dict() + tune_params["tile_m"] = [32, 64, 128, 256] + tune_params["tile_n"] = [32, 64, 128, 256] + tune_params["tile_k"] = [16, 32, 64, 128] + + # each block computes one tile of C, so the tile sizes act as block sizes to compute the grid + results, env = tune_kernel("matmul", __file__, (M, N), args, tune_params, answer=answer, + atol=M * 2**(-11), block_size_names=["tile_m", "tile_n"], + lang="generic_python", call_function=call_cutile) + + +def tune_swizzled(M, N, K): + A = torch.randn(M, K, device="cuda", dtype=torch.float16) + B = torch.randn(K, N, device="cuda", dtype=torch.float16) + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + + args = [A, B, C] + answer = [None, None, (A @ B).cpu()] + + tune_params = dict() + tune_params["tile_m"] = [32, 64, 128, 256] + tune_params["tile_n"] = [32, 64, 128, 256] + tune_params["tile_k"] = [16, 32, 64, 128] + tune_params["group_size_m"] = [1, 2, 4, 8, 16] + + # 1D grid with one block per tile of C + results, env = tune_kernel("matmul_swizzled", __file__, M * N, args, tune_params, answer=answer, + atol=M * 2**(-11), block_size_names=["tile_m"], grid_div_x=["tile_m", "tile_n"], + strategy="bayes_opt", strategy_options={"max_fevals": 100}, + lang="generic_python", call_function=call_cutile) + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + run(M, N, K) + tune(M, N, K) + tune_swizzled(M, N, K) diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 5f26b8236..5acaa8e7a 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -296,6 +296,7 @@ def classify_compile_exception(self, e): r"\bnum_ctas\b", r"no valid warp partition", r"no valid schedule", + r"power of (two|2)", # many tile based DSLs require power of two tile sizes # Generic resource exhaustion r"too many resources", @@ -304,6 +305,9 @@ def classify_compile_exception(self, e): r"exceeds maximum", r"exceeds limit", + # Compilation taking too long, usually caused by large tile sizes + r"exceeded timeout", + # Errors reported by the compiler toolchain r"\bptxas\b", r"\bnvcc error", diff --git a/test/context.py b/test/context.py index dd54b3462..298303652 100644 --- a/test/context.py +++ b/test/context.py @@ -107,6 +107,7 @@ def python_dsl_present(module, subpackage=None): triton_present = python_dsl_present("triton") tilus_present = python_dsl_present("tilus") tilelang_present = python_dsl_present("tilelang") +cutile_present = python_dsl_present("cuda", "tile") try: import pyatf @@ -161,6 +162,7 @@ def python_dsl_present(module, subpackage=None): skip_if_no_triton = pytest.mark.skipif(not triton_present, reason="Triton not installed") skip_if_no_tilus = pytest.mark.skipif(not tilus_present, reason="Tilus not installed") skip_if_no_tilelang = pytest.mark.skipif(not tilelang_present, reason="TileLang not installed") +skip_if_no_cutile = pytest.mark.skipif(not cutile_present, reason="cuTile not installed") def skip_backend(backend: str): diff --git a/test/generic_python_kernels/cutile_vector_add.py b/test/generic_python_kernels/cutile_vector_add.py new file mode 100644 index 000000000..7f702ad39 --- /dev/null +++ b/test/generic_python_kernels/cutile_vector_add.py @@ -0,0 +1,27 @@ +import cuda.tile as ct +import torch + +kernel_name = "vector_add" + + +@ct.kernel +def vector_add(c, a, b, block_size_x: ct.Constant[int]): + pid = ct.bid(0) + a_tile = ct.load(a, index=(pid,), shape=(block_size_x,)) + b_tile = ct.load(b, index=(pid,), shape=(block_size_x,)) + ct.store(c, index=(pid,), tile=a_tile + b_tile) + + +def arguments(c, a, b, n): + # tiles are loaded with bounds checks, so n is not needed + return [c, a, b] + + +def tune_params(n): + # each block processes one tile, so the tile size is the block size used to compute the grid + return {"block_size_x": [128, 256]} + + +def call_function(kernel_function, args, kwargs, grid): + # ct.launch only takes positional arguments, the tunable kernel argument block_size_x comes last + ct.launch(torch.cuda.current_stream(), grid, kernel_function, (*args, *kwargs.values())) diff --git a/test/test_generic_python_dsls.py b/test/test_generic_python_dsls.py index 6c88bb82e..51e7d3a01 100644 --- a/test/test_generic_python_dsls.py +++ b/test/test_generic_python_dsls.py @@ -11,6 +11,7 @@ from .context import ( skip_if_no_cupy, skip_if_no_cute, + skip_if_no_cutile, skip_if_no_numba_cuda, skip_if_no_taichi, skip_if_no_tilelang, @@ -30,6 +31,7 @@ pytest.param("triton", marks=[skip_if_no_torch, skip_if_no_triton]), pytest.param("tilus", marks=[skip_if_no_torch, skip_if_no_tilus]), pytest.param("tilelang", marks=[skip_if_no_torch, skip_if_no_tilelang]), + pytest.param("cutile", marks=[skip_if_no_torch, skip_if_no_cutile]), ] @@ -90,6 +92,8 @@ def wrapped(error, wrapper=None): RuntimeError("error[CUDA_LAUNCH_INVALID_CONFIG]: CUDA launch failed: cudaErrorInvalidValue (1)"), # CuTe RuntimeError("error: cudaErrorInvalidConfiguration (error code: 9)"), # CuTe TypeError("Invalid configuration, CuTe compile failed with TypeError"), # resource message wins over type + RuntimeError("Invalid argument \"shape\" of load(): Dimension #0 of shape (96,) is not a power of two"), # cuTile + RuntimeError("`tileiras` compiler exceeded timeout 30s. Using a smaller tile size may reduce compilation time."), wrapped(RuntimeError("out of resource: shared memory")), # message only in the chained exception ] @@ -99,6 +103,7 @@ def wrapped(error, wrapper=None): AttributeError("module 'triton.language' has no attribute 'lod'"), TypeError("vector_add() takes 4 positional arguments but 5 were given"), RuntimeError("Type mismatch: float32 and int32"), + RuntimeError("Undefined variable undefined_name used"), # cuTile wrapped(NameError("name 'x' is not defined")), wrapped(TypeError("unsupported operand")), ] From f656a7d1194ba16350e2fab8e09dafc74abb7f70 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 14:16:03 +0200 Subject: [PATCH 27/44] detect Python DSLs and select default call functions Kernels in Python files are now detected as lang="generic_python", and when no call_function is passed, the DSL of the kernel is detected from its decorators or base classes, resolved through the imports in the source file, to select a default call function. Detection only parses the source, so none of the optional DSLs are imported. - move detect_language from util.py to utils/language_detection.py, together with is_python_file and detect_python_dsl - move the call functions from the examples to utils/call_functions.py - examples no longer pass lang or call_function, except where they launch the kernel in a custom way - test modules use the default call functions, add tests for language and DSL detection and for not importing DSLs during detection --- CHANGELOG.md | 1 + examples/python/call_functions.py | 125 ---------------- examples/python/cute_vec_add.py | 3 +- examples/python/cutile_vec_add.py | 4 +- examples/python/matmul/cupy_matmul.py | 6 - examples/python/matmul/cute_matmul.py | 10 +- examples/python/matmul/cutile_matmul.py | 10 +- examples/python/matmul/numba_matmul.py | 8 -- examples/python/matmul/taichi_matmul.py | 6 - examples/python/matmul/tilelang_matmul.py | 8 -- examples/python/matmul/tilus_matmul.py | 8 -- examples/python/matmul/triton_matmul.py | 8 -- examples/python/matmul/warp_matmul.py | 6 - examples/python/numba_vec_add.py | 3 - examples/python/tilelang_vec_add.py | 4 +- examples/python/tilus_vec_add.py | 5 - examples/python/triton_vec_add.py | 6 +- examples/python/warp_vec_add.py | 2 - kernel_tuner/interface.py | 8 +- kernel_tuner/kernel_sources/kernel_source.py | 14 +- .../kernel_sources/kernel_source_fn.py | 33 +++-- kernel_tuner/util.py | 14 -- kernel_tuner/utils/call_functions.py | 134 ++++++++++++++++++ kernel_tuner/utils/language_detection.py | 122 ++++++++++++++++ test/generic_python_kernels/__init__.py | 3 +- .../cupyx_vector_add.py | 7 - .../generic_python_kernels/cute_vector_add.py | 13 -- .../cutile_vector_add.py | 6 - .../numba_vector_add.py | 6 - .../taichi_vector_add.py | 4 - .../tilelang_vector_add.py | 4 - .../tilus_vector_add.py | 4 - .../triton_vector_add.py | 4 - .../generic_python_kernels/warp_vector_add.py | 8 -- test/test_generic_python_dsls.py | 12 +- test/test_kernel_source_fn.py | 2 +- test/test_language_detection.py | 120 ++++++++++++++++ test/test_lazy_imports.py | 19 +++ test/test_util_functions.py | 19 --- 39 files changed, 450 insertions(+), 329 deletions(-) delete mode 100644 examples/python/call_functions.py create mode 100644 kernel_tuner/utils/call_functions.py create mode 100644 kernel_tuner/utils/language_detection.py create mode 100644 test/test_language_detection.py diff --git a/CHANGELOG.md b/CHANGELOG.md index aafcc53c7..d24e52063 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). ## Unreleased - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers - Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, cuTile, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` +- Detect kernels in Python files as `lang="generic_python"`, and detect the DSL of Python kernels to select a default `call_function` when none is given - Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend - Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel - Fix output verification failing for PyTorch tensors as expected answers diff --git a/examples/python/call_functions.py b/examples/python/call_functions.py deleted file mode 100644 index 63dc7b6e5..000000000 --- a/examples/python/call_functions.py +++ /dev/null @@ -1,125 +0,0 @@ -import torch - -def call_tilus(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) - - -def call_triton(kernel_function, args, kwargs, grid, threads, params): - if "num_warps" in params.keys(): - kwargs["num_warps"] = params["num_warps"] - if "num_stages" in params.keys(): - kwargs["num_stages"] = params["num_stages"] - - kernel_function[grid](*args, **kwargs) - - -def call_tilelang(kernel_function, args, kwargs): - compiled_kernel = kernel_function(**kwargs) - compiled_kernel(*args) - - -def call_numba(kernel_function, args, kwargs, grid, threads): - from numba import cuda - - numba_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - numba_args.append(cuda.as_cuda_array(arg)) - else: - numba_args.append(arg) - - kernel_function[grid, threads](*numba_args, **kwargs) - - -def call_cupyx(kernel_function, args, kwargs, grid, threads): - import cupy as cp - - cupy_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - cupy_args.append(cp.from_dlpack(arg)) - else: - cupy_args.append(arg) - kernel_function(grid, threads, tuple(cupy_args)) - - -def call_cute(kernel_function, args, kwargs, grid, threads, params): - import cutlass.cute as cute - from cutlass.cute.runtime import from_dlpack - - # Initialize cache if it does not exist - if not hasattr(call_cute, "custom_cache"): - call_cute.custom_cache = {} - - # Convert Torch tensors to CuTe tensors with correct layout - cute_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - arg_ = from_dlpack(arg) - cute_args.append(arg_) - else: - cute_args.append(arg) - - # Form cache key from tuning parameters - param_keys = sorted(params.keys()) - cache_str = type(kernel_function).__name__ - for k in param_keys: - cache_str += "_" + str(params[k]) - - # Check if kernel exists in cache. Otherwise, compile and save - if cache_str in call_cute.custom_cache: - compiled_kernel = call_cute.custom_cache[cache_str] - else: - compiled_kernel = cute.compile(kernel_function, *cute_args) - call_cute.custom_cache[cache_str] = compiled_kernel - - compiled_kernel(*cute_args, **kwargs) - - -def call_taichi(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) - - -def call_warp(kernel_function, args, kwargs, grid, threads, params): - import warp as wp - - # Convert Torch tensors to Warp args - warp_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - warp_args.append(wp.from_torch(arg)) - else: - warp_args.append(arg) - - # Check if block_dim is in the tuning parameters. Otherwise, use - # the computed thread block dimensions. - if 'block_dim' in params.keys(): - threads_per_block = params['block_dim'] - else: - threads_per_block = threads[0] * threads[1] * threads[2] - - # Check if dim is in the tuning parameters. Otherwise, compute from - # grid and threads. - if 'dim' in params.keys(): - dimensions = params['dim'] - else: - dimensions = [grid[i] * threads[i] for i in range(len(grid))] - - # launch kernel - wp.launch( - kernel_function, - dim=dimensions, - inputs=warp_args, - block_dim=threads_per_block - ) - - -def call_cutile(kernel_function, args, kwargs, grid): - import cuda.tile as ct - - # ct.launch only accepts positional arguments, tunable parameters that are also - # kernel arguments are passed in kwargs and should be the last kernel arguments. - # Large tiles can take very long to compile, so limit the compile time, the - # resulting timeout error makes Kernel Tuner skip the configuration. - with ct.compiler_timeout(60): - ct.launch(torch.cuda.current_stream(), grid, kernel_function, (*args, *kwargs.values())) diff --git a/examples/python/cute_vec_add.py b/examples/python/cute_vec_add.py index ecd2386e4..36baeb0f1 100644 --- a/examples/python/cute_vec_add.py +++ b/examples/python/cute_vec_add.py @@ -4,7 +4,6 @@ import cutlass.cute as cute from kernel_tuner import tune_kernel -from call_functions import call_cute @cute.kernel def vec_add_kernel( @@ -50,7 +49,7 @@ def main(): answer = [None, None, (a+b).cpu(), None] tune_kernel("vec_add", __file__, size, args, tune_params, answer=answer, - lang="generic_python", call_function=call_cute, verbose=True) + verbose=True) if __name__ == "__main__": diff --git a/examples/python/cutile_vec_add.py b/examples/python/cutile_vec_add.py index d81d344fe..003ae2798 100644 --- a/examples/python/cutile_vec_add.py +++ b/examples/python/cutile_vec_add.py @@ -2,7 +2,6 @@ import torch from kernel_tuner import tune_kernel -from call_functions import call_cutile @ct.kernel @@ -26,8 +25,7 @@ def tune(): # Each block processes one tile, so the tile size acts as the block size to compute the grid results, env = tune_kernel("vec_add", __file__, size, args, tune_params, answer=answer, - block_size_names=["tile_size"], lang="generic_python", - call_function=call_cutile) + block_size_names=["tile_size"]) if __name__ == "__main__": diff --git a/examples/python/matmul/cupy_matmul.py b/examples/python/matmul/cupy_matmul.py index be95fad98..300bd4c2f 100644 --- a/examples/python/matmul/cupy_matmul.py +++ b/examples/python/matmul/cupy_matmul.py @@ -3,11 +3,7 @@ import numpy as np from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_cupyx # noqa: E402 @jit.rawkernel() @@ -62,8 +58,6 @@ def tune(M, N, K): tune_params=tune_params, answer=[None, None, A.dot(B), None, None, None], atol=1e-1, - call_function=call_cupyx, - lang="generic_python", restrictions=restrictions, verbose=True, ) diff --git a/examples/python/matmul/cute_matmul.py b/examples/python/matmul/cute_matmul.py index 681218b40..a84f76d59 100644 --- a/examples/python/matmul/cute_matmul.py +++ b/examples/python/matmul/cute_matmul.py @@ -8,11 +8,7 @@ from cutlass.cute.runtime import from_dlpack from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_cute # noqa: E402 # might need export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH to work @@ -94,8 +90,8 @@ def tune_naive_matmul(M, N, K): tune_params["block_size_y"] = [2**i for i in range(1, 10)] restrictions = ["block_size_x * block_size_y >= 32", "block_size_x * block_size_y <= 1024"] - results, env = tune_kernel("matmul", __file__, size, args, tune_params, lang="generic_python", - call_function=call_cute, answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) + results, env = tune_kernel("matmul", __file__, size, args, tune_params, + answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) @@ -1057,7 +1053,7 @@ def tune_optimized(M, N, K, L=1): results, env = tune_kernel("TensorOpGemm", __file__, M * N, args, tune_params, verbose=True, restrictions=restrictions, #strategy="bayes_opt", - lang="generic_python", call_function=call_cute_custom, answer=[None, None, c_ref.cpu()], atol=M * 2**(-10)) + call_function=call_cute_custom, answer=[None, None, c_ref.cpu()], atol=M * 2**(-10)) diff --git a/examples/python/matmul/cutile_matmul.py b/examples/python/matmul/cutile_matmul.py index cc6d34ea1..19d3b54d7 100644 --- a/examples/python/matmul/cutile_matmul.py +++ b/examples/python/matmul/cutile_matmul.py @@ -2,11 +2,7 @@ import torch from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_cutile # noqa: E402 @ct.kernel @@ -75,8 +71,7 @@ def tune(M, N, K): # each block computes one tile of C, so the tile sizes act as block sizes to compute the grid results, env = tune_kernel("matmul", __file__, (M, N), args, tune_params, answer=answer, - atol=M * 2**(-11), block_size_names=["tile_m", "tile_n"], - lang="generic_python", call_function=call_cutile) + atol=M * 2**(-11), block_size_names=["tile_m", "tile_n"]) def tune_swizzled(M, N, K): @@ -96,8 +91,7 @@ def tune_swizzled(M, N, K): # 1D grid with one block per tile of C results, env = tune_kernel("matmul_swizzled", __file__, M * N, args, tune_params, answer=answer, atol=M * 2**(-11), block_size_names=["tile_m"], grid_div_x=["tile_m", "tile_n"], - strategy="bayes_opt", strategy_options={"max_fevals": 100}, - lang="generic_python", call_function=call_cutile) + strategy="bayes_opt", strategy_options={"max_fevals": 100}) if __name__ == "__main__": diff --git a/examples/python/matmul/numba_matmul.py b/examples/python/matmul/numba_matmul.py index 69d0585f9..345fd7b1f 100644 --- a/examples/python/matmul/numba_matmul.py +++ b/examples/python/matmul/numba_matmul.py @@ -2,11 +2,7 @@ from numba import cuda, float32 from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_numba # noqa: E402 # Source: https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication @@ -192,11 +188,9 @@ def tune_basic(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=answer, atol=atol, restrictions=restrictions, - call_function=call_numba, ) @@ -236,7 +230,6 @@ def tune_optimized(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=answer, atol=atol, restrictions = restrictions, @@ -245,7 +238,6 @@ def tune_optimized(M, N, K): grid_div_y = ["BLK_N"], strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, - call_function=call_numba, ) diff --git a/examples/python/matmul/taichi_matmul.py b/examples/python/matmul/taichi_matmul.py index 4992ed086..df19a52a7 100644 --- a/examples/python/matmul/taichi_matmul.py +++ b/examples/python/matmul/taichi_matmul.py @@ -2,11 +2,7 @@ import taichi as ti from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_taichi # noqa: E402 ti.init(arch=ti.gpu) @@ -53,8 +49,6 @@ def tune(M, N, K): arguments=args, tune_params=tune_params, answer=answer, - lang="generic_python", - call_function=call_taichi, atol=M * 2**(-11), block_size_names = ["BLOCK_DIM"], ) diff --git a/examples/python/matmul/tilelang_matmul.py b/examples/python/matmul/tilelang_matmul.py index 0e7658639..a4b08f6b6 100644 --- a/examples/python/matmul/tilelang_matmul.py +++ b/examples/python/matmul/tilelang_matmul.py @@ -5,11 +5,7 @@ import itertools from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_tilelang # noqa: E402 # https://github.com/tile-ai/tilelang/tree/main/examples/gemm # num_threads and num_stages added as variables to enable tuning. @@ -262,12 +258,10 @@ def tune_basic(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, C_ref.cpu()], atol=M * 2**(-11), strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, - call_function=call_tilelang, verbose=True, restrictions=restrictions, ) @@ -304,12 +298,10 @@ def tune_opt(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, C_ref.cpu()], atol=M * 2**(-11), strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, - call_function=call_tilelang, verbose=True, restrictions=restrictions, ) diff --git a/examples/python/matmul/tilus_matmul.py b/examples/python/matmul/tilus_matmul.py index 7a75197e7..798e973e3 100644 --- a/examples/python/matmul/tilus_matmul.py +++ b/examples/python/matmul/tilus_matmul.py @@ -5,11 +5,7 @@ from tilus.utils import cdiv from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_tilus # noqa: E402 # This kernel is copied from the Tilus project: @@ -211,12 +207,10 @@ def tune_basic(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, None, None, None, C_ref.cpu()], atol=M * 2**(-11), strategy = "bayes_opt", strategy_options = {"max_fevals": 200}, - call_function=call_tilus, ) @@ -246,13 +240,11 @@ def tune_opt(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, None, None, None, C_ref.cpu()], atol=M * 2**(-11), restrictions=restrictions, strategy = "bayes_opt", strategy_options = {"max_fevals": 200}, - call_function=call_tilus, ) diff --git a/examples/python/matmul/triton_matmul.py b/examples/python/matmul/triton_matmul.py index 9ffd5c282..429735855 100644 --- a/examples/python/matmul/triton_matmul.py +++ b/examples/python/matmul/triton_matmul.py @@ -3,11 +3,7 @@ import triton.language as tl from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_triton # noqa: E402 @@ -231,13 +227,11 @@ def tune_basic(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, C_ref.cpu(), None, None, None], atol=M * 2**(-11), block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, - call_function=call_triton, ) @@ -273,7 +267,6 @@ def tune_opt(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, C_ref.cpu(), None, None, None, None, None, None, None, None, None], atol=M * 2**(-11), block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], @@ -281,7 +274,6 @@ def tune_opt(M, N, K): restrictions = restrictions, strategy = "bayes_opt", strategy_options = {"max_fevals": 100}, - call_function=call_triton, ) diff --git a/examples/python/matmul/warp_matmul.py b/examples/python/matmul/warp_matmul.py index 051eb506f..ef10f6b1b 100644 --- a/examples/python/matmul/warp_matmul.py +++ b/examples/python/matmul/warp_matmul.py @@ -2,11 +2,7 @@ import warp as wp from kernel_tuner import tune_kernel -import sys -from pathlib import Path -sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # for call_functions.py -from call_functions import call_warp # noqa: E402 wp.init() wp.config.enable_backward = False @@ -71,9 +67,7 @@ def tune(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=answer, - call_function=call_warp, ) diff --git a/examples/python/numba_vec_add.py b/examples/python/numba_vec_add.py index 3bc90de86..df39d6f70 100644 --- a/examples/python/numba_vec_add.py +++ b/examples/python/numba_vec_add.py @@ -2,7 +2,6 @@ import numpy as np from kernel_tuner import tune_kernel -from call_functions import call_numba @cuda.jit @@ -33,9 +32,7 @@ def tune(): problem_size=N, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, c_expect], - call_function=call_numba, ) if __name__ == "__main__": diff --git a/examples/python/tilelang_vec_add.py b/examples/python/tilelang_vec_add.py index 38b4ab309..32d389d24 100644 --- a/examples/python/tilelang_vec_add.py +++ b/examples/python/tilelang_vec_add.py @@ -3,7 +3,6 @@ import torch from kernel_tuner import tune_kernel -from call_functions import call_tilelang @tilelang.jit # infers target from tensors at first call @@ -37,8 +36,7 @@ def tune(): answer = [None, None, (A + B).cpu()] - res, env = tune_kernel("add", __file__, N, args, tune_params, lang="generic_python", - call_function=call_tilelang, answer=answer) + res, env = tune_kernel("add", __file__, N, args, tune_params, answer=answer) if __name__ == "__main__": tune() \ No newline at end of file diff --git a/examples/python/tilus_vec_add.py b/examples/python/tilus_vec_add.py index ca215503a..e509b307e 100644 --- a/examples/python/tilus_vec_add.py +++ b/examples/python/tilus_vec_add.py @@ -4,7 +4,6 @@ import torch from kernel_tuner import tune_kernel, run_kernel -from call_functions import call_tilus @@ -63,9 +62,7 @@ def tune(size): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, None, c_expect.cpu()], - call_function=call_tilus, verbose=True, ) @@ -85,8 +82,6 @@ def run(size): problem_size=size, arguments=args, params={"block_size_x": 32, "num_warps": 4}, - lang="generic_python", - call_function=call_tilus, ) c_expect = c_expect.cpu() diff --git a/examples/python/triton_vec_add.py b/examples/python/triton_vec_add.py index c4f6bb643..e82642317 100644 --- a/examples/python/triton_vec_add.py +++ b/examples/python/triton_vec_add.py @@ -4,7 +4,6 @@ import triton.language as tl from kernel_tuner import tune_kernel, run_kernel -from call_functions import call_triton @triton.jit def add_op(x, y): @@ -43,8 +42,7 @@ def tune(): tune_params["block_size_x"] = [2**i for i in range(11)] - result = run_kernel("add_kernel", __file__, size, args, {"block_size_x": 256}, - lang="generic_python", call_function=call_triton) + result = run_kernel("add_kernel", __file__, size, args, {"block_size_x": 256}) assert np.allclose(c_expect.cpu(), result[2]) @@ -54,9 +52,7 @@ def tune(): problem_size=size, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, c_expect.cpu(), None], - call_function=call_triton, ) diff --git a/examples/python/warp_vec_add.py b/examples/python/warp_vec_add.py index 0ccb07a96..aee5a96b3 100644 --- a/examples/python/warp_vec_add.py +++ b/examples/python/warp_vec_add.py @@ -3,7 +3,6 @@ import torch from kernel_tuner import tune_kernel -from call_functions import call_warp wp.init() @@ -62,7 +61,6 @@ def tune(): problem_size=n, arguments=args, tune_params=tune_params, - lang="generic_python", answer=[None, None, c_expect, None], call_function=call_warp, ) diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index fc1b736ba..062ba2314 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -158,7 +158,8 @@ def __deepcopy__(self, _): """Specifies the language used for GPU kernels. The kernel_tuner automatically detects the language, but if it fails, you may specify the language using this argument, currently supported: "CUDA", "CuPy", - "nvcuda", "OpenCL", "HIP", "C", or "Generic_Python.""", + "nvcuda", "OpenCL", "HIP", "C", or "Generic_Python". Kernels in Python + files (.py) are detected as "Generic_Python".""", "string", ), ), @@ -315,7 +316,10 @@ def __deepcopy__(self, _): "call_function", ( """When the language Generic Python is used, a call function that calls the kernel in the Python - DSL must be specified. The function must take the following positional arguments: + DSL. If not specified, the DSL of the kernel is detected from its decorators or base classes and + a default call function from kernel_tuner.utils.call_functions is used. Supported DSLs are + Triton, Numba, CuPy (cupyx.jit), Warp, Taichi, CuTe, Tilus, TileLang, and cuTile. + The function must take the following positional arguments: :kernel_function: the callable function with the tuning parameters inserted. :args: list of kernel arguments, as provided by the user in the argument. :kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, diff --git a/kernel_tuner/kernel_sources/kernel_source.py b/kernel_tuner/kernel_sources/kernel_source.py index 51fb75d9b..bdba5d88b 100644 --- a/kernel_tuner/kernel_sources/kernel_source.py +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -1,8 +1,7 @@ -import kernel_tuner.util as util - from abc import abstractmethod from kernel_tuner.language import Language +from kernel_tuner.utils.language_detection import detect_language, is_python_file from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData @@ -19,8 +18,13 @@ class KernelSourceFactory(type): the __init__ call of the KernelSource class to initalize some common variables. ''' def __call__(cls, kernel_name, kernel_sources, lang, defines=None, call_function=None): - if lang == None: - language = None + first_source = kernel_sources[0] if isinstance(kernel_sources, list) and kernel_sources else kernel_sources + if lang is None and is_python_file(first_source): + # kernels in Python files are written in a Python DSL + lang = Language.GENERIC_PYTHON.value + + if lang is None: + language = None else: try: language = Language(lang.upper()) @@ -62,7 +66,7 @@ def __init__(self, kernel_name, kernel_sources, lang, defines=None): if callable(self.kernel_sources[0]): raise TypeError("Please specify language when using a code generator function") kernel_string = self.get_kernel_string(0) - language = util.detect_language(kernel_string) + language = detect_language(kernel_string) self.lang = Language(language.upper()) else: try: diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index fd4bfa3ad..79395f437 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -13,7 +13,9 @@ from kernel_tuner.language import Language from kernel_tuner.kernel_sources.kernel_source import KernelSource from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData -from kernel_tuner.util import get_kernel_ast, get_arg_names +from kernel_tuner.util import get_kernel_ast, get_arg_names, normalize_call_function +from kernel_tuner.utils.call_functions import DEFAULT_CALL_FUNCTIONS, get_default_call_function +from kernel_tuner.utils.language_detection import detect_python_dsl class KernelSourceFn(KernelSource): @@ -24,8 +26,9 @@ class KernelSourceFn(KernelSource): must be a path to the file where the kernel with kernel_name lives. The kernel can be decorated by a JIT decorator. - A call function to specify how the kernel should be launched must be supplied. The call - function must take the following arguments: + A call function specifies how the kernel should be launched. If no call function is supplied, + the DSL of the kernel is detected and a default call function from kernel_tuner.utils.call_functions + is used. The call function must take the following arguments: - kernel_function: the callable function with the tuning parameters inserted. - args: list of kernel arguments, as provided by the user in the argument. - kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, @@ -40,16 +43,16 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function if isinstance(kernel_source, list): raise ValueError("KernelSourceFn only supports a single kernel source") - if self.lang == Language.GENERIC_PYTHON: + if not isinstance(kernel_name, str): + raise TypeError("kernel_name should be a string, got ", type(kernel_name)) + + if self.lang == Language.GENERIC_PYTHON: if call_function is None: - raise ValueError("call_function must be supplied for language Generic Python") + call_function = self._default_call_function(kernel_name, kernel_source) if not callable(call_function): raise TypeError(f"call_function of type {type(call_function)} is not a callable object.") self.call_function = call_function - if not isinstance(kernel_name, str): - raise TypeError("kernel_name should be a string, got ", type(kernel_name)) - source_ast = get_kernel_ast(kernel_name, kernel_source) if isinstance(source_ast, tuple): # Class based kernel self.source_tree = source_ast[0] @@ -63,6 +66,20 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function self.dependencies = self._find_dependencies(kernel_source) + @staticmethod + def _default_call_function(kernel_name, kernel_source): + """Select the default call function for the DSL the kernel is written in.""" + dsl = detect_python_dsl(kernel_name, kernel_source) + if dsl is None: + raise ValueError( + f"Could not detect the Python DSL of kernel {kernel_name}, please pass a call_function. " + f"Kernels are detected by their decorators or base classes, supported DSLs are " + f"{list(DEFAULT_CALL_FUNCTIONS)}." + ) + logging.debug(f"detected Python DSL {dsl} for kernel {kernel_name}") + return normalize_call_function(get_default_call_function(dsl)) + + def prepare_kernel_instance(self, kernel_options, params, grid, threads): ''' Given the dictionary of tuning parameter values for this configuration, diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index 9320fba62..a16c8df5f 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -559,20 +559,6 @@ def delete_temp_file(filename): raise e -def detect_language(kernel_string): - """Attempt to detect language from the kernel_string.""" - kernel_string = kernel_string.lower().strip() - if "__global__" in kernel_string: - lang = "CUDA" - elif "__kernel" in kernel_string: - lang = "OpenCL" - elif any(token in kernel_string for token in ["@cuda", "@kernel", "@device_code"]): - lang = "Julia" - else: - lang = "C" - return lang - - def get_best_config(results, objective, objective_higher_is_better=False): """Returns the best configuration from a list of results according to some objective.""" func = max if objective_higher_is_better else min diff --git a/kernel_tuner/utils/call_functions.py b/kernel_tuner/utils/call_functions.py new file mode 100644 index 000000000..3c60f0ee3 --- /dev/null +++ b/kernel_tuner/utils/call_functions.py @@ -0,0 +1,134 @@ +"""Default call functions that launch kernels written in Python DSLs, used with ``lang="generic_python"``. + +A call function receives the kernel function with the tunable parameters of the current configuration +inserted, the kernel arguments as prepared by Kernel Tuner (PyTorch CUDA tensors and Python scalars), +and keyword arguments for the tunable parameters that are also arguments of the kernel. Optionally, it +also receives the grid and thread block dimensions computed by Kernel Tuner and the tunable parameters +of the configuration, see :func:`kernel_tuner.util.normalize_call_function`. + +When ``call_function`` is not passed to ``tune_kernel`` or ``run_kernel``, the DSL of the kernel is +detected and the matching call function from this module is used. + +The DSLs are optional dependencies, so they are only imported when a call function is used. +""" + + +def _convert_tensors(args, convert): + """Convert the PyTorch tensors in args using convert, leave other arguments as they are.""" + import torch + + return [convert(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] + + +def call_triton(kernel_function, args, kwargs, grid, threads, params): + """Launch a Triton kernel, num_warps and num_stages are taken from the tunable parameters if present.""" + for option in ("num_warps", "num_stages"): + if option in params: + kwargs[option] = params[option] + kernel_function[grid](*args, **kwargs) + + +def call_numba(kernel_function, args, kwargs, grid, threads): + """Launch a Numba CUDA kernel with the grid and thread block dimensions computed by Kernel Tuner.""" + from numba import cuda + + numba_args = _convert_tensors(args, cuda.as_cuda_array) + kernel_function[grid, threads](*numba_args, **kwargs) + + +def call_cupyx(kernel_function, args, kwargs, grid, threads): + """Launch a CuPy JIT (cupyx.jit.rawkernel) kernel, which only accepts positional arguments.""" + import cupy as cp + + cupy_args = _convert_tensors(args, cp.from_dlpack) + kernel_function(grid, threads, (*cupy_args, *kwargs.values())) + + +def call_warp(kernel_function, args, kwargs, grid, threads, params): + """Launch a Warp kernel. + + The launch dimensions are taken from the tunable parameter ``dim`` if present, otherwise the total + number of threads is used. The thread block size is taken from the tunable parameter ``block_dim`` if + present, otherwise the thread block dimensions computed by Kernel Tuner are used. + """ + import warp as wp + + wp.init() + warp_args = _convert_tensors(args, wp.from_torch) + block_dim = params.get("block_dim", threads[0] * threads[1] * threads[2]) + if "dim" in params: + dim = params["dim"] + else: + # launch with as many dimensions as the problem, kernels using i = wp.tid() need a 1D launch + dim = [g * t for g, t in zip(grid, threads)] + while len(dim) > 1 and dim[-1] == 1: + dim.pop() + wp.launch(kernel_function, dim=dim, inputs=[*warp_args, *kwargs.values()], block_dim=block_dim) + + +def call_taichi(kernel_function, args, kwargs): + """Launch a Taichi kernel, Taichi determines the launch configuration itself.""" + kernel_function(*args, **kwargs) + + +def call_cute(kernel_function, args, kwargs): + """Compile and launch a CuTe DSL kernel. + + The kernel is compiled once per configuration. Kernel Tuner compiles and benchmarks the configurations + one by one, so only the most recently compiled kernel is kept. + """ + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + cute_args = _convert_tensors(args, from_dlpack) + cached_function, compiled_kernel = getattr(call_cute, "_last_compiled", (None, None)) + if cached_function is not kernel_function: + compiled_kernel = cute.compile(kernel_function, *cute_args) + call_cute._last_compiled = (kernel_function, compiled_kernel) + compiled_kernel(*cute_args, **kwargs) + + +def call_tilus(kernel_function, args, kwargs): + """Launch a Tilus kernel, Tilus kernels are scripts that determine the launch configuration themselves.""" + kernel_function(*args, **kwargs) + + +def call_tilelang(kernel_function, args, kwargs): + """Launch a TileLang kernel, the tunable parameters are passed to the kernel factory function.""" + compiled_kernel = kernel_function(**kwargs) + compiled_kernel(*args) + + +def call_cutile(kernel_function, args, kwargs, grid): + """Launch a cuTile kernel on the current PyTorch stream. + + ``ct.launch`` only accepts positional arguments, so the tunable parameters that are also kernel + arguments are passed after the other arguments: these should be the last arguments of the kernel. + Large tiles can take very long to compile, so the compile time is limited to 60 seconds, the + resulting timeout error makes Kernel Tuner skip the configuration. + """ + import cuda.tile as ct + import torch + + with ct.compiler_timeout(60): + ct.launch(torch.cuda.current_stream(), grid, kernel_function, (*args, *kwargs.values())) + + +DEFAULT_CALL_FUNCTIONS = { + "triton": call_triton, + "numba": call_numba, + "cupyx": call_cupyx, + "warp": call_warp, + "taichi": call_taichi, + "cute": call_cute, + "tilus": call_tilus, + "tilelang": call_tilelang, + "cutile": call_cutile, +} + + +def get_default_call_function(dsl): + """Return the default call function for a DSL, as detected by detect_python_dsl.""" + if dsl not in DEFAULT_CALL_FUNCTIONS: + raise ValueError(f"No default call function for DSL {dsl}, supported are {list(DEFAULT_CALL_FUNCTIONS)}") + return DEFAULT_CALL_FUNCTIONS[dsl] diff --git a/kernel_tuner/utils/language_detection.py b/kernel_tuner/utils/language_detection.py new file mode 100644 index 000000000..6d97027f5 --- /dev/null +++ b/kernel_tuner/utils/language_detection.py @@ -0,0 +1,122 @@ +"""Functions to detect the programming language of a kernel, and the DSL used by Python kernels. + +The DSL of a Python kernel is detected from the source code alone: the decorators (or base +classes) of the kernel are resolved using the imports in the source file. This way, none of +the DSLs, which are optional dependencies, have to be imported to detect them. +""" + +import ast +from pathlib import Path + +# Maps the fully qualified names of the decorators and base classes of Python kernels to their DSL. +# A name matches if it is equal to, or starts with, one of these prefixes. +PYTHON_DSL_PREFIXES = { + "triton": "triton", + "numba.cuda": "numba", + "cupyx.jit": "cupyx", + "warp": "warp", + "taichi": "taichi", + "cutlass.cute": "cute", + "tilus": "tilus", + "tilelang": "tilelang", + "cuda.tile": "cutile", +} + + +def detect_language(kernel_string): + """Attempt to detect language from the kernel_string.""" + kernel_string = kernel_string.lower().strip() + if "__global__" in kernel_string: + lang = "CUDA" + elif "__kernel" in kernel_string: + lang = "OpenCL" + elif any(token in kernel_string for token in ["@cuda", "@kernel", "@device_code"]): + lang = "Julia" + else: + lang = "C" + return lang + + +def is_python_file(kernel_source): + """Return True if kernel_source is the path to an existing Python file.""" + if not isinstance(kernel_source, (str, Path)): + return False + try: + path = Path(kernel_source) + return path.suffix == ".py" and path.is_file() + except (OSError, ValueError): # for example, source code strings that are too long to be a path + return False + + +def _import_aliases(tree): + """Map the names bound by the import statements in tree to the fully qualified names they refer to.""" + aliases = {} + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.asname: + aliases[alias.asname] = alias.name + else: + # "import a.b" binds "a", attribute access then resolves the rest + root = alias.name.split(".")[0] + aliases[root] = root + elif isinstance(node, ast.ImportFrom) and node.module and node.level == 0: + for alias in node.names: + aliases[alias.asname or alias.name] = f"{node.module}.{alias.name}" + return aliases + + +def _qualified_name(node, aliases): + """Return the fully qualified name of a decorator or base class expression, or None.""" + # decorators with arguments, such as @jit.rawkernel(), are calls of the decorator + if isinstance(node, ast.Call): + node = node.func + # indexing, such as ct.Constant[int], is not a decorator or base class we recognize + parts = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if not isinstance(node, ast.Name): + return None + parts.append(aliases.get(node.id, node.id)) + return ".".join(reversed(parts)) + + +def _dsl_of_name(name): + """Return the DSL for a fully qualified name, using the longest matching prefix.""" + matches = [prefix for prefix in PYTHON_DSL_PREFIXES if name == prefix or name.startswith(prefix + ".")] + if not matches: + return None + return PYTHON_DSL_PREFIXES[max(matches, key=len)] + + +def detect_python_dsl(kernel_name, filepath): + """Detect the DSL of a Python kernel from its decorators or base classes. + + :param kernel_name: name of the kernel function or class + :type kernel_name: string + + :param filepath: path to the Python file that contains the kernel + :type filepath: string or Path + + :returns: the name of the DSL, one of the values of PYTHON_DSL_PREFIXES, or None if not detected + :rtype: string or None + """ + from kernel_tuner.util import get_kernel_ast, read_file + + kernel_ast = get_kernel_ast(kernel_name, filepath) + aliases = _import_aliases(ast.parse(read_file(filepath), filename=str(filepath))) + + # class based kernels (Tilus) are recognized by their base class or the decorators of __call__ + if isinstance(kernel_ast, tuple): + class_node, call_node = kernel_ast + candidates = class_node.bases + class_node.decorator_list + call_node.decorator_list + else: + candidates = kernel_ast.decorator_list + + for node in candidates: + name = _qualified_name(node, aliases) + dsl = name and _dsl_of_name(name) + if dsl: + return dsl + return None diff --git a/test/generic_python_kernels/__init__.py b/test/generic_python_kernels/__init__.py index 57af35627..6862f41d0 100644 --- a/test/generic_python_kernels/__init__.py +++ b/test/generic_python_kernels/__init__.py @@ -6,5 +6,6 @@ - ``kernel_name``: name of the kernel function or class in the module - ``arguments(c, a, b, n)``: the kernel arguments in the order the kernel expects them - ``tune_params(n)``: a small set of tunable parameters -- ``call_function``: the call function that launches the kernel + +The kernels are launched with the default call functions for their DSL. """ diff --git a/test/generic_python_kernels/cupyx_vector_add.py b/test/generic_python_kernels/cupyx_vector_add.py index d4d1f725e..79466c549 100644 --- a/test/generic_python_kernels/cupyx_vector_add.py +++ b/test/generic_python_kernels/cupyx_vector_add.py @@ -1,5 +1,3 @@ -import cupy as cp -import torch from cupyx import jit kernel_name = "vector_add" @@ -18,8 +16,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256]} - - -def call_function(kernel_function, args, kwargs, grid, threads): - cupy_args = [cp.from_dlpack(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] - kernel_function(grid, threads, tuple(cupy_args)) diff --git a/test/generic_python_kernels/cute_vector_add.py b/test/generic_python_kernels/cute_vector_add.py index 0d5e60cf9..157589f27 100644 --- a/test/generic_python_kernels/cute_vector_add.py +++ b/test/generic_python_kernels/cute_vector_add.py @@ -1,6 +1,4 @@ import cutlass.cute as cute -import torch -from cutlass.cute.runtime import from_dlpack kernel_name = "vector_add" @@ -30,14 +28,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256]} - - -# kernels compiled with cute.compile, indexed by the kernel function of each configuration -compiled_kernels = {} - - -def call_function(kernel_function, args, kwargs): - cute_args = [from_dlpack(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] - if kernel_function not in compiled_kernels: - compiled_kernels[kernel_function] = cute.compile(kernel_function, *cute_args) - compiled_kernels[kernel_function](*cute_args, **kwargs) diff --git a/test/generic_python_kernels/cutile_vector_add.py b/test/generic_python_kernels/cutile_vector_add.py index 7f702ad39..71562d16a 100644 --- a/test/generic_python_kernels/cutile_vector_add.py +++ b/test/generic_python_kernels/cutile_vector_add.py @@ -1,5 +1,4 @@ import cuda.tile as ct -import torch kernel_name = "vector_add" @@ -20,8 +19,3 @@ def arguments(c, a, b, n): def tune_params(n): # each block processes one tile, so the tile size is the block size used to compute the grid return {"block_size_x": [128, 256]} - - -def call_function(kernel_function, args, kwargs, grid): - # ct.launch only takes positional arguments, the tunable kernel argument block_size_x comes last - ct.launch(torch.cuda.current_stream(), grid, kernel_function, (*args, *kwargs.values())) diff --git a/test/generic_python_kernels/numba_vector_add.py b/test/generic_python_kernels/numba_vector_add.py index 57299bada..32aaadcf0 100644 --- a/test/generic_python_kernels/numba_vector_add.py +++ b/test/generic_python_kernels/numba_vector_add.py @@ -1,4 +1,3 @@ -import torch from numba import cuda kernel_name = "vector_add" @@ -17,8 +16,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256]} - - -def call_function(kernel_function, args, kwargs, grid, threads): - numba_args = [cuda.as_cuda_array(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] - kernel_function[grid, threads](*numba_args, **kwargs) diff --git a/test/generic_python_kernels/taichi_vector_add.py b/test/generic_python_kernels/taichi_vector_add.py index 1169e20a8..5761b0928 100644 --- a/test/generic_python_kernels/taichi_vector_add.py +++ b/test/generic_python_kernels/taichi_vector_add.py @@ -28,7 +28,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256]} - - -def call_function(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) diff --git a/test/generic_python_kernels/tilelang_vector_add.py b/test/generic_python_kernels/tilelang_vector_add.py index 8471aafe0..f52675bc3 100644 --- a/test/generic_python_kernels/tilelang_vector_add.py +++ b/test/generic_python_kernels/tilelang_vector_add.py @@ -26,7 +26,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256], "n": [n]} - - -def call_function(kernel_function, args, kwargs): - kernel_function(**kwargs)(*args) diff --git a/test/generic_python_kernels/tilus_vector_add.py b/test/generic_python_kernels/tilus_vector_add.py index 69fd9e3d1..2b8ada442 100644 --- a/test/generic_python_kernels/tilus_vector_add.py +++ b/test/generic_python_kernels/tilus_vector_add.py @@ -31,7 +31,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256], "num_warps": [4]} - - -def call_function(kernel_function, args, kwargs): - kernel_function(*args, **kwargs) diff --git a/test/generic_python_kernels/triton_vector_add.py b/test/generic_python_kernels/triton_vector_add.py index 03a50ef29..015aef23a 100644 --- a/test/generic_python_kernels/triton_vector_add.py +++ b/test/generic_python_kernels/triton_vector_add.py @@ -19,7 +19,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256]} - - -def call_function(kernel_function, args, kwargs, grid): - kernel_function[grid](*args, **kwargs) diff --git a/test/generic_python_kernels/warp_vector_add.py b/test/generic_python_kernels/warp_vector_add.py index e30629a11..df45bd119 100644 --- a/test/generic_python_kernels/warp_vector_add.py +++ b/test/generic_python_kernels/warp_vector_add.py @@ -1,8 +1,5 @@ -import torch import warp as wp -wp.init() - kernel_name = "vector_add" @@ -19,8 +16,3 @@ def arguments(c, a, b, n): def tune_params(n): return {"block_size_x": [128, 256]} - - -def call_function(kernel_function, args, kwargs, grid, threads): - warp_args = [wp.from_torch(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] - wp.launch(kernel_function, dim=grid[0] * threads[0], inputs=warp_args, block_dim=threads[0]) diff --git a/test/test_generic_python_dsls.py b/test/test_generic_python_dsls.py index 51e7d3a01..9a759c0f8 100644 --- a/test/test_generic_python_dsls.py +++ b/test/test_generic_python_dsls.py @@ -50,16 +50,8 @@ def test_generic_python_dsl(dsl): answer = [a + b if arg is c else None for arg in args] tune_params = kernel.tune_params(n) - results, _ = tune_kernel( - kernel.kernel_name, - kernel.__file__, - n, - args, - tune_params, - lang="generic_python", - call_function=kernel.call_function, - answer=answer, - ) + # the language and call function are detected from the kernel source + results, _ = tune_kernel(kernel.kernel_name, kernel.__file__, n, args, tune_params, answer=answer) num_configs = np.prod([len(v) for v in tune_params.values()]) assert len(results) == num_configs diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py index 9749dc6e4..f9b852f61 100644 --- a/test/test_kernel_source_fn.py +++ b/test/test_kernel_source_fn.py @@ -65,7 +65,7 @@ def test_initiation(): ''' Test invalid KernelSourceFn initations ''' - with pytest.raises(ValueError, match=r"call_function must be supplied for language .*"): + with pytest.raises(ValueError, match=r"Could not detect the Python DSL .*, please pass a call_function"): KernelSource("mock_kernel", KS_FILE, "generic_python") with pytest.raises(FileNotFoundError, match=r"Kernel source file .* not found"): diff --git a/test/test_language_detection.py b/test/test_language_detection.py new file mode 100644 index 000000000..f7e3245db --- /dev/null +++ b/test/test_language_detection.py @@ -0,0 +1,120 @@ +"""Tests for detecting the language of kernels and the DSL of Python kernels. + +The DSL of Python kernels is detected from the source code only, so these tests do not require +any of the DSLs to be installed. +""" + +import textwrap + +import pytest + +from kernel_tuner.kernel_sources.kernel_source import KernelSource +from kernel_tuner.language import Language +from kernel_tuner.utils.call_functions import DEFAULT_CALL_FUNCTIONS, call_triton, get_default_call_function +from kernel_tuner.utils.language_detection import detect_language, detect_python_dsl, is_python_file + + +def write_source(tmp_path, source, name="kernel.py"): + path = tmp_path / name + path.write_text(textwrap.dedent(source)) + return path + + +def test_detect_language1(): + kernel_string = "__global__ void vector_add( ... );" + lang = detect_language(kernel_string) + assert lang == "CUDA" + + +def test_detect_language2(): + kernel_string = "__kernel void vector_add( ... );" + lang = detect_language(kernel_string) + assert lang == "OpenCL" + + +def test_detect_language3(): + kernel_string = "blabla" + lang = detect_language(kernel_string) + assert lang == "C" + + +def test_is_python_file(tmp_path): + assert is_python_file(write_source(tmp_path, "x = 1")) + assert is_python_file(str(write_source(tmp_path, "x = 1"))) + assert not is_python_file(write_source(tmp_path, "__global__ void k() {}", name="kernel.cu")) + assert not is_python_file(str(tmp_path / "does_not_exist.py")) + assert not is_python_file("__global__ void k() {}") + assert not is_python_file("x" * 10000) # too long to be a path + assert not is_python_file(lambda params: "") + + +dsl_sources = [ + ("import triton\n@triton.jit\ndef k(): pass", "triton"), + ("import triton\n@triton.autotune(configs=[])\n@triton.jit\ndef k(): pass", "triton"), + ("from numba import cuda\n@cuda.jit\ndef k(): pass", "numba"), + ("import numba.cuda\n@numba.cuda.jit\ndef k(): pass", "numba"), + ("from numba import jit\n@jit\ndef k(): pass", None), # Numba for CPUs is not supported + ("from cupyx import jit\n@jit.rawkernel()\ndef k(): pass", "cupyx"), + ("import warp as wp\n@wp.kernel()\ndef k(): pass", "warp"), + ("import taichi as ti\n@ti.kernel\ndef k(): pass", "taichi"), + ("import cutlass.cute as cute\n@cute.jit\ndef k(): pass", "cute"), + ("import cutlass\n@cutlass.cute.kernel\ndef k(): pass", "cute"), + ("import cutlass.cute as cute\nclass k:\n @cute.jit\n def __call__(self): pass", "cute"), + ("import tilus\nclass k(tilus.Script):\n def __call__(self): pass", "tilus"), + ("from tilus import Script\nclass k(Script):\n def __call__(self): pass", "tilus"), + ("import tilelang\n@tilelang.jit\ndef k(): pass", "tilelang"), + ("import cuda.tile as ct\n@ct.kernel\ndef k(): pass", "cutile"), + ("from cuda import tile\n@tile.kernel\ndef k(): pass", "cutile"), + ("import cuda.tile\n@cuda.tile.kernel\ndef k(): pass", "cutile"), + ("def k(): pass", None), + ("import functools\n@functools.cache\ndef k(): pass", None), + ("import triton\ndef helper(): pass\n@triton.jit\ndef other(): pass\ndef k(): pass", None), +] + + +@pytest.mark.parametrize("source, dsl", dsl_sources, ids=[s for s, _ in dsl_sources]) +def test_detect_python_dsl(tmp_path, source, dsl): + assert detect_python_dsl("k", write_source(tmp_path, source)) == dsl + + +def test_every_dsl_has_a_default_call_function(): + detected_dsls = {dsl for _, dsl in dsl_sources if dsl is not None} + assert detected_dsls == set(DEFAULT_CALL_FUNCTIONS) + with pytest.raises(ValueError): + get_default_call_function("unknown_dsl") + + +def test_kernel_source_detects_generic_python_and_call_function(tmp_path): + path = write_source(tmp_path, "import triton\n@triton.jit\ndef k(): pass") + ks = KernelSource("k", str(path), None) + assert ks.lang == Language.GENERIC_PYTHON + # call_triton accepts all optional arguments, so normalizing it returns the function itself + assert ks.call_function is call_triton + + +def test_kernel_source_user_call_function_takes_precedence(tmp_path): + path = write_source(tmp_path, "import triton\n@triton.jit\ndef k(): pass") + + def my_call_function(kernel_function, args, kwargs, grid, threads, params): + pass + + ks = KernelSource("k", path, None, call_function=my_call_function) + assert ks.lang == Language.GENERIC_PYTHON + assert ks.call_function is my_call_function + + # a call function also makes undecorated kernels usable + path = write_source(tmp_path, "def k(): pass", name="plain.py") + ks = KernelSource("k", path, "generic_python", call_function=my_call_function) + assert ks.call_function is my_call_function + + +def test_kernel_source_unknown_dsl_requires_call_function(tmp_path): + path = write_source(tmp_path, "def k(): pass") + with pytest.raises(ValueError, match="Could not detect the Python DSL"): + KernelSource("k", path, None) + + +def test_kernel_source_non_python_file_is_not_generic_python(tmp_path): + path = write_source(tmp_path, "__global__ void k(float *a) {}", name="kernel.cu") + ks = KernelSource("k", str(path), None) + assert ks.lang == Language.CUDA diff --git a/test/test_lazy_imports.py b/test/test_lazy_imports.py index 7faf2c716..4f7f7fb17 100644 --- a/test/test_lazy_imports.py +++ b/test/test_lazy_imports.py @@ -45,3 +45,22 @@ def test_backend_modules_not_loaded_on_import(): """ loaded = json.loads(_run_isolated(code)) assert loaded == [] + + +def test_python_dsls_not_imported_by_detection(tmp_path): + kernel_file = tmp_path / "kernel.py" + kernel_file.write_text( + "import triton\nimport warp\nimport taichi\nimport tilus\nimport tilelang\nimport cuda.tile as ct\n" + "import cutlass.cute\nfrom numba import cuda\nfrom cupyx import jit\n\n@ct.kernel\ndef k(a): pass\n" + ) + code = f""" +import json, sys +from kernel_tuner.kernel_sources.kernel_source import KernelSource +from kernel_tuner.utils.language_detection import detect_python_dsl +dsl = detect_python_dsl("k", {str(kernel_file)!r}) +ks = KernelSource("k", {str(kernel_file)!r}, None) +dsls = ["triton", "numba", "cupy", "cupyx", "warp", "taichi", "cutlass", "tilus", "tilelang", "cuda.tile"] +print(json.dumps({{"dsl": dsl, "lang": str(ks.lang), "loaded": [m for m in dsls if m in sys.modules]}})) +""" + result = json.loads(_run_isolated(code)) + assert result == {"dsl": "cutile", "lang": "GENERIC_PYTHON", "loaded": []} diff --git a/test/test_util_functions.py b/test/test_util_functions.py index 2dc0ff16b..aed8e97f4 100644 --- a/test/test_util_functions.py +++ b/test/test_util_functions.py @@ -240,25 +240,6 @@ def test_check_restrictions(): assert answer == e -def test_detect_language1(): - kernel_string = "__global__ void vector_add( ... );" - lang = detect_language(kernel_string) - assert lang == "CUDA" - - -def test_detect_language2(): - kernel_string = "__kernel void vector_add( ... );" - lang = detect_language(kernel_string) - assert lang == "OpenCL" - - -def test_detect_language3(): - kernel_string = "blabla" - lang = detect_language(kernel_string) - assert lang == "C" - - -@skip_if_no_pycuda def test_get_device_interface1(): lang = "PYCUDA" dev = core.DeviceInterface(core.KernelSource("", "", lang=lang)) From fc0328643ecec200b1b4c70142bb00ae2172ff3c Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 14:16:05 +0200 Subject: [PATCH 28/44] document Python backend and Python DSL examples - add Python backend section to backends.rst, and correct which CUDA backend is the default - add README.rst for the Python DSL examples and link the DSL examples from examples/README.rst - add the language_detection and call_functions modules to design.rst --- doc/source/backends.rst | 112 +++++++++++++++++++++++++++++- doc/source/design.rst | 10 +++ examples/README.rst | 8 ++- examples/python/README.rst | 136 +++++++++++++++++++++++++++++++++++++ 4 files changed, 261 insertions(+), 5 deletions(-) create mode 100644 examples/python/README.rst diff --git a/doc/source/backends.rst b/doc/source/backends.rst index d132a9014..80953c00e 100644 --- a/doc/source/backends.rst +++ b/doc/source/backends.rst @@ -5,8 +5,8 @@ Backends ======== -Kernel Tuner implements multiple backends for CUDA, one for OpenCL, one for HIP, and a generic -Compiler backend. +Kernel Tuner implements multiple backends for CUDA, one for OpenCL, one for HIP, a generic +Compiler backend, and a Python backend for kernels written in Python-based DSLs. Selecting a backend is in most cases automatic and is done based on the kernel's programming language, but sometimes you'll want to specifically choose a backend. @@ -15,7 +15,8 @@ language, but sometimes you'll want to specifically choose a backend. CUDA Backends ------------- -PyCUDA is default CUDA backend in Kernel Tuner. It is comparable in feature completeness with CuPy. +For ``lang="CUDA"``, Kernel Tuner uses the CUDA-Python backend if cuda-python is installed, otherwise +CuPy, and otherwise PyCUDA. PyCUDA is comparable in feature completeness with CuPy. Because the HIP kernel language is identical to the CUDA kernel language, HIP is included here as well. To use HIP on nvidia GPUs, see https://github.com/jatinx/hip-on-nv. @@ -63,3 +64,108 @@ used to compile the kernels. Compiler used, "nvcc", "nvrtc", "nvrtc", "hiprtc" +Python Backend +-------------- + +The Python backend tunes kernels written in Python-based embedded domain-specific languages (DSLs). +Kernel Tuner supports the DSLs listed in the table below, examples for each DSL can be found in +`examples/python `__. + +To tune a Python kernel, pass the name of the kernel function (or class) as ``kernel_name`` and the path +to the Python file that contains the kernel as ``kernel_source``. Kernel Tuner selects the Python backend +for kernels in Python files automatically, or you can select it with ``lang="generic_python"``. +The Python backend requires PyTorch with CUDA support and currently only supports NVIDIA GPUs. +The DSLs are optional dependencies, which are only imported when they are used. + +.. code-block:: python + + import triton + import triton.language as tl + + @triton.jit + def vector_add(c_ptr, a_ptr, b_ptr, n, block_size_x: tl.constexpr): + offsets = tl.program_id(axis=0) * block_size_x + tl.arange(0, block_size_x) + mask = offsets < n + a = tl.load(a_ptr + offsets, mask=mask) + b = tl.load(b_ptr + offsets, mask=mask) + tl.store(c_ptr + offsets, a + b, mask=mask) + + ... + + results, env = tune_kernel("vector_add", __file__, n, [c, a, b, n], {"block_size_x": [128, 256, 512]}) + + +Call functions +~~~~~~~~~~~~~~ + +Every DSL has its own way to launch a kernel. Kernel Tuner uses a *call function* to launch the kernel +for each configuration. When ``call_function`` is not passed to ``tune_kernel`` or ``run_kernel``, +Kernel Tuner detects the DSL of the kernel from its decorators or base classes, and uses the default call +function for that DSL from :mod:`kernel_tuner.utils.call_functions`: + +.. csv-table:: Supported Python DSLs + :header: DSL, Python package, Detected by, Default launch + :widths: auto + + Triton, "triton", "``@triton.jit``", "``kernel[grid]``, ``num_warps`` and ``num_stages`` are taken from the tunable parameters" + Numba, "numba-cuda", "``@numba.cuda.jit``", "``kernel[grid, threads]``" + CuPy, "cupy", "``@cupyx.jit.rawkernel``", "``kernel(grid, threads, args)``" + Warp, "warp-lang", "``@warp.kernel``", "``wp.launch``, tunable parameters ``dim`` and ``block_dim`` set the launch dimensions" + Taichi, "taichi", "``@taichi.kernel``", "``kernel(*args)``" + CuTe, "nvidia-cutlass-dsl", "``@cutlass.cute.jit``", "``cute.compile`` once per configuration" + Tilus, "tilus", "subclass of ``tilus.Script``", "``kernel(*args)``" + TileLang, "tilelang", "``@tilelang.jit``", "tunable parameters are passed to the kernel factory" + cuTile, "cuda-tile", "``@cuda.tile.kernel``", "``ct.launch`` with a compile timeout of 60 seconds" + +If a kernel needs to be launched differently, you can pass your own call function. The call function +receives the following positional arguments: + +- ``kernel_function``: the kernel with the values of the tunable parameters inserted +- ``args``: the list of kernel arguments, as prepared by Kernel Tuner +- ``kwargs``: a dictionary with the tunable parameters that are also arguments of the kernel + +And optionally, in this order, the arguments ``grid`` and ``threads`` with the grid and thread block +dimensions computed by Kernel Tuner, and ``params``, a dictionary with the tunable parameters of the +configuration. For example: + +.. code-block:: python + + def call_function(kernel_function, args, kwargs, grid): + kernel_function[grid](*args, **kwargs) + + +Tunable parameters +~~~~~~~~~~~~~~~~~~ + +Kernel Tuner inserts the values of the tunable parameters into the kernel source code of each +configuration. Tunable parameters that are arguments of the kernel are passed to the call function in +``kwargs``. Other occurrences of tunable parameters in the kernel, and in functions defined in the same +file that are called by the kernel, are replaced by their values: reading a variable (or class attribute) +with the name of a tunable parameter returns its value, and assignments to such variables are +overridden with the value. This makes it possible to tune values that are defined inside a kernel, as well +as attributes of class-based kernels. + +Only the imports, the kernel, and the functions called by the kernel are included in the source code of +each configuration. Other module-level statements and variables in the kernel source file are not +available to the kernel, unless they are tunable parameters. + +Kernel Tuner computes the grid dimensions from the problem size and the block size parameters, in the +same way as for CUDA kernels. For tile-based DSLs, use ``block_size_names`` to indicate the tunable +parameters that determine the tile size, so that the grid has one block per tile. The maximum number of +threads per block is not checked by the Python backend, as tile-based DSLs determine the number of +threads themselves. + + +Arguments and errors +~~~~~~~~~~~~~~~~~~~~ + +Kernel arguments can be NumPy arrays, PyTorch tensors, or scalars. Kernel Tuner copies NumPy arrays and +PyTorch tensors to separate PyTorch tensors on the GPU, so the arguments passed by the user are not +modified, and resets outputs between kernel runs. The kernels are benchmarked using CUDA events on the +current PyTorch stream. Compiling a kernel is done by launching it once, so the compile time reported by +Kernel Tuner excludes the time of that kernel run. + +Configurations that cannot be compiled or launched because they use too many resources, are too large +to compile within the time limit, or use tile sizes that the DSL does not support, are skipped. Other +errors, such as errors in the kernel code, stop the tuning process. Constant, shared, and texture memory +arguments are not supported by the Python backend. diff --git a/doc/source/design.rst b/doc/source/design.rst index 8d8b9b357..21c48f3b3 100644 --- a/doc/source/design.rst +++ b/doc/source/design.rst @@ -150,3 +150,13 @@ kernel_tuner.util .. automodule:: kernel_tuner.util :members: + +kernel_tuner.utils.language_detection +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: kernel_tuner.utils.language_detection + :members: + +kernel_tuner.utils.call_functions +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. automodule:: kernel_tuner.utils.call_functions + :members: diff --git a/examples/README.rst b/examples/README.rst index e2749882d..a26fbc3df 100644 --- a/examples/README.rst +++ b/examples/README.rst @@ -4,6 +4,10 @@ Kernel Tuner Examples Most of the examples show how to use Kernel Tuner to tune a CUDA, OpenCL, or C kernel, while demonstrating a particular usecase of Kernel Tuner. +The examples in `examples/python `__ show how to tune kernels written in +Python-based DSLs, such as Triton, Numba, CuPy (``cupyx.jit``), Warp, Taichi, CuTe, Tilus, TileLang, and cuTile. +These are described in more detail in the `README of the Python DSL examples `__. + Except for `test\_vector\_add.py `__ and `test\_vector\_add_parameterized.py `__, which show how to write tests for GPU kernels with Kernel Tuner. @@ -18,7 +22,7 @@ Below we list the example applications and the features they illustrate. Vector Add ---------- -[`CUDA `__] [`CUDA-C++ `__] [`OpenCL `__] [`C `__] [`Fortran `__] [`OpenACC-C++ `__] [`OpenACC-Fortran `__] +[`CUDA `__] [`CUDA-C++ `__] [`OpenCL `__] [`C `__] [`Fortran `__] [`OpenACC-C++ `__] [`OpenACC-Fortran `__] [`Numba `__] [`Triton `__] [`Warp `__] [`CuTe `__] [`Tilus `__] [`TileLang `__] [`cuTile `__] - use Kernel Tuner to tune a simple kernel Stencil @@ -28,7 +32,7 @@ Stencil Matrix Multiplication --------------------- -[`CUDA `__] [`OpenCL `__] +[`CUDA `__] [`OpenCL `__] [`CuPy `__] [`Numba `__] [`Taichi `__] [`Warp `__] [`Triton `__] [`Tilus `__] [`TileLang `__] [`CuTe `__] [`cuTile `__] - pass a filename instead of a string with code - use 2-dimensional thread blocks and tiling in both dimensions - tell Kernel Tuner to compute the grid dimensions for 2D thread blocks with tiling diff --git a/examples/python/README.rst b/examples/python/README.rst new file mode 100644 index 000000000..5ba463f6b --- /dev/null +++ b/examples/python/README.rst @@ -0,0 +1,136 @@ +Kernel Tuner Examples for Python DSLs +===================================== + +These examples show how to use Kernel Tuner to tune kernels written in Python-based embedded +domain-specific languages (DSLs), such as Triton, Numba, CuPy (``cupyx.jit``), Warp, Taichi, CuTe, +Tilus, TileLang, and cuTile. Kernel Tuner tunes these kernels using its Python backend, see the +`documentation on backends `__ +for more information. + +To run an example, install PyTorch with CUDA support and the DSL used by the example, for example +``pip install triton``, ``pip install numba-cuda``, ``pip install cupy-cuda13x``, ``pip install warp-lang``, +``pip install taichi``, ``pip install nvidia-cutlass-dsl``, ``pip install tilus``, ``pip install tilelang``, +or ``pip install cuda-tile``. Kernel Tuner only imports a DSL when it is used. + +In all examples, the kernel is passed to ``tune_kernel`` as the name of the kernel function (or +class) and the path to the Python file that contains it. Kernel Tuner detects that the kernel is +written in Python and which DSL it uses, and launches the kernel with a default call function for +that DSL. There is no need to pass the ``lang`` or ``call_function`` options, unless you want to +launch the kernel in a different way, as shown in the Warp vector add and CuTe matrix multiplication +examples. + +.. note:: + + Please do not use the examples as performance benchmarks. + The examples here are created specifically to highlight certain features in Kernel Tuner. + Please contact the developers if you are interested in benchmarking Kernel Tuner. + +Below we list the example applications and the features they illustrate. + +Vector Add +---------- + +Numba +~~~~~ +[`numba_vec_add.py `__] + - tune the thread block size of a Numba CUDA kernel just like a CUDA kernel + - pass NumPy arrays as kernel arguments, which Kernel Tuner copies to the GPU + +Triton +~~~~~~ +[`triton_vec_add.py `__] + - tune a tunable parameter that is also a kernel argument (a ``tl.constexpr``) + - pass PyTorch tensors that already live on the GPU as kernel arguments + - use ``run_kernel`` to run a single configuration of a Python kernel + +Tilus +~~~~~ +[`tilus_vec_add.py `__] + - tune a class-based kernel, where the tunable parameters are attributes of the class + - tune the number of warps per thread block + +TileLang +~~~~~~~~ +[`tilelang_vec_add.py `__] + - tune the arguments of a kernel factory function decorated with ``@tilelang.jit`` + - pass a fixed value to the kernel factory using a tunable parameter with a single value + +CuTe +~~~~ +[`cute_vec_add.py `__] + - tune a value that is assigned inside the ``@cute.jit`` function that launches the kernel + +cuTile +~~~~~~ +[`cutile_vec_add.py `__] + - tune the tile size of a tile-based kernel + - use the tile size as block size, so Kernel Tuner computes a grid with one block per tile + +Warp +~~~~ +[`warp_vec_add.py `__] + - tune the amount of work per thread, a value that is assigned inside the kernel + - use a helper function (``@wp.func``) that is called by the kernel + - pass a user-defined call function to control how the kernel is launched + +Matrix Multiplication +--------------------- + +CuPy +~~~~ +[`cupy_matmul.py `__] + - tune a ``cupyx.jit.rawkernel`` with 2-dimensional thread blocks + - use the restrictions option to limit the search to valid thread block sizes + +Numba +~~~~~ +[`numba_matmul.py `__] + - tune a basic kernel and a kernel that uses shared memory tiling and register blocking + - tell Kernel Tuner to compute the grid dimensions for 2D thread blocks with tiling + - use restrictions that relate multiple tunable parameters to each other + - use Bayesian Optimization to search a large search space + +Taichi +~~~~~~ +[`taichi_matmul.py `__] + - tune the block size that Taichi uses to parallelize a loop (``ti.loop_config``) + +Warp +~~~~ +[`warp_matmul.py `__] + - use the special tunable parameters ``block_dim`` and ``dim`` to set the thread block size and + launch dimensions of a Warp kernel + +Triton +~~~~~~ +[`triton_matmul.py `__] + - tune a basic kernel and the optimized kernel from the Triton tutorials + - tune the Triton compiler options ``num_warps`` and ``num_stages`` + - use Bayesian Optimization to search a large search space + +Tilus +~~~~~ +[`tilus_matmul.py `__] + - tune the tile sizes and number of warps of class-based kernels + - use Bayesian Optimization to search a large search space + +TileLang +~~~~~~~~ +[`tilelang_matmul.py `__] + - tune the tile sizes, number of threads, and number of software pipelining stages + - use Bayesian Optimization to search a large search space + +CuTe +~~~~ +[`cute_matmul.py `__] + - tune a naive kernel and the Ampere tensor core GEMM from the CuTe DSL examples + - tune a class-based kernel, where the tunable parameters are attributes of the class + - use restrictions to exclude configurations that exceed the shared memory or that the kernel does not support + - pass a user-defined call function that creates the tensor layouts required by the kernel + +cuTile +~~~~~~ +[`cutile_matmul.py `__] + - tune the tile sizes of a tile-based kernel that uses tensor cores + - tune the size of groups of tiles to improve L2 cache reuse, using a 1-dimensional grid + - tell Kernel Tuner to compute a grid with one block per tile of the output matrix From cfbb70c8dfe0f7f72d529792d06055701568bd34 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 22:38:44 +0200 Subject: [PATCH 29/44] add parallel_compile option to compile kernels in parallel tune_kernel(parallel_compile=True|n) uses the new ParallelCompileRunner, which compiles the configurations the strategy evaluates at once in parallel threads, and then verifies and benchmarks them one after the other, so compilation does not affect the benchmarks. - split compilation into Backend.build(), which is thread-safe, and Backend.load(), implemented by the nvcuda, pycuda, cupy, and generic_python backends - add DeviceInterface.build_kernel() and let compile_and_benchmark() take a prebuilt kernel instance and build - compile Numba, Warp, cuTile, and CuTe kernels in worker processes, as these cannot compile in parallel threads; the main process loads the compiled kernels from the DSL disk caches, or from kernels exported with export_to_c for CuTe, at most 8 processes by default - raise errors from worker processes with their exception chain, so they are classified as before, and fall back to compiling in the main process when kernels cannot be compiled in worker processes - make normalized call functions picklable - name temporary kernel modules after their file, so worker processes import them under the same name - document parallel compilation in parallel.rst --- CHANGELOG.md | 2 + doc/source/design.rst | 6 + doc/source/parallel.rst | 50 +++ kernel_tuner/backends/backend.py | 17 + kernel_tuner/backends/cupy.py | 42 ++- kernel_tuner/backends/generic_python.py | 135 ++++++-- .../backends/generic_python_processes.py | 325 ++++++++++++++++++ kernel_tuner/backends/nvcuda.py | 70 ++-- kernel_tuner/backends/pycuda.py | 71 +++- kernel_tuner/core.py | 57 ++- kernel_tuner/interface.py | 45 ++- .../kernel_sources/kernel_source_fn.py | 26 +- kernel_tuner/runners/parallel_compile.py | 123 +++++++ kernel_tuner/runners/sequential.py | 33 +- kernel_tuner/util.py | 25 +- test/test_generic_python_dsls.py | 89 ++++- test/test_generic_python_functions.py | 2 +- test/test_parallel_compile.py | 149 ++++++++ test/test_util_functions.py | 12 + 19 files changed, 1158 insertions(+), 121 deletions(-) create mode 100644 kernel_tuner/backends/generic_python_processes.py create mode 100644 kernel_tuner/runners/parallel_compile.py create mode 100644 test/test_parallel_compile.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d24e52063..fd3850b79 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ This project adheres to [Semantic Versioning](http://semver.org/). - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers - Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, cuTile, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` - Detect kernels in Python files as `lang="generic_python"`, and detect the DSL of Python kernels to select a default `call_function` when none is given +- Add `parallel_compile` option to compile kernels in parallel threads before benchmarking them one after the other, supported by the CUDA backends and all Python DSLs. Numba, Warp, cuTile, and CuTe kernels are compiled in worker processes instead of threads +- Fix PyCUDA backend raising an error instead of skipping configurations that use too much shared memory - Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend - Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel - Fix output verification failing for PyTorch tensors as expected answers diff --git a/doc/source/design.rst b/doc/source/design.rst index 21c48f3b3..ab43de72a 100644 --- a/doc/source/design.rst +++ b/doc/source/design.rst @@ -101,6 +101,12 @@ kernel_tuner.runners.sequential.SimulationRunner :special-members: __init__ :members: +kernel_tuner.runners.parallel_compile.ParallelCompileRunner +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +.. autoclass:: kernel_tuner.runners.parallel_compile.ParallelCompileRunner + :special-members: __init__ + :members: + Backends -------- diff --git a/doc/source/parallel.rst b/doc/source/parallel.rst index 0fdd557a4..2f7daff3f 100644 --- a/doc/source/parallel.rst +++ b/doc/source/parallel.rst @@ -148,3 +148,53 @@ Finally, launch your application using: .. code-block:: RAY_ADDRESS=145.184.221.164:6379 python my_tuning_script.py + + +Parallel Compilation +-------------------- + +Compiling kernels can take a significant part of the tuning time, in particular for kernels written in +Python-based DSLs, where compiling a single configuration can take seconds to minutes. Independently of +parallel tuning with Ray, Kernel Tuner can compile kernels in parallel threads on the local machine, by passing +the ``parallel_compile`` argument to ``tune_kernel``: + +.. code-block:: python + + kernel_tuner.tune_kernel( + "vector_add", + kernel_string, + size, + args, + tune_params, + parallel_compile=True, + ) + +If ``parallel_compile`` is set to ``True``, Kernel Tuner uses as many threads as there are CPU cores, or it can be +set to an integer ``n`` to use exactly ``n`` threads. When the optimization strategy evaluates multiple +configurations at once, Kernel Tuner first compiles all these configurations in parallel, and then verifies and +benchmarks them one after the other. Kernels are never benchmarked while other kernels are being compiled, so the +compilation does not affect the benchmark results. As with parallel tuning, strategies with maximum or limited +parallelism benefit from parallel compilation, while strategies without parallelism compile one configuration at a +time. + +Parallel compilation is supported by the CUDA backends (CUDA-Python, CuPy, and PyCUDA) and by the Python backend +for all supported DSLs. Other backends compile the kernels one at a time. +``parallel_compile`` cannot be combined with ``parallel`` or ``simulation_mode``. + +Numba, Warp, cuTile, and CuTe cannot compile kernels in parallel threads: the first three compile one kernel at a +time because of a global compiler lock, and CuTe kernels can only be launched from the thread that compiled them. +These kernels are compiled in worker processes instead, at most one per thread, and at most 8 when +``parallel_compile`` is set to ``True``. A worker process compiles the +kernel by running it once, after which Kernel Tuner loads the compiled kernel from disk. Numba and CuTe kernels are +stored in a temporary cache that Kernel Tuner removes, Warp and cuTile kernels are stored in the kernel caches of +these DSLs. Note that: + +* Starting the worker processes takes several seconds, so compiling in worker processes only pays off when + compiling a kernel takes long, or when many kernels are compiled. +* Each worker process creates its own CUDA context on the GPU and a copy of the kernel arguments, which uses GPU + memory. Use a smaller number for ``parallel_compile`` if the GPU runs out of memory. +* The worker processes are started with the ``spawn`` method of ``multiprocessing``, which imports the script that + calls ``tune_kernel`` in each worker process. Protect the code that starts tuning with + ``if __name__ == "__main__":``, as the examples do. A ``call_function`` must be defined at the top level of a + module, so that it can be sent to the worker processes. When kernels cannot be compiled in worker processes, + Kernel Tuner prints a warning and compiles them one at a time. diff --git a/kernel_tuner/backends/backend.py b/kernel_tuner/backends/backend.py index 4b4e4aaba..65f483499 100644 --- a/kernel_tuner/backends/backend.py +++ b/kernel_tuner/backends/backend.py @@ -46,6 +46,23 @@ def compile(self, kernel_instance): """This method must implement the compilation of a kernel into a callable function.""" pass + def build(self, kernel_instance): + """Compile a kernel without loading it onto the device, the result is passed to load(). + + Kernel Tuner can call build() for multiple kernels in parallel threads, so implementations must be + thread-safe and must not modify the state of the backend. Backends that do not implement build() + compile the kernel in load() instead. + """ + return None + + def load(self, kernel_instance, build_result): + """Load a kernel compiled by build() onto the device and return a callable function. + + Called from the main thread, right before the kernel is verified and benchmarked. By default, this + compiles the kernel, for backends that do not implement build(). + """ + return self.compile(kernel_instance) + @abstractmethod def start_event(self): """This method must implement the recording of the start of a measurement.""" diff --git a/kernel_tuner/backends/cupy.py b/kernel_tuner/backends/cupy.py index 91554c0e4..1482572e0 100644 --- a/kernel_tuner/backends/cupy.py +++ b/kernel_tuner/backends/cupy.py @@ -115,16 +115,23 @@ def ready_argument_list(self, arguments): def compile(self, kernel_instance): """Call the CUDA compiler to compile the kernel, return the device function. - :param kernel_name: The name of the kernel to be compiled, used to lookup the - function after compilation. - :type kernel_name: string - - :param kernel_string: The CUDA kernel code that contains the function `kernel_name` - :type kernel_string: string + :param kernel_instance: The kernel instance, containing the name and code of the kernel + :type kernel_instance: kernel_tuner.core.KernelInstance :returns: An CUDA kernel that can be called directly. :rtype: cupy.RawKernel """ + return self.load(kernel_instance, self.build(kernel_instance)) + + def build(self, kernel_instance): + """Compile the kernel into a module using NVRTC, this is thread-safe. + + :param kernel_instance: The kernel instance, containing the name and code of the kernel + :type kernel_instance: kernel_tuner.core.KernelInstance + + :returns: The compiled module + :rtype: cupy.RawModule + """ kernel_string = kernel_instance.kernel_string kernel_name = kernel_instance.name @@ -135,11 +142,26 @@ def compile(self, kernel_instance): options = tuple(compiler_options) - self.current_module = cp.RawModule( - code=kernel_string, options=options, name_expressions=[kernel_name] - ) + module = cp.RawModule(code=kernel_string, options=options, name_expressions=[kernel_name]) + # CuPy compiles lazily, compile now on the device of this backend (the current device is per thread) + with self.dev: + module.compile() + return module + + def load(self, kernel_instance, build_result): + """Return the device function from a module compiled by build(). - self.func = self.current_module.get_function(kernel_name) + :param kernel_instance: The kernel instance + :type kernel_instance: kernel_tuner.core.KernelInstance + + :param build_result: The module returned by build() + :type build_result: cupy.RawModule + + :returns: An CUDA kernel that can be called directly. + :rtype: cupy.RawKernel + """ + self.current_module = build_result + self.func = self.current_module.get_function(kernel_instance.name) self.num_regs = self.func.num_regs return self.func diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 5acaa8e7a..dd99f03a8 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -1,10 +1,17 @@ import logging -import inspect import re import builtins +import threading import numpy as np from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.backends.generic_python_processes import ( + PROCESS_BUILD_DSLS, + BuildProcessPool, + dsl_compile_hooks, + instantiate_kernel, +) +from kernel_tuner.kernel_sources.kernel_source_fn import kernel_module_name from kernel_tuner.observers.generic_python import GenericPythonRuntimeObserver try: @@ -13,6 +20,16 @@ torch = None +# DSLs whose kernels are not compiled in parallel threads. Numba, Warp, and cuTile compile one kernel at a +# time because of a global compiler lock, so parallel threads are not faster (and even slower for Warp). +# CuTe keeps the state of compiled kernels per thread, so these cannot be launched from another thread. +# When build_processes is set, these are compiled in worker processes instead, see generic_python_processes. +SERIAL_COMPILE_DSLS = PROCESS_BUILD_DSLS + +# Returned by build() for kernels that were compiled by a worker process, load() loads these from disk +PROCESS_BUILT = "process_built" + + class GenericPythonFunctions(GPUBackend): """Class that groups the Python DSL functions on maintains state about the device.""" @@ -59,10 +76,16 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None self.units = {"time": "ms", "power": "s,mW", "energy": "J"} - # Variables to be filled in at compile time, needed for running the kernel after compilation. + # Set when a kernel is loaded, needed for running the kernel after compilation. self.call_function = None self.signature = None - self.gpu_kwargs = None + + # The maximum number of worker processes used by build() for DSLs that cannot compile kernels in + # parallel threads. Set by runners that compile in parallel, by default these kernels are compiled + # by load() in the main process. + self.build_processes = 0 + self._build_pool = None + self._build_pool_lock = threading.Lock() super().__init__(device=device, iterations=iterations, compiler_options=compiler_options, observers=observers) @@ -107,7 +130,7 @@ def ready_argument_list(self, arguments): def compile(self, kernel_instance, gpu_args=None): - """Compile the kernel by executing it once. This enforces that the kernel is cached. + """Compile the kernel by executing it once. This enforces that the kernel is cached. :param kernel_instance: The kernel instance containing information such as the kernel_source, grid, threads, params, etc. @@ -119,45 +142,99 @@ def compile(self, kernel_instance, gpu_args=None): :returns: A kernel that can be called with the user-defined call function. :rtype: callable """ + return self.load(kernel_instance, self.build(kernel_instance, gpu_args), gpu_args) + + def build(self, kernel_instance, gpu_args=None): + """Compile the kernel by executing it once, this can be called from multiple threads. + + The DSL caches the compiled kernel, so the kernel is not compiled again when it is benchmarked. + DSLs that compile kernels one at a time, even from multiple threads, are compiled in a worker + process if build_processes is set, otherwise nothing is done and the kernel is compiled by load(). + + :returns: The kernel function, PROCESS_BUILT if the kernel was compiled by a worker process, + or None if the kernel is compiled by load(). + """ + if kernel_instance.kernel_source.dsl in SERIAL_COMPILE_DSLS: + if not self.build_processes: + return None + kernel_source = kernel_instance.kernel_source + kwargs = self._kernel_kwargs(kernel_source.signature, kernel_instance.params) + pool = self._get_build_pool(gpu_args) + if pool.build(kernel_instance, gpu_args, kernel_source.call_function, kwargs): + return PROCESS_BUILT + return None + # the current device is set per thread, and would otherwise be the default device + torch.cuda.set_device(self.device_id) + return self._compile_by_running(kernel_instance, gpu_args) + + def _get_build_pool(self, gpu_args): + with self._build_pool_lock: + if self._build_pool is None: + self._build_pool = BuildProcessPool(self.device_id, gpu_args, self.build_processes) + return self._build_pool + + def shutdown_build_processes(self): + """Stop the worker processes started by build(), if any.""" + with self._build_pool_lock: + pool, self._build_pool = self._build_pool, None + if pool is not None: + pool.shutdown() + + def load(self, kernel_instance, build_result, gpu_args=None): + """Return the kernel function compiled by build(), or compile it now if build() did not. + + :returns: A kernel that can be called with the user-defined call function. + :rtype: callable + """ + # These are the same for every configuration, but are needed to run the kernel + self.call_function = kernel_instance.kernel_source.call_function + self.signature = kernel_instance.kernel_source.signature + + if build_result is None: + return self._compile_by_running(kernel_instance, gpu_args) + if build_result is PROCESS_BUILT: + return self._compile_by_running(kernel_instance, gpu_args, load_from_worker=True) + return build_result + + def _compile_by_running(self, kernel_instance, gpu_args, load_from_worker=False): + """Compile the kernel by running it once with the user-defined call function. + + :param load_from_worker: Load the kernel that a worker process compiled and stored on disk. + :type load_from_worker: bool + """ if kernel_instance.kernel_fn is None: raise ValueError("kernel_fn is None, currently Generic Python only supports callable kernel_source") if gpu_args is None: raise ValueError("gpu_args is None, Generic Python needs gpu args to compile the kernel") - # The first time we compile, we also set the call function and the signature - # We need this later to run the kernel. - if self.call_function is None or self.signature is None: - self.call_function = kernel_instance.kernel_source.call_function - self.signature = kernel_instance.kernel_source.signature - - grid = kernel_instance.grid - threads = kernel_instance.threads - params = kernel_instance.params - # If the kernel source is a class, we use the __call__ function as the kernel_function. - if inspect.isclass(kernel_instance.kernel_fn): - kernel_function = kernel_instance.kernel_fn() - elif callable(kernel_instance.kernel_fn): - kernel_function = kernel_instance.kernel_fn - else: - raise TypeError("kernel function is not a class or function") + kernel_function = instantiate_kernel(kernel_instance.kernel_fn) - # Tuning params can contain kernel arguments. In such cases, create keyword arguments with - # the values of the tuning params. - self.gpu_kwargs = {} - if params is not None: - for arg_name in self.signature: - if arg_name in params: - self.gpu_kwargs[arg_name] = params[arg_name] + kernel_source = kernel_instance.kernel_source + params = kernel_instance.params + kwargs = self._kernel_kwargs(kernel_source.signature, params) # Call the user-defined call function in order to compile the kernel. self.synchronize() - self.call_function(kernel_function, gpu_args, self.gpu_kwargs, grid, threads, params) + grid, threads = kernel_instance.grid, kernel_instance.threads + if load_from_worker: + module_name = kernel_module_name(kernel_instance.temp_files[0]) + with dsl_compile_hooks(kernel_source.dsl, kernel_function, module_name, self._build_pool.export_dir, False): + kernel_source.call_function(kernel_function, gpu_args, kwargs, grid, threads, params) + else: + kernel_source.call_function(kernel_function, gpu_args, kwargs, grid, threads, params) self.synchronize() return kernel_function + @staticmethod + def _kernel_kwargs(signature, params): + """Tuning params can be kernel arguments, these are passed as keyword arguments to the kernel.""" + if params is None: + return {} + return {arg_name: params[arg_name] for arg_name in signature if arg_name in params} + def start_event(self): """Records the event that marks the start of a measurement.""" self.start.record() @@ -192,7 +269,7 @@ def run_kernel(self, func, gpu_args, threads, grid, params=None): configuration :type params: dict """ - self.call_function(func, gpu_args, self.gpu_kwargs, grid, threads, params) + self.call_function(func, gpu_args, self._kernel_kwargs(self.signature, params), grid, threads, params) def synchronize(self): diff --git a/kernel_tuner/backends/generic_python_processes.py b/kernel_tuner/backends/generic_python_processes.py new file mode 100644 index 000000000..c697dd8e2 --- /dev/null +++ b/kernel_tuner/backends/generic_python_processes.py @@ -0,0 +1,325 @@ +"""Compile Python DSL kernels in worker processes, for DSLs that cannot compile kernels in parallel threads. + +Numba, Warp, and cuTile hold a global lock while compiling, and CuTe kernels can only be launched from the +thread that compiled them. However, these DSLs can store compiled kernels on disk. A worker process imports +the same kernel module file as the main process, and compiles the kernel by running it once, after which the +main process loads the compiled kernel from disk instead of compiling it: + +- Numba: caching is enabled on the kernel, the cache files are removed once the main process loaded them. +- Warp and cuTile: their kernel caches are enabled by default. +- CuTe: ``cute.compile`` does not use the file cache of CuTe, so the compiled kernel is exported with + ``export_to_c`` and loaded with ``cutlass.runtime.load_module``. + +If a kernel cannot be compiled in a worker process, for example because the call function cannot be sent to +the worker, the main process compiles the kernel instead. +""" + +import glob +import importlib.util +import inspect +import itertools +import logging +import multiprocessing +import os +import pickle +import shutil +import sys +import tempfile +import threading +import weakref +from collections import namedtuple +from concurrent.futures import ProcessPoolExecutor +from concurrent.futures.process import BrokenProcessPool +from contextlib import contextmanager + +from kernel_tuner.kernel_sources.kernel_source_fn import kernel_module_name + +# DSLs whose kernels are compiled in worker processes when compiling in parallel +PROCESS_BUILD_DSLS = {"numba", "warp", "cutile", "cute"} + +BuildJob = namedtuple( + "BuildJob", + ["module_path", "module_name", "kernel_name", "dsl", "call_function", "args", "kwargs", "grid", "threads", "params", + "export_dir"], +) + + +@contextmanager +def dsl_compile_hooks(dsl, kernel_function, module_name, export_dir, export): + """Store the kernel compiled while running the call function on disk, or load it from disk. + + :param export: True in a worker process, to store the compiled kernel on disk. False in the main process, + to load the kernel stored by the worker process. + :type export: bool + """ + if dsl == "numba" and hasattr(kernel_function, "enable_caching"): + kernel_function.enable_caching() + yield + if not export: + _remove_numba_cache_files(kernel_function) + elif dsl == "cute": + with _cute_export_hook(module_name, export_dir, export): + yield + else: + yield + + +def _remove_numba_cache_files(kernel_function): + """Remove the cache files of a Numba kernel, once it is loaded they are no longer needed.""" + try: + cache = kernel_function._cache + for path in glob.glob(os.path.join(cache.cache_path, glob.escape(cache._impl.filename_base) + ".*")): + os.remove(path) + except Exception as e: + logging.debug(f"could not remove Numba cache files: {e}") + + +@contextmanager +def _cute_export_hook(module_name, export_dir, export): + """Replace cute.compile to export compiled kernels, or to load exported kernels instead of compiling them.""" + import cutlass.cute as cute + import cutlass.runtime + + original_compile = cute.compile + counter = itertools.count() + keys = [] + + def compile_hook(*args, **kwargs): + # call functions compile the kernels in the same order in the worker and in the main process + key = f"{module_name}_{next(counter)}" + keys.append(key) + object_file = os.path.join(export_dir, key + ".o") + if not export and os.path.exists(object_file): + try: + return cutlass.runtime.load_module(object_file)[key] + except Exception as e: + logging.debug(f"could not load exported CuTe kernel {object_file}, compiling it instead: {e}") + compiled = original_compile(*args, **kwargs) + if export: + try: + compiled.export_to_c(export_dir, key, function_prefix=key) + except Exception as e: + logging.debug(f"could not export CuTe kernel {key}: {e}") + return compiled + + cute.compile = compile_hook + try: + yield + finally: + cute.compile = original_compile + if not export: + for key in keys: + for path in glob.glob(os.path.join(export_dir, glob.escape(key) + ".*")): + os.remove(path) + + +def instantiate_kernel(kernel_fn): + """Class-based kernels are instantiated, their __call__ method is the kernel.""" + if inspect.isclass(kernel_fn): + return kernel_fn() + if callable(kernel_fn): + return kernel_fn + raise TypeError("kernel function is not a class or function") + + +def _to_host(arg): + """Copy PyTorch GPU tensors to the host, to send them to a worker process.""" + import torch + + if isinstance(arg, torch.Tensor) and arg.is_cuda: + return arg.cpu() + return arg + + +def _to_device(arg): + import torch + + if isinstance(arg, torch.Tensor): + return arg.to("cuda") + return arg + + +def _picklable_exception_chain(e): + """Return the exception and the exceptions it was raised from, as they can be sent to the main process. + + The chain is followed like Python does when printing a traceback. Exceptions that cannot be pickled are + replaced by a RuntimeError with the name of the exception type and its message. + """ + chain = [] + while e is not None and all(e is not seen for seen in chain): + chain.append(e) + e = e.__cause__ if e.__cause__ is not None else (None if e.__suppress_context__ else e.__context__) + picklable = [] + for exc in chain: + try: + pickle.loads(pickle.dumps(exc)) + picklable.append(exc) + except Exception: + picklable.append(RuntimeError(f"{type(exc).__name__}: {exc}")) + return picklable + + +def _raise_exception_chain(chain): + for exc, cause in zip(chain, chain[1:]): + exc.__cause__ = cause + raise chain[0] + + +# state of a worker process, set by _init_worker +_worker = {} + + +def _init_worker(device_id, args): + import torch + + torch.cuda.set_device(device_id) + _worker["args"] = [_to_device(arg) for arg in args] + _worker["broken"] = False + + +def _cuda_is_healthy(): + import torch + + try: + torch.cuda.synchronize() + return True + except Exception: + return False + + +def _build_in_worker(job): + """Compile a kernel in a worker process by running it once. + + :returns: None if the kernel was compiled, or the exception chain of the error raised while compiling. + """ + import torch + + if _worker["broken"]: + # the CUDA context of this process is unusable after a sticky error, such as an illegal memory access, + # exit and let the main process start a new worker + os._exit(1) + + try: + args = list(_worker["args"]) + for index, arg in job.args.items(): + args[index] = _to_device(arg) + + spec = importlib.util.spec_from_file_location(job.module_name, job.module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[job.module_name] = module + spec.loader.exec_module(module) + kernel_function = instantiate_kernel(getattr(module, job.kernel_name)) + + with dsl_compile_hooks(job.dsl, kernel_function, job.module_name, job.export_dir, export=True): + torch.cuda.synchronize() + job.call_function(kernel_function, args, job.kwargs, job.grid, job.threads, job.params) + torch.cuda.synchronize() + return None + except Exception as e: + _worker["broken"] = not _cuda_is_healthy() + return _picklable_exception_chain(e) + finally: + sys.modules.pop(job.module_name, None) + + +class BuildProcessPool: + """A pool of worker processes that compile kernels, so that the main process can load them from disk.""" + + def __init__(self, device_id, gpu_args, num_processes): + """Create the pool, the worker processes are started when kernels are built. + + :param gpu_args: The kernel arguments, these are copied to each worker process once. + :type gpu_args: list + + :param num_processes: The maximum number of worker processes. + :type num_processes: int + """ + self.device_id = device_id + self.gpu_args = list(gpu_args) + self.num_processes = num_processes + self.export_dir = tempfile.mkdtemp(prefix="kernel_tuner_build_") + # also remove the exported kernels when tuning stops with an error, and shutdown() is not called + weakref.finalize(self, shutil.rmtree, self.export_dir, ignore_errors=True) + self.disabled = False + self.any_succeeded = False + self._executor = None + self._lock = threading.Lock() + + def _get_executor(self): + with self._lock: + if self._executor is None: + # CUDA cannot be used in forked processes + self._executor = ProcessPoolExecutor( + max_workers=self.num_processes, + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_worker, + initargs=(self.device_id, [_to_host(arg) for arg in self.gpu_args]), + ) + return self._executor + + def build(self, kernel_instance, gpu_args, call_function, kwargs): + """Compile the kernel of kernel_instance in a worker process, this is thread-safe. + + Errors raised while compiling the kernel are raised again in this process. + + :returns: True if the kernel was compiled by a worker process, False if it must be compiled by the + main process instead. + :rtype: bool + """ + if self.disabled: + return False + + # only send the arguments that differ from the arguments the workers already have, such as + # scalars and arguments that depend on the configuration + args = {i: _to_host(arg) for i, arg in enumerate(gpu_args) if arg is not self.gpu_args[i]} + job = BuildJob( + module_path=kernel_instance.temp_files[0], + module_name=kernel_module_name(kernel_instance.temp_files[0]), + kernel_name=kernel_instance.name, + dsl=kernel_instance.kernel_source.dsl, + call_function=call_function, + args=args, + kwargs=kwargs, + grid=kernel_instance.grid, + threads=kernel_instance.threads, + params=kernel_instance.params, + export_dir=self.export_dir, + ) + + executor = self._get_executor() + try: + error_chain = executor.submit(_build_in_worker, job).result() + except BrokenProcessPool as e: + self._handle_broken_pool(executor, e) + return False + except Exception as e: + # usually the job could not be pickled, for example because the call function is a lambda + logging.warning(f"Could not compile kernels in worker processes, compiling in this process instead: {e}") + self.disabled = True + return False + + if error_chain is not None: + _raise_exception_chain(error_chain) + self.any_succeeded = True + return True + + def _handle_broken_pool(self, executor, error): + """Start new worker processes after a worker exited, unless worker processes never worked at all.""" + with self._lock: + if not self.any_succeeded: + # for example, the worker processes cannot import the script that calls tune_kernel + logging.warning( + f"Could not compile kernels in worker processes, compiling in this process instead: {error}" + ) + self.disabled = True + if self._executor is executor: + self._executor = None + executor.shutdown(wait=False, cancel_futures=True) + + def shutdown(self): + """Stop the worker processes and remove the exported kernels.""" + with self._lock: + executor, self._executor = self._executor, None + if executor is not None: + executor.shutdown(wait=True, cancel_futures=True) + shutil.rmtree(self.export_dir, ignore_errors=True) diff --git a/kernel_tuner/backends/nvcuda.py b/kernel_tuner/backends/nvcuda.py index cdfcbc1ef..48a820592 100644 --- a/kernel_tuner/backends/nvcuda.py +++ b/kernel_tuner/backends/nvcuda.py @@ -259,16 +259,23 @@ def ready_argument_list(self, arguments): def compile(self, kernel_instance): """Call the CUDA compiler to compile the kernel, return the device function. - :param kernel_name: The name of the kernel to be compiled, used to lookup the - function after compilation. - :type kernel_name: string - - :param kernel_string: The CUDA kernel code that contains the function `kernel_name` - :type kernel_string: string + :param kernel_instance: The kernel instance, containing the name and code of the kernel + :type kernel_instance: kernel_tuner.core.KernelInstance :returns: A kernel that can be launched by the CUDA runtime :rtype: """ + return self.load(kernel_instance, self.build(kernel_instance)) + + def build(self, kernel_instance): + """Compile the kernel to PTX using NVRTC, this is thread-safe and does not use the device. + + :param kernel_instance: The kernel instance, containing the name and code of the kernel + :type kernel_instance: kernel_tuner.core.KernelInstance + + :returns: The PTX and the lowered name of the kernel inside the PTX + :rtype: tuple(bytes, bytes) + """ kernel_string = kernel_instance.kernel_string kernel_name = kernel_instance.name expression_name = str.encode(kernel_name) @@ -311,26 +318,10 @@ def compile(self, kernel_instance): err = nvrtc.nvrtcGetPTX(program, buff) cuda_error_check(err) - # Load the module - err, self.current_module = driver.cuModuleLoadData(np.char.array(buff)) - if err == driver.CUresult.CUDA_ERROR_INVALID_PTX: - raise SkippableFailure("uses too much shared data") - else: - cuda_error_check(err) - - # First, get the "lowered" name of the kernel (i.e., the name inside the PTX). - # After, we can use the lowered name to lookup the kernel in the module. + # Get the "lowered" name of the kernel (i.e., the name inside the PTX), which is + # used to lookup the kernel in the module after loading it. err, lowered_name = nvrtc.nvrtcGetLoweredName(program, expression_name) cuda_error_check(err) - err, self.func = driver.cuModuleGetFunction( - self.current_module, lowered_name - ) - cuda_error_check(err) - - # get the number of registers per thread used in this kernel - num_regs = driver.cuFuncGetAttribute(driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, self.func) - assert num_regs[0] == 0, f"Retrieving number of registers per thread unsuccesful: code {num_regs[0]}" - self.num_regs = num_regs[1] except RuntimeError as re: _, n = nvrtc.nvrtcGetProgramLogSize(program) @@ -339,6 +330,37 @@ def compile(self, kernel_instance): print(log.decode("utf-8")) raise re + return buff, lowered_name + + def load(self, kernel_instance, build_result): + """Load a kernel compiled by build() onto the device and return the device function. + + :param kernel_instance: The kernel instance + :type kernel_instance: kernel_tuner.core.KernelInstance + + :param build_result: The PTX and lowered name returned by build() + :type build_result: tuple(bytes, bytes) + + :returns: A kernel that can be launched by the CUDA runtime + :rtype: + """ + ptx, lowered_name = build_result + + # Load the module + err, self.current_module = driver.cuModuleLoadData(np.char.array(ptx)) + if err == driver.CUresult.CUDA_ERROR_INVALID_PTX: + raise SkippableFailure("uses too much shared data") + else: + cuda_error_check(err) + + err, self.func = driver.cuModuleGetFunction(self.current_module, lowered_name) + cuda_error_check(err) + + # get the number of registers per thread used in this kernel + num_regs = driver.cuFuncGetAttribute(driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, self.func) + assert num_regs[0] == 0, f"Retrieving number of registers per thread unsuccesful: code {num_regs[0]}" + self.num_regs = num_regs[1] + return self.func def start_event(self): diff --git a/kernel_tuner/backends/pycuda.py b/kernel_tuner/backends/pycuda.py index dca9fa1bf..c5a008950 100644 --- a/kernel_tuner/backends/pycuda.py +++ b/kernel_tuner/backends/pycuda.py @@ -3,6 +3,7 @@ from __future__ import print_function import logging +import threading import numpy as np @@ -170,27 +171,40 @@ def ready_argument_list(self, arguments): def compile(self, kernel_instance): """Call the CUDA compiler to compile the kernel, return the device function. - :param kernel_name: The name of the kernel to be compiled, used to lookup the - function after compilation. - :type kernel_name: string - - :param kernel_string: The CUDA kernel code that contains the function `kernel_name` - :type kernel_string: string + :param kernel_instance: The kernel instance, containing the name and code of the kernel + :type kernel_instance: kernel_tuner.core.KernelInstance :returns: An CUDA kernel that can be called directly. :rtype: pycuda.driver.Function """ + return self.load(kernel_instance, self.build(kernel_instance)) + + def build(self, kernel_instance): + """Compile the kernel with nvcc and load it into a module, this is thread-safe. + + PyCUDA links and loads the module as part of compilation, so the context of this backend + is made current for the calling thread while building the module. + + :param kernel_instance: The kernel instance, containing the name and code of the kernel + :type kernel_instance: kernel_tuner.core.KernelInstance + + :returns: The compiled module + :rtype: pycuda.compiler.SourceModule + """ kernel_string = kernel_instance.kernel_string kernel_name = kernel_instance.name - try: - no_extern_c = 'extern "C"' in kernel_string + no_extern_c = 'extern "C"' in kernel_string - compiler_options = ["-Xcompiler=-Wall"] - if self.compiler_options: - compiler_options += self.compiler_options + compiler_options = ["-Xcompiler=-Wall"] + if self.compiler_options: + compiler_options += self.compiler_options - self.current_module = self.source_mod( + other_thread = threading.current_thread() is not threading.main_thread() + if other_thread: + self.context.push() + try: + return self.source_mod( kernel_string, options=compiler_options + ["-e", kernel_name], arch=("compute_" + self.cc) if self.cc != "00" else None, @@ -198,16 +212,39 @@ def compile(self, kernel_instance): cache_dir=False, no_extern_c=no_extern_c, ) - - self.func = self.current_module.get_function(kernel_name) - if not isinstance(self.func, str): - self.num_regs = self.func.num_regs - return self.func except drv.CompileError as e: if "uses too much shared data" in e.stderr: raise SkippableFailure("uses too much shared data") else: raise e + except drv.LogicError as e: + # DynamicSourceModule compiles to PTX, the PTX is compiled to machine code when it is linked, + # at which point the CUDA driver rejects kernels that use too much shared memory + if "cuLinkComplete failed: device kernel image is invalid" in str(e): + raise SkippableFailure("uses too much shared data") + else: + raise e + finally: + if other_thread: + drv.Context.pop() + + def load(self, kernel_instance, build_result): + """Return the device function from a module compiled by build(). + + :param kernel_instance: The kernel instance + :type kernel_instance: kernel_tuner.core.KernelInstance + + :param build_result: The module returned by build() + :type build_result: pycuda.compiler.SourceModule + + :returns: An CUDA kernel that can be called directly. + :rtype: pycuda.driver.Function + """ + self.current_module = build_result + self.func = self.current_module.get_function(kernel_instance.name) + if not isinstance(self.func, str): + self.num_regs = self.func.num_regs + return self.func def start_event(self): """Records the event that marks the start of a measurement.""" diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 2b82ab246..8fb372a28 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -51,6 +51,10 @@ def _get_cupy(): ) +KernelBuild = namedtuple("KernelBuild", ["result", "error", "time"]) +"""The result of DeviceInterface.build_kernel: the build result, the exception raised, and the build time in ms.""" + + class KernelInstance(_KernelInstance): """Class that represents the specific parameterized instance of a kernel.""" @@ -489,7 +493,15 @@ def check_kernel_output(self, func, gpu_args, instance, answer, atol, verify, ve print("expected: ", answer, "\ngot: ", result_host) raise RuntimeError("Kernel result verification failed for: " + util.get_config_string(instance.params)) - def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, to): + def compile_and_benchmark( + self, kernel_source, gpu_args, params, kernel_options, to, instance=None, build=None, delete_temp_files=True + ): + """Compile, verify, and benchmark a kernel configuration. + + To compile multiple configurations in parallel, runners can first create the kernel instance with + create_kernel_instance() and build it with build_kernel(), and pass these as ``instance`` and ``build``. + Pass ``delete_temp_files=False`` to keep the temporary files of the instance, to use it again. + """ # reset previous timers last_compilation_time = None last_verification_time = None @@ -502,7 +514,8 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, instance_string = util.get_instance_string(params) logging.debug("compile_and_benchmark " + instance_string) - instance = self.create_kernel_instance(kernel_source, kernel_options, params, verbose) + if instance is None: + instance = self.create_kernel_instance(kernel_source, kernel_options, params, verbose) if isinstance(instance, util.ErrorConfig): result["__error__"] = util.InvalidConfig() else: @@ -512,7 +525,7 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, try: # compile the kernel start_compilation = time.perf_counter() - func = self.compile_kernel(instance, verbose, gpu_args) + func = self.compile_kernel(instance, verbose, gpu_args, build) if not func: result["__error__"] = util.CompilationFailedConfig() else: @@ -528,6 +541,9 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, # stop compilation stopwatch and convert to milliseconds last_compilation_time = 1000 * (time.perf_counter() - start_compilation) + if build is not None: + # the kernel was built before, possibly in parallel with other kernels + last_compilation_time += build.time # test kernel for correctness if func and (to.answer or to.verify or self.output_observers): @@ -553,7 +569,8 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, raise e # clean up any temporary files, if no error occurred - instance.delete_temp_files() + if delete_temp_files: + instance.delete_temp_files() # For Python DSLs, the compilation time also includes one kernel run, so we subtract the runtime if self.lang.upper() == "GENERIC_PYTHON" and last_compilation_time and "time" in result: @@ -567,20 +584,46 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, return result - def compile_kernel(self, instance, verbose, gpu_args=None): + def build_kernel(self, instance, gpu_args=None): + """Compile the kernel for this specific instance without loading it, this is thread-safe. + + This is the part of the compilation that can be done in parallel for multiple kernel instances. + Exceptions are not raised, but returned as part of the result, so they can be handled when the + kernel is loaded by compile_kernel(). + + :returns: The result of the backend's build method, any exception it raised, and the build time in ms. + :rtype: KernelBuild + """ + start = time.perf_counter() + result, error = None, None + try: + if self.lang.upper() == "GENERIC_PYTHON": + result = self.dev.build(instance, gpu_args) + else: + result = self.dev.build(instance) + except Exception as e: + error = e + return KernelBuild(result, error, 1000 * (time.perf_counter() - start)) + + def compile_kernel(self, instance, verbose, gpu_args=None, build=None): """Compile the kernel for this specific instance. Generic Python kernels are compiled by running them once, which requires ``gpu_args``. + If the kernel was already built by build_kernel(), pass the result as ``build`` to only load it. """ logging.debug("compile_kernel " + instance.name) # compile kernel_string into device func func = None try: + if build is None: + build = self.build_kernel(instance, gpu_args) + if build.error is not None: + raise build.error if self.lang.upper() == "GENERIC_PYTHON": - func = self.dev.compile(instance, gpu_args) + func = self.dev.load(instance, build.result, gpu_args) else: - func = self.dev.compile(instance) + func = self.dev.load(instance, build.result) except Exception as e: # compiles may fail because certain kernel configurations use too # much shared memory for example, the desired behavior is to simply diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index 062ba2314..417fbe0e0 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -316,18 +316,15 @@ def __deepcopy__(self, _): "call_function", ( """When the language Generic Python is used, a call function that calls the kernel in the Python - DSL. If not specified, the DSL of the kernel is detected from its decorators or base classes and - a default call function from kernel_tuner.utils.call_functions is used. Supported DSLs are - Triton, Numba, CuPy (cupyx.jit), Warp, Taichi, CuTe, Tilus, TileLang, and cuTile. - The function must take the following positional arguments: - :kernel_function: the callable function with the tuning parameters inserted. - :args: list of kernel arguments, as provided by the user in the argument. - :kwargs: dictionary of kernel keyword arguments. If a tuning parameter is in the kernel signature, - the tuning parameter will be added as a keyword argument. - Optionally, the following arguments can be used. The order and name of the arguments must match. - :grid: the launch grid (tuple with 3 values), as computed by KernelTuner - :threads: the thread block size (tuple with 3 values), as computed by KernelTuner - :params: dictionary with the values of the tuning params for the specific configuration.""", + DSL. If not specified, the DSL of the kernel is detected from its decorators or base classes and + a default call function from kernel_tuner.utils.call_functions is used. Supported DSLs are + Triton, Numba, CuPy (cupyx.jit), Warp, Taichi, CuTe, Tilus, TileLang, and cuTile. + The call function receives three positional arguments: kernel_function, the kernel with the + values of the tunable parameters inserted; args, the list of kernel arguments; and kwargs, a + dictionary with the tunable parameters that are also arguments of the kernel. Optionally, the + call function can take the arguments grid and threads, the grid and thread block dimensions + computed by Kernel Tuner, and params, a dictionary with the tunable parameters of the + configuration, in this order and with these names.""", "function", ), ), @@ -532,6 +529,18 @@ def __deepcopy__(self, _): "int|bool", ), ), + ( + "parallel_compile", + ( + """Set to `True` or an integer to compile kernels in parallel threads on the local machine, before + benchmarking them one after the other. If set to an integer, this will be the number of threads, + otherwise the number of CPU cores is used. Kernels are compiled in parallel when the search + strategy evaluates multiple configurations at once, such as brute_force, random_sample, and + population-based strategies. Numba, Warp, cuTile, and CuTe kernels are compiled in worker + processes instead of threads. Cannot be combined with `parallel` or `simulation_mode`.""", + "int|bool", + ), + ), ( "observers", ( @@ -651,6 +660,7 @@ def tune_kernel( metrics=None, simulation_mode=None, parallel=None, + parallel_compile=None, observers=None, objective=None, objective_higher_is_better=None, @@ -776,7 +786,9 @@ def tune_kernel( # Create runner if parallel and simulation_mode: raise ValueError("Enabling `parallel` and `simulation_mode` together is not supported") - elif simulation_mode: + if parallel_compile and (parallel or simulation_mode): + raise ValueError("Enabling `parallel_compile` together with `parallel` or `simulation_mode` is not supported") + if simulation_mode: from kernel_tuner.runners.simulation import SimulationRunner runner = SimulationRunner(kernelsource, kernel_options, device_options, iterations, observers) @@ -799,6 +811,13 @@ def tune_kernel( runner = ParallelRunner( kernelsource, kernel_options, device_options, tuning_options, iterations, observers, num_workers=num_workers ) + elif parallel_compile: + from kernel_tuner.runners.parallel_compile import ParallelCompileRunner + + num_workers = None if parallel_compile is True else parallel_compile + runner = ParallelCompileRunner( + kernelsource, kernel_options, device_options, iterations, observers, num_workers=num_workers + ) else: from kernel_tuner.runners.sequential import SequentialRunner diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 79395f437..d0bbd7a27 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -1,7 +1,7 @@ import inspect import ast import copy -import uuid +import os import sys import logging @@ -18,6 +18,15 @@ from kernel_tuner.utils.language_detection import detect_python_dsl +def kernel_module_name(temp_file_path): + """Return the name of the module created for the kernel source in temp_file_path. + + Worker processes that compile kernels import the same file under the same name, the name of + the module is part of the key under which some DSLs cache compiled kernels. + """ + return os.path.splitext(os.path.basename(temp_file_path))[0] + + class KernelSourceFn(KernelSource): """ Class that holds the Python-function-based kernel sources. @@ -46,9 +55,12 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function if not isinstance(kernel_name, str): raise TypeError("kernel_name should be a string, got ", type(kernel_name)) + # the DSL is also used to decide whether kernels can be compiled in parallel threads + self.dsl = detect_python_dsl(kernel_name, kernel_source) + if self.lang == Language.GENERIC_PYTHON: if call_function is None: - call_function = self._default_call_function(kernel_name, kernel_source) + call_function = self._default_call_function(kernel_name, self.dsl) if not callable(call_function): raise TypeError(f"call_function of type {type(call_function)} is not a callable object.") self.call_function = call_function @@ -67,9 +79,8 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function @staticmethod - def _default_call_function(kernel_name, kernel_source): + def _default_call_function(kernel_name, dsl): """Select the default call function for the DSL the kernel is written in.""" - dsl = detect_python_dsl(kernel_name, kernel_source) if dsl is None: raise ValueError( f"Could not detect the Python DSL of kernel {kernel_name}, please pass a call_function. " @@ -146,11 +157,12 @@ def apply_params_to_source_fn(self, params): #print(new_source) - # Create a unique module name and write new source to it. - module_name = f'temp_kernel_module_{uuid.uuid4().hex}' - with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as temp_file: + # Write the new source to a uniquely named module, the module is named after the file. + prefix = 'temp_kernel_module_' + with tempfile.NamedTemporaryFile(mode='w', prefix=prefix, suffix='.py', delete=False) as temp_file: temp_file.write(new_source) temp_file_path = temp_file.name + module_name = kernel_module_name(temp_file_path) # Register the module in sys.modules before executing it spec = importlib.util.spec_from_file_location(module_name, temp_file_path) diff --git a/kernel_tuner/runners/parallel_compile.py b/kernel_tuner/runners/parallel_compile.py new file mode 100644 index 000000000..1ac32a5b7 --- /dev/null +++ b/kernel_tuner/runners/parallel_compile.py @@ -0,0 +1,123 @@ +"""A runner that compiles configurations in parallel threads, and benchmarks them one after the other.""" + +import logging +import os +from concurrent.futures import ThreadPoolExecutor + +from kernel_tuner.runners.sequential import SequentialRunner +from kernel_tuner.util import ErrorConfig, Timer + +# The maximum number of worker processes used to compile kernels when the number of workers is not specified. +# Each worker process creates its own CUDA context and copy of the kernel arguments on the GPU. +DEFAULT_MAX_BUILD_PROCESSES = 8 + + +class ParallelCompileRunner(SequentialRunner): + """ParallelCompileRunner compiles configurations in parallel on the local machine. + + The configurations passed to run() are processed in chunks. The kernels of all configurations in a chunk + are compiled in parallel threads, and once all compilations have finished, the configurations are + verified and benchmarked one after the other, exactly like the SequentialRunner does. This way, the + compilation of other kernels does not affect the benchmarks. + + Compilation is done in threads, which only reduces compile time if the compiler does not hold Python's + global interpreter lock, for example because it runs as a separate process, or in C/C++ code that releases + the lock. Backends that do not implement DeviceInterface.build_kernel compile the kernels when they are + loaded, in which case the configurations are compiled one after the other. + + Python DSLs that cannot compile kernels in parallel threads (Numba, Warp, cuTile, and CuTe) compile the + kernels in worker processes instead, from which the main process loads the compiled kernels. + """ + + def __init__(self, kernel_source, kernel_options, device_options, iterations, observers, num_workers=None): + """Instantiate the ParallelCompileRunner. + + :param num_workers: The number of threads used for compilation, by default the number of CPU cores. + :type num_workers: int + """ + super().__init__(kernel_source, kernel_options, device_options, iterations, observers) + self.num_workers = num_workers or os.cpu_count() or 1 + # Compiling a few kernels per thread balances the load over the threads, while limiting + # the number of compiled kernels that are kept in memory until they are benchmarked. + self.chunk_size = 4 * self.num_workers + + # some backends compile kernels that cannot be compiled in threads in worker processes instead, + # each worker process uses GPU memory, so the default number of processes is limited + if hasattr(self.dev.dev, "build_processes"): + self.dev.dev.build_processes = num_workers or min(self.num_workers, DEFAULT_MAX_BUILD_PROCESSES) + + def shutdown(self): + """Stop the worker processes that the backend may have started to compile kernels.""" + if hasattr(self.dev.dev, "shutdown_build_processes"): + self.dev.dev.shutdown_build_processes() + + def run(self, parameter_space, tuning_options): + """Compile the configurations in parameter_space in parallel, and benchmark them one after the other. + + :param parameter_space: The parameter space as an iterable. + :type parameter_space: iterable + + :param tuning_options: A dictionary with all options regarding the tuning process. + :type tuning_options: kernel_tuner.interface.Options + + :returns: A list of dictionaries for executed kernel configurations and their execution times. + :rtype: dict() + """ + logging.debug("parallel compile runner started for " + self.kernel_options.kernel_name) + + elements = list(parameter_space) + results = [] + worker_time = 0 + warmup_time = 0 + for start in range(0, len(elements), self.chunk_size): + chunk = elements[start : start + self.chunk_size] + + build_timer = Timer() + prebuilt = self._build(chunk, tuning_options) + worker_time += build_timer.get() + + chunk_results, chunk_worker_time, chunk_warmup_time = self._evaluate(chunk, tuning_options, prebuilt) + results += chunk_results + worker_time += chunk_worker_time + warmup_time += chunk_warmup_time + + self._add_strategy_and_framework_time(results, worker_time, warmup_time) + return results + + def _build(self, chunk, tuning_options): + """Build the kernels of the configurations in chunk that will be evaluated, in parallel. + + :returns: The kernel instance and build of each configuration, indexed by its position in the chunk. + :rtype: dict(int: tuple(KernelInstance, KernelBuild)) + """ + # Only build configurations that will be evaluated within the budget of function evaluations, + # and that are not in the cache. A time limit may still be reached while benchmarking. + budget = tuning_options.budget + if budget.is_done(): + return {} + remaining = budget.get_evaluations_remaining() + + instances = {} + futures = {} + with ThreadPoolExecutor(max_workers=self.num_workers) as pool: + for index, element in enumerate(chunk): + if index >= remaining: + break + x_int = ",".join([str(i) for i in element]) + if tuning_options.cache and x_int in tuning_options.cache: + continue + + # kernel instances are created in the main thread, as this may involve writing files + params = dict(zip(tuning_options.tune_params.keys(), element)) + instance = self.dev.create_kernel_instance( + self.kernel_source, self.kernel_options, params, tuning_options.verbose + ) + instances[index] = instance + if isinstance(instance, ErrorConfig): + continue + gpu_args = self.dev.preprocess_gpu_arguments(self.gpu_args, params) + futures[index] = pool.submit(self.dev.build_kernel, instance, gpu_args) + + # all builds have finished when the pool is shut down at the end of the with statement + return {index: (instance, futures[index].result() if index in futures else None) + for index, instance in instances.items()} diff --git a/kernel_tuner/runners/sequential.py b/kernel_tuner/runners/sequential.py index c987524f7..285bb7510 100644 --- a/kernel_tuner/runners/sequential.py +++ b/kernel_tuner/runners/sequential.py @@ -64,12 +64,26 @@ def run(self, parameter_space, tuning_options): """ logging.debug("sequential runner started for " + self.kernel_options.kernel_name) + results, worker_time, warmup_time = self._evaluate(parameter_space, tuning_options) + self._add_strategy_and_framework_time(results, worker_time, warmup_time) + return results + + def _evaluate(self, parameter_space, tuning_options, prebuilt=None): + """Evaluate the configurations in parameter_space one after the other. + + :param prebuilt: Kernel instances and builds for configurations that were built before, indexed + by the position of the configuration in parameter_space, see ParallelCompileRunner. + :type prebuilt: dict(int: tuple(KernelInstance, KernelBuild)) + + :returns: The results, the time spent by the worker in seconds, and the warmup time in seconds. + """ + prebuilt = prebuilt or {} results = [] worker_time = 0 warmup_time = 0 # iterate over parameter space - for element in parameter_space: + for index, element in enumerate(parameter_space): # If the time limit is exceeded, just skip this element. Add `None` to # indicate to CostFunc that no result is available for this config. if tuning_options.budget.is_done(): @@ -86,23 +100,30 @@ def run(self, parameter_space, tuning_options): cache_entry = tuning_options.cache[x_int] params.update(copy_without_benchmark_timings(cache_entry)) else: + instance, build = prebuilt.get(index, (None, None)) + # attempt to warmup the GPU by running the first config in the parameter space and ignoring the result if not self.warmed_up: warmup_timer = Timer() + # keep the temporary files of a prebuilt instance, it is used again below self.dev.compile_and_benchmark( - self.kernel_source, self.gpu_args, params, self.kernel_options, tuning_options + self.kernel_source, self.gpu_args, params, self.kernel_options, tuning_options, + instance, build, delete_temp_files=instance is None, ) self.warmed_up = True warmup_time = warmup_timer.get() result = self.dev.compile_and_benchmark( - self.kernel_source, self.gpu_args, params, self.kernel_options, tuning_options + self.kernel_source, self.gpu_args, params, self.kernel_options, tuning_options, instance, build ) # Collect total time spent by worker in seconds worker_time += ( result["compile_time"] + result["verification_time"] + result["benchmark_time"] ) / 1000 + if build is not None: + # time spent building in parallel is accounted for by the runner that built the kernel + worker_time -= build.time / 1000 assert check_result_type(result) @@ -129,6 +150,10 @@ def run(self, parameter_space, tuning_options): # all visited configurations are added to results to provide a trace for optimization strategies results.append(params) + return results, worker_time, warmup_time + + def _add_strategy_and_framework_time(self, results, worker_time, warmup_time): + """Amortize the time spent by the strategy and the framework over the results.""" # Count the number of valid results num_valid_results = sum(bool(r) for r in results) @@ -146,5 +171,3 @@ def run(self, parameter_space, tuning_options): # Time must be in ms result["strategy_time"] = 1000 * strategy_time / num_valid_results result["framework_time"] = 1000 * framework_time / num_valid_results - - return results diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index a16c8df5f..5476aa330 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -1170,6 +1170,21 @@ def has_kw_argument(func, name): return lambda answer, result_host, atol: v(answer, result_host) +class _NormalizedCallFunction: + """Call function that accepts grid, threads and params, and passes only the ones the user function takes. + + This is a class instead of a closure, so that it can be pickled to compile kernels in worker processes. + """ + + def __init__(self, function, optional_args): + self.function = function + self.optional_args = optional_args + + def __call__(self, kernel_function, args, kwargs, grid, threads, params): + optional = {"grid": grid, "threads": threads, "params": params} + return self.function(kernel_function, args, kwargs, *(optional[name] for name in self.optional_args)) + + def normalize_call_function(v): """Normalize a user-specified call function for language Generic Python. @@ -1184,16 +1199,16 @@ def has_kw_argument(func, name): sig = signature(func) return name in sig.parameters - if v is None: - return None + if v is None or isinstance(v, _NormalizedCallFunction): + return v if has_kw_argument(v, "grid"): if has_kw_argument(v, "threads"): if has_kw_argument(v, "params"): return v - return lambda kernel_function, args, kwargs, grid, threads, params: v(kernel_function, args, kwargs, grid, threads) - return lambda kernel_function, args, kwargs, grid, threads, params: v(kernel_function, args, kwargs, grid) - return lambda kernel_function, args, kwargs, grid, threads, params: v(kernel_function, args, kwargs) + return _NormalizedCallFunction(v, ("grid", "threads")) + return _NormalizedCallFunction(v, ("grid",)) + return _NormalizedCallFunction(v, ()) diff --git a/test/test_generic_python_dsls.py b/test/test_generic_python_dsls.py index 9a759c0f8..d291a2a5a 100644 --- a/test/test_generic_python_dsls.py +++ b/test/test_generic_python_dsls.py @@ -1,12 +1,19 @@ """End-to-end tests for tuning kernels written in Python DSLs with lang="generic_python".""" import importlib +import pickle import numpy as np import pytest from kernel_tuner import tune_kernel from kernel_tuner.backends.generic_python import GenericPythonFunctions +from kernel_tuner.backends.generic_python_processes import ( + PROCESS_BUILD_DSLS, + BuildProcessPool, + _picklable_exception_chain, + _raise_exception_chain, +) from .context import ( skip_if_no_cupy, @@ -35,8 +42,8 @@ ] -@pytest.mark.parametrize("dsl", dsls) -def test_generic_python_dsl(dsl): +def tune_vector_add(dsl, parallel_compile): + """Tune the vector_add kernel of a DSL and check that every configuration is valid.""" # the kernel modules import their DSL, so only import them when the DSL is installed kernel = importlib.import_module(f".generic_python_kernels.{dsl}_vector_add", package=__package__) if getattr(kernel, "requires_skip", None): @@ -51,13 +58,89 @@ def test_generic_python_dsl(dsl): tune_params = kernel.tune_params(n) # the language and call function are detected from the kernel source - results, _ = tune_kernel(kernel.kernel_name, kernel.__file__, n, args, tune_params, answer=answer) + results, _ = tune_kernel( + kernel.kernel_name, kernel.__file__, n, args, tune_params, answer=answer, parallel_compile=parallel_compile + ) num_configs = np.prod([len(v) for v in tune_params.values()]) assert len(results) == num_configs # verification raises an exception on wrong results, so every configuration must have a valid time failed = [result for result in results if not isinstance(result.get("time"), float)] assert not failed, f"configurations failed: {failed}" + return results + + +@pytest.mark.parametrize("parallel_compile", [None, 4], ids=["sequential", "parallel_compile"]) +@pytest.mark.parametrize("dsl", dsls) +def test_generic_python_dsl(dsl, parallel_compile): + tune_vector_add(dsl, parallel_compile) + + +@pytest.mark.parametrize("dsl", [dsl for dsl in dsls if dsl.values[0] in PROCESS_BUILD_DSLS]) +def test_generic_python_dsl_process_builds(dsl, monkeypatch): + """DSLs that cannot compile in parallel threads compile in worker processes, the main process loads them.""" + built = [] + original_build = BuildProcessPool.build + + def recording_build(self, *args, **kwargs): + built.append(original_build(self, *args, **kwargs)) + return built[-1] + + monkeypatch.setattr(BuildProcessPool, "build", recording_build) + results = tune_vector_add(dsl, parallel_compile=2) + + # the first configuration is benchmarked twice, for warming up the GPU, but only built once + assert len(built) == len(results) + assert all(built), "some kernels were not compiled by a worker process" + + +@skip_if_no_torch +@skip_if_no_numba_cuda +def test_generic_python_process_builds_fallback(caplog): + """Kernels are compiled in the main process when the call function cannot be sent to worker processes.""" + from .generic_python_kernels import numba_vector_add as kernel + + n = 4096 + a = np.random.randn(n).astype(np.float32) + b = np.random.randn(n).astype(np.float32) + c = np.zeros_like(a) + + def call_numba_local(kernel_function, args, kwargs, grid, threads): + from numba import cuda + + kernel_function[grid, threads](*[cuda.as_cuda_array(arg) for arg in args[:3]], args[3]) + + results, _ = tune_kernel( + kernel.kernel_name, kernel.__file__, n, [c, a, b, n], {"block_size_x": [128, 256, 512]}, + answer=[a + b, None, None, None], call_function=call_numba_local, parallel_compile=2, + ) + assert all(isinstance(result.get("time"), float) for result in results) + assert "Could not compile kernels in worker processes" in caplog.text + + +def test_worker_exception_chain(): + """Errors raised in worker processes are raised in the main process with the exceptions they came from.""" + + class Unpicklable(Exception): + def __init__(self, message, handle): + super().__init__(message) + self.handle = lambda: handle + + error = wrapped(Unpicklable("too many registers", 1), RuntimeError("compilation failed")) + chain = pickle.loads(pickle.dumps(_picklable_exception_chain(error))) + assert [type(e) for e in chain] == [RuntimeError, RuntimeError] + assert str(chain[1]) == "Unpicklable: too many registers" + + with pytest.raises(RuntimeError) as raised: + _raise_exception_chain(chain) + assert GenericPythonFunctions.classify_compile_exception(None, raised.value) == "resource_error" + + # exception types are kept when possible, they are used to classify the error + chain = pickle.loads(pickle.dumps(_picklable_exception_chain(wrapped(NameError("name 'x' is not defined"))))) + with pytest.raises(RuntimeError) as raised: + _raise_exception_chain(chain) + assert isinstance(raised.value.__cause__, NameError) + assert GenericPythonFunctions.classify_compile_exception(None, raised.value) == "user_error" def wrapped(error, wrapper=None): diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py index b8cddea88..003bf59ae 100644 --- a/test/test_generic_python_functions.py +++ b/test/test_generic_python_functions.py @@ -122,7 +122,7 @@ def test_gpu_kwargs(): gpu_args = dev.ready_argument_list(args) callable_fn = dev.compile_kernel(instance, verbose=False, gpu_args=gpu_args) - kwargs = dev.dev.gpu_kwargs + kwargs = dev.dev._kernel_kwargs(ks.signature, params) assert kwargs["mock_param"] == 64 assert callable_fn(*args, **kwargs) == 64 diff --git a/test/test_parallel_compile.py b/test/test_parallel_compile.py new file mode 100644 index 000000000..92caed690 --- /dev/null +++ b/test/test_parallel_compile.py @@ -0,0 +1,149 @@ +"""Tests for compiling kernels in parallel threads with the ParallelCompileRunner (parallel_compile option).""" + +import threading + +import numpy as np +import pytest + +from kernel_tuner import tune_kernel +from kernel_tuner.core import DeviceInterface + +from .context import skip_backend, skip_if_no_gcc + +CUDA_KERNEL = """ +extern "C" __global__ void vector_add(float *c, float *a, float *b, int n) { + __shared__ float buffer[block_size_x * shared_factor]; + int i = blockIdx.x * block_size_x + threadIdx.x; + buffer[threadIdx.x * shared_factor] = 1.0f; + if (i < n) { + c[i] = a[i] + b[i] * buffer[threadIdx.x * shared_factor]; + } +} +""" + +C_KERNEL = """ +#include +float vector_add(float *c, float *a, float *b, int n) { + for (int i = 0; i < n; i += block_size_x) { + for (int j = i; j < i + block_size_x && j < n; j++) { + c[j] = a[j] + b[j]; + } + } + return 0.0f; +} +""" + + +def vector_add_args(size=10000): + a = np.random.randn(size).astype(np.float32) + b = np.random.randn(size).astype(np.float32) + c = np.zeros_like(b) + return [c, a, b, np.int32(size)], [a + b, None, None, None] + + +def valid_configs(results): + return sorted( + tuple(sorted((k, v) for k, v in r.items() if k in ("block_size_x", "shared_factor"))) + for r in results + if isinstance(r.get("time"), float) + ) + + +@pytest.mark.parametrize("lang", ["NVCUDA", "PYCUDA", "CUPY"]) +def test_parallel_compile_cuda(lang, monkeypatch): + skip_backend(lang) + args, answer = vector_add_args() + # 64 * 512 floats is 128 KB of shared memory, too much for any GPU, so those configurations are skipped + tune_params = {"block_size_x": [32, 64, 128, 256, 512], "shared_factor": [1, 64]} + + # record which threads the kernels are built in + build_threads = [] + original_build_kernel = DeviceInterface.build_kernel + + def recording_build_kernel(self, instance, gpu_args=None): + build_threads.append(threading.current_thread()) + return original_build_kernel(self, instance, gpu_args) + + monkeypatch.setattr(DeviceInterface, "build_kernel", recording_build_kernel) + + sequential, _ = tune_kernel("vector_add", CUDA_KERNEL, 10000, args, tune_params, lang=lang, answer=answer) + assert all(t is threading.main_thread() for t in build_threads) + build_threads.clear() + + parallel, _ = tune_kernel( + "vector_add", CUDA_KERNEL, 10000, args, tune_params, lang=lang, answer=answer, parallel_compile=4 + ) + assert len(build_threads) == 10 + assert all(t is not threading.main_thread() for t in build_threads) + + # the same configurations are valid, and the ones that use too much shared memory are skipped + assert len(parallel) == len(sequential) == 10 + assert valid_configs(parallel) == valid_configs(sequential) + assert 0 < len(valid_configs(parallel)) < 10 + assert all(r["compile_time"] > 0 for r in parallel if isinstance(r.get("time"), float)) + + +@skip_if_no_gcc +def test_parallel_compile_backend_without_build(): + """The C backend does not implement build(), so kernels are compiled when they are loaded.""" + args, answer = vector_add_args() + tune_params = {"block_size_x": [1, 2, 4, 8, 16, 32]} + results, _ = tune_kernel( + "vector_add", C_KERNEL, 10000, args, tune_params, lang="C", answer=answer, parallel_compile=3 + ) + assert len(valid_configs(results)) == 6 + + +@skip_if_no_gcc +def test_parallel_compile_chunks_and_budget(monkeypatch): + args, answer = vector_add_args() + tune_params = {"block_size_x": list(range(1, 41))} + + built = [] + original_build_kernel = DeviceInterface.build_kernel + + def counting_build_kernel(self, instance, gpu_args=None): + built.append(instance.params["block_size_x"]) + return original_build_kernel(self, instance, gpu_args) + + monkeypatch.setattr(DeviceInterface, "build_kernel", counting_build_kernel) + + # with 2 threads, the 40 configurations are processed in chunks of 8 + results, _ = tune_kernel( + "vector_add", C_KERNEL, 10000, args, tune_params, lang="C", answer=answer, parallel_compile=2 + ) + assert len(valid_configs(results)) == 40 + assert sorted(built) == list(range(1, 41)) + + # only configurations within the budget are built + built.clear() + results, _ = tune_kernel( + "vector_add", C_KERNEL, 10000, args, tune_params, lang="C", answer=answer, parallel_compile=2, + strategy="random_sample", strategy_options={"max_fevals": 11}, + ) + assert len(results) == 11 + assert len(built) == 11 + assert sorted(built) == sorted(r["block_size_x"] for r in results) + + +@pytest.mark.parametrize("lang", ["NVCUDA", "PYCUDA", "CUPY"]) +def test_parallel_compile_error_is_raised(lang, tmp_path, monkeypatch): + skip_backend(lang) + # the kernel source is written to the current directory when compilation fails + monkeypatch.chdir(tmp_path) + args, _ = vector_add_args() + kernel = CUDA_KERNEL.replace("buffer[threadIdx.x * shared_factor] = 1.0f;", "undefined_variable = 1.0f;") + tune_params = {"block_size_x": [32, 64], "shared_factor": [1]} + with pytest.raises(Exception): + tune_kernel("vector_add", kernel, 10000, args, tune_params, lang=lang, parallel_compile=2, quiet=True) + + +def test_parallel_compile_option_combinations(): + args, _ = vector_add_args() + tune_params = {"block_size_x": [32]} + with pytest.raises(ValueError, match="parallel_compile"): + tune_kernel("vector_add", C_KERNEL, 10000, args, tune_params, lang="C", parallel_compile=2, parallel=2) + with pytest.raises(ValueError, match="parallel_compile"): + tune_kernel( + "vector_add", C_KERNEL, 10000, args, tune_params, lang="C", parallel_compile=2, simulation_mode=True + ) diff --git a/test/test_util_functions.py b/test/test_util_functions.py index aed8e97f4..c5ce94287 100644 --- a/test/test_util_functions.py +++ b/test/test_util_functions.py @@ -3,6 +3,7 @@ import json import os +import pickle import warnings import datetime @@ -644,6 +645,17 @@ def kernel(x): assert v(kernel, (1,), {}, grid=1, threads=2, params=3) == 2 + +def call_with_grid(kernel_function, args, kwargs, grid): + return kernel_function(*args, grid) + + +def test_normalize_call_function_picklable(): + """Normalized call functions can be pickled, to send them to worker processes that compile kernels.""" + v = pickle.loads(pickle.dumps(normalize_call_function(call_with_grid))) + assert v(lambda x, grid: x + grid, (1,), {}, grid=2, threads=3, params=4) == 3 + assert normalize_call_function(v) is v + class MockRunner: simulation_mode = False From fc7879083c5c3a94bb51b51c05a4256fbd2d9db5 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 30 Sep 2026 23:09:00 +0200 Subject: [PATCH 30/44] restore skip_if_no_pycuda on test_get_device_interface1 The decorator was accidentally removed together with the detect_language tests in f656a7d1, which made the test fail on CI where PyCUDA is not installed. --- test/test_util_functions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_util_functions.py b/test/test_util_functions.py index c5ce94287..ed9b966a4 100644 --- a/test/test_util_functions.py +++ b/test/test_util_functions.py @@ -241,6 +241,7 @@ def test_check_restrictions(): assert answer == e +@skip_if_no_pycuda def test_get_device_interface1(): lang = "PYCUDA" dev = core.DeviceInterface(core.KernelSource("", "", lang=lang)) From 94d69122a49801f685a6131e57084eadf30833eb Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Thu, 1 Oct 2026 23:11:26 +0200 Subject: [PATCH 31/44] stop leaving temporary kernel files in /tmp Import temporary kernel modules with a loader that does not write .pyc files, both in the main process and in the build workers. Every configuration gets its own module that is imported only once, so the cached bytecode was never used and only accumulated in /tmp/__pycache__. The tests in test_kernel_source_fn.py and test_generic_python_functions.py now write their temporary kernel modules to their pytest tmp_path. --- .../backends/generic_python_processes.py | 8 ++----- .../kernel_sources/kernel_source_fn.py | 24 +++++++++++++++---- test/test_generic_python_functions.py | 2 +- test/test_kernel_source_fn.py | 14 +++++++++++ 4 files changed, 36 insertions(+), 12 deletions(-) diff --git a/kernel_tuner/backends/generic_python_processes.py b/kernel_tuner/backends/generic_python_processes.py index c697dd8e2..945c233c4 100644 --- a/kernel_tuner/backends/generic_python_processes.py +++ b/kernel_tuner/backends/generic_python_processes.py @@ -15,7 +15,6 @@ """ import glob -import importlib.util import inspect import itertools import logging @@ -32,7 +31,7 @@ from concurrent.futures.process import BrokenProcessPool from contextlib import contextmanager -from kernel_tuner.kernel_sources.kernel_source_fn import kernel_module_name +from kernel_tuner.kernel_sources.kernel_source_fn import import_kernel_module, kernel_module_name # DSLs whose kernels are compiled in worker processes when compiling in parallel PROCESS_BUILD_DSLS = {"numba", "warp", "cutile", "cute"} @@ -204,10 +203,7 @@ def _build_in_worker(job): for index, arg in job.args.items(): args[index] = _to_device(arg) - spec = importlib.util.spec_from_file_location(job.module_name, job.module_path) - module = importlib.util.module_from_spec(spec) - sys.modules[job.module_name] = module - spec.loader.exec_module(module) + module = import_kernel_module(job.module_name, job.module_path) kernel_function = instantiate_kernel(getattr(module, job.kernel_name)) with dsl_compile_hooks(job.dsl, kernel_function, job.module_name, job.export_dir, export=True): diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index d0bbd7a27..b5e7160e7 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -6,6 +6,7 @@ import logging import tempfile +import importlib.machinery import importlib.util from typing import Any @@ -27,6 +28,23 @@ def kernel_module_name(temp_file_path): return os.path.splitext(os.path.basename(temp_file_path))[0] +class _NoBytecodeLoader(importlib.machinery.SourceFileLoader): + """Source loader that does not write .pyc files, kernel modules are imported only once.""" + + def set_data(self, path, data, *, _mode=0o666): + pass + + +def import_kernel_module(module_name, temp_file_path): + """Import the kernel module in temp_file_path under module_name and register it in sys.modules.""" + loader = _NoBytecodeLoader(module_name, temp_file_path) + spec = importlib.util.spec_from_file_location(module_name, temp_file_path, loader=loader) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + class KernelSourceFn(KernelSource): """ Class that holds the Python-function-based kernel sources. @@ -164,11 +182,7 @@ def apply_params_to_source_fn(self, params): temp_file_path = temp_file.name module_name = kernel_module_name(temp_file_path) - # Register the module in sys.modules before executing it - spec = importlib.util.spec_from_file_location(module_name, temp_file_path) - temp_module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = temp_module - spec.loader.exec_module(temp_module) + temp_module = import_kernel_module(module_name, temp_file_path) new_fn = getattr(temp_module, self.kernel_name) return new_fn, temp_file_path diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py index 003bf59ae..495457c15 100644 --- a/test/test_generic_python_functions.py +++ b/test/test_generic_python_functions.py @@ -1,5 +1,5 @@ from .context import skip_if_no_torch -from .test_kernel_source_fn import call_mock +from .test_kernel_source_fn import call_mock, temp_kernel_dir # noqa: F401, temp_kernel_dir is an autouse fixture from kernel_tuner.core import DeviceInterface, KernelInstance from kernel_tuner.kernel_sources.kernel_source import KernelSource import numpy as np diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py index f9b852f61..e7025ad88 100644 --- a/test/test_kernel_source_fn.py +++ b/test/test_kernel_source_fn.py @@ -1,4 +1,5 @@ import os +import tempfile from pathlib import Path import pytest import inspect @@ -11,6 +12,14 @@ from kernel_tuner.kernel_sources.kernel_source_str import KernelSourceStr KS_FILE = Path(__file__).resolve() + + +@pytest.fixture(autouse=True) +def temp_kernel_dir(tmp_path, monkeypatch): + """Write the temporary kernel modules of a test to its tmp_path, instead of leaving them in /tmp.""" + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + return tmp_path + # Helper functions -------------------------------------------- def normalize_ast(src: str): @@ -169,4 +178,9 @@ def test_dependencies(): assert res == 128 +def test_no_bytecode(temp_kernel_dir): + ks = KernelSourceFn("mock_kernel", KS_FILE, "generic_python", call_function=call_mock) + _, temp_path = ks.apply_params_to_source_fn({"mock_param": 128}) + assert Path(temp_path).parent == temp_kernel_dir + assert not (temp_kernel_dir / "__pycache__").exists() From 10baa85ea11bbd8902aa3c3d976a291dae1a5b91 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Mon, 5 Oct 2026 10:49:19 +0200 Subject: [PATCH 32/44] fix parallel compilation of Taichi, Tilus, and Numba float16 kernels - compile Taichi kernels one at a time in the main thread, Taichi does not support compiling kernels in parallel threads - compile Tilus kernels in worker processes, Tilus keeps the state of the kernel it is compiling in global variables - compile Numba kernels that cannot be cached, such as float16 kernels, in the main process instead of worker processes - keep the message of exceptions that lose it when pickled, such as nvJitLinkError, so that compile errors from workers are classified correctly --- CHANGELOG.md | 2 +- doc/source/parallel.rst | 14 ++- kernel_tuner/backends/generic_python.py | 8 +- .../backends/generic_python_processes.py | 88 ++++++++++++++----- kernel_tuner/interface.py | 5 +- kernel_tuner/runners/parallel_compile.py | 5 +- .../numba_vector_add_fp16.py | 11 +++ test/test_generic_python_dsls.py | 38 ++++++++ 8 files changed, 140 insertions(+), 31 deletions(-) create mode 100644 test/generic_python_kernels/numba_vector_add_fp16.py diff --git a/CHANGELOG.md b/CHANGELOG.md index fd3850b79..b4a7ab848 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers - Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, cuTile, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` - Detect kernels in Python files as `lang="generic_python"`, and detect the DSL of Python kernels to select a default `call_function` when none is given -- Add `parallel_compile` option to compile kernels in parallel threads before benchmarking them one after the other, supported by the CUDA backends and all Python DSLs. Numba, Warp, cuTile, and CuTe kernels are compiled in worker processes instead of threads +- Add `parallel_compile` option to compile kernels in parallel threads before benchmarking them one after the other, supported by the CUDA backends and all Python DSLs. Numba, Warp, cuTile, CuTe, and Tilus kernels are compiled in worker processes instead of threads, Taichi kernels one at a time in the main thread - Fix PyCUDA backend raising an error instead of skipping configurations that use too much shared memory - Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend - Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel diff --git a/doc/source/parallel.rst b/doc/source/parallel.rst index 2f7daff3f..649c5398a 100644 --- a/doc/source/parallel.rst +++ b/doc/source/parallel.rst @@ -181,16 +181,19 @@ Parallel compilation is supported by the CUDA backends (CUDA-Python, CuPy, and P for all supported DSLs. Other backends compile the kernels one at a time. ``parallel_compile`` cannot be combined with ``parallel`` or ``simulation_mode``. -Numba, Warp, cuTile, and CuTe cannot compile kernels in parallel threads: the first three compile one kernel at a -time because of a global compiler lock, and CuTe kernels can only be launched from the thread that compiled them. +Numba, Warp, cuTile, CuTe, and Tilus cannot compile kernels in parallel threads: Numba, Warp, and cuTile compile +one kernel at a time because of a global compiler lock, CuTe kernels can only be launched from the thread that +compiled them, and Tilus keeps the state of the kernel it is compiling in global variables. These kernels are compiled in worker processes instead, at most one per thread, and at most 8 when ``parallel_compile`` is set to ``True``. A worker process compiles the kernel by running it once, after which Kernel Tuner loads the compiled kernel from disk. Numba and CuTe kernels are -stored in a temporary cache that Kernel Tuner removes, Warp and cuTile kernels are stored in the kernel caches of -these DSLs. Note that: +stored in a temporary cache that Kernel Tuner removes, Warp, cuTile, and Tilus kernels are stored in the kernel +caches of these DSLs. Note that: * Starting the worker processes takes several seconds, so compiling in worker processes only pays off when compiling a kernel takes long, or when many kernels are compiled. +* Numba cannot store some kernels in its cache, such as kernels that use float16. Kernel Tuner then prints a + warning and compiles these kernels one at a time. * Each worker process creates its own CUDA context on the GPU and a copy of the kernel arguments, which uses GPU memory. Use a smaller number for ``parallel_compile`` if the GPU runs out of memory. * The worker processes are started with the ``spawn`` method of ``multiprocessing``, which imports the script that @@ -198,3 +201,6 @@ these DSLs. Note that: ``if __name__ == "__main__":``, as the examples do. A ``call_function`` must be defined at the top level of a module, so that it can be sent to the worker processes. When kernels cannot be compiled in worker processes, Kernel Tuner prints a warning and compiles them one at a time. + +Taichi kernels cannot be compiled in parallel threads either. Taichi compiles kernels quickly, so these are +compiled one at a time in the main thread. diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index dd99f03a8..6df2ea720 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -23,8 +23,12 @@ # DSLs whose kernels are not compiled in parallel threads. Numba, Warp, and cuTile compile one kernel at a # time because of a global compiler lock, so parallel threads are not faster (and even slower for Warp). # CuTe keeps the state of compiled kernels per thread, so these cannot be launched from another thread. +# Tilus keeps the state of the kernel it is compiling in global variables, so it can compile one kernel at a time. # When build_processes is set, these are compiled in worker processes instead, see generic_python_processes. -SERIAL_COMPILE_DSLS = PROCESS_BUILD_DSLS +# Taichi keeps global compiler state that does not support compiling in threads, and has to set up its runtime +# in the main thread. Taichi compiles quickly, so its kernels are always compiled by load() in the main thread. +MAIN_THREAD_DSLS = {"taichi"} +SERIAL_COMPILE_DSLS = PROCESS_BUILD_DSLS | MAIN_THREAD_DSLS # Returned by build() for kernels that were compiled by a worker process, load() loads these from disk PROCESS_BUILT = "process_built" @@ -155,7 +159,7 @@ def build(self, kernel_instance, gpu_args=None): or None if the kernel is compiled by load(). """ if kernel_instance.kernel_source.dsl in SERIAL_COMPILE_DSLS: - if not self.build_processes: + if not self.build_processes or kernel_instance.kernel_source.dsl in MAIN_THREAD_DSLS: return None kernel_source = kernel_instance.kernel_source kwargs = self._kernel_kwargs(kernel_source.signature, kernel_instance.params) diff --git a/kernel_tuner/backends/generic_python_processes.py b/kernel_tuner/backends/generic_python_processes.py index 945c233c4..9f8ca300d 100644 --- a/kernel_tuner/backends/generic_python_processes.py +++ b/kernel_tuner/backends/generic_python_processes.py @@ -1,12 +1,14 @@ """Compile Python DSL kernels in worker processes, for DSLs that cannot compile kernels in parallel threads. -Numba, Warp, and cuTile hold a global lock while compiling, and CuTe kernels can only be launched from the -thread that compiled them. However, these DSLs can store compiled kernels on disk. A worker process imports -the same kernel module file as the main process, and compiles the kernel by running it once, after which the -main process loads the compiled kernel from disk instead of compiling it: +Numba, Warp, and cuTile hold a global lock while compiling, CuTe kernels can only be launched from the +thread that compiled them, and Tilus keeps the state of the kernel it is compiling in global variables. +However, these DSLs can store compiled kernels on disk. A worker process imports the same kernel module file +as the main process, and compiles the kernel by running it once, after which the main process loads the +compiled kernel from disk instead of compiling it: - Numba: caching is enabled on the kernel, the cache files are removed once the main process loaded them. -- Warp and cuTile: their kernel caches are enabled by default. + Numba cannot cache some kernels, such as kernels that use float16, these are compiled by the main process. +- Warp, cuTile, and Tilus: their kernel caches are enabled by default. - CuTe: ``cute.compile`` does not use the file cache of CuTe, so the compiled kernel is exported with ``export_to_c`` and loaded with ``cutlass.runtime.load_module``. @@ -34,7 +36,7 @@ from kernel_tuner.kernel_sources.kernel_source_fn import import_kernel_module, kernel_module_name # DSLs whose kernels are compiled in worker processes when compiling in parallel -PROCESS_BUILD_DSLS = {"numba", "warp", "cutile", "cute"} +PROCESS_BUILD_DSLS = {"numba", "warp", "cutile", "cute", "tilus"} BuildJob = namedtuple( "BuildJob", @@ -50,17 +52,42 @@ def dsl_compile_hooks(dsl, kernel_function, module_name, export_dir, export): :param export: True in a worker process, to store the compiled kernel on disk. False in the main process, to load the kernel stored by the worker process. :type export: bool + + :returns: A list to which the errors are added that prevented storing the compiled kernel on disk. + :rtype: list(Exception) """ + not_stored = [] if dsl == "numba" and hasattr(kernel_function, "enable_caching"): kernel_function.enable_caching() - yield - if not export: + if export: + _catch_numba_cache_errors(kernel_function, not_stored) + yield not_stored + # in a worker, Numba may have written the index of its cache before it failed to store the kernel + if not export or not_stored: _remove_numba_cache_files(kernel_function) elif dsl == "cute": with _cute_export_hook(module_name, export_dir, export): - yield + yield not_stored else: - yield + yield not_stored + + +def _catch_numba_cache_errors(kernel_function, not_stored): + """Let the kernel compile when Numba cannot store it in its cache, and add the error to not_stored. + + Numba cannot cache some kernels, such as kernels that are linked with other files, which includes kernels + that use float16. + """ + cache = kernel_function._cache + save_overload = cache.save_overload + + def save_overload_or_record_error(sig, data): + try: + save_overload(sig, data) + except Exception as e: + not_stored.append(e) + + cache.save_overload = save_overload_or_record_error def _remove_numba_cache_files(kernel_function): @@ -141,8 +168,10 @@ def _to_device(arg): def _picklable_exception_chain(e): """Return the exception and the exceptions it was raised from, as they can be sent to the main process. - The chain is followed like Python does when printing a traceback. Exceptions that cannot be pickled are - replaced by a RuntimeError with the name of the exception type and its message. + The chain is followed like Python does when printing a traceback. Exceptions that cannot be pickled, or + that lose part of their message when pickled, are replaced by a RuntimeError with the name of the exception + type and its message. For example, nvJitLinkError only pickles its status, not the error log of the linker + that its message includes, which is needed to classify the error. """ chain = [] while e is not None and all(e is not seen for seen in chain): @@ -151,10 +180,12 @@ def _picklable_exception_chain(e): picklable = [] for exc in chain: try: - pickle.loads(pickle.dumps(exc)) - picklable.append(exc) + if str(pickle.loads(pickle.dumps(exc))) == str(exc): + picklable.append(exc) + continue except Exception: - picklable.append(RuntimeError(f"{type(exc).__name__}: {exc}")) + pass + picklable.append(RuntimeError(f"{type(exc).__name__}: {exc}")) return picklable @@ -186,10 +217,15 @@ def _cuda_is_healthy(): return False +class KernelNotStored(str): + """Returned by a worker process that compiled a kernel that the DSL could not store on disk, with the reason.""" + + def _build_in_worker(job): """Compile a kernel in a worker process by running it once. - :returns: None if the kernel was compiled, or the exception chain of the error raised while compiling. + :returns: None if the kernel was compiled, KernelNotStored if the kernel was compiled but could not be stored + on disk, or the exception chain of the error raised while compiling. """ import torch @@ -206,10 +242,12 @@ def _build_in_worker(job): module = import_kernel_module(job.module_name, job.module_path) kernel_function = instantiate_kernel(getattr(module, job.kernel_name)) - with dsl_compile_hooks(job.dsl, kernel_function, job.module_name, job.export_dir, export=True): + with dsl_compile_hooks(job.dsl, kernel_function, job.module_name, job.export_dir, export=True) as not_stored: torch.cuda.synchronize() job.call_function(kernel_function, args, job.kwargs, job.grid, job.threads, job.params) torch.cuda.synchronize() + if not_stored: + return KernelNotStored(f"{type(not_stored[0]).__name__}: {not_stored[0]}") return None except Exception as e: _worker["broken"] = not _cuda_is_healthy() @@ -284,7 +322,7 @@ def build(self, kernel_instance, gpu_args, call_function, kwargs): executor = self._get_executor() try: - error_chain = executor.submit(_build_in_worker, job).result() + result = executor.submit(_build_in_worker, job).result() except BrokenProcessPool as e: self._handle_broken_pool(executor, e) return False @@ -294,8 +332,18 @@ def build(self, kernel_instance, gpu_args, call_function, kwargs): self.disabled = True return False - if error_chain is not None: - _raise_exception_chain(error_chain) + if isinstance(result, KernelNotStored): + # the other configurations of the kernel can most likely not be stored either + with self._lock: + if not self.disabled: + logging.warning( + "Could not store kernels compiled by worker processes, compiling in this process instead: " + + result + ) + self.disabled = True + return False + if result is not None: + _raise_exception_chain(result) self.any_succeeded = True return True diff --git a/kernel_tuner/interface.py b/kernel_tuner/interface.py index 417fbe0e0..1d64fa9b0 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -536,8 +536,9 @@ def __deepcopy__(self, _): benchmarking them one after the other. If set to an integer, this will be the number of threads, otherwise the number of CPU cores is used. Kernels are compiled in parallel when the search strategy evaluates multiple configurations at once, such as brute_force, random_sample, and - population-based strategies. Numba, Warp, cuTile, and CuTe kernels are compiled in worker - processes instead of threads. Cannot be combined with `parallel` or `simulation_mode`.""", + population-based strategies. Numba, Warp, cuTile, CuTe, and Tilus kernels are compiled in + worker processes instead of threads, Taichi kernels one at a time. Cannot be combined with + `parallel` or `simulation_mode`.""", "int|bool", ), ), diff --git a/kernel_tuner/runners/parallel_compile.py b/kernel_tuner/runners/parallel_compile.py index 1ac32a5b7..3d5904c2f 100644 --- a/kernel_tuner/runners/parallel_compile.py +++ b/kernel_tuner/runners/parallel_compile.py @@ -25,8 +25,9 @@ class ParallelCompileRunner(SequentialRunner): the lock. Backends that do not implement DeviceInterface.build_kernel compile the kernels when they are loaded, in which case the configurations are compiled one after the other. - Python DSLs that cannot compile kernels in parallel threads (Numba, Warp, cuTile, and CuTe) compile the - kernels in worker processes instead, from which the main process loads the compiled kernels. + Python DSLs that cannot compile kernels in parallel threads (Numba, Warp, cuTile, CuTe, and Tilus) compile + the kernels in worker processes instead, from which the main process loads the compiled kernels. Taichi kernels + are compiled one at a time in the main thread. """ def __init__(self, kernel_source, kernel_options, device_options, iterations, observers, num_workers=None): diff --git a/test/generic_python_kernels/numba_vector_add_fp16.py b/test/generic_python_kernels/numba_vector_add_fp16.py new file mode 100644 index 000000000..f07a78856 --- /dev/null +++ b/test/generic_python_kernels/numba_vector_add_fp16.py @@ -0,0 +1,11 @@ +from numba import cuda + +kernel_name = "vector_add" + + +# Numba links float16 kernels with other files, which means that Numba cannot store them in its cache +@cuda.jit +def vector_add(c, a, b, n): + i = cuda.grid(1) + if i < n: + c[i] = a[i] + b[i] diff --git a/test/test_generic_python_dsls.py b/test/test_generic_python_dsls.py index d291a2a5a..f0cee30e8 100644 --- a/test/test_generic_python_dsls.py +++ b/test/test_generic_python_dsls.py @@ -118,6 +118,25 @@ def call_numba_local(kernel_function, args, kwargs, grid, threads): assert "Could not compile kernels in worker processes" in caplog.text +@skip_if_no_torch +@skip_if_no_numba_cuda +def test_generic_python_process_builds_not_stored(caplog): + """Kernels are compiled in the main process when the DSL cannot store the kernels compiled by workers.""" + from .generic_python_kernels import numba_vector_add_fp16 as kernel + + n = 4096 + a = np.random.randn(n).astype(np.float16) + b = np.random.randn(n).astype(np.float16) + c = np.zeros_like(a) + + results, _ = tune_kernel( + kernel.kernel_name, kernel.__file__, n, [c, a, b, n], {"block_size_x": [128, 256, 512]}, + answer=[a + b, None, None, None], atol=1e-2, parallel_compile=2, + ) + assert all(isinstance(result.get("time"), float) for result in results) + assert "Could not store kernels compiled by worker processes" in caplog.text + + def test_worker_exception_chain(): """Errors raised in worker processes are raised in the main process with the exceptions they came from.""" @@ -135,6 +154,14 @@ def __init__(self, message, handle): _raise_exception_chain(chain) assert GenericPythonFunctions.classify_compile_exception(None, raised.value) == "resource_error" + # exceptions that lose their message when pickled, like nvJitLinkError, are replaced to keep the message + error = LosesMessage(4, "\nnvJitLink error log: ptxas error : Entry function uses too much shared data") + chain = pickle.loads(pickle.dumps(_picklable_exception_chain(wrapped(error)))) + assert "too much shared data" in str(chain[1]) + with pytest.raises(RuntimeError) as raised: + _raise_exception_chain(chain) + assert GenericPythonFunctions.classify_compile_exception(None, raised.value) == "resource_error" + # exception types are kept when possible, they are used to classify the error chain = pickle.loads(pickle.dumps(_picklable_exception_chain(wrapped(NameError("name 'x' is not defined"))))) with pytest.raises(RuntimeError) as raised: @@ -143,6 +170,17 @@ def __init__(self, message, handle): assert GenericPythonFunctions.classify_compile_exception(None, raised.value) == "user_error" +class LosesMessage(Exception): + """Like nvJitLinkError, only the status is pickled, not the error log that is part of the message.""" + + def __init__(self, status, log=""): + super().__init__(f"ERROR_PTX_COMPILE ({status}){log}") + self.status = status + + def __reduce__(self): + return (type(self), (self.status,)) + + def wrapped(error, wrapper=None): """Raise error and wrap it the way DSLs and call functions do, returning the wrapper.""" wrapper = wrapper or RuntimeError("wrapper without details") From 96dbccb7dfb9982d8752b400c17974ea955c8290 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Mon, 5 Oct 2026 10:49:20 +0200 Subject: [PATCH 33/44] require Python 3.12.2 or newer Python versions before 3.12.2 can raise "AST constructor recursion depth mismatch" when code is parsed in parallel threads (CPython gh-106905), which happens when compiling Python DSL kernels with parallel_compile. Also replace tomli with the built-in tomllib. --- .github/workflows/docs-on-release.yml | 4 ++-- .github/workflows/docs.yml | 2 +- CHANGELOG.md | 1 + INSTALL.rst | 2 +- doc/source/dev-environment.rst | 8 ++++---- noxfile.py | 6 +++--- pyproject.toml | 5 +---- test/test_toml_file.py | 4 ++-- 8 files changed, 15 insertions(+), 17 deletions(-) diff --git a/.github/workflows/docs-on-release.yml b/.github/workflows/docs-on-release.yml index da5b05cdd..649de4475 100644 --- a/.github/workflows/docs-on-release.yml +++ b/.github/workflows/docs-on-release.yml @@ -12,10 +12,10 @@ jobs: environment: dev_environment runs-on: ubuntu-latest steps: - - name: Get Python 3.10 + - name: Get Python 3.12 uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.12" - name: Checkout uses: actions/checkout@v4 with: diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 3d53495ff..5295f5dcc 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -15,7 +15,7 @@ jobs: steps: - uses: actions/setup-python@v4 with: - python-version: "3.10" + python-version: "3.12" - name: Checkout uses: actions/checkout@v4 with: diff --git a/CHANGELOG.md b/CHANGELOG.md index b4a7ab848..c65fc29e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ All notable changes to this project will be documented in this file. This project adheres to [Semantic Versioning](http://semver.org/). ## Unreleased +- Require Python 3.12.2 or newer, Python 3.10 and 3.11 are no longer supported - Fix `np.shuffle` crash and wrong return type in `Searchspace` distributed random and LHS samplers - Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, cuTile, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` - Detect kernels in Python files as `lang="generic_python"`, and detect the DSL of Python kernels to select a default `call_function` when none is given diff --git a/INSTALL.rst b/INSTALL.rst index 8e938676f..4d1ffbac5 100644 --- a/INSTALL.rst +++ b/INSTALL.rst @@ -20,7 +20,7 @@ Linux users could type the following to download and install Python 3 using Mini wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh bash Miniconda3-latest-Linux-x86_64.sh -You are of course also free to use your own Python installation, and the Kernel Tuner is developed to be fully compatible with Python 3.10 and newer. +You are of course also free to use your own Python installation, and the Kernel Tuner is developed to be fully compatible with Python 3.12.2 and newer. Installing Python Packages -------------------------- diff --git a/doc/source/dev-environment.rst b/doc/source/dev-environment.rst index 275f40039..c4480413f 100644 --- a/doc/source/dev-environment.rst +++ b/doc/source/dev-environment.rst @@ -27,9 +27,9 @@ Steps with :bash:`sudo` access (e.g. on a local device): * After installation, restart your shell. #. Install the required Python versions: * On some systems, additional packages may be needed to build Python versions. For example on Ubuntu: :bash:`sudo apt install build-essential zlib1g-dev libncurses5-dev libgdbm-dev libnss3-dev libssl-dev libreadline-dev libffi-dev libsqlite3-dev wget libbz2-dev liblzma-dev lzma`. - * Install the Python versions with: :bash:`pyenv install 3.9 3.10 3.11 3.12 3.13`. The reason we're installing all these versions as opposed to just one, is so we can test against all supported Python versions. -#. Set the Python versions so they can be found: :bash:`pyenv local 3.9 3.10 3.11 3.12 3.13` (replace :bash:`local` with :bash:`global` when not using the virtualenv). -#. Setup a local virtual environment in the folder: :bash:`pyenv virtualenv 3.11 kerneltuner` (or whatever environment name and Python version you prefer). + * Install the Python versions with: :bash:`pyenv install 3.12 3.13 3.14`. The reason we're installing all these versions as opposed to just one, is so we can test against all supported Python versions. +#. Set the Python versions so they can be found: :bash:`pyenv local 3.12 3.13 3.14` (replace :bash:`local` with :bash:`global` when not using the virtualenv). +#. Setup a local virtual environment in the folder: :bash:`pyenv virtualenv 3.12 kerneltuner` (or whatever environment name and Python version you prefer). #. `Install Poetry `__. * Use :bash:`curl -sSL https://install.python-poetry.org | python3 -` to install Poetry. * Make sure to add Poetry to :bash:`PATH` as instructed at the end of the installation. @@ -63,7 +63,7 @@ Steps without :bash:`sudo` access (e.g. on a cluster): * Exit the shell and re-enter to make sure Conda is available. :bash:`cd` to the kernel tuner directory. * [Optional] if you have limited user folder space, the Pip cache can be pointed elsewhere with the environment variable :bash:`PIP_CACHE_DIR`. The cache location can be checked with :bash:`pip cache dir`. On Linux, to point the entire :bash:`~/.cache` default elsewhere, use the :bash:`XDG_CACHE_HOME` environment variable. * [Optional] update Conda if available before continuing: :bash:`conda update -n base -c conda-forge conda`. -#. Setup a virtual environment: :bash:`conda create --name kerneltuner python=3.11` (or whatever Python version and environment name you prefer). +#. Setup a virtual environment: :bash:`conda create --name kerneltuner python=3.12` (or whatever Python version and environment name you prefer). #. Activate the virtual environment: :bash:`conda activate kerneltuner`. * [Optional] to use the correct environment by default, execute :bash:`conda config --set auto_activate_base false`, and add `conda activate kerneltuner` to your :bash:`.bash_profile` or :bash:`.bashrc`. #. Make sure that non-Python dependencies are loaded if applicable, such as CUDA, OpenCL or HIP. On most clusters it is possible to load (or unload) modules (e.g. CUDA, OpenCL / ROCM). For more information, see :ref:`Installation `. diff --git a/noxfile.py b/noxfile.py index f93da3e5a..cd290b427 100644 --- a/noxfile.py +++ b/noxfile.py @@ -21,7 +21,7 @@ # set the test parameters verbose = False -python_versions_to_test = ["3.14", "3.13", "3.12", "3.11"] +python_versions_to_test = ["3.14", "3.13", "3.12"] nox.options.stop_on_first_error = True nox.options.error_on_missing_interpreters = True nox.options.default_venv_backend = "virtualenv" @@ -72,9 +72,9 @@ def create_settings(session: Session) -> None: # obtain workspace level settings from the 'noxsettings.toml' file if settings_file_path.exists(): with settings_file_path.open(mode="rb") as fp: - import tomli + import tomllib - nox_settings = tomli.load(fp) + nox_settings = tomllib.load(fp) venvbackend = nox_settings["venvbackend"] envdir = nox_settings["envdir"] assert venvbackend in venvbackend_values, ( diff --git a/pyproject.toml b/pyproject.toml index 6977f175c..00bf3dd54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,15 +41,13 @@ classifiers = [ "Topic :: Scientific/Engineering", "Topic :: Software Development", "Topic :: System :: Distributed Computing", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14" ] # ATTENTION: if anything is changed here, run `poetry update` -requires-python = ">=3.10,<4" # <4 is because of hip-python # NOTE when changing the Python versions, also change the test versions in the Noxfile and GitHub Actions +requires-python = ">=3.12.2,<4" # <4 is because of hip-python # NOTE when changing the Python versions, also change the test versions in the Noxfile and GitHub Actions dependencies = [ "numpy>=2.0.0", "scipy>=1.14.1", # Python >=3.13 needs scipy >=1.14 @@ -130,7 +128,6 @@ nox = "^2024.4.15" nox-poetry = "^1.0.3" ruff = "^0.15.0" pep440 = "^0.1.2" -tomli = "^2.0.1" # held back by Python <= 3.10, can be replaced by built-in [tomllib](https://docs.python.org/3.11/library/tomllib.html) from Python 3.11 onwards scikit-optimize = "0.10.2" # development dependencies are unused for now, as this is already covered by test and docs diff --git a/test/test_toml_file.py b/test/test_toml_file.py index 29b353ae9..f8c035036 100644 --- a/test/test_toml_file.py +++ b/test/test_toml_file.py @@ -2,13 +2,13 @@ from pathlib import Path -import tomli +import tomllib package_root = Path(".").parent.parent pyproject_toml_path = package_root / "pyproject.toml" assert pyproject_toml_path.exists() with pyproject_toml_path.open(mode="rb") as fp: - pyproject = tomli.load(fp) + pyproject = tomllib.load(fp) project = pyproject["project"] if "project" in pyproject else pyproject["tool"]["poetry"] From f62ce4a35e122780eb3a58ecece978192836b2d9 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Mon, 5 Oct 2026 09:25:52 +0000 Subject: [PATCH 34/44] skip skopt strategy test when scikit-optimize is not installed --- test/context.py | 8 ++++++++ test/strategies/test_strategies.py | 4 +++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/test/context.py b/test/context.py index 298303652..2d017033f 100644 --- a/test/context.py +++ b/test/context.py @@ -116,6 +116,13 @@ def python_dsl_present(module, subpackage=None): except ImportError: pyatf_present = False +try: + import skopt + + skopt_present = True +except ImportError: + skopt_present = False + try: import pymoo pymoo_present = True @@ -152,6 +159,7 @@ def python_dsl_present(module, subpackage=None): ) skip_if_no_hip = pytest.mark.skipif(not hip_present, reason="No HIP Python found") skip_if_no_pyatf = pytest.mark.skipif(not pyatf_present, reason="PyATF not installed") +skip_if_no_skopt = pytest.mark.skipif(not skopt_present, reason="scikit-optimize not installed") skip_if_no_methodology = pytest.mark.skipif(not methodology_present, reason="Autotuning Methodology not found") skip_if_no_pymoo = pytest.mark.skipif(not pymoo_present, reason="No PyMOO found") skip_if_no_torch = pytest.mark.skipif(not gen_python_torch_present, reason="Torch not installed or no CUDA device") diff --git a/test/strategies/test_strategies.py b/test/strategies/test_strategies.py index 854bdb649..f12595a6b 100644 --- a/test/strategies/test_strategies.py +++ b/test/strategies/test_strategies.py @@ -8,7 +8,7 @@ from kernel_tuner.util import InvalidConfig from kernel_tuner.interface import strategy_map -from ..context import skip_if_no_bayesopt_botorch, skip_if_no_bayesopt_gpytorch, skip_if_no_pyatf +from ..context import skip_if_no_bayesopt_botorch, skip_if_no_bayesopt_gpytorch, skip_if_no_pyatf, skip_if_no_skopt cache_filename = Path(__file__).parent / "test_cache_file.json" @@ -51,6 +51,8 @@ def vector_add(): strategies.append(pytest.param(s, marks=skip_if_no_bayesopt_botorch)) elif 'pyatf' in s.lower(): strategies.append(pytest.param(s, marks=skip_if_no_pyatf)) + elif s.lower() == 'skopt': + strategies.append(pytest.param(s, marks=skip_if_no_skopt)) else: strategies.append(s) From 8f332421136585e95d3973a57038042851ac6716 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Mon, 5 Oct 2026 11:11:43 +0000 Subject: [PATCH 35/44] keep only the basic matmul examples and use the same matrix size in all of them The optimized matmul kernels move to a separate benchmark suite. Compute the references of the CuPy, Numba, Taichi, and Warp examples in float32 so that the larger matrices can be verified quickly. --- examples/python/README.rst | 17 +- examples/python/matmul/cupy_matmul.py | 8 +- examples/python/matmul/cute_matmul.py | 970 ---------------------- examples/python/matmul/cutile_matmul.py | 45 - examples/python/matmul/numba_matmul.py | 180 +--- examples/python/matmul/taichi_matmul.py | 8 +- examples/python/matmul/tilelang_matmul.py | 218 ----- examples/python/matmul/tilus_matmul.py | 140 ---- examples/python/matmul/triton_matmul.py | 172 ---- examples/python/matmul/warp_matmul.py | 7 +- 10 files changed, 22 insertions(+), 1743 deletions(-) diff --git a/examples/python/README.rst b/examples/python/README.rst index 5ba463f6b..b53d33e2f 100644 --- a/examples/python/README.rst +++ b/examples/python/README.rst @@ -85,10 +85,8 @@ CuPy Numba ~~~~~ [`numba_matmul.py `__] - - tune a basic kernel and a kernel that uses shared memory tiling and register blocking - - tell Kernel Tuner to compute the grid dimensions for 2D thread blocks with tiling - - use restrictions that relate multiple tunable parameters to each other - - use Bayesian Optimization to search a large search space + - tune a basic kernel with 2-dimensional thread blocks + - use the restrictions option to limit the search to valid thread block sizes Taichi ~~~~~~ @@ -104,14 +102,14 @@ Warp Triton ~~~~~~ [`triton_matmul.py `__] - - tune a basic kernel and the optimized kernel from the Triton tutorials + - tune a basic tiled kernel that uses ``tl.dot`` - tune the Triton compiler options ``num_warps`` and ``num_stages`` - use Bayesian Optimization to search a large search space Tilus ~~~~~ [`tilus_matmul.py `__] - - tune the tile sizes and number of warps of class-based kernels + - tune the tile sizes and number of warps of a class-based kernel - use Bayesian Optimization to search a large search space TileLang @@ -123,14 +121,11 @@ TileLang CuTe ~~~~ [`cute_matmul.py `__] - - tune a naive kernel and the Ampere tensor core GEMM from the CuTe DSL examples - - tune a class-based kernel, where the tunable parameters are attributes of the class - - use restrictions to exclude configurations that exceed the shared memory or that the kernel does not support - - pass a user-defined call function that creates the tensor layouts required by the kernel + - tune a naive kernel, where the thread block size is set in the ``@cute.jit`` function that launches it + - use the restrictions option to limit the search to valid thread block sizes cuTile ~~~~~~ [`cutile_matmul.py `__] - tune the tile sizes of a tile-based kernel that uses tensor cores - - tune the size of groups of tiles to improve L2 cache reuse, using a 1-dimensional grid - tell Kernel Tuner to compute a grid with one block per tile of the output matrix diff --git a/examples/python/matmul/cupy_matmul.py b/examples/python/matmul/cupy_matmul.py index 300bd4c2f..d3900b09f 100644 --- a/examples/python/matmul/cupy_matmul.py +++ b/examples/python/matmul/cupy_matmul.py @@ -34,7 +34,7 @@ def run(M, N, K): # Correctness verification c_ref = cp.matmul(a, b) - assert cp.allclose(c, c_ref, rtol=1e-2, atol=1e-1) + assert cp.allclose(c, c_ref, rtol=1e-2, atol=M * 2**(-11)) print("Succes") @@ -56,14 +56,14 @@ def tune(M, N, K): problem_size=size, arguments=args, tune_params=tune_params, - answer=[None, None, A.dot(B), None, None, None], - atol=1e-1, + answer=[None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16), None, None, None], + atol=M * 2**(-11), restrictions=restrictions, verbose=True, ) if __name__ == "__main__": - M, N, K = 1024, 1024, 1024 + M, N, K = 4096, 4096, 4096 tune(M, N, K) \ No newline at end of file diff --git a/examples/python/matmul/cute_matmul.py b/examples/python/matmul/cute_matmul.py index a84f76d59..f5e8f47e8 100644 --- a/examples/python/matmul/cute_matmul.py +++ b/examples/python/matmul/cute_matmul.py @@ -1,10 +1,7 @@ -import math import torch -from typing import Tuple import cutlass import cutlass.cute as cute -import cutlass.utils as utils from cutlass.cute.runtime import from_dlpack from kernel_tuner import tune_kernel @@ -94,974 +91,7 @@ def tune_naive_matmul(M, N, K): answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) - -# This kernel is copied from the CuTe DSL project: -# https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/ampere/tensorop_gemm.py -# -# Original example: tensorop_gemm.py -# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# Modifications in file: -# - Tuning paramters made more explicit - -## Optimized matmul ========================================================== -# Copyright (c) 2025 - 2026 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. - - - -""" -A dense GEMM (C = A * B) example for the NVIDIA Ampere architecture using CUTE DSL. -- Matrix A is MxKxL, L is batch dimension, A can be row-major("K") or column-major("M") -- Matrix B is NxKxL, L is batch dimension, B can be row-major("N") or column-major("K") -- Matrix C is MxNxL, L is batch dimension, C can be row-major("N") or column-major("M") - -This GEMM kernel supports the following features: - - Utilizes Ampere's tensor cores for matrix multiply-accumulate (MMA) operations - - Threadblock rasterization to improve data re-use - - Supports multi-stage pipeline to overlap computation and memory access - - Implements shared memory buffering for epilogue to increase coalesced global memory access - -This GEMM works as follows: -1. Load A and B matrices from global memory (GMEM) to shared memory (SMEM) using asynchronous copies. -2. Perform matrix multiply-accumulate (MMA) operations. -3. Store results from registers (RMEM) to shared memory (SMEM), then to global memory (GMEM). - -The Ampere tensor core instruction used operates as follows: -- Read matrix A from SMEM -- Read matrix B from SMEM -- Perform MMA operation and store the result in Accumulator(register) - -To run this example: - -.. code-block:: bash - - python examples/ampere/tensorop_gemm.py \ - --mnkl 8192,8192,8192,1 --atom_layout_mnk 2,2,1 \ - --ab_dtype Float16 \ - --c_dtype Float16 --acc_dtype Float32 \ - --a_major m --b_major n --c_major n - -The above example command computes with M=8192, N=8192, K=8192, -batch_count=1. The atom layout's shape is 2x2x1 and the input, mma -accumulator, and output data type are set as fp16, fp32 and fp16, -respectively. - -To collect performance with NCU profiler: - -.. code-block:: bash - - ncu python examples/ampere/tensorop_gemm.py \ - --mnkl 8192,8192,8192,1 --atom_layout_mnk 2,2,1 \ - --ab_dtype Float16 \ - --c_dtype Float16 --acc_dtype Float32 \ - --a_major m --b_major n --c_major n \ - --skip_ref_check --iterations 2 - -Constraints: -* Supported input and output data types: fp16 -* Support accumulator data types: f32 -* Default tile shape is set to be 128x128x32 -* Atom layout's MNK shape is set so that tile shape can be divided by MMA - instruction shape -* The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, - i.e, number of elements is a multiple of 8 -""" - - -class TensorOpGemm: - def __init__(self): - self.ab_dtype = cutlass.Float16 - self.c_dtype = cutlass.Float16 - self.acc_dtype = cutlass.Float32 - self.bM = 128 # extracted from cta_tiler for KT support - self.bN = 128 # extracted from cta_tiler for KT support - self.bK = 32 # extracted from cta_tiler for KT support - self.cta_tiler = (self.bM, self.bN, self.bK) - self.num_stages = 3 - atom_lay_M = 2 # extracted from atom_layout_mnk for KT support - atom_lay_N = 2 # extracted from atom_layout_mnk for KT support - atom_lay_K = 1 # extracted from atom_layout_mnk for KT support - self.atom_layout_mnk = (atom_lay_M, atom_lay_N, atom_lay_K) # moved from init parameter to here for KT support - self.num_threads = atom_lay_M * atom_lay_N * atom_lay_K * 32 - - self.mma_inst_shape = (16, 8, 16) - mmaM, mmaN, mmaK = self.mma_inst_shape - - assert self.bM % (atom_lay_M * mmaM) == 0, ( - "bM must be divisible by MMA instruction" - ) - assert self.bN % (atom_lay_N * mmaN) == 0, ( - "bN must be divisible by MMA instruction" - ) - assert atom_lay_K == 1, "this example does not support atom layout K > 1" - assert self.bK % mmaK == 0, "bK must be divisible by MMA instruction" - assert self.num_stages >= 3, "num_stages must be greater than or equal to 3" - - @cute.jit - def __call__( - self, - mA: cute.Tensor, - mB: cute.Tensor, - mC: cute.Tensor, - epilogue_op: cutlass.Constexpr = lambda x: x, - ): - # The grid divides the problems's M, N, and L dimensions by the - # respective modes of the tile shape (bM, bN, 1). The K dimension is - # handled within a block via a multistage process. - - self.a_major_mode = utils.LayoutEnum.from_tensor(mA) - self.b_major_mode = utils.LayoutEnum.from_tensor(mB) - self.c_major_mode = utils.LayoutEnum.from_tensor(mC) - - # /////////////////////////////////////////////////////////////////////////////// - # Shared memory layout: - # /////////////////////////////////////////////////////////////////////////////// - - # Creates a layout with the size required for the provided tile - # size and num stages (stages are used for K dimension) that is also - # sectioned into 64x8 or 8x32 layout atoms. The swizzle is set so that - # the atom for the shared memory -> register copy does not encounter - # bank conflicts - - # assume the input is 16B align - ab_copy_bits = 128 - sA_layout = self._make_smem_layout_AB( - mA.element_type, - self.a_major_mode, - ab_copy_bits, - (self.cta_tiler[0], self.cta_tiler[2], self.num_stages), - ) - sB_layout = self._make_smem_layout_AB( - mB.element_type, - self.b_major_mode, - ab_copy_bits, - (self.cta_tiler[1], self.cta_tiler[2], self.num_stages), - ) - - # Creates a similar layout but without num_stages or layout atoms - sC_layout = self._make_smem_layout_C( - mC.element_type, - self.c_major_mode, - ab_copy_bits, - (self.cta_tiler[0], self.cta_tiler[1]), - ) - - # /////////////////////////////////////////////////////////////////////////////// - # Tiled copy: - # The majorness of tA/tB/tC follows the majorness of gA/gB/gC, - # enabling merged accesses to global memory for faster data - # transfer between global and shared memory. - # /////////////////////////////////////////////////////////////////////////////// - - # Create a copy atom for a global to shared memory asynchronous copy - atom_async_copy = cute.make_copy_atom( - cute.nvgpu.cpasync.CopyG2SOp( - cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL - ), - mA.element_type, - num_bits_per_copy=ab_copy_bits, - ) - - # Create thread layouts for tiled copy from the copy atom where the - # thread layout simply follows the leading dimension of the tensor - tiled_copy_A = self._make_gmem_tiled_copy_AB( - atom_async_copy, mA.element_type, self.a_major_mode, ab_copy_bits - ) - tiled_copy_B = self._make_gmem_tiled_copy_AB( - atom_async_copy, mB.element_type, self.b_major_mode, ab_copy_bits - ) - - # Creates a synchronous copy atom and thread layouts for the epilogue - c_copy_bits = 128 - atom_sync_copy = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - mC.element_type, - num_bits_per_copy=c_copy_bits, - ) - tiled_copy_C = self._make_gmem_tiled_copy_C( - atom_sync_copy, mC.element_type, self.c_major_mode, c_copy_bits - ) - - # /////////////////////////////////////////////////////////////////////////////// - # Tiled MMA - # /////////////////////////////////////////////////////////////////////////////// - - # Creates a mma atom with 16x8x16 shape for MNK - op = cute.nvgpu.warp.MmaF16BF16Op( - self.ab_dtype, self.acc_dtype, self.mma_inst_shape - ) - - permutation_mnk = ( - self.atom_layout_mnk[0] * self.mma_inst_shape[0], - # if atom layout's N-mode is 1, to leverage the largest coalesced - # shared memory -> register copy, set the tiled mma's N mode to 16 - self.atom_layout_mnk[1] * self.mma_inst_shape[1] * 2, - self.atom_layout_mnk[2] * self.mma_inst_shape[2], - ) - - # Created a tiled mma that tiles the atom according to specified layout. - # For a 2x2x1 atom layout, the mma atom is duplicated 4 times, twice - # across M and twice across N - tC = cute.make_layout(self.atom_layout_mnk) - tiled_mma = cute.make_tiled_mma( - op, - tC, - permutation_mnk=permutation_mnk, - ) - - # grid_dim: ((m + BLK_M - 1) // BLK_M, (n + BLK_N - 1) // BLK_N, l) - grid_dim = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) - - # Add threadblock rasterization to improve re-use of data - raster_factor = 1 - grid_dim_n = cute.size(grid_dim[1]) - # Thresholds picked so that it doesn't cause too many no-op CTAs - if grid_dim_n > 5: - raster_factor = 8 - elif grid_dim_n > 2: - raster_factor = 4 - elif grid_dim_n > 1: - raster_factor = 2 - rasterization_remap_grid_dim = ( - cute.size(grid_dim[0]) * raster_factor, - (cute.size(grid_dim[1]) + raster_factor - 1) // raster_factor, - cute.size(grid_dim[2]), - ) - - self.kernel( - mA, - mB, - mC, - sA_layout, - sB_layout, - sC_layout, - tiled_copy_A, - tiled_copy_B, - tiled_copy_C, - tiled_mma, - raster_factor, - epilogue_op, - ).launch( - grid=rasterization_remap_grid_dim, - block=[self.num_threads, 1, 1], - ) - - @cute.kernel - def kernel( - self, - mA: cute.Tensor, - mB: cute.Tensor, - mC: cute.Tensor, - sA_layout: cute.ComposedLayout, - sB_layout: cute.ComposedLayout, - sC_layout: cute.ComposedLayout, - tiled_copy_A: cute.TiledCopy, - tiled_copy_B: cute.TiledCopy, - tiled_copy_C: cute.TiledCopy, - tiled_mma: cute.TiledMma, - rasterization_factor: cutlass.Int32, - epilogue_op: cutlass.Constexpr = lambda x: x, - ): - # Thread index, block index - tidx, _, _ = cute.arch.thread_idx() - bidx, bidy, bidz = cute.arch.block_idx() - grid_dim = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) - offset_tile_x, offset_tile_y = self.raster_tile( - bidx, bidy, rasterization_factor - ) - # Early exit if CTA is out of range - if grid_dim[0] <= offset_tile_x or grid_dim[1] <= offset_tile_y: - pass - else: - tiler_coord = (offset_tile_x, offset_tile_y, None) - - # /////////////////////////////////////////////////////////////////////////////// - # Get the appropriate tiles for this thread block. - # gA: (BLK_M, BLK_N, k), gB: (BLK_N, BLK_K, k), gC: (BLK_M, BLK_N) - # /////////////////////////////////////////////////////////////////////////////// - gA = cute.local_tile( - mA[None, None, bidz], - tiler=self.cta_tiler, - coord=tiler_coord, - proj=(1, None, 1), - ) - gB = cute.local_tile( - mB[None, None, bidz], - tiler=self.cta_tiler, - coord=tiler_coord, - proj=(None, 1, 1), - ) - gC = cute.local_tile( - mC[None, None, bidz], - tiler=self.cta_tiler, - coord=tiler_coord, - proj=(1, 1, None), - ) - - # By default, if the tensor k mode does not divide into the tile k - # size, then last tiles in the k dimension are irregular. - # Instead, make the first tiles irregular when k is irregular. - # This allows us to handle the irregular tile first to avoid - # checking for this condition within the mainloop. - - # residual_k is a negative number indicating the amount needed to - # shift the pointer by in dimension k - residual_k = cute.size(mA, mode=[1]) - cutlass.Int32(self.bK) * cute.size( - gA, mode=[2] - ) - - # move the pointer of gA/gB in the `-k` direction - gA = cute.domain_offset((0, residual_k, 0), gA) - gB = cute.domain_offset((0, residual_k, 0), gB) - # input is 16B aligned - gA = cute.make_tensor(gA.iterator.align(16), gA.layout) - gB = cute.make_tensor(gB.iterator.align(16), gB.layout) - - # Construct identity layout for sA and sB (mirrors global tensors, - # used for predication only) - mcA = cute.make_identity_tensor(mA.layout.shape) - mcB = cute.make_identity_tensor(mB.layout.shape) - cA = cute.local_tile( - mcA[None, None, bidz], - tiler=self.cta_tiler, - coord=tiler_coord, - proj=(1, None, 1), - ) - cB = cute.local_tile( - mcB[None, None, bidz], - tiler=self.cta_tiler, - coord=tiler_coord, - proj=(None, 1, 1), - ) - - cA = cute.domain_offset((0, residual_k, 0), cA) - cB = cute.domain_offset((0, residual_k, 0), cB) - - # /////////////////////////////////////////////////////////////////////////////// - # Create shared memory buffers and get the appropriate fragments for this thread. - # sA: (BLK_M, BLK_K, PIPE) , sB: (BLK_N, BLK_K, PIPE) - # tAgA: (CPY, CPY_M, CPY_K, k) , tBgB: (CPY, CPY_N, CPY_K, k) - # tAsA: (CPY, CPY_M, CPY_K, PIPE) , tBsB: (CPY, CPY_N, CPY_K, PIPE) - # /////////////////////////////////////////////////////////////////////////////// - @cute.struct - class SharedStorageAB: - a: cute.struct.Align[ - cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)], - 16, - ] - b: cute.struct.Align[ - cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)], - 16, - ] - - @cute.struct - class SharedStorageC: - c: cute.struct.Align[ - cute.struct.MemRange[mC.element_type, cute.cosize(sC_layout)], - 16, - ] - - # Shared memory buffer - smem = cutlass.utils.SmemAllocator() - # Shared memory allocated for operations with A, B will be - # overwritten for operations on C. This is to improve performance - # by reducing the size of shared memory requested by each block - storage = smem.allocate( - max(SharedStorageAB.size_in_bytes(), SharedStorageC.size_in_bytes()), - byte_alignment=16, - ) - sA = SharedStorageAB(storage).a.get_tensor(sA_layout) - sB = SharedStorageAB(storage).b.get_tensor(sB_layout) - sC = SharedStorageC(storage).c.get_tensor(sC_layout) - - thr_copy_A = tiled_copy_A.get_slice(tidx) - thr_copy_B = tiled_copy_B.get_slice(tidx) - thr_copy_C = tiled_copy_C.get_slice(tidx) - tAgA = thr_copy_A.partition_S(gA) - tAsA = thr_copy_A.partition_D(sA) - tBgB = thr_copy_B.partition_S(gB) - tBsB = thr_copy_B.partition_D(sB) - tCsC_epilogue = thr_copy_C.partition_S(sC) - tCgC_epilogue = thr_copy_C.partition_D(gC) - - # Repeat the partitioning with identity layouts - tAcA = thr_copy_A.partition_S(cA) - tBcB = thr_copy_B.partition_S(cB) - - # /////////////////////////////////////////////////////////////////////////////// - # Predicate: Mark indices that need to copy when problem_shape isn't a multiple - # of tile_shape - # /////////////////////////////////////////////////////////////////////////////// - - # For predication over the tensors A (M/K), B (N/K), and (in the - # epilogue) C (M/N), we will compute it in a fashion similar to an - # outer product. The predication along one of the dimensions is - # evaluated and stored in a predication tensor. Then, the - # predication for the remaining dimension is handled later via an - # if/else branch at the copy. - # For A and B, predication booleans along M/N are stored in a - # predication tensor and along K is handled via a if/else branch. - - # Allocate predicate tensors for M and N. Predication is checked - # at the granularity of a copy atom, so the predicate tensor does not - # need separate booleans for individual elements within a copy - # atom (for example, the elements of tAgA.shape[0][0].) - tApA = cute.make_rmem_tensor( - cute.make_layout( - ( - tAgA.shape[0][1], - cute.size(tAgA, mode=[1]), - cute.size(tAgA, mode=[2]), - ), - stride=(cute.size(tAgA, mode=[1]), 1, 0), - ), - cutlass.Boolean, - ) - tBpB = cute.make_rmem_tensor( - cute.make_layout( - ( - tBsB.shape[0][1], - cute.size(tBsB, mode=[1]), - cute.size(tBsB, mode=[2]), - ), - stride=(cute.size(tBsB, mode=[1]), 1, 0), - ), - cutlass.Boolean, - ) - # Set predicates for M/N bounds - for rest_v in range(tApA.shape[0]): - for m in range(tApA.shape[1]): - tApA[rest_v, m, 0] = cute.elem_less( - tAcA[(0, rest_v), m, 0, 0][0], mA.shape[0] - ) - for rest_v in range(tBpB.shape[0]): - for n in range(tBpB.shape[1]): - tBpB[rest_v, n, 0] = cute.elem_less( - tBcB[(0, rest_v), n, 0, 0][0], mB.shape[0] - ) - - # /////////////////////////////////////////////////////////////////////////////// - # Prefetch Prologue - # /////////////////////////////////////////////////////////////////////////////// - # Clear the smem tiles to account for predicated off loads - tAsA.fill(0) - tBsB.fill(0) - cute.arch.sync_threads() - # Start async loads for the first k-tile. Here we take care of the k residue - # via if/else check along the k dimension. Because we shifted the identity tensor - # by the residue_k and because the identity tensor is a coord tensor, the - # values of any identity tensor element that is poison is less than -1 - num_smem_stages = cute.size(tAsA, mode=[3]) - k_tile_count = cute.size(tAgA, mode=[3]) - k_tile_index = cutlass.Int32(0) - - for k in range(tApA.shape[2]): - if cute.elem_less(cutlass.Int32(-1), tAcA[0, 0, k, 0][1]): - cute.copy( - tiled_copy_A, - tAgA[None, None, k, k_tile_index], - tAsA[None, None, k, 0], - pred=tApA[None, None, k], - ) - for k in range(tBpB.shape[2]): - if cute.elem_less(cutlass.Int32(-1), tBcB[0, 0, k, 0][1]): - cute.copy( - tiled_copy_B, - tBgB[None, None, k, k_tile_index], - tBsB[None, None, k, 0], - pred=tBpB[None, None, k], - ) - k_tile_index = k_tile_index + 1 - cute.arch.cp_async_commit_group() - - # Start async loads for rest of the k-tiles - for k_tile in range(1, num_smem_stages - 1): - if k_tile == k_tile_count: - tApA.fill(0) - tBpB.fill(0) - cute.copy( - tiled_copy_A, - tAgA[None, None, None, k_tile_index], - tAsA[None, None, None, k_tile], - pred=tApA, - ) - cute.copy( - tiled_copy_B, - tBgB[None, None, None, k_tile_index], - tBsB[None, None, None, k_tile], - pred=tBpB, - ) - k_tile_index = k_tile_index + 1 - cute.arch.cp_async_commit_group() - - # /////////////////////////////////////////////////////////////////////////////// - # Tile MMA compute thread partitions and allocate accumulators - # /////////////////////////////////////////////////////////////////////////////// - thr_mma = tiled_mma.get_slice(tidx) - tCsA = thr_mma.partition_A(sA) - tCsB = thr_mma.partition_B(sB) - tCsC = thr_mma.partition_C(sC) - tCgC = thr_mma.partition_C(gC) - tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) - tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) - tCrC = tiled_mma.make_fragment_C(tCgC) - # Clear the accumulator - tCrC.fill(0.0) - - # /////////////////////////////////////////////////////////////////////////////// - # Copy Atom A/B retiling - # /////////////////////////////////////////////////////////////////////////////// - - # Create the copy atoms for the copy from shared memory to register - atom_copy_s2r_A = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp( - self.a_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 - ), - mA.element_type, - ) - atom_copy_s2r_B = cute.make_copy_atom( - cute.nvgpu.warp.LdMatrix8x8x16bOp( - self.b_major_mode != utils.LayoutEnum.ROW_MAJOR, 4 - ), - mB.element_type, - ) - - # Creates the tiled copy so that it matches the thread-value layout - # expected by the tiled mma - tiled_copy_s2r_A = cute.make_tiled_copy_A(atom_copy_s2r_A, tiled_mma) - tiled_copy_s2r_B = cute.make_tiled_copy_B(atom_copy_s2r_B, tiled_mma) - - thr_copy_ldmatrix_A = tiled_copy_s2r_A.get_slice(tidx) - thr_copy_ldmatrix_B = tiled_copy_s2r_B.get_slice(tidx) - tCsA_copy_view = thr_copy_ldmatrix_A.partition_S(sA) - tCrA_copy_view = thr_copy_ldmatrix_A.retile(tCrA) - tCsB_copy_view = thr_copy_ldmatrix_B.partition_S(sB) - tCrB_copy_view = thr_copy_ldmatrix_B.retile(tCrB) - - # Current pipe index in smem to read from / write to - smem_pipe_read = 0 - smem_pipe_write = num_smem_stages - 1 - - tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read] - tCsB_p = tCsB_copy_view[None, None, None, smem_pipe_read] - - # /////////////////////////////////////////////////////////////////////////////// - # PREFETCH register pipeline - # /////////////////////////////////////////////////////////////////////////////// - num_k_block = cute.size(tCrA, mode=[2]) - if num_k_block > 1: - # Wait until our first prefetched tile is loaded in - cute.arch.cp_async_wait_group(num_smem_stages - 2) - cute.arch.sync_threads() - # Prefetch the first k-block rmem from the first k-tile - cute.copy( - tiled_copy_s2r_A, - tCsA_p[None, None, 0], - tCrA_copy_view[None, None, 0], - ) - cute.copy( - tiled_copy_s2r_B, - tCsB_p[None, None, 0], - tCrB_copy_view[None, None, 0], - ) - - # /////////////////////////////////////////////////////////////////////////////// - # Mainloop - # 1. Shared memory pipeline (gmem -> smem): - # The default smem pipeline depth is 3, meaning that for shared - # memory buffers, we allocate three times the size described by the - # CTA tiler. We prefetch 2 of these buffers before entering the main - # loop. Considering only the transfer from global memory to shared - # memory, the general structure of the mainloop is: - # (1) copy k-tile from gmem to smem; - # (2) perform gemm computation on k-tile; - # (3) wait for the next copy to finish. - # The `cute.arch.cp_async_wait_group(num_smem_stages - 2)` command - # waits for the number of unfinished 'copy' to be <= 1. The advantage - # of this approach is that it allows for simultaneous production - # (i.e., step (1)) and consumption (i.e., step (2)) of smem. - # A common misconception is to prefetch N buffers and rewrite - # the pipeline logic to wait on N-1 pending copies. The disadvantage - # of this approach is that it requires fully consuming a buffer in - # order to open an empty buffer for the next copy. - # 2. Register pipeline (smem -> register): - # Similarly, the register pipeline produces i+1, consumes i, and - # produces i+2... Notably, i and i+1 do not use the same register, - # eliminating dependencies on the same register for better parallelism. - # 3. Combining the smem and register pipelines results in the mainloop. - # /////////////////////////////////////////////////////////////////////////////// - for k_tile in range(k_tile_count): - for k_block in cutlass.range(num_k_block, unroll_full=True): - if k_block == num_k_block - 1: - tCsA_p = tCsA_copy_view[None, None, None, smem_pipe_read] - tCsB_p = tCsB_copy_view[None, None, None, smem_pipe_read] - cute.arch.cp_async_wait_group(num_smem_stages - 2) - cute.arch.sync_threads() - - # Load A, B from shared memory to registers for k_block + 1 - k_block_next = (k_block + 1) % num_k_block # static - cute.copy( - tiled_copy_s2r_A, - tCsA_p[None, None, k_block_next], - tCrA_copy_view[None, None, k_block_next], - ) - cute.copy( - tiled_copy_s2r_B, - tCsB_p[None, None, k_block_next], - tCrB_copy_view[None, None, k_block_next], - ) - - # Fetch next A: To better interleave global memory access and compute - # instructions, we intentionally use the sequence: copy A, perform GEMM, - # then copy B. - if k_block == 0: - if k_tile + num_smem_stages - 1 < k_tile_count: - cute.copy( - tiled_copy_A, - tAgA[None, None, None, k_tile_index], - tAsA[None, None, None, smem_pipe_write], - pred=tApA, - ) - - # Thread-level register gemm for k_block - cute.gemm( - tiled_mma, - tCrC, - tCrA[None, None, k_block], - tCrB[None, None, k_block], - tCrC, - ) - - # Fetch next B and update smem pipeline read/write - if k_block == 0: - if k_tile + num_smem_stages - 1 < k_tile_count: - cute.copy( - tiled_copy_B, - tBgB[None, None, None, k_tile_index], - tBsB[None, None, None, smem_pipe_write], - pred=tBpB, - ) - k_tile_index = k_tile_index + 1 - cute.arch.cp_async_commit_group() - smem_pipe_write = smem_pipe_read - smem_pipe_read = smem_pipe_read + 1 - if smem_pipe_read == num_smem_stages: - smem_pipe_read = 0 - - # Sync before epilogue - cute.arch.cp_async_wait_group(0) - cute.arch.sync_threads() - - # /////////////////////////////////////////////////////////////////////////////// - # Epilogue with fusion - # /////////////////////////////////////////////////////////////////////////////// - tCrD = cute.make_fragment_like(tCrC, self.c_dtype) - tCrD[None] = epilogue_op(tCrC.load()).to(self.c_dtype) - - # Copy results of D back to shared memory - cute.autovec_copy(tCrD, tCsC) - - # Create coord tensor for C - ceilM, ceilN, _ = cute.ceil_div(mC.shape, (self.bM, self.bN, 1)) - mcC = cute.make_identity_tensor( - ( - cute.size(ceilM) * self.cta_tiler[0], - cute.size(ceilN) * self.cta_tiler[1], - 1, - ) - ) - cC = cute.local_tile( - mcC[None, None, bidz], - tiler=self.cta_tiler, - coord=tiler_coord, - proj=(1, 1, None), - ) - tCcC = thr_copy_C.partition_S(cC) - - tCrC_epilogue = cute.make_fragment_like(tCsC_epilogue) - # Wait for all writes to shared memory to finish before starting copies - # using the new layouts - cute.arch.sync_threads() - cute.autovec_copy(tCsC_epilogue, tCrC_epilogue) - - # Create predication tensor for m - tCpC = cute.make_rmem_tensor( - cute.make_layout( - ( - tCgC_epilogue.shape[0][1], - cute.size(tCgC_epilogue, mode=[1]), - cute.size(tCgC_epilogue, mode=[2]), - ), - stride=(cute.size(tCgC_epilogue, mode=[1]), 1, 0), - ), - cutlass.Boolean, - ) - for rest_v in range(tCpC.shape[0]): - for m in range(tCpC.shape[1]): - tCpC[rest_v, m, 0] = cute.elem_less( - tCcC[(0, rest_v), m, 0][0], mC.shape[0] - ) - - # Copy to global memory using better vectorization - for rest_v in range(tCpC.shape[0]): - for n in range(tCpC.shape[2]): - if cute.elem_less(tCcC[(0, rest_v), 0, n][1], mC.shape[1]): - cute.copy( - tiled_copy_C, - tCrC_epilogue[None, None, n], - tCgC_epilogue[None, None, n], - pred=tCpC[None, None, n], - ) - return - - def _make_smem_layout_AB(self, dtype, major_mode, copy_bits, smem_tiler): - major_mode_size = ( - smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] - ) - major_mode_size = 64 if major_mode_size >= 64 else major_mode_size - - swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) - swizzle_bits = min(swizzle_bits, 3) - - layout_atom_outer = ( - cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) - if major_mode == utils.LayoutEnum.ROW_MAJOR - else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) - ) - layout_atom = cute.make_composed_layout( - cute.make_swizzle(swizzle_bits, 3, 3), - 0, - layout_atom_outer, - ) - layout = cute.tile_to_shape(layout_atom, smem_tiler, (0, 1, 2)) - return layout - - def _make_smem_layout_C(self, dtype, major_mode, copy_bits, smem_tiler): - major_mode_size = ( - smem_tiler[1] if major_mode == utils.LayoutEnum.ROW_MAJOR else smem_tiler[0] - ) - - swizzle_bits = int(math.log2(major_mode_size * dtype.width // copy_bits)) - swizzle_bits = min(swizzle_bits, 3) - - layout_atom_outer = ( - cute.make_layout((8, major_mode_size), stride=(major_mode_size, 1)) - if major_mode == utils.LayoutEnum.ROW_MAJOR - else cute.make_layout((major_mode_size, 8), stride=(1, major_mode_size)) - ) - layout_atom = cute.make_composed_layout( - cute.make_swizzle(swizzle_bits, 3, 4), - 0, - layout_atom_outer, - ) - - # Due to the thread layout of the mma, remove swizzle in C to - # prevent shared memory fragments owned by an single thread from - # holding swizzles - if major_mode == utils.LayoutEnum.COL_MAJOR: - layout_atom = cute.make_composed_layout( - cute.make_swizzle(0, 3, 4), 0, layout_atom_outer - ) - layout = cute.tile_to_shape( - layout_atom, - smem_tiler, - (0, 1), - ) - return layout - - def _make_gmem_tiled_copy_AB(self, atom_copy, dtype, major_mode, copy_bits): - copy_elems = copy_bits // dtype.width - shape_dim_1 = cute.size(self.bK) // copy_elems - # thread layout for copy - thread_layout = cute.make_layout( - (self.num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) - ) - if major_mode != utils.LayoutEnum.ROW_MAJOR: - shape_dim_0 = cute.size(self.bM) // copy_elems - thread_layout = cute.make_layout( - (shape_dim_0, self.num_threads // shape_dim_0), stride=(1, shape_dim_0) - ) - # Value layout for copy - value_layout = ( - cute.make_layout((1, copy_elems)) - if major_mode == utils.LayoutEnum.ROW_MAJOR - else cute.make_layout((copy_elems, 1)) - ) - return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) - - def _make_gmem_tiled_copy_C(self, atom_copy, dtype, major_mode, copy_bits): - copy_elems = copy_bits // dtype.width - shape_dim_1 = cute.size(self.bN) // copy_elems - # thread layout for copy - thread_layout = cute.make_layout( - (self.num_threads // shape_dim_1, shape_dim_1), stride=(shape_dim_1, 1) - ) - if major_mode != utils.LayoutEnum.ROW_MAJOR: - shape_dim_0 = cute.size(self.bM) // copy_elems - thread_layout = cute.make_layout( - (shape_dim_0, self.num_threads // shape_dim_0), stride=(1, shape_dim_0) - ) - value_layout = ( - cute.make_layout((1, copy_elems)) - if major_mode == utils.LayoutEnum.ROW_MAJOR - else cute.make_layout((copy_elems, 1)) - ) - return cute.make_tiled_copy_tv(atom_copy, thread_layout, value_layout) - - def raster_tile(self, i, j, f): - new_i = i // f - new_j = (i % f) + (j * f) - return (new_i, new_j) - - -# Need a custum call function for the optimized kernel because of the tensor layout -def call_cute_custom(kernel_function, args, kwargs, grid, threads, params): - # Initialize cache if it does not exist - if not hasattr(call_cute_custom, "custom_cache"): - call_cute_custom.custom_cache = {} - - # Convert Torch tensors to CuTe tensors with correct layout - cute_args = [] - for arg in args: - if isinstance(arg, torch.Tensor): - cute_tensor = ( - from_dlpack(arg, assumed_align=16) - .mark_layout_dynamic(leading_dim=1) - .mark_compact_shape_dynamic( - mode=1, - stride_order=(2, 0, 1), - divisibility=8, - ) - ) - cute_args.append(cute_tensor) - else: - cute_args.append(arg) - - # Form cache key from tuning parameters - param_keys = sorted(params.keys()) - cache_str = type(kernel_function).__name__ - for k in param_keys: - cache_str += "_" + str(params[k]) - - # Check if kernel exists in cache. Otherwise, compile and save - if cache_str in call_cute_custom.custom_cache: - compiled_kernel = call_cute_custom.custom_cache[cache_str] - else: - # Wrap in try-except because CuTe gives TypeError for certain block sizes. This - # stops the tuning process and we do not want this - try: - compiled_kernel = cute.compile(kernel_function, *cute_args) - call_cute_custom.custom_cache[cache_str] = compiled_kernel - except Exception as e: - raise RuntimeError(f"Invalid configuration, CuTe compile failed with {type(e).__name__}") from e - - - - compiled_kernel(*cute_args) - - - -def run_optimized(M, N, K, L=1): - a = torch.randn((L, M, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) - b = torch.randn((L, N, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) # b is transposed - c = torch.empty((L, M, N), device="cuda", dtype=torch.float16).permute(1, 2, 0) - - c_ref = torch.matmul( - a.permute(2, 0, 1), - b.permute(2, 1, 0) - ).permute(1, 2, 0) - - def create_cute_tensor(torch_tensor): - cute_tensor = ( - from_dlpack(torch_tensor, assumed_align=16) - .mark_layout_dynamic(leading_dim=1) - .mark_compact_shape_dynamic( - mode=1, - stride_order=(2, 0, 1), - divisibility=8, - ) - ) - return cute_tensor - - a_, b_, c_ = create_cute_tensor(a), create_cute_tensor(b), create_cute_tensor(c) - - gemm = TensorOpGemm() - compiled_kernel = cute.compile(gemm, a_, b_, c_) - compiled_kernel(a_, b_, c_) - - assert torch.allclose(c_ref, c, atol= M * 2**(-11), rtol=1e-2) - print("Sucess") - - -def tune_optimized(M, N, K, L=1): - a = torch.randn((L, M, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) - b = torch.randn((L, N, K), device="cuda", dtype=torch.float16).permute(1, 2, 0) # b is transposed - c = torch.empty((L, M, N), device="cuda", dtype=torch.float16).permute(1, 2, 0) - - c_ref = torch.matmul( - a.permute(2, 0, 1), - b.permute(2, 1, 0) - ).permute(1, 2, 0) - - - args = [a, b, c] - tune_params = { - "bM":[16, 32, 64, 128, 256], - "bN":[32, 64, 128, 256], - "bK": [16, 32, 64, 128], # must be divisable by 16 - "num_stages": [3, 4, 5], # restricted to >= 3 by the kernel - "atom_lay_M": [1, 2, 4], - "atom_lay_N": [1, 2], - } - - restrictions = [ - "bM % (atom_lay_M * 16) == 0", "bN % (atom_lay_N * 8) == 0", # layout - "atom_lay_M * atom_lay_N * 32 <= 1024", # number of threads - "num_stages * (bK * (bM + bN)) * 2 <= 49152", # SMEM - "bM >= atom_lay_M * 32", # ensure each atom sees at least 2 MMA tiles per dimension - "bN >= atom_lay_N * 16", # ensure each atom sees at least 2 MMA tiles per dimension - # the swizzled smem layout of C only supports these tiles, others fail in autovec_copy - "bN >= 128", "bN <= 128 * atom_lay_N", - ] - - - - results, env = tune_kernel("TensorOpGemm", __file__, M * N, args, tune_params, verbose=True, restrictions=restrictions, #strategy="bayes_opt", - call_function=call_cute_custom, answer=[None, None, c_ref.cpu()], atol=M * 2**(-10)) - - - if __name__ == "__main__": m, n, k = 4096, 4096, 4096 run_naive_matmul(m, n, k) tune_naive_matmul(m, n, k) - - mnkl = (4096, 4096, 4096, 1) - tune_optimized(m, n, k) - run_optimized(m, n, k) diff --git a/examples/python/matmul/cutile_matmul.py b/examples/python/matmul/cutile_matmul.py index 19d3b54d7..9ae9ef3e0 100644 --- a/examples/python/matmul/cutile_matmul.py +++ b/examples/python/matmul/cutile_matmul.py @@ -4,7 +4,6 @@ from kernel_tuner import tune_kernel - @ct.kernel def matmul(A, B, C, tile_m: ct.Constant[int], tile_n: ct.Constant[int], tile_k: ct.Constant[int]): bidx = ct.bid(0) @@ -20,29 +19,6 @@ def matmul(A, B, C, tile_m: ct.Constant[int], tile_n: ct.Constant[int], tile_k: ct.store(C, index=(bidx, bidy), tile=ct.astype(acc, C.dtype)) -@ct.kernel -def matmul_swizzled(A, B, C, tile_m: ct.Constant[int], tile_n: ct.Constant[int], tile_k: ct.Constant[int], - group_size_m: ct.Constant[int]): - # map the 1D block index to a tile of C in groups of group_size_m rows of tiles, - # so consecutive blocks reuse the same tiles of B, which improves L2 cache reuse - bid = ct.bid(0) - num_bid_m = ct.cdiv(A.shape[0], tile_m) - num_bid_n = ct.cdiv(B.shape[1], tile_n) - num_bid_in_group = group_size_m * num_bid_n - first_bid_m = (bid // num_bid_in_group) * group_size_m - group_size = ct.minimum(num_bid_m - first_bid_m, group_size_m) - bidx = first_bid_m + (bid % num_bid_in_group) % group_size - bidy = (bid % num_bid_in_group) // group_size - - acc = ct.full((tile_m, tile_n), 0, dtype=ct.float32) - for k in range(ct.num_tiles(A, axis=1, shape=(tile_m, tile_k))): - a = ct.load(A, index=(bidx, k), shape=(tile_m, tile_k), padding_mode=ct.PaddingMode.ZERO) - b = ct.load(B, index=(k, bidy), shape=(tile_k, tile_n), padding_mode=ct.PaddingMode.ZERO) - acc = ct.mma(a, b, acc) - - ct.store(C, index=(bidx, bidy), tile=ct.astype(acc, C.dtype)) - - def run(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -74,28 +50,7 @@ def tune(M, N, K): atol=M * 2**(-11), block_size_names=["tile_m", "tile_n"]) -def tune_swizzled(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.zeros(M, N, device="cuda", dtype=torch.float16) - - args = [A, B, C] - answer = [None, None, (A @ B).cpu()] - - tune_params = dict() - tune_params["tile_m"] = [32, 64, 128, 256] - tune_params["tile_n"] = [32, 64, 128, 256] - tune_params["tile_k"] = [16, 32, 64, 128] - tune_params["group_size_m"] = [1, 2, 4, 8, 16] - - # 1D grid with one block per tile of C - results, env = tune_kernel("matmul_swizzled", __file__, M * N, args, tune_params, answer=answer, - atol=M * 2**(-11), block_size_names=["tile_m"], grid_div_x=["tile_m", "tile_n"], - strategy="bayes_opt", strategy_options={"max_fevals": 100}) - - if __name__ == "__main__": M, N, K = 4096, 4096, 4096 run(M, N, K) tune(M, N, K) - tune_swizzled(M, N, K) diff --git a/examples/python/matmul/numba_matmul.py b/examples/python/matmul/numba_matmul.py index 345fd7b1f..d742b1224 100644 --- a/examples/python/matmul/numba_matmul.py +++ b/examples/python/matmul/numba_matmul.py @@ -1,10 +1,9 @@ import numpy as np -from numba import cuda, float32 +from numba import cuda from kernel_tuner import tune_kernel - # Source: https://nvidia.github.io/numba-cuda/user/examples.html#matrix-multiplication @cuda.jit def matmul(A, B, C): @@ -16,96 +15,12 @@ def matmul(A, B, C): C[i, j] = tmp -# Translated to Numba-CUDA from https://github.com/cupy/cupy/blob/main/examples/gemm/sgemm.cu -@cuda.jit -def optimized_matmul(M, N, K, A, B, C): - DIM_X = 16 - DIM_Y = 16 - BLK_M = 64 - BLK_N = 64 - BLK_K = 16 - THR_M = 4 # Should be equal to BLK_M / DIM_X - THR_N = 4 # Should be equal to BLK_N / DIM_Y - - # thread indices - idx = cuda.threadIdx.x - idy = cuda.threadIdx.y - - blx = cuda.blockIdx.x - bly = cuda.blockIdx.y - - # shared memory - sA = cuda.shared.array((BLK_K, BLK_M), np.float16) - sB = cuda.shared.array((BLK_N, BLK_K), np.float16) - - # registers - rC = cuda.local.array((THR_N, THR_M), float32) - rA = cuda.local.array(THR_M, float32) - rB = cuda.local.array(THR_N, float32) - - # init accumulator - for n in range(THR_N): - for m in range(THR_M): - rC[n][m] = 0.0 - - # global indices - base_row = blx * BLK_M - base_col = bly * BLK_N - - # loop over K tiles - for kk in range(0, K, BLK_K): - - # load A tile into shared memory - for i in range(idy, BLK_K, DIM_Y): - for j in range(idx, BLK_M, DIM_X): - row = base_row + j - col = kk + i - if row < M and col < K: - sA[i, j] = A[row, col] - else: - sA[i, j] = 0.0 - - # load B tile - for i in range(idy, BLK_N, DIM_Y): - for j in range(idx, BLK_K, DIM_X): - row = kk + j - col = base_col + i - if row < K and col < N: - sB[i, j] = B[row, col] - else: - sB[i, j] = 0.0 - - cuda.syncthreads() - - # compute - for k in range(BLK_K): - for m in range(THR_M): - rA[m] = float32(sA[k, m * DIM_X + idx]) - - for n in range(THR_N): - rB[n] = float32(sB[n * DIM_Y + idy, k]) - - for n in range(THR_N): - for m in range(THR_M): - rC[n][m] += rA[m] * rB[n] - - cuda.syncthreads() - - # write back - for n in range(THR_N): - col = base_col + n * DIM_Y + idy - for m in range(THR_M): - row = base_row + m * DIM_X + idx - if row < M and col < N: - C[row, col] = np.float16(rC[n][m]) - - def run_basic(M, N, K): # create numpy arrays A = np.random.rand(M, K).astype(np.float16) B = np.random.rand(K, N).astype(np.float16) C = np.zeros((M, N), dtype=np.float16) - C_ref = A @ B + C_ref = (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) # copy to GPU A_d = cuda.to_device(A) @@ -133,38 +48,6 @@ def run_basic(M, N, K): print("Succes") -def run_optimized(M, N, K): - # inputs - A = np.random.rand(M, K).astype(np.float16) - B = np.random.rand(K, N).astype(np.float16) - C = np.zeros((M, N), dtype=np.float16) - C_ref = A @ B - - # move to GPU - dA = cuda.to_device(A) - dB = cuda.to_device(B) - dC = cuda.to_device(C) - - threads_per_block = (16, 16) - blocks_per_grid = ( - (M + 63) // 64, # BLK_M = 64 - (N + 63) // 64, # BLK_N = 64 - ) - - # launch - optimized_matmul[blocks_per_grid, threads_per_block](M, N, K, dA, dB, dC) - cuda.synchronize() - - # copy result back - C_result = dC.copy_to_host() - - # check - np.testing.assert_allclose(C_result, C_ref, rtol=1e-2, atol=M * 2**(-11)) - print("Succes") - - - - def tune_basic(M, N, K): # create inputs as normal, but do not copy to device A = np.random.rand(M, K).astype(np.float16) @@ -179,7 +62,7 @@ def tune_basic(M, N, K): restrictions = ["block_size_x * block_size_y <= 1024"] - answer = [None, None, A @ B] + answer = [None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16)] atol = M * 2**(-11) results, env = tune_kernel( @@ -194,63 +77,8 @@ def tune_basic(M, N, K): ) -def tune_optimized(M, N, K): - # create inputs as normal, but do not copy to device - A = np.random.rand(M, K).astype(np.float16) - B = np.random.rand(K, N).astype(np.float16) - C = np.zeros((M, N), dtype=np.float16) - - size = (M, N) - args = [M, N, K, A, B, C] - tune_params = { - "DIM_X": [2**i for i in range(1, 10)], - "DIM_Y": [2**i for i in range(1, 10)], - "BLK_M": [2**i for i in range(1, 10)], - "BLK_N": [2**i for i in range(1, 10)], - "BLK_K": [2**i for i in range(1, 10)], - "THR_M": [2**i for i in range(1, 9)], # Restricted to BLK_M / DIM_X - "THR_N": [2**i for i in range(1, 9)], # Restricted to BLK_N / DIM_Y - } - - answer = [None, None, None, None, None, A @ B] - atol = M * 2**(-11) - - restrictions = [ - "BLK_M % DIM_X == 0", - "BLK_N % DIM_Y == 0", - "THR_M == BLK_M / DIM_X", - "THR_N == BLK_N / DIM_Y", - "DIM_X * DIM_Y <= 1024", - "DIM_X * DIM_Y >= 32", - ] - - results, env = tune_kernel( - kernel_name="optimized_matmul", - kernel_source=__file__, - problem_size=size, - arguments=args, - tune_params=tune_params, - answer=answer, - atol=atol, - restrictions = restrictions, - block_size_names = ["DIM_X", "DIM_Y"], - grid_div_x = ["BLK_M"], - grid_div_y = ["BLK_N"], - strategy = "bayes_opt", - strategy_options = {"max_fevals": 100}, - ) - - - if __name__ == "__main__": - M, N, K = 1024, 1024, 1024 + M, N, K = 4096, 4096, 4096 run_basic(M, N, K) tune_basic(M, N, K) - - run_optimized(M, N, K) - tune_optimized(M, N, K) - - - - \ No newline at end of file diff --git a/examples/python/matmul/taichi_matmul.py b/examples/python/matmul/taichi_matmul.py index df19a52a7..3568eb00f 100644 --- a/examples/python/matmul/taichi_matmul.py +++ b/examples/python/matmul/taichi_matmul.py @@ -23,7 +23,7 @@ def run(M, N, K): A = np.random.rand(M, K).astype(np.float16) B = np.random.rand(K, N).astype(np.float16) C = np.zeros((M, N), dtype=np.float16) - C_ref = A @ B + C_ref = (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) matmul(A, B, C) @@ -40,7 +40,7 @@ def tune(M, N, K): args = [A, B, C] tune_params = {"BLOCK_DIM": [2**i for i in range(5, 11)]} - answer = [None, None, A @ B] + answer = [None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16)] results, env = tune_kernel( kernel_name="matmul", @@ -54,8 +54,8 @@ def tune(M, N, K): ) if __name__ == "__main__": - run(1024, 1024, 1024) - tune(1024, 1024, 1024) + run(4096, 4096, 4096) + tune(4096, 4096, 4096) diff --git a/examples/python/matmul/tilelang_matmul.py b/examples/python/matmul/tilelang_matmul.py index a4b08f6b6..740b27c1e 100644 --- a/examples/python/matmul/tilelang_matmul.py +++ b/examples/python/matmul/tilelang_matmul.py @@ -2,8 +2,6 @@ import tilelang import tilelang.language as T -import itertools - from kernel_tuner import tune_kernel @@ -39,153 +37,6 @@ def gemm( return gemm - -# This kernel is copied from the TileLang project: -# https://github.com/tile-ai/tilelang/tree/main/examples/gemm/example_gemm.py -# -# Original example: example_gemm.py -# Copyright (c) the TileLang authors -# -# Modifications in file: -# - num_threads and num_stages added as metaparameters -# - Removed annotated memory layout for tiles A and B -# - dummy parameter to trigger fresh compilation when timing repeated tuning -@tilelang.jit -def matmul_opt(M, N, K, block_M, block_N, block_K, num_threads=128, num_stages=3, dummy=0, dtype=T.float16, accum_dtype=T.float): - @T.prim_func - def main( - A: T.Tensor((M, K), dtype), - B: T.Tensor((K, N), dtype), - C: T.Tensor((M, N), dtype), - ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): - # Allocate shared and local fragments - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - - # Enable swizzle-based rasterization for better L2 locality - panel_size = 10 - T.use_swizzle(panel_size=panel_size, enable=True) - - # Clear the local accumulation buffer - T.clear(C_local) - - # Pipelined iteration over K dimension - for idx in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - # Copy tile of A - T.copy(A[by * block_M, idx * block_K], A_shared) - - # Parallel copy tile of B - for ko, j in T.Parallel(block_K, block_N): - B_shared[ko, j] = B[idx * block_K + ko, bx * block_N + j] - - # Perform local GEMM on the shared-memory tiles - T.gemm(A_shared, B_shared, C_local) - - # Copy the result tile back - T.copy(C_local, C[by * block_M, bx * block_N]) - - return main - - - - -def get_configs(kt=False): - """ - Generate a list of kernel tuning configuration dictionaries for a tiled matrix-multiply. - This function is used for tuning experiments with the built-in tuner - - Returns: - List[dict]: A list of configuration dictionaries - """ - - block_M = [32, 64, 128, 256] - block_N = [32, 64, 128, 256] - block_K = [16, 32, 64, 128] - num_threads = [256]#[64, 128, 256, 512] - num_stages = [3]#[0, 1, 2, 3, 4] - panel_size = [10]#[4, 6, 8, 10] - - - _configs = list( - itertools.product( - block_M, - block_N, - block_K, - num_threads, - num_stages, - panel_size, - ) - ) - - if kt: - configs = { - "block_M": block_M, - "block_N": block_N, - "block_K": block_K, - "num_stages": num_stages, - "num_threads": num_threads, - "panel_size": panel_size, - } - else: - configs = [ - { - "block_M": c[0], - "block_N": c[1], - "block_K": c[2], - "num_threads": c[3], - "num_stages": c[4], - "panel_size": c[5], - } - for c in _configs - ] - return configs - - - -# For autotuning experiment -# https://github.com/tile-ai/tilelang/blob/main/examples/gemm/example_gemm_autotune.py -@tilelang.autotune(configs=get_configs(), warmup=1, rep=32) -@tilelang.jit -def matmul_opt_autotune(M, N, K, block_M, block_N, block_K, num_threads=128, num_stages=3, panel_size=10, dummy=0, dtype=T.float16, accum_dtype=T.float): - @T.prim_func - def main( - A: T.Tensor((M, K), dtype), - B: T.Tensor((K, N), dtype), - C: T.Tensor((M, N), dtype), - ): - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=num_threads) as (bx, by): - # Allocate shared and local fragments - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - - # Enable swizzle-based rasterization for better L2 locality - T.use_swizzle(panel_size=panel_size, enable=True) - - # Clear the local accumulation buffer - T.clear(C_local) - - # Pipelined iteration over K dimension - for idx in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages): - # Copy tile of A - T.copy(A[by * block_M, idx * block_K], A_shared) - - # Parallel copy tile of B - for ko, j in T.Parallel(block_K, block_N): - B_shared[ko, j] = B[idx * block_K + ko, bx * block_N + j] - - # Perform local GEMM on the shared-memory tiles - T.gemm(A_shared, B_shared, C_local) - - # Copy the result tile back - T.copy(C_local, C[by * block_M, bx * block_N]) - - return main - - - def run_basic(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -205,28 +56,6 @@ def run_basic(M, N, K): print("Succes") -def run_opt(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.zeros(M, N, device="cuda", dtype=torch.float16) - C_ref = A @ B - - kernel = matmul_opt( - M, N, K, - dummy=0, - block_M=64, - block_N=64, - block_K=32, - num_stages=3, - num_threads=128, - ) - - kernel(A, B, C) - - assert torch.allclose(C, C_ref, atol=M * 2 **(-11), rtol=1e-2) - print("Succes") - - def tune_basic(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -267,54 +96,7 @@ def tune_basic(M, N, K): ) -def tune_opt(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.empty((M, N), device='cuda', dtype=torch.float16) - C_ref = A @ B - - size = (M, N) - - args = [A, B, C] - - tune_params = dict() - tune_params["block_M"] = [2**i for i in range(5, 9)] - tune_params["block_N"] = [2**i for i in range(5, 9)] - tune_params["block_K"] = [2**i for i in range(4, 8)] - tune_params["num_threads"] = [64, 128, 256, 512] - tune_params["num_stages"] = [0, 1, 2, 3, 4] - tune_params["panel_size"] = [4, 6, 8, 10] - tune_params["M"] = [M] - tune_params["N"] = [N] - tune_params["K"] = [K] - - restrictions = [ - "2 * num_stages * block_K * (block_M + block_N) <= 49152" - ] - - results, env = tune_kernel( - kernel_name="matmul_opt", - kernel_source=__file__, - problem_size=size, - arguments=args, - tune_params=tune_params, - answer=[None, None, C_ref.cpu()], - atol=M * 2**(-11), - strategy = "bayes_opt", - strategy_options = {"max_fevals": 100}, - verbose=True, - restrictions=restrictions, - ) - - - - if __name__ == "__main__": M, N, K = 4096, 4096, 4096 run_basic(M, N, K) - run_opt(M, N, K) - tune_basic(M, N, K) - tune_opt(M, N, K) - - diff --git a/examples/python/matmul/tilus_matmul.py b/examples/python/matmul/tilus_matmul.py index 798e973e3..657690a4e 100644 --- a/examples/python/matmul/tilus_matmul.py +++ b/examples/python/matmul/tilus_matmul.py @@ -7,7 +7,6 @@ from kernel_tuner import tune_kernel - # This kernel is copied from the Tilus project: # https://github.com/NVIDIA/tilus/blob/main/examples/matmul/matmul_v0.py # @@ -69,90 +68,6 @@ def __call__( self.store_global(gc, acc_f16, offsets=[offset_m, offset_n]) - - -# This kernel is copied from the Tilus project: -# https://nvidia.github.io/tilus/stable/tutorials/matmul-ampere/matmul/matmul_v4.html -# -# Original example: matmul_v4.py -# Copyright (c) the Tilus authors -# -# Modifications in file: -# - Removed auto-tuning decorators -# - Added default values (None) for the __init__ parameters -class MatmulOpt(tilus.Script): - def __init__(self, num_warps=None, block_m=None, block_n=None, block_k=None, num_stages=None): - super().__init__() - self.block_m = block_m - self.block_n = block_n - self.block_k = block_k - self.num_warps = num_warps - self.num_stages = num_stages - - def __call__( - self, - m_size: int32, n_size: int, k_size: int, - a_ptr: ~float16, b_ptr: ~float16, c_ptr: ~float16, - ): - self.attrs.blocks = [cdiv(m_size, self.block_m), cdiv(n_size, self.block_n)] - self.attrs.warps = self.num_warps - - block_m, block_n, block_k = self.block_m, self.block_n, self.block_k - offset_m: int32 = block_m * self.blockIdx.x - offset_n: int32 = block_n * self.blockIdx.y - - ga = self.global_view(a_ptr, dtype=float16, shape=[m_size, k_size]) - gb = self.global_view(b_ptr, dtype=float16, shape=[k_size, n_size]) - sa = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_m, block_k]) - sb = self.shared_tensor(dtype=float16, shape=[self.num_stages, block_k, block_n]) - acc = self.register_tensor(dtype=float32, shape=[block_m, block_n], init=0.0) - - for stage in range(self.num_stages - 1): - offset_k = stage * self.block_k - self.copy_async(src=ga, dst=sa[stage], offsets=[offset_m, offset_k]) - self.copy_async(src=gb, dst=sb[stage], offsets=[offset_k, offset_n]) - self.copy_async_commit_group() - - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - current_stage: int32 = 0 - preload_stage: int32 = self.num_stages - 1 - for offset_k in self.range(0, k_size, block_k, unroll=self.num_stages): - # computation for current tile - a = self.load_shared(sa[current_stage]) - b = self.load_shared(sb[current_stage]) - self.dot(a, b, acc, out=acc) - - # preload the next tile of A and B into shared memory - preload_offset_k = offset_k + (self.num_stages - 1) * block_k - self.copy_async( - src=ga, - dst=sa[preload_stage], - offsets=[offset_m, preload_offset_k], - ) - self.copy_async( - src=gb, - dst=sb[preload_stage], - offsets=[preload_offset_k, offset_n], - ) - self.copy_async_commit_group() - - # update the stage - current_stage = (current_stage + 1) % self.num_stages - preload_stage = (preload_stage + 1) % self.num_stages - self.copy_async_wait_group(n=self.num_stages - 2) - self.sync() - - self.free_shared(sa) - self.free_shared(sb) - - casted_acc = self.cast(acc, dtype=float16) - gc = self.global_view(c_ptr, dtype=float16, shape=[m_size, n_size]) - self.store_global(gc, casted_acc, offsets=[offset_m, offset_n]) - - - def run_basic(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -168,22 +83,6 @@ def run_basic(M, N, K): print("Succes") -def run_optimized(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.empty((M, N), device='cuda', dtype=torch.float16) - C_ref = A @ B - - matmul = MatmulOpt(num_warps=4, block_m=64, block_n=64, block_k=16, num_stages=3) - torch.cuda.synchronize() - - matmul(M, N, K, A, B, C) - torch.cuda.synchronize() - - torch.testing.assert_close(C_ref, C, atol=M * 2**(-11), rtol=1e-2) - print("Succes") - - def tune_basic(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -214,46 +113,7 @@ def tune_basic(M, N, K): ) -def tune_opt(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.empty((M, N), device='cuda', dtype=torch.float16) - C_ref = A @ B - - size = (M, N) - - args = [M, N, K, A, B, C] - - tune_params = dict() - tune_params["block_m"] = [2**i for i in range(5, 9)] - tune_params["block_n"] = [2**i for i in range(5, 9)] - tune_params["block_k"] = [2**i for i in range(4, 8)] - tune_params["num_warps"] = [2, 4, 8, 16] - tune_params["num_stages"] = [2, 3, 4, 5] - - # Shared memory restriction - restrictions = ["2 * num_stages * block_k * (block_m + block_n) <= 49152"] - - results, env = tune_kernel( - kernel_name="MatmulOpt", - kernel_source=__file__, - problem_size=size, - arguments=args, - tune_params=tune_params, - answer=[None, None, None, None, None, C_ref.cpu()], - atol=M * 2**(-11), - restrictions=restrictions, - strategy = "bayes_opt", - strategy_options = {"max_fevals": 200}, - ) - - - - - if __name__ == "__main__": M, N, K = 4096, 4096, 4096 run_basic(M, N, K) - run_optimized(M, N, K) tune_basic(M, N, K) - tune_opt(M, N, K) \ No newline at end of file diff --git a/examples/python/matmul/triton_matmul.py b/examples/python/matmul/triton_matmul.py index 429735855..e43fb7c6f 100644 --- a/examples/python/matmul/triton_matmul.py +++ b/examples/python/matmul/triton_matmul.py @@ -5,8 +5,6 @@ from kernel_tuner import tune_kernel - - @triton.jit def matmul_basic( A_ptr, B_ptr, C_ptr, @@ -47,103 +45,6 @@ def matmul_basic( tl.store(c_ptrs, acc, mask=mask_c) - -# This kernel is copied from the Triton project: -# https://github.com/triton-lang/triton/blob/main/python/tutorials/03-matrix-multiplication.py -# -# Original example: 03-matrix-multiplication.py -# Copyright (c) the Triton authors -# -# Modifications in file: -# - Removed auto-tuning decorators -# - Removed activation function -@triton.jit -def matmul_opt( - # Pointers to matrices - a_ptr, b_ptr, c_ptr, - # Matrix dimensions - M, N, K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. `stride_am` is how much to increase `a_ptr` - # by to get the element one row down (A has M rows). - stride_am, stride_ak, # - stride_bk, stride_bn, # - stride_cm, stride_cn, - # Meta-parameters - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, # - GROUP_SIZE_M: tl.constexpr, # -): - """Kernel for computing the matmul C = A x B. - A has shape (M, K), B has shape (K, N) and C has shape (M, N) - """ - # ----------------------------------------------------------- - # Map program ids `pid` to the block of C it should compute. - # This is done in a grouped ordering to promote L2 data reuse. - # See above `L2 Cache Optimizations` section for details. - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - # ----------------------------------------------------------- - # Add some integer bound assumptions. - # This helps to guide integer analysis in the backend to optimize - # load/store offset address calculation - tl.assume(pid_m >= 0) - tl.assume(pid_n >= 0) - tl.assume(stride_am > 0) - tl.assume(stride_ak > 0) - tl.assume(stride_bn > 0) - tl.assume(stride_bk > 0) - tl.assume(stride_cm > 0) - tl.assume(stride_cn > 0) - - # ---------------------------------------------------------- - # Create pointers for the first blocks of A and B. - # We will advance this pointer as we move in the K direction - # and accumulate - # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers - # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers - # See above `Pointer Arithmetic` section for details - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - # ----------------------------------------------------------- - # Iterate to compute a block of the C matrix. - # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block - # of fp32 values for higher accuracy. - # `accumulator` will be converted back to fp16 after the loop. - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - # Load the next block of A and B, generate a mask by checking the K dimension. - # If it is out of bounds, set it to 0. - a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) - # We accumulate along the K dimension. - accumulator = tl.dot(a, b, accumulator) - # Advance the ptrs to the next K block. - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - c = accumulator.to(tl.float16) - - # ----------------------------------------------------------- - # Write back the block of the output matrix C with masks. - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=c_mask) - - - def run_basic(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -172,36 +73,6 @@ def run_basic(M, N, K): print("Passed") -def run_opt(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.empty((M, N), device='cuda', dtype=torch.float16) - C_ref = A @ B - - BLOCK_SIZE_M = 64 - BLOCK_SIZE_N = 64 - BLOCK_SIZE_K = 32 - GROUP_SIZE_M = 32 - - grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),) - - matmul_opt[grid]( - A, B, C, - M, N, K, - A.stride(0), A.stride(1), - B.stride(0), B.stride(1), - C.stride(0), C.stride(1), - BLOCK_SIZE_M=BLOCK_SIZE_M, - BLOCK_SIZE_N=BLOCK_SIZE_N, - BLOCK_SIZE_K=BLOCK_SIZE_K, - GROUP_SIZE_M=GROUP_SIZE_M, - ) - - assert torch.allclose(C, C_ref, rtol=1e-2, atol= M * 2**(-11)) - - print("Passed") - - def tune_basic(M, N, K): A = torch.randn(M, K, device="cuda", dtype=torch.float16) B = torch.randn(K, N, device="cuda", dtype=torch.float16) @@ -235,53 +106,10 @@ def tune_basic(M, N, K): ) -def tune_opt(M, N, K): - A = torch.randn(M, K, device="cuda", dtype=torch.float16) - B = torch.randn(K, N, device="cuda", dtype=torch.float16) - C = torch.empty((M, N), device='cuda', dtype=torch.float16) - C_ref = A @ B - - size = M * N - - args = [A, B, C, - M, N, K, - A.stride(0), A.stride(1), - B.stride(0), B.stride(1), - C.stride(0), C.stride(1), - ] - - tune_params = dict() - tune_params["BLOCK_SIZE_M"] = [2**i for i in range(5, 9)] - tune_params["BLOCK_SIZE_N"] = [2**i for i in range(5, 9)] - tune_params["BLOCK_SIZE_K"] = [2**i for i in range(4, 8)] - tune_params["GROUP_SIZE_M"] = [4, 6, 8, 10] - tune_params["num_warps"] = [2, 4, 8, 16] - tune_params["num_stages"] = [1, 2, 3, 4, 5] - tune_params["M_dim"] = [M] - - restrictions = ["GROUP_SIZE_M * BLOCK_SIZE_M < M_dim"] # Otherwise grouped ordering has no effect - - results, env = tune_kernel( - kernel_name="matmul_opt", - kernel_source=__file__, - problem_size=size, - arguments=args, - tune_params=tune_params, - answer=[None, None, C_ref.cpu(), None, None, None, None, None, None, None, None, None], - atol=M * 2**(-11), - block_size_names = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], - grid_div_x = ["BLOCK_SIZE_M", "BLOCK_SIZE_N"], - restrictions = restrictions, - strategy = "bayes_opt", - strategy_options = {"max_fevals": 100}, - ) if __name__ == "__main__": M, N, K = 4096, 4096, 4096 run_basic(M, N, K) - #run_opt(M, N, K) - tune_basic(M, N, K) - #tune_opt(M, N, K) diff --git a/examples/python/matmul/warp_matmul.py b/examples/python/matmul/warp_matmul.py index ef10f6b1b..660de9727 100644 --- a/examples/python/matmul/warp_matmul.py +++ b/examples/python/matmul/warp_matmul.py @@ -41,7 +41,7 @@ def run_gemm(M, N, K): wp.launch(gemm, dim=(M, N), inputs=[A_wp, B_wp, C_wp]) - np.testing.assert_allclose(C_wp.numpy(), A @ B, rtol=1e-2, atol=M * 2**(-11)) + np.testing.assert_allclose(C_wp.numpy(), (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16), rtol=1e-2, atol=M * 2**(-11)) print("Succes") @@ -59,7 +59,7 @@ def tune(M, N, K): tune_params["dim"] = [size] args = [A, B, C] - answer = [None, None, A @ B] + answer = [None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16)] results, env = tune_kernel( kernel_name="gemm", @@ -68,12 +68,13 @@ def tune(M, N, K): arguments=args, tune_params=tune_params, answer=answer, + atol=M * 2**(-11), ) if __name__ == "__main__": - M, N, K = 1024, 1024, 1024 + M, N, K = 4096, 4096, 4096 run_gemm(M, N, K) tune(M, N, K) From a7dac17d501e230c75554097e43f53011c9acc5f Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Tue, 6 Oct 2026 12:11:39 +0200 Subject: [PATCH 36/44] add autotune decorator for kernels in Python DSLs kernel_tuner.autotune tunes a kernel written in a Python DSL inside an application, in the spirit of triton.autotune. The decorator is placed on top of the decorator of the DSL and the kernel is launched with the launch syntax of the DSL. When the kernel is launched with a new tuning key, it is tuned with tune_kernel on copies of the arguments, after which the best configuration is launched directly. - supports Triton, Numba, CuPy, Warp, Taichi, CuTe DSL, Tilus, TileLang, and cuTile, with launch() helpers for cuTile and Warp - search space as tune_params or a list of configs (dicts or triton.Config), brute_force by default, other tune_kernel options are passed on, including parallel_compile - optional verification with a reference function - tuning results can be stored in wisdom files in the format of Kernel Launcher, so that a new process does not tune again - kernel copies created for each configuration no longer contain the autotune decorator --- CHANGELOG.md | 1 + doc/source/backends.rst | 55 ++ doc/source/user-api.rst | 2 + examples/python/README.rst | 48 ++ examples/python/cutile_vec_add_autotune.py | 46 ++ examples/python/numba_vec_add_autotune.py | 41 ++ examples/python/tilelang_vec_add_autotune.py | 41 ++ examples/python/triton_matmul_autotune.py | 58 ++ examples/python/triton_vec_add_autotune.py | 65 ++ examples/python/warp_vec_add_autotune.py | 45 ++ kernel_tuner/__init__.py | 2 + kernel_tuner/decorator.py | 629 ++++++++++++++++++ .../kernel_sources/kernel_source_fn.py | 24 +- kernel_tuner/utils/wisdom.py | 146 ++++ .../autotuned/__init__.py | 10 + .../autotuned/cupyx_kernel.py | 24 + .../autotuned/cute_kernel.py | 29 + .../autotuned/cutile_kernel.py | 27 + .../autotuned/numba_kernel.py | 24 + .../autotuned/taichi_kernel.py | 29 + .../autotuned/tilelang_kernel.py | 26 + .../autotuned/tilus_kernel.py | 35 + .../autotuned/triton_kernel.py | 20 + .../autotuned/warp_kernel.py | 25 + test/test_decorator.py | 328 +++++++++ 25 files changed, 1778 insertions(+), 2 deletions(-) create mode 100644 examples/python/cutile_vec_add_autotune.py create mode 100644 examples/python/numba_vec_add_autotune.py create mode 100644 examples/python/tilelang_vec_add_autotune.py create mode 100644 examples/python/triton_matmul_autotune.py create mode 100644 examples/python/triton_vec_add_autotune.py create mode 100644 examples/python/warp_vec_add_autotune.py create mode 100644 kernel_tuner/decorator.py create mode 100644 kernel_tuner/utils/wisdom.py create mode 100644 test/generic_python_kernels/autotuned/__init__.py create mode 100644 test/generic_python_kernels/autotuned/cupyx_kernel.py create mode 100644 test/generic_python_kernels/autotuned/cute_kernel.py create mode 100644 test/generic_python_kernels/autotuned/cutile_kernel.py create mode 100644 test/generic_python_kernels/autotuned/numba_kernel.py create mode 100644 test/generic_python_kernels/autotuned/taichi_kernel.py create mode 100644 test/generic_python_kernels/autotuned/tilelang_kernel.py create mode 100644 test/generic_python_kernels/autotuned/tilus_kernel.py create mode 100644 test/generic_python_kernels/autotuned/triton_kernel.py create mode 100644 test/generic_python_kernels/autotuned/warp_kernel.py create mode 100644 test/test_decorator.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c65fc29e7..530170a9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ This project adheres to [Semantic Versioning](http://semver.org/). - Add support for tuning kernels written in Python DSLs, such as Triton, Numba, Warp, Taichi, Tilus, TileLang, CuTe, cuTile, and CuPy (cupyx.jit), using `lang="generic_python"` and a `call_function` - Detect kernels in Python files as `lang="generic_python"`, and detect the DSL of Python kernels to select a default `call_function` when none is given - Add `parallel_compile` option to compile kernels in parallel threads before benchmarking them one after the other, supported by the CUDA backends and all Python DSLs. Numba, Warp, cuTile, CuTe, and Tilus kernels are compiled in worker processes instead of threads, Taichi kernels one at a time in the main thread +- Add `kernel_tuner.autotune` decorator that tunes Python DSL kernels inside an application when they are launched with a new tuning key, like `triton.autotune`, with optional verification and wisdom files in the format of Kernel Launcher to store tuning results - Fix PyCUDA backend raising an error instead of skipping configurations that use too much shared memory - Add support for PyTorch tensors and other objects implementing the CUDA Array Interface as kernel arguments in the cuda-python (nvcuda) backend - Change PyCUDA backend to copy PyTorch CUDA tensors to a separate device allocation, the user's tensors are no longer modified by the kernel diff --git a/doc/source/backends.rst b/doc/source/backends.rst index 80953c00e..82d887729 100644 --- a/doc/source/backends.rst +++ b/doc/source/backends.rst @@ -169,3 +169,58 @@ Configurations that cannot be compiled or launched because they use too many res to compile within the time limit, or use tile sizes that the DSL does not support, are skipped. Other errors, such as errors in the kernel code, stop the tuning process. Constant, shared, and texture memory arguments are not supported by the Python backend. + + +Autotuning kernels with a decorator +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Instead of calling ``tune_kernel``, a Python kernel can also be tuned inside an application with the +:func:`kernel_tuner.autotune` decorator, which works like ``triton.autotune``. The decorator is placed on +top of the decorator of the DSL, and the kernel is launched as usual. When the kernel is launched with a new +*tuning key*, Kernel Tuner tunes the kernel on copies of the arguments, after which the best configuration is +launched. Later launches with the same key launch the best configuration directly: + +.. code-block:: python + + from kernel_tuner import autotune + + @autotune(tune_params={"block_size_x": [128, 256, 512, 1024], "num_warps": [2, 4, 8]}, key=["n"]) + @triton.jit + def vector_add(c_ptr, a_ptr, b_ptr, n, block_size_x: tl.constexpr): + ... + + vector_add[lambda meta: (triton.cdiv(n, meta["block_size_x"]),)](c, a, b, n) + +The tuning key consists of the values of the arguments named in ``key`` and the data types of all array +arguments, ``key`` can also be a function of the kernel arguments. Instead of ``tune_params``, you can pass a list +of configurations with ``configs``, as dictionaries or ``triton.Config`` objects, and only these configurations +are benchmarked. The search strategy is ``brute_force`` by default, other options of ``tune_kernel``, such as +``strategy``, ``restrictions``, ``iterations``, and ``parallel_compile``, are passed on to ``tune_kernel``. + +The kernel is launched with the launch syntax of its DSL. The grid and thread block dimensions can be functions of +a dictionary with the kernel arguments by name and the tunable parameters, as in Triton. In DSLs where the thread +block size is part of the launch, such as Numba and CuPy, pass ``grid`` and ``threads`` functions to the decorator +instead, these replace the launch dimensions of the launch. + +.. csv-table:: Launching autotuned kernels + :header: DSL, Launch + :widths: auto + + Triton, "``kernel[grid](*args)``" + Numba, "``kernel[grid, block](*args)``" + CuPy, "``kernel[grid, block](*args)`` or ``kernel(grid, block, args)``" + Warp, "``kernel[dim](*args)`` or ``kernel.launch(dim, inputs, ...)``, a tunable parameter ``block_dim`` sets the block size" + Taichi, "``kernel(*args)``" + CuTe, "``kernel(*args)``" + Tilus, "``Script()(*args)``, constructors with arguments are not supported" + TileLang, "``factory(*factory_args)(*args)``, the factory arguments are part of the tuning key" + cuTile, "``kernel[grid](*args)`` or ``kernel.launch(stream, grid, args)``" + +Outputs are not verified by default. To verify the outputs of every configuration, pass a ``reference`` function +that receives the kernel arguments and returns the expected outputs, as a dictionary with the names of the output +arguments or a list with a value for every argument (None for inputs), and optionally a tolerance ``atol``. + +With ``wisdom``, the tuning results are stored in a directory with *wisdom files*, one per kernel, in the format of +`Kernel Launcher `__. When a kernel is launched with a tuning key +for which the wisdom file has results on the same GPU, the best stored configuration is used without tuning. +Wisdom files can therefore also be produced ahead of time and shipped with an application. diff --git a/doc/source/user-api.rst b/doc/source/user-api.rst index e646f5bb1..9e9a6f175 100644 --- a/doc/source/user-api.rst +++ b/doc/source/user-api.rst @@ -13,6 +13,8 @@ This file provides all the details you need about how to call the Kernel Tuner's .. autofunction:: kernel_tuner.run_kernel +.. autofunction:: kernel_tuner.autotune + .. autofunction:: kernel_tuner.store_results .. autofunction:: kernel_tuner.create_device_targets diff --git a/examples/python/README.rst b/examples/python/README.rst index b53d33e2f..a1259a7ca 100644 --- a/examples/python/README.rst +++ b/examples/python/README.rst @@ -129,3 +129,51 @@ cuTile [`cutile_matmul.py `__] - tune the tile sizes of a tile-based kernel that uses tensor cores - tell Kernel Tuner to compute a grid with one block per tile of the output matrix + +Autotuning while the application runs +------------------------------------- + +These examples use the ``kernel_tuner.autotune`` decorator instead of ``tune_kernel``, which works like +``triton.autotune``: the kernel is launched as usual, and tuned when it is launched for a new tuning key. + +Triton vector add +~~~~~~~~~~~~~~~~~ +[`triton_vec_add_autotune.py `__] + - add the decorator on top of ``@triton.jit`` and launch the kernel with a grid function, as in Triton + - tune again for every new vector size, the tuning key + - replace ``triton.autotune`` by passing its list of ``triton.Config`` objects as ``configs`` + +Numba vector add +~~~~~~~~~~~~~~~~ +[`numba_vec_add_autotune.py `__] + - tune the thread block size of a kernel that is launched with ``kernel[grid, block]`` + - let the decorator compute the launch dimensions with ``grid`` and ``threads`` functions + - pass Numba device arrays and verify the output of every configuration with a ``reference`` function + +Warp vector add +~~~~~~~~~~~~~~~ +[`warp_vec_add_autotune.py `__] + - launch the kernel with ``kernel.launch(dim, inputs=[...])``, like ``wp.launch`` + - tune a constant in the kernel together with the thread block size ``block_dim`` + - compute the launch dimensions from the kernel arguments and the tunable parameters + +cuTile vector add +~~~~~~~~~~~~~~~~~ +[`cutile_vec_add_autotune.py `__] + - launch the kernel with ``kernel.launch(stream, grid, args)``, like ``ct.launch`` + - use a function of the kernel arguments as the tuning key + - store the tuning results in a wisdom file, so that the next run does not tune again + +TileLang vector add +~~~~~~~~~~~~~~~~~~~ +[`tilelang_vec_add_autotune.py `__] + - add the decorator on top of a kernel factory decorated with ``@tilelang.jit`` + - tune an argument of the factory, the other arguments of the factory are part of the tuning key + +Triton matrix multiplication +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +[`triton_matmul_autotune.py `__] + - tune a kernel with the ``kernel_tuner.autotune`` decorator instead of ``tune_kernel``, like ``triton.autotune`` + - tune again when the kernel is launched with a new matrix size, the tuning key + - verify the outputs of every configuration against a reference function + - store the tuning results in wisdom files, so that the next run does not tune again diff --git a/examples/python/cutile_vec_add_autotune.py b/examples/python/cutile_vec_add_autotune.py new file mode 100644 index 000000000..3b924cde9 --- /dev/null +++ b/examples/python/cutile_vec_add_autotune.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python +"""Autotune a cuTile vector add kernel with the Kernel Tuner autotune decorator. + +The decorated kernel is launched like ct.launch, but with kernel.launch(stream, grid, args). The tuning key is a +function of the kernel arguments, here the length of the vectors. + +The tuning results are stored in wisdom files in the directory "wisdom", in the format of Kernel Launcher. When the +example runs again, the kernel is not tuned again, but the best configuration from the wisdom file is launched. +""" + +import cuda.tile as ct +import torch + +from kernel_tuner import autotune + + +@autotune( + tune_params={"tile_size": [128, 256, 512, 1024, 2048]}, + key=lambda c, a, b: c.shape[0], + wisdom="wisdom", +) +@ct.kernel +def vector_add(c, a, b, tile_size: ct.Constant[int]): + pid = ct.bid(0) + a_tile = ct.load(a, index=(pid,), shape=(tile_size,)) + b_tile = ct.load(b, index=(pid,), shape=(tile_size,)) + ct.store(c, index=(pid,), tile=a_tile + b_tile) + + +def grid(meta): + """One block per tile, meta contains the kernel arguments and the tunable parameters.""" + return (ct.cdiv(meta["c"].shape[0], meta["tile_size"]),) + + +if __name__ == "__main__": + n = 1 << 24 + a = torch.randn(n, device="cuda") + b = torch.randn(n, device="cuda") + c = torch.empty_like(a) + + # instead of ct.launch(stream, grid, vector_add, (c, a, b, tile_size)), the tile size is chosen by the decorator + vector_add.launch(torch.cuda.current_stream(), grid, (c, a, b)) + + torch.testing.assert_close(c, a + b) + tuned = "tuned" if vector_add.tuning_results else "read from the wisdom file" + print(f"best configuration ({tuned}):", vector_add.best_configs) diff --git a/examples/python/numba_vec_add_autotune.py b/examples/python/numba_vec_add_autotune.py new file mode 100644 index 000000000..69d8dea43 --- /dev/null +++ b/examples/python/numba_vec_add_autotune.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python +"""Autotune a Numba CUDA vector add kernel with the Kernel Tuner autotune decorator. + +In Numba, the thread block size is part of the launch: kernel[grid, block](*args). Because the block size is the +tunable parameter, the decorator computes the launch dimensions with the grid and threads functions, which +receive the kernel arguments and the tunable parameters of a configuration. These replace the launch dimensions +of the launch, so the kernel can be launched with any placeholder dimensions. + +The reference function makes Kernel Tuner verify the output of every configuration while tuning. +""" + +import numpy as np +from numba import cuda + +from kernel_tuner import autotune + + +@autotune( + tune_params={"block_size_x": [32, 64, 128, 256, 512, 1024]}, + key=["n"], + grid=lambda p: ((p["n"] + p["block_size_x"] - 1) // p["block_size_x"],), + threads=lambda p: (p["block_size_x"],), + reference=lambda c, a, b, n: {"c": a.copy_to_host() + b.copy_to_host()}, +) +@cuda.jit +def vector_add(c, a, b, n): + i = cuda.grid(1) + if i < n: + c[i] = a[i] + b[i] + + +if __name__ == "__main__": + n = 10_000_000 + a = cuda.to_device(np.random.randn(n).astype(np.float32)) + b = cuda.to_device(np.random.randn(n).astype(np.float32)) + c = cuda.device_array_like(a) + + vector_add[1, 1](c, a, b, n) # the launch dimensions are computed by the decorator + + assert np.allclose(c.copy_to_host(), a.copy_to_host() + b.copy_to_host()) + print("best configuration:", vector_add.best_configs) diff --git a/examples/python/tilelang_vec_add_autotune.py b/examples/python/tilelang_vec_add_autotune.py new file mode 100644 index 000000000..91905622c --- /dev/null +++ b/examples/python/tilelang_vec_add_autotune.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python +"""Autotune a TileLang vector add kernel with the Kernel Tuner autotune decorator. + +TileLang kernels are created by a kernel factory decorated with @tilelang.jit. The decorator is placed on top of +the factory, the tunable parameters are arguments of the factory. Calling the factory with the other arguments +returns a kernel that is tuned when it is first called. The arguments of the factory are part of the tuning key. +""" + +import tilelang +import tilelang.language as T +import torch + +from kernel_tuner import autotune + + +@autotune(tune_params={"block_size": [128, 256, 512, 1024]}) +@tilelang.jit +def vector_add(n: int, block_size: int = 128): + @T.prim_func + def kernel( + C: T.Tensor((n,), "float32"), + A: T.Tensor((n,), "float32"), + B: T.Tensor((n,), "float32"), + ): + with T.Kernel(T.ceildiv(n, block_size), threads=block_size) as bx: + for i in T.Parallel(block_size): + C[bx * block_size + i] = A[bx * block_size + i] + B[bx * block_size + i] + + return kernel + + +if __name__ == "__main__": + for n in (1 << 16, 1 << 24): + a = torch.randn(n, device="cuda") + b = torch.randn(n, device="cuda") + c = torch.empty_like(a) + + vector_add(n)(c, a, b) # block_size is chosen by the decorator, tuned once for each n + + torch.testing.assert_close(c, a + b) + print("best configurations:", vector_add.best_configs) diff --git a/examples/python/triton_matmul_autotune.py b/examples/python/triton_matmul_autotune.py new file mode 100644 index 000000000..64ee23b58 --- /dev/null +++ b/examples/python/triton_matmul_autotune.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python +"""Tune a Triton matrix multiplication with the autotune decorator, while the application runs. + +The kernel is tuned when it is first launched for a new matrix size (the tuning key), after which the best +configuration is launched. The tuning results are stored in wisdom files in the directory "wisdom", so running +this example again launches the best configurations without tuning. +""" + +import torch +import triton +import triton.language as tl + +from kernel_tuner import autotune + + +@autotune( + tune_params={ + "BLOCK_M": [32, 64, 128], + "BLOCK_N": [32, 64, 128], + "BLOCK_K": [32, 64], + "num_warps": [4, 8], + }, + restrictions=["BLOCK_M * BLOCK_N <= 128 * 64"], + key=["M", "N", "K"], + reference=lambda a, b, c, M, N, K: {"c_ptr": (a.float() @ b.float()).half()}, + atol=0.5, # float16 outputs of large matrices differ by up to a few float16 steps + wisdom="wisdom", +) +@triton.jit +def matmul(a_ptr, b_ptr, c_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K) + a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :], + mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0) + b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :], + mask=(offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0) + acc += tl.dot(a, b) + tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc.to(tl.float16), + mask=(offs_m[:, None] < M) & (offs_n[None, :] < N)) + + +def grid(meta): + return (triton.cdiv(meta["M"], meta["BLOCK_M"]), triton.cdiv(meta["N"], meta["BLOCK_N"])) + + +if __name__ == "__main__": + for size in (1024, 2048): + a = torch.randn(size, size, device="cuda", dtype=torch.float16) + b = torch.randn(size, size, device="cuda", dtype=torch.float16) + c = torch.zeros(size, size, device="cuda", dtype=torch.float16) + matmul[grid](a, b, c, size, size, size) # tunes on the first launch for this size + torch.testing.assert_close(c, (a.float() @ b.float()).half(), atol=1e-1, rtol=1e-2) + print(f"{size}x{size}: best configuration {matmul.best_configs[(size, size, size, *['torch.float16'] * 3)]}") diff --git a/examples/python/triton_vec_add_autotune.py b/examples/python/triton_vec_add_autotune.py new file mode 100644 index 000000000..d3bdad172 --- /dev/null +++ b/examples/python/triton_vec_add_autotune.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python +"""Autotune a Triton vector add kernel with the Kernel Tuner autotune decorator. + +The decorator is placed on top of @triton.jit and the kernel is launched as usual. When the kernel is launched for +a new vector size (the tuning key), Kernel Tuner tunes it on copies of the arguments and launches the best +configuration. Later launches with the same size launch the best configuration directly. + +The second kernel shows that the decorator can replace triton.autotune: it accepts the same list of triton.Config +objects, and only benchmarks these configurations. +""" + +import torch +import triton +import triton.language as tl + +from kernel_tuner import autotune + + +@autotune( + tune_params={"BLOCK_SIZE": [128, 256, 512, 1024], "num_warps": [2, 4, 8]}, + key=["n"], # tune again for every new vector size +) +@triton.jit +def vector_add(c_ptr, a_ptr, b_ptr, n, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n + a = tl.load(a_ptr + offsets, mask=mask) + b = tl.load(b_ptr + offsets, mask=mask) + tl.store(c_ptr + offsets, a + b, mask=mask) + + +# the same kernel, with the configurations of an existing triton.autotune decorator +@autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 256}, num_warps=4), + triton.Config({"BLOCK_SIZE": 1024}, num_warps=8), + ], + key=["n"], +) +@triton.jit +def vector_add_configs(c_ptr, a_ptr, b_ptr, n, BLOCK_SIZE: tl.constexpr): + offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n + a = tl.load(a_ptr + offsets, mask=mask) + b = tl.load(b_ptr + offsets, mask=mask) + tl.store(c_ptr + offsets, a + b, mask=mask) + + +def grid(meta): + """As in Triton, the grid is a function of the kernel arguments and the configuration.""" + return (triton.cdiv(meta["n"], meta["BLOCK_SIZE"]),) + + +if __name__ == "__main__": + for kernel in (vector_add, vector_add_configs): + for n in (10_000, 10_000_000, 10_000): + a = torch.randn(n, device="cuda") + b = torch.randn(n, device="cuda") + c = torch.empty_like(a) + kernel[grid](c, a, b, n) # tunes on the first launch for each n + torch.testing.assert_close(c, a + b) + + print(f"{kernel.name}: tuned {len(kernel.tuning_results)} times, best configurations:") + for key, config in kernel.best_configs.items(): + print(f" n={key[0]}: {config}") diff --git a/examples/python/warp_vec_add_autotune.py b/examples/python/warp_vec_add_autotune.py new file mode 100644 index 000000000..03157a6a0 --- /dev/null +++ b/examples/python/warp_vec_add_autotune.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python +"""Autotune a Warp vector add kernel with the Kernel Tuner autotune decorator. + +Two parameters are tuned: the number of elements per thread, a constant in the kernel that Kernel Tuner replaces +by the value of each configuration, and the thread block size, which the decorator passes to wp.launch as +block_dim. The decorated kernel is launched like wp.launch, but with kernel.launch(...). Because the number of +threads depends on the elements per thread, the launch dimensions are a function of the kernel arguments and the +tunable parameters. +""" + +import numpy as np +import warp as wp + +from kernel_tuner import autotune + +wp.config.log_level = wp.LOG_WARNING # do not print a message for every compiled module +wp.init() + +work_per_thread = 1 # replaced by the value of the tunable parameter + + +@autotune(tune_params={"work_per_thread": [1, 2, 4, 8], "block_dim": [128, 256, 512]}, key=["n"]) +@wp.kernel +def vector_add(c: wp.array(dtype=float), a: wp.array(dtype=float), b: wp.array(dtype=float), n: int): + base = wp.tid() * work_per_thread + for i in range(work_per_thread): + if base + i < n: + c[base + i] = a[base + i] + b[base + i] + + +def dim(meta): + return (meta["n"] + meta["work_per_thread"] - 1) // meta["work_per_thread"] + + +if __name__ == "__main__": + n = 10_000_000 + a = wp.array(np.random.randn(n).astype(np.float32), dtype=float, device="cuda") + b = wp.array(np.random.randn(n).astype(np.float32), dtype=float, device="cuda") + c = wp.zeros(n, dtype=float, device="cuda") + + vector_add.launch(dim, inputs=[c, a, b, n]) # instead of wp.launch(vector_add, dim, inputs=[...]) + # or launch the kernel like a Triton kernel: vector_add[dim](c, a, b, n) + + assert np.allclose(c.numpy(), a.numpy() + b.numpy()) + print("best configuration:", vector_add.best_configs) diff --git a/kernel_tuner/__init__.py b/kernel_tuner/__init__.py index 3823d71b4..e3126e48f 100644 --- a/kernel_tuner/__init__.py +++ b/kernel_tuner/__init__.py @@ -1,6 +1,7 @@ from importlib.metadata import version from kernel_tuner.interface import run_kernel, tune_kernel, tune_kernel_T1, tune_cache +from kernel_tuner.decorator import autotune __version__ = version(__package__) @@ -12,6 +13,7 @@ "tune_cache", "tune_kernel_T1", "__version__", + "autotune" ] diff --git a/kernel_tuner/decorator.py b/kernel_tuner/decorator.py new file mode 100644 index 000000000..6b7f2c742 --- /dev/null +++ b/kernel_tuner/decorator.py @@ -0,0 +1,629 @@ +"""Autotune kernels written in Python-embedded DSLs with a decorator, in the spirit of ``triton.autotune``. + +The ``autotune`` decorator is placed on top of the decorator of the DSL. The decorated kernel is launched with the +launch syntax of the DSL. When the kernel is launched for a new tuning key, Kernel Tuner tunes the kernel on copies +of the arguments, after which the best configuration is launched. Later launches with the same key directly launch +the best configuration:: + + from kernel_tuner import autotune + + @autotune(tune_params={"BLOCK_SIZE": [128, 256, 512, 1024]}, key=["n"]) + @triton.jit + def vector_add(c_ptr, a_ptr, b_ptr, n, BLOCK_SIZE: tl.constexpr): + ... + + vector_add[lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)](c, a, b, n) + +How a kernel is launched depends on the DSL: + +- Triton: ``kernel[grid](*args)`` +- Numba and CuPy: ``kernel[grid, block](*args)``, CuPy also ``kernel(grid, block, args)`` +- cuTile: ``kernel[grid](*args)`` or ``kernel.launch(stream, grid, args)``, like ``ct.launch`` +- Warp: ``kernel[dim](*args)`` or ``kernel.launch(dim, inputs, ...)``, like ``wp.launch`` +- Taichi and CuTe DSL: ``kernel(*args)`` +- Tilus: ``Script()(*args)``, constructors with arguments are not supported +- TileLang: ``factory(*factory_args)(*args)``, the factory arguments are part of the tuning key + +The grid and thread block dimensions can be functions of a dictionary with the kernel arguments (by name) and the +tunable parameters of the configuration, as in Triton. They can also be given to the decorator, which is needed +when the launch dimensions depend on the tunable parameters in DSLs that do not support functions, such as Numba. + +Tuning results can be stored in wisdom files, the format used by Kernel Launcher, see +:mod:`kernel_tuner.utils.wisdom`. A new process then uses the stored results instead of tuning again. + +Tuning copies the arguments of the kernel, so the arguments of the user are not modified while tuning. A possible +extension is to support the ``restore_value`` and ``reset_to_zero`` options of ``triton.autotune``, to tune on the +arguments of the user instead of copies. +""" + +import atexit +import inspect +import logging +import math +import os +import threading +import warnings + +from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn +from kernel_tuner.util import delete_temp_file, get_arg_names, get_kernel_ast +from kernel_tuner.utils import wisdom +from kernel_tuner.utils.language_detection import detect_python_dsl + +# attributes in which the DSLs keep the Python function that was decorated, in order of preference +_WRAPPED_FUNCTION_ATTRIBUTES = ("fn", "py_func", "_pyfunc", "_func", "orig_func", "func", "__wrapped__") + +# launch options of Triton that can be tunable parameters +_TRITON_LAUNCH_OPTIONS = ("num_warps", "num_stages", "num_ctas", "maxnreg") + + +def _unwrap(kernel): + """Return the Python function or class that a DSL decorator wraps. + + Some DSLs return an object that keeps the function in an attribute, others return a wrapper function made with + functools.wraps, which keeps the function in ``__wrapped__``. + """ + for _ in range(5): + if inspect.isclass(kernel): + return kernel + for attribute in _WRAPPED_FUNCTION_ATTRIBUTES: + inner = getattr(kernel, attribute, None) + if inner is not None and inner is not kernel and (callable(inner) or hasattr(inner, "orig_func")): + kernel = inner + break + else: + if inspect.isfunction(kernel): + return kernel + break + raise TypeError(f"Could not find the Python function decorated by {type(kernel).__name__}") + + +def _to_torch(arg): + """Convert GPU arrays of other libraries to PyTorch tensors that share the memory, which Kernel Tuner copies.""" + import numpy as np + import torch + + if isinstance(arg, (torch.Tensor, np.ndarray, np.generic)) or not hasattr(arg, "shape"): + return arg + if hasattr(arg, "__dlpack__"): + return torch.from_dlpack(arg) + if hasattr(arg, "__cuda_array_interface__"): + return torch.as_tensor(arg, device="cuda") + return arg + + +def _expected_like(value, argument): + """Convert an expected output to the array type of the argument passed to tune_kernel, on the host. + + Kernel Tuner compares the outputs on the host, and requires the expected output to be of the same type. + """ + import numpy as np + import torch + + value = _to_torch(value) # GPU arrays of other libraries + if isinstance(argument, np.ndarray): + return value.detach().cpu().numpy() if isinstance(value, torch.Tensor) else np.asarray(value) + if isinstance(argument, torch.Tensor): + return value.detach().cpu() if isinstance(value, torch.Tensor) else torch.as_tensor(np.asarray(value)) + return value + + +def _dtype_name(dtype): + """Return a short name for the data type of an array, Warp uses classes as data types.""" + return dtype.__name__ if isinstance(dtype, type) else str(dtype) + + +def _convert_tensors(args, convert): + """Convert the PyTorch tensors in args, used to pass the copies made by Kernel Tuner to the DSL.""" + import torch + + return [convert(arg) if isinstance(arg, torch.Tensor) else arg for arg in args] + + +class _Launch: + """How the user launched the kernel: launch dimensions and options of the launch syntax of the DSL.""" + + __slots__ = ("grid", "threads", "extra", "options") + + def __init__(self, grid=None, threads=None, extra=(), options=None): + self.grid = grid + self.threads = threads + self.extra = extra # remaining positional launch options, such as the stream in Numba + self.options = options or {} # keyword launch options, such as the stream in cuTile + + def key(self): + """Launch options that select a different kernel are part of the tuning key.""" + return tuple(self.options.get("factory", ())) + + +# launch functions per DSL -------------------------------------------------------------------------------------- +# These launch the kernel of a configuration. They receive the kernel, the arguments, the tunable parameters that +# are kernel arguments (kwargs), the resolved grid and thread block dimensions, the tunable parameters, and a +# dictionary in which they can keep state per configuration, such as a compiled kernel. + + +def _launch_triton(kernel, args, kwargs, grid, threads, launch, params, state): + for option in _TRITON_LAUNCH_OPTIONS: + if option in params: + kwargs[option] = params[option] + return kernel[grid](*args, **kwargs) + + +def _launch_numba(kernel, args, kwargs, grid, threads, launch, params, state): + from numba import cuda + + return kernel[(grid, threads, *launch.extra)](*_convert_tensors(args, cuda.as_cuda_array), **kwargs) + + +def _launch_cupyx(kernel, args, kwargs, grid, threads, launch, params, state): + import cupy as cp + + args = _convert_tensors(args, cp.from_dlpack) + return kernel(grid, threads, (*args, *kwargs.values()), **launch.options) + + +def _launch_warp(kernel, args, kwargs, grid, threads, launch, params, state): + import warp as wp + + wp.init() + options = dict(launch.options) + if "block_dim" in params: + options["block_dim"] = params["block_dim"] + elif threads is not None: + options["block_dim"] = math.prod(threads) if isinstance(threads, (tuple, list)) else threads + inputs = [*_convert_tensors(args, wp.from_torch), *kwargs.values()] + return wp.launch(kernel, dim=grid, inputs=inputs, **options) + + +def _launch_taichi(kernel, args, kwargs, grid, threads, launch, params, state): + return kernel(*args, **kwargs) + + +def _launch_cute(kernel, args, kwargs, grid, threads, launch, params, state): + import cutlass.cute as cute + from cutlass.cute.runtime import from_dlpack + + args = _convert_tensors(args, from_dlpack) + if "compiled" not in state: + state["compiled"] = cute.compile(kernel, *args) + return state["compiled"](*args, **kwargs) + + +def _launch_tilus(kernel, args, kwargs, grid, threads, launch, params, state): + if "script" not in state: + # Kernel Tuner passes an instance while tuning, the dispatch table holds the class + state["script"] = kernel() if inspect.isclass(kernel) else kernel + return state["script"](*args, **kwargs) + + +def _launch_tilelang(kernel, args, kwargs, grid, threads, launch, params, state): + if "compiled" not in state: + factory_args, factory_kwargs = launch.options["factory"] + state["compiled"] = kernel(*factory_args, **{**dict(factory_kwargs), **kwargs}) + return state["compiled"](*args) + + +def _launch_cutile(kernel, args, kwargs, grid, threads, launch, params, state): + import cuda.tile as ct + import torch + + stream = launch.options.get("stream") or torch.cuda.current_stream() + with ct.compiler_timeout(60): + return ct.launch(stream, grid, kernel, (*args, *kwargs.values())) + + +_LAUNCH_FUNCTIONS = { + "triton": _launch_triton, + "numba": _launch_numba, + "cupyx": _launch_cupyx, + "warp": _launch_warp, + "taichi": _launch_taichi, + "cute": _launch_cute, + "tilus": _launch_tilus, + "tilelang": _launch_tilelang, + "cutile": _launch_cutile, +} + +# DSLs that launch kernels with a subscript, kernel[grid](*args), or with Triton-style calling +_SUBSCRIPT_DSLS = {"triton", "numba", "cupyx", "cutile", "warp"} + + +def _named_arguments(argument_names, args, kwargs, spec): + """Return the kernel arguments by name, for TileLang these are the arguments of the kernel factory.""" + if "factory" in spec.options: + factory_args, factory_kwargs = spec.options["factory"] + return {**dict(zip(argument_names, factory_args)), **dict(factory_kwargs)} + return {**dict(zip(argument_names, args)), **kwargs} + + +def _launch_dimensions(grid, threads, spec, argument_names, args, kwargs, params): + """Return the grid and thread block dimensions, those of the decorator take precedence over those of the launch. + + Dimensions can be functions of a dictionary with the kernel arguments by name and the tunable parameters. + """ + grid = grid if grid is not None else spec.grid + threads = threads if threads is not None else spec.threads + if callable(grid) or callable(threads): + meta = {**_named_arguments(argument_names, args, kwargs, spec), **params} + grid = grid(meta) if callable(grid) else grid + threads = threads(meta) if callable(threads) else threads + return grid, threads + + +class _TuningCallFunction: + """The call function used by tune_kernel, it launches each configuration like the user launched the kernel. + + The launch dimensions are computed from the copies of the arguments that Kernel Tuner passes. Instances can be + pickled when the launch dimensions can, so that parallel_compile can compile kernels in worker processes. + """ + + def __init__(self, dsl, argument_names, grid, threads, spec, extra_kwargs): + self.dsl = dsl + self.argument_names = argument_names + self.grid = grid + self.threads = threads + self.spec = spec + self.extra_kwargs = extra_kwargs + self._states = {} + + def __call__(self, kernel_function, args, kwargs, grid, threads, params): + grid, threads = _launch_dimensions( + self.grid, self.threads, self.spec, self.argument_names, args, self.extra_kwargs, params + ) + state = self._states.setdefault(id(kernel_function), {}) + launch_function = _LAUNCH_FUNCTIONS[self.dsl] + launch_function(kernel_function, args, {**self.extra_kwargs, **kwargs}, grid, threads, self.spec, params, state) + + def __getstate__(self): + state = self.__dict__.copy() + state["_states"] = {} # compiled kernels are not sent to worker processes + return state + + +def _configs_to_search_space(configs): + """Convert a list of configurations to tunable parameters and a restriction that only allows these configurations. + + :param configs: Dictionaries with the values of the tunable parameters, or ``triton.Config`` objects. + :type configs: list(dict or triton.Config) + + :returns: The tunable parameters and a restriction function. + :rtype: tuple(dict, callable) + """ + dicts = [] + for config in configs: + if isinstance(config, dict): + dicts.append(dict(config)) + elif hasattr(config, "kwargs") and hasattr(config, "num_warps"): # triton.Config + if getattr(config, "pre_hook", None) is not None: + raise ValueError("triton.Config objects with a pre_hook are not supported") + values = dict(config.kwargs) + values["num_warps"] = config.num_warps + values["num_stages"] = config.num_stages + # only add the other launch options when they are set, the defaults are left to Triton + for option, default in (("num_ctas", 1), ("maxnreg", None)): + if getattr(config, option, default) != default: + values[option] = getattr(config, option) + dicts.append(values) + else: + raise TypeError(f"Configurations must be dictionaries or triton.Config objects, got {type(config)}") + if not dicts: + raise ValueError("configs must contain at least one configuration") + + names = list(dicts[0]) + for values in dicts: + for option in ("num_ctas", "maxnreg"): # Triton options that only some configurations set + if option in names and option not in values: + values[option] = {"num_ctas": 1}.get(option) + if set(values) != set(names): + raise ValueError(f"All configurations must set the same parameters, got {list(values)} and {names}") + + tune_params = {name: list(dict.fromkeys(values[name] for values in dicts)) for name in names} + allowed = {tuple(values[name] for name in names) for values in dicts} + + def only_given_configs(params): + return tuple(params[name] for name in names) in allowed + + return tune_params, only_given_configs + + +class _Entry: + """The kernel of the best configuration for a tuning key, with everything needed to launch it.""" + + __slots__ = ("kernel", "params", "kwargs", "state") + + def __init__(self, kernel, params, kwargs): + self.kernel = kernel + self.params = params + self.kwargs = kwargs # the tunable parameters that are kernel arguments + self.state = {} + + +class AutotunedKernel: + """A kernel that is tuned when it is launched for a new tuning key, created by :func:`autotune`.""" + + def __init__(self, kernel, tune_params, restrictions, key, grid, threads, reference, atol, wisdom_dir, + tune_options): + """Wrap a kernel decorated by a DSL, see :func:`autotune` for the arguments.""" + self.kernel = kernel + function = _unwrap(kernel) + self.name = function.__name__ + self.source_file = inspect.getsourcefile(function) + self.dsl = detect_python_dsl(self.name, self.source_file) + if self.dsl not in _LAUNCH_FUNCTIONS: + raise ValueError( + f"Could not detect the Python DSL of kernel {self.name}, supported are {list(_LAUNCH_FUNCTIONS)}" + ) + self._launch_function = _LAUNCH_FUNCTIONS[self.dsl] + + kernel_ast = get_kernel_ast(self.name, self.source_file) + self.signature = get_arg_names(kernel_ast[1] if isinstance(kernel_ast, tuple) else kernel_ast) + + self.tune_params = tune_params + self.restrictions = restrictions + self.key = key + self.grid = grid + self.threads = threads + self.reference = reference + self.atol = atol + self.wisdom_dir = wisdom_dir + self.tune_options = tune_options + + # arguments of the kernel that the user passes, the tunable parameters are passed by Kernel Tuner + self.argument_names = [name for name in self.signature if name not in tune_params] + self._dispatch = {} + self._lock = threading.Lock() + self.tuning_results = {} # tuning key -> results of tune_kernel + + @property + def best_configs(self): + """The best configuration for each tuning key that the kernel was launched with.""" + return {key: entry.params for key, entry in self._dispatch.items()} + + # launch syntax of the DSLs --------------------------------------------------------------------------------- + + def __getitem__(self, launch): + """Launch the kernel with a subscript: kernel[grid](*args) or kernel[grid, block](*args).""" + if self.dsl not in _SUBSCRIPT_DSLS: + raise TypeError(f"{self.dsl} kernels are launched by calling them, not with kernel[grid](...)") + if self.dsl in ("numba", "cupyx") and isinstance(launch, tuple) and len(launch) >= 2: + spec = _Launch(launch[0], launch[1], extra=tuple(launch[2:])) + else: + spec = _Launch(grid=launch) + return lambda *args, **kwargs: self._launch(args, kwargs, spec) + + def __call__(self, *args, **kwargs): + """Launch the kernel by calling it, or return the script instance or kernel of Tilus and TileLang.""" + if self.dsl == "cupyx": + grid, block, kernel_args = args + return self._launch(tuple(kernel_args), {}, _Launch(grid, block, options=kwargs)) + if self.dsl == "tilus": + if args or kwargs: + raise TypeError("Tilus scripts with constructor arguments are not supported by autotune") + return _BoundKernel(self, _Launch()) + if self.dsl == "tilelang": + # the factory arguments select the kernel, so they are part of the tuning key + factory = (tuple(args), tuple(sorted(kwargs.items()))) + return _BoundKernel(self, _Launch(options={"factory": factory})) + if self.dsl in ("taichi", "cute"): + return self._launch(args, kwargs, _Launch()) + raise TypeError(f"{self.dsl} kernels are launched with kernel[grid](...), see kernel_tuner.decorator") + + def launch(self, *args, **kwargs): + """Launch the kernel like the launch function of the DSL: ``ct.launch`` for cuTile, ``wp.launch`` for Warp. + + For cuTile: ``kernel.launch(stream, grid, args)``. + For Warp: ``kernel.launch(dim, inputs, outputs=(), **options)``, the options are passed to ``wp.launch``. + """ + if self.dsl == "cutile": + stream, grid, kernel_args = args + return self._launch(tuple(kernel_args), {}, _Launch(grid=grid, options={"stream": stream})) + if self.dsl == "warp": + if len(args) < 1 and "dim" not in kwargs: + raise TypeError("launch() needs the launch dimensions dim") + dim = args[0] if args else kwargs.pop("dim") + inputs = args[1] if len(args) > 1 else kwargs.pop("inputs", ()) + outputs = kwargs.pop("outputs", ()) + return self._launch((*inputs, *outputs), {}, _Launch(grid=dim, options=kwargs)) + raise TypeError(f"launch() is only available for cuTile and Warp kernels, not for {self.dsl}") + + # tuning and dispatch --------------------------------------------------------------------------------------- + + def _tuning_key(self, args, kwargs, spec): + """Return the tuning key of a launch. + + The key contains the values of the arguments named in key, the data types of all arrays, and the arguments + of TileLang kernel factories. + """ + if callable(self.key): + return (self.key(*args, **kwargs), *spec.key()) + key = [] + if self.key: + named = _named_arguments(self.argument_names, args, kwargs, spec) + key.extend(named.get(name) for name in self.key) + key.extend(_dtype_name(arg.dtype) for arg in (*args, *kwargs.values()) if hasattr(arg, "dtype")) + return (*key, *spec.key()) + + def _launch(self, args, kwargs, spec): + key = self._tuning_key(args, kwargs, spec) + entry = self._dispatch.get(key) + if entry is None: + with self._lock: + entry = self._dispatch.get(key) + if entry is None: + entry = self._dispatch[key] = self._select(key, args, kwargs, spec) + grid, threads = _launch_dimensions(self.grid, self.threads, spec, self.argument_names, args, kwargs, + entry.params) + return self._launch_function( + entry.kernel, list(args), {**kwargs, **entry.kwargs}, grid, threads, spec, entry.params, entry.state + ) + + def _select(self, key, args, kwargs, spec): + """Select the best configuration for the tuning key from the wisdom file, or by tuning the kernel.""" + import torch + + device_name = torch.cuda.get_device_name() + tunable_parameters = list(self.tune_params) + params = None + filename = wisdom.wisdom_file(self.wisdom_dir, self.name) if self.wisdom_dir else None + if filename: + records = wisdom.read_wisdom(filename, tunable_parameters) + params = wisdom.best_wisdom_config(records, key, device_name, tunable_parameters) + if params is None: + results = self._tune(args, kwargs, spec) + self.tuning_results[key] = results + valid = [result for result in results if isinstance(result.get("time"), float)] + if not valid: + raise RuntimeError(f"No valid configuration found while tuning {self.name} for key {key}") + best = min(valid, key=lambda result: result["time"]) + params = {name: best[name] for name in tunable_parameters} + if filename: + wisdom.write_wisdom(filename, self.name, tunable_parameters, key, device_name, valid) + return self._entry(params) + + def _entry(self, params): + """Create the kernel of a configuration, in the same way as Kernel Tuner does while tuning.""" + kernel_source = KernelSourceFn(self.name, self.source_file, "generic_python", call_function=_not_called) + source_params = {name: value for name, value in params.items() if name not in self.signature} + kernel, temp_file = kernel_source.apply_params_to_source_fn(source_params) + # DSLs may read the source file of a kernel again when they compile it for new argument types + atexit.register(delete_temp_file, temp_file) + kwargs = {name: params[name] for name in self.signature if name in params} + return _Entry(kernel, params, kwargs) + + def _answer(self, args, kwargs, kernel_args): + """Return the expected outputs computed by the reference function, in the form of tune_kernel's answer.""" + if self.reference is None: + return None + expected = self.reference(*args, **kwargs) + if isinstance(expected, dict): + if self.dsl == "tilelang": + raise ValueError("reference must return a list for TileLang kernels, the arguments have no names") + names = self.argument_names + unknown = [name for name in expected if name not in names[: len(kernel_args)]] + if unknown: + raise ValueError(f"reference returned outputs for unknown arguments {unknown}, arguments are {names}") + answer = [expected.get(name) for name in names[: len(kernel_args)]] + else: + answer = list(expected) + [None] * (len(kernel_args) - len(expected)) + return [None if value is None else _expected_like(value, arg) for value, arg in zip(answer, kernel_args)] + + def _tune(self, args, kwargs, spec): + """Tune the kernel with tune_kernel on copies of the arguments.""" + from kernel_tuner.interface import tune_kernel + + # kernel arguments passed by name are passed in the order of the signature + kernel_args = list(args) + for name in self.argument_names[len(args):]: + if name not in kwargs: + break + kernel_args.append(kwargs[name]) + extra_kwargs = {name: value for name, value in kwargs.items() if name not in self.argument_names} + # streams cannot be sent to worker processes, configurations are benchmarked on the current stream + launch_options = {name: value for name, value in spec.options.items() if name != "stream"} + extra = (0, *spec.extra[1:]) if spec.extra else () + tuning_spec = _Launch(spec.grid, spec.threads, extra, launch_options) + call_function = _TuningCallFunction( + self.dsl, self.argument_names, self.grid, self.threads, tuning_spec, extra_kwargs + ) + + kernel_args = [_to_torch(arg) for arg in kernel_args] + options = dict(self.tune_options) + options.setdefault("strategy", "brute_force") + options.setdefault("quiet", True) + restrictions = options.pop("restrictions", self.restrictions) + with warnings.catch_warnings(): + # the launch dimensions are computed by the call function, not from block size parameters + warnings.simplefilter("ignore", UserWarning) + results, _ = tune_kernel( + self.name, + self.source_file, + 1, + kernel_args, + self.tune_params, + restrictions=restrictions, + lang="generic_python", + call_function=call_function, + answer=self._answer(args, kwargs, kernel_args), + atol=self.atol, + **options, + ) + logging.debug(f"tuned {self.name}: {len(results)} configurations") + return results + + +class _BoundKernel: + """A Tilus script instance or a TileLang kernel returned by a factory, launched by calling it.""" + + def __init__(self, autotuned, spec): + self._autotuned = autotuned + self._spec = spec + + def __call__(self, *args, **kwargs): + return self._autotuned._launch(args, kwargs, self._spec) + + +def _not_called(kernel_function, args, kwargs): + """Kernels created for dispatch are launched by the launch functions of the decorator.""" + raise RuntimeError("not called") + + +def autotune(tune_params=None, *, configs=None, restrictions=None, key=None, grid=None, threads=None, + reference=None, atol=1e-6, wisdom=None, **tune_options): + """Tune a kernel written in a Python-embedded DSL when it is launched for a new tuning key. + + :param tune_params: The tunable parameters and their values, as in tune_kernel. + :type tune_params: dict(str: list) + + :param configs: Instead of tune_params, a list of configurations: dictionaries with the values of the tunable + parameters or ``triton.Config`` objects. Only these configurations are benchmarked. + :type configs: list(dict or triton.Config) + + :param restrictions: Restrictions on the tunable parameters, as in tune_kernel. + :type restrictions: list(str) or callable + + :param key: The names of the kernel arguments whose values are part of the tuning key, or a function of the + kernel arguments that returns the key. The data types of the array arguments are always part of the key. + When the kernel is launched with a new key, the kernel is tuned again. + :type key: list(str) or callable + + :param grid: The grid dimensions, or a function of a dictionary with the kernel arguments and tunable + parameters that returns the grid dimensions. Overrides the grid of the launch. + :type grid: tuple or callable + + :param threads: The thread block dimensions, or a function like grid. Overrides the block of the launch. + :type threads: tuple or callable + + :param reference: A function of the kernel arguments that returns the expected outputs, as a dictionary with + the names of the output arguments, or as a list with one value per argument (None for inputs). The outputs + of every configuration are compared with these while tuning. By default the outputs are not verified. + :type reference: callable + + :param atol: The tolerance used to compare the outputs with those of the reference function. + :type atol: float + + :param wisdom: A directory in which the tuning results are stored, in wisdom files named after the kernel. + Configurations stored for the same tuning key and GPU are used instead of tuning the kernel again. + :type wisdom: str + + :param tune_options: Other options passed to tune_kernel, such as strategy, strategy_options, iterations, + parallel_compile, or verbose. The strategy is brute_force by default. + + :returns: A decorator that returns an AutotunedKernel. + """ + if (tune_params is None) == (configs is None): + raise ValueError("Pass either tune_params or configs to autotune") + if configs is not None: + tune_params, only_given_configs = _configs_to_search_space(configs) + if restrictions is None: + restrictions = [only_given_configs] + elif callable(restrictions): + restrictions = [restrictions, only_given_configs] + else: + restrictions = [*restrictions, only_given_configs] + tune_options.setdefault("quiet", not tune_options.get("verbose", False)) + wisdom_dir = os.fspath(wisdom) if wisdom is not None else None + + def decorator(kernel): + return AutotunedKernel(kernel, tune_params, restrictions, key, grid, threads, reference, atol, wisdom_dir, + tune_options) + + return decorator diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index b5e7160e7..6517eaba0 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -16,7 +16,26 @@ from kernel_tuner.kernel_sources.model.prepared_kernel_source_data import PreparedKernelSourceData from kernel_tuner.util import get_kernel_ast, get_arg_names, normalize_call_function from kernel_tuner.utils.call_functions import DEFAULT_CALL_FUNCTIONS, get_default_call_function -from kernel_tuner.utils.language_detection import detect_python_dsl +from kernel_tuner.utils.language_detection import _import_aliases, _qualified_name, detect_python_dsl + +# Decorators of Kernel Tuner itself, these are not copied into the module created for each configuration +KERNEL_TUNER_DECORATORS = {"kernel_tuner.autotune", "kernel_tuner.decorator.autotune"} + + +def remove_kernel_tuner_decorators(node, source_file): + """Remove the decorators of Kernel Tuner from a kernel function or class node. + + The module created for each configuration contains a copy of the kernel with its decorators. The decorator + that autotunes the kernel must not be copied, as it would tune each copy again, and its arguments refer to + names that are not copied into that module. + """ + with open(source_file, "r") as f: + aliases = _import_aliases(ast.parse(f.read(), filename=source_file)) + node.decorator_list = [ + decorator + for decorator in node.decorator_list + if _qualified_name(decorator, aliases) not in KERNEL_TUNER_DECORATORS + ] def kernel_module_name(temp_file_path): @@ -88,8 +107,9 @@ def __init__(self, kernel_name, kernel_source, lang, defines=None, call_function self.source_tree = source_ast[0] self.signature = get_arg_names(source_ast[1]) else: - self.source_tree = source_ast + self.source_tree = source_ast self.signature = get_arg_names(source_ast) + remove_kernel_tuner_decorators(self.source_tree, kernel_source) self.kernel_fn = self.source_tree # This is where we will store the transformed source. self.import_nodes = self._find_import_nodes(kernel_source) diff --git a/kernel_tuner/utils/wisdom.py b/kernel_tuner/utils/wisdom.py new file mode 100644 index 000000000..5c2fe2bda --- /dev/null +++ b/kernel_tuner/utils/wisdom.py @@ -0,0 +1,146 @@ +"""Read and write wisdom files, which store tuning results in the format of Kernel Launcher. + +Kernel Launcher (https://github.com/KernelTuner/kernel_launcher) stores the tuning results of a kernel in a +wisdom file named ``.wisdom``. The file is in JSON Lines format: the first line is a header that lists the +names of the tunable parameters, and every following line is a record with the values of the tunable parameters +(``config``), the ``problem_size``, the objective (``time``), and the ``environment`` in which the configuration +was benchmarked, which includes the name of the GPU (``device_name``). + +Kernel Launcher expects the problem size to be a list of integers. The tuning keys of the autotune decorator can +also contain other values, such as data types, so each record also stores the complete tuning key in ``key``. +Kernel Launcher ignores this field. +""" + +import json +import logging +import os +import tempfile + +WISDOM_VERSION = "1.0" +WISDOM_OBJECTIVE = "time" + + +def wisdom_file(directory, name): + """Return the path of the wisdom file for the kernel with the given name.""" + return os.path.join(directory, f"{name}.wisdom") + + +def normalize_key(key): + """Convert a tuning key to the JSON representation that is stored in wisdom files.""" + return json.loads(json.dumps(key, default=str)) + + +def problem_size_of_key(key): + """Return the integers in a tuning key, which Kernel Launcher uses as the problem size.""" + sizes = [] + + def collect(value): + if isinstance(value, bool): + return + if isinstance(value, int): + sizes.append(value) + elif isinstance(value, (list, tuple)): + for item in value: + collect(item) + + collect(normalize_key(key)) + return sizes + + +def read_wisdom(filename, tunable_parameters): + """Read the records of a wisdom file. + + :param filename: Path of the wisdom file. + :type filename: str + + :param tunable_parameters: Names of the tunable parameters. Records are only returned if the file stores the + same tunable parameters, otherwise the file belongs to a different version of the kernel. + :type tunable_parameters: list(str) + + :returns: The records in the file, an empty list if the file does not exist or does not match. + :rtype: list(dict) + """ + if not os.path.isfile(filename): + return [] + with open(filename) as handle: + lines = [line for line in handle if line.strip()] + if not lines: + return [] + header = json.loads(lines[0]) + if header.get("tunable_parameters") != list(tunable_parameters): + logging.warning( + f"Ignoring wisdom file {filename}, it stores the tunable parameters {header.get('tunable_parameters')} " + f"instead of {list(tunable_parameters)}" + ) + return [] + return [json.loads(line) for line in lines[1:]] + + +def best_wisdom_config(records, key, device_name, tunable_parameters): + """Return the configuration with the lowest time for the tuning key on the device, or None. + + :returns: A dictionary with the values of the tunable parameters, or None if no record matches. + :rtype: dict or None + """ + key = normalize_key(key) + matching = [ + record + for record in records + if record.get("key") == key and record.get("environment", {}).get("device_name") == device_name + ] + if not matching: + return None + best = min(matching, key=lambda record: record[WISDOM_OBJECTIVE]) + return dict(zip(tunable_parameters, best["config"])) + + +def write_wisdom(filename, name, tunable_parameters, key, device_name, results): + """Store the tuning results for a tuning key in a wisdom file. + + Records for the same tuning key and device are replaced, other records are kept. + + :param name: The name of the kernel, stored as the key of the wisdom file. + :type name: str + + :param results: The valid results of tune_kernel, these contain the tunable parameters and the time. + :type results: list(dict) + """ + directory = os.path.dirname(os.path.abspath(filename)) + os.makedirs(directory, exist_ok=True) + if os.path.isfile(filename) and os.path.getsize(filename) > 0: + with open(filename) as handle: + stored_parameters = json.loads(handle.readline()).get("tunable_parameters") + if stored_parameters != list(tunable_parameters): + # do not overwrite the results of another version of the kernel + logging.warning(f"Not storing tuning results in {filename}, it stores other tunable parameters") + return + key = normalize_key(key) + records = [ + record + for record in read_wisdom(filename, tunable_parameters) + if record.get("key") != key or record.get("environment", {}).get("device_name") != device_name + ] + problem_size = problem_size_of_key(key) + for result in results: + records.append( + { + "config": [result[param] for param in tunable_parameters], + "problem_size": problem_size, + WISDOM_OBJECTIVE: result[WISDOM_OBJECTIVE], + "environment": {"device_name": device_name}, + "key": key, + } + ) + + header = { + "tunable_parameters": list(tunable_parameters), + "version": WISDOM_VERSION, + "objective": WISDOM_OBJECTIVE, + "key": name, + } + lines = [json.dumps(header)] + [json.dumps(record, default=str) for record in records] + # write to a temporary file first, so that a process that reads the file never sees a partial file + with tempfile.NamedTemporaryFile("w", dir=directory, suffix=".tmp", delete=False) as handle: + handle.write("\n".join(lines) + "\n") + temp_name = handle.name + os.replace(temp_name, filename) diff --git a/test/generic_python_kernels/autotuned/__init__.py b/test/generic_python_kernels/autotuned/__init__.py new file mode 100644 index 000000000..88e0d6e02 --- /dev/null +++ b/test/generic_python_kernels/autotuned/__init__.py @@ -0,0 +1,10 @@ +"""Vector add kernels in Python DSLs decorated with kernel_tuner.autotune, used by test_decorator.py. + +Each DSL lives in its own module, because Kernel Tuner copies all imports of the kernel source file into the +module it generates for each configuration. Every module provides: + +- ``AUTOTUNE``: the arguments of the autotune decorator of the kernel +- ``vector_add``: the decorated kernel +- ``launch(kernel, c, a, b, n)``: launches an autotuned kernel with the launch syntax of the DSL, the arrays are + PyTorch CUDA tensors that are converted to the array type of the DSL where needed +""" diff --git a/test/generic_python_kernels/autotuned/cupyx_kernel.py b/test/generic_python_kernels/autotuned/cupyx_kernel.py new file mode 100644 index 000000000..2a070f01d --- /dev/null +++ b/test/generic_python_kernels/autotuned/cupyx_kernel.py @@ -0,0 +1,24 @@ +import cupy as cp +from cupyx import jit + +from kernel_tuner import autotune + +AUTOTUNE = dict( + tune_params={"block_size_x": [128, 256]}, + key=["n"], + grid=lambda p: ((p["n"] + p["block_size_x"] - 1) // p["block_size_x"],), + threads=lambda p: (p["block_size_x"],), +) + + +@autotune(**AUTOTUNE) +@jit.rawkernel() +def vector_add(c, a, b, n): + i = jit.blockIdx.x * jit.blockDim.x + jit.threadIdx.x + if i < n: + c[i] = a[i] + b[i] + + +def launch(kernel, c, a, b, n): + c, a, b = (cp.from_dlpack(x) for x in (c, a, b)) + kernel((1,), (1,), (c, a, b, n)) # the launch dimensions are replaced by those of the decorator diff --git a/test/generic_python_kernels/autotuned/cute_kernel.py b/test/generic_python_kernels/autotuned/cute_kernel.py new file mode 100644 index 000000000..ab3c38e5f --- /dev/null +++ b/test/generic_python_kernels/autotuned/cute_kernel.py @@ -0,0 +1,29 @@ +import cutlass.cute as cute + +from kernel_tuner import autotune + +AUTOTUNE = dict(tune_params={"block_size_x": [128, 256]}, key=["n"]) + + +@cute.kernel +def vector_add_kernel(gC: cute.Tensor, gA: cute.Tensor, gB: cute.Tensor, n: cute.Int32): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + bdim, _, _ = cute.arch.block_dim() + i = bdim * bidx + tidx + if i < n: + gC[i] = gA[i] + gB[i] + + +@autotune(**AUTOTUNE) +@cute.jit +def vector_add(mC: cute.Tensor, mA: cute.Tensor, mB: cute.Tensor, n: cute.Int32): + block_size_x = 128 + vector_add_kernel(mC, mA, mB, n).launch( + grid=(cute.ceil_div(n, block_size_x), 1, 1), + block=(block_size_x, 1, 1), + ) + + +def launch(kernel, c, a, b, n): + kernel(c, a, b, n) diff --git a/test/generic_python_kernels/autotuned/cutile_kernel.py b/test/generic_python_kernels/autotuned/cutile_kernel.py new file mode 100644 index 000000000..87e94427e --- /dev/null +++ b/test/generic_python_kernels/autotuned/cutile_kernel.py @@ -0,0 +1,27 @@ +import cuda.tile as ct +import torch + +from kernel_tuner import autotune + +AUTOTUNE = dict(tune_params={"block_size_x": [128, 256]}, key=lambda c, a, b: (c.shape[0],)) + + +@autotune(**AUTOTUNE) +@ct.kernel +def vector_add(c, a, b, block_size_x: ct.Constant[int]): + pid = ct.bid(0) + a_tile = ct.load(a, index=(pid,), shape=(block_size_x,)) + b_tile = ct.load(b, index=(pid,), shape=(block_size_x,)) + ct.store(c, index=(pid,), tile=a_tile + b_tile) + + +def grid(meta): + return (ct.cdiv(meta["c"].shape[0], meta["block_size_x"]),) + + +def launch(kernel, c, a, b, n): + kernel.launch(torch.cuda.current_stream(), grid, (c, a, b)) + + +def launch_subscript(kernel, c, a, b, n): + kernel[grid](c, a, b) diff --git a/test/generic_python_kernels/autotuned/numba_kernel.py b/test/generic_python_kernels/autotuned/numba_kernel.py new file mode 100644 index 000000000..43dcd5349 --- /dev/null +++ b/test/generic_python_kernels/autotuned/numba_kernel.py @@ -0,0 +1,24 @@ +from numba import cuda + +from kernel_tuner import autotune + +# the thread block size is a launch dimension in Numba, so the decorator computes the launch dimensions +AUTOTUNE = dict( + tune_params={"block_size_x": [128, 256]}, + key=["n"], + grid=lambda p: ((p["n"] + p["block_size_x"] - 1) // p["block_size_x"],), + threads=lambda p: (p["block_size_x"],), +) + + +@autotune(**AUTOTUNE) +@cuda.jit +def vector_add(c, a, b, n): + i = cuda.grid(1) + if i < n: + c[i] = a[i] + b[i] + + +def launch(kernel, c, a, b, n): + c, a, b = (cuda.as_cuda_array(x) for x in (c, a, b)) + kernel[1, 1](c, a, b, n) # the launch dimensions are replaced by those of the decorator diff --git a/test/generic_python_kernels/autotuned/taichi_kernel.py b/test/generic_python_kernels/autotuned/taichi_kernel.py new file mode 100644 index 000000000..1cc48b6cb --- /dev/null +++ b/test/generic_python_kernels/autotuned/taichi_kernel.py @@ -0,0 +1,29 @@ +import taichi as ti + +from kernel_tuner import autotune + +ti.init(arch=ti.cuda) + +# Taichi silently falls back to the CPU when it cannot use CUDA +requires_skip = None if ti.lang.impl.current_cfg().arch == ti.cuda else "Taichi cannot use CUDA" + +AUTOTUNE = dict(tune_params={"block_size_x": [128, 256]}, key=["n"]) + +block_size_x = 128 # replaced by the value of the tunable parameter + + +@autotune(**AUTOTUNE) +@ti.kernel +def vector_add( + c: ti.types.ndarray(dtype=ti.f32, ndim=1), + a: ti.types.ndarray(dtype=ti.f32, ndim=1), + b: ti.types.ndarray(dtype=ti.f32, ndim=1), + n: ti.i32, +): + ti.loop_config(block_dim=block_size_x) + for i in range(n): + c[i] = a[i] + b[i] + + +def launch(kernel, c, a, b, n): + kernel(c, a, b, n) diff --git a/test/generic_python_kernels/autotuned/tilelang_kernel.py b/test/generic_python_kernels/autotuned/tilelang_kernel.py new file mode 100644 index 000000000..324b2417f --- /dev/null +++ b/test/generic_python_kernels/autotuned/tilelang_kernel.py @@ -0,0 +1,26 @@ +import tilelang +import tilelang.language as T + +from kernel_tuner import autotune + +AUTOTUNE = dict(tune_params={"block_size_x": [128, 256]}, key=["n"]) + + +@autotune(**AUTOTUNE) +@tilelang.jit +def vector_add(n: int, block_size_x: int = 128): + @T.prim_func + def vector_add_kernel( + C: T.Tensor((n,), "float32"), + A: T.Tensor((n,), "float32"), + B: T.Tensor((n,), "float32"), + ): + with T.Kernel(T.ceildiv(n, block_size_x), threads=block_size_x) as bx: + for i in T.Parallel(block_size_x): + C[bx * block_size_x + i] = A[bx * block_size_x + i] + B[bx * block_size_x + i] + + return vector_add_kernel + + +def launch(kernel, c, a, b, n): + kernel(n)(c, a, b) diff --git a/test/generic_python_kernels/autotuned/tilus_kernel.py b/test/generic_python_kernels/autotuned/tilus_kernel.py new file mode 100644 index 000000000..05f8c56a0 --- /dev/null +++ b/test/generic_python_kernels/autotuned/tilus_kernel.py @@ -0,0 +1,35 @@ +import tilus +from tilus import float32, int32 +from tilus.utils import cdiv + +from kernel_tuner import autotune + +AUTOTUNE = dict(tune_params={"block_size_x": [128, 256], "num_warps": [4]}, key=["n"]) + + +@autotune(**AUTOTUNE) +class VectorAdd(tilus.Script): + def __init__(self): + super().__init__() + self.block_size_x = 128 + self.num_warps = 4 + + def __call__(self, c_ptr: ~float32, a_ptr: ~float32, b_ptr: ~float32, n: int32): + self.attrs.blocks = [cdiv(n, self.block_size_x)] + self.attrs.warps = self.num_warps + + offset: int32 = self.block_size_x * self.blockIdx.x + ga = self.global_view(a_ptr, dtype=float32, shape=[n]) + gb = self.global_view(b_ptr, dtype=float32, shape=[n]) + gc = self.global_view(c_ptr, dtype=float32, shape=[n]) + + a = self.load_global(ga, offsets=[offset], shape=[self.block_size_x]) + b = self.load_global(gb, offsets=[offset], shape=[self.block_size_x]) + self.store_global(gc, a + b, offsets=[offset]) + + +vector_add = VectorAdd + + +def launch(kernel, c, a, b, n): + kernel()(c, a, b, n) diff --git a/test/generic_python_kernels/autotuned/triton_kernel.py b/test/generic_python_kernels/autotuned/triton_kernel.py new file mode 100644 index 000000000..f8017d37e --- /dev/null +++ b/test/generic_python_kernels/autotuned/triton_kernel.py @@ -0,0 +1,20 @@ +import triton +import triton.language as tl + +from kernel_tuner import autotune + +AUTOTUNE = dict(tune_params={"block_size_x": [128, 256], "num_warps": [2, 4]}, key=["n"]) + + +@autotune(**AUTOTUNE) +@triton.jit +def vector_add(c_ptr, a_ptr, b_ptr, n, block_size_x: tl.constexpr): + offsets = tl.program_id(axis=0) * block_size_x + tl.arange(0, block_size_x) + mask = offsets < n + a = tl.load(a_ptr + offsets, mask=mask) + b = tl.load(b_ptr + offsets, mask=mask) + tl.store(c_ptr + offsets, a + b, mask=mask) + + +def launch(kernel, c, a, b, n): + kernel[lambda meta: (triton.cdiv(n, meta["block_size_x"]),)](c, a, b, n) diff --git a/test/generic_python_kernels/autotuned/warp_kernel.py b/test/generic_python_kernels/autotuned/warp_kernel.py new file mode 100644 index 000000000..7de5a1c33 --- /dev/null +++ b/test/generic_python_kernels/autotuned/warp_kernel.py @@ -0,0 +1,25 @@ +import warp as wp + +from kernel_tuner import autotune + +wp.init() + +AUTOTUNE = dict(tune_params={"block_dim": [128, 256]}, key=["n"]) + + +@autotune(**AUTOTUNE) +@wp.kernel +def vector_add(c: wp.array(dtype=float), a: wp.array(dtype=float), b: wp.array(dtype=float), n: int): + i = wp.tid() + if i < n: + c[i] = a[i] + b[i] + + +def launch(kernel, c, a, b, n): + c, a, b = (wp.from_torch(x) for x in (c, a, b)) + kernel.launch(n, inputs=[c, a, b, n]) + + +def launch_subscript(kernel, c, a, b, n): + c, a, b = (wp.from_torch(x) for x in (c, a, b)) + kernel[n](c, a, b, n) diff --git a/test/test_decorator.py b/test/test_decorator.py new file mode 100644 index 000000000..cd3e2b60d --- /dev/null +++ b/test/test_decorator.py @@ -0,0 +1,328 @@ +"""Tests for kernel_tuner.decorator.autotune and the wisdom files in which it stores tuning results.""" + +import importlib +import json + +import pytest + +import kernel_tuner +from kernel_tuner import autotune +from kernel_tuner.decorator import AutotunedKernel, _configs_to_search_space +from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn +from kernel_tuner.utils import wisdom + +from .context import ( + skip_if_no_cupy, + skip_if_no_cute, + skip_if_no_cutile, + skip_if_no_numba_cuda, + skip_if_no_taichi, + skip_if_no_tilelang, + skip_if_no_tilus, + skip_if_no_torch, + skip_if_no_triton, + skip_if_no_warp, +) + +# every DSL needs Torch with a CUDA device, which the generic_python backend uses for memory and timing +dsls = [ + pytest.param("triton", marks=[skip_if_no_torch, skip_if_no_triton]), + pytest.param("numba", marks=[skip_if_no_torch, skip_if_no_numba_cuda]), + pytest.param("cupyx", marks=[skip_if_no_torch, skip_if_no_cupy]), + pytest.param("warp", marks=[skip_if_no_torch, skip_if_no_warp]), + pytest.param("taichi", marks=[skip_if_no_torch, skip_if_no_taichi]), + pytest.param("cute", marks=[skip_if_no_torch, skip_if_no_cute]), + pytest.param("tilus", marks=[skip_if_no_torch, skip_if_no_tilus]), + pytest.param("tilelang", marks=[skip_if_no_torch, skip_if_no_tilelang]), + pytest.param("cutile", marks=[skip_if_no_torch, skip_if_no_cutile]), +] + + +def kernel_module(dsl): + """Import the module with the autotuned kernel of a DSL, these import their DSL.""" + module = importlib.import_module(f".generic_python_kernels.autotuned.{dsl}_kernel", package=__package__) + if getattr(module, "requires_skip", None): + pytest.skip(module.requires_skip) + return module + + +def fresh_kernel(module, **options): + """Decorate the DSL kernel of a module again, so that every test starts without tuning results.""" + return autotune(**{**module.AUTOTUNE, **options})(module.vector_add.kernel) + + +def vectors(n): + import torch + + a = torch.randn(n, device="cuda") + b = torch.randn(n, device="cuda") + c = torch.zeros(n, device="cuda") + return c, a, b + + +def launch_and_check(module, kernel, n, launch=None): + import torch + + c, a, b = vectors(n) + (launch or module.launch)(kernel, c, a, b, n) + torch.cuda.synchronize() + assert torch.allclose(c, a + b) + + +# tests that do not need a GPU ---------------------------------------------------------------------------------- + + +def test_configs_to_search_space_dicts(): + tune_params, restriction = _configs_to_search_space([{"x": 1, "y": 8}, {"x": 4, "y": 16}, {"x": 4, "y": 8}]) + assert tune_params == {"x": [1, 4], "y": [8, 16]} + assert restriction({"x": 1, "y": 8}) + assert restriction({"x": 4, "y": 16}) + assert not restriction({"x": 1, "y": 16}) + + +def test_configs_to_search_space_triton_configs(): + class Config: # the attributes of triton.Config that are used + def __init__(self, kwargs, num_warps=4, num_stages=3, num_ctas=1, maxnreg=None, pre_hook=None): + self.kwargs = kwargs + self.num_warps = num_warps + self.num_stages = num_stages + self.num_ctas = num_ctas + self.maxnreg = maxnreg + self.pre_hook = pre_hook + + tune_params, restriction = _configs_to_search_space( + [Config({"BLOCK": 64}, num_warps=2), Config({"BLOCK": 128}, num_stages=4)] + ) + assert tune_params == {"BLOCK": [64, 128], "num_warps": [2, 4], "num_stages": [3, 4]} + assert restriction({"BLOCK": 64, "num_warps": 2, "num_stages": 3}) + assert not restriction({"BLOCK": 64, "num_warps": 4, "num_stages": 3}) + + # launch options that only some configurations set get the default of Triton for the others + tune_params, _ = _configs_to_search_space([Config({"BLOCK": 64}, num_ctas=2), Config({"BLOCK": 128})]) + assert tune_params["num_ctas"] == [2, 1] + + with pytest.raises(ValueError): + _configs_to_search_space([Config({"BLOCK": 64}, pre_hook=lambda args: None)]) + + +def test_configs_to_search_space_errors(): + with pytest.raises(ValueError): + _configs_to_search_space([{"x": 1}, {"y": 2}]) + with pytest.raises(ValueError): + _configs_to_search_space([]) + with pytest.raises(TypeError): + _configs_to_search_space([(1, 2)]) + + +def test_autotune_is_top_level(): + """The decorator is available next to tune_kernel.""" + import kernel_tuner.decorator + + assert kernel_tuner.autotune is kernel_tuner.decorator.autotune + assert "autotune" in kernel_tuner.__all__ + + +def test_autotune_needs_tune_params_or_configs(): + with pytest.raises(ValueError): + autotune() + with pytest.raises(ValueError): + autotune({"x": [1]}, configs=[{"x": 1}]) + + +def test_wisdom_files(tmp_path): + filename = wisdom.wisdom_file(tmp_path, "vector_add") + params = ["block_size_x", "num_warps"] + results = [{"block_size_x": 128, "num_warps": 4, "time": 2.0}, {"block_size_x": 256, "num_warps": 4, "time": 1.0}] + key = (4096, "torch.float32") + wisdom.write_wisdom(filename, "vector_add", params, key, "GPU A", results) + wisdom.write_wisdom(filename, "vector_add", params, key, "GPU B", results[:1]) + + # the file follows the format of Kernel Launcher + with open(filename) as handle: + lines = [json.loads(line) for line in handle] + assert lines[0] == {"tunable_parameters": params, "version": "1.0", "objective": "time", "key": "vector_add"} + assert lines[1]["config"] == [128, 4] + assert lines[1]["problem_size"] == [4096] + assert lines[1]["environment"] == {"device_name": "GPU A"} + + records = wisdom.read_wisdom(filename, params) + assert len(records) == 3 + assert wisdom.best_wisdom_config(records, key, "GPU A", params) == {"block_size_x": 256, "num_warps": 4} + assert wisdom.best_wisdom_config(records, key, "GPU B", params) == {"block_size_x": 128, "num_warps": 4} + assert wisdom.best_wisdom_config(records, key, "GPU C", params) is None + assert wisdom.best_wisdom_config(records, (8192, "torch.float32"), "GPU A", params) is None + + # results for the same key and GPU are replaced, others are kept + wisdom.write_wisdom(filename, "vector_add", params, key, "GPU A", [dict(results[0], time=0.5)]) + records = wisdom.read_wisdom(filename, params) + assert len(records) == 2 + assert wisdom.best_wisdom_config(records, key, "GPU A", params) == {"block_size_x": 128, "num_warps": 4} + + # a file of a kernel with other tunable parameters is neither used nor overwritten + assert wisdom.read_wisdom(filename, ["block_size_x"]) == [] + wisdom.write_wisdom(filename, "vector_add", ["block_size_x"], key, "GPU A", [{"block_size_x": 64, "time": 1.0}]) + assert len(wisdom.read_wisdom(filename, params)) == 2 + + +@pytest.mark.parametrize( + "import_statement, decorator", + [ + ("from kernel_tuner import autotune", "autotune"), + ("import kernel_tuner", "kernel_tuner.autotune"), + ("import kernel_tuner as kt", "kt.autotune"), + ("from kernel_tuner.decorator import autotune", "autotune"), + ("import kernel_tuner.decorator as ktd", "ktd.autotune"), + ], +) +def test_kernel_copies_drop_autotune_decorator(tmp_path, import_statement, decorator): + """The modules created for each configuration contain the kernel without the autotune decorator.""" + source = tmp_path / "kernel.py" + source.write_text( + "import functools\n" + f"{import_statement}\n" + "\n" + f"@{decorator}(tune_params={{'x': [1, 2]}})\n" + "@functools.lru_cache\n" + "def kernel(a):\n" + " return a * x\n" + ) + kernel_source = KernelSourceFn("kernel", str(source), "generic_python", call_function=lambda f, args, kwargs: None) + # only the other decorators are kept, x is replaced by the value of the tunable parameter + assert len(kernel_source.source_tree.decorator_list) == 1 + kernel_function, _ = kernel_source.apply_params_to_source_fn({"x": 3}) + assert not isinstance(kernel_function, AutotunedKernel) + assert kernel_function(2) == 6 + + +# tests that tune kernels on a GPU ------------------------------------------------------------------------------ + + +@pytest.mark.parametrize("dsl", dsls) +def test_autotune_dsl(dsl): + """The kernel is tuned when it is launched with a new tuning key, and the best configuration is launched.""" + module = kernel_module(dsl) + kernel = fresh_kernel(module) + num_configs = len([config for config in _product(kernel.tune_params)]) + + launch_and_check(module, kernel, 5000) + assert len(kernel.tuning_results) == 1 + assert len(next(iter(kernel.tuning_results.values()))) == num_configs + + # the same key is not tuned again, a new key is + launch_and_check(module, kernel, 5000) + assert len(kernel.tuning_results) == 1 + launch_and_check(module, kernel, 7000) + assert len(kernel.tuning_results) == 2 + for params in kernel.best_configs.values(): + assert all(params[name] in values for name, values in kernel.tune_params.items()) + + +@pytest.mark.parametrize("dsl", [dsl for dsl in dsls if dsl.values[0] in ("warp", "cutile")]) +def test_autotune_triton_style_launch(dsl): + """cuTile and Warp kernels can also be launched like Triton kernels, besides with launch().""" + module = kernel_module(dsl) + kernel = fresh_kernel(module) + launch_and_check(module, kernel, 5000, launch=module.launch_subscript) + launch_and_check(module, kernel, 5000) + assert len(kernel.tuning_results) == 1 + + +@pytest.mark.parametrize("dsl", dsls) +def test_autotune_parallel_compile(dsl, monkeypatch): + """Kernels are compiled in parallel while tuning, in worker processes for the DSLs that need these.""" + from kernel_tuner.backends.generic_python_processes import PROCESS_BUILD_DSLS, BuildProcessPool + + built = [] + original_build = BuildProcessPool.build + + def recording_build(self, *args, **kwargs): + built.append(original_build(self, *args, **kwargs)) + return built[-1] + + monkeypatch.setattr(BuildProcessPool, "build", recording_build) + module = kernel_module(dsl) + kernel = fresh_kernel(module, parallel_compile=2) + launch_and_check(module, kernel, 5000) + assert len(next(iter(kernel.tuning_results.values()))) == len(list(_product(kernel.tune_params))) + + # the call function can be sent to worker processes, unless the launch dimensions are lambdas (Numba here) + if dsl in PROCESS_BUILD_DSLS and dsl != "numba": + assert built and all(built) + + +def _product(tune_params): + import itertools + + return itertools.product(*tune_params.values()) + + +@skip_if_no_torch +@skip_if_no_triton +def test_autotune_wisdom(tmp_path, monkeypatch): + """Tuning results are stored in a wisdom file, which a new kernel object uses instead of tuning again.""" + module = kernel_module("triton") + kernel = fresh_kernel(module, wisdom=tmp_path) + launch_and_check(module, kernel, 5000) + filename = wisdom.wisdom_file(tmp_path, "vector_add") + records = wisdom.read_wisdom(filename, list(kernel.tune_params)) + assert len(records) == 4 + assert all(record["problem_size"] == [5000] for record in records) + + def no_tuning(*args, **kwargs): + raise AssertionError("the kernel is tuned although the wisdom file has results") + + new_kernel = fresh_kernel(module, wisdom=tmp_path) + monkeypatch.setattr(new_kernel, "_tune", no_tuning) + launch_and_check(module, new_kernel, 5000) + assert new_kernel.best_configs == kernel.best_configs + + +@skip_if_no_torch +@skip_if_no_triton +def test_autotune_triton_configs(): + """A list of triton.Config objects can be used instead of tune_params, only these are benchmarked.""" + import triton + + module = kernel_module("triton") + configs = [triton.Config({"block_size_x": 128}, num_warps=2), triton.Config({"block_size_x": 256}, num_warps=8)] + kernel = autotune(configs=configs, key=["n"])(module.vector_add.kernel) + launch_and_check(module, kernel, 5000) + results = next(iter(kernel.tuning_results.values())) + assert sorted((r["block_size_x"], r["num_warps"]) for r in results) == [(128, 2), (256, 8)] + + +@skip_if_no_torch +@skip_if_no_triton +def test_autotune_reference(): + """The outputs are verified with the reference function, wrong outputs stop tuning with an error.""" + module = kernel_module("triton") + kernel = fresh_kernel(module, reference=lambda c, a, b, n: {"c_ptr": a + b}, atol=1e-5) + launch_and_check(module, kernel, 5000) + + wrong = fresh_kernel(module, reference=lambda c, a, b, n: [a - b], atol=1e-5) + with pytest.raises(RuntimeError, match="verification failed"): + launch_and_check(module, wrong, 5000) + + +@skip_if_no_torch +@skip_if_no_numba_cuda +def test_autotune_reference_numpy(): + """Expected outputs computed with NumPy from the arrays of another DSL are compared with the outputs.""" + module = kernel_module("numba") + kernel = fresh_kernel(module, reference=lambda c, a, b, n: {"c": a.copy_to_host() + b.copy_to_host()}) + launch_and_check(module, kernel, 5000) + + +@skip_if_no_torch +@skip_if_no_triton +def test_autotune_keyword_arguments(): + """Kernel arguments can be passed by name, also while tuning.""" + import torch + import triton + + module = kernel_module("triton") + kernel = fresh_kernel(module) + c, a, b = vectors(5000) + kernel[lambda meta: (triton.cdiv(meta["n"], meta["block_size_x"]),)](c, a, b_ptr=b, n=5000) + torch.cuda.synchronize() + assert torch.allclose(c, a + b) From 64a9e961c47edfa3c97329a5b3794ab3821ed9d0 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Tue, 6 Oct 2026 15:53:11 +0200 Subject: [PATCH 37/44] explain why _NoBytecodeLoader.set_data is empty --- kernel_tuner/kernel_sources/kernel_source_fn.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/kernel_tuner/kernel_sources/kernel_source_fn.py b/kernel_tuner/kernel_sources/kernel_source_fn.py index 6517eaba0..5652598a0 100644 --- a/kernel_tuner/kernel_sources/kernel_source_fn.py +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -51,6 +51,9 @@ class _NoBytecodeLoader(importlib.machinery.SourceFileLoader): """Source loader that does not write .pyc files, kernel modules are imported only once.""" def set_data(self, path, data, *, _mode=0o666): + # Intentionally empty: SourceFileLoader only calls set_data to write the bytecode cache (.pyc) of a module + # into a __pycache__ directory next to the source file. Kernel modules are temporary files that are imported + # once, so writing their bytecode would only leave __pycache__ directories behind in the temp directory. pass From d7dcc0243b333fa515575b07c724ca8a4e43190f Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Tue, 6 Oct 2026 17:50:23 +0200 Subject: [PATCH 38/44] use a constant for the shared memory error message --- kernel_tuner/backends/nvcuda.py | 4 ++-- kernel_tuner/backends/pycuda.py | 8 ++++---- kernel_tuner/core.py | 2 +- kernel_tuner/util.py | 6 ++++++ 4 files changed, 13 insertions(+), 7 deletions(-) diff --git a/kernel_tuner/backends/nvcuda.py b/kernel_tuner/backends/nvcuda.py index 48a820592..33598e4c6 100644 --- a/kernel_tuner/backends/nvcuda.py +++ b/kernel_tuner/backends/nvcuda.py @@ -6,7 +6,7 @@ from kernel_tuner.backends.backend import GPUBackend, get_device_array, is_host_array from kernel_tuner.observers.nvcuda import CudaRuntimeObserver -from kernel_tuner.util import SkippableFailure +from kernel_tuner.util import SHARED_MEMORY_ERROR, SkippableFailure from kernel_tuner.utils.nvcuda import cuda_error_check, to_valid_nvrtc_gpu_arch_cc, find_cuda_home, _check def preload_python_nvrtc(): @@ -349,7 +349,7 @@ def load(self, kernel_instance, build_result): # Load the module err, self.current_module = driver.cuModuleLoadData(np.char.array(ptx)) if err == driver.CUresult.CUDA_ERROR_INVALID_PTX: - raise SkippableFailure("uses too much shared data") + raise SkippableFailure(SHARED_MEMORY_ERROR) else: cuda_error_check(err) diff --git a/kernel_tuner/backends/pycuda.py b/kernel_tuner/backends/pycuda.py index c5a008950..6d8959087 100644 --- a/kernel_tuner/backends/pycuda.py +++ b/kernel_tuner/backends/pycuda.py @@ -10,7 +10,7 @@ from kernel_tuner.backends.backend import GPUBackend, get_device_array, is_host_array from kernel_tuner.observers.nvml import nvml # noqa F401 from kernel_tuner.observers.pycuda import PyCudaRuntimeObserver -from kernel_tuner.util import SkippableFailure +from kernel_tuner.util import SHARED_MEMORY_ERROR, SkippableFailure # embedded in try block to be able to generate documentation # and run tests without pycuda installed @@ -213,15 +213,15 @@ def build(self, kernel_instance): no_extern_c=no_extern_c, ) except drv.CompileError as e: - if "uses too much shared data" in e.stderr: - raise SkippableFailure("uses too much shared data") + if SHARED_MEMORY_ERROR in e.stderr: + raise SkippableFailure(SHARED_MEMORY_ERROR) else: raise e except drv.LogicError as e: # DynamicSourceModule compiles to PTX, the PTX is compiled to machine code when it is linked, # at which point the CUDA driver rejects kernels that use too much shared memory if "cuLinkComplete failed: device kernel image is invalid" in str(e): - raise SkippableFailure("uses too much shared data") + raise SkippableFailure(SHARED_MEMORY_ERROR) else: raise e finally: diff --git a/kernel_tuner/core.py b/kernel_tuner/core.py index 8fb372a28..5e0b63f5b 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -636,7 +636,7 @@ def compile_kernel(self, instance, verbose, gpu_args=None, build=None): reason = f"\n{e}" else: shared_mem_error_messages = [ - "uses too much shared data", + util.SHARED_MEMORY_ERROR, "local memory limit exceeded", r"local memory \(\d+\) exceeds limit \(\d+\)", ] diff --git a/kernel_tuner/util.py b/kernel_tuner/util.py index 5476aa330..db3a4fbbc 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -116,6 +116,12 @@ class SkippableFailure(Exception): """Exception used to raise when compiling or launching a kernel fails for a reason that can be expected.""" +# Part of the error message of ptxas for kernels that use too much shared memory. Backends also raise +# SkippableFailure with this message when a kernel uses too much shared memory, so that these configurations +# are recognized as skippable by DeviceInterface.compile_kernel. +SHARED_MEMORY_ERROR = "uses too much shared data" + + class StopCriterionReached(Exception): """Exception thrown when a stop criterion has been reached.""" From 765a0baa4054f5c25bef1261145c231f4294679e Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 7 Oct 2026 10:08:26 +0200 Subject: [PATCH 39/44] use separate names for the launch dimensions computed by the tuning call function --- kernel_tuner/decorator.py | 14 ++++++++++++-- test/test_decorator.py | 34 +++++++++++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/kernel_tuner/decorator.py b/kernel_tuner/decorator.py index 6b7f2c742..aad26233b 100644 --- a/kernel_tuner/decorator.py +++ b/kernel_tuner/decorator.py @@ -266,12 +266,22 @@ def __init__(self, dsl, argument_names, grid, threads, spec, extra_kwargs): self._states = {} def __call__(self, kernel_function, args, kwargs, grid, threads, params): - grid, threads = _launch_dimensions( + """Launch the kernel of a configuration, see :func:`kernel_tuner.util.normalize_call_function`. + + The grid and threads computed by Kernel Tuner are not used: the decorator calls tune_kernel with a problem + size of 1, because the launch dimensions are those of the user's launch or the decorator, which can depend + on the kernel arguments and the tunable parameters. The arguments are part of the signature nonetheless, + as Kernel Tuner only passes params to call functions that have arguments named grid, threads, and params. + """ + launch_grid, launch_threads = _launch_dimensions( self.grid, self.threads, self.spec, self.argument_names, args, self.extra_kwargs, params ) state = self._states.setdefault(id(kernel_function), {}) launch_function = _LAUNCH_FUNCTIONS[self.dsl] - launch_function(kernel_function, args, {**self.extra_kwargs, **kwargs}, grid, threads, self.spec, params, state) + launch_function( + kernel_function, args, {**self.extra_kwargs, **kwargs}, launch_grid, launch_threads, self.spec, params, + state + ) def __getstate__(self): state = self.__dict__.copy() diff --git a/test/test_decorator.py b/test/test_decorator.py index cd3e2b60d..79269f236 100644 --- a/test/test_decorator.py +++ b/test/test_decorator.py @@ -7,7 +7,7 @@ import kernel_tuner from kernel_tuner import autotune -from kernel_tuner.decorator import AutotunedKernel, _configs_to_search_space +from kernel_tuner.decorator import AutotunedKernel, _configs_to_search_space, _Launch, _TuningCallFunction from kernel_tuner.kernel_sources.kernel_source_fn import KernelSourceFn from kernel_tuner.utils import wisdom @@ -122,6 +122,38 @@ def test_autotune_is_top_level(): assert "autotune" in kernel_tuner.__all__ +def _grid_of_meta(meta): + return (meta["n"] // meta["block_size_x"],) + + +def test_tuning_call_function(monkeypatch): + """The tuning call function launches kernels like the user, ignoring the launch dimensions of Kernel Tuner.""" + import pickle + + from kernel_tuner import decorator + from kernel_tuner.util import normalize_call_function + + launches = [] + + def recording_launch(kernel, args, kwargs, grid, threads, launch, params, state): + launches.append((kernel, args, kwargs, grid, threads)) + + monkeypatch.setitem(decorator._LAUNCH_FUNCTIONS, "triton", recording_launch) + call_function = _TuningCallFunction("triton", ["c", "a", "b", "n"], None, (128, 1, 1), _Launch(_grid_of_meta), {}) + + # Kernel Tuner passes params only to call functions with arguments named grid, threads, and params + assert normalize_call_function(call_function) is call_function + + params = {"block_size_x": 256} + call_function("kernel", [1, 2, 3, 4096], {"block_size_x": 256}, (1, 1, 1), (1, 1, 1), params) + assert launches == [("kernel", [1, 2, 3, 4096], {"block_size_x": 256}, (16,), (128, 1, 1))] + + # the call function can be sent to worker processes when its launch dimensions can be pickled + copy = pickle.loads(pickle.dumps(call_function)) + copy("kernel", [1, 2, 3, 2048], {}, (1, 1, 1), (1, 1, 1), params) + assert launches[-1][3] == (8,) + + def test_autotune_needs_tune_params_or_configs(): with pytest.raises(ValueError): autotune() From c58ee867545d7b144c318648b6345655419b9fa1 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 7 Oct 2026 09:09:57 +0000 Subject: [PATCH 40/44] keep the locked clocks when changing one of the clocks with locked clocks With use_locked_clocks, setting the graphics clock locked the memory clock at its current reading and vice versa. The current reading can differ from the locked clock (9501 MHz for a memory clock locked at 10001 MHz on an RTX 6000 Ada, or a low idle graphics clock), so tuning nvml_gr_clock silently changed the memory clock, and tuning both nvml_gr_clock and nvml_mem_clock made the setters overwrite each other's clock with stale values. --- kernel_tuner/observers/nvml.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/kernel_tuner/observers/nvml.py b/kernel_tuner/observers/nvml.py index eb5bbaf42..86ed7e696 100644 --- a/kernel_tuner/observers/nvml.py +++ b/kernel_tuner/observers/nvml.py @@ -265,7 +265,8 @@ def gr_clock(self, new_clock): """Control the graphics clock (may require permission), only values compatible with the memory clock can be set directly.""" if self.use_locked_clocks: if new_clock != self.locked_gr_clock: - self.set_clocks(self.mem_clock, new_clock) + # keep the locked memory clock, the current memory clock can differ from the locked clock + self.set_clocks(self.locked_mem_clock or self.mem_clock, new_clock) else: # if using applications clocks if new_clock != pynvml.nvmlDeviceGetApplicationsClock(self.dev, pynvml.NVML_CLOCK_GRAPHICS): @@ -280,7 +281,8 @@ def mem_clock(self): def mem_clock(self, new_clock): if self.use_locked_clocks: if new_clock != self.locked_mem_clock: - self.set_clocks(new_clock, self.gr_clock) + # keep the locked graphics clock, the current graphics clock can differ from the locked clock + self.set_clocks(new_clock, self.locked_gr_clock or self.gr_clock) # if using applications clocks else: if new_clock != pynvml.nvmlDeviceGetApplicationsClock(self.dev, pynvml.NVML_CLOCK_MEM): From 8a55c81528b74b509b34dc321198a2820a2fa411 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Wed, 7 Oct 2026 10:04:10 +0000 Subject: [PATCH 41/44] use the device argument in the generic Python backend The backend used the current CUDA device instead of the device passed to tune_kernel. It now makes the requested device the current device, on which the DSLs launch their kernels, and copies tensor arguments to that device. --- kernel_tuner/backends/generic_python.py | 11 +++++--- test/test_generic_python_functions.py | 35 +++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index 6df2ea720..c7238283f 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -56,6 +56,9 @@ def __init__(self, device=0, iterations=7, compiler_options=None, observers=None logging.error("Unable to import Torch") raise ImportError("Unable to import Torch") + # make the requested device the current device: the DSLs launch kernels on the current device, and + # PyTorch creates streams, events, and tensors on it + torch.cuda.set_device(device) self.device_id = torch.cuda.current_device() self.device_properties = torch.cuda.get_device_properties(self.device_id) self.name = torch.cuda.get_device_name(self.device_id) @@ -115,12 +118,12 @@ def ready_argument_list(self, arguments): if arg.dim() == 0: # Scalar tensor, convert to Python scalar torch_args.append(arg.item()) else: - if arg.is_cuda: # already on GPU, need deep copy to not overwrite - torch_args.append(arg.clone()) + if arg.is_cuda: # already on GPU, need deep copy to not overwrite, on the tuned device + torch_args.append(arg.to(self.device_id, copy=True)) else: # Copy from CPU to GPU - torch_args.append(arg.contiguous().to("cuda")) + torch_args.append(arg.contiguous().to(self.device_id)) elif isinstance(arg, np.ndarray): # Convert Numpy CPU array to Torch GPU Tensor - torch_args.append(torch.from_numpy(arg).to("cuda")) + torch_args.append(torch.from_numpy(arg).to(self.device_id)) elif isinstance(arg, np.generic): # Numpy scalar, convert to Python scalar torch_args.append(arg.item()) elif isinstance(arg, (int, float, bool, str)): # Is already Python diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py index 495457c15..04b9614c2 100644 --- a/test/test_generic_python_functions.py +++ b/test/test_generic_python_functions.py @@ -3,6 +3,7 @@ from kernel_tuner.core import DeviceInterface, KernelInstance from kernel_tuner.kernel_sources.kernel_source import KernelSource import numpy as np +import pytest from pathlib import Path import os @@ -127,3 +128,37 @@ def test_gpu_kwargs(): assert kwargs["mock_param"] == 64 assert callable_fn(*args, **kwargs) == 64 + + +@skip_if_no_torch +def test_ready_argument_list_device_and_strides(): + ks, _, _ = get_context() + dev = DeviceInterface(ks, device=0) + # a CUDA tensor with permuted strides, a CPU tensor, and a NumPy array + permuted = torch.randn(4, 6, device="cuda").t() + cpu = torch.randn(5) + array = np.arange(3, dtype=np.float32) + + gpu_args = dev.ready_argument_list([permuted, cpu, array]) + + assert all(arg.device == torch.device("cuda", 0) for arg in gpu_args) + # the copy keeps the layout of the CUDA tensor + assert gpu_args[0].stride() == permuted.stride() + assert gpu_args[0] is not permuted and torch.equal(gpu_args[0], permuted) + + +@skip_if_no_torch +def test_device_argument(): + if torch.cuda.device_count() < 2: + pytest.skip("needs two GPUs") + ks, args, _ = get_context() # b is created on the default device 0 + current = torch.cuda.current_device() + try: + dev = DeviceInterface(ks, device=1) + assert dev.dev.device_id == 1 + assert torch.cuda.current_device() == 1 + gpu_args = dev.ready_argument_list(args) + assert gpu_args[1].device == torch.device("cuda", 1) + assert torch.equal(gpu_args[1].cpu(), args[1].cpu()) + finally: + torch.cuda.set_device(current) From c7fc6750d4d8adf6445b88ab80618915062933b3 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Thu, 8 Oct 2026 13:28:44 +0200 Subject: [PATCH 42/44] trying to get the CI working again --- .github/workflows/test-python-package.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-python-package.yml b/.github/workflows/test-python-package.yml index 60b51c324..d1a151032 100644 --- a/.github/workflows/test-python-package.yml +++ b/.github/workflows/test-python-package.yml @@ -21,7 +21,7 @@ jobs: strategy: matrix: - os: [ubuntu-latest, macos-13] + os: [ubuntu-latest, macos-latest] steps: - uses: actions/checkout@v7 From 496eaa7698209cfbabdd93d164613492063fd12a Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Thu, 8 Oct 2026 15:04:16 +0000 Subject: [PATCH 43/44] exclude the host launch time from the measured time of Python DSL kernels The start event was recorded on an idle GPU before the call function, so the measured time included the time the host needs to convert the arguments and launch the kernel: 5 to 64 microseconds for the GEMM kernels of the DSLs, up to 8% of a 0.8 ms kernel, and a source of noise. A sleep kernel before the start event now keeps the GPU busy until the kernel has been queued. --- kernel_tuner/backends/generic_python.py | 15 ++++++++++++++- test/test_generic_python_functions.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/kernel_tuner/backends/generic_python.py b/kernel_tuner/backends/generic_python.py index c7238283f..9762ab879 100644 --- a/kernel_tuner/backends/generic_python.py +++ b/kernel_tuner/backends/generic_python.py @@ -34,6 +34,11 @@ PROCESS_BUILT = "process_built" +# Number of GPU clock cycles that the sleep kernel in start_event() keeps the GPU busy before a measurement: 0.6 ms at +# 3 GHz, much longer than the time needed to launch a kernel from Python +LAUNCH_SLEEP_CYCLES = 2_000_000 + + class GenericPythonFunctions(GPUBackend): """Class that groups the Python DSL functions on maintains state about the device.""" @@ -243,7 +248,15 @@ def _kernel_kwargs(signature, params): return {arg_name: params[arg_name] for arg_name in signature if arg_name in params} def start_event(self): - """Records the event that marks the start of a measurement.""" + """Records the event that marks the start of a measurement. + + A sleep kernel before the event keeps the GPU busy while the host launches the kernel that is measured. The + start event completes when the sleep kernel finishes, after the kernel has been queued, so the measured time + does not include the time the host needs to launch the kernel (converting the arguments and launching the + kernel takes tens of microseconds for some DSLs). On an idle GPU, the event would complete immediately and + the launch time would be measured as part of the kernel. + """ + torch.cuda._sleep(LAUNCH_SLEEP_CYCLES) self.start.record() def stop_event(self): diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py index 04b9614c2..19eada417 100644 --- a/test/test_generic_python_functions.py +++ b/test/test_generic_python_functions.py @@ -162,3 +162,21 @@ def test_device_argument(): assert torch.equal(gpu_args[1].cpu(), args[1].cpu()) finally: torch.cuda.set_device(current) + + +@skip_if_no_torch +def test_start_event_excludes_launch_time(): + import time + + ks, _, _ = get_context() + dev = DeviceInterface(ks).dev + x = torch.ones(16, device="cuda") + x.add_(1) # load the kernel, as Kernel Tuner runs a kernel before it is benchmarked + torch.cuda.synchronize() + dev.start_event() + time.sleep(0.0005) # the host takes 0.5 ms to launch the kernel + x.add_(1) + dev.stop_event() + torch.cuda.synchronize() + # only the kernel is measured, not the time before it was launched + assert dev.start.elapsed_time(dev.end) < 0.25 From 7f94c70c8d1fe7383c766b570ac28a7a2d1ee721 Mon Sep 17 00:00:00 2001 From: Ben van Werkhoven Date: Fri, 9 Oct 2026 11:13:05 +0200 Subject: [PATCH 44/44] fix two typos in the backend documentation --- doc/source/backends.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/source/backends.rst b/doc/source/backends.rst index 99bf7db95..2008304da 100644 --- a/doc/source/backends.rst +++ b/doc/source/backends.rst @@ -18,13 +18,13 @@ CUDA Backends Kernel Tuner automatically selects a CUDA backend when it detects CUDA as the language of the kernel to be tuned. Depending on whether the required dependencies are installed it first checks -if it could use the ``cuda-python`` backend, then ``cupy``, then PyCUDA. This behavior can be overwritten +if it could use the ``cuda-python`` backend, then ``cupy``, then PyCUDA. This behavior can be overridden by manually selecting a backend with the ``lang=`` option. Because the HIP kernel language is identical to the CUDA kernel language, HIP is included here as well. To use HIP on nvidia GPUs, see https://github.com/jatinx/hip-on-nv. HIP must be indicated manually using the ``lang=`` option. -Kernel inputs and outputs can be handeld slightly differently by different backends, but the default +Kernel inputs and outputs can be handled slightly differently by different backends, but the default way to pass arguments to the kernel is through Numpy arrays, which is supported by all backends. For example, while the PyCUDA backend expects all inputs and outputs to be Numpy arrays, the CuPy backend also supports cupy arrays as input and output arguments for the kernels. This gives the user more control