def gen_rhs(self, comp: Expr, shift: List[int], offset: List[int], rel=None, dimrels=None):
"""
:param comp: The expression to print
:param shift: The shift of scatter
:param offset: Scattered from
:param rel: Added offset when using loops
:return:
"""
from st.expr import ConstRef
self.printer.shift = shift[:]
self.printer.offset = offset[:]
self.printer.rel = ConstRef("vit")
self.printer.dimrels = dimrels
if rel:
self.printer.rel += rel
return self.printer.print_str(comp)
def prologue(self, toplevel: CodeBlock):
from st.expr import ConstRef
self.TILE_NAME = [ConstRef(i) for i in self.TILE_NAME]
self.backend.codegen.LAYOUTREL = True
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
input = Grid("in", 3)
output = Grid("out", 3)
c = []
for cc in range(125):
c.append(ConstRef('c[{}]'.format(cc)))
calc = c[0] * input(i, j, k) + c[1] * input(i - 3, j, k) + c[2] * input(
i - 2, j, k
) + c[3] * input(i - 1, j, k) + c[4] * input(i + 1, j, k) + c[5] * input(
i + 2, j, k) + c[6] * input(i + 3, j, k) + c[7] * input(
i, j - 3, k) + c[8] * input(i, j - 2, k) + c[9] * input(
i, j - 1, k) + c[10] * input(i, j + 1, k) + c[11] * input(
i, j + 2, k) + c[12] * input(i, j + 3, k) + c[13] * input(
i, j, k - 3) + c[14] * input(i, j, k - 2) + c[15] * input(
i, j, k - 1) + c[16] * input(
i, j, k + 1) + c[17] * input(
i, j, k + 2) + c[18] * input(i, j, k + 3)
output(i, j, k).assign(calc)
STENCIL = [output]
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
input = Grid("in", 3)
output = Grid("out", 3)
a0 = ConstRef("MPI_A0")
a1 = ConstRef("MPI_A1")
a2 = ConstRef("MPI_A2")
a3 = ConstRef("MPI_A3")
a4 = ConstRef("MPI_A4")
# Express computation
# output[i, j, k] is assumed
calc = a0 * input(i, j, k) + \
a1 * input(i + 1, j, k) + \
a1 * input(i - 1, j, k) + \
a1 * input(i, j + 1, k) + \
a1 * input(i, j - 1, k) + \
a1 * input(i, j, k + 1) + \
a1 * input(i, j, k - 1) + \
a2 * input(i + 2, j, k) + \
a2 * input(i - 2, j, k) + \
a2 * input(i, j + 2, k) + \
a2 * input(i, j - 2, k) + \
a2 * input(i, j, k + 2) + \
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
input = Grid("in", 3)
output = Grid("out", 3)
alpha = ConstRef("MPI_ALPHA")
beta = ConstRef("MPI_BETA")
# Express computation
# output[i, j, k] is assumed
calc = alpha * input(i, j, k) + \
beta * input(i + 1, j, k) + \
beta * input(i - 1, j, k) + \
beta * input(i, j + 1, k) + \
beta * input(i, j - 1, k) + \
beta * input(i, j, k + 1) + \
beta * input(i, j, k - 1)
output(i, j, k).assign(calc)
STENCIL = [output]
from st.expr import Index, ConstRef, If
from st.grid import Grid
from st.func import Func
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
maxfunc = Func("max", 2)
# Declare grid
input = Grid("bIn", 3)
output = Grid("bOut", 3)
param = [
ConstRef("coeff[0]"),
ConstRef("coeff[1]"),
ConstRef("coeff[2]"),
ConstRef("coeff[3]"),
ConstRef("coeff[4]"),
ConstRef("coeff[5]"),
ConstRef("coeff[6]")
]
zero = ConstRef("0.0")
# Express computation
# output[i, j, k] is assumed
calc = param[0] * maxfunc(input(i, j, k), zero) + \
param[1] * maxfunc(input(i + 1, j, k), zero) + \
param[2] * maxfunc(input(i - 1, j, k), zero) + \
param[3] * maxfunc(input(i, j + 1, k), zero) + \
qv = Grid("qv", 3)
qw = Grid("qw", 3)
qpres = Grid("qpres", 3)
flux_irho = Grid('flux_irho', 3)
flux_imx = Grid("flux_imx", 3)
flux_imy = Grid("flux_imy", 3)
flux_imz = Grid("flux_imz", 3)
flux_iene = Grid("flux_iene", 3)
cons_imx = Grid("cons_imx", 3)
cons_imy = Grid("cons_imy", 3)
cons_imz = Grid("cons_imz", 3)
cons_iene = Grid("cons_iene", 3)
dxinv0 = ConstRef('dxinv0')
dxinv1 = ConstRef('dxinv1')
dxinv2 = ConstRef('dxinv2')
ALP = ConstRef('ALP')
BET = ConstRef('BET')
GAM = ConstRef('GAM')
DEL = ConstRef('DEL')
qum4 = qu(i - 4, j, k)
qum3 = qu(i - 3, j, k)
qum2 = qu(i - 2, j, k)
qum1 = qu(i - 1, j, k)
qup1 = qu(i + 1, j, k)
qup2 = qu(i + 2, j, k)
qup3 = qu(i + 3, j, k)
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
input = Grid("in", 3)
output = Grid("out", 3)
# Symmetries of 125pt stencil imply 10 constants, permuted +/-
# 0 0 0 - 1 of these
a0 = ConstRef("MPI_C0")
# 0 0 1 - 6 of these
a1 = ConstRef("MPI_C1")
# 0 0 2 - 6 of these
a2 = ConstRef("MPI_C2")
# 0 1 1 - 12 of these
a3 = ConstRef("MPI_C3")
# 0 1 2 - 24 of these
a4 = ConstRef("MPI_C4")
# 0 2 2 - 12 of these
a5 = ConstRef("MPI_C5")
# 1 1 1 - 8 of these
a6 = ConstRef("MPI_C6")
# 1 1 2 - 24 of these
a7 = ConstRef("MPI_C7")
# 1 2 2 - 24 of these
a8 = ConstRef("MPI_C8")
# 2 2 2 - 8 of these
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
Ac = Grid("Ac", 3)
Ap = Grid("Ap", 3)
Dinv = Grid("Dinv", 3)
RHS = Grid("RHS", 3)
output = Grid("out", 3)
c1 = ConstRef('c1')
c2 = ConstRef('c2')
h2inv = ConstRef('h2inv')
calc = Ac(i, j, k) + c1 * (Ac(i, j, k) + Ap(i, j, k)) + c2 * Dinv(i, j, k) * (
RHS(i, j, k) +
(Ac(i, j, k) + h2inv *
(0.03 * (Ac(i - 1, j - 1, k - 1) + Ac(i + 1, j - 1, k - 1) +
Ac(i - 1, j + 1, k - 1) + Ac(i + 1, j + 1, k - 1) +
Ac(i - 1, j - 1, k + 1) + Ac(i + 1, j - 1, k + 1) +
Ac(i - 1, j + 1, k + 1) + Ac(i + 1, j + 1, k + 1)) + 0.1 *
(Ac(i, j - 1, k - 1) + Ac(i - 1, j, k - 1) + Ac(i + 1, j, k - 1) +
Ac(i, j + 1, k - 1) + Ac(i - 1, j - 1, k) + Ac(i + 1, j - 1, k) +
Ac(i - 1, j + 1, k) + Ac(i + 1, j + 1, k) + Ac(i, j - 1, k + 1) +
Ac(i - 1, j, k + 1) + Ac(i + 1, j, k + 1) + Ac(i, j + 1, k + 1)) + 0.46 *
(Ac(i, j, k - 1) + Ac(i, j - 1, k) + Ac(i - 1, j, k) + Ac(i + 1, j, k) +
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
input = Grid("bIn", 3)
output = Grid("bOut", 3)
param = [
ConstRef("coeff[0]"),
ConstRef("coeff[1]"),
ConstRef("coeff[2]"),
ConstRef("coeff[3]"),
ConstRef("coeff[4]"),
ConstRef("coeff[5]"),
ConstRef("coeff[6]")
]
# Express computation
# output[i, j, k] is assumed
calc = param[0] * input(i, j, k) + \
param[1] * input(i + 1, j, k) + \
param[2] * input(i - 1, j, k) + \
param[3] * input(i, j + 1, k) + \
param[4] * input(i, j - 1, k) + \
param[5] * input(i, j, k + 1) + \
param[6] * input(i, j, k - 1)
output(i, j, k).assign(calc)
j = Index(1)
k = Index(2)
# Declare grid
ux = Grid("ux", 3)
vx = Grid("vx", 3)
wx = Grid("wx", 3)
uy = Grid("uy", 3)
vy = Grid("vy", 3)
wy = Grid("wy", 3)
uz = Grid("uz", 3)
vz = Grid("vz", 3)
wz = Grid("wz", 3)
q_1 = Grid("q_1", 3)
q_2 = Grid("q_2", 3)
q_3 = Grid("q_3", 3)
q_5 = Grid("q_5", 3)
diffflux1 = Grid("diffflux1", 3)
diffflux2 = Grid("diffflux2", 3)
diffflux3 = Grid("diffflux3", 3)
diffflux4 = Grid("diffflux4", 3)
dxinv0 = ConstRef("dxinv0")
dxinv1 = ConstRef("dxinv1")
dxinv2 = ConstRef("dxinv2")
from st.expr import Index, ConstRef
from st.grid import Grid
# Declare indices
i = Index(0)
j = Index(1)
k = Index(2)
# Declare grid
input = Grid("in", 3)
output = Grid("out", 3)
a0 = ConstRef("MPI_B0")
a1 = ConstRef("MPI_B1")
a2 = ConstRef("MPI_B2")
# Express computation
# output[i, j, k] is assumed
calc = a0 * input(i, j, k) + \
a1 * input(i + 1, j, k) + \
a1 * input(i - 1, j, k) + \
a1 * input(i, j + 1, k) + \
a1 * input(i, j - 1, k) + \
a1 * input(i, j, k + 1) + \
a1 * input(i, j, k - 1) + \
a2 * input(i + 2, j, k) + \
a2 * input(i - 2, j, k) + \
a2 * input(i, j + 2, k) + \
a2 * input(i, j - 2, k) + \
a2 * input(i, j, k + 2) + \
a2 * input(i, j, k - 2)
output(i, j, k).assign(calc)