def test_matmul_input_validation():
namespace = pt.Namespace()
a = pt.make_placeholder(namespace,
name="a",
shape=(10, 10),
dtype=np.float)
b = pt.make_placeholder(namespace,
name="b",
shape=(20, 10),
dtype=np.float)
with pytest.raises(ValueError):
a @ b
c = pt.make_placeholder(namespace, name="c", shape=(), dtype=np.float)
with pytest.raises(ValueError):
c @ c
pt.make_size_param(namespace, "n")
d = pt.make_placeholder(namespace,
name="d",
shape="(n, n)",
dtype=np.float)
d @ d
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)
ns = pt.Namespace()
pt.make_size_param(ns, "nelements")
pt.make_size_param(ns, "nnodes")
u = pt.make_placeholder(ns,
name="u",
shape="(nelements, nnodes)",
dtype=np.float64)
op = AdvectionOperator(discr,
c=1,
flux_type="central",
dg_ops=DGOps1D(discr, ns))
result = op.apply(u)
prog = pt.generate_loopy(result, target=pt.PyOpenCLTarget(queue))
print(prog.program)
u = np.sin(discr.nodes())
print(prog(u=u.reshape(nelements, nnodes))[1][0])
def test_call_loopy_shape_inference():
from pytato.loopy import call_loopy
from pytato.utils import are_shapes_equal
import loopy as lp
knl = lp.make_kernel(
["{[i, j]: 0<=i<(2*n + 3*m + 2) and 0<=j<(6*n + 4*m + 3)}",
"{[ii, jj]: 0<=ii<m and 0<=jj<n}"],
"""
<> tmp = sum([i, j], A[i, j])
out[ii, jj] = tmp*(ii + jj)
""", lang_version=(2018, 2))
# {{{ variant 1
A = pt.make_placeholder(name="x", shape=(20, 37), dtype=np.float64) # noqa: N806
y = call_loopy(knl, {"A": A})["out"]
assert are_shapes_equal(y.shape, (4, 3))
# }}}
# {{{ variant 2
n1 = pt.make_size_param("n1")
n2 = pt.make_size_param("n2")
A = pt.make_placeholder(name="x", # noqa: N806
shape=(4*n1 + 6*n2 + 2, 12*n1 + 8*n2 + 3),
dtype=np.float64)
y = call_loopy(knl, {"A": A})["out"]
assert are_shapes_equal(y.shape, (2*n2, 2*n1))
class DGOps1D(AbstractDGOps1D):
@AbstractDGOps1D.array_ops.getter
def array_ops(self):
return pt
def __init__(self, discr):
self.discr = discr
nelements = pt.make_size_param("nelements")
nnodes = pt.make_size_param("nnodes")
# {{{ DG data
@cached_property
def normals(self):
return pt.make_data_wrapper(self.discr.normals, shape=(self.nelements, 2))
@cached_property
def interp_mat(self):
return pt.make_data_wrapper(self.discr.interp, shape=(2, self.nnodes))
@cached_property
def inv_mass_mat(self):
return pt.make_data_wrapper(self.discr.inv_mass, shape=(self.nnodes,
self.nnodes))
@cached_property
def stiffness_mat(self):
return pt.make_data_wrapper(self.discr.stiffness, shape=(self.nnodes,
self.nnodes))
@cached_property
def face_mass_mat(self):
return pt.make_data_wrapper(self.discr.face_mass, shape=(self.nnodes, 2))
# }}}
def interp(self, vec):
return pt.matmul(self.interp_mat, vec.T).T
def inv_mass(self, vec):
return pt.matmul(self.inv_mass_mat, vec.T).T
def stiffness(self, vec):
return pt.matmul(self.stiffness_mat, vec.T).T
def face_mass(self, vec):
return pt.matmul(self.face_mass_mat, vec.T).T
def face_swap(self, vec):
return pt.stack(
(
pt.roll(vec[:, 1], +1),
pt.roll(vec[:, 0], -1)),
axis=1)
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_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)
x_in = np.arange(1., 10.).reshape(3, 3)
assert_allclose_to_numpy(pt.roll(x, shift=shift, axis=axis), queue,
{x: x_in})
def test_asciidag():
n = pt.make_size_param("n")
array = pt.make_placeholder(name="array", shape=n, dtype=np.float64)
stack = pt.stack([array, 2*array, array + 6])
y = stack @ stack.T
from pytato import get_ascii_graph
res = get_ascii_graph(y, use_color=False)
ref_str = r"""* Inputs
*-. Placeholder
|\ \
* | | IndexLambda
| |/
|/|
| * IndexLambda
|/
* Stack
|\
* | AxisPermutation
|/
* Einsum
* Outputs
"""
assert res == ref_str
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 test_size_param(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
namespace = pt.Namespace()
n = pt.make_size_param(namespace, name="n")
pt.make_placeholder(namespace, name="x", shape="(n,)", dtype=np.int)
prog = pt.generate_loopy(n, target=pt.PyOpenCLTarget(queue))
x_in = np.array([1, 2, 3, 4, 5])
_, (n_out, ) = prog(x=x_in)
assert n_out == 5
def test_data_wrapper(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
# Without name/shape
namespace = pt.Namespace()
x_in = np.array([1, 2, 3, 4, 5])
x = pt.make_data_wrapper(namespace, x_in)
prog = pt.generate_loopy(x, target=pt.PyOpenCLTarget(queue))
_, (x_out, ) = prog()
assert (x_out == x_in).all()
# With name/shape
namespace = pt.Namespace()
x_in = np.array([[1, 2], [3, 4], [5, 6]])
pt.make_size_param(namespace, "n")
x = pt.make_data_wrapper(namespace, x_in, name="x", shape="(n, 2)")
prog = pt.generate_loopy(x, target=pt.PyOpenCLTarget(queue))
_, (x_out, ) = prog()
assert (x_out == x_in).all()
def main():
ns = pt.Namespace()
pt.make_size_param(ns, "n")
array = pt.make_placeholder(ns, name="array", shape="n", dtype=np.float)
stack = pt.stack([array, 2*array, array + 6])
ns.assign("stack", stack)
result = stack @ stack.T
dot_code = pt.get_dot_graph(result)
with open(GRAPH_DOT, "w") as outf:
outf.write(dot_code)
logger.info("wrote '%s'", GRAPH_DOT)
dot_path = shutil.which("dot")
if dot_path is not None:
subprocess.run([dot_path, "-Tsvg", GRAPH_DOT, "-o", GRAPH_SVG], check=True)
logger.info("wrote '%s'", GRAPH_SVG)
else:
logger.info("'dot' executable not found; cannot convert to SVG")
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)
u_initial = np.sin(discr.nodes())
ns = pt.Namespace()
pt.make_size_param(ns, "nelements")
pt.make_size_param(ns, "nnodes")
u = pt.make_placeholder(ns,
name="u",
shape="(nelements, nnodes)",
dtype=np.float64)
op = AdvectionOperator(discr,
c=1,
flux_type=flux_type,
dg_ops=DGOps1D(discr, ns))
result = op.apply(u)
prog = pt.generate_loopy(result, target=pt.PyOpenCLTarget(queue))
u = rk4(
lambda t, y: prog(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_matmul_input_validation():
a = pt.make_placeholder(name="a", shape=(10, 10), dtype=np.float64)
b = pt.make_placeholder(name="b", shape=(20, 10), dtype=np.float64)
with pytest.raises(ValueError):
a @ b
c = pt.make_placeholder(name="c", shape=(), dtype=np.float64)
with pytest.raises(ValueError):
c @ c
n = pt.make_size_param("n")
d = pt.make_placeholder(name="d", shape=(n, n), dtype=np.float64)
d @ d
def test_toposortmapper():
n = pt.make_size_param("n")
array = pt.make_placeholder(name="array", shape=n, dtype=np.float64)
stack = pt.stack([array, 2*array, array + 6])
y = stack @ stack.T
tm = pt.transform.TopoSortMapper()
tm(y)
from pytato.array import (AxisPermutation, IndexLambda,
Placeholder, Einsum, SizeParam, Stack)
assert isinstance(tm.topological_order[0], SizeParam)
assert isinstance(tm.topological_order[1], Placeholder)
assert isinstance(tm.topological_order[2], IndexLambda)
assert isinstance(tm.topological_order[3], IndexLambda)
assert isinstance(tm.topological_order[4], Stack)
assert isinstance(tm.topological_order[5], AxisPermutation)
assert isinstance(tm.topological_order[6], Einsum)
import numpy as np
import pytato as pt
n = pt.make_size_param("n")
a = pt.make_placeholder(name="a", shape=(n, n), dtype=np.float64)
a2a = a @ (2 * a)
aat = a @ a.T
result = pt.DictOfNamedArrays({"a2a": a2a, "aat": aat})
# {{{ execute
import pyopencl as cl
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
prg = pt.generate_loopy(result, cl_device=queue.device)
a = np.random.randn(20, 20)
_, out = prg(queue, a=a)
assert np.allclose(out["a2a"], a @ (2 * a))
assert np.allclose(out["aat"], a @ a.T)
# }}}
# {{{ generate OpenCL code
prg = pt.generate_loopy(result)
import loopy as lp
print(lp.generate_code_v2(prg.program).device_code())
