Newer
Older
Andreas Klöckner
committed
knl = lp.tag_inames(knl, dict(j="ilp"))
knl = lp.preprocess_kernel(knl, ctx.devices[0])
for k in lp.generate_loop_schedules(knl):
assert k.temporary_variables["a"].shape == (16,)
# }}}
Andreas Klöckner
committed
def test_write_parameter(ctx_factory):
dtype = np.float32
ctx = ctx_factory()
knl = lp.make_kernel(
Andreas Klöckner
committed
"{[i,j]: 0<=i,j<n }",
"""
a = sum((i,j), i*j)
b = sum(i, sum(j, i*j))
n = 15
""",
[
lp.GlobalArg("a", dtype, shape=()),
lp.GlobalArg("b", dtype, shape=()),
lp.ValueArg("n", np.int32, approximately=1000),
],
assumptions="n>=1")
Andreas Klöckner
committed
import pytest
with pytest.raises(RuntimeError):
Andreas Klöckner
committed
lp.CompiledKernel(ctx, knl).get_code()
Andreas Klöckner
committed
def test_arg_shape_guessing(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
Andreas Klöckner
committed
"{[i,j]: 0<=i,j<n }",
"""
a = 1.5 + sum((i,j), i*j)
b[i, j] = i*j
c[i+j, j] = b[j,i]
""",
[
lp.GlobalArg("a", shape=lp.auto),
lp.GlobalArg("b", shape=lp.auto),
lp.GlobalArg("c", shape=lp.auto),
Andreas Klöckner
committed
lp.ValueArg("n"),
],
assumptions="n>=1")
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code())
Andreas Klöckner
committed
def test_arg_guessing(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n }",
"""
a = 1.5 + sum((i,j), i*j)
b[i, j] = i*j
c[i+j, j] = b[j,i]
""",
assumptions="n>=1")
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code())
def test_arg_guessing_with_reduction(ctx_factory):
#logging.basicConfig(level=logging.DEBUG)
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n }",
"""
a = 1.5 + sum((i,j), i*j)
d = 1.5 + sum((i,j), b[i,j])
b[i, j] = i*j
c[i+j, j] = b[j,i]
""",
assumptions="n>=1")
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code())
def test_nonlinear_index(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n }",
"""
a[i*i] = 17
""",
[
lp.GlobalArg("a", shape="n"),
lp.ValueArg("n"),
],
assumptions="n>=1")
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code())
def test_triangle_domain(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n and i <= j}",
"a[i,j] = 17",
assumptions="n>=1")
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code())
def test_offsets_and_slicing(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
knl = lp.make_kernel(
"{[i,j]: 0<=i<n and 0<=j<m }",
"""
b[i,j] = 2*a[i,j]
""",
assumptions="n>=1 and m>=1",
default_offset=lp.auto)
knl = lp.tag_data_axes(knl, "a,b", "stride:auto,stride:1")
cknl = lp.CompiledKernel(ctx, knl)
a_full = cl.clrandom.rand(queue, (n, n), np.float64)
a_full_h = a_full.get()
b_full = cl.clrandom.rand(queue, (n, n), np.float64)
b_full_h = b_full.get()
a_sub = (slice(3, 10), slice(5, 10))
a = a_full[a_sub]
b_sub = (slice(3+3, 10+3), slice(5+4, 10+4))
b = b_full[b_sub]
b_full_h[b_sub] = 2*a_full_h[a_sub]
print(cknl.get_highlighted_code({"a": a.dtype}))
import numpy.linalg as la
assert la.norm(b_full.get() - b_full_h) < 1e-13
Andreas Klöckner
committed
def test_vector_ilp_with_prefetch(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
Andreas Klöckner
committed
"{ [i]: 0<=i<n }",
"out[i] = 2*a[i]",
[
# Tests that comma'd arguments interoperate with
Andreas Klöckner
committed
# argument guessing.
lp.GlobalArg("out,a", np.float32, shape=lp.auto),
"..."
])
knl = lp.split_iname(knl, "i", 128, inner_tag="l.0")
knl = lp.split_iname(knl, "i_outer", 4, outer_tag="g.0", inner_tag="ilp")
knl = lp.add_prefetch(knl, "a", ["i_inner", "i_outer_inner"])
cknl = lp.CompiledKernel(ctx, knl)
cknl.cl_kernel_info()
Andreas Klöckner
committed
import re
Andreas Klöckner
committed
assert len(list(re.finditer("barrier", code))) == 1
def test_convolution(ctx_factory):
knl = lp.make_kernel(
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
"{ [iimg, ifeat, icolor, im_x, im_y, f_x, f_y]: \
-f_w <= f_x,f_y <= f_w \
and 0 <= im_x < im_w and 0 <= im_y < im_h \
and 0<=iimg<=nimgs and 0<=ifeat<nfeats and 0<=icolor<ncolors \
}",
"""
out[iimg, ifeat, im_x, im_y] = sum((f_x, f_y, icolor), \
img[iimg, f_w+im_x-f_x, f_w+im_y-f_y, icolor] \
* f[ifeat, f_w+f_x, f_w+f_y, icolor])
""",
[
lp.GlobalArg("f", dtype, shape=lp.auto),
lp.GlobalArg("img", dtype, shape=lp.auto),
lp.GlobalArg("out", dtype, shape=lp.auto),
"..."
],
assumptions="f_w>=1 and im_w, im_h >= 2*f_w+1 and nfeats>=1 and nimgs>=0",
flags="annotate_inames",
defines=dict(ncolors=3))
f_w = 3
knl = lp.fix_parameters(knl, f_w=f_w)
ref_knl = knl
def variant_0(knl):
#knl = lp.split_iname(knl, "im_x", 16, inner_tag="l.0")
knl = lp.set_loop_priority(knl, "iimg,im_x,im_y,ifeat,f_x,f_y")
return knl
def variant_1(knl):
knl = lp.split_iname(knl, "im_x", 16, inner_tag="l.0")
knl = lp.set_loop_priority(knl, "iimg,im_x_outer,im_y,ifeat,f_x,f_y")
return knl
def variant_2(knl):
knl = lp.split_iname(knl, "im_x", 16, outer_tag="g.0", inner_tag="l.0")
knl = lp.split_iname(knl, "im_y", 16, outer_tag="g.1", inner_tag="l.1")
knl = lp.tag_inames(knl, dict(ifeat="g.2"))
knl = lp.add_prefetch(knl, "f[ifeat,:,:,:]")
knl = lp.add_prefetch(knl, "img", "im_x_inner, im_y_inner, f_x, f_y")
return knl
for variant in [
variant_2
]:
lp.auto_test_vs_ref(ref_knl, ctx, variant(knl),
parameters=dict(
im_w=128, im_h=128, f_w=f_w,
))
def test_convolution_with_nonzero_base(ctx_factory):
# This is kept alive as a test for domains that don't start at zero.
# These are a bad idea for split_iname, which places its origin at zero
# and therefore produces a first block that is odd-sized.
#
# Therefore, for real tests, check test_convolution further up.
ctx = ctx_factory()
dtype = np.float32
knl = lp.make_kernel(
"{ [iimg, ifeat, icolor, im_x, im_y, f_x, f_y]: \
-f_w <= f_x,f_y <= f_w \
and f_w <= im_x < im_w-f_w and f_w <= im_y < im_h-f_w \
and 0<=iimg<=nimgs and 0<=ifeat<nfeats and 0<=icolor<ncolors \
}",
out[iimg, ifeat, im_x-f_w, im_y-f_w] = sum((f_x, f_y, icolor), \
img[iimg, im_x-f_x, im_y-f_y, icolor] \
* f[ifeat, f_w+f_x, f_w+f_y, icolor])
""",
[
lp.GlobalArg("f", dtype, shape=lp.auto),
lp.GlobalArg("img", dtype, shape=lp.auto),
lp.GlobalArg("out", dtype, shape=lp.auto),
"..."
],
assumptions="f_w>=1 and im_w, im_h >= 2*f_w+1 and nfeats>=1 and nimgs>=0",
flags="annotate_inames",
def variant_0(knl):
#knl = lp.split_iname(knl, "im_x", 16, inner_tag="l.0")
knl = lp.set_loop_priority(knl, "iimg,im_x,im_y,ifeat,f_x,f_y")
return knl
def variant_1(knl):
knl = lp.split_iname(knl, "im_x", 16, inner_tag="l.0")
knl = lp.set_loop_priority(knl, "iimg,im_x_outer,im_y,ifeat,f_x,f_y")
for variant in [
variant_0,
]:
lp.auto_test_vs_ref(ref_knl, ctx, variant(knl),
parameters=dict(
def test_c_instruction(ctx_factory):
#logging.basicConfig(level=logging.DEBUG)
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n }",
[
lp.CInstruction("i", """
x = sin((float) i);
""", assignees="x"),
"a[i*i] = x",
],
[
lp.GlobalArg("a", shape="n"),
lp.ValueArg("n"),
lp.TemporaryVariable("x", np.float32),
],
assumptions="n>=1")
knl = lp.split_iname(knl, "i", 128, outer_tag="g.0", inner_tag="l.0")
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code())
def test_dependent_domain_insn_iname_finding(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel([
"{[isrc_box]: 0<=isrc_box<nsrc_boxes}",
"{[isrc,idim]: isrc_start<=isrc<isrc_end and 0<=idim<dim}",
],
"""
<> src_ibox = source_boxes[isrc_box]
<> isrc_start = box_source_starts[src_ibox]
<> isrc_end = isrc_start+box_source_counts_nonchild[src_ibox]
<> strength = strengths[isrc] {id=set_strength}
""",
[
lp.GlobalArg("box_source_starts,box_source_counts_nonchild",
None, shape=None),
lp.GlobalArg("strengths",
None, shape="nsources"),
assert "isrc_box" in knl.insn_inames("set_strength")
print(lp.CompiledKernel(ctx, knl).get_highlighted_code(
dict(
source_boxes=np.int32,
box_source_starts=np.int32,
box_source_counts_nonchild=np.int32,
strengths=np.float64,
def test_inames_deps_from_write_subscript(ctx_factory):
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n}",
"""
<> src_ibox = source_boxes[i]
<int32> something = 5
a[src_ibox] = sum(j, something) {id=myred}
""",
[
lp.GlobalArg("box_source_starts,box_source_counts_nonchild,a",
None, shape=None),
"..."])
assert "i" in knl.insn_inames("myred")
knl = lp.make_kernel(
"{[i,j,k]: 0<=i,j,k<n}",
"""
b = sum((i,j,k), a[i,j,k])
""",
[
lp.GlobalArg("box_source_starts,box_source_counts_nonchild,a",
None, shape=None),
"..."])
knl = lp.split_reduction_outward(knl, "j,k")
def test_modulo_indexing(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,j]: 0<=i<n and 0<=j<5}",
"""
b[i] = sum(j, a[(i+j)%n])
""",
[
lp.GlobalArg("a", None, shape="n"),
"..."
]
)
print(knl)
print(lp.CompiledKernel(ctx, knl).get_highlighted_code(
def test_rob_stroud_bernstein(ctx_factory):
ctx = ctx_factory()
# NOTE: tmp would have to be zero-filled beforehand
knl = lp.make_kernel(
"{[el, i2, alpha1,alpha2]: \
0 <= el < nels and \
0 <= i2 < nqp1d and \
0 <= alpha1 <= deg and 0 <= alpha2 <= deg-alpha1 }",
"""
Andreas Klöckner
committed
<> xi = qpts[1, i2] {inames=+el}
Andreas Klöckner
committed
<> aind = 0 {id=aind_init,inames=+i2:el}
tmp[el,alpha1,i2] = tmp[el,alpha1,i2] + w * coeffs[aind] \
Andreas Klöckner
committed
{id=write_tmp,inames=+alpha2}
w = w * r * ( deg - alpha1 - alpha2 ) / (1 + alpha2) \
{id=update_w,dep=init_w:write_tmp}
aind = aind + 1 \
{id=aind_incr,\
dep=aind_init:write_tmp:update_w, \
Andreas Klöckner
committed
inames=+el:i2:alpha1:alpha2}
# Must declare coeffs to have "no" shape, to keep loopy
# from trying to figure it out the shape automatically.
lp.GlobalArg("coeffs", None, shape=None),
"..."
],
knl = lp.fix_parameters(knl, nqp1d=7, deg=4)
knl = lp.split_iname(knl, "el", 16, inner_tag="l.0")
knl = lp.split_iname(knl, "el_outer", 2, outer_tag="g.0", inner_tag="ilp",
slabs=(0, 1))
knl = lp.tag_inames(knl, dict(i2="l.1", alpha1="unr", alpha2="unr"))
print(lp.CompiledKernel(ctx, knl).get_highlighted_code(
dict(
qpts=np.float32,
coeffs=np.float32,
tmp=np.float32,
def test_rob_stroud_bernstein_full(ctx_factory):
#logging.basicConfig(level=logging.DEBUG)
ctx = ctx_factory()
# NOTE: result would have to be zero-filled beforehand
knl = lp.make_kernel(
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
"{[el, i2, alpha1,alpha2, i1_2, alpha1_2, i2_2]: \
0 <= el < nels and \
0 <= i2 < nqp1d and \
0 <= alpha1 <= deg and 0 <= alpha2 <= deg-alpha1 and\
\
0 <= i1_2 < nqp1d and \
0 <= alpha1_2 <= deg and \
0 <= i2_2 < nqp1d \
}",
"""
<> xi = qpts[1, i2] {inames=+el}
<> s = 1-xi
<> r = xi/s
<> aind = 0 {id=aind_init,inames=+i2:el}
<> w = s**(deg-alpha1) {id=init_w}
<> tmp[alpha1,i2] = tmp[alpha1,i2] + w * coeffs[aind] \
{id=write_tmp,inames=+alpha2}
w = w * r * ( deg - alpha1 - alpha2 ) / (1 + alpha2) \
{id=update_w,dep=init_w:write_tmp}
aind = aind + 1 \
{id=aind_incr,\
dep=aind_init:write_tmp:update_w, \
inames=+el:i2:alpha1:alpha2}
<> xi2 = qpts[0, i1_2] {dep=aind_incr,inames=+el}
<> s2 = 1-xi2
<> r2 = xi2/s2
<> w2 = s2**deg
result[el, i1_2, i2_2] = result[el, i1_2, i2_2] + \
w2 * tmp[alpha1_2, i2_2] \
{inames=el:alpha1_2:i1_2:i2_2}
w2 = w2 * r2 * (deg-alpha1_2) / (1+alpha1_2)
""",
[
# Must declare coeffs to have "no" shape, to keep loopy
# from trying to figure it out the shape automatically.
lp.GlobalArg("coeffs", None, shape=None),
"..."
],
assumptions="deg>=0 and nels>=1"
)
knl = lp.fix_parameters(knl, nqp1d=7, deg=4)
if 0:
knl = lp.split_iname(knl, "el", 16, inner_tag="l.0")
knl = lp.split_iname(knl, "el_outer", 2, outer_tag="g.0", inner_tag="ilp",
slabs=(0, 1))
knl = lp.tag_inames(knl, dict(i2="l.1", alpha1="unr", alpha2="unr"))
from pickle import dumps, loads
knl = loads(dumps(knl))
knl = lp.CompiledKernel(ctx, knl).get_highlighted_code(
dict(
qpts=np.float32,
tmp=np.float32,
coeffs=np.float32,
result=np.float32,
))
@pytest.mark.parametrize("vec_len", [2, 3, 4, 8, 16])
def test_vector_types(ctx_factory, vec_len):
knl = lp.make_kernel(
"{ [i,j]: 0<=i<n and 0<=j<vec_len }",
"out[i,j] = 2*a[i,j]",
[
lp.GlobalArg("a", np.float32, shape=lp.auto),
lp.GlobalArg("out", np.float32, shape=lp.auto),
"..."
knl = lp.fix_parameters(knl, vec_len=vec_len)
ref_knl = knl
knl = lp.tag_data_axes(knl, "out", "c,vec")
knl = lp.tag_inames(knl, dict(j="unr"))
knl = lp.split_iname(knl, "i", 128, outer_tag="g.0", inner_tag="l.0")
lp.auto_test_vs_ref(ref_knl, ctx, knl,
parameters=dict(
n=20000
),
fills_entire_output=False)
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
def test_tag_data_axes(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{ [i,j,k]: 0<=i,j,k<n }",
"out[i,j,k] = 15")
ref_knl = knl
with pytest.raises(lp.LoopyError):
lp.tag_data_axes(knl, "out", "N1,N0,N5")
with pytest.raises(lp.LoopyError):
lp.tag_data_axes(knl, "out", "N1,N0,c")
knl = lp.tag_data_axes(knl, "out", "N1,N0,N2")
knl = lp.tag_inames(knl, dict(j="g.0", i="g.1"))
lp.auto_test_vs_ref(ref_knl, ctx, knl,
parameters=dict(n=20))
def test_conditional(ctx_factory):
#logging.basicConfig(level=logging.DEBUG)
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
knl = lp.make_kernel(
"{ [i,j]: 0<=i,j<n }",
"""
<> my_a = a[i,j] {id=read_a}
<> a_less_than_zero = my_a < 0 {dep=read_a,inames=i:j}
my_a = 2*my_a {id=twice_a,dep=read_a,if=a_less_than_zero}
my_a = my_a+1 {id=aplus,dep=twice_a,if=a_less_than_zero}
out[i,j] = 2*my_a {dep=aplus}
""",
[
lp.GlobalArg("a", np.float32, shape=lp.auto),
lp.GlobalArg("out", np.float32, shape=lp.auto),
"..."
])
ref_knl = knl
lp.auto_test_vs_ref(ref_knl, ctx, knl,
parameters=dict(
n=200
))
Andreas Klöckner
committed
def test_ilp_loop_bound(ctx_factory):
# The salient bit of this test is that a joint bound on (outer, inner)
# from a split occurs in a setting where the inner loop has been ilp'ed.
# In 'normal' parallel loops, the inner index is available for conditionals
# throughout. In ILP'd loops, not so much.
ctx = ctx_factory()
knl = lp.make_kernel(
Andreas Klöckner
committed
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
"{ [i,j,k]: 0<=i,j,k<n }",
"""
out[i,k] = sum(j, a[i,j]*b[j,k])
""",
[
lp.GlobalArg("a,b", np.float32, shape=lp.auto),
"...",
],
assumptions="n>=1")
ref_knl = knl
knl = lp.set_loop_priority(knl, "j,i,k")
knl = lp.split_iname(knl, "k", 4, inner_tag="ilp")
lp.auto_test_vs_ref(ref_knl, ctx, knl,
parameters=dict(
n=200
))
def test_arg_shape_uses_assumptions(ctx_factory):
# If arg shape determination does not use assumptions, then it won't find a
# static shape for out, which is at least 1 x 1 in size, but otherwise of
# size n x n.
lp.make_kernel(
"{ [i,j]: 0<=i,j<n }",
"""
out[i,j] = 2*a[i,j]
out[0,0] = 13.0
""", assumptions="n>=1")
def test_slab_decomposition_does_not_double_execute(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
knl = lp.make_kernel(
"{ [i]: 0<=i<n }",
ref_knl = knl
for outer_tag in ["for", "g.0"]:
knl = ref_knl
knl = lp.split_iname(knl, "i", 4, slabs=(0, 1), inner_tag="unr",
outer_tag=outer_tag)
knl = lp.set_loop_priority(knl, "i_outer")
a = cl.clrandom.rand(queue, 20, np.float32)
a_ref = a.copy()
a_knl = a.copy()
knl = lp.set_options(knl, "write_cl")
knl(queue, a=a_knl)
ref_knl(queue, a=a_ref)
queue.finish()
assert (a_ref == a_knl).get().all()
def test_multiple_writes_to_local_temporary(ctx_factory):
# Loopy would previously only handle barrier insertion correctly if exactly
# one instruction wrote to each local temporary. This tests that multiple
# writes are OK.
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i,e]: 0<=i<5 and 0<=e<nelements}",
"""
<> temp[i, 0] = 17
temp[i, 1] = 15
""")
knl = lp.tag_inames(knl, dict(i="l.0"))
knl = lp.preprocess_kernel(knl, ctx.devices[0])
for k in lp.generate_loop_schedules(knl):
code, _ = lp.generate_code(k)
def test_fd_demo(ctx_factory):
knl = lp.make_kernel(
"{[i,j]: 0<=i,j<n}",
"result[i,j] = u[i, j]**2 + -1 + (-4)*u[i + 1, j + 1] \
+ u[i + 1 + 1, j + 1] + u[i + 1 + -1, j + 1] \
+ u[i + 1, j + 1 + 1] + u[i + 1, j + 1 + -1]")
knl = lp.split_iname(knl,
"i", 16, outer_tag="g.1", inner_tag="l.1")
knl = lp.split_iname(knl,
"j", 16, outer_tag="g.0", inner_tag="l.0")
knl = lp.add_prefetch(knl, "u",
["i_inner", "j_inner"],
fetch_bounding_box=True)
#n = 1000
#u = cl.clrandom.rand(queue, (n+2, n+2), dtype=np.float32)
knl = lp.set_options(knl, write_cl=True)
knl = lp.add_and_infer_dtypes(knl, dict(u=np.float32))
knl = lp.preprocess_kernel(knl)
knl = lp.get_one_scheduled_kernel(knl)
code, inf = lp.generate_code(knl)
assert "double" not in code
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
def test_make_copy_kernel(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
intermediate_format = "f,f,sep"
a1 = np.random.randn(1024, 4, 3)
cknl1 = lp.make_copy_kernel(intermediate_format)
cknl1 = lp.fix_parameters(cknl1, n2=3)
cknl1 = lp.set_options(cknl1, write_cl=True)
evt, a2 = cknl1(queue, input=a1)
cknl2 = lp.make_copy_kernel("c,c,c", intermediate_format)
cknl2 = lp.fix_parameters(cknl2, n2=3)
evt, a3 = cknl2(queue, input=a2)
assert (a1 == a3).all()
def test_set_arg_order():
knl = lp.make_kernel(
"{ [i,j]: 0<=i,j<n }",
"out[i,j] = a[i]*b[j]")
knl = lp.set_argument_order(knl, "out,a,n,b")
def test_affine_map_inames():
knl = lp.make_kernel(
"{[e, i,j,n]: 0<=e<E and 0<=i,j,n<N}",
"rhsQ[e, n+i, j] = rhsQ[e, n+i, j] - D[i, n]*x[i,j]")
knl = lp.affine_map_inames(knl,
"i", "i0",
"i0 = n+i")
print(knl)
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
def test_poisson(ctx_factory):
# Stolen from Peter Coogan and Rob Kirby for FEM assembly
ctx = ctx_factory()
nbf = 5
nqp = 5
sdim = 3
knl = lp.make_kernel(
"{ [c,i,j,k,ell,ell2]: \
0 <= c < nels and \
0 <= i < nbf and \
0 <= j < nbf and \
0 <= k < nqp and \
0 <= ell < sdim and \
0 <= ell2 < sdim }",
"""
dpsi(bf,k0,dir) := sum(ell2, DFinv[c,ell2,dir] * DPsi[bf,k0,ell2] )
Ael[c,i,j] = J[c] * w[k] * sum(ell, dpsi(i,k,ell) * dpsi(j,k,ell))
""",
assumptions="nels>=1 and nbf >= 1 and nels mod 4 = 0")
knl = lp.fix_parameters(knl, nbf=nbf, sdim=sdim, nqp=nqp)
ref_knl = knl
knl = lp.set_loop_priority(knl, ["c", "j", "i", "k"])
def variant_1(knl):
knl = lp.precompute(knl, "dpsi", "i,k,ell", default_tag='for')
knl = lp.set_loop_priority(knl, "c,i,j")
return knl
def variant_2(knl):
knl = lp.precompute(knl, "dpsi", "i,ell", default_tag='for')
knl = lp.set_loop_priority(knl, "c,i,j")
return knl
def add_types(knl):
return lp.add_and_infer_dtypes(knl, dict(
w=np.float32,
J=np.float32,
DPsi=np.float32,
DFinv=np.float32,
))
for variant in [
#variant_1,
variant_2
]:
knl = variant(knl)
lp.auto_test_vs_ref(
add_types(ref_knl), ctx, add_types(knl),
parameters=dict(n=5, nels=15, nbf=5, sdim=2, nqp=7))
if __name__ == "__main__":
if len(sys.argv) > 1:
exec(sys.argv[1])
else:
from py.test.cmdline import main
main([__file__])
# vim: foldmethod=marker