MAINT, BUG, DOC: various issues with functor workunit flexibility

[tylerjereddy]

While working on gh-177, I spent quite some time dealing with weird/unhelpful error messages and/or limitations that arise with functors that don't arise with isolated workunits. For example,

import pykokkos as pk


@pk.functor
class Workload_gh_173:
    def __init__(self, s: int):
        self.size: pk.int = s


    @pk.workunit
    def store_result(self, i: int, result: pk.View1D[pk.int64]):
        result[i] = self.size


def test_gh_173():
    w = Workload_gh_173(900)
    result = pk.View([1], dtype=pk.int64)
    pk.parallel_for(1, w.store_result, result=result)

Fails after python runtests.py -t tests/test_regressions.py with a long traceback that isn't particularly helpful in identifying the problem:

_____________________________________________________________________________________________________________________________________________________________________________________ test_gh_173 ______________________________________________________________________________________________________________________________________________________________________________________

self = <pykokkos.core.translators.static.StaticTranslator object at 0x7f8d33942a70>, source = (['class Workload_gh_173:\n', '    def __init__(self, s: int):\n', '        self.size: pk.int = s\n', '\n', '\n', '    @pk.workunit\n', ...], 4)

    def translate_workunits(self, source: Tuple[List[str], int]) -> Tuple[Dict[cppast.DeclRefExpr, Tuple[str, cppast.MethodDecl]], bool]:
        """
        Translate the workunits
    
        :param source: the python source code of the workload
        :returns: a tuple of a dictionary mapping from workload name
            to a tuple of operation name and source, and a boolean
            indicating whether the workunit has a call to pk.rand()
        """
    
        node_visitor = WorkunitVisitor(
            {}, source, self.pk_members.views, self.pk_members.pk_workunits,
            self.pk_members.fields, self.pk_members.pk_functions,
            self.pk_members.classtype_methods, self.pk_import, debug=True)
    
        workunits: Dict[cppast.DeclRefExpr, Tuple[str, cppast.MethodDecl]] = {}
    
        has_rand_call: bool = False
        for n, w in self.pk_members.pk_workunits.items():
            try:
>               workunits[n] = node_visitor.visit(w)

pykokkos/core/translators/static.py:215: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <pykokkos.core.visitors.workunit_visitor.WorkunitVisitor object at 0x7f8d33942830>, node = <ast.FunctionDef object at 0x7f8d33943580>

    def visit(self, node):
        """Visit a node."""
        method = 'visit_' + node.__class__.__name__
        visitor = getattr(self, method, self.generic_visit)
>       return visitor(node)

/usr/lib/python3.10/ast.py:410: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <pykokkos.core.visitors.workunit_visitor.WorkunitVisitor object at 0x7f8d33942830>, node = <ast.FunctionDef object at 0x7f8d33943580>

    def visit_FunctionDef(self, node: ast.FunctionDef) -> Union[str, Tuple[str, cppast.MethodDecl]]:
        if self.is_nested_call(node):
            params: List[cppast.ParmVarDecl] = [a for a in self.visit(node.args)]
            body = cppast.CompoundStmt([self.visit(b) for b in node.body])
    
            workunit = cppast.LambdaExpr("[&]", params, body)
            self.nested_work_units[node.name] = workunit
    
            return ""
    
        else:
            operation: Optional[str] = self.get_operation_type(node)
            if operation is None:
>               self.error(node.args, "Incorrect types in workunit definition")

pykokkos/core/visitors/workunit_visitor.py:36: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <pykokkos.core.visitors.workunit_visitor.WorkunitVisitor object at 0x7f8d33942830>, node = <ast.arguments object at 0x7f8d33943550>, message = 'Incorrect types in workunit definition'

    def error(self, node, message):
>       visitors_util.error(self.src, self.debug, node, message)

pykokkos/core/visitors/pykokkos_visitor.py:677: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

src = (['class Workload_gh_173:\n', '    def __init__(self, s: int):\n', '        self.size: pk.int = s\n', '\n', '\n', '    @pk.workunit\n', ...], 4), debug = True, node = <ast.arguments object at 0x7f8d33943550>, message = 'Incorrect types in workunit definition'

    def error(src, debug: bool, node, message) -> None:
        if hasattr(node, "lineno"):
            print(f"\n\033[31m\033[01mError on line {node.lineno} \033[0m: {message}")
        else:
            print(f"\n\033[31m\033[01mError\033[0m: {message}")
    
        if debug:
            print("DEBUG AST:")
            pretty_print(node)
    
        if hasattr(node, "lineno"):
            print(src[0][node.lineno - src[1] - 1], end="")
            err_len = node.end_col_offset - node.col_offset if node.end_col_offset else 1
            print(" " * node.col_offset + "^" * err_len)
    
>       sys.exit("PyKokkos: Translation failed")
E       SystemExit: PyKokkos: Translation failed

pykokkos/core/visitors/visitors_util.py:141: SystemExit

During handling of the above exception, another exception occurred:

    def test_gh_173():
        w = Workload_gh_173(900)
        result = pk.View([1], dtype=pk.int64)
>       pk.parallel_for(1, w.store_result, result=result)

tests/test_regressions.py:18: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
pykokkos/interface/parallel_dispatch.py:171: in parallel_for
    func, args = runtime_singleton.runtime.run_workunit(
pykokkos/core/runtime.py:97: in run_workunit
    members: Optional[PyKokkosMembers] = self.precompile_workunit(workunit,policy.space)
pykokkos/core/runtime.py:68: in precompile_workunit
    members: Optional[PyKokkosMembers] = self.compiler.compile_object(module_setup, space, km.is_uvm_enabled())
pykokkos/core/compiler.py:149: in compile_object
    self.compile_entity(module_setup.main, module_setup, entity, parser.get_classtypes(), space, force_uvm, members)
pykokkos/core/compiler.py:191: in compile_entity
    functor, bindings, cast = translator.translate(entity, classtypes)
pykokkos/core/translators/static.py:77: in translate
    workunits, has_rand_call = self.translate_workunits(source)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <pykokkos.core.translators.static.StaticTranslator object at 0x7f8d33942a70>, source = (['class Workload_gh_173:\n', '    def __init__(self, s: int):\n', '        self.size: pk.int = s\n', '\n', '\n', '    @pk.workunit\n', ...], 4)

    def translate_workunits(self, source: Tuple[List[str], int]) -> Tuple[Dict[cppast.DeclRefExpr, Tuple[str, cppast.MethodDecl]], bool]:
        """
        Translate the workunits
    
        :param source: the python source code of the workload
        :returns: a tuple of a dictionary mapping from workload name
            to a tuple of operation name and source, and a boolean
            indicating whether the workunit has a call to pk.rand()
        """
    
        node_visitor = WorkunitVisitor(
            {}, source, self.pk_members.views, self.pk_members.pk_workunits,
            self.pk_members.fields, self.pk_members.pk_functions,
            self.pk_members.classtype_methods, self.pk_import, debug=True)
    
        workunits: Dict[cppast.DeclRefExpr, Tuple[str, cppast.MethodDecl]] = {}
    
        has_rand_call: bool = False
        for n, w in self.pk_members.pk_workunits.items():
            try:
                workunits[n] = node_visitor.visit(w)
                has_rand_call = has_rand_call or node_visitor.has_rand_call
                if node_visitor.has_rand_call:
                    workunit: cppast.MethodDecl = workunits[n][1]
                    self.add_rand_pool_state(workunit)
                    node_visitor.has_rand_call = False
            except:
                print(f"Translation of {w} failed")
>               sys.exit(1)
E               SystemExit: 1

pykokkos/core/translators/static.py:223: SystemExit
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- Captured stdout call ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

Error: Incorrect types in workunit definition
DEBUG AST:
arguments(
    posonlyargs=[],
    args=[
        arg(arg='self'),
        arg(
            arg='i',
            annotation=Name(id='int', ctx=Load())),
        arg(
            arg='result',
            annotation=Subscript(
                value=Attribute(
                    value=Name(id='pk', ctx=Load()),
                    attr='View1D',
                    ctx=Load()),
                slice=Attribute(
                    value=Name(id='pk', ctx=Load()),
                    attr='int64',
                    ctx=Load()),
                ctx=Load()))],
    kwonlyargs=[],
    kw_defaults=[],
    defaults=[])
Translation of <ast.FunctionDef object at 0x7f8d33943580> failed
=============================================================================================================================================================================== short test summary info ================================================================================================================================================================================
FAILED tests/test_regressions.py::test_gh_173 - SystemExit: 1

Conversely (and confusingly), if I extract the workunit outside of the functor, this type of construction works just fine:

import pykokkos as pk


@pk.workunit
def store_result(i: int, result: pk.View1D[pk.int64]):
    result[i] = 0


def test_gh_173():
    result = pk.View([1], dtype=pk.int64)
    pk.parallel_for(1, store_result, result=result)

I'm inclined to say this is a bug because the error message is just not clear enough to explain why this is allowed in one case and not the other.

To make matters worse, I get a different error message with this kind of construct:

import pykokkos as pk


@pk.functor
class Workload_gh_173:
    def __init__(self, s: int):
        self.size: pk.int = s


    @pk.workunit
    def store_result(self, i: int, j: int):
        printf("i:", i)


def test_gh_173():
    w = Workload_gh_173(900)
    pk.parallel_for(1, w.store_result, j=3)

The error message is basically nonsense to me:

    def test_gh_173():
        w = Workload_gh_173(900)
>       pk.parallel_for(1, w.store_result, j=3)

tests/test_regressions.py:17: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
pykokkos/interface/parallel_dispatch.py:171: in parallel_for
    func, args = runtime_singleton.runtime.run_workunit(
pykokkos/core/runtime.py:97: in run_workunit
    members: Optional[PyKokkosMembers] = self.precompile_workunit(workunit,policy.space)
pykokkos/core/runtime.py:68: in precompile_workunit
    members: Optional[PyKokkosMembers] = self.compiler.compile_object(module_setup, space, km.is_uvm_enabled())
pykokkos/core/compiler.py:149: in compile_object
    self.compile_entity(module_setup.main, module_setup, entity, parser.get_classtypes(), space, force_uvm, members)
pykokkos/core/compiler.py:197: in compile_entity
    cpp_setup.compile(output_dir, functor, cast, bindings, space, force_uvm, self.get_compiler())
pykokkos/core/cpp_setup.py:67: in compile
    self.invoke_script(output_dir, space, enable_uvm, compiler)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <pykokkos.core.cpp_setup.CppSetup object at 0x7efc0276eb30>, output_dir = PosixPath('pk_cpp/home/tyler/github_projects/pykokkos/runtests/test_regressions_Workload_gh_173/OpenMP'), space = <ExecutionSpace.OpenMP: 'OpenMP'>, enable_uvm = False, compiler = 'g++'

    def invoke_script(self, output_dir: Path, space: ExecutionSpace, enable_uvm: bool, compiler: str) -> None:
        """
        Invoke the compilation script
    
        :param output_dir: the base directory
        :param space: the execution space of the workload
        :param enable_uvm: whether to enable CudaUVMSpace
        :param compiler: what compiler to use
        """
    
        view_space: str = "Kokkos::HostSpace"
        if space is ExecutionSpace.Cuda:
            if enable_uvm:
                view_space = "Kokkos::CudaUVMSpace"
        if space is ExecutionSpace.HIP:
            if enable_uvm:
                view_space = "Kokkos::Experimental::HIPManagedSpace"
    
        space_value: str
        if space.value == "HIP":
            space_value = "Experimental::HIP"
        else:
            space_value = space.value
    
        view_layout: str = str(get_default_layout(get_default_memory_space(space)))
        view_layout = view_layout.split(".")[-1]
        view_layout = f"Kokkos::{view_layout}"
    
        precision: str = km.get_default_precision().__name__.split(".")[-1]
        lib_path: Path
        include_path: Path
        compiler_path: Path
        lib_path, include_path, compiler_path = self.get_kokkos_paths(space, compiler)
        compute_capability: str = self.get_cuda_compute_capability(compiler)
        lib_suffix: str = self.get_kokkos_lib_suffix(space)
    
        command: List[str] = [f"./{self.script}",
                              compiler,             # What compiler to use
                              self.module_file,     # Compilation target
                              space_value,          # Execution space
                              view_space,           # Argument views memory space
                              view_layout,          # Argument views memory layout
                              precision,            # Default real precision
                              str(lib_path),        # Path to Kokkos install lib/ directory
                              str(include_path),    # Path to Kokkos install include/ directory
                              compute_capability,   # Device compute capability
                              lib_suffix,           # The libkokkos* suffix identifying the gpu
                              str(compiler_path)]   # The path to the compiler to use
        compile_result = subprocess.run(command, cwd=output_dir, capture_output=True, check=False)
    
        if compile_result.returncode != 0:
            print(compile_result.stderr.decode("utf-8"))
            print(f"C++ compilation in {output_dir} failed")
>           sys.exit(1)
E           SystemExit: 1

pykokkos/core/cpp_setup.py:257: SystemExit
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- Captured stdout call ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
In file included from /home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/Kokkos_OpenMP.hpp:236,
                 from /home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/decl/Kokkos_Declare_OPENMP.hpp:49,
                 from /home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/KokkosCore_Config_DeclareBackend.hpp:47,
                 from /home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/Kokkos_Core.hpp:57,
                 from ./bindings.cpp:3:
/home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/OpenMP/Kokkos_OpenMP_Parallel.hpp: In instantiation of ‘static std::enable_if_t<((! std::is_void<typename Kokkos::RangePolicy<Traits ...>::work_tag>::value) && std::is_same<Enable, typename Kokkos::RangePolicy<Traits ...>::work_tag>::value)> Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::exec_work(const FunctorType&, Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::Member) [with Enable = pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result; FunctorType = pk_functor_Workload_gh_173<Kokkos::OpenMP>; Traits = {Kokkos::OpenMP, pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result}; std::enable_if_t<((! std::is_void<typename Kokkos::RangePolicy<Traits ...>::work_tag>::value) && std::is_same<Enable, typename Kokkos::RangePolicy<Traits ...>::work_tag>::value)> = void; typename Kokkos::RangePolicy<Traits ...>::work_tag = pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result; Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::Member = long unsigned int]’:
/home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/OpenMP/Kokkos_OpenMP_Parallel.hpp:91:16:   required from ‘static void Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::exec_range(const FunctorType&, Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::Member, Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::Member) [with FunctorType = pk_functor_Workload_gh_173<Kokkos::OpenMP>; Traits = {Kokkos::OpenMP, pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result}; Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::Member = long unsigned int]’
/home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/OpenMP/Kokkos_OpenMP_Parallel.hpp:138:17:   required from ‘void Kokkos::Impl::ParallelFor<FunctorType, Kokkos::RangePolicy<Traits ...>, Kokkos::OpenMP>::execute() const [with FunctorType = pk_functor_Workload_gh_173<Kokkos::OpenMP>; Traits = {Kokkos::OpenMP, pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result}]’
/home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/Kokkos_Parallel.hpp:176:18:   required from ‘void Kokkos::parallel_for(const string&, const ExecPolicy&, const FunctorType&) [with ExecPolicy = Kokkos::RangePolicy<Kokkos::OpenMP, pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result>; FunctorType = pk_functor_Workload_gh_173<Kokkos::OpenMP>; Enable = void; std::string = std::__cxx11::basic_string<char>]’
./bindings.cpp:12:313:   required from here
/home/tyler/python_310_pykokkos_work/lib/python3.10/site-packages/pykokkos_base-0.0.7-py3.10-linux-x86_64.egg/include/kokkos/OpenMP/Kokkos_OpenMP_Parallel.hpp:106:12: error: no match for call to ‘(const pk_functor_Workload_gh_173<Kokkos::OpenMP>) (Kokkos::Impl::ParallelFor<pk_functor_Workload_gh_173<Kokkos::OpenMP>, Kokkos::RangePolicy<Kokkos::OpenMP, pk_functor_Workload_gh_173<Kokkos::OpenMP>::store_result>, Kokkos::OpenMP>::WorkTag, const Member&)’
  106 |     functor(WorkTag{}, iwork);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~
In file included from ./bindings.cpp:9:
../functor.hpp:5:301: note: candidate: ‘void pk_functor_Workload_gh_173<ExecSpace>::operator()(const pk_functor_Workload_gh_173<ExecSpace>::store_result&, int32_t, int32_t) const [with ExecSpace = Kokkos::OpenMP; int32_t = int]’
    5 | template <class ExecSpace>struct pk_functor_Workload_gh_173{struct store_result{};int32_t size; pk_functor_Workload_gh_173(int32_t size, int32_t pk_randpool_num_states, int32_t pk_randpool_seed){ this->size= size; }; pk_functor_Workload_gh_173(int32_t size){ this->size= size; };KOKKOS_FUNCTION void operator()(const store_result& , int32_t i, int32_t j)const{ printf("i:",i); };};
      |                                                                                                                                                                                                                                                                                                             ^~~~~~~~
../functor.hpp:5:301: note:   candidate expects 3 arguments, 2 provided
/home/tyler/.linuxbrew/opt/binutils/bin/ld: cannot find ./bindings.cpp.o: No such file or directory
collect2: error: ld returned 1 exit status

C++ compilation in pk_cpp/home/tyler/github_projects/pykokkos/runtests/test_regressions_Workload_gh_173/OpenMP failed
=============================================================================================================================================================================== short test summary info ================================================================================================================================================================================
FAILED tests/test_regressions.py::test_gh_173 - SystemExit: 1

And if I extract the workunit from the functor, things work "ok" again, so I find this arrangement of behaviors really hard to reason about:

import pykokkos as pk


@pk.workunit
def store_result(i: int, j: int):
    printf("i:", i)


def test_gh_173():
    pk.parallel_for(1, store_result, j=3)

[JBludau]

From the messages it looks like the second, unused argument is not passed correctly when the workunit is inside a functor ...