feat: kt now accepts HIP Python DeviceArray type

Author: MiloLurati

kernel_tuner/backends/compiler.py

@@ -26,6 +26,17 @@
 except ImportError:
     cp = None
 
+try:
+    from jip import hip
+except ImportError:
+    hip = None
+
+try:
+    from hip._util.types import DeviceArray
+except ImportError:
+    Pointer = Exception # using Exception here as a type that will never be among kernel arguments
+    DeviceArray = Exception
+
 
 def is_cupy_array(array):
     """Check if something is a cupy array.
@@ -145,9 +156,9 @@ def ready_argument_list(self, arguments):
         ctype_args = [None for _ in arguments]
 
         for i, arg in enumerate(arguments):
-            if not (isinstance(arg, (np.ndarray, np.number)) or is_cupy_array(arg)):
-                raise TypeError(f"Argument is not numpy or cupy ndarray or numpy scalar but a {type(arg)}")
-            dtype_str = str(arg.dtype)
+            if not (isinstance(arg, (np.ndarray, np.number, DeviceArray)) or is_cupy_array(arg)):
+                raise TypeError(f"Argument is not numpy or cupy ndarray or numpy scalar or HIP Python DeviceArray but a {type(arg)}")
+            dtype_str = arg.typestr if isinstance(arg, DeviceArray) else str(arg.dtype)
             if isinstance(arg, np.ndarray):
                 if dtype_str in dtype_map.keys():
                     # In numpy <= 1.15, ndarray.ctypes.data_as does not itself keep a reference
@@ -156,13 +167,20 @@ def ready_argument_list(self, arguments):
                     # (This changed in numpy > 1.15.)
                     # data_ctypes = data.ctypes.data_as(C.POINTER(dtype_map[dtype_str]))
                     data_ctypes = arg.ctypes.data_as(C.POINTER(dtype_map[dtype_str]))
+                    numpy_arg = arg
                 else:
                     raise TypeError("unknown dtype for ndarray")
             elif isinstance(arg, np.generic):
                 data_ctypes = dtype_map[dtype_str](arg)
+                numpy_arg = arg
             elif is_cupy_array(arg):
                 data_ctypes = C.c_void_p(arg.data.ptr)
-            ctype_args[i] = Argument(numpy=arg, ctypes=data_ctypes)
+                numpy_arg = arg
+            elif isinstance(arg, DeviceArray):
+                data_ctypes = arg.as_c_void_p()
+                numpy_arg = None
+
+            ctype_args[i] = Argument(numpy=numpy_arg, ctypes=data_ctypes)
         return ctype_args
 
     def compile(self, kernel_instance):
@@ -380,6 +398,12 @@ def memcpy_dtoh(self, dest, src):
         :param src: An Argument for some memory allocation
         :type src: Argument
         """
+        # If src.numpy is None, it means we're dealing with a HIP Python DeviceArray
+        if src.numpy is None:
+            # Skip memory copies for HIP Python DeviceArray
+            # This is because DeviceArray manages its own memory and donesn't need
+            # explicit copies like numpy arrays do
+            return
         if isinstance(dest, np.ndarray) and is_cupy_array(src.numpy):
             # Implicit conversion to a NumPy array is not allowed.
             value = src.numpy.get()
