Batch PyKokkos Kernels

[kennykos]

The example script launches B small kernels with requisite overhead
that could be grouped into one larger kernel to reduce launch overhead.
It would be nice if there was pykokkos functionality to indicate kernel batching
while keeping the current structure of the outer for loop.

import cupy as cp
import pykokkos as pk

@pk.workunit
def work(wid, a): 
    a[wid] = a[wid] + 1 

def main():
    B = 10
    N = 10
    a = cp.ones((B, N)) 
    pk.set_default_space(pk.Cuda)
    for batch in range(B):
        pk.parallel_for("work", 10, work, a=a[batch])
    print(a)

main()

[gliga]

This would be similar to our work on kernel fusion: https://users.ece.utexas.edu/~gligoric/papers/AlAwarETAL25PyFuser.pdf
Have you tried running this with fusion on?

[IvanGrigorik]

Indeed, fuser works
The problem that I and Gabriel noted that there is no documentation on how to use it at all.
Here is a quick proof that the fuser works and how to enable it:

Initial code:

import cupy as cp
import pykokkos as pk

@pk.workunit
def work(wid, a): 
    a[wid] = a[wid] + 1 

def main():
    B = 100000
    N = 10
    a = cp.ones((B, N)) 
    pk.set_default_space(pk.Cuda)
    for batch in range(B):
        pk.parallel_for("work", 10, work, a=a[batch])
    print(a)

main()

Time required with no fusion (unset PK_FUSION)

real    0m27.213s
user    0m35.134s
sys     0m0.990s

Time required with fusion (export PK_FUSION="naive")

real    0m14.840s
user    0m22.729s
sys     0m1.136s

Both times shows second-time program execution

[gliga]

To the best of my knowledge, we are working on updating our documentation. This can be something to be documented as well.