def to_codegen_result(codegen_state, insn_id, domain, check_inames,
required_preds, ast):
from loopy.codegen.bounds import get_bounds_checks
from loopy.symbolic import constraint_to_expr
bounds_checks = get_bounds_checks(domain,
check_inames,
codegen_state.implemented_domain,
overapproximate=False)
bounds_check_set = isl.Set.universe(domain.get_space()) \
.add_constraints(bounds_checks)
bounds_check_set, new_implemented_domain = isl.align_two(
bounds_check_set, codegen_state.implemented_domain)
new_implemented_domain = new_implemented_domain & bounds_check_set
if bounds_check_set.is_empty():
return None
condition_exprs = [constraint_to_expr(cns) for cns in bounds_checks]
condition_exprs.extend(required_preds -
codegen_state.implemented_predicates)
if condition_exprs:
from pymbolic.primitives import LogicalAnd
from pymbolic.mapper.stringifier import PREC_NONE
ast = codegen_state.ast_builder.emit_if(
codegen_state.expression_to_code_mapper(
LogicalAnd(tuple(condition_exprs)), PREC_NONE), ast)
return CodeGenerationResult.new(codegen_state, insn_id, ast,
new_implemented_domain)
def realize_conditional(self, node, context_cond=None):
scope = self.scope_stack[-1]
cond_name = intern("loopy_cond%d" % self.condition_id_counter)
self.condition_id_counter += 1
assert cond_name not in scope.type_map
scope.type_map[cond_name] = np.int32
from pymbolic import var
cond_var = var(cond_name)
self.add_expression_instruction(cond_var,
self.parse_expr(node, node.expr))
cond_expr = cond_var
if context_cond is not None:
from pymbolic.primitives import LogicalAnd
cond_expr = LogicalAnd((cond_var, context_cond))
self.conditions_data.append((context_cond, cond_var))
else:
self.conditions_data.append((None, cond_var))
self.conditions.append(cond_expr)
def to_codegen_result(codegen_state, insn_id, domain, check_inames,
required_preds, ast):
chk_domain = isl.Set.from_basic_set(domain)
chk_domain = chk_domain.remove_redundancies()
chk_domain = codegen_state.kernel.cache_manager.eliminate_except(
chk_domain, check_inames, (dim_type.set, ))
chk_domain, new_implemented_domain = _get_new_implemented_domain(
codegen_state.kernel, chk_domain, codegen_state.implemented_domain)
if chk_domain.is_empty():
return None
condition_exprs = []
if not chk_domain.plain_is_universe():
from loopy.symbolic import set_to_cond_expr
condition_exprs.append(set_to_cond_expr(chk_domain))
condition_exprs.extend(required_preds -
codegen_state.implemented_predicates)
if condition_exprs:
from pymbolic.primitives import LogicalAnd
from pymbolic.mapper.stringifier import PREC_NONE
ast = codegen_state.ast_builder.emit_if(
codegen_state.expression_to_code_mapper(
LogicalAnd(tuple(condition_exprs)), PREC_NONE), ast)
return CodeGenerationResult.new(codegen_state, insn_id, ast,
new_implemented_domain)
def parse_postfix(self, pstate, min_precedence, left_exp):
from pymbolic.parser import (
_PREC_CALL, _PREC_COMPARISON, _openpar,
_PREC_LOGICAL_OR, _PREC_LOGICAL_AND)
from pymbolic.primitives import (
Comparison, LogicalAnd, LogicalOr)
next_tag = pstate.next_tag()
if next_tag is _openpar and _PREC_CALL > min_precedence:
raise TranslationError("parenthesis operator only works on names")
elif next_tag in self.COMP_MAP and _PREC_COMPARISON > min_precedence:
pstate.advance()
left_exp = Comparison(
left_exp,
self.COMP_MAP[next_tag],
self.parse_expression(pstate, _PREC_COMPARISON))
did_something = True
elif next_tag is _and and _PREC_LOGICAL_AND > min_precedence:
pstate.advance()
left_exp = LogicalAnd((left_exp,
self.parse_expression(pstate, _PREC_LOGICAL_AND)))
did_something = True
elif next_tag is _or and _PREC_LOGICAL_OR > min_precedence:
pstate.advance()
left_exp = LogicalOr((left_exp,
self.parse_expression(pstate, _PREC_LOGICAL_OR)))
did_something = True
else:
left_exp, did_something = ExpressionParserBase.parse_postfix(
self, pstate, min_precedence, left_exp)
return left_exp, did_something
def flat_LogicalAnd(*children): # noqa
from pymbolic.primitives import LogicalAnd
result = []
for child in children:
if isinstance(child, LogicalAnd):
result.extend(child.children)
else:
result.append(child)
return LogicalAnd(tuple(result))
def wrap_in_if(codegen_state, condition_exprs, inner):
if condition_exprs:
from pymbolic.primitives import LogicalAnd
from pymbolic.mapper.stringifier import PREC_NONE
cur_ast = inner.current_ast(codegen_state)
return inner.with_new_ast(
codegen_state,
codegen_state.ast_builder.emit_if(
codegen_state.expression_to_code_mapper(
LogicalAnd(tuple(condition_exprs)), PREC_NONE), cur_ast))
return inner
def construct_else_condition(self):
context_cond, prev_cond = self.conditions_data.pop()
if prev_cond is None:
raise RuntimeError("else if may not follow else")
self.conditions.pop()
from pymbolic.primitives import LogicalNot, LogicalAnd
else_expr = LogicalNot(prev_cond)
if context_cond is not None:
else_expr = LogicalAnd((else_expr, context_cond))
return else_expr
def parse_postfix(self, pstate, min_precedence, left_exp):
from pymbolic.parser import (_PREC_CALL, _PREC_COMPARISON, _openpar,
_PREC_LOGICAL_OR, _PREC_LOGICAL_AND)
from pymbolic.primitives import (Comparison, LogicalAnd, LogicalOr)
next_tag = pstate.next_tag()
if next_tag is _openpar and _PREC_CALL > min_precedence:
raise TranslationError("parenthesis operator only works on names")
elif next_tag in self.COMP_MAP and _PREC_COMPARISON > min_precedence:
pstate.advance()
left_exp = Comparison(
left_exp, self.COMP_MAP[next_tag],
self.parse_expression(pstate, _PREC_COMPARISON))
did_something = True
elif next_tag is _and and _PREC_LOGICAL_AND > min_precedence:
pstate.advance()
left_exp = LogicalAnd(
(left_exp, self.parse_expression(pstate, _PREC_LOGICAL_AND)))
did_something = True
elif next_tag is _or and _PREC_LOGICAL_OR > min_precedence:
pstate.advance()
left_exp = LogicalOr(
(left_exp, self.parse_expression(pstate, _PREC_LOGICAL_OR)))
did_something = True
else:
left_exp, did_something = ExpressionParserBase.parse_postfix(
self, pstate, min_precedence, left_exp)
if isinstance(left_exp,
tuple) and min_precedence < self._PREC_FUNC_ARGS:
# this must be a complex literal
if len(left_exp) != 2:
raise TranslationError("complex literals must have "
"two entries")
r, i = left_exp
dtype = (r.dtype.type(0) + i.dtype.type(0))
if dtype == np.float32:
dtype = np.complex64
else:
dtype = np.complex128
left_exp = dtype(float(r) + float(i) * 1j)
return left_exp, did_something
def make_spectra_knl(self, is_real, rank_shape):
from pymbolic import var, parse
indices = i, j, k = parse("i, j, k")
momenta = [var("momenta_"+xx) for xx in ("x", "y", "z")]
ksq = sum((dk_i * mom[ii])**2
for mom, dk_i, ii in zip(momenta, self.dk, indices))
kmag = var("sqrt")(ksq)
bin_expr = var("round")(kmag / self.bin_width)
if is_real:
from pymbolic.primitives import If, Comparison, LogicalAnd
nyq = self.grid_shape[-1] / 2
condition = LogicalAnd((Comparison(momenta[2][k], ">", 0),
Comparison(momenta[2][k], "<", nyq)))
count = If(condition, 2, 1)
else:
count = 1
fk = var("fk")[i, j, k]
weight_expr = count * kmag**(var("k_power")) * var("abs")(fk)**2
histograms = {"spectrum": (bin_expr, weight_expr)}
args = [
lp.GlobalArg("fk", self.cdtype, shape=("Nx", "Ny", "Nz"),
offset=lp.auto),
lp.GlobalArg("momenta_x", self.rdtype, shape=("Nx",)),
lp.GlobalArg("momenta_y", self.rdtype, shape=("Ny",)),
lp.GlobalArg("momenta_z", self.rdtype, shape=("Nz",)),
lp.ValueArg("k_power", self.rdtype),
...
]
from pystella.histogram import Histogrammer
return Histogrammer(self.decomp, histograms, self.num_bins,
self.rdtype, args=args, rank_shape=rank_shape)
def parse_postfix(self, pstate, min_precedence, left_exp):
import pymbolic.primitives as primitives
did_something = False
next_tag = pstate.next_tag()
if next_tag is _openpar and _PREC_CALL > min_precedence:
pstate.advance()
args, kwargs = self.parse_arglist(pstate)
if kwargs:
left_exp = primitives.CallWithKwargs(left_exp, args, kwargs)
else:
left_exp = primitives.Call(left_exp, args)
did_something = True
elif next_tag is _openbracket and _PREC_CALL > min_precedence:
pstate.advance()
pstate.expect_not_end()
left_exp = primitives.Subscript(left_exp, self.parse_expression(pstate))
pstate.expect(_closebracket)
pstate.advance()
did_something = True
elif next_tag is _if and _PREC_IF > min_precedence:
from pymbolic.primitives import If
then_expr = left_exp
pstate.advance()
pstate.expect_not_end()
condition = self.parse_expression(pstate, _PREC_LOGICAL_OR)
pstate.expect(_else)
pstate.advance()
else_expr = self.parse_expression(pstate)
left_exp = If(condition, then_expr, else_expr)
did_something = True
elif next_tag is _dot and _PREC_CALL > min_precedence:
pstate.advance()
pstate.expect(_identifier)
left_exp = primitives.Lookup(left_exp, pstate.next_str())
pstate.advance()
did_something = True
elif next_tag is _plus and _PREC_PLUS > min_precedence:
pstate.advance()
right_exp = self.parse_expression(pstate, _PREC_PLUS)
if isinstance(left_exp, primitives.Sum):
left_exp = primitives.Sum(left_exp.children + (right_exp,))
else:
left_exp = primitives.Sum((left_exp, right_exp))
did_something = True
elif next_tag is _minus and _PREC_PLUS > min_precedence:
pstate.advance()
right_exp = self.parse_expression(pstate, _PREC_PLUS)
if isinstance(left_exp, primitives.Sum):
left_exp = primitives.Sum(left_exp.children + ((-right_exp),)) # noqa pylint:disable=invalid-unary-operand-type
else:
left_exp = primitives.Sum((left_exp, -right_exp)) # noqa pylint:disable=invalid-unary-operand-type
did_something = True
elif next_tag is _times and _PREC_TIMES > min_precedence:
pstate.advance()
right_exp = self.parse_expression(pstate, _PREC_PLUS)
if isinstance(left_exp, primitives.Product):
left_exp = primitives.Product(left_exp.children + (right_exp,))
else:
left_exp = primitives.Product((left_exp, right_exp))
did_something = True
elif next_tag is _floordiv and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp = primitives.FloorDiv(
left_exp, self.parse_expression(pstate, _PREC_TIMES))
did_something = True
elif next_tag is _over and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp = primitives.Quotient(
left_exp, self.parse_expression(pstate, _PREC_TIMES))
did_something = True
elif next_tag is _modulo and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp = primitives.Remainder(
left_exp, self.parse_expression(pstate, _PREC_TIMES))
did_something = True
elif next_tag is _power and _PREC_POWER > min_precedence:
pstate.advance()
left_exp = primitives.Power(
left_exp, self.parse_expression(pstate, _PREC_TIMES))
did_something = True
elif next_tag is _and and _PREC_LOGICAL_AND > min_precedence:
pstate.advance()
from pymbolic.primitives import LogicalAnd
left_exp = LogicalAnd((
left_exp,
self.parse_expression(pstate, _PREC_LOGICAL_AND)))
did_something = True
elif next_tag is _or and _PREC_LOGICAL_OR > min_precedence:
pstate.advance()
from pymbolic.primitives import LogicalOr
left_exp = LogicalOr((
left_exp,
self.parse_expression(pstate, _PREC_LOGICAL_OR)))
did_something = True
elif next_tag is _bitwiseor and _PREC_BITWISE_OR > min_precedence:
pstate.advance()
from pymbolic.primitives import BitwiseOr
left_exp = BitwiseOr((
left_exp,
self.parse_expression(pstate, _PREC_BITWISE_OR)))
did_something = True
elif next_tag is _bitwisexor and _PREC_BITWISE_XOR > min_precedence:
pstate.advance()
from pymbolic.primitives import BitwiseXor
left_exp = BitwiseXor((
left_exp,
self.parse_expression(pstate, _PREC_BITWISE_XOR)))
did_something = True
elif next_tag is _bitwiseand and _PREC_BITWISE_AND > min_precedence:
pstate.advance()
from pymbolic.primitives import BitwiseAnd
left_exp = BitwiseAnd((
left_exp,
self.parse_expression(pstate, _PREC_BITWISE_AND)))
did_something = True
elif next_tag is _rightshift and _PREC_SHIFT > min_precedence:
pstate.advance()
from pymbolic.primitives import RightShift
left_exp = RightShift(
left_exp,
self.parse_expression(pstate, _PREC_SHIFT))
did_something = True
elif next_tag is _leftshift and _PREC_SHIFT > min_precedence:
pstate.advance()
from pymbolic.primitives import LeftShift
left_exp = LeftShift(
left_exp,
self.parse_expression(pstate, _PREC_SHIFT))
did_something = True
elif next_tag in self._COMP_TABLE and _PREC_COMPARISON > min_precedence:
pstate.advance()
from pymbolic.primitives import Comparison
left_exp = Comparison(
left_exp,
self._COMP_TABLE[next_tag],
self.parse_expression(pstate, _PREC_COMPARISON))
did_something = True
elif next_tag is _colon and _PREC_SLICE >= min_precedence:
pstate.advance()
expr_pstate = pstate.copy()
assert not isinstance(left_exp, primitives.Slice)
from pytools.lex import ParseError
try:
next_expr = self.parse_expression(expr_pstate, _PREC_SLICE)
except ParseError:
# no expression follows, too bad.
left_exp = primitives.Slice((left_exp, None,))
else:
left_exp = _join_to_slice(left_exp, next_expr)
pstate.assign(expr_pstate)
did_something = True
elif next_tag is _comma and _PREC_COMMA > min_precedence:
# The precedence makes the comma left-associative.
pstate.advance()
if pstate.is_at_end() or pstate.next_tag() is _closepar:
if isinstance(left_exp, (tuple, list)) \
and not isinstance(left_exp, FinalizedContainer):
# left_expr is a container with trailing commas
pass
else:
left_exp = (left_exp,)
else:
new_el = self.parse_expression(pstate, _PREC_COMMA)
if isinstance(left_exp, (tuple, list)) \
and not isinstance(left_exp, FinalizedContainer):
left_exp = left_exp + (new_el,)
else:
left_exp = (left_exp, new_el)
did_something = True
return left_exp, did_something
def __init__(self, fft, effective_k, dk, dx):
self.fft = fft
if not callable(effective_k):
if effective_k != 0:
from pystella.derivs import FirstCenteredDifference
h = effective_k
effective_k = FirstCenteredDifference(h).get_eigenvalues
else:
def effective_k(k, dx): # pylint: disable=function-redefined
return k
queue = self.fft.sub_k["momenta_x"].queue
sub_k = list(x.get().astype("int") for x in self.fft.sub_k.values())
eff_mom_names = ("eff_mom_x", "eff_mom_y", "eff_mom_z")
self.eff_mom = {}
for mu, (name, kk) in enumerate(zip(eff_mom_names, sub_k)):
eff_k = effective_k(dk[mu] * kk.astype(fft.rdtype), dx[mu])
eff_k[abs(sub_k[mu]) == fft.grid_shape[mu] // 2] = 0.
eff_k[sub_k[mu] == 0] = 0.
import pyopencl.array as cla
self.eff_mom[name] = cla.to_device(queue, eff_k)
from pymbolic import var, parse
from pymbolic.primitives import If, Comparison, LogicalAnd
from pystella import Field
indices = parse("i, j, k")
eff_k = tuple(
var(array)[mu] for array, mu in zip(eff_mom_names, indices))
fabs, sqrt, conj = parse("fabs, sqrt, conj")
kmag = sqrt(sum(kk**2 for kk in eff_k))
from pystella import ElementWiseMap
vector = Field("vector", shape=(3, ))
vector_T = Field("vector_T", shape=(3, ))
kvec_zero = LogicalAnd(
tuple(Comparison(fabs(eff_k[mu]), "<", 1e-14) for mu in range(3)))
# note: write all output via private temporaries to allow for in-place
div = var("div")
div_insn = [(div, sum(eff_k[mu] * vector[mu] for mu in range(3)))]
self.transversify_knl = ElementWiseMap(
{
vector_T[mu]: If(kvec_zero, 0,
vector[mu] - eff_k[mu] / kmag**2 * div)
for mu in range(3)
},
tmp_instructions=div_insn,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
import loopy as lp
def assign(asignee, expr, **kwargs):
default = dict(within_inames=frozenset(("i", "j", "k")),
no_sync_with=[("*", "any")])
default.update(kwargs)
return lp.Assignment(asignee, expr, **default)
kmag, Kappa = parse("kmag, Kappa")
eps_insns = [
assign(kmag, sqrt(sum(kk**2 for kk in eff_k))),
assign(Kappa, sqrt(sum(kk**2 for kk in eff_k[:2])))
]
zero = fft.cdtype.type(0)
kx_ky_zero = LogicalAnd(
tuple(Comparison(fabs(eff_k[mu]), "<", 1e-10) for mu in range(2)))
kz_nonzero = Comparison(fabs(eff_k[2]), ">", 1e-10)
eps = var("eps")
eps_insns.extend([
assign(
eps[0],
If(kx_ky_zero, If(kz_nonzero, fft.cdtype.type(1 / 2**.5),
zero),
(eff_k[0] * eff_k[2] / kmag - 1j * eff_k[1]) / Kappa /
2**.5)),
assign(
eps[1],
If(kx_ky_zero,
If(kz_nonzero, fft.cdtype.type(1j / 2**(1 / 2)),
zero), (eff_k[1] * eff_k[2] / kmag + 1j * eff_k[0]) /
Kappa / 2**.5)),
assign(eps[2], If(kx_ky_zero, zero, -Kappa / kmag / 2**.5))
])
plus, minus, lng = Field("plus"), Field("minus"), Field("lng")
plus_tmp, minus_tmp = parse("plus_tmp, minus_tmp")
pol_isns = [(plus_tmp,
sum(vector[mu] * conj(eps[mu]) for mu in range(3))),
(minus_tmp, sum(vector[mu] * eps[mu] for mu in range(3)))]
args = [
lp.TemporaryVariable("kmag"),
lp.TemporaryVariable("Kappa"),
lp.TemporaryVariable("eps", shape=(3, )), ...
]
self.vec_to_pol_knl = ElementWiseMap(
{
plus: plus_tmp,
minus: minus_tmp
},
tmp_instructions=eps_insns + pol_isns,
args=args,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
vector_tmp = var("vector_tmp")
vec_insns = [(vector_tmp[mu], plus * eps[mu] + minus * conj(eps[mu]))
for mu in range(3)]
self.pol_to_vec_knl = ElementWiseMap(
{vector[mu]: vector_tmp[mu]
for mu in range(3)},
tmp_instructions=eps_insns + vec_insns,
args=args,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
ksq = sum(kk**2 for kk in eff_k)
lng_rhs = If(kvec_zero, 0, -div / ksq * 1j)
self.vec_decomp_knl = ElementWiseMap(
{
plus: plus_tmp,
minus: minus_tmp,
lng: lng_rhs
},
tmp_instructions=eps_insns + pol_isns + div_insn,
args=args,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
lng_rhs = If(kvec_zero, 0, -div / ksq**.5 * 1j)
self.vec_decomp_knl_times_abs_k = ElementWiseMap(
{
plus: plus_tmp,
minus: minus_tmp,
lng: lng_rhs
},
tmp_instructions=eps_insns + pol_isns + div_insn,
args=args,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
from pystella.sectors import tensor_index as tid
eff_k_hat = tuple(kk / sqrt(sum(kk**2 for kk in eff_k))
for kk in eff_k)
hij = Field("hij", shape=(6, ))
hij_TT = Field("hij_TT", shape=(6, ))
Pab = var("P")
Pab_insns = [(Pab[tid(a, b)], (If(Comparison(a, "==", b), 1, 0) -
eff_k_hat[a - 1] * eff_k_hat[b - 1]))
for a in range(1, 4) for b in range(a, 4)]
hij_TT_tmp = var("hij_TT_tmp")
TT_insns = [(hij_TT_tmp[tid(a, b)],
sum((Pab[tid(a, c)] * Pab[tid(d, b)] -
Pab[tid(a, b)] * Pab[tid(c, d)] / 2) * hij[tid(c, d)]
for c in range(1, 4) for d in range(1, 4)))
for a in range(1, 4) for b in range(a, 4)]
# note: where conditionals (branch divergence) go can matter:
# this kernel is twice as fast when putting the branching in the global
# write, rather than when setting hij_TT_tmp
write_insns = [(hij_TT[tid(a,
b)], If(kvec_zero, 0, hij_TT_tmp[tid(a,
b)]))
for a in range(1, 4) for b in range(a, 4)]
self.tt_knl = ElementWiseMap(
write_insns,
tmp_instructions=Pab_insns + TT_insns,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
tensor_to_pol_insns = {
plus:
sum(hij[tid(c, d)] * conj(eps[c - 1]) * conj(eps[d - 1])
for c in range(1, 4) for d in range(1, 4)),
minus:
sum(hij[tid(c, d)] * eps[c - 1] * eps[d - 1] for c in range(1, 4)
for d in range(1, 4))
}
self.tensor_to_pol_knl = ElementWiseMap(
tensor_to_pol_insns,
tmp_instructions=eps_insns,
args=args,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
pol_to_tensor_insns = {
hij[tid(a, b)]: (plus * eps[a - 1] * eps[b - 1] +
minus * conj(eps[a - 1]) * conj(eps[b - 1]))
for a in range(1, 4) for b in range(a, 4)
}
self.pol_to_tensor_knl = ElementWiseMap(
pol_to_tensor_insns,
tmp_instructions=eps_insns,
args=args,
lsize=(32, 1, 1),
rank_shape=fft.shape(True),
)
def _add_statement(self, stmt):
stmt_id = self.next_statement_id()
read_variables = set(stmt.get_read_variables())
written_variables = set(stmt.get_written_variables())
# Add the global execution state as an implicitly read variable.
read_variables.add(self._EXECUTION_STATE)
# Build the condition attribute.
if not self._conditional_expression_stack:
condition = True
elif len(self._conditional_expression_stack) == 1:
condition = self._conditional_expression_stack[0]
else:
from pymbolic.primitives import LogicalAnd
condition = LogicalAnd(tuple(self._conditional_expression_stack))
from dagrt.utils import get_variables
read_variables |= get_variables(condition)
is_non_assignment = (not isinstance(
stmt, (Assign, AssignImplicit, AssignFunctionCall)))
# We regard all non-assignments as having potential external side
# effects (i.e., writing to EXECUTION_STATE). To keep the global
# variables in a well-defined state, ensure that all updates to global
# variables have happened before a non-assignment.
if is_non_assignment:
from dagrt.utils import is_state_variable
read_variables |= {
var
for var in self._seen_var_names if is_state_variable(var)
}
written_variables.add(self._EXECUTION_STATE)
depends_on = set()
# Ensure this statement happens after the last write of all the
# variables it reads or writes.
for var in read_variables | written_variables:
writer = self._writer_map.get(var, None)
if writer is not None:
depends_on.add(writer)
# Ensure this statement happens after the last read(s) of the variables
# it writes to.
for var in written_variables:
readers = self._reader_map.get(var, set())
depends_on |= readers
# Keep the graph sparse by clearing the readers set.
readers.clear()
for var in written_variables:
self._writer_map[var] = stmt_id
for var in read_variables:
# reader_map should ignore reads that happen before writes, so
# ignore if this statement also reads *var*.
if var in written_variables:
continue
self._reader_map.setdefault(var, set()).add(stmt_id)
stmt = stmt.copy(id=stmt_id,
condition=condition,
depends_on=frozenset(depends_on))
self.statements.append(stmt)
self._seen_var_names |= read_variables | written_variables
def parse_postfix(self, pstate, min_precedence, left_exp):
import pymbolic.primitives as primitives
did_something = False
next_tag = pstate.next_tag()
if next_tag is _openpar and _PREC_CALL > min_precedence:
pstate.advance()
args, kwargs = self.parse_arglist(pstate)
if kwargs:
left_exp = primitives.CallWithKwargs(left_exp, args, kwargs)
else:
left_exp = primitives.Call(left_exp, args)
did_something = True
elif next_tag is _openbracket and _PREC_CALL > min_precedence:
pstate.advance()
pstate.expect_not_end()
left_exp = primitives.Subscript(left_exp,
self.parse_expression(pstate))
pstate.expect(_closebracket)
pstate.advance()
did_something = True
elif next_tag is _dot and _PREC_CALL > min_precedence:
pstate.advance()
pstate.expect(_identifier)
left_exp = primitives.Lookup(left_exp, pstate.next_str())
pstate.advance()
did_something = True
elif next_tag is _plus and _PREC_PLUS > min_precedence:
pstate.advance()
left_exp += self.parse_expression(pstate, _PREC_PLUS)
did_something = True
elif next_tag is _minus and _PREC_PLUS > min_precedence:
pstate.advance()
left_exp -= self.parse_expression(pstate, _PREC_PLUS)
did_something = True
elif next_tag is _times and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp *= self.parse_expression(pstate, _PREC_TIMES)
did_something = True
elif next_tag is _floordiv and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp //= self.parse_expression(pstate, _PREC_TIMES)
did_something = True
elif next_tag is _over and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp /= self.parse_expression(pstate, _PREC_TIMES)
did_something = True
elif next_tag is _modulo and _PREC_TIMES > min_precedence:
pstate.advance()
left_exp %= self.parse_expression(pstate, _PREC_TIMES)
did_something = True
elif next_tag is _power and _PREC_POWER > min_precedence:
pstate.advance()
left_exp **= self.parse_expression(pstate, _PREC_POWER)
did_something = True
elif next_tag is _and and _PREC_LOGICAL_AND > min_precedence:
pstate.advance()
from pymbolic.primitives import LogicalAnd
left_exp = LogicalAnd(
(left_exp, self.parse_expression(pstate, _PREC_LOGICAL_AND)))
did_something = True
elif next_tag is _or and _PREC_LOGICAL_OR > min_precedence:
pstate.advance()
from pymbolic.primitives import LogicalOr
left_exp = LogicalOr(
(left_exp, self.parse_expression(pstate, _PREC_LOGICAL_OR)))
did_something = True
elif next_tag is _bitwiseor and _PREC_BITWISE_OR > min_precedence:
pstate.advance()
from pymbolic.primitives import BitwiseOr
left_exp = BitwiseOr(
(left_exp, self.parse_expression(pstate, _PREC_BITWISE_OR)))
did_something = True
elif next_tag is _bitwisexor and _PREC_BITWISE_XOR > min_precedence:
pstate.advance()
from pymbolic.primitives import BitwiseXor
left_exp = BitwiseXor(
(left_exp, self.parse_expression(pstate, _PREC_BITWISE_XOR)))
did_something = True
elif next_tag is _bitwiseand and _PREC_BITWISE_AND > min_precedence:
pstate.advance()
from pymbolic.primitives import BitwiseAnd
left_exp = BitwiseAnd(
(left_exp, self.parse_expression(pstate, _PREC_BITWISE_AND)))
did_something = True
elif next_tag is _rightshift and _PREC_SHIFT > min_precedence:
pstate.advance()
from pymbolic.primitives import RightShift
left_exp = RightShift(left_exp,
self.parse_expression(pstate, _PREC_SHIFT))
did_something = True
elif next_tag is _leftshift and _PREC_SHIFT > min_precedence:
pstate.advance()
from pymbolic.primitives import LeftShift
left_exp = LeftShift(left_exp,
self.parse_expression(pstate, _PREC_SHIFT))
did_something = True
elif next_tag in self._COMP_TABLE and _PREC_COMPARISON > min_precedence:
pstate.advance()
from pymbolic.primitives import Comparison
left_exp = Comparison(
left_exp, self._COMP_TABLE[next_tag],
self.parse_expression(pstate, _PREC_COMPARISON))
did_something = True
elif next_tag is _colon and _PREC_SLICE >= min_precedence:
pstate.advance()
expr_pstate = pstate.copy()
assert not isinstance(left_exp, primitives.Slice)
from pytools.lex import ParseError
try:
next_expr = self.parse_expression(expr_pstate, _PREC_SLICE)
except ParseError:
# no expression follows, too bad.
left_exp = primitives.Slice((
left_exp,
None,
))
else:
left_exp = _join_to_slice(left_exp, next_expr)
pstate.assign(expr_pstate)
elif next_tag is _comma and _PREC_COMMA > min_precedence:
# The precedence makes the comma left-associative.
pstate.advance()
if pstate.is_at_end() or pstate.next_tag() is _closepar:
left_exp = (left_exp, )
else:
new_el = self.parse_expression(pstate, _PREC_COMMA)
if isinstance(left_exp, tuple) \
and not isinstance(left_exp, FinalizedTuple):
left_exp = left_exp + (new_el, )
else:
left_exp = (left_exp, new_el)
did_something = True
return left_exp, did_something
def __call__(self, dtype, operand1, operand2, callables_table, target):
from pymbolic.primitives import LogicalAnd
return LogicalAnd((operand1, operand2)), callables_table
