def test_small_batched_matvec(ctx_factory):
dtype = np.float32
ctx = ctx_factory()
order = "C"
K = 9997 # noqa
Np = 36 # noqa
knl = lp.make_kernel(
"{[i,j,k]: 0<=k<K and 0<= i,j < %d}" % Np,
["result[k, i] = sum(j, d[i, j]*f[k, j])"], [
lp.GlobalArg("d", dtype, shape=(Np, Np), order=order),
lp.GlobalArg("f", dtype, shape=("K", Np), order=order),
lp.GlobalArg("result", dtype, shape=("K", Np), order=order),
lp.ValueArg("K", np.int32, approximately=1000),
],
name="batched_matvec",
assumptions="K>=1")
seq_knl = knl
align_bytes = 64
knl = lp.add_prefetch(knl, 'd[:,:]', default_tag="l.auto")
pad_mult = lp.find_padding_multiple(knl, "f", 0, align_bytes)
knl = lp.split_array_dim(knl, ("f", 0), pad_mult)
knl = lp.add_padding(knl, "f", 0, align_bytes)
lp.auto_test_vs_ref(seq_knl,
ctx,
knl,
op_count=[K * 2 * Np**2 / 1e9],
op_label=["GFlops"],
parameters=dict(K=K))
def test_small_batched_matvec(ctx_factory):
dtype = np.float32
ctx = ctx_factory()
order = "C"
K = 9997 # noqa
Np = 36 # noqa
knl = lp.make_kernel(
"{[i,j,k]: 0<=k<K and 0<= i,j < %d}" % Np,
[
"result[k, i] = sum(j, d[i, j]*f[k, j])"
],
[
lp.GlobalArg("d", dtype, shape=(Np, Np), order=order),
lp.GlobalArg("f", dtype, shape=("K", Np), order=order),
lp.GlobalArg("result", dtype, shape=("K", Np), order=order),
lp.ValueArg("K", np.int32, approximately=1000),
], name="batched_matvec", assumptions="K>=1")
seq_knl = knl
align_bytes = 64
knl = lp.add_prefetch(knl, 'd[:,:]')
pad_mult = lp.find_padding_multiple(knl, "f", 0, align_bytes)
knl = lp.split_array_dim(knl, ("f", 0), pad_mult)
knl = lp.add_padding(knl, "f", 0, align_bytes)
lp.auto_test_vs_ref(seq_knl, ctx, knl,
op_count=[K*2*Np**2/1e9], op_label=["GFlops"],
parameters=dict(K=K))
def variant_fancy_padding(knl):
knl = lp.tag_inames(knl, dict(n="l.0"))
pad_mult = lp.find_padding_multiple(knl, "u", 1, 32)
arg_names = [
prefix+name
for name in ["u", "v", "w", "p"]
for prefix in ["", "rhs"]]
knl = lp.split_array_dim(knl, [(nm, 0) for nm in arg_names], pad_mult)
return knl
def variant_fancy_padding(knl):
knl = lp.tag_inames(knl, dict(n="l.0"))
pad_mult = lp.find_padding_multiple(knl, "u", 1, 32)
arg_names = [
prefix+name
for name in ["u", "v", "w", "p"]
for prefix in ["", "rhs"]]
knl = lp.split_array_dim(knl, [(nm, 0) for nm in arg_names], pad_mult)
return knl