Not able to infer ``atan2`` for ``CudaTarget``

[mitkotak]

Here's the MWE

import loopy as lp
import numpy as np
knl = lp.make_kernel(
        domains = "{ [_0]: 0<=_0<10 }",
        instructions = "out[_0] = atan2(a[_0], b[_0])",
        kernel_data = [lp.GlobalArg("out", shape=lp.auto, dtype="float64"),
            lp.GlobalArg("a", shape=lp.auto, dtype="float64"),
            lp.GlobalArg("b", shape=lp.auto, dtype="float64")],
        target=lp.CudaTarget(),
        lang_version=(2018, 2))
knl_string = lp.generate_code_v2(knl).device_code()

Here's the generated kernel

---------------------------------------------------------------------------
KERNEL: loopy_kernel
---------------------------------------------------------------------------
ARGUMENTS:
out: type: np:dtype('float64'), shape: (10), dim_tags: (N0:stride:1) out aspace: global
a: type: np:dtype('float64'), shape: (10), dim_tags: (N0:stride:1) in aspace: global
b: type: np:dtype('float64'), shape: (10), dim_tags: (N0:stride:1) in aspace: global
---------------------------------------------------------------------------
DOMAINS:
{ [_0] : 0 <= _0 <= 9 }
---------------------------------------------------------------------------
INAME TAGS:
_0: None
---------------------------------------------------------------------------
INSTRUCTIONS:
for _0
  out[_0] = atan2(a[_0], b[_0])  {id=insn}
end _0
---------------------------------------------------------------------------

Here's the error trace

Traceback (most recent call last):
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/tools.py", line 922, in wrapper
    result = transform_cache[cache_key]
  File "/home/mitak2/env/lib/python3.10/site-packages/pytools/persistent_dict.py", line 583, in __getitem__
    return self.fetch(key, _stacklevel=1)
  File "/home/mitak2/env/lib/python3.10/site-packages/pytools/persistent_dict.py", line 702, in fetch
    raise NoSuchEntryError(key)
pytools.persistent_dict.NoSuchEntryError: ((TranslationUnit(callables_table=pmap({'loopy_kernel': CallableKernel(name='loopy_kernel', subkernel=LoopKernel(domains=[BasicSet("{ [_0] : 0 <= _0 <= 9 }")], instructions=[Assignment(within_inames=frozenset({'_0'}), temp_var_type=Optional(), id='insn', predicates=frozenset(), priority=0, groups=frozenset(), depends_on_is_final=False, assignee=Subscript(Variable('out'), Variable('_0')), atomicity=(), depends_on=frozenset(), tags=frozenset(), expression=Call(Variable('atan2'), (Subscript(Variable('a'), Variable('_0')), Subscript(Variable('b'), Variable('_0')))), conflicts_with_groups=frozenset(), within_inames_is_final=False, no_sync_with=frozenset())], args=[<out: type: np:dtype('float64'), shape: (10), dim_tags: (N0:stride:1) out aspace: global>, <a: type: np:dtype('float64'), shape: (10), dim_tags: (N0:stride:1) in aspace: global>, <b: type: np:dtype('float64'), shape: (10), dim_tags: (N0:stride:1) in aspace: global>], assumptions=BasicSet("{  :  }"), temporary_variables={}, inames={'_0': Iname(name='_0', tags=frozenset())}, substitutions={}, options=Options(disable_global_barriers=False, enforce_array_accesses_within_bounds=True, no_numpy=False, trace_assignments=False, edit_code=False, check_dep_resolution=True, enforce_variable_access_ordered=True, annotate_inames=False, skip_arg_checks=False, return_dict=False, allow_terminal_colors=True, build_options=[], write_code=False, write_wrapper=False, insert_gbarriers=False, trace_assignment_values=False, cl_exec_manage_array_events=True), target=<loopy.target.cuda.CudaTarget object at 0x7f51df2b0250>, tags=frozenset(), state=<KernelState.INITIAL: 0>, name='loopy_kernel', preambles=(), preamble_generators=(), symbol_manglers=(), linearization=None, iname_slab_increments=immutables.Map({}), loop_priority=frozenset(), applied_iname_rewrites=(), index_dtype=np:dtype('int32'), silenced_warnings=[], overridden_get_grid_sizes_for_insn_ids=None), arg_id_to_descr=None, arg_id_to_dtype=None)}), func_id_to_in_knl_callable_mappers=[], entrypoints=frozenset({'loopy_kernel'}), target=<loopy.target.cuda.CudaTarget object at 0x7f51df2b0250>),), {})

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/mitak2/env/lib/python3.10/site-packages/pymbolic/mapper/__init__.py", line 242, in __call__
    return self._cache[cache_key]
KeyError: (<class 'pymbolic.primitives.Call'>, Call(ResolvedFunction(Variable('atan2')), (Subscript(Variable('a'), Variable('_0')), Subscript(Variable('b'), Variable('_0')))), (), ())

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/mitak2/cudagraph_pytato/loopy_atan2.py", line 14, in <module>
    knl_string = lp.generate_code_v2(knl).device_code()
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/codegen/__init__.py", line 597, in generate_code_v2
    program = preprocess_program(program)
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/tools.py", line 938, in wrapper
    result = func(*args, **kwargs)
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/preprocess.py", line 811, in preprocess_program
    program = infer_unknown_types(program, expect_completion=False)
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/type_inference.py", line 1053, in infer_unknown_types
    new_callable, clbl_inf_ctx = program.callables_table[e].with_types(
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/kernel/function_interface.py", line 726, in with_types
    infer_unknown_types_for_a_single_kernel(
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/type_inference.py", line 989, in infer_unknown_types_for_a_single_kernel
    type_inf_mapper(insn.expression,
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/type_inference.py", line 208, in __call__
    result = super().__call__(
  File "/home/mitak2/env/lib/python3.10/site-packages/pymbolic/mapper/__init__.py", line 244, in __call__
    result = super().rec(expr, *args, **kwargs)
  File "/home/mitak2/env/lib/python3.10/site-packages/pymbolic/mapper/__init__.py", line 153, in __call__
    return method(expr, *args, **kwargs)
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/type_inference.py", line 440, in map_call
    .with_types(arg_id_to_dtype,
  File "/home/mitak2/env/lib/python3.10/site-packages/loopy/kernel/function_interface.py", line 519, in with_types
    raise LoopyError("No type inference information present for "
loopy.diagnostic.LoopyError: No type inference information present for the function atan2