Skip to content

Kernel callables can't inline functions with ValueArg

I'm not sure the exact status of all the kernel callable changes. But we're getting very close to switching in Firedrake, and so here we are.

The following, I think ought to work, but does not:

import loopy as lp
from loopy.transform.callable import _match_caller_callee_argument_dimension_
callee = lp.make_function("{[i] : 0 <= i < 6}", "a[i] = a[i] + j", 
                          [lp.GlobalArg("a", dtype="double", shape=(6, ), is_output_only=False),
                           lp.ValueArg("j", dtype="int")],
                          name="callee")
caller = lp.make_kernel("{[j] : 0 <= j < 2}", "callee(a[:], j)",
                        [lp.GlobalArg("a", dtype="double", shape=(6, ), is_output_only=False)],
                        name="caller", target=lp.CTarget())

registered = lp.register_callable_kernel(caller, callee)
inlined = _match_caller_callee_argument_dimension_(registered, callee.name)
inlined = lp.inline_callable_kernel(inlined, callee.name)


AttributeError                            Traceback (most recent call last)
<ipython-input-42-823de0802ea3> in <module>
     10 
     11 registered = lp.register_callable_kernel(caller, callee)
---> 12 inlined = _match_caller_callee_argument_dimension_(registered, callee.name)
     13 inlined = lp.inline_callable_kernel(inlined, callee.name)

~/Documents/work/src/firedrake/src/loopy/loopy/transform/callable.py in _match_caller_callee_argument_dimension_(program, callee_function_name)
    737             callee_function_name].subkernel
    738     new_callee_kernel = _match_caller_callee_argument_dimension_for_single_kernel(
--> 739             caller_knl, old_callee_knl)
    740 
    741     new_callables_table = program.callables_table.copy()

~/Documents/work/src/firedrake/src/loopy/loopy/transform/callable.py in _match_caller_callee_argument_dimension_for_single_kernel(caller_knl, callee_knl)
    643                 _shape_1_if_empty(
    644                     par.get_array_arg_descriptor(caller_knl).shape)
--> 645                 for par in parameters]
    646         kw_to_pos, pos_to_kw = get_kw_pos_association(callee_knl)
    647         for i in range(len(parameters), len(parameters)+len(kw_parameters)):

~/Documents/work/src/firedrake/src/loopy/loopy/transform/callable.py in <listcomp>(.0)
    643                 _shape_1_if_empty(
    644                     par.get_array_arg_descriptor(caller_knl).shape)
--> 645                 for par in parameters]
    646         kw_to_pos, pos_to_kw = get_kw_pos_association(callee_knl)
    647         for i in range(len(parameters), len(parameters)+len(kw_parameters)):

AttributeError: 'Variable' object has no attribute 'get_array_arg_descriptor'

It looks like the inlining assumes that all arguments to an inlineable kernel are ArrayArgs, but here, one is a ValueArg.