from typing import Tuple import pykokkos as pk from parse_args import parse_args @pk.workload(y=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), x=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), A=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight)) class Workload: def __init__(self, N: int, M: int, E: int, nrepeat: int, fill: bool): self.N: int = N self.M: int = M self.E: int = E self.nrepeat: int = nrepeat self.y: pk.View2D[pk.double] = pk.View([E, N], pk.double, layout=pk.Layout.LayoutRight) self.x: pk.View2D[pk.double] = pk.View([E, M], pk.double, layout=pk.Layout.LayoutRight) self.A: pk.View3D[pk.double] = pk.View([E, N, M], pk.double, layout=pk.Layout.LayoutRight) if fill: self.y.fill(1) self.x.fill(1) self.A.fill(1) else:
from typing import Tuple import pykokkos as pk from parse_args import parse_args @pk.functor(y=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), x=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), A=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight)) class Workload: def __init__(self, N: int, M: int, E: int, fill: bool): self.N: int = N self.M: int = M self.y: pk.View2D[pk.double] = pk.View([E, N], pk.double, layout=pk.Layout.LayoutRight) self.x: pk.View2D[pk.double] = pk.View([E, M], pk.double, layout=pk.Layout.LayoutRight) self.A: pk.View3D[pk.double] = pk.View([E, N, M], pk.double, layout=pk.Layout.LayoutRight) if fill: self.y.fill(1) self.x.fill(1) self.A.fill(1) else: for e in range(E): for i in range(N):
import pykokkos as pk import argparse import sys @pk.workload(inp=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), out=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight)) class main: def __init__(self, iterations, n, tile_size, star, radius): self.iterations: int = iterations self.n: int = n self.tile_size: int = tile_size self.star: int = star self.radius: int = radius self.inp: pk.View2D[pk.double] = pk.View([self.n, self.n], pk.double, layout=pk.Layout.LayoutRight) self.out: pk.View2D[pk.double] = pk.View([self.n, self.n], pk.double, layout=pk.Layout.LayoutRight) self.norm: float = 0 self.stencil_time: float = 0 @pk.main def run(self): t: int = tile_size r: int = radius
import pykokkos as pk import argparse import sys @pk.workload(A=pk.ViewTypeInfo(layout=pk.LayoutRight), B=pk.ViewTypeInfo(layout=pk.LayoutRight)) class main: def __init__(self, iterations, order, tile_size, permute): self.iterations: int = iterations self.order: int = order self.tile_size: int = tile_size self.permute: int = permute self.A: pk.View2D[pk.double] = pk.View([self.order, self.order], pk.double, layout=pk.LayoutRight) self.B: pk.View2D[pk.double] = pk.View([self.order, self.order], pk.double, layout=pk.LayoutRight) self.abserr: float = 0 self.transpose_time: float = 0 self.addit: float = (self.iterations) * (0.5 * (self.iterations - 1)) @pk.main def run(self): pk.parallel_for( pk.MDRangePolicy([0, 0], [self.order, self.order], [self.tile_size, self.tile_size]), self.init)
vz: float = ans - 0.5 mass_i_sqrt = math.sqrt(mass_i) v[i][0] = vx / mass_i_sqrt v[i][1] = vy / mass_i_sqrt v[i][2] = vz / mass_i_sqrt total_mass += mass_i total_momentum_x += mass_i * v[i][0] total_momentum_y += mass_i * v[i][1] total_momentum_z += mass_i * v[i][2] return total_mass, total_momentum_x, total_momentum_y, total_momentum_z @pk.functor(x=pk.ViewTypeInfo(layout=pk.LayoutRight)) class init_system: def __init__(self, s: System, ix_start: int, ix_end: int, iy_start: int, iy_end: int, iz_start: int, iz_end: int, lattice_constant: float, basis: pk.View2D[pk.double]): self.ix_start: int = ix_start self.ix_end: int = ix_end self.iy_start: int = iy_start self.iy_end: int = iy_end self.iz_start: int = iz_start self.iz_end: int = iz_end self.lattice_constant: float = lattice_constant self.basis: pk.View2D[pk.double] = basis self.x: pk.View2D[pk.double] = s.x self.type: pk.View1D[pk.int32] = s.type
import pykokkos as pk @pk.workload(subview=pk.ViewTypeInfo(trait=pk.Unmanaged)) class Workload: def __init__(self): self.view: pk.View2D[pk.int32] = pk.View([10, 10], pk.int32) self.subview: pk.View1D[pk.int32] = self.view[3, 2:5] @pk.main def run(self): pk.parallel_for(10, self.work) @pk.workunit def work(self, i: int): self.view[i][i] = 1 @pk.callback def callback(self) -> None: print(self.view) if __name__ == "__main__": pk.execute(pk.ExecutionSpace.Default, Workload())
import math import random import pykokkos as pk @pk.workload(count=pk.ViewTypeInfo(trait=pk.Atomic)) class SimpleAtomics: def __init__(self, n): self.N: int = n self.data: pk.View1D[pk.int32] = pk.View([n], pk.int32) self.result: pk.View1D[pk.int32] = pk.View([n], pk.int32) self.count: pk.View1D[pk.int32] = pk.View([1], pk.int32, trait=pk.Trait.Atomic) for i in range(n): self.data[i] = random.randint(0, n) @pk.main def run(self): pk.parallel_for(self.N, self.findprimes) @pk.callback def results(self): for i in range(int(self.count[0])): print(int(self.result[i]), end=", ") print("\nFound", int(self.count[0]), "prime numbers in", self.N, "random numbers")
import unittest import pykokkos as pk # Tests translation of KOKKOS_FUNCTIONS to C++ @pk.functor(subview_1=pk.ViewTypeInfo(trait=pk.Unmanaged), subview_2=pk.ViewTypeInfo(trait=pk.Unmanaged)) class KokkosFunctionsTestReduceFunctor: def __init__(self, threads: int, i_1: int, i_2: int, f_1: float, f_2: float, b_1: bool): self.threads: int = threads self.i_1: int = i_1 self.i_2: int = i_2 self.f_1: float = f_1 self.f_2: float = f_2 self.b_1: bool = b_1 self.view1D: pk.View1D[pk.int32] = pk.View([threads], pk.int32) self.view2D: pk.View2D[pk.int32] = pk.View([threads, threads], pk.int32) self.view3D: pk.View3D[pk.int32] = pk.View([threads, threads, threads], pk.int32) self.subview_1: pk.View1D[pk.int32] = self.view1D[threads // 2:] self.subview_2: pk.View2D[pk.int32] = self.view2D[threads // 2:, :] @pk.function def return_constant(self) -> int: return 5
import pykokkos as pk from integrator import Integrator from system import System from types_h import t_x, t_v, t_f, t_type, t_mass, t_id @pk.workload(x=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight)) class InitialIntegrateFunctor: def __init__( self, x: t_x, v: t_v, f: t_f, t: t_type, mass: t_mass, i: t_id, dtf: float, dtv: float, step: int, N_local: int ): self.x: pk.View2D[pk.double] = x self.v: pk.View2D[pk.double] = v self.f: pk.View2D[pk.double] = f self.type: pk.View1D[pk.int32] = t self.mass: pk.View1D[pk.double] = mass self.id: pk.View1D[pk.int32] = i self.dtf: float = dtf self.dtv: float = dtv self.step: int = step self.N_local: int = N_local @pk.main def run(self) -> None: pk.parallel_for("IntegratorNVE::initial_integrate", self.N_local, self.integrate)
class Particle: def __init__(self): self.x: float = 0.0 self.y: float = 0.0 self.z: float = 0.0 self.vx: float = 0.0 self.vy: float = 0.0 self.vz: float = 0.0 self.mass: float = 0.0 self.q: float = 0.0 self.id: int = 0 self.type: int = 0 @pk.workload(x=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), pack_indicies_all=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), pack_indicies=pk.ViewTypeInfo(trait=pk.Unmanaged)) class CommSerial(Comm): def __init__(self, s: System, comm_depth: float): super().__init__(s, comm_depth) # copied from System self.domain_x: float = 0.0 self.domain_y: float = 0.0 self.domain_z: float = 0.0 self.sub_domain_x: float = 0.0 self.sub_domain_y: float = 0.0 self.sub_domain_z: float = 0.0
class MyView2D(pk.View): def __init__(self, x: int, y: int, dtype: pk.DataTypeClass = pk.int32): super().__init__([x, y], dtype) class MyView3D(pk.View): def __init__(self, x: int, y: int, z: int, dtype: pk.DataTypeClass = pk.int32): super().__init__([x, y, z], dtype) # Tests for correctness of pk.View @pk.functor(subview1D=pk.ViewTypeInfo(trait=pk.Unmanaged), subview2D=pk.ViewTypeInfo(trait=pk.Unmanaged), subview3D=pk.ViewTypeInfo(trait=pk.Unmanaged)) class ViewsTestFunctor: def __init__(self, threads: int, i_1: int, i_2: int, i_3: int, i_4: int): self.threads: int = threads self.i_1: int = i_1 self.i_2: int = i_2 self.i_3: int = i_3 self.i_4: int = i_4 self.view1D: pk.View1D[pk.int32] = pk.View([i_1], pk.int32) self.view2D: pk.View2D[pk.int32] = pk.View([i_1, i_2], pk.int32) self.view3D: pk.View3D[pk.int32] = pk.View([i_1, i_2, i_3], pk.int32) self.myView1D: pk.View1D[int] = MyView1D(i_1, int)
from typing import Tuple import pykokkos as pk from parse_args import parse_args @pk.workunit( y=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), x=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight), A=pk.ViewTypeInfo(layout=pk.Layout.LayoutRight)) def yAx(team_member: pk.TeamMember, acc: pk.Acc[float], N:int, M: int, y: pk.View2D[pk.double], x: pk.View2D[pk.double], A: pk.View3D[pk.double]): e: int = team_member.league_rank() def team_reduce(j: int, team_acc: pk.Acc[float]): def vector_reduce(i: int, vector_acc: pk.Acc[float]): vector_acc += A[e][j][i] * x[e][i] tempM: float = pk.parallel_reduce( pk.ThreadVectorRange(team_member, M), vector_reduce) team_acc += y[e][j] * tempM tempN: float = pk.parallel_reduce( pk.TeamThreadRange(team_member, N), team_reduce) def single_closure(): nonlocal acc acc += tempN pk.single(pk.PerTeam(team_member), single_closure)