@@ -397,6 +421,12 @@ def memcpy_htod(self, dest, src):
         :param src: A numpy or cupy array containing the source data
         :type src: np.ndarray or cupy.ndarray
         """
+        # If src.numpy is None, it means we're dealing with a HIP Python DeviceArray
+        if dest.numpy is None:
+            # Skip memory copies for HIP Python DeviceArray
+            # This is because DeviceArray manages its own memory and donesn't need
+            # explicit copies like numpy arrays do
+            return
         if isinstance(dest.numpy, np.ndarray) and is_cupy_array(src):
             # Implicit conversion to a NumPy array is not allowed.
             value = src.get()

kernel_tuner/core.py

@@ -29,6 +29,11 @@
 except ImportError:
     torch = util.TorchPlaceHolder()
 
+try:
+    from hip._util.types import DeviceArray
+except ImportError:
+    DeviceArray = Exception # using Exception here as a type that will never be among kernel arguments
+
 _KernelInstance = namedtuple(
     "_KernelInstance",
     [
@@ -495,7 +500,7 @@ def check_kernel_output(
 
             should_sync = [answer[i] is not None for i, arg in enumerate(instance.arguments)]
         else:
-            should_sync = [isinstance(arg, (np.ndarray, cp.ndarray, torch.Tensor)) for arg in instance.arguments]
+            should_sync = [isinstance(arg, (np.ndarray, cp.ndarray, torch.Tensor, DeviceArray)) for arg in instance.arguments]
 
         # re-copy original contents of output arguments to GPU memory, to overwrite any changes
         # by earlier kernel runs
@@ -659,7 +664,7 @@ def compile_kernel(self, instance, verbose):
                         f"skipping config {util.get_instance_string(instance.params)} reason: too much shared memory used"
                     )
             else:
-                logging.debug("compile_kernel failed due to error: " + str(e))
+                print("compile_kernel failed due to error: " + error_message)
                 print("Error while compiling:", instance.name)
                 raise e
         return func

kernel_tuner/util.py

@@ -102,6 +102,10 @@ class StopCriterionReached(Exception):
     "block_size_z",
 ]
 
+try:
+    from hip._util.types import DeviceArray
+except ImportError:
+    DeviceArray = Exception # using Exception here as a type that will never be among kernel arguments
 
 def check_argument_type(dtype, kernel_argument):
     """Check if the numpy.dtype matches the type used in the code."""
@@ -125,59 +129,59 @@ def check_argument_type(dtype, kernel_argument):
         return any([substr in kernel_argument for substr in types_map[dtype]])
     return False  # unknown dtype. do not throw exception to still allow kernel to run.
 
-
 def check_argument_list(kernel_name, kernel_string, args):
-    """Raise an exception if a kernel arguments do not match host arguments."""
+    """Raise an exception if kernel arguments do not match host arguments."""
     kernel_arguments = list()
     collected_errors = list()
+    
     for iterator in re.finditer(kernel_name + "[ \n\t]*" + r"\(", kernel_string):
         kernel_start = iterator.end()
         kernel_end = kernel_string.find(")", kernel_start)
         if kernel_start != 0:
             kernel_arguments.append(kernel_string[kernel_start:kernel_end].split(","))
+
     for arguments_set, arguments in enumerate(kernel_arguments):
         collected_errors.append(list())
         if len(arguments) != len(args):
             collected_errors[arguments_set].append(
                 "Kernel and host argument lists do not match in size."
             )
             continue
+
         for i, arg in enumerate(args):
             kernel_argument = arguments[i]
 
-            # Fix to deal with tunable arguments
+            # Handle tunable arguments
             if isinstance(arg, Tunable):
                 continue
-
-            if not isinstance(arg, (np.ndarray, np.generic, cp.ndarray, torch.Tensor)):
+                
+            # Handle numpy arrays and other array types
+            if not isinstance(arg, (np.ndarray, np.generic, cp.ndarray, torch.Tensor, DeviceArray)):
                 raise TypeError(
-                    "Argument at position "
-                    + str(i)
-                    + " of type: "
-                    + str(type(arg))
-                    + " should be of type np.ndarray or numpy scalar"
+                    f"Argument at position {i} of type: {type(arg)} should be of type "
+                    "np.ndarray, numpy scalar, or HIP Python DeviceArray type"
                 )
 
             correct = True
-            if isinstance(arg, np.ndarray) and "*" not in kernel_argument:
-                correct = False  # array is passed to non-pointer kernel argument
+            if isinstance(arg, np.ndarray):
+                if "*" not in kernel_argument:
+                    correct = False
+            
+            if isinstance(arg, DeviceArray):
+                str_dtype = str(np.dtype(arg.typestr))
+            else:
+                str_dtype = str(arg.dtype)
 
-            if correct and check_argument_type(str(arg.dtype), kernel_argument):
+            if correct and check_argument_type(str_dtype, kernel_argument):
                 continue
-
+            
             collected_errors[arguments_set].append(
-                "Argument at position "
-                + str(i)
-                + " of dtype: "
-                + str(arg.dtype)
-                + " does not match "
-                + kernel_argument
-                + "."
+                f"Argument at position {i} of dtype: {str_dtype} does not match {kernel_argument}."
             )
+
         if not collected_errors[arguments_set]:
-            # We assume that if there is a possible list of arguments that matches with the provided one
-            # it is the right one
             return
+
     for errors in collected_errors:
         warnings.warn(errors[0], UserWarning)
 

test/test_file_utils.py

@@ -1,10 +1,18 @@
 import json
 
 import pytest
+import ctypes
 from jsonschema import validate
+import numpy as np
+import warnings
+try:
+    from hip import hip
+except:
+    hip = None
 
 from kernel_tuner.file_utils import output_file_schema, store_metadata_file, store_output_file
-from kernel_tuner.util import delete_temp_file
+from kernel_tuner.util import delete_temp_file, check_argument_list
+from .context import skip_if_no_hip
 
 from .test_runners import cache_filename, env, tune_kernel  # noqa: F401
 
@@ -55,3 +63,33 @@ def test_store_metadata_file():
     finally:
         # clean up
         delete_temp_file(filename)
+
+def hip_check(call_result):
+    err = call_result[0]
+    result = call_result[1:]
+    if len(result) == 1:
+        result = result[0]
+    if isinstance(err, hip.hipError_t) and err != hip.hipError_t.hipSuccess:
+        raise RuntimeError(str(err))
+    return result
+
+@skip_if_no_hip
+def test_check_argument_list_device_array():
+    """Test check_argument_list with DeviceArray"""
+    float_kernel = """
+    __global__ void simple_kernel(float* input) {
+        // kernel code
+    }
+    """
+    host_array = np.ones((100,), dtype=np.float32)
+    num_bytes = host_array.size * host_array.itemsize
+    device_array = hip_check(hip.hipMalloc(num_bytes))
+    device_array.configure(
+        typestr="float32",
+        shape=host_array.shape,
+        itemsize=host_array.itemsize
+    )
+    
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        check_argument_list("simple_kernel", float_kernel, [device_array])