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 5f8e66ef3..530170a9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,17 @@ 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 +- 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 +- 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/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/backends.rst b/doc/source/backends.rst index 09d3b1574..2008304da 100644 --- a/doc/source/backends.rst +++ b/doc/source/backends.rst @@ -6,8 +6,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. @@ -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 @@ -74,3 +74,163 @@ points to the root directory of the CUDA installation. For example, by adding 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. + + +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/design.rst b/doc/source/design.rst index 8d8b9b357..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 -------- @@ -150,3 +156,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/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/doc/source/parallel.rst b/doc/source/parallel.rst index 0fdd557a4..649c5398a 100644 --- a/doc/source/parallel.rst +++ b/doc/source/parallel.rst @@ -148,3 +148,59 @@ 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, 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, 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 + 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. + +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/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/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/__init__.py b/examples/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/python/README.rst b/examples/python/README.rst new file mode 100644 index 000000000..a1259a7ca --- /dev/null +++ b/examples/python/README.rst @@ -0,0 +1,179 @@ +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 with 2-dimensional thread blocks + - use the restrictions option to limit the search to valid thread block sizes + +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 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 a class-based kernel + - 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, 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 + - 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/cute_vec_add.py b/examples/python/cute_vec_add.py new file mode 100644 index 000000000..36baeb0f1 --- /dev/null +++ b/examples/python/cute_vec_add.py @@ -0,0 +1,57 @@ +import torch + +import cutlass +import cutlass.cute as cute + +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 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) + + args = [a, b, c, size] + 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, + verbose=True) + + +if __name__ == "__main__": + main() + diff --git a/examples/python/cutile_vec_add.py b/examples/python/cutile_vec_add.py new file mode 100644 index 000000000..003ae2798 --- /dev/null +++ b/examples/python/cutile_vec_add.py @@ -0,0 +1,32 @@ +import cuda.tile as ct +import torch + +from kernel_tuner import tune_kernel + + +@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"]) + + +if __name__ == "__main__": + tune() 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/matmul/cupy_matmul.py b/examples/python/matmul/cupy_matmul.py new file mode 100644 index 000000000..d3900b09f --- /dev/null +++ b/examples/python/matmul/cupy_matmul.py @@ -0,0 +1,69 @@ +import cupy as cp +from cupyx import jit +import numpy as np + +from kernel_tuner import tune_kernel + + + +@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 + + +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) + + # block / grid configuration + block = (16, 16) + grid = ((N + block[0] - 1) // block[0], (M + block[1] - 1) // block[1]) + + # launch kernel + gemm[grid, block](a, b, c, M, N, K) + cp.cuda.Device().synchronize() + + # Correctness verification + c_ref = cp.matmul(a, b) + assert cp.allclose(c, c_ref, rtol=1e-2, atol=M * 2**(-11)) + + print("Succes") + + +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(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=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + 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 = 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 new file mode 100644 index 000000000..f5e8f47e8 --- /dev/null +++ b/examples/python/matmul/cute_matmul.py @@ -0,0 +1,97 @@ +import torch + +import cutlass +import cutlass.cute as cute +from cutlass.cute.runtime import from_dlpack + +from kernel_tuner import tune_kernel + + +# might need export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH to work + +## Basic Matmul ================================================================ + +@cute.kernel +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() + + # Global indices + n = bx * bdx + tx + m = by * bdy + ty + + if m < M and n < N: + acc = cutlass.Float32(0.0) + + for k in range(K): + a = gA[m, k].to(cutlass.Float32) + b = gB[k, n].to(cutlass.Float32) + acc += a * b + + 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) + + M, K = mA.shape + _, N = mB.shape + + grid = ( + (N + block[0] - 1) // block[0], + (M + block[1] - 1) // block[1], + 1, + ) + + kernel = naive_matmul_kernel(mA, mB, mC, M, K, N) + + 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) + b = torch.randn(K, N, device="cuda", dtype=torch.float16) + c = torch.zeros(M, N, device="cuda", dtype=torch.float16) + + args = [a, b, c] + size = M * N + answer = [None, None, (a @ b).cpu()] + 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"] + + results, env = tune_kernel("matmul", __file__, size, args, tune_params, + answer=answer, atol=M * 2 **(-11), restrictions=restrictions, verbose=False) + + +if __name__ == "__main__": + m, n, k = 4096, 4096, 4096 + run_naive_matmul(m, n, k) + tune_naive_matmul(m, n, k) diff --git a/examples/python/matmul/cutile_matmul.py b/examples/python/matmul/cutile_matmul.py new file mode 100644 index 000000000..9ae9ef3e0 --- /dev/null +++ b/examples/python/matmul/cutile_matmul.py @@ -0,0 +1,56 @@ +import cuda.tile as ct +import torch + +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) + 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)) + + +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"]) + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + run(M, N, K) + tune(M, N, K) diff --git a/examples/python/matmul/numba_matmul.py b/examples/python/matmul/numba_matmul.py new file mode 100644 index 000000000..d742b1224 --- /dev/null +++ b/examples/python/matmul/numba_matmul.py @@ -0,0 +1,84 @@ +import numpy as np +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): + 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 + + +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.astype(np.float32) @ B.astype(np.float32)).astype(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 + 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, 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) + + size = (M, N) + args = [A, B, C] + 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"] + + answer = [None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16)] + atol = M * 2**(-11) + + results, env = tune_kernel( + kernel_name="matmul", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + answer=answer, + atol=atol, + restrictions=restrictions, + ) + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + + run_basic(M, N, K) + tune_basic(M, N, K) diff --git a/examples/python/matmul/taichi_matmul.py b/examples/python/matmul/taichi_matmul.py new file mode 100644 index 000000000..3568eb00f --- /dev/null +++ b/examples/python/matmul/taichi_matmul.py @@ -0,0 +1,61 @@ +import numpy as np +import taichi as ti + +from kernel_tuner import tune_kernel + + +ti.init(arch=ti.gpu) + +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)): + 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): + sum += A[i, k] * B[k, j] + C[i, j] = ti.cast(sum, ti.f16) + + +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.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) + + 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): + 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 = {"BLOCK_DIM": [2**i for i in range(5, 11)]} + + answer = [None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16)] + + results, env = tune_kernel( + kernel_name="matmul", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + answer=answer, + atol=M * 2**(-11), + block_size_names = ["BLOCK_DIM"], + ) + +if __name__ == "__main__": + run(4096, 4096, 4096) + tune(4096, 4096, 4096) + + + diff --git a/examples/python/matmul/tilelang_matmul.py b/examples/python/matmul/tilelang_matmul.py new file mode 100644 index 000000000..740b27c1e --- /dev/null +++ b/examples/python/matmul/tilelang_matmul.py @@ -0,0 +1,102 @@ +import torch +import tilelang +import tilelang.language as T + +from kernel_tuner import tune_kernel + + +# 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 + def gemm( + A: T.Tensor((M, K), dtype), + B: T.Tensor((K, N), dtype), + C: T.Tensor((M, N), dtype), + ): + 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) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + T.clear(C_local) + + # 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=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) + + T.copy(C_local, C[by * block_M, bx * block_N]) + + return gemm + + +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 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(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] + + 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, + 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) + tune_basic(M, N, K) diff --git a/examples/python/matmul/tilus_matmul.py b/examples/python/matmul/tilus_matmul.py new file mode 100644 index 000000000..657690a4e --- /dev/null +++ b/examples/python/matmul/tilus_matmul.py @@ -0,0 +1,119 @@ +import torch + +import tilus +from tilus import float16, float32, int32 +from tilus.utils import cdiv + +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 +# +# Original example: matmul_v0.py +# Copyright (c) the Tilus authors +class MatmulBasic(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, 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 + cdiv(n_size, self.block_n), # the y dimension size of the grid + ] + 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 + 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 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 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"] = [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] + + + results, env = tune_kernel( + kernel_name="MatmulBasic", + 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), + strategy = "bayes_opt", + strategy_options = {"max_fevals": 200}, + ) + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + run_basic(M, N, K) + tune_basic(M, N, K) diff --git a/examples/python/matmul/triton_matmul.py b/examples/python/matmul/triton_matmul.py new file mode 100644 index 000000000..e43fb7c6f --- /dev/null +++ b/examples/python/matmul/triton_matmul.py @@ -0,0 +1,115 @@ +import torch +import triton +import triton.language as tl + +from kernel_tuner import tune_kernel + + +@triton.jit +def matmul_basic( + A_ptr, B_ptr, C_ptr, + M, N, K, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, +): + m = tl.program_id(0) + n = tl.program_id(1) + + # 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) + + # Accumulator for C tile + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + # Loop over K in chunks + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + offs_k = k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + + # 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] + + mask_a = (offs_m[:, None] < M) & (offs_k[None, :] < K) + mask_b = (offs_k[:, None] < K) & (offs_n[None, :] < N) + + 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) + + +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, + 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 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, + ] + + 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["num_warps"] = [2, 4, 8, 16] + 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, + 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}, + ) + + + + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + run_basic(M, N, K) + tune_basic(M, N, K) diff --git a/examples/python/matmul/warp_matmul.py b/examples/python/matmul/warp_matmul.py new file mode 100644 index 000000000..660de9727 --- /dev/null +++ b/examples/python/matmul/warp_matmul.py @@ -0,0 +1,89 @@ +import numpy as np +import warp as wp + +from kernel_tuner import tune_kernel + + +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) + + +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) + 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(gemm, dim=(M, N), inputs=[A_wp, B_wp, C_wp]) + + 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") + + +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) + C = np.zeros((M, N), dtype=np.float16) + + size = (M, N) + + tune_params = dict() + tune_params["block_dim"] = [2**i for i in range(5, 11)] + tune_params["dim"] = [size] + + args = [A, B, C] + answer = [None, None, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16)] + + results, env = tune_kernel( + kernel_name="gemm", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + answer=answer, + atol=M * 2**(-11), + ) + + + +if __name__ == "__main__": + M, N, K = 4096, 4096, 4096 + + run_gemm(M, N, K) + tune(M, N, K) + + + + + + + + + diff --git a/examples/python/numba_vec_add.py b/examples/python/numba_vec_add.py new file mode 100644 index 000000000..df39d6f70 --- /dev/null +++ b/examples/python/numba_vec_add.py @@ -0,0 +1,39 @@ +from numba import cuda +import numpy as np + +from kernel_tuner import tune_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 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)]} + + + results, env = tune_kernel( + kernel_name="f", + kernel_source=__file__, + problem_size=N, + arguments=args, + tune_params=tune_params, + answer=[None, None, c_expect], + ) + +if __name__ == "__main__": + tune() \ No newline at end of file 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.py b/examples/python/tilelang_vec_add.py new file mode 100644 index 000000000..32d389d24 --- /dev/null +++ b/examples/python/tilelang_vec_add.py @@ -0,0 +1,42 @@ +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 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", __file__, N, args, tune_params, answer=answer) + +if __name__ == "__main__": + tune() \ No newline at end of file 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/tilus_vec_add.py b/examples/python/tilus_vec_add.py new file mode 100644 index 000000000..e509b307e --- /dev/null +++ b/examples/python/tilus_vec_add.py @@ -0,0 +1,97 @@ +import tilus +from tilus import float32, int32 +from tilus.utils import cdiv +import torch + +from kernel_tuner import tune_kernel, run_kernel + + + + +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 elements per block + self.num_warps = num_warps # number of warps 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_x)] + self.attrs.warps = self.num_warps + + # 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=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]) + + 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 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() + c_expect = a + b + + + 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( + kernel_name="VecAddV", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + answer=[None, None, None, c_expect.cpu()], + verbose=True, + ) + + +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() + c_expect = a + b + + + args = [size, a, b, c] + + results = run_kernel( + kernel_name="VecAddV", + kernel_source=__file__, + problem_size=size, + arguments=args, + params={"block_size_x": 32, "num_warps": 4}, + ) + + c_expect = c_expect.cpu() + + assert torch.allclose(results[-1], c_expect) + + + +if __name__ == "__main__": + size = 10000000 + tune(size) + run(size) + 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.py b/examples/python/triton_vec_add.py new file mode 100644 index 000000000..e82642317 --- /dev/null +++ b/examples/python/triton_vec_add.py @@ -0,0 +1,60 @@ +import numpy as np +import torch +import triton +import triton.language as tl + +from kernel_tuner import tune_kernel, run_kernel + +@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 tune(): + 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(11)] + + + result = run_kernel("add_kernel", __file__, size, args, {"block_size_x": 256}) + assert np.allclose(c_expect.cpu(), result[2]) + + + results, env = tune_kernel( + kernel_name="add_kernel", + kernel_source=__file__, + problem_size=size, + arguments=args, + tune_params=tune_params, + answer=[None, None, c_expect.cpu(), None], + ) + + +if __name__ == "__main__": + tune() 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.py b/examples/python/warp_vec_add.py new file mode 100644 index 000000000..aee5a96b3 --- /dev/null +++ b/examples/python/warp_vec_add.py @@ -0,0 +1,70 @@ +import warp as wp +import numpy as np +import torch + +from kernel_tuner import tune_kernel + + +wp.init() + + +@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 = 8 + 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]) + + +def call_warp(kernel_function, args, kwargs, grid, threads, params): + warp_args = [] + for arg in args: + if isinstance(arg, torch.Tensor): + warp_args.append(wp.from_torch(arg)) + else: + warp_args.append(arg) + # 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) + + +def tune(): + n = 1024 + + # Create host arrays + 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)] + tune_params["size"] = [n] + args = [a, b, c, n] + + results, env = tune_kernel( + kernel_name="vec_add", + kernel_source=__file__, + problem_size=n, + arguments=args, + tune_params=tune_params, + answer=[None, None, c_expect, None], + call_function=call_warp, + ) + + +if __name__ == "__main__": + tune() \ No newline at end of file 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/backends/backend.py b/kernel_tuner/backends/backend.py index 6063dbb43..65f483499 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.""" @@ -17,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 new file mode 100644 index 000000000..9762ab879 --- /dev/null +++ b/kernel_tuner/backends/generic_python.py @@ -0,0 +1,477 @@ +import logging +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: + import torch +except ImportError: + 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. +# 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. +# 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" + + +# 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.""" + + 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. + + :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 + """ + + if not torch: + 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) + self.max_threads = 10**18 # 'inf' to support tile based programming models, which can use less threads then the size of a tile. + + env = {} + 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"} + + # Set when a kernel is loaded, needed for running the kernel after compilation. + self.call_function = None + self.signature = 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) + + def ready_argument_list(self, 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, 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(self.device_id)) + elif isinstance(arg, np.ndarray): # Convert Numpy CPU array to Torch GPU Tensor + 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 + torch_args.append(arg) + elif type(arg) in vars(builtins).values(): + torch_args.append(arg) + else: + 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 + + :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 + """ + 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 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) + 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") + + # If the kernel source is a class, we use the __call__ function as the kernel_function. + kernel_function = instantiate_kernel(kernel_instance.kernel_fn) + + 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() + 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. + + 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): + """Records the event that marks the end of a measurement.""" + self.end.record() + + 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, 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) + + :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 + """ + self.call_function(func, gpu_args, self._kernel_kwargs(self.signature, params), grid, threads, params) + + + 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, 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. + + 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 + + :returns: "resource_error" , "user_error" or "unknown" + :rtype: string + """ + + RESOURCE_PATTERNS = ( + # Shared memory + r"shared memory", + r"\bsmem\b", + r"too much shared", + + # Registers / occupancy + r"too many registers", + r"register spill", + r"out of registers", + r"register usage", + + # Launch configuration + 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", + r"power of (two|2)", # many tile based DSLs require power of two tile sizes + + # Generic resource exhaustion + 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", + + # 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", + r"\bllvm error", + r"\bmlir error", + r"lowering failed", + ) + + USER_ERROR_PATTERNS = ( + # Undefined / missing symbols + r"not defined", + r"undefined variable", + r"without definition", + r"unknown variable", + r"\bunbound\b", + + # Type / shape errors + r"type mismatch", + r"invalid type", + r"cannot convert", + r"expected .* but got", + r"incompatible types", + + # Indexing / bounds + r"index out of bounds", + r"out of bounds access", + r"invalid index", + + # IR / AST construction + r"failed to build", + r"invalid expression", + r"malformed", + r"illegal operation", + + # Frontend components + r"\btranspiler\b", + r"\bfrontend\b", + r"\bast\b", + ) + + # Errors in the kernel code that no tuning parameter value can cause (IndentationError is a SyntaxError) + CODE_ERROR_TYPES = (NameError, UnboundLocalError, SyntaxError) + + # 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) + + if any(isinstance(exc, CODE_ERROR_TYPES) for exc in chain): + return "user_error" + + if match_any(RESOURCE_PATTERNS, msg): + return "resource_error" + + if any(isinstance(exc, USER_ERROR_TYPES) for exc in chain): + return "user_error" + + if match_any(USER_ERROR_PATTERNS, msg): + return "user_error" + + return "unknown" diff --git a/kernel_tuner/backends/generic_python_processes.py b/kernel_tuner/backends/generic_python_processes.py new file mode 100644 index 000000000..9f8ca300d --- /dev/null +++ b/kernel_tuner/backends/generic_python_processes.py @@ -0,0 +1,369 @@ +"""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, 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. + 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``. + +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 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 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", "tilus"} + +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 + + :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() + 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 not_stored + else: + 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): + """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, 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): + 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: + if str(pickle.loads(pickle.dumps(exc))) == str(exc): + picklable.append(exc) + continue + except Exception: + pass + 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 + + +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, 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 + + 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) + + 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) 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() + 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: + result = 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 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 + + 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 c5f9b577d..33598e4c6 100644 --- a/kernel_tuner/backends/nvcuda.py +++ b/kernel_tuner/backends/nvcuda.py @@ -4,9 +4,9 @@ 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.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(): @@ -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) @@ -255,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) @@ -307,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) @@ -335,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(SHARED_MEMORY_ERROR) + 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): @@ -454,12 +480,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 +497,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 1fc8c6ab1..6d8959087 100644 --- a/kernel_tuner/backends/pycuda.py +++ b/kernel_tuner/backends/pycuda.py @@ -3,13 +3,14 @@ from __future__ import print_function import logging +import threading 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 SHARED_MEMORY_ERROR, SkippableFailure # embedded in try block to be able to generate documentation # and run tests without pycuda installed @@ -20,8 +21,7 @@ except ImportError: class PyCudaPlaceHolder: - def __init__(self): - self.PointerHolderBase = object + pass drv = PyCudaPlaceHolder() pycuda_available = False @@ -35,23 +35,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 +140,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 +149,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)) @@ -189,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, @@ -217,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") + 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(SHARED_MEMORY_ERROR) + 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.""" @@ -368,16 +386,13 @@ def memset(self, allocation, value, size): def memcpy_dtoh(self, 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: 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): """Perform a host to device memory copy. @@ -385,11 +400,15 @@ def memcpy_htod(self, dest, src): :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 4a7156962..5e0b63f5b 100644 --- a/kernel_tuner/core.py +++ b/kernel_tuner/core.py @@ -20,6 +20,7 @@ def _get_cupy(): import kernel_tuner.util as util from kernel_tuner.accuracy import Tunable from kernel_tuner.backends.backend import GPUBackend +from kernel_tuner.kernel_sources.kernel_source import KernelSource # noqa: F401 (re-exported as core.KernelSource) from kernel_tuner.observers.observer import BenchmarkObserver, ContinuousObserver, OutputObserver, PrologueObserver from kernel_tuner.observers.tegra import TegraObserver @@ -40,6 +41,7 @@ def _get_cupy(): "name", "kernel_source", "kernel_string", + "kernel_fn", "temp_files", "threads", "grid", @@ -49,16 +51,34 @@ 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.""" + def __new__(cls, *args, **kwargs): + # 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) < 9: + 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 [] # already done during compilation + temp_filename = util.get_temp_filename(suffix=self.kernel_source.get_suffix()) util.write_file(temp_filename, self.kernel_string) ret = [temp_filename] @@ -66,177 +86,8 @@ 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", "JULIA": ".jl"} - 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. +# KernelSource has been moved to kernel_sources/kernel_source_str.py to support Generic Python - 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, lang=self.lang) - else: - logging.debug("Checking of arguments list not supported yet for code generators.") - - def infer_julia_backend(self): - """Infer the Julia backend from the kernel source.""" - backend = None - if self.lang.upper() != "JULIA": - return backend - - kernel_string = self.get_kernel_string(0) - if kernel_string.find("using CUDA") != -1: - backend = "cuda" - elif kernel_string.find("using ROCBackend") != -1: - backend = "amd" - elif kernel_string.find("using oneAPI") != -1: - backend = "intel" - elif kernel_string.find("using Metal") != -1: - backend = "metal" - else: - raise ValueError("Could not infer Julia backend from kernel source, provide it as a `compiler_option`") - return backend def instantiate_observer(observer, args): @@ -321,6 +172,7 @@ def __init__( logging.debug("DeviceInterface instantiated, lang=%s", lang) + # Ensure observers is a list observers = observers or [] @@ -371,9 +223,13 @@ def __init__( backend = CompilerFunctions backend_options["compiler"] = compiler backend_options["observers"] = observers + elif lang.upper() == "GENERIC_PYTHON": + from kernel_tuner.backends.generic_python import GenericPythonFunctions + + backend = GenericPythonFunctions else: raise NotImplementedError( - "Sorry, support for languages other than CUDA, OpenCL, HIP, C, Julia, Fortran is not implemented yet" + "Sorry, support for languages other than CUDA, OpenCL, HIP, C, Julia, Fortran and Generic Python is not implemented yet" ) if issubclass(backend, GPUBackend): @@ -427,17 +283,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 self.lang.upper() == "GENERIC_PYTHON": # Generic Python needs params to launch + 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): @@ -445,7 +307,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() @@ -457,10 +319,12 @@ def benchmark_default(self, func, gpu_args, threads, grid, result): 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()) - 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() @@ -468,7 +332,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() @@ -520,15 +384,15 @@ 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) duration = 1 for obs in self.continuous_observers: obs.results = result duration = max(duration, obs.continuous_duration) if len(self.continuous_observers) > 0: - 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 @@ -565,6 +429,7 @@ def benchmark(self, func, gpu_args, instance, verbose, objective, skip_nvml_sett def check_kernel_output(self, func, gpu_args, instance, answer, atol, verify, verbose): """Runs the kernel once and checks the result against answer.""" logging.debug("check_kernel_output") + cp = _get_cupy() # get the answer for this parameter configuration if isinstance(answer, Tunable): @@ -572,7 +437,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) @@ -587,7 +452,6 @@ def check_kernel_output(self, func, gpu_args, instance, answer, atol, verify, ve answer[i] is not None or util.is_julia_array(arg) 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, cp.ndarray, torch.Tensor, DeviceArray) + cupy_ndarray) @@ -600,7 +464,9 @@ def check_kernel_output(self, func, gpu_args, instance, answer, atol, verify, ve self.dev.refresh_memory(gpu_args, instance.arguments, should_sync) # run the kernel + self.dev.synchronize() check = self.run_kernel(func, gpu_args, instance) + self.dev.synchronize() if not check: # runtime failure occurred that should be ignored, skip correctness check return @@ -624,9 +490,18 @@ def check_kernel_output(self, func, gpu_args, instance, answer, atol, verify, ve 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): + 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 @@ -639,8 +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: @@ -650,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) + func = self.compile_kernel(instance, verbose, gpu_args, build) if not func: result["__error__"] = util.CompilationFailedConfig() else: @@ -666,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): @@ -691,7 +569,12 @@ 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: + last_compilation_time -= result["time"] result["compile_time"] = last_compilation_time or 0 result["verification_time"] = last_verification_time or 0 @@ -701,30 +584,68 @@ def compile_and_benchmark(self, kernel_source, gpu_args, params, kernel_options, return result - def compile_kernel(self, instance, verbose): - """Compile the kernel for this specific instance.""" + 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: - func = self.dev.compile(instance) + 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.load(instance, build.result, gpu_args) + else: + 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 # 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 self.lang.upper() == "GENERIC_PYTHON": + # 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 = [ + util.SHARED_MEMORY_ERROR, + "local memory limit exceeded", + r"local memory \(\d+\) exceeds limit \(\d+\)", + ] + skippable = any(re.search(msg, error_message) for msg in shared_mem_error_messages) + reason = "too much shared memory used" + if skippable: + 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 much shared memory used" - ) + print(f"skipping config {util.get_instance_string(instance.params)} reason: {reason}") else: print("compile_kernel failed due to error: " + error_message) print("Error while compiling:", instance.name) @@ -767,29 +688,39 @@ 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.upper() != "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() # 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, ) - # check for templated kernel - if kernel_source.lang in ["CUDA"] and "<" in name and ">" in name: - kernel_string, name = wrap_templated_kernel(kernel_string, name) + # templated CUDA kernels are wrapped by KernelSourceStr.prepare_kernel_instance + name = instance_data.kernel_name + kernel_string = instance_data.kernel_str # Preprocess GPU arguments. Require for handling `Tunable` arguments 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.""" @@ -826,6 +757,7 @@ def ready_argument_list(self, arguments): return gpu_args + def run_kernel(self, func, gpu_args, instance): """Run a compiled kernel instance on a device.""" logging.debug("run_kernel %s", instance.name) @@ -833,7 +765,7 @@ def run_kernel(self, func, gpu_args, instance): 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") @@ -842,10 +774,21 @@ def run_kernel(self, func, gpu_args, instance): logging.debug("encountered unexpected runtime failure: " + str(e)) raise e return True + def retrieve_results_to_host(self, arguments: list, should_sync: list[bool], gpu_args, answer: list): """Retrieve results from device to host memory for all arguments that should be synchronized.""" result_host = [] + if self.lang.upper() == "GENERIC_PYTHON": + # arguments are either Torch tensors or built-in Python types + for i, arg in enumerate(gpu_args): + if not should_sync[i]: + result_host.append(None) + elif isinstance(arg, torch.Tensor): + result_host.append(arg.cpu()) + else: + result_host.append(arg) + return result_host for i, arg in enumerate(arguments): if not should_sync[i]: result_host.append(None) @@ -855,13 +798,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? @@ -984,7 +927,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." diff --git a/kernel_tuner/decorator.py b/kernel_tuner/decorator.py new file mode 100644 index 000000000..aad26233b --- /dev/null +++ b/kernel_tuner/decorator.py @@ -0,0 +1,639 @@ +"""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): + """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}, launch_grid, launch_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/interface.py b/kernel_tuner/interface.py index c1d72608e..c2357a2dc 100644 --- a/kernel_tuner/interface.py +++ b/kernel_tuner/interface.py @@ -40,6 +40,8 @@ import kernel_tuner.util as util from kernel_tuner.accuracy import Tunable from kernel_tuner.file_utils import get_input_file, get_t4_metadata, get_t4_results, import_class_from_file +from kernel_tuner.kernel_sources.kernel_source import KernelSource +from kernel_tuner.language import Language from kernel_tuner.searchspace import Searchspace from kernel_tuner.util import get_objective_defaults @@ -129,9 +131,10 @@ 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, 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 @@ -155,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", or "C".""", + "nvcuda", "OpenCL", "HIP", "C", or "Generic_Python". Kernels in Python + files (.py) are detected as "Generic_Python".""", "string", ), ), @@ -202,7 +206,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", ), ), @@ -306,6 +312,22 @@ def __deepcopy__(self, _): "dict", ), ), + ( + "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 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", + ), + ), ] ) @@ -507,6 +529,19 @@ 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, 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", + ), + ), ( "observers", ( @@ -626,17 +661,20 @@ def tune_kernel( metrics=None, simulation_mode=None, parallel=None, + parallel_compile=None, observers=None, objective=None, objective_higher_is_better=None, objectives=None, seed=None, + call_function=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 = core.KernelSource(kernel_name, kernel_source, lang, defines) + kernelsource = KernelSource(kernel_name, kernel_source, lang, defines, util.normalize_call_function(call_function)) # Convert Julia types if lang is not None and lang.upper() == "JULIA": @@ -720,6 +758,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 # (a falsy strategy means the default brute_force strategy is used) strategy = strategy or "brute_force" @@ -745,7 +787,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) @@ -768,6 +812,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 @@ -966,11 +1017,12 @@ def run_kernel( block_size_names=None, quiet=False, log=None, + call_function=None, ): 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, util.normalize_call_function(call_function)) if lang is not None and lang.upper() == "JULIA": params = util.julia_list_of_pairs_to_dict(params) @@ -987,6 +1039,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) @@ -1006,8 +1062,8 @@ def run_kernel( # see if the kernel arguments have correct type util.check_argument_list(instance.name, instance.kernel_string, arguments, lang=lang) - # compile the kernel - func = dev.compile_kernel(instance, False) + # compile the kernel (Generic Python kernels need the arguments to compile) + func = dev.compile_kernel(instance, False, gpu_args) if func is None: raise RuntimeError("cannot compile kernel, too much shared memory used") @@ -1029,16 +1085,27 @@ def run_kernel( if not dev.run_kernel(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(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]) + return results 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..bdba5d88b --- /dev/null +++ b/kernel_tuner/kernel_sources/kernel_source.py @@ -0,0 +1,86 @@ +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 + + +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): + 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()) + 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.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 ks_str: + return KernelSourceStr(kernel_name, kernel_sources, lang, defines) + else: + 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) + + + +class KernelSource(metaclass=KernelSourceFactory): + + 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) + language = detect_language(kernel_string) + self.lang = Language(language.upper()) + else: + 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: + raise NotImplementedError("create_kernel_instance not implemented") + + @abstractmethod + def check_argument_lists(self, kernel_name, arguments): + raise NotImplementedError("check_argument_lists not implemented") + + 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..5652598a0 --- /dev/null +++ b/kernel_tuner/kernel_sources/kernel_source_fn.py @@ -0,0 +1,369 @@ +import inspect +import ast +import copy +import os +import sys +import logging + +import tempfile +import importlib.machinery +import importlib.util + +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 +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 _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): + """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 _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 + + +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. + + 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 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, + 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): + super().__init__(kernel_name, kernel_source, lang, defines) + if isinstance(kernel_source, list): + raise ValueError("KernelSourceFn only supports a single kernel source") + + 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, 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 + + 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) + 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) + self.dependencies = self._find_dependencies(kernel_source) + + + @staticmethod + def _default_call_function(kernel_name, dsl): + """Select the default call function for the DSL the kernel is written in.""" + 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, + generate a kernel instance with these tuning parameters inserted. kernel_options, + grid and threads are not needed for Python kernels. + ''' + # 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( + 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): + ''' + 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 to store the transformed AST in. + new_module = ast.Module(body=[], type_ignores=[]) + + # Add imports + new_module.body.extend(self.import_nodes) + + # 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) # apply tuning params to dependencies + new_module.body.append(transformed_dep_node) + + # Transform main kernel + source_tree_copy = copy.deepcopy(self.source_tree) + transformed_tree = transformer.visit(source_tree_copy) + + # Add transformed main kernel to new module + new_module.body.append(transformed_tree) + + # Fix locations and generate source + ast.fix_missing_locations(new_module) + new_source = ast.unparse(new_module) + + #print(new_source) + + # 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) + + temp_module = import_kernel_module(module_name, temp_file_path) + 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, 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. + ''' + with open(filepath, "r") as f: + source_code = f.read() + tree = ast.parse(source_code, filename=filepath) + + # 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 + ''' + 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 + + 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 + + + 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) + ): + 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) + + + 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..30282f58b --- /dev/null +++ b/kernel_tuner/kernel_sources/kernel_source_str.py @@ -0,0 +1,193 @@ +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 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), + 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 == Language.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 == Language.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 = {Language.CUDA: ".cu", Language.OPENCL: ".cl", Language.C: ".c", Language.JULIA: ".jl"} + 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, lang=self.lang) + else: + logging.debug("Checking of arguments list not supported yet for code generators.") + + def infer_julia_backend(self): + """Infer the Julia backend from the kernel source.""" + backend = None + if self.lang.upper() != "JULIA": + return backend + + kernel_string = self.get_kernel_string(0) + if kernel_string.find("using CUDA") != -1: + backend = "cuda" + elif kernel_string.find("using ROCBackend") != -1: + backend = "amd" + elif kernel_string.find("using oneAPI") != -1: + backend = "intel" + elif kernel_string.find("using Metal") != -1: + backend = "metal" + else: + raise ValueError("Could not infer Julia backend from kernel source, provide it as a `compiler_option`") + return backend 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..9cf4f9644 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 = {} diff --git a/kernel_tuner/language.py b/kernel_tuner/language.py new file mode 100644 index 000000000..605fd668f --- /dev/null +++ b/kernel_tuner/language.py @@ -0,0 +1,29 @@ +from enum import Enum + + +class Language(str, Enum): + """Languages (and explicitly selectable backends) supported by Kernel Tuner. + + Members are also strings, so code that compares ``lang`` to plain strings + (e.g. ``lang.upper() == "CUDA"``) keeps working. + """ + + CUDA = "CUDA" + PYCUDA = "PYCUDA" + CUPY = "CUPY" + NVCUDA = "NVCUDA" + OPENCL = "OPENCL" + HIP = "HIP" + C = "C" + FORTRAN = "FORTRAN" + JULIA = "JULIA" + HYPERTUNER = "HYPERTUNER" + GENERIC_PYTHON = "GENERIC_PYTHON" + + def __str__(self): + return self.value + + @classmethod + def is_valid(cls, value: str) -> bool: + """Test if a language is valid in Kernel Tuner framework.""" + return value in cls._value2member_map_ diff --git a/kernel_tuner/observers/generic_python.py b/kernel_tuner/observers/generic_python.py new file mode 100644 index 000000000..62dfa20c7 --- /dev/null +++ b/kernel_tuner/observers/generic_python.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 GenericPythonRuntimeObserver(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/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): diff --git a/kernel_tuner/runners/parallel_compile.py b/kernel_tuner/runners/parallel_compile.py new file mode 100644 index 000000000..3d5904c2f --- /dev/null +++ b/kernel_tuner/runners/parallel_compile.py @@ -0,0 +1,124 @@ +"""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, 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): + """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 a917b68fa..051ec6de9 100644 --- a/kernel_tuner/util.py +++ b/kernel_tuner/util.py @@ -117,6 +117,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.""" @@ -163,6 +169,9 @@ def check_argument_type(dtype, kernel_argument): def check_argument_list(kernel_name, kernel_string, args, lang=None): """Raise an exception if kernel arguments do not match host arguments.""" + if kernel_string is None: + return + kernel_arguments = list() collected_errors = list() @@ -557,20 +566,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 @@ -765,6 +760,69 @@ 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 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 + (ast.ClassDef node, ast.FunctionDef node) in case the kernel is represented as the __call__ function + of a class (for Tilus support). + """ + 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") + + source = read_file(filepath) + tree = ast.parse(source, filename=str(filepath)) + + # 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 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}") + + raise ValueError(f"Kernel {kernel_name} not found in {filepath}") + + +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): @@ -1119,6 +1177,48 @@ 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. + + 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 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 _NormalizedCallFunction(v, ("grid", "threads")) + return _NormalizedCallFunction(v, ("grid",)) + return _NormalizedCallFunction(v, ()) + + + def parse_restrictions( restrictions: list[str], tune_params: dict, monolithic=False, format=None ) -> list[tuple[Union[Constraint, str], list[str]]]: 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/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/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/context.py b/test/context.py index fc8b44b9f..2d017033f 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 @@ -79,6 +80,35 @@ except ImportError: bayes_opt_gpytorch_present = False +try: + import torch + 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") +cutile_present = python_dsl_present("cuda", "tile") + try: import pyatf @@ -86,6 +116,13 @@ except ImportError: pyatf_present = False +try: + import skopt + + skopt_present = True +except ImportError: + skopt_present = False + try: import pymoo pymoo_present = True @@ -122,12 +159,22 @@ ) 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") +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") +skip_if_no_cutile = pytest.mark.skipif(not cutile_present, reason="cuTile not installed") 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") @@ -145,3 +192,5 @@ def skip_backend(backend: str): pytest.skip("HIP Python not installed") 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 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..6862f41d0 --- /dev/null +++ b/test/generic_python_kernels/__init__.py @@ -0,0 +1,11 @@ +"""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 + +The kernels are launched with the default call functions for their DSL. +""" 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/generic_python_kernels/cupyx_vector_add.py b/test/generic_python_kernels/cupyx_vector_add.py new file mode 100644 index 000000000..79466c549 --- /dev/null +++ b/test/generic_python_kernels/cupyx_vector_add.py @@ -0,0 +1,18 @@ +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]} 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..157589f27 --- /dev/null +++ b/test/generic_python_kernels/cute_vector_add.py @@ -0,0 +1,30 @@ +import cutlass.cute as cute + +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]} 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..71562d16a --- /dev/null +++ b/test/generic_python_kernels/cutile_vector_add.py @@ -0,0 +1,21 @@ +import cuda.tile as ct + +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]} 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..32aaadcf0 --- /dev/null +++ b/test/generic_python_kernels/numba_vector_add.py @@ -0,0 +1,18 @@ +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]} 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/generic_python_kernels/taichi_vector_add.py b/test/generic_python_kernels/taichi_vector_add.py new file mode 100644 index 000000000..5761b0928 --- /dev/null +++ b/test/generic_python_kernels/taichi_vector_add.py @@ -0,0 +1,30 @@ +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]} 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..f52675bc3 --- /dev/null +++ b/test/generic_python_kernels/tilelang_vector_add.py @@ -0,0 +1,28 @@ +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]} 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..2b8ada442 --- /dev/null +++ b/test/generic_python_kernels/tilus_vector_add.py @@ -0,0 +1,33 @@ +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]} 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..015aef23a --- /dev/null +++ b/test/generic_python_kernels/triton_vector_add.py @@ -0,0 +1,21 @@ +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]} 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..df45bd119 --- /dev/null +++ b/test/generic_python_kernels/warp_vector_add.py @@ -0,0 +1,18 @@ +import warp as wp + +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]} 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) diff --git a/test/test_backend.py b/test/test_backend.py index d56bdcce0..5f07f1d63 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -5,8 +5,9 @@ skip_if_no_opencl, skip_if_no_pycuda, skip_if_no_julia, + skip_if_no_torch, ) -from kernel_tuner.backends import backend, compiler, cupy, nvcuda, opencl, pycuda, julia +from kernel_tuner.backends import backend, compiler, cupy, nvcuda, opencl, pycuda, julia, generic_python class WrongBackend(backend.Backend): @@ -51,3 +52,8 @@ def test_pycuda_backend(): @skip_if_no_julia def test_julia_backend(): dev = julia.JuliaFunctions() + + +@skip_if_no_torch +def test_generic_python_backend(): + dev = generic_python.GenericPythonFunctions() diff --git a/test/test_decorator.py b/test/test_decorator.py new file mode 100644 index 000000000..79269f236 --- /dev/null +++ b/test/test_decorator.py @@ -0,0 +1,360 @@ +"""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, _Launch, _TuningCallFunction +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 _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() + 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) diff --git a/test/test_generic_python_dsls.py b/test/test_generic_python_dsls.py new file mode 100644 index 000000000..f0cee30e8 --- /dev/null +++ b/test/test_generic_python_dsls.py @@ -0,0 +1,246 @@ +"""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, + 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("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.param("cutile", marks=[skip_if_no_torch, skip_if_no_cutile]), +] + + +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): + 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) + + # 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, 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 + + +@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.""" + + 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" + + # 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: + _raise_exception_chain(chain) + assert isinstance(raised.value.__cause__, NameError) + 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") + 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 + 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 +] + +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"), + RuntimeError("Undefined variable undefined_name used"), # cuTile + 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" diff --git a/test/test_generic_python_functions.py b/test/test_generic_python_functions.py new file mode 100644 index 000000000..19eada417 --- /dev/null +++ b/test/test_generic_python_functions.py @@ -0,0 +1,182 @@ +from .context import skip_if_no_torch +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 +import pytest +from pathlib import Path +import os + +try: + import torch + torch_present = True +except ImportError: + pass + +KS_FILE = os.path.join(Path(__file__).resolve().parent, "test_kernel_source_fn.py") + +# 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", KS_FILE, "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) + + # 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 +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", KS_FILE, "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._kernel_kwargs(ks.signature, params) + + 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) + + +@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 diff --git a/test/test_kernel_source_fn.py b/test/test_kernel_source_fn.py new file mode 100644 index 000000000..e7025ad88 --- /dev/null +++ b/test/test_kernel_source_fn.py @@ -0,0 +1,186 @@ +import os +import tempfile +from pathlib import Path +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 + +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): + 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 + +@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", 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) + assert isinstance(ks_str, KernelSourceStr) + + +def test_initiation(): + ''' + Test invalid KernelSourceFn initations + ''' + 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"): + KernelSource("mock_kernel", "This is a string Kernel", "generic_python", call_function=call_mock) + + 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"): + 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", KS_FILE, "generic_python", call_function="not a function") + + + +def test_param_subsitution(): + params = {"mock_param": 128} + 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) + 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", KS_FILE, "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", 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 = """ +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("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", 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. + + 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() 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_observers.py b/test/test_observers.py index 1ad8f8454..bafb3d9ec 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): 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_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"] 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), "