KeyError in with_descrs

[isuruf]

Traceback

    knl = self.get_optimized_kernel(**kwargs)
sumpy/e2e.py:650: in get_optimized_kernel
    knl = self.get_kernel(result_dtype)
sumpy/e2e.py:645: in get_kernel
    loopy_knl = lp.inline_callable_kernel(loopy_knl, name)
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/transform/callable.py:500: in inline_callable_kernel
    translation_unit = infer_arg_descr(translation_unit)
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/preprocess.py:499: in infer_arg_descr
    new_callable, clbl_inf_ctx = program.callables_table[e].with_descrs(
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/kernel/function_interface.py:805: in with_descrs
    subkernel, clbl_inf_ctx = traverse_to_infer_arg_descr(subkernel,
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/preprocess.py:457: in traverse_to_infer_arg_descr
    arg_descr_inf_mapper.map_kernel(kernel))
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/preprocess.py:425: in map_kernel
    new_insns.append(insn.with_transformed_expressions(mapper))
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/kernel/instruction.py:1119: in with_transformed_expressions
    expression = f(self.expression)
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/preprocess.py:408: in __call__
    return IdentityMapper.__call__(self, expr,
../pymbolic/pymbolic/mapper/__init__.py:138: in __call__
    return method(expr, *args, **kwargs)
/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/preprocess.py:387: in map_call
    new_clbl, self.clbl_inf_ctx = clbl.with_descrs(arg_id_to_descr,
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = CallableKernel(name='m2l_preprocess', arg_id_to_descr=None, subkernel=LoopKernel(linearization=None, preambles=[], loc...t("{  :  }"), symbol_manglers=[], tags=frozenset(), loop_priority=frozenset(), substitutions={}), arg_id_to_dtype=None)
arg_id_to_descr = {-1: ArrayArgDescriptor(dim_tags=(FixedStrideArrayDimTag(target_axis=0, layout_nesting_level=None, stride=1),), addres...r(dim_tags=(FixedStrideArrayDimTag(target_axis=0, layout_nesting_level=None, stride=1),), address_space=2, shape=(3,))}
clbl_inf_ctx = CallablesInferenceContext(callables={'m2l_preprocess_multipole': CallableKernel(name='m2l_preprocess_multipole', arg_i...iqueNameGenerator object at 0x7fe99c5234f0>, new_entrypoints=frozenset(), renames=defaultdict(<class 'frozenset'>, {}))

    def with_descrs(self, arg_id_to_descr, clbl_inf_ctx):
    
        # arg_id_to_descr expressions provided are from the caller's namespace,
        # need to register
    
        kw_to_pos, pos_to_kw = get_kw_pos_association(self.subkernel)
    
        kw_to_callee_idx = {arg.name: i
                            for i, arg in enumerate(self.subkernel.args)}
    
        new_args = self.subkernel.args[:]
    
        for arg_id, descr in arg_id_to_descr.items():
            if isinstance(arg_id, int):
>               arg_id = pos_to_kw[arg_id]
E               KeyError: 1

/home/idf2/miniforge3/lib/python3.10/site-packages/loopy/kernel/function_interface.py:767: KeyError

Kernel:

KERNEL: m2l_preprocess_multipole
---------------------------------------------------------------------------
ARGUMENTS:
nsrc_boxes: ValueArg, type: np:dtype('int32')
preprocessed_src_expansions: type: <auto/runtime>, shape: (nsrc_boxes, 9), dim_tags: (N1:stride:9, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
src_expansions: type: <auto/runtime>, shape: (nsrc_boxes, 3), dim_tags: (N1:stride:3, N0:stride:1), offset: <class 'loopy.kernel.data.auto'> aspace: global
src_rscale: ValueArg, type: <auto/runtime>
---------------------------------------------------------------------------
DOMAINS:
[nsrc_boxes] -> { [isrc_box] : 0 <= isrc_box < nsrc_boxes }
{ [isrc_coeff] : 0 <= isrc_coeff <= 2 }
{ [itgt_coeff, itgt_coeff_0, itgt_coeff_1] : 0 <= itgt_coeff <= 8 and 0 <= itgt_coeff_0 <= 8 and 0 <= itgt_coeff_1 <= 8 }
---------------------------------------------------------------------------
INAME TAGS:
isrc_box: None
isrc_coeff: None
itgt_coeff: None
itgt_coeff_0: None
itgt_coeff_1: None
---------------------------------------------------------------------------
TEMPORARIES:
coeffs: type: <auto/runtime>, shape: (9), dim_tags: (N0:stride:1) scope:auto
---------------------------------------------------------------------------
INSTRUCTIONS:
   for isrc_box, itgt_coeff_0
↱      coeffs[itgt_coeff_0] = 0  {id=init}
│    end itgt_coeff_0
└↱   [itgt_coeff]: coeffs[itgt_coeff] = m2l_preprocess([itgt_coeff]: coeffs[itgt_coeff], [isrc_coeff]: src_expansions[isrc_box, isrc_coeff])  {id=update}
 │   for itgt_coeff_1
 └     preprocessed_src_expansions[isrc_box, itgt_coeff_1] = coeffs[itgt_coeff_1]  {id=insn}
   end isrc_box, itgt_coeff_1
---------------------------------------------------------------------------
---------------------------------------------------------------------------
KERNEL: m2l_preprocess
---------------------------------------------------------------------------
ARGUMENTS:
coeff: type: <auto/runtime>, shape: (9), dim_tags: (N0:stride:1) aspace: global
src_coeffs: type: <auto/runtime>, shape: (3), dim_tags: (N0:stride:1) aspace: global
---------------------------------------------------------------------------
DOMAINS:
{ [icoeff] : 0 <= icoeff <= 8 }
---------------------------------------------------------------------------
INAME TAGS:
icoeff: None
---------------------------------------------------------------------------
INSTRUCTIONS:
↱  coeff[0] = src_coeffs[0]  {id=coeff_insn_0}
└↱ coeff[3] = src_coeffs[1]  {id=coeff_insn_1}
 └ coeff[1] = src_coeffs[2]  {id=coeff_insn_2}
--------------------------------------------------------------------------

Pickled file m2l_preprocess.knl.tar.gz

[inducer]

I think the issue here is a bad error message. coeff is an output argument of m2l_preprocess, but the call site doesn't have it as an output argument.

• [ ] We should catch these inconsistencies before they cause confusion.
• [ ] We don't currently document kernel callables very well, including input and output arguments. We should fix that.