Skip to content

Initialize multiple levels of stratch memory. #290

@kennykos

Description

@kennykos

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions