TheRock/rocm-libraries/projects/rocprim/tuning/tuning_scripts/tuning_device_merge.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 typing import List, Optional, OrderedDict
import numpy as np
from hip import hip
import sys

sys.path.append("../")
from utils import hip_check, TYPE_CONFIGS, Parser
from tuning_scripts.base_tuner import BaseTuner

KEY_TYPES = ["rocprim::int128_t", "int64_t", "int", "short", "int8_t", "double", "float", "rocprim::half"]
VALUE_TYPES = ["rocprim::int128_t", "int64_t", "int", "short", "int8_t"]

DEFAULT_BYTES = 1024 * 1024 * 32 * 4
DEFAULT_ITERATIONS = 32
DEFAULT_MAX_FEVAL = 100
DEFAULT_OUTPUT_DIR = "../output"
DEFAULT_STRATEGY = "dual_annealing"


class Tuner(BaseTuner):
    def __init__(self, **kwargs):
        args = {
            "algo_full_name": "device_merge",
            "bytes_size": DEFAULT_BYTES,
            "iterations": DEFAULT_ITERATIONS,
            "output_dir": DEFAULT_OUTPUT_DIR,
            "include_default_config": True,
            "max_fevals": DEFAULT_MAX_FEVAL,
            "save_metadata": False,
            "strategy": DEFAULT_STRATEGY,
            "simulation_mode": False,
            "arch_name": None,
        }

        filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
        args.update(filtered_kwargs)
        super().__init__(**args)

    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."""
        params = OrderedDict()
        params["block_size_x"] = [64 * i for i in range(1, 17)]
        params["ipt"] = [1, 2] + [4 * i for i in range(1, 65)]
        return params

    def _get_restrictions(
        self, key_type: str, value_type: Optional[str] = None
    ) -> List[str]:
        """Constraints for what parameter combinations are valid during tuning"""
        size = self.bytes_size // TYPE_CONFIGS[key_type].size
        element_size = TYPE_CONFIGS[key_type].size
        if value_type:
            element_size += TYPE_CONFIGS[value_type].size

        def validate(params):
            block_size = params['block_size_x']
            ipt = params['ipt']

            # Total size constraint
            if block_size * ipt > size:
                return False

            # Memory size constraint
            if block_size * ipt * element_size > 65536:
                return False

            # Block size constraint
            if block_size > 1024:
                return False

            # Items per thread constraint - high ipts don't perform well
            if ipt >= block_size:
                return False

            # High ipts on gfx1030 cause HSA_STATUS_ERROR_INVALID_ISA
            if params.get('arch_name') == 'gfx1030' and ipt > 28:
                return False

            return True

        return validate

    def _get_grid_div_x(self):
        """Return the grid_div_x parameter for kernel_tuner.tune_kernel()"""
        return ["block_size_x"]

    def _allocate_device_memory(self, key_type: str, value_type: Optional[str] = None):
        """
        Allocates device memory for merge operation.

        Args:
            key_type: Type of the keys
            value_type: Optional type of the values

        Returns:
            Tuple containing:
            - List of device pointers and other arguments for kernel
            - Dictionary of allocated memory pointers for cleanup
            - Total bytes allocated
            - Total size of arrays
        """
        key_info = TYPE_CONFIGS[key_type]
        size = self.bytes_size // key_info.size
        size1 = size // 2
        size2 = size - size1

        allocated_mem = {}
        if self.simulation_mode:
            kernel_args = [
                np.array([]),
                np.array([]),
                np.array([]),
                np.array([]),
                np.array([]),
            ]
            if value_type:
                kernel_args.extend([np.array([]), np.array([]), np.array([])])

        else:
            # Generate and sort input data
            if key_type in ["rocprim::int128_t", "rocprim::uint128_t"]:
                keys_input1 = np.zeros(size1, dtype=key_info.numpy_type)
                keys_input2 = np.zeros(size2, dtype=key_info.numpy_type)
                keys_input1['low'] = np.sort(np.random.rand(size1)).astype(keys_input1['low'][0])
                keys_input2['low'] = np.sort(np.random.rand(size2)).astype(keys_input2['low'][0])
            else:
                keys_input1 = np.sort(np.random.rand(size1)).astype(key_info.numpy_type)
                keys_input2 = np.sort(np.random.rand(size2)).astype(key_info.numpy_type)

            # Allocate device memory for keys and configure arrays
            allocated_mem["d_keys_input1"] = hip_check(
                hip.hipMalloc(size1 * key_info.size)
            )
            allocated_mem["d_keys_input2"] = hip_check(
                hip.hipMalloc(size2 * key_info.size)
            )
            allocated_mem["d_keys_output"] = hip_check(
                hip.hipMalloc(size * key_info.size)
            )

            # Copy key data to device
            hip_check(
                hip.hipMemcpy(
                    allocated_mem["d_keys_input1"],
                    keys_input1,
                    size1 * key_info.size,
                    hip.hipMemcpyKind.hipMemcpyHostToDevice,
                )
            )
            hip_check(
                hip.hipMemcpy(
                    allocated_mem["d_keys_input2"],
                    keys_input2,
                    size2 * key_info.size,
                    hip.hipMemcpyKind.hipMemcpyHostToDevice,
                )
            )

            # Create initial argument list for size query
            kernel_args = [
                allocated_mem["d_keys_input1"],
                allocated_mem["d_keys_input2"],
                allocated_mem["d_keys_output"],
            ]

            # Handle value arrays if needed
            if value_type:
                value_info = TYPE_CONFIGS[value_type]

                # Generate and sort input data
                if value_type in ["rocprim::int128_t", "rocprim::uint128_t"]:
                    values_input1 = np.zeros(size1, dtype=value_info.numpy_type)
                    values_input2 = np.zeros(size2, dtype=value_info.numpy_type)
                    values_input1['low'] = np.sort(np.random.rand(size1)).astype(values_input1['low'][0])
                    values_input2['low'] = np.sort(np.random.rand(size2)).astype(values_input2['low'][0])
                else:
                    values_input1 = np.sort(np.random.rand(size1)).astype(value_info.numpy_type)
                    values_input2 = np.sort(np.random.rand(size2)).astype(value_info.numpy_type)

                allocated_mem["d_values_input1"] = hip_check(
                    hip.hipMalloc(size1 * value_info.size)
                )
                allocated_mem["d_values_input2"] = hip_check(
                    hip.hipMalloc(size2 * value_info.size)
                )
                allocated_mem["d_values_output"] = hip_check(
                    hip.hipMalloc(size * value_info.size)
                )

                # Copy value data to device
                hip_check(
                    hip.hipMemcpy(
                        allocated_mem["d_values_input1"],
                        values_input1,
                        size1 * value_info.size,
                        hip.hipMemcpyKind.hipMemcpyHostToDevice,
                    )
                )
                hip_check(
                    hip.hipMemcpy(
                        allocated_mem["d_values_input2"],
                        values_input2,
                        size2 * value_info.size,
                        hip.hipMemcpyKind.hipMemcpyHostToDevice,
                    )
                )

                kernel_args.extend(
                    [
                        allocated_mem["d_values_input1"],
                        allocated_mem["d_values_input2"],
                        allocated_mem["d_values_output"],
                    ]
                )

        # Add sizes to arguments
        kernel_args.extend([np.uint64(size1), np.uint64(size2)])

        return kernel_args, allocated_mem

    def tune_all(self) -> None:
        """Tune for all key type and value type combinations"""
        for key_type in KEY_TYPES:
            self.tune_type(key_type)
            for value_type in VALUE_TYPES:
                self.tune_type(key_type, value_type)


if __name__ == "__main__":
    parser = Parser.get_parser(
        DEFAULT_BYTES,
        DEFAULT_ITERATIONS,
        DEFAULT_MAX_FEVAL,
        DEFAULT_OUTPUT_DIR,
        DEFAULT_STRATEGY,
    )
    args = parser.parse_args()

    tuner = Tuner(
        algo_full_name="device_merge",
        bytes_size=args.size,
        iterations=args.iterations,
        output_dir=args.output_dir,
        include_default_config=args.include_default_config,
        max_fevals=args.max_fevals,
        save_metadata=args.save_metadata,
        strategy=args.strategy,
        simulation_mode=args.simulation_mode,
        arch_name=args.arch_name,
        seed=args.seed,
    )
    tuner.tune_all()