Currently, it appears that parallel_reduce only allows for a sum reduction between threads. Kokkos has several options, would it be possible to access these through pykokkos?
To be concrete, I am trying to make a simple kernel that finds the maximum of an array, an example is below. I am able to use each thread's accumulator to find the max of the values that thread sees, but upon reduction those values are summed. It would be useful here to have other reduction operations (in this case min/max, but I can see products being very useful as well)
import pykokkos as pk
import cupy as cp
@pk.workunit
def array_max(wid, maxval, arr):
if maxval < arr[wid]:
maxval = arr[wid]
def main():
N = 10
arr = cp.random.uniform(low=0.0, high=10.0, size=N)
m = pk.parallel_reduce("max", N, array_max, arr=arr)
print(m) # currently prints the sum of thread maxes
main()
As discussed with @kennykos
Currently, it appears that
parallel_reduceonly allows for a sum reduction between threads. Kokkos has several options, would it be possible to access these through pykokkos?To be concrete, I am trying to make a simple kernel that finds the maximum of an array, an example is below. I am able to use each thread's accumulator to find the max of the values that thread sees, but upon reduction those values are summed. It would be useful here to have other reduction operations (in this case min/max, but I can see products being very useful as well)
As discussed with @kennykos