def fix_ops_acc(expr, args):
if expr.is_Symbol or expr.is_Number:
return expr
if expr.is_Indexed:
return Indexed(
expr.base,
Macro('OPS_ACC%d(%s)' %
(args.index(expr.name), expr.indices[0].name)))
else:
for i in expr.args:
return expr.func(*[fix_ops_acc(i, args) for i in expr.args])
def _(f, szs, sregistry):
assert len(szs) == len(f.dimensions) - 1
pname = sregistry.make_name(prefix='%sL' % f.name)
cbk = lambda i, pname=pname: FIndexed(i, pname)
expr = sum([MacroArgument(d0.name)*szs[d1]
for d0, d1 in zip(f.dimensions, f.dimensions[1:])])
expr += MacroArgument(f.dimensions[-1].name)
expr = Indexed(IndexedData(f.name, None, f), expr)
define = DefFunction(pname, f.dimensions)
header = (ccode(define), ccode(expr))
return header, cbk
def new_ops_arg(self, indexed):
"""
Create an :class:`Indexed` node using OPS representation.
Parameters
----------
indexed : :class:`Indexed`
Indexed object using devito representation.
Returns
-------
:class:`Indexed`
Indexed node using OPS representation.
"""
# Build the OPS arg identifier
time_index = split_affine(indexed.indices[TimeFunction._time_position])
ops_arg_id = '%s%s%s' % (indexed.name, time_index.var,
time_index.shift)
if ops_arg_id not in self.ops_args:
# Create the indexed object
ops_arg = Array(name=ops_arg_id,
dimensions=[Dimension(name=namespace['ops_acc'])],
dtype=indexed.dtype)
self.ops_args[ops_arg_id] = ops_arg
else:
ops_arg = self.ops_args[ops_arg_id]
# Get the space indices
space_indices = [
e for i, e in enumerate(indexed.indices)
if i != TimeFunction._time_position
]
# Define the Macro used in OPS arg index
access_macro = Macro(
'OPS_ACC%d(%s)' % (list(self.ops_args).index(ops_arg_id), ','.join(
str(split_affine(i).shift) for i in space_indices)))
# Create Indexed object representing the OPS arg access
new_indexed = Indexed(ops_arg.indexed, access_macro)
return new_indexed
def promote_scalar_expressions(exprs, shape, indices, onstack):
"""
Transform a collection of scalar expressions into tensor expressions.
"""
processed = []
# Fist promote the LHS
mapper = {}
for k, v in FlowGraph(exprs).items():
if v.is_scalar:
# Create a new function symbol
data = Array(name=k.name, shape=shape,
dimensions=indices, onstack=onstack)
indexed = Indexed(data.indexed, *indices)
mapper[k] = indexed
processed.append(Eq(indexed, v.rhs))
else:
processed.append(Eq(k, v.rhs))
# Propagate the transformed LHS through the expressions
processed = [Eq(n.lhs, n.rhs.xreplace(mapper)) for n in processed]
return processed
def _eliminate_inter_stencil_redundancies(self, cluster, template,
**kwargs):
"""
Search for redundancies across the expressions and expose them
to the later stages of the optimisation pipeline by introducing
new temporaries of suitable rank.
Two type of redundancies are sought:
* Time-invariants, and
* Across different space points
Examples
========
Let ``t`` be the time dimension, ``x, y, z`` the space dimensions. Then:
1) temp = (a[x,y,z]+b[x,y,z])*c[t,x,y,z]
>>>
ti[x,y,z] = a[x,y,z] + b[x,y,z]
temp = ti[x,y,z]*c[t,x,y,z]
2) temp1 = 2.0*a[x,y,z]*b[x,y,z]
temp2 = 3.0*a[x,y,z+1]*b[x,y,z+1]
>>>
ti[x,y,z] = a[x,y,z]*b[x,y,z]
temp1 = 2.0*ti[x,y,z]
temp2 = 3.0*ti[x,y,z+1]
"""
if cluster.is_sparse:
return cluster
# For more information about "aliases", refer to collect.__doc__
mapper, aliases = collect(cluster.exprs)
# Redundancies will be stored in space-varying temporaries
g = cluster.trace
indices = g.space_indices
time_invariants = {v.rhs: g.time_invariant(v) for v in g.values()}
# Find the candidate expressions
processed = []
candidates = OrderedDict()
for k, v in g.items():
# Cost check (to keep the memory footprint under control)
naliases = len(mapper.get(v.rhs, []))
cost = estimate_cost(v, True) * naliases
if cost >= self.thresholds['min-cost-alias'] and\
(naliases > 1 or time_invariants[v.rhs]):
candidates[v.rhs] = k
else:
processed.append(v)
# Create alias Clusters and all necessary substitution rules
# for the new temporaries
alias_clusters = ClusterGroup()
rules = OrderedDict()
for origin, alias in aliases.items():
if all(i not in candidates for i in alias.aliased):
continue
# Construct an iteration space suitable for /alias/
intervals, sub_iterators, directions = cluster.ispace.args
intervals = [
Interval(i.dim, *alias.relaxed_diameter.get(i.dim, i.limits))
for i in cluster.ispace.intervals
]
ispace = IterationSpace(intervals, sub_iterators, directions)
# Optimization: perhaps we can lift the cluster outside the time dimension
if all(time_invariants[i] for i in alias.aliased):
ispace = ispace.project(lambda i: not i.is_Time)
# Build a symbolic function for /alias/
intervals = ispace.intervals
halo = [(abs(intervals[i].lower), abs(intervals[i].upper))
for i in indices]
function = Array(name=template(), dimensions=indices, halo=halo)
access = tuple(i - intervals[i].lower for i in indices)
expression = Eq(Indexed(function.indexed, *access), origin)
# Construct a data space suitable for /alias/
mapper = detect_accesses(expression)
parts = {
k: IntervalGroup(build_intervals(v)).add(intervals)
for k, v in mapper.items() if k
}
dspace = DataSpace([i.zero() for i in intervals], parts)
# Create a new Cluster for /alias/
alias_clusters.append(Cluster([expression], ispace, dspace))
# Add substitution rules
for aliased, distance in alias.with_distance:
access = [
i - intervals[i].lower + j for i, j in distance
if i in indices
]
temporary = Indexed(function.indexed, *tuple(access))
rules[candidates[aliased]] = temporary
rules[aliased] = temporary
# Group clusters together if possible
alias_clusters = groupby(alias_clusters).finalize()
alias_clusters.sort(key=lambda i: i.is_dense)
# Switch temporaries in the expression trees
processed = [e.xreplace(rules) for e in processed]
return alias_clusters + [cluster.rebuild(processed)]
def _eliminate_inter_stencil_redundancies(self, cluster, template,
**kwargs):
"""
Search for redundancies across the expressions and expose them
to the later stages of the optimisation pipeline by introducing
new temporaries of suitable rank.
Two type of redundancies are sought:
* Time-invariants, and
* Across different space points
Examples
========
Let ``t`` be the time dimension, ``x, y, z`` the space dimensions. Then:
1) temp = (a[x,y,z]+b[x,y,z])*c[t,x,y,z]
>>>
ti[x,y,z] = a[x,y,z] + b[x,y,z]
temp = ti[x,y,z]*c[t,x,y,z]
2) temp1 = 2.0*a[x,y,z]*b[x,y,z]
temp2 = 3.0*a[x,y,z+1]*b[x,y,z+1]
>>>
ti[x,y,z] = a[x,y,z]*b[x,y,z]
temp1 = 2.0*ti[x,y,z]
temp2 = 3.0*ti[x,y,z+1]
"""
if cluster.is_sparse:
return cluster
# For more information about "aliases", refer to collect.__doc__
mapper, aliases = collect(cluster.exprs)
# Redundancies will be stored in space-varying temporaries
g = cluster.trace
indices = g.space_indices
time_invariants = {v.rhs: g.time_invariant(v) for v in g.values()}
# Template for captured redundancies
shape = tuple(i.symbolic_extent for i in indices)
make = lambda i: Array(
name=template(i), shape=shape, dimensions=indices).indexed
# Find the candidate expressions
processed = []
candidates = OrderedDict()
for k, v in g.items():
# Cost check (to keep the memory footprint under control)
naliases = len(mapper.get(v.rhs, []))
cost = estimate_cost(v, True) * naliases
if cost >= self.thresholds['min-cost-alias'] and\
(naliases > 1 or time_invariants[v.rhs]):
candidates[v.rhs] = k
else:
processed.append(Eq(k, v.rhs))
# Create temporaries capturing redundant computation
expressions = []
stencils = []
rules = OrderedDict()
for c, (origin, alias) in enumerate(aliases.items()):
if all(i not in candidates for i in alias.aliased):
continue
# Build alias expression
function = make(c)
expressions.append(Eq(Indexed(function, *indices), origin))
# Build substitution rules
for aliased, distance in alias.with_distance:
coordinates = [
sum([i, j]) for i, j in distance.items() if i in indices
]
temporary = Indexed(function, *tuple(coordinates))
rules[candidates[aliased]] = temporary
rules[aliased] = temporary
# Build cluster stencil
stencil = alias.anti_stencil.anti(cluster.stencil)
if all(time_invariants[i] for i in alias.aliased):
# Optimization: drop time dimension if time-invariant and the
# alias involves a complex calculation
stencil = stencil.section(g.time_indices)
stencils.append(stencil)
# Create the alias clusters
alias_clusters = clusterize(expressions, stencils)
alias_clusters = sorted(alias_clusters, key=lambda i: i.is_dense)
# Switch temporaries in the expression trees
processed = [e.xreplace(rules) for e in processed]
return alias_clusters + [cluster.rebuild(processed)]
def linearize_accesses(iet, cache, sregistry):
"""
Turn Indexeds into FIndexeds and create the necessary access Macros.
"""
# Find all objects amenable to linearization
symbol_names = {i.name for i in FindSymbols('indexeds').visit(iet)}
functions = [f for f in FindSymbols().visit(iet)
if ((f.is_DiscreteFunction or f.is_Array) and
f.ndim > 1 and
f.name in symbol_names)]
functions = sorted(functions, key=lambda f: len(f.dimensions), reverse=True)
# Find unique sizes (unique -> minimize necessary registers)
mapper = DefaultOrderedDict(list)
for f in functions:
if f not in cache:
# NOTE: the outermost dimension is unnecessary
for d in f.dimensions[1:]:
# TODO: same grid + same halo => same padding, however this is
# never asserted throughout the compiler yet... maybe should do
# it when in debug mode at `prepare_arguments` time, ie right
# before jumping to C?
mapper[(d, f._size_halo[d], getattr(f, 'grid', None))].append(f)
# Build all exprs such as `x_fsz0 = u_vec->size[1]`
imapper = DefaultOrderedDict(list)
for (d, halo, _), v in mapper.items():
name = sregistry.make_name(prefix='%s_fsz' % d.name)
s = Symbol(name=name, dtype=np.int32, is_const=True)
try:
expr = DummyExpr(s, v[0]._C_get_field(FULL, d).size, init=True)
except AttributeError:
assert v[0].is_Array
expr = DummyExpr(s, v[0].symbolic_shape[d], init=True)
for f in v:
imapper[f].append((d, s))
cache[f].stmts0.append(expr)
# Build all exprs such as `y_slc0 = y_fsz0*z_fsz0`
built = {}
mapper = DefaultOrderedDict(list)
for f, v in imapper.items():
for n, (d, _) in enumerate(v):
expr = prod(list(zip(*v[n:]))[1])
try:
stmt = built[expr]
except KeyError:
name = sregistry.make_name(prefix='%s_slc' % d.name)
s = Symbol(name=name, dtype=np.int32, is_const=True)
stmt = built[expr] = DummyExpr(s, expr, init=True)
mapper[f].append(stmt.write)
cache[f].stmts1.append(stmt)
mapper.update([(f, []) for f in functions if f not in mapper])
# Build defines. For example:
# `define uL(t, x, y, z) u[(t)*t_slice_sz + (x)*x_slice_sz + (y)*y_slice_sz + (z)]`
headers = []
findexeds = {}
for f, szs in mapper.items():
if cache[f].cbk is not None:
# Perhaps we've already built an access macro for `f` through another efunc
findexeds[f] = cache[f].cbk
else:
assert len(szs) == len(f.dimensions) - 1
pname = sregistry.make_name(prefix='%sL' % f.name)
expr = sum([MacroArgument(d.name)*s for d, s in zip(f.dimensions, szs)])
expr += MacroArgument(f.dimensions[-1].name)
expr = Indexed(IndexedData(f.name, None, f), expr)
define = DefFunction(pname, f.dimensions)
headers.append((ccode(define), ccode(expr)))
cache[f].cbk = findexeds[f] = lambda i, pname=pname: FIndexed(i, pname)
# Build "functional" Indexeds. For example:
# `u[t2, x+8, y+9, z+7] => uL(t2, x+8, y+9, z+7)`
mapper = {}
for n in FindNodes(Expression).visit(iet):
subs = {}
for i in retrieve_indexed(n.expr):
try:
subs[i] = findexeds[i.function](i)
except KeyError:
pass
mapper[n] = n._rebuild(expr=uxreplace(n.expr, subs))
# Put together all of the necessary exprs for `y_fsz0`, ..., `y_slc0`, ...
stmts0 = filter_ordered(flatten(cache[f].stmts0 for f in functions))
if stmts0:
stmts0.append(BlankLine)
stmts1 = filter_ordered(flatten(cache[f].stmts1 for f in functions))
if stmts1:
stmts1.append(BlankLine)
iet = Transformer(mapper).visit(iet)
body = iet.body._rebuild(body=tuple(stmts0) + tuple(stmts1) + iet.body.body)
iet = iet._rebuild(body=body)
return iet, headers