Generic Support for Python eDSLs

.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:

.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:

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)

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
 --------------------------

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 <https://github.com/kerneltuner/kernel_tuner/tree/master/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 <https://github.com/KernelTuner/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.

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:

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 <https://python-poetry.org/docs/#installing-with-the-official-installer>`__. 
     * 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 <installation>`.