Fix CudaCallable type inference never being called

loopy/target/cuda.py

@@ -26,9 +26,11 @@
 THE SOFTWARE.
 """
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
+from constantdict import constantdict
+from typing_extensions import Self, override
 
 from cgen import Const, Declarator, Generable, Pointer
 from pymbolic import var
@@ -46,14 +48,15 @@
 from loopy.kernel.function_interface import ScalarCallable
 from loopy.target.c import CFamilyASTBuilder, CFamilyTarget
 from loopy.target.c.codegen.expression import ExpressionToCExpressionMapper
-from loopy.types import NumpyType
+from loopy.types import LoopyType, NumpyType
 
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Mapping, Sequence
 
     from loopy.codegen import CodeGenerationState
     from loopy.codegen.result import CodeGenerationResult
+    from loopy.translation_unit import CallablesInferenceContext
 
 
 # {{{ vector types
@@ -144,68 +147,94 @@ def _register_vector_types(dtype_registry):
 
 
 class CudaCallable(ScalarCallable):
+    """
+    Records information about functions provided by the CUDA runtime.
+    """
+
+    @override
+    def with_types(self,
+                   arg_id_to_dtype: Mapping[int | str, LoopyType],
+                   clbl_inf_ctx: CallablesInferenceContext,
+               ) -> tuple[Self, CallablesInferenceContext]:
+        name = self.name
 
-    def cuda_with_types(self, arg_id_to_dtype, callables_table):
+        for id in arg_id_to_dtype:
+            if not isinstance(id, int):
+                raise LoopyError(f"'{name}' can take only positional arguments")
 
-        name = self.name
+        arg_num_to_dtype = {cast("int", id): t for id, t in arg_id_to_dtype.items()}
 
         if name in _CUDA_SPECIFIC_FUNCTIONS:
             num_args = _CUDA_SPECIFIC_FUNCTIONS[name]
 
             # {{{ sanity checks
 
-            for id, dtype in arg_id_to_dtype.items():
+            for id, dtype in arg_num_to_dtype.items():
                 if not -1 <= id < num_args:
                     raise LoopyError("%s can take only %d arguments." % (name,
                             num_args))
 
-                if dtype is not None and dtype.kind == "c":
+                if dtype is not None and dtype.is_complex():
                     raise LoopyTypeError(
                         f"'{name}' does not support complex arguments.")
 
             # }}}
 
             for i in range(num_args):
-                if i not in arg_id_to_dtype or arg_id_to_dtype[i] is None:
+                if arg_num_to_dtype.get(i) is None:
                     # the types provided aren't mature enough to specialize the
                     # callable
                     return (
-                            self.copy(arg_id_to_dtype=arg_id_to_dtype),
-                            callables_table)
+                            self.copy(
+                                arg_id_to_dtype=constantdict(arg_num_to_dtype)),
+                            clbl_inf_ctx)
 
             dtype = np.result_type(*[
-                    dtype.numpy_dtype for id, dtype in arg_id_to_dtype.items()
+                    dtype.numpy_dtype for id, dtype in arg_num_to_dtype.items()
                     if id >= 0])
 
             updated_arg_id_to_dtype = {id: NumpyType(dtype)
                     for id in range(-1, num_args)}
 
             return (
                     self.copy(name_in_target=name,
-                        arg_id_to_dtype=updated_arg_id_to_dtype),
-                    callables_table)
+                        arg_id_to_dtype=constantdict(updated_arg_id_to_dtype)),
+                    clbl_inf_ctx)
 
         if name == "dot":
             # CUDA dot function:
             # Performs dot product. Input types: vector and return type: scalar.
             for i in range(2):
-                if i not in arg_id_to_dtype or arg_id_to_dtype[i] is None:
+                if arg_num_to_dtype.get(i) is None:
                     # the types provided aren't mature enough to specialize the
                     # callable
                     return (
-                            self.copy(arg_id_to_dtype=arg_id_to_dtype),
-                            callables_table)
+                            self.copy(
+                                arg_id_to_dtype=constantdict(arg_num_to_dtype)),
+                            clbl_inf_ctx)
+
+            input_dtype = arg_num_to_dtype[0]
+            fields = input_dtype.numpy_dtype.fields
+            if fields is None:
+                raise LoopyTypeError(
+                    f"'{name}' requires vector-typed arguments, got "
+                    f"{input_dtype}")
 
-            input_dtype = arg_id_to_dtype[0]
+            # CUDA's vector types name their first component 'x'.
+            scalar_dtype = fields["x"][0]
 
-            scalar_dtype, _offset, _field_name = input_dtype.fields["x"]
-            return_dtype = scalar_dtype
-            return self.copy(arg_id_to_dtype={0: input_dtype, 1: input_dtype,
-                                              -1: return_dtype})
+            return (
+                    self.copy(name_in_target=name,
+                        arg_id_to_dtype=constantdict({
+                            -1: NumpyType(scalar_dtype),
+                            0: input_dtype,
+                            1: input_dtype
+                            })),
+                    clbl_inf_ctx)
 
         return (
-                self.copy(arg_id_to_dtype=arg_id_to_dtype),
-                callables_table)
+                self.copy(arg_id_to_dtype=constantdict(arg_num_to_dtype)),
+                clbl_inf_ctx)
 
 
 def get_cuda_callables():

test/test_target.py

@@ -968,6 +968,53 @@ def test_bounds_check_with_shape_only_param():
     lp.generate_code_v2(knl)
 
 
+def test_cuda_specific_callables():
+    # CudaCallable's type inference was spelled 'cuda_with_types' and was
+    # therefore never called, leaving rsqrt/atan2/dot unusable on CudaTarget.
+    from loopy.diagnostic import LoopyTypeError
+    from loopy.target.cuda import vec as cuda_vec
+
+    knl = lp.make_kernel(
+        "{[i]: 0<=i<n}",
+        "out[i] = rsqrt(a[i]) + atan2(a[i], b[i])",
+        [lp.GlobalArg("a,b,out", np.float32, shape="n"),
+         lp.ValueArg("n", np.int32)],
+        target=lp.CudaTarget())
+    code = lp.generate_code_v2(knl).device_code()
+    assert "rsqrt(" in code
+    assert "atan2(" in code
+
+    # 'dot' returns the scalar type of its vector-typed arguments.
+    knl = lp.make_kernel(
+        "{[i]: 0<=i<n}",
+        "out[i] = dot(a[i], b[i])",
+        [lp.GlobalArg("a,b", cuda_vec.float4, shape="n"),
+         lp.GlobalArg("out", shape="n"),
+         lp.ValueArg("n", np.int32)],
+        target=lp.CudaTarget())
+    assert (lp.infer_unknown_types(knl)["loopy_kernel"].arg_dict["out"].dtype
+            == lp.to_loopy_type(np.float32))
+    assert "dot(" in lp.generate_code_v2(knl).device_code()
+
+    knl = lp.make_kernel(
+        "{[i]: 0<=i<n}",
+        "out[i] = rsqrt(a[i])",
+        [lp.GlobalArg("a,out", np.complex64, shape="n"),
+         lp.ValueArg("n", np.int32)],
+        target=lp.CudaTarget())
+    with pytest.raises(LoopyTypeError):
+        lp.generate_code_v2(knl)
+
+    knl = lp.make_kernel(
+        "{[i]: 0<=i<n}",
+        "out[i] = atan2(a[i], b[i], a[i])",
+        [lp.GlobalArg("a,b,out", np.float32, shape="n"),
+         lp.ValueArg("n", np.int32)],
+        target=lp.CudaTarget())
+    with pytest.raises(LoopyError, match="can take only 2 arguments"):
+        lp.generate_code_v2(knl)
+
+
 if __name__ == "__main__":
     import sys
     if len(sys.argv) > 1: