Auto-tuning Julia kernels

.github/workflows/test-python-package.yml

@@ -24,16 +24,26 @@ jobs:
                 os: [ubuntu-latest, macos-latest]
 
         steps:
-            - uses: actions/checkout@v4
+            - uses: actions/checkout@v7
             - name: Setup Nox
               uses: fjwillemsen/setup-nox2@v3.0.0
             - name: Setup Poetry
               uses: Gr1N/setup-poetry@v9
             - run: poetry self add poetry-plugin-export
+            - uses: julia-actions/setup-julia@v3
+              with:
+                  version: '1.11'   # when changed, also see `require_julia` in noxfile.py and the Julia version in Project.toml
             - name: Run tests with Nox
               run: |
+                  rm -rf .nox
                   pip install nox-poetry
-                  nox -- skip-gpu github-action
+                  nox -- skip-gpu skip-julia github-action
+            - name: Run Julia tests with Nox
+              run: |
+                  rm -rf .nox
+                  pip install nox-poetry
+                  nox --session tests-3.14 -- skip-gpu github-action
+            #   [ -d "~" ] && mkdir -p ~/.julia/registries
             # - name: Upload Coverage report to CodeCov
             #   uses: codecov/codecov-action@v3
             #   with:

.vscode/extensions.json

@@ -1,14 +1,15 @@
 {
-  // See https://go.microsoft.com/fwlink/?LinkId=827846 to learn about workspace recommendations.
-  // Extension identifier format: ${publisher}.${name}. Example: vscode.csharp
-  // List of extensions which should be recommended for users of this workspace.
-  "recommendations": [
-    "ms-python.python",
-    "ms-python.black-formatter",
-    "charliermarsh.ruff",
-    "bungcip.better-toml",
-    "njpwerner.autodocstring",
-  ],
-  // List of extensions recommended by VS Code that should not be recommended for users of this workspace.
-  "unwantedRecommendations": []
+    // See https://go.microsoft.com/fwlink/?LinkId=827846 to learn about workspace recommendations.
+    // Extension identifier format: ${publisher}.${name}. Example: vscode.csharp
+    // List of extensions which should be recommended for users of this workspace.
+    "recommendations": [
+        "ms-python.python",
+        "ms-python.black-formatter",
+        "charliermarsh.ruff",
+        "bungcip.better-toml",
+        "njpwerner.autodocstring",
+        "julialang.language-julia",
+    ],
+    // List of extensions recommended by VS Code that should not be recommended for users of this workspace.
+    "unwantedRecommendations": []
 }

.vscode/settings.json

@@ -17,14 +17,16 @@
     "black-formatter.args": [
         "--config=pyproject.toml"
     ],
-    "ruff.args": [
-        "--config=pyproject.toml"
-    ],
+    "ruff.configuration": "pyproject.toml",
     "autoDocstring.docstringFormat": "google-notypes",
     "esbonio.sphinx.confDir": "",
     "python.testing.pytestArgs": [
         "test"
     ],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,
-}
+    "sonarlint.connectedMode.project": {
+        "projectKey": "KernelTuner_kernel_tuner",
+        "connectionId": "kerneltuner",
+    }
+}

CONTRIBUTING.rst

@@ -26,7 +26,7 @@ Before creating a pull request please ensure the following:
 
 * You are working in an up-to-date development environment
 * You are a human developer. We are not interested in purely AI generated code contributions.
-* You have written unit tests to test your additions and all unit tests pass (run :bash:`nox`). If you do not have the required hardware, you can run :bash:`nox -- skip-gpu`, or :bash:`skip-cuda`, :bash:`skip-hip`, :bash:`skip-opencl`.
+* You have written unit tests to test your additions and all unit tests pass (run :bash:`nox`). If you do not have the required hardware, you can run :bash:`nox -- skip-gpu`, or :bash:`skip-cuda`, :bash:`skip-hip`, :bash:`skip-opencl`, :bash:`skip-julia`.
 * The examples still work and produce the same (or better) results
 * An entry about the change or addition is created in :bash:`CHANGELOG.md`
 

doc/source/dev-environment.rst

@@ -40,7 +40,7 @@ Steps with :bash:`sudo` access (e.g. on a local device):
     * Activate the environment with :bash:`pyenv activate kerneltuner`.
     * Make sure :bash:`which python` and :bash:`which pip` point to the expected Python location and version. 
     * Update Pip with :bash:`pip install --upgrade pip`.
