feat(typing): return object from kernel_get_work_group_info

Author: alexfikl

pyopencl/_monkeypatch.py

@@ -259,7 +259,10 @@ def kernel_set_arg_types(self: _cl.Kernel, arg_types):
                     devs=self.context.devices))
 
 
-def kernel_get_work_group_info(self: _cl.Kernel, param: int, device: _cl.Device):
+def kernel_get_work_group_info(
+        self: _cl.Kernel,
+        param: int,
+        device: _cl.Device) -> object:
     try:
         wg_info_cache = self._wg_info_cache
     except AttributeError:

pyopencl/array.py

@@ -244,9 +244,10 @@ def kernel_runner(out: Array, *args: P.args, **kwargs: P.kwargs) -> cl.Event:
         assert queue is not None
 
         knl = kernel_getter(out, *args, **kwargs)
-        work_group_info = cast("int", knl.get_work_group_info(
+        work_group_info = knl.get_work_group_info(
             cl.kernel_work_group_info.WORK_GROUP_SIZE,
-            queue.device))
+            queue.device)
+        assert isinstance(work_group_info, int)
         gs, ls = out._get_sizes(queue, work_group_info)
 
         knl_args = (out, *args, out.size)
@@ -2706,10 +2707,11 @@ def make_func_for_chunk_size(chunk_size):
         if start_i + chunk_size > vec_count:
             knl = make_func_for_chunk_size(vec_count-start_i)
 
-        gs, ls = dest_indices._get_sizes(queue,
-                knl.get_work_group_info(
-                    cl.kernel_work_group_info.WORK_GROUP_SIZE,
-                    queue.device))
+        work_group_info = knl.get_work_group_info(
+            cl.kernel_work_group_info.WORK_GROUP_SIZE,
+            queue.device)
+        assert isinstance(work_group_info, int)
+        gs, ls = dest_indices._get_sizes(queue, work_group_info)
 
         wait_for_this = (
             *wait_for,

pyopencl/characterize/__init__.py

@@ -24,8 +24,6 @@
 """
 
 
-from typing import cast
-
 from pytools import memoize
 
 import pyopencl as cl
@@ -70,9 +68,9 @@ def reasonable_work_group_size_multiple(
         }
         """)
     prg.build()
-    return cast("int", prg.knl.get_work_group_info(
+    return prg.knl.get_work_group_info(
             cl.kernel_work_group_info.PREFERRED_WORK_GROUP_SIZE_MULTIPLE,
-            dev))
+            dev)
 
 
 def nv_compute_capability(dev: cl.Device):

pyopencl/elementwise.py

@@ -327,6 +327,7 @@ def __call__(self, *args, **kwargs) -> cl.Event:
         max_wg_size = kernel.get_work_group_info(
                 cl.kernel_work_group_info.WORK_GROUP_SIZE,
                 queue.device)
+        assert isinstance(max_wg_size, int)
 
         if range_ is not None:
             start = range_.start

pyopencl/reduction.py

@@ -329,7 +329,7 @@ def __init__(
 
         dtype_out = self.dtype_out = np.dtype(dtype_out)
 
-        max_group_size = None
+        max_group_size: int | None = None
         trip_count = 0
 
         while True:
@@ -342,6 +342,7 @@ def __init__(
             kernel_max_wg_size = self.stage_1_inf.kernel.get_work_group_info(
                     cl.kernel_work_group_info.WORK_GROUP_SIZE,
                     ctx.devices[0])
+            assert isinstance(kernel_max_wg_size, int)
 
             if self.stage_1_inf.group_size <= kernel_max_wg_size:
                 break

pyopencl/scan.py

@@ -26,7 +26,7 @@
 import logging
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 
@@ -1303,9 +1303,10 @@ def _finish_setup_impl(self) -> None:
             # at the desired work group size? Building it is the
             # only way to find out.
             kernel_max_wg_size = min(
-                    candidate_scan_info.kernel.get_work_group_info(
-                        cl.kernel_work_group_info.WORK_GROUP_SIZE,
-                        dev)
+                    cast("int",
+                         candidate_scan_info.kernel.get_work_group_info(
+                             cl.kernel_work_group_info.WORK_GROUP_SIZE,
+                             dev))
                     for dev in self.devices)
 
             if candidate_scan_info.wg_size <= kernel_max_wg_size: