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)
