def test_same_placeholder_name_raises():
from pytato.diagnostic import NameClashError
x = pt.make_placeholder(name="arr", shape=(10, 4), dtype=float)
y = pt.make_placeholder(name="arr", shape=(10, 4), dtype=float)
with pytest.raises(NameClashError):
pt.generate_loopy(x+y)
n1 = pt.make_size_param("n")
n2 = pt.make_size_param("n")
x = pt.make_placeholder(name="arr", shape=(n1, n2), dtype=float)
with pytest.raises(NameClashError):
pt.generate_loopy(2*x)
def test_advection_convergence(order, flux_type):
errors = []
hs = []
import pytato as pt
import pyopencl as cl
cl_ctx = cl.create_some_context()
queue = cl.CommandQueue(cl_ctx)
for nelements in (8, 12, 16, 20):
discr = DGDiscr1D(0, 2 * np.pi, nelements=nelements, nnodes=order)
dg_ops = DGOps1D(discr)
u_initial = np.sin(discr.nodes())
u = pt.make_placeholder(name="u",
shape=(dg_ops.nelements, dg_ops.nnodes),
dtype=np.float64)
op = AdvectionOperator(discr, c=1, flux_type=flux_type, dg_ops=dg_ops)
result = op.apply(u)
prog = pt.generate_loopy(result, cl_device=queue.device)
u = rk4(lambda t, y: prog(queue, u=y.reshape(nelements, order))[1][0].
reshape(-1),
u_initial,
t_initial=0,
t_final=np.pi,
dt=0.01)
u_ref = -u_initial
hs.append(discr.h)
errors.append(integrate(discr, (u - u_ref)**2)**0.5)
eoc, _ = np.polyfit(np.log(hs), np.log(errors), 1)
assert eoc >= order - 0.1, eoc
def test_zero_length_arrays():
x = pt.make_placeholder("x", shape=(0, 4), dtype=float)
y = 2*x
assert y.shape == (0, 4)
knl = pt.generate_loopy(y).kernel
assert all(dom.is_empty() for dom in knl.domains if dom.total_dim() != 0)
def test_make_placeholder_noname():
x = pt.make_placeholder("x", shape=(10, 4), dtype=float)
y = 2*x
knl = pt.generate_loopy(y).kernel
assert x.name in knl.arg_dict
assert x.name in knl.get_read_variables()
def test_make_placeholder_noname():
ns = pt.Namespace()
x = pt.make_placeholder(ns, shape=(10, 4), dtype=float)
y = 2 * x
knl = pt.generate_loopy(y).program
assert x.name in knl.arg_dict
assert x.name in knl.get_read_variables()
def to_numpy(self, array):
import pytato as pt
prg = pt.generate_loopy(array).program
prog_kwargs = {
arg_name: self.ns[arg_name].data
for arg_name in prg.arg_dict if arg_name in self.ns
}
evt, (cl_array, ) = prg(self.queue, **prog_kwargs)
return cl_array.get(queue=self.queue)
def generate_code_for_partition(partition: GraphPartition) \
-> Dict[PartId, BoundProgram]:
"""Return a mapping of partition identifiers to their
:class:`pytato.target.BoundProgram`."""
from pytato import generate_loopy
part_id_to_prg = {}
for part in partition.parts.values():
d = DictOfNamedArrays(
{var_name: partition.var_name_to_result[var_name]
for var_name in part.output_names
})
part_id_to_prg[part.pid] = generate_loopy(d)
return part_id_to_prg
def test_transpose(ctx_factory):
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
shape = (2, 8)
from numpy.random import default_rng
rng = default_rng()
x_in = rng.random(size=shape)
namespace = pt.Namespace()
x = pt.make_data_wrapper(namespace, x_in)
prog = pt.generate_loopy(x.T, target=pt.PyOpenCLTarget(queue))
_, (x_out,) = prog()
assert (x_out == x_in.T).all()
def test_roll(ctx_factory, shift, axis):
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
namespace = pt.Namespace()
pt.make_size_param(namespace, "n")
x = pt.make_placeholder(namespace, name="x", shape=("n", "n"), dtype=np.float)
prog = pt.generate_loopy(
pt.roll(x, shift=shift, axis=axis),
target=pt.PyOpenCLTarget(queue))
x_in = np.arange(1., 10.).reshape(3, 3)
_, (x_out,) = prog(x=x_in)
assert (x_out == np.roll(x_in, shift=shift, axis=axis)).all()
def freeze(self, array):
import pytato as pt
import pyopencl.array as cla
if isinstance(array, pt.Placeholder):
cl_array = cla.empty(self.queue,
shape=array.shape,
dtype=array.dtype)
return pt.make_data_wrapper(self.ns, cl_array)
prg = pt.generate_loopy(array).program
prog_kwargs = {
arg_name: self.ns[arg_name].data
for arg_name in prg.arg_dict if arg_name in self.ns
and isinstance(self.ns[arg_name], pt.array.DataWrapper)
}
evt, (cl_array, ) = prg(self.queue, **prog_kwargs)
return pt.make_data_wrapper(self.ns, cl_array)
def test_axis_permutation(ctx_factory, axes):
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
ndim = len(axes)
shape = (3, 4, 5)[:ndim]
from numpy.random import default_rng
rng = default_rng()
x_in = rng.random(size=shape)
namespace = pt.Namespace()
x = pt.make_data_wrapper(namespace, x_in)
prog = pt.generate_loopy(
pt.transpose(x, axes),
target=pt.PyOpenCLTarget(queue))
_, (x_out,) = prog()
assert (x_out == np.transpose(x_in, axes)).all()
def assert_allclose_to_numpy(expr: Array,
queue: cl.CommandQueue,
parameters: Dict[Placeholder, Any] = {}):
"""
Raises an :class:`AssertionError`, if there is a discrepancy between *expr*
evaluated lazily via :mod:`pytato` and eagerly via :mod:`numpy`.
:arg queue: An instance of :class:`pyopencl.CommandQueue` to which the
generated kernel must be enqueued.
"""
np_result = NumpyBasedEvaluator(numpy, expr.namespace, parameters)(expr)
prog = pt.generate_loopy(expr, target=pt.PyOpenCLTarget(queue))
evt, (pt_result, ) = prog(
**{placeholder.name: data
for placeholder, data in parameters.items()})
assert pt_result.shape == np_result.shape
assert pt_result.dtype == np_result.dtype
numpy.testing.assert_allclose(np_result, pt_result)
def test_stack(ctx_factory, input_dims):
cl_ctx = ctx_factory()
queue = cl.CommandQueue(cl_ctx)
shape = (2, 2, 2)[:input_dims]
from numpy.random import default_rng
rng = default_rng()
x_in = rng.random(size=shape)
y_in = rng.random(size=shape)
namespace = pt.Namespace()
x = pt.make_data_wrapper(namespace, x_in)
y = pt.make_data_wrapper(namespace, y_in)
for axis in range(0, 1 + input_dims):
prog = pt.generate_loopy(
pt.stack((x, y), axis=axis),
target=pt.PyOpenCLTarget(queue))
_, (out,) = prog()
assert (out == np.stack((x_in, y_in), axis=axis)).all()
def main():
import pytato as pt
import pyopencl as cl
cl_ctx = cl.create_some_context()
queue = cl.CommandQueue(cl_ctx)
nelements = 20
nnodes = 3
discr = DGDiscr1D(0, 2 * np.pi, nelements=nelements, nnodes=nnodes)
dg_ops = DGOps1D(discr)
op = AdvectionOperator(discr, c=1, flux_type="central", dg_ops=dg_ops)
u = pt.make_placeholder(name="u",
shape=(dg_ops.nelements, dg_ops.nnodes),
dtype=np.float64)
result = op.apply(u)
prog = pt.generate_loopy(result, cl_device=queue.device)
print(prog.program)
u = np.sin(discr.nodes())
print(prog(queue, u=u.reshape(nelements, nnodes))[1][0])
def main():
x_in = np.random.randn(2, 2)
x = pt.make_data_wrapper(x_in)
y = pt.stack([x @ x.T, 2 * x, 42 + x])
y = y + 55
tm = TopoSortMapper()
tm(y)
from functools import partial
pfunc = partial(get_partition_id, tm.topological_order)
# Find the partitions
outputs = pt.DictOfNamedArrays({"out": y})
partition = find_partition(outputs, pfunc)
# Show the partitions
from pytato.visualization import get_dot_graph_from_partition
get_dot_graph_from_partition(partition)
# Execute the partitions
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
prg_per_partition = generate_code_for_partition(partition)
context = execute_partition(partition, prg_per_partition, queue)
final_res = [context[k] for k in outputs.keys()]
# Execute the unpartitioned code for comparison
prg = pt.generate_loopy(y)
_, (out, ) = prg(queue)
np.testing.assert_allclose([out], final_res)
print("Partitioning test succeeded.")