-#. Install the project, dependencies and extras: :bash:`poetry install --with test,docs -E cuda -E opencl -E hip`, leaving out :bash:`-E cuda`, :bash:`-E opencl` or :bash:`-E hip` if this does not apply on your system. To go all-out, use :bash:`--all-extras`
+#. Install the project, dependencies and extras: :bash:`poetry install --with test,docs -E cuda -E opencl -E hip -E julia`, leaving out :bash:`-E cuda`, :bash:`-E opencl` etc. if this does not apply on your system. To go all-out, use :bash:`--all-extras`
     * Depending on the environment, it may be necessary or convenient to install extra packages such as :bash:`cupy-cuda11x` / :bash:`cupy-cuda12x`, and :bash:`cuda-python`. These are currently not defined as dependencies for kernel-tuner, but can be part of tests.
     * Do not forget to make sure the paths are set correctly. If you're using CUDA, the desired CUDA version should be in :bash:`$PATH`, :bash:`$LD_LIBARY_PATH` and :bash:`$CPATH`.
     * Re-open the shell for changes to take effect.
@@ -72,8 +72,9 @@ Steps without :bash:`sudo` access (e.g. on a cluster):
 #. `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.
     * Add the poetry export plugin with :bash:`poetry self add poetry-plugin-export`. 
-#. Install the project, dependencies and extras: :bash:`poetry install --with test,docs -E cuda -E opencl -E hip`, leaving out :bash:`-E cuda`, :bash:`-E opencl` or :bash:`-E hip` if this does not apply on your system. To go all-out, use :bash:`--all-extras`.
+#. Install the project, dependencies and extras: :bash:`poetry install --with test,docs -E cuda -E opencl -E hip -E julia`, leaving out :bash:`-E cuda`, :bash:`-E opencl` etc. if this does not apply on your system. To go all-out, use :bash:`--all-extras`.
     * If you run into "keyring" or other seemingly weird issues, this is a known issue with Poetry on some systems. Do: :bash:`pip install keyring`, :bash:`python3 -m keyring --disable`.
+    * Kernel Tuner has a dependency on Python-Constraint, which provides binaries for most systems. On some older systems, these binaries may not be compatible (on Linux usually due to an outdated LDD version, check with :bash:`ldd --version`, must be >=2.35). Without binaries, your system will try to build it yourself, so make sure build tools are available (e.g. with :bash:`module load gcc/13.2.0`). 
     * Depending on the environment, it may be necessary or convenient to install extra packages such as :bash:`cupy-cuda11x` / :bash:`cupy-cuda12x`, and :bash:`cuda-python`. These are currently not defined as dependencies for kernel-tuner, but can be part of tests.
     * Verify that your development environment has no missing installs or updates with :bash:`poetry install --sync --dry-run --with test`. 
 #. Check if the environment is setup correctly by running :bash:`pytest`. All tests should pass, except if you're not on a GPU node, or one or more extras has been left out in the previous step, then these tests will skip gracefully.
@@ -102,9 +103,10 @@ If you do not have fully compatible hardware or environment, you can use the fol
 * :bash:`nox -- skip-cuda` to skip tests involving CUDA.
 * :bash:`nox -- skip-hip` to skip tests involving HIP.
 * :bash:`nox -- skip-opencl` to skip tests involving OpenCL.
-* :bash:`nox -- skip-gpu` to skip all tests on the GPU (the same as :bash:`nox -- skip-cuda skip-hip skip-opencl`), especially helpful if you don't have a GPU locally. 
+* :bash:`nox -- skip-julia` to skip tests involving Julia.
+* :bash:`nox -- skip-gpu` to skip all tests on the GPU (the same as :bash:`nox -- skip-cuda skip-hip skip-opencl`), especially helpful if you don't have a GPU in your system. 
 
-Contributions you make to the Kernel Tuner should not break any of the tests even if you cannot run them locally!
+Contributions you make to Kernel Tuner should not break any of the tests even if you cannot run them locally!
 
 Running with :bash:`pytest` will test against your local Python version and PIP packages. 
 In this case, tests that require PyCuda and/or a CUDA capable GPU will be skipped automatically if these are not installed/present. 

kernel_tuner/accuracy.py

@@ -1,19 +1,25 @@
+"""Module for accuracy measurement and tunable parameters."""
+
+import logging
+import re
 from collections import UserDict
 from typing import Dict
+
 import numpy as np
-import logging
-import re
 
 from kernel_tuner.observers import OutputObserver
 
 
 class Tunable(UserDict):
+    """``Tunable`` can be used as a parameter value dependent input argument when tuning kernels."""
+
     def __init__(self, param_key: str, arrays: Dict):
-        """The ``Tunable`` object can be used as an input argument when tuning
-        kernels. It is a container that holds several arrays internally and
+        """The ``Tunable`` object can be used as an input argument when tuning kernels.
+
+        It is a container that holds several arrays internally and
         selects one array during benchmarking based on the value of a tunable parameter.
 
-        Example
+        Example:
         -------
         Consider this example::
 
@@ -37,6 +43,7 @@
         self.param_key = param_key
 
     def select_for_configuration(self, params):
+        """Select the array for the given configuration."""
         if callable(self.param_key):
             option = self.param_key(params)
         elif self.param_key in params:
@@ -46,13 +53,16 @@
 
         if option not in self.data:
             list = ", ".join(map(str, self.data.keys()))
-            raise KeyError(
-                f"'{option}' is not a valid parameter value, should be one of: {list}"
-            )
+            raise KeyError(f"'{option}' is not a valid parameter value, should be one of: {list}")
 
-        return self.data[option]
+        # continue recursively until we find a non-Tunable
+        el = self.data[option]
+        if isinstance(el, Tunable):
+            return el.select_for_configuration(params)
+        else:
+            return el
 
-    def __call__(self, params):
+    def __call__(self, params):  # noqa: D102
         return self.select_for_configuration(params)
 
 
@@ -70,6 +80,7 @@
     if dtype is None:
         try:
             from ml_dtypes import bfloat16
+
             dtype = bfloat16
         except ImportError:
             pass
@@ -78,6 +89,7 @@
     if dtype is None:
         try:
             from jax.numpy import bfloat16
+
             dtype = bfloat16
         except ImportError:
             pass
@@ -86,6 +98,7 @@
     if dtype is None:
         try:
             from tensorflow import bfloat16
+
             dtype = bfloat16.as_numpy_dtype
         except ImportError:
             pass
@@ -100,9 +113,9 @@
 
 
 def _to_float_dtype(x: str) -> np.dtype:
-    """Convert a string to a numpy data type (``dtype``). This function recognizes
-    common names (such as ``f16`` or ``kfloat``), and uses ``np.dtype(x)`` as a
-    fallback.
+    """Convert a string to a numpy data type (``dtype``).
+
+    This function recognizes common names (such as ``f16`` or ``kfloat``), and uses ``np.dtype(x)`` as a fallback.
     """
     if isinstance(x, str):
         x = x.lower()
@@ -123,16 +136,17 @@
 
 
 class TunablePrecision(Tunable):
-    def __init__(
-        self, param_key: str, array: np.ndarray, dtypes: Dict[str, np.dtype] = None
-    ):
-        """The ``Tunable`` object can be used as an input argument when tuning
-        kernels. It is a container that internally holds several arrays
+    """``TunablePrecision`` can be used as a precision-level dependent input argument when tuning kernels."""
+
+    def __init__(self, param_key: str, array: np.ndarray, dtypes: Dict[str, np.dtype] = None):
+        """The ``Tunable`` object can be used as an input argument when tuning kernels.
+
+        It is a container that internally holds several arrays
         containing the same data, but stored in using different levels of
         precision. During benchamrking, one array is selected based on the value
         of the tunable parameter called ``param_key``.
 
-        Example
+        Example:
         -------
         Consider this example::
 
@@ -156,7 +170,6 @@
             if bfloat16 is not None:
                 dtypes["bfloat16"] = bfloat16
 
-
         # If dtype is a list, convert it to a dictionary
         if isinstance(dtypes, (list, tuple)):
             dtypes = dict((name, _to_float_dtype(name)) for name in dtypes)
@@ -197,7 +210,6 @@
     The value of `EPS` is used for relative errors to prevent division by zero.
     ``
     """
-
     # Prepocess the provided name:
     # - convert to lowercase
     # - remove the word "error"
@@ -278,17 +290,13 @@
         raise ValueError(f"invalid error metric provided: {user_key}")
 
     # cast both arguments to f64 before passing them to the metric
-    return lambda a, b: metric(
-        a.astype(np.float64, copy=False), b.astype(np.float64, copy=False)
-    )
+    return lambda a, b: metric(a.astype(np.float64, copy=False), b.astype(np.float64, copy=False))
 
 
 class AccuracyObserver(OutputObserver):
-    """``AccuracyObserver`` measures the error on the output produced by a kernel
-    by comparing the output against a reference output.
+    """``AccuracyObserver`` measures the error on the output produced by a kernel by comparing to a reference output.
 
-    By default, it uses the root mean-squared error (RMSE) and uses the
-    metric name ``"error"``.
+    By default, it uses the root mean-squared error (RMSE) and uses the metric name ``"error"``.
     """
 
     def __init__(self, metric=None, key="error", *, atol=1e-8):
@@ -303,7 +311,6 @@
         :param atol: The tolerance used in relative metrics to prevent
                      division by zero. It is ignored by absolute error metrics.
         """
-
         # Default metric is RMSE
         if not metric:
             metric = "rmse"
@@ -317,6 +324,7 @@
         self.result = None
 
     def process_output(self, answers, outputs):
+        """Process the output produced by the kernel and compare it to the reference answers."""
         errors = []
 
         for answer, output in zip(answers, outputs):
@@ -326,4 +334,5 @@
         self.result = max(errors)
 
     def get_results(self):
+        """Get the results produced by this observer."""
         return dict([(self.key, self.result)])