pykokkos icon indicating copy to clipboard operation
pykokkos copied to clipboard

Initialize multiple levels of stratch memory.

Open kennykos opened this issue 5 months ago • 0 comments

Kokkos supports calling multiple levels of scratch memory as in the code below

policy.set_scratch_size(0, PerTeam(bytes0))
      .set_scratch_size(1, PerTeam(bytes1));

but pykokkos throws an error when I call

pk.TeamPolicy(league_size, team_size).set_scratch_size(0, pk.PerTeam(bytes0))
                                     .set_scratch_size(1, pk.PerTeam(bytes1))

Error:

warmup run
Traceback (most recent call last):
  File "/work/09661/gkk345/ls6/3dcapsules/python/development/gpu_timing.py", line 312, in <module>
    run()
  File "/work/09661/gkk345/ls6/3dcapsules/python/development/gpu_timing.py", line 106, in run
    p2g(gpu_pre)
  File "/work/09661/gkk345/ls6/3dcapsules/python/development/particle2grid/_p2g.py", line 28, in p2g
    pk.execute(
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/parallel_dispatch.py", line 253, in execute
    runtime_singleton.runtime.run_workload(space, workload)
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/runtime.py", line 58, in run_workload
    members: PyKokkosMembers = self.compiler.compile_object(module_setup, space, km.is_uvm_enabled(), None, None, None, set())
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/compiler.py", line 185, in compile_object
    members = self.extract_members(entity, classtypes)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/compiler.py", line 364, in extract_members
    members.extract(entity, classtypes)
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/translators/members.py", line 100, in extract
    self.reduction_result_queue, self.timer_result_queue = self.get_queues(source, name, pk_import)
                                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/translators/members.py", line 158, in get_queues
    node_visitor.visit(main)
  File "/work/09661/gkk345/ls6/miniconda3/envs/pyse/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/visitors/kokkosmain_visitor.py", line 38, in visit_FunctionDef
    run_body += serializer.serialize(self.visit(statement))
                                     ^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/miniconda3/envs/pyse/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/visitors/pykokkos_visitor.py", line 510, in visit_Expr
    call: cppast.CallExpr = self.visit(node.value)
                            ^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/miniconda3/envs/pyse/lib/python3.11/ast.py", line 418, in visit
    return visitor(node)
           ^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/visitors/kokkosmain_visitor.py", line 298, in visit_Call
    policy = self.add_space_to_policy(policy)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/core/visitors/kokkosmain_visitor.py", line 424, in add_space_to_policy
    policy_constructor.add_template_param(space)
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'MemberCallExpr' object has no attribute 'add_template_param'

Machine: GPU-a100, cuda 11.4

kennykos avatar Sep 06 '24 18:09 kennykos