Kokkos Tools in Python

Annotations

If you have used Kokkos::pushRegion(...), Kokkos::popRegion()`, etc. in your Kokkos C++ code, all of these features are available. The primary difference with respect to using Kokkos Tools in Python is that Python does not support argument references, e.g.:

uint32_t idx;
Kokkos::Tools::createProfileSection("mysection", &idx);
// ...

in C/C++, the tools callback can assign a value to idx but in Python, any changes to the arguments in Python are not propagated to the calling function. Thus, in Python, these tools must return the value:

idx = kokkos.tools.create_profile_section("mysection")

The full set of annotation functions currently available are shown below:

import kokkos

print("Profile library loaded: ".format(kokkos.tools.profile_library_loaded()))

kokkos.tools.push_region("myregion")
# ...
kokkos.tools.pop_region()

idx = kokkos.tools.create_profile_section("mysection")
kokkos.tools.start_section(idx)
# ...
kokkos.tools.stop_section(idx)
kokkos.tools.destroy_profile_section(idx)

kokkos.tools.mark_event("myevent")

kokkos.tools.declare_metadata("dogs", "good")
kokkos.tools.declare_metadata("cats", "good")
kokkos.tools.declare_metadata("fish", "eh")

Callbacks

Defining Callbacks

Python functions can be used to generate tool callbacks. In general, all the function signatures as the same, with a few exceptions, which are noted below by return statements and """RETURN INDEX""" instead of pass:

def print_help(arg0):
    pass


def parse_args(argc, argv):
    pass


def init(seq, ver, deviceCount, deviceInfo):
    pass


def finalize():
    pass


def begin_parallel_for(name, devid):
    """RETURN INDEX"""
    return index


def end_parallel_for(kernid):
    pass


def begin_parallel_reduce(name, devid):
    """RETURN INDEX"""
    return index


def end_parallel_reduce(kernid):
    pass


def begin_parallel_scan(name, devid):
    """RETURN INDEX"""
    return index


def end_parallel_scan(kernid):
    pass


def begin_fence(name, devid):
    """RETURN INDEX"""
    return index


def end_fence(kernid):
    pass


def push_region(name):
    pass


def pop_region():
    pass


def allocate_data(handle, label, ptr, size):
    pass


def deallocate_data(handle, label, ptr, size):
    pass


def create_profile_section(name):
    """RETURN INDEX"""
    return index


def start_profile_section(secid):
    pass


def stop_profile_section(secid):
    pass


def destroy_profile_section(secid):
    pass


def profile_event(name):
    pass


def begin_deep_copy(dst_handle, dst_name, dst_ptr, src_handle, src_name, src_ptr, size):
    pass


def end_deep_copy():
    pass

Assigning callbacks

kokkos.tools.set_parse_args_callback(parse_args)
kokkos.tools.set_print_help_callback(print_help)

kokkos.tools.set_init_callback(init)
kokkos.tools.set_finalize_callback(finalize)

kokkos.tools.set_begin_parallel_for_callback(begin_parallel_for)
kokkos.tools.set_end_parallel_for_callback(end_parallel_for)

kokkos.tools.set_begin_parallel_reduce_callback(begin_parallel_reduce)
kokkos.tools.set_end_parallel_reduce_callback(end_parallel_reduce)

kokkos.tools.set_begin_parallel_scan_callback(begin_parallel_scan)
kokkos.tools.set_end_parallel_scan_callback(end_parallel_scan)

kokkos.tools.set_begin_fence_callback(begin_fence)
kokkos.tools.set_end_fence_callback(end_fence)

kokkos.tools.set_push_region_callback(push_region)
kokkos.tools.set_pop_region_callback(pop_region)

kokkos.tools.set_allocate_data_callback(allocate_data)
kokkos.tools.set_deallocate_data_callback(deallocate_data)

kokkos.tools.set_create_profile_section_callback(create_profile_section)
kokkos.tools.set_start_profile_section_callback(start_profile_section)
kokkos.tools.set_stop_profile_section_callback(stop_profile_section)
kokkos.tools.set_destroy_profile_section_callback(destroy_profile_section)

kokkos.tools.set_profile_event_callback(profile_event)

kokkos.tools.set_begin_deep_copy_callback(begin_deep_copy)
kokkos.tools.set_end_deep_copy_callback(end_deep_copy)

Note, this can be highly condensed with Python if callbacks are named appropriately:

import sys
for itr in ("parse_args", "print_help", "init", "finalize",
            "push_region", "pop_region", "allocate_data", "deallocate_data",
            "create_profile_section", "start_profile_section",
            "stop_profile_section", "destroy_profile_section",
            "begin_parallel_for", "begin_parallel_reduce",
            "begin_parallel_scan", "begin_fence", "begin_deep_copy",
            "end_parallel_for", "end_parallel_reduce",
            "end_parallel_scan", "end_fence", "end_deep_copy",
            "profile_event",):
    func = getattr(sys.modules[__name__], f"{itr}")
    getattr(kokkos.tools, f"set_{itr}_callback")(func)