TheRock/rocm-libraries/projects/rocprim/tuning/tuning_scripts/base_tuner.py

#!/usr/bin/env python3

# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

from abc import ABC, abstractmethod
from typing import List, Optional, OrderedDict, Dict
from kernel_tuner.file_utils import store_metadata_file, store_output_file
import kernel_tuner
import json
from pathlib import Path
import numpy as np
from jinja2 import Environment, FileSystemLoader
from utils import hip_check, TYPE_CONFIGS, ConfigParser
from hip import hip

"""
The following base class is used when implementing the tuning for new algorithms
using Kernel Tuner. For a detailed example, check out the tuning of device_merge.

The Tuner class manages the complete tuning workflow, handling parameter configurations,
device memory, result collection, and creates the necessary C++
wrapper code that interfaces between Kernel Tuner and rocPRIM templates. Since we are
tuning templated functions in rocPRIM rather than direct kernels, we need to create C
extern wrappers that return a float representing execution time. This wrapping is
necessary because we cannot know the C++ mangled name of the compiled templated function
ahead of time. For each combination of key-value types, the Tuner creates a
unique header file. We then use Kernel Tuner's C compiler functionality to compile and
benchmark these wrappers, rather than using its HIP backend. This is because we're
tuning complete algorithm functions that may internally launch one or more kernels,
rather than tuning individual kernels directly. The wrapper files serve as the
interface between Kernel Tuner and the rocPRIM algorithm implementations, allowing us to
measure and optimize performance across the entire algorithm execution. A base template
is available to use with the different algorithms. This template can be constomized
by implementing the different jinja blocks when extending the base template.

For each type combination, the tuning workflow in the Tuner class prepares the specific
parameters needed for tune_kernel, including tuning parameters like block sizes and
items per thread, function arguments from the wrapper, problem size specifications, and
parameter space restrictions. It also handles compiler configuration and strategy
options. A key feature integrated into the tuning process is the ability to
incorporate previous best configurations from older tuning runs. This ensures
continuity and improvement - new tuning results will be either equivalent to or better
than previous ones. This is achieved by parsing the existing configuration header file
to extract relevant configs for the current type combination being tuned. All this
happens with the ConfigParser class in config_parser.py. These configurations are added
to the parameter space, and if the chosen strategy doesn't naturally explore them, they
are benchmarked separately to ensure inclusion in the final results.

The final output consists of JSON files containing the tuning results, organized by
architecture, algorithm name, and data type combinations. These files include time,
performance metrics, parameter configurations, and optional metadata for analysis. One
particularly useful aspect of using Kernel Tuner is its ability to pause and resume
tuning sessions without losing previous progress, thanks to the caching system that
utilizes these JSON output files.

To adapt this framework for a new algorithm, you'll need to implement the Tuner and a
extend the jinja base template for your specific algorithm requirements. The device_merge
implementation serves as a comprehensive example.

For detailed information:
Kernel Tuner documentation: https://kerneltuner.github.io/kernel_tuner/stable/contents.html
ROCm HIP Python Wrapper: https://rocm.docs.amd.com/projects/hip-python/en/latest/index.html
"""


class BaseTuner(ABC):
    def __init__(
        self,
        algo_full_name: str,
        bytes_size: int,
        iterations: int,
        output_dir: str,
        include_default_config: bool,
        max_fevals: int,
        save_metadata: bool,
        strategy: str,
        simulation_mode: bool,
        arch_name: str,
        seed: int,
    ):
        """Initialize the tuner with configuration parameters.

        Args:
            algo_full_name: The full algorithm name, with type and name
            bytes_size: Size in bytes for the problem/data to tune with
            iterations: Number of iterations Kernel Tuner will run for each config
            output_dir: Directory path to store tuning results and metadata
            include_default_config: Whether to also test default configurations
            max_fevals: Maximum number of function evaluations for the tuning strategy
            save_metadata: Whether to save additional tuning metadata
            strategy: Tuning strategy to use (e.g. "dual_annealing", "brute_force")
        """
        self.algo_full_name = algo_full_name
        algo_types = ["warp", "block", "device"]
        algo_full_name_parts = algo_full_name.split("_")
        if algo_full_name_parts[0] in algo_types:
            self.algo_type = algo_full_name_parts[0]
            self.algo_name = "_".join(algo_full_name_parts[1:])
        else:
            raise ValueError(
                f"Invalid algo_type in {algo_full_name}. Must start with one of {algo_types}"
            )
        self.device_id = 0
        self.bytes_size = bytes_size
        self.iterations = iterations
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.simulation_mode = simulation_mode
        if not self.simulation_mode:
            self.device_properties = hip.hipDeviceProp_t()
            hip.hipGetDeviceProperties(self.device_properties, self.device_id)
            self.arch_name = self.device_properties.gcnArchName.decode().split(":")[0]
        else:
            assert arch_name, "Simulation mode requires --arch-name to be specified"
            self.arch_name = arch_name
        self.include_default_config = include_default_config
        self.max_fevals = max_fevals
        self.save_metadata = save_metadata
        self.strategy = strategy
        self.seed = None if seed == -1 else seed

        if self.include_default_config:
            self.config_parser = ConfigParser(
                self.algo_name, self._get_tune_params().keys()
            )
            with open(
                f"../../rocprim/include/rocprim/device/detail/config/{self.algo_full_name}.hpp"
            ) as f:
                self.default_configs = self.config_parser.parse_header(f.read())
        else:
            self.config_parser = None
            self.default_configs = None

    @abstractmethod
    def _get_tune_params(self) -> OrderedDict:
        """Returns tuning parameters and their possible values as an OrderedDict.
        Each parameter maps to a list of valid values to explore during tuning."""
        pass

    @abstractmethod
    def _get_restrictions(
        self, key_type: str, value_type: Optional[str] = None
    ) -> List[str]:
        """Define constraints for what parameter combinations are valid during tuning.

        Two options:
        1. Return list of string expressions that must all evaluate to True
            Example: ["block_size_x * ipt <= size", "block_size_x <= 1024"]
            These expressions are evaluated as boolean conditions by the Kernel Tuner.
            Each string must be a valid Python expression using the parameter names.

        2. Return a callable (function/lambda) that takes a params dictionary as input
            and returns True/False to validate the configuration.

            See Kernel Tuner's documentation for more details
        """
        pass

    @abstractmethod
    def _allocate_device_memory(self):
        """
        DeviceArray wrapper must be used to properly manage GPU memory allocation/deallocation
        and configure array properties. See HIP Python wrapper documentation for examples.
        """
        pass

    @abstractmethod
    def _get_grid_div_x(self):
        """Return the grid_div_x parameter for kernel_tuner.tune_kernel()"""
        pass

    @abstractmethod
    def tune_all(self) -> None:
        """Call tune_type for all key type and value type combinations"""
        pass

    def _get_problem_size(self, key_type: str, value_type: Optional[str] = None):
        return self.bytes_size // TYPE_CONFIGS[key_type].size

    def _free_device_memory(self, allocated_mem: Dict) -> None:
        """Free all device memory allocated with DeviceArray wrapper"""
        if not self.simulation_mode:
            for ptr in self.allocated_mem.values():
                if ptr is not None:
                    hip_check(hip.hipFree(ptr))

    def _get_metrics(self, key_type: str, value_type: Optional[str] = None) -> Dict:
        """Default metrics for performance measurement."""
        total_bytes = self._get_problem_size(key_type, value_type) * (
            TYPE_CONFIGS[key_type].size
            + (TYPE_CONFIGS[value_type].size if value_type else 0)
        )
        return {"GB/s": lambda p: (total_bytes / (p["time"] * 1e6))}

    def _run_default_config(
        self, tune_kernel_args: Dict, key_type: str, value_type: Optional[str] = None
    ) -> None:
        """Runs the default configuration if enabled."""
        if not self.include_default_config or self.simulation_mode:
            return

        print("Running default config")
        default_tune_params = self.config_parser.get_default_config(
            self.arch_name, key_type, value_type
        )

        default_args = tune_kernel_args.copy()
        default_args.update(
            {"tune_params": default_tune_params, "strategy": "brute_force"}
        )

        kernel_tuner.tune_kernel(**default_args)

    def tune_type(
        self,
        key_type: str,
        value_type: Optional[str] = None,
    ) -> None:
        """Performs auto-tuning for a specific key type and optional value type combination."""
        print(f"\nTuning for {key_type} {value_type if value_type else ''}")
        print(f"Using size: {self.bytes_size} bytes, iterations: {self.iterations}")
        strategy_print_message = (
            f"Using strategy: {self.strategy if self.strategy else 'brute_force'}"
        )
        strategy_print_message += (
            f", max fevals: {self.max_fevals}" if self.strategy != "brute_force" else ""
        )
        print(strategy_print_message)

        try:
            tune_kernel_args = self._get_base_tune_kernel_args(key_type, value_type)
            # Run main tuning
            results, _ = kernel_tuner.tune_kernel(**tune_kernel_args)
            # Run default config if enabled
            self._run_default_config(tune_kernel_args, key_type, value_type)

            cache_file_path = self._get_cache_file_path(key_type, value_type)
            self._save_output(cache_file_path, key_type, value_type, results)

        except Exception as e:
            print(f"Failed tuning for {key_type} {value_type if value_type else ''}")
            print(f"Error: {str(e)}")
            raise
        finally:
            self._free_device_memory(self.allocated_mem)

    def generate_wrapper(
        self,
        tune_params: OrderedDict,
        key_type: str,
        value_type: Optional[str] = None,
        output_file: bool = False,
    ) -> str:
        """Generate wrapper code using Jinja2 template inheritance."""
        env = Environment(
            loader=FileSystemLoader("templates"), trim_blocks=True, lstrip_blocks=True
        )

        template = env.get_template(f"{self.algo_full_name}_wrapper_template")
        context = {
            "algo_type": self.algo_type,
            "algo_name": self.algo_name,
            "config_params": ", ".join(tune_params.keys()),
            "key_type": key_type,
            "value_type": value_type,
        }

        content = template.render(**context)

        if output_file:
            base_name = f"device_{self.algo_name}_wrapper"
            file_name = (
                f"{base_name}_{value_type}.hpp" if value_type else f"{base_name}.hpp"
            )
            with open(file_name, "w") as f:
                f.write(content)

        return content

    def _get_base_tune_kernel_args(
        self,
        key_type: str,
        value_type: Optional[str] = None,
    ) -> Dict:
        """Returns base arguments for kernel_tuner.tune_kernel()."""
        np.random.seed(self.seed)
        args, allocated_mem = self._allocate_device_memory(key_type, value_type)
        self.allocated_mem = allocated_mem

        wrapper_string = self.generate_wrapper(
            tune_params=self._get_tune_params(),
            key_type=key_type,
            value_type=value_type,
        )

        tune_kernel_args = {
            "kernel_name": f"{self.algo_full_name}_wrapper",
            "kernel_source": wrapper_string,
            "problem_size": self._get_problem_size(key_type, value_type),
            "arguments": args,
            "tune_params": self._get_tune_params(),
            "strategy": self.strategy,
            "grid_div_x": self._get_grid_div_x(),
            "metrics": self._get_metrics(key_type, value_type),
            "iterations": self.iterations,
            "cache": str(self._get_cache_file_path(key_type, value_type)),
            "lang": "C",
            "compiler": "hipcc",
            "compiler_options": self._get_compiler_options(),
            "restrictions": self._get_restrictions(key_type, value_type),
            "verbose": False,
            "log": False,
            "device": self.device_id,
            "simulation_mode": self.simulation_mode,
        }

        if self.strategy != "brute_force":
            tune_kernel_args["strategy_options"] = {"max_fevals": self.max_fevals}

        return tune_kernel_args

    def _get_cache_file_name(self, key_type: str, value_type: str = None):
        """Return the name of the cache file based on algo name, arch name and key value types"""
        cache_file_path = f'{self.algo_full_name}_{self.arch_name}_{key_type.replace("rocprim::", "")}'
        if value_type:
            cache_file_path += f"_{value_type.replace("rocprim::", "")}"
        cache_file_path += "_cache.json"

        return cache_file_path

    def _get_cache_file_path(self, key_type: str, value_type: str = None):
        "Return the path of the cache file"
        return self.output_dir / self._get_cache_file_name(key_type, value_type)

    def _save_output(
        self,
        cache_file: str,
        key_type,
        value_type,
        results: Dict = None,
    ):
        """Save tuning results and metadata to files."""
        if self.simulation_mode:
            assert results
            cache_file = Path(
                f"../simulated_output/{self.strategy}_fevals{self.max_fevals}"
            )
            cache_file.mkdir(parents=True, exist_ok=True)
            cache_file = cache_file / self._get_cache_file_name(
                self.algo_full_name, key_type, value_type
            )
            store_output_file(str(cache_file), results, self.tune_params)

        with open(cache_file, "r") as f:
            cache_dict = json.load(f)
        if "arch_name" not in cache_dict:
            with open(cache_file, "r") as f:
                content = f.read()

            new_content = "{\n"
            new_content += f'"arch_name": "{self.arch_name}",\n'
            new_content += f'"algo_name": "{self.algo_full_name}",\n'
            new_content += f'"key_type": "{key_type}",\n'
            value_type_string = f"{value_type}" if value_type else "empty_type"
            new_content += f'"value_type": "{value_type_string}",'
            new_content += content.lstrip()[1:]

            with open(cache_file, "w") as f:
                f.write(new_content)

        if self.save_metadata:
            store_metadata_file(cache_file.replace("cache", "metadata"))

    def _get_compiler_options(self) -> List[str]:
        """Returns a list with all compiler options to pass to Kernel Tuner"""
        return [
            "-fPIC",
            "-std=c++17",
            "-I/opt/rocm/include",
            "-I../../rocprim/include",
            "-I../../build/rocprim/include/rocprim",
            "-Wno-#pragma-messages",
            f"--offload-arch={self.arch_name}",
        ]