def test_matmul(ctx_factory, buffer_inames):
ctx = ctx_factory()
if (buffer_inames and
ctx.devices[0].platform.name == "Portable Computing Language"):
pytest.skip("crashes on pocl")
logging.basicConfig(level=logging.INFO)
fortran_src = """
subroutine dgemm(m,n,ell,a,b,c)
implicit none
real*8 a(m,ell),b(ell,n),c(m,n)
integer m,n,k,i,j,ell
do j = 1,n
do i = 1,m
do k = 1,ell
c(i,j) = c(i,j) + b(k,j)*a(i,k)
end do
end do
end do
end subroutine
"""
knl, = lp.parse_fortran(fortran_src)
assert len(knl.domains) == 1
ref_knl = knl
knl = lp.split_iname(knl, "i", 16,
outer_tag="g.0", inner_tag="l.1")
knl = lp.split_iname(knl, "j", 8,
outer_tag="g.1", inner_tag="l.0")
knl = lp.split_iname(knl, "k", 32)
knl = lp.assume(knl, "n mod 32 = 0")
knl = lp.assume(knl, "m mod 32 = 0")
knl = lp.assume(knl, "ell mod 16 = 0")
knl = lp.extract_subst(knl, "a_acc", "a[i1,i2]", parameters="i1, i2")
knl = lp.extract_subst(knl, "b_acc", "b[i1,i2]", parameters="i1, i2")
knl = lp.precompute(knl, "a_acc", "k_inner,i_inner",
precompute_outer_inames='i_outer, j_outer, k_outer',
default_tag="l.auto")
knl = lp.precompute(knl, "b_acc", "j_inner,k_inner",
precompute_outer_inames='i_outer, j_outer, k_outer',
default_tag="l.auto")
knl = lp.buffer_array(knl, "c", buffer_inames=buffer_inames,
init_expression="0", store_expression="base+buffer")
lp.auto_test_vs_ref(ref_knl, ctx, knl, parameters=dict(n=128, m=128, ell=128))
def test_precompute_some_exist(ctx_factory):
fortran_src = """
subroutine dgemm(m,n,ell,a,b,c)
implicit none
real*8 a(m,ell),b(ell,n),c(m,n)
integer m,n,k,i,j,ell
do j = 1,n
do i = 1,m
do k = 1,ell
c(i,j) = c(i,j) + b(k,j)*a(i,k)
end do
end do
end do
end subroutine
"""
knl = lp.parse_fortran(fortran_src)
assert len(knl["dgemm"].domains) == 1
knl = lp.split_iname(knl, "i", 8, outer_tag="g.0", inner_tag="l.1")
knl = lp.split_iname(knl, "j", 8, outer_tag="g.1", inner_tag="l.0")
knl = lp.split_iname(knl, "k", 8)
knl = lp.assume(knl, "n mod 8 = 0")
knl = lp.assume(knl, "m mod 8 = 0")
knl = lp.assume(knl, "ell mod 8 = 0")
knl = lp.extract_subst(knl, "a_acc", "a[i1,i2]", parameters="i1, i2")
knl = lp.extract_subst(knl, "b_acc", "b[i1,i2]", parameters="i1, i2")
knl = lp.precompute(knl,
"a_acc",
"k_inner,i_inner",
precompute_inames="ktemp,itemp",
precompute_outer_inames="i_outer, j_outer, k_outer",
default_tag="l.auto")
knl = lp.precompute(knl,
"b_acc",
"j_inner,k_inner",
precompute_inames="itemp,k2temp",
precompute_outer_inames="i_outer, j_outer, k_outer",
default_tag="l.auto")
ref_knl = knl
ctx = ctx_factory()
lp.auto_test_vs_ref(ref_knl,
ctx,
knl,
parameters=dict(n=128, m=128, ell=128))
def test_assume(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel("{[i]: 0<=i<n}",
"a[i] = a[i] + 1",
[lp.GlobalArg("a", np.float32, shape="n"), "..."],
target=lp.PyOpenCLTarget(ctx.devices[0]))
knl = lp.split_iname(knl, "i", 16)
knl = lp.prioritize_loops(knl, "i_outer,i_inner")
knl = lp.assume(knl, "n mod 16 = 0")
knl = lp.assume(knl, "n > 10")
code = lp.generate_code_v2(knl).device_code()
assert "if" not in code
def test_matmul(ctx_factory, buffer_inames):
logging.basicConfig(level=logging.INFO)
fortran_src = """
subroutine dgemm(m,n,l,a,b,c)
implicit none
real*8 a(m,l),b(l,n),c(m,n)
integer m,n,k,i,j,l
do j = 1,n
do i = 1,m
do k = 1,l
c(i,j) = c(i,j) + b(k,j)*a(i,k)
end do
end do
end do
end subroutine
"""
knl, = lp.parse_fortran(fortran_src)
assert len(knl.domains) == 1
ref_knl = knl
knl = lp.split_iname(knl, "i", 16, outer_tag="g.0", inner_tag="l.1")
knl = lp.split_iname(knl, "j", 8, outer_tag="g.1", inner_tag="l.0")
knl = lp.split_iname(knl, "k", 32)
knl = lp.assume(knl, "n mod 32 = 0")
knl = lp.assume(knl, "m mod 32 = 0")
knl = lp.assume(knl, "l mod 16 = 0")
knl = lp.extract_subst(knl, "a_acc", "a[i1,i2]", parameters="i1, i2")
knl = lp.extract_subst(knl, "b_acc", "b[i1,i2]", parameters="i1, i2")
knl = lp.precompute(knl, "a_acc", "k_inner,i_inner")
knl = lp.precompute(knl, "b_acc", "j_inner,k_inner")
knl = lp.buffer_array(knl,
"c",
buffer_inames=buffer_inames,
init_expression="0",
store_expression="base+buffer")
ctx = ctx_factory()
lp.auto_test_vs_ref(ref_knl,
ctx,
knl,
parameters=dict(n=128, m=128, l=128))
def test_check_bounds_with_caller_assumptions(ctx_factory):
import islpy as isl
from loopy.diagnostic import LoopyIndexError
arange = lp.make_function("{[i]: 0<=i<n}",
"""
y[i] = i
""",
name="arange")
knl = lp.make_kernel(
"{[i]: 0<=i<20}",
"""
[i]: Y[i] = arange(N)
""",
[lp.GlobalArg("Y", shape=(20, )),
lp.ValueArg("N", dtype=np.int32)],
name="epoint")
knl = lp.merge([knl, arange])
with pytest.raises(LoopyIndexError):
lp.generate_code_v2(knl)
knl = knl.with_kernel(
lp.assume(knl.default_entrypoint, isl.BasicSet("[N] -> { : N <= 20}")))
lp.auto_test_vs_ref(knl, ctx_factory(), parameters={"N": 15})
def test_ispc_streaming_stores():
stream_dtype = np.float32
index_dtype = np.int32
knl = lp.make_kernel(
"{[i]: 0<=i<n}",
"a[i] = b[i] + scalar * c[i]",
target=lp.ISPCTarget(), index_dtype=index_dtype,
name="stream_triad")
vars = ["a", "b", "c", "scalar"]
knl = lp.assume(knl, "n>0")
knl = lp.split_iname(
knl, "i", 2**18, outer_tag="g.0", slabs=(0, 1))
knl = lp.split_iname(knl, "i_inner", 8, inner_tag="l.0")
knl = lp.tag_instructions(knl, "!streaming_store")
knl = lp.add_and_infer_dtypes(knl, {
var: stream_dtype
for var in vars
})
knl = lp.set_argument_order(knl, vars + ["n"])
knl = lp.preprocess_kernel(knl)
knl = lp.get_one_scheduled_kernel(knl)
lp.generate_code_v2(knl).all_code()
def test_ispc_streaming_stores():
stream_dtype = np.float32
index_dtype = np.int32
knl = lp.make_kernel(
"{[i]: 0<=i<n}",
"a[i] = b[i] + scalar * c[i]",
target=lp.ISPCTarget(), index_dtype=index_dtype,
name="stream_triad")
vars = ["a", "b", "c", "scalar"]
knl = lp.assume(knl, "n>0")
knl = lp.split_iname(
knl, "i", 2**18, outer_tag="g.0", slabs=(0, 1))
knl = lp.split_iname(knl, "i_inner", 8, inner_tag="l.0")
knl = lp.tag_instructions(knl, "!streaming_store")
knl = lp.add_and_infer_dtypes(knl, {
var: stream_dtype
for var in vars
})
knl = lp.set_argument_order(knl, vars + ["n"])
knl = lp.preprocess_kernel(knl)
knl = lp.get_one_scheduled_kernel(knl)
lp.generate_code_v2(knl).all_code()
def test_assume(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel("{[i]: 0<=i<n}", "a[i] = a[i] + 1",
[lp.GlobalArg("a", np.float32, shape="n"), "..."])
knl = lp.split_iname(knl, "i", 16)
knl = lp.set_loop_priority(knl, "i_outer,i_inner")
knl = lp.assume(knl, "n mod 16 = 0")
knl = lp.assume(knl, "n > 10")
knl = lp.preprocess_kernel(knl, ctx.devices[0])
kernel_gen = lp.generate_loop_schedules(knl)
for gen_knl in kernel_gen:
print(gen_knl)
compiled = lp.CompiledKernel(ctx, gen_knl)
print(compiled.get_code())
assert "if" not in compiled.get_code()
def test_matmul(ctx_factory, buffer_inames):
logging.basicConfig(level=logging.INFO)
fortran_src = """
subroutine dgemm(m,n,l,a,b,c)
implicit none
real*8 a(m,l),b(l,n),c(m,n)
integer m,n,k,i,j,l
do j = 1,n
do i = 1,m
do k = 1,l
c(i,j) = c(i,j) + b(k,j)*a(i,k)
end do
end do
end do
end subroutine
"""
knl, = lp.parse_fortran(fortran_src)
assert len(knl.domains) == 1
ref_knl = knl
knl = lp.split_iname(knl, "i", 16,
outer_tag="g.0", inner_tag="l.1")
knl = lp.split_iname(knl, "j", 8,
outer_tag="g.1", inner_tag="l.0")
knl = lp.split_iname(knl, "k", 32)
knl = lp.assume(knl, "n mod 32 = 0")
knl = lp.assume(knl, "m mod 32 = 0")
knl = lp.assume(knl, "l mod 16 = 0")
knl = lp.extract_subst(knl, "a_acc", "a[i1,i2]", parameters="i1, i2")
knl = lp.extract_subst(knl, "b_acc", "b[i1,i2]", parameters="i1, i2")
knl = lp.precompute(knl, "a_acc", "k_inner,i_inner")
knl = lp.precompute(knl, "b_acc", "j_inner,k_inner")
knl = lp.buffer_array(knl, "c", buffer_inames=buffer_inames,
init_expression="0", store_expression="base+buffer")
ctx = ctx_factory()
lp.auto_test_vs_ref(ref_knl, ctx, knl, parameters=dict(n=128, m=128, l=128))
def transform(knl, vars, stream_dtype):
vars = [v.strip() for v in vars.split(",")]
knl = lp.assume(knl, "n>0")
knl = lp.split_iname(knl, "i", 2**18, outer_tag="g.0", slabs=(0, 1))
knl = lp.split_iname(knl, "i_inner", 8, inner_tag="l.0")
knl = lp.add_and_infer_dtypes(knl, {var: stream_dtype for var in vars})
knl = lp.set_argument_order(knl, vars + ["n"])
return knl
def test_precompute_some_exist(ctx_factory):
fortran_src = """
subroutine dgemm(m,n,ell,a,b,c)
implicit none
real*8 a(m,ell),b(ell,n),c(m,n)
integer m,n,k,i,j,ell
do j = 1,n
do i = 1,m
do k = 1,ell
c(i,j) = c(i,j) + b(k,j)*a(i,k)
end do
end do
end do
end subroutine
"""
knl, = lp.parse_fortran(fortran_src)
assert len(knl.domains) == 1
knl = lp.split_iname(knl, "i", 8,
outer_tag="g.0", inner_tag="l.1")
knl = lp.split_iname(knl, "j", 8,
outer_tag="g.1", inner_tag="l.0")
knl = lp.split_iname(knl, "k", 8)
knl = lp.assume(knl, "n mod 8 = 0")
knl = lp.assume(knl, "m mod 8 = 0")
knl = lp.assume(knl, "ell mod 8 = 0")
knl = lp.extract_subst(knl, "a_acc", "a[i1,i2]", parameters="i1, i2")
knl = lp.extract_subst(knl, "b_acc", "b[i1,i2]", parameters="i1, i2")
knl = lp.precompute(knl, "a_acc", "k_inner,i_inner",
precompute_inames="ktemp,itemp",
default_tag="l.auto")
knl = lp.precompute(knl, "b_acc", "j_inner,k_inner",
precompute_inames="itemp,k2temp",
default_tag="l.auto")
ref_knl = knl
ctx = ctx_factory()
lp.auto_test_vs_ref(ref_knl, ctx, knl, parameters=dict(n=128, m=128, ell=128))
def test_assume(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i]: 0<=i<n}",
"a[i] = a[i] + 1",
[lp.GlobalArg("a", np.float32, shape="n"), "..."])
knl = lp.split_iname(knl, "i", 16)
knl = lp.set_loop_priority(knl, "i_outer,i_inner")
knl = lp.assume(knl, "n mod 16 = 0")
knl = lp.assume(knl, "n > 10")
knl = lp.preprocess_kernel(knl, ctx.devices[0])
kernel_gen = lp.generate_loop_schedules(knl)
for gen_knl in kernel_gen:
print(gen_knl)
compiled = lp.CompiledKernel(ctx, gen_knl)
print(compiled.get_code())
assert "if" not in compiled.get_code()
def vanilla():
k = lp.make_kernel(
"{ [i] : k <= i < n}", """
a[i] = a[i] + 1
""", [
lp.ValueArg("k", dtype="int32"),
lp.ValueArg("n", dtype="int32"),
lp.GlobalArg("a", shape=(None, ), dtype="int32")
])
k = lp.assume(k, "k >= 0 and n >= k")
return k
def test_fd_1d(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel("{[i]: 0<=i<n}", "result[i] = u[i+1]-u[i]")
knl = lp.add_and_infer_dtypes(knl, {"u": np.float32})
ref_knl = knl
knl = lp.split_iname(knl, "i", 16)
knl = lp.extract_subst(knl, "u_acc", "u[j]", parameters="j")
knl = lp.precompute(knl, "u_acc", "i_inner", default_tag="for")
knl = lp.assume(knl, "n mod 16 = 0")
lp.auto_test_vs_ref(ref_knl, ctx, knl, parameters=dict(n=2048))
def test_rename_argument_with_assumptions():
import islpy as isl
knl = lp.make_kernel("{[i]: 0<=i<n_old}", """
y[i] = 2.0f
""")
knl = lp.assume(knl, "n_old=10")
knl = lp.rename_argument(knl, "n_old", "n_new")
assumptions = knl["loopy_kernel"].assumptions
assert "n_old" not in assumptions.get_var_dict()
assert "n_new" in assumptions.get_var_dict()
assert ((assumptions
& isl.BasicSet("[n_new]->{: n_new=10}")) == assumptions)
def transform(knl, vars, stream_dtype):
vars = [v.strip() for v in vars.split(",")]
knl = lp.assume(knl, "n>0")
knl = lp.split_iname(
knl, "i", 2**18, outer_tag="g.0", slabs=(0, 1))
knl = lp.split_iname(knl, "i_inner", 8, inner_tag="l.0")
knl = lp.add_and_infer_dtypes(knl, {
var: stream_dtype
for var in vars
})
knl = lp.set_argument_order(knl, vars + ["n"])
return knl
def test_numba_cuda_target():
knl = lp.make_kernel("{[i,j,k]: 0<=i,j<M and 0<=k<N}",
"D[i,j] = sqrt(sum(k, (X[i, k]-X[j, k])**2))",
target=lp.NumbaCudaTarget())
knl = lp.assume(knl, "M>0")
knl = lp.split_iname(knl, "i", 16, outer_tag='g.0')
knl = lp.split_iname(knl, "j", 128, inner_tag='l.0', slabs=(0, 1))
knl = lp.add_prefetch(knl, "X[i,:]")
knl = lp.fix_parameters(knl, N=3)
knl = lp.prioritize_loops(knl, "i_inner,j_outer")
knl = lp.tag_inames(knl, "k:unr")
knl = lp.tag_array_axes(knl, "X", "N0,N1")
knl = lp.add_and_infer_dtypes(knl, {"X": np.float32})
print(lp.generate_code_v2(knl).all_code())
def test_numba_cuda_target():
knl = lp.make_kernel(
"{[i,j,k]: 0<=i,j<M and 0<=k<N}",
"D[i,j] = sqrt(sum(k, (X[i, k]-X[j, k])**2))",
target=lp.NumbaCudaTarget())
knl = lp.assume(knl, "M>0")
knl = lp.split_iname(knl, "i", 16, outer_tag='g.0')
knl = lp.split_iname(knl, "j", 128, inner_tag='l.0', slabs=(0, 1))
knl = lp.add_prefetch(knl, "X[i,:]", default_tag="l.auto")
knl = lp.fix_parameters(knl, N=3)
knl = lp.prioritize_loops(knl, "i_inner,j_outer")
knl = lp.tag_inames(knl, "k:unr")
knl = lp.tag_array_axes(knl, "X", "N0,N1")
knl = lp.add_and_infer_dtypes(knl, {"X": np.float32})
print(lp.generate_code_v2(knl).all_code())
def test_fd_1d(ctx_factory):
ctx = ctx_factory()
knl = lp.make_kernel(
"{[i]: 0<=i<n}",
"result[i] = u[i+1]-u[i]")
knl = lp.add_and_infer_dtypes(knl, {"u": np.float32})
ref_knl = knl
knl = lp.split_iname(knl, "i", 16)
knl = lp.extract_subst(knl, "u_acc", "u[j]", parameters="j")
knl = lp.precompute(knl, "u_acc", "i_inner", default_tag="for")
knl = lp.assume(knl, "n mod 16 = 0")
lp.auto_test_vs_ref(
ref_knl, ctx, knl,
parameters=dict(n=2048))
def test_integer_associativity():
knl = lp.make_kernel(
"{[i] : 0<=i<arraylen}", """
e := (i // (ncomp * elemsize))
d := ((i // elemsize) % ncomp)
s := (i % elemsize)
v[i] = u[ncomp * indices[(s) + elemsize*(e)] + (d)]
""")
knl = lp.add_and_infer_dtypes(knl, {
"u": np.float64,
"elemsize, ncomp, indices": np.int32
})
import islpy as isl
knl = lp.assume(
knl,
isl.BasicSet("[elemsize, ncomp] -> "
"{ : elemsize>= 0 and ncomp >= 0}"))
print(lp.generate_code_v2(knl).device_code())
assert ("u[ncomp * indices[i % elemsize + elemsize "
"* loopy_floor_div_int32(i, ncomp * elemsize)] "
"+ loopy_mod_pos_b_int32(i / elemsize, ncomp)]"
in lp.generate_code_v2(knl).device_code())
def test_fuzz_expression_code_gen(ctx_factory, expr_type, random_seed):
from pymbolic import evaluate
def get_numpy_type(x):
if expr_type in ["real", "complex"]:
if isinstance(x, (complex, np.complexfloating)):
return np.complex128
else:
return np.float64
elif expr_type in ["int", "int_nonneg"]:
return np.int64
else:
raise ValueError("unknown expr_type: %s" % expr_type)
from random import seed
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
seed(random_seed)
data = []
instructions = []
ref_values = {}
if expr_type in ["real", "complex"]:
result_type = np.complex128
elif expr_type in ["int", "int_nonneg"]:
result_type = np.int64
else:
assert False
var_names = []
fuzz_iter = iter(generate_random_fuzz_examples(expr_type))
count = 0
while True:
if count == 10:
break
i, expr, var_values = next(fuzz_iter)
var_name = "expr%d" % i
print(expr)
#assert_parse_roundtrip(expr)
if expr_type in ["int", "int_nonneg"]:
result_type_iinfo = np.iinfo(np.int32)
bceval_mapper = BoundsCheckingEvaluationMapper(
var_values,
lbound=result_type_iinfo.min,
ubound=result_type_iinfo.max)
print(expr)
try:
ref_values[var_name] = bceval_mapper(expr)
except BoundsCheckError:
print(expr)
print("BOUNDS CHECK FAILED")
continue
else:
try:
ref_values[var_name] = evaluate(expr, var_values)
except ZeroDivisionError:
continue
count += 1
data.append(lp.GlobalArg(var_name, result_type, shape=()))
data.extend([
lp.TemporaryVariable(name, get_numpy_type(val))
for name, val in var_values.items()
])
instructions.extend([
lp.Assignment(name,
get_numpy_type(val)(val))
for name, val in var_values.items()
])
instructions.append(lp.Assignment(var_name, expr))
if expr_type == "int_nonneg":
var_names.extend(var_values)
knl = lp.make_kernel("{ : }", instructions, data, seq_dependencies=True)
import islpy as isl
knl = lp.assume(
knl,
isl.BasicSet(
"[%s] -> { : %s}" %
(", ".join(var_names), " and ".join("%s >= 0" % name
for name in var_names))))
knl = lp.set_options(knl, return_dict=True)
print(knl)
evt, lp_values = knl(queue, out_host=True)
for name, ref_value in ref_values.items():
lp_value = lp_values[name]
if expr_type in ["real", "complex"]:
err = abs(ref_value - lp_value) / abs(ref_value)
elif expr_type in ["int", "int_nonneg"]:
err = abs(ref_value - lp_value)
else:
assert False
if abs(err) > 1e-10:
print(80 * "-")
print(knl)
print(80 * "-")
print(lp.generate_code_v2(knl).device_code())
print(80 * "-")
print(f"WRONG: {name} rel error={err:g}")
print("reference=%r" % ref_value)
print("loopy=%r" % lp_value)
print(80 * "-")
1 / 0
print(lp.generate_code_v2(knl).device_code())
def set_up_volume_loop(kernel, Nq): # noqa
kernel = lp.fix_parameters(kernel, Nq=Nq)
kernel = lp.prioritize_loops(kernel, "e,k,j,i")
kernel = lp.tag_inames(kernel, dict(e="g.0", j="l.1", i="l.0"))
kernel = lp.assume(kernel, "elements >= 1")
return kernel
def test_gnuma_horiz_kernel(ctx_factory, ilp_multiple, Nq, opt_level):
ctx = ctx_factory()
filename = "strongVolumeKernels.f90"
with open(filename, "r") as sourcef:
source = sourcef.read()
source = source.replace("datafloat", "real*4")
hsv_r, hsv_s = [
knl
for knl in lp.parse_fortran(source, filename, auto_dependencies=False)
if "KernelR" in knl.name or "KernelS" in knl.name
]
hsv_r = lp.tag_instructions(hsv_r, "rknl")
hsv_s = lp.tag_instructions(hsv_s, "sknl")
hsv = lp.fuse_kernels([hsv_r, hsv_s], ["_r", "_s"])
#hsv = hsv_s
from gnuma_loopy_transforms import (fix_euler_parameters,
set_q_storage_format,
set_D_storage_format)
hsv = lp.fix_parameters(hsv, Nq=Nq)
hsv = lp.set_loop_priority(hsv, "e,k,j,i")
hsv = lp.tag_inames(hsv, dict(e="g.0", j="l.1", i="l.0"))
hsv = lp.assume(hsv, "elements >= 1")
hsv = fix_euler_parameters(hsv, p_p0=1, p_Gamma=1.4, p_R=1)
for name in ["Q", "rhsQ"]:
hsv = set_q_storage_format(hsv, name)
hsv = set_D_storage_format(hsv)
#hsv = lp.add_prefetch(hsv, "volumeGeometricFactors")
ref_hsv = hsv
if opt_level == 0:
tap_hsv = hsv
hsv = lp.add_prefetch(hsv, "D[:,:]")
if opt_level == 1:
tap_hsv = hsv
# turn the first reads into subst rules
local_prep_var_names = set()
for insn in lp.find_instructions(hsv, "tag:local_prep"):
assignee, = insn.assignee_var_names()
local_prep_var_names.add(assignee)
hsv = lp.assignment_to_subst(hsv, assignee)
# precompute fluxes
hsv = lp.assignment_to_subst(hsv, "JinvD_r")
hsv = lp.assignment_to_subst(hsv, "JinvD_s")
r_fluxes = lp.find_instructions(hsv, "tag:compute_fluxes and tag:rknl")
s_fluxes = lp.find_instructions(hsv, "tag:compute_fluxes and tag:sknl")
if ilp_multiple > 1:
hsv = lp.split_iname(hsv, "k", 2, inner_tag="ilp")
ilp_inames = ("k_inner", )
flux_ilp_inames = ("kk", )
else:
ilp_inames = ()
flux_ilp_inames = ()
rtmps = []
stmps = []
flux_store_idx = 0
for rflux_insn, sflux_insn in zip(r_fluxes, s_fluxes):
for knl_tag, insn, flux_inames, tmps, flux_precomp_inames in [
("rknl", rflux_insn, (
"j",
"n",
), rtmps, (
"jj",
"ii",
)),
("sknl", sflux_insn, (
"i",
"n",
), stmps, (
"ii",
"jj",
)),
]:
flux_var, = insn.assignee_var_names()
print(insn)
reader, = lp.find_instructions(
hsv,
"tag:{knl_tag} and reads:{flux_var}".format(knl_tag=knl_tag,
flux_var=flux_var))
hsv = lp.assignment_to_subst(hsv, flux_var)
flux_store_name = "flux_store_%d" % flux_store_idx
flux_store_idx += 1
tmps.append(flux_store_name)
hsv = lp.precompute(hsv,
flux_var + "_subst",
flux_inames + ilp_inames,
temporary_name=flux_store_name,
precompute_inames=flux_precomp_inames +
flux_ilp_inames,
default_tag=None)
if flux_var.endswith("_s"):
hsv = lp.tag_array_axes(hsv, flux_store_name, "N0,N1,N2?")
else:
hsv = lp.tag_array_axes(hsv, flux_store_name, "N1,N0,N2?")
n_iname = "n_" + flux_var.replace("_r", "").replace("_s", "")
if n_iname.endswith("_0"):
n_iname = n_iname[:-2]
hsv = lp.rename_iname(hsv,
"n",
n_iname,
within="id:" + reader.id,
existing_ok=True)
hsv = lp.tag_inames(hsv, dict(ii="l.0", jj="l.1"))
for iname in flux_ilp_inames:
hsv = lp.tag_inames(hsv, {iname: "ilp"})
hsv = lp.alias_temporaries(hsv, rtmps)
hsv = lp.alias_temporaries(hsv, stmps)
if opt_level == 2:
tap_hsv = hsv
for prep_var_name in local_prep_var_names:
if prep_var_name.startswith("Jinv") or "_s" in prep_var_name:
continue
hsv = lp.precompute(
hsv, lp.find_one_rule_matching(hsv, prep_var_name + "_*subst*"))
if opt_level == 3:
tap_hsv = hsv
hsv = lp.add_prefetch(hsv, "Q[ii,jj,k,:,:,e]", sweep_inames=ilp_inames)
if opt_level == 4:
tap_hsv = hsv
tap_hsv = lp.tag_inames(
tap_hsv, dict(Q_dim_field_inner="unr", Q_dim_field_outer="unr"))
hsv = lp.buffer_array(hsv,
"rhsQ",
ilp_inames,
fetch_bounding_box=True,
default_tag="for",
init_expression="0",
store_expression="base + buffer")
if opt_level == 5:
tap_hsv = hsv
tap_hsv = lp.tag_inames(
tap_hsv,
dict(rhsQ_init_field_inner="unr",
rhsQ_store_field_inner="unr",
rhsQ_init_field_outer="unr",
rhsQ_store_field_outer="unr",
Q_dim_field_inner="unr",
Q_dim_field_outer="unr"))
# buffer axes need to be vectorized in order for this to work
hsv = lp.tag_array_axes(hsv, "rhsQ_buf", "c?,vec,c")
hsv = lp.tag_array_axes(hsv, "Q_fetch", "c?,vec,c")
hsv = lp.tag_array_axes(hsv, "D_fetch", "f,f")
hsv = lp.tag_inames(hsv, {
"Q_dim_k": "unr",
"rhsQ_init_k": "unr",
"rhsQ_store_k": "unr"
},
ignore_nonexistent=True)
if opt_level == 6:
tap_hsv = hsv
tap_hsv = lp.tag_inames(
tap_hsv,
dict(rhsQ_init_field_inner="unr",
rhsQ_store_field_inner="unr",
rhsQ_init_field_outer="unr",
rhsQ_store_field_outer="unr",
Q_dim_field_inner="unr",
Q_dim_field_outer="unr"))
hsv = lp.tag_inames(
hsv,
dict(rhsQ_init_field_inner="vec",
rhsQ_store_field_inner="vec",
rhsQ_init_field_outer="unr",
rhsQ_store_field_outer="unr",
Q_dim_field_inner="vec",
Q_dim_field_outer="unr"))
if opt_level == 7:
tap_hsv = hsv
hsv = lp.collect_common_factors_on_increment(
hsv, "rhsQ_buf", vary_by_axes=(0, ) if ilp_multiple > 1 else ())
if opt_level >= 8:
tap_hsv = hsv
hsv = tap_hsv
if 1:
print("OPS")
op_poly = lp.get_op_poly(hsv)
print(lp.stringify_stats_mapping(op_poly))
print("MEM")
gmem_poly = lp.sum_mem_access_to_bytes(lp.get_gmem_access_poly(hsv))
print(lp.stringify_stats_mapping(gmem_poly))
hsv = lp.set_options(hsv,
cl_build_options=[
"-cl-denorms-are-zero",
"-cl-fast-relaxed-math",
"-cl-finite-math-only",
"-cl-mad-enable",
"-cl-no-signed-zeros",
])
hsv = hsv.copy(name="horizontalStrongVolumeKernel")
results = lp.auto_test_vs_ref(ref_hsv,
ctx,
hsv,
parameters=dict(elements=300),
quiet=True)
elapsed = results["elapsed_wall"]
print("elapsed", elapsed)
def set_up_volume_loop(kernel, Nq): # noqa
kernel = lp.fix_parameters(kernel, Nq=Nq)
kernel = lp.prioritize_loops(kernel, "e,k,j,i")
kernel = lp.tag_inames(kernel, dict(e="g.0", j="l.1", i="l.0"))
kernel = lp.assume(kernel, "elements >= 1")
return kernel
def test_gnuma_horiz_kernel(ctx_factory, ilp_multiple, Nq, opt_level):
ctx = ctx_factory()
filename = "strongVolumeKernels.f90"
with open(filename, "r") as sourcef:
source = sourcef.read()
source = source.replace("datafloat", "real*4")
hsv_r, hsv_s = [
knl for knl in lp.parse_fortran(source, filename, auto_dependencies=False)
if "KernelR" in knl.name or "KernelS" in knl.name
]
hsv_r = lp.tag_instructions(hsv_r, "rknl")
hsv_s = lp.tag_instructions(hsv_s, "sknl")
hsv = lp.fuse_kernels([hsv_r, hsv_s], ["_r", "_s"])
#hsv = hsv_s
from gnuma_loopy_transforms import (
fix_euler_parameters,
set_q_storage_format, set_D_storage_format)
hsv = lp.fix_parameters(hsv, Nq=Nq)
hsv = lp.set_loop_priority(hsv, "e,k,j,i")
hsv = lp.tag_inames(hsv, dict(e="g.0", j="l.1", i="l.0"))
hsv = lp.assume(hsv, "elements >= 1")
hsv = fix_euler_parameters(hsv, p_p0=1, p_Gamma=1.4, p_R=1)
for name in ["Q", "rhsQ"]:
hsv = set_q_storage_format(hsv, name)
hsv = set_D_storage_format(hsv)
#hsv = lp.add_prefetch(hsv, "volumeGeometricFactors")
ref_hsv = hsv
if opt_level == 0:
tap_hsv = hsv
hsv = lp.add_prefetch(hsv, "D[:,:]")
if opt_level == 1:
tap_hsv = hsv
# turn the first reads into subst rules
local_prep_var_names = set()
for insn in lp.find_instructions(hsv, "tag:local_prep"):
assignee, = insn.assignee_var_names()
local_prep_var_names.add(assignee)
hsv = lp.assignment_to_subst(hsv, assignee)
# precompute fluxes
hsv = lp.assignment_to_subst(hsv, "JinvD_r")
hsv = lp.assignment_to_subst(hsv, "JinvD_s")
r_fluxes = lp.find_instructions(hsv, "tag:compute_fluxes and tag:rknl")
s_fluxes = lp.find_instructions(hsv, "tag:compute_fluxes and tag:sknl")
if ilp_multiple > 1:
hsv = lp.split_iname(hsv, "k", 2, inner_tag="ilp")
ilp_inames = ("k_inner",)
flux_ilp_inames = ("kk",)
else:
ilp_inames = ()
flux_ilp_inames = ()
rtmps = []
stmps = []
flux_store_idx = 0
for rflux_insn, sflux_insn in zip(r_fluxes, s_fluxes):
for knl_tag, insn, flux_inames, tmps, flux_precomp_inames in [
("rknl", rflux_insn, ("j", "n",), rtmps, ("jj", "ii",)),
("sknl", sflux_insn, ("i", "n",), stmps, ("ii", "jj",)),
]:
flux_var, = insn.assignee_var_names()
print(insn)
reader, = lp.find_instructions(hsv,
"tag:{knl_tag} and reads:{flux_var}"
.format(knl_tag=knl_tag, flux_var=flux_var))
hsv = lp.assignment_to_subst(hsv, flux_var)
flux_store_name = "flux_store_%d" % flux_store_idx
flux_store_idx += 1
tmps.append(flux_store_name)
hsv = lp.precompute(hsv, flux_var+"_subst", flux_inames + ilp_inames,
temporary_name=flux_store_name,
precompute_inames=flux_precomp_inames + flux_ilp_inames,
default_tag=None)
if flux_var.endswith("_s"):
hsv = lp.tag_data_axes(hsv, flux_store_name, "N0,N1,N2?")
else:
hsv = lp.tag_data_axes(hsv, flux_store_name, "N1,N0,N2?")
n_iname = "n_"+flux_var.replace("_r", "").replace("_s", "")
if n_iname.endswith("_0"):
n_iname = n_iname[:-2]
hsv = lp.rename_iname(hsv, "n", n_iname, within="id:"+reader.id,
existing_ok=True)
hsv = lp.tag_inames(hsv, dict(ii="l.0", jj="l.1"))
for iname in flux_ilp_inames:
hsv = lp.tag_inames(hsv, {iname: "ilp"})
hsv = lp.alias_temporaries(hsv, rtmps)
hsv = lp.alias_temporaries(hsv, stmps)
if opt_level == 2:
tap_hsv = hsv
for prep_var_name in local_prep_var_names:
if prep_var_name.startswith("Jinv") or "_s" in prep_var_name:
continue
hsv = lp.precompute(hsv,
lp.find_one_rule_matching(hsv, prep_var_name+"_*subst*"))
if opt_level == 3:
tap_hsv = hsv
hsv = lp.add_prefetch(hsv, "Q[ii,jj,k,:,:,e]", sweep_inames=ilp_inames)
if opt_level == 4:
tap_hsv = hsv
tap_hsv = lp.tag_inames(tap_hsv, dict(
Q_dim_field_inner="unr",
Q_dim_field_outer="unr"))
hsv = lp.buffer_array(hsv, "rhsQ", ilp_inames,
fetch_bounding_box=True, default_tag="for",
init_expression="0", store_expression="base + buffer")
if opt_level == 5:
tap_hsv = hsv
tap_hsv = lp.tag_inames(tap_hsv, dict(
rhsQ_init_field_inner="unr", rhsQ_store_field_inner="unr",
rhsQ_init_field_outer="unr", rhsQ_store_field_outer="unr",
Q_dim_field_inner="unr",
Q_dim_field_outer="unr"))
# buffer axes need to be vectorized in order for this to work
hsv = lp.tag_data_axes(hsv, "rhsQ_buf", "c?,vec,c")
hsv = lp.tag_data_axes(hsv, "Q_fetch", "c?,vec,c")
hsv = lp.tag_data_axes(hsv, "D_fetch", "f,f")
hsv = lp.tag_inames(hsv,
{"Q_dim_k": "unr", "rhsQ_init_k": "unr", "rhsQ_store_k": "unr"},
ignore_nonexistent=True)
if opt_level == 6:
tap_hsv = hsv
tap_hsv = lp.tag_inames(tap_hsv, dict(
rhsQ_init_field_inner="unr", rhsQ_store_field_inner="unr",
rhsQ_init_field_outer="unr", rhsQ_store_field_outer="unr",
Q_dim_field_inner="unr",
Q_dim_field_outer="unr"))
hsv = lp.tag_inames(hsv, dict(
rhsQ_init_field_inner="vec", rhsQ_store_field_inner="vec",
rhsQ_init_field_outer="unr", rhsQ_store_field_outer="unr",
Q_dim_field_inner="vec",
Q_dim_field_outer="unr"))
if opt_level == 7:
tap_hsv = hsv
hsv = lp.collect_common_factors_on_increment(hsv, "rhsQ_buf",
vary_by_axes=(0,) if ilp_multiple > 1 else ())
if opt_level >= 8:
tap_hsv = hsv
hsv = tap_hsv
if 1:
print("OPS")
op_poly = lp.get_op_poly(hsv)
print(lp.stringify_stats_mapping(op_poly))
print("MEM")
gmem_poly = lp.sum_mem_access_to_bytes(lp.get_gmem_access_poly(hsv))
print(lp.stringify_stats_mapping(gmem_poly))
hsv = lp.set_options(hsv, cl_build_options=[
"-cl-denorms-are-zero",
"-cl-fast-relaxed-math",
"-cl-finite-math-only",
"-cl-mad-enable",
"-cl-no-signed-zeros",
])
hsv = hsv.copy(name="horizontalStrongVolumeKernel")
results = lp.auto_test_vs_ref(ref_hsv, ctx, hsv, parameters=dict(elements=300),
quiet=True)
elapsed = results["elapsed_wall"]
print("elapsed", elapsed)
