def expression_hash(self):
from firedrake.slate.slac.utils import traverse_dags
hashdata = []
for op in traverse_dags([self]):
if isinstance(op, AssembledVector):
data = (
type(op).__name__,
op.arg_function_spaces[0].ufl_element().
_ufl_signature_data_(),
)
elif isinstance(op, Block):
data = (
type(op).__name__,
op._indices,
)
elif isinstance(op, BlockAssembledVector):
data = (type(op).__name__, op._indices, op._original_function,
op._function)
elif isinstance(op, Factorization):
data = (
type(op).__name__,
op.decomposition,
)
elif isinstance(op, Tensor):
data = (op.form.signature(), )
elif isinstance(op, (UnaryOp, BinaryOp)):
data = (type(op).__name__, )
else:
raise ValueError("Unhandled type %r" % type(op))
hashdata.append(data + (op.prec, ))
hashdata = "".join("%s" % (s, ) for s in hashdata)
return hashlib.sha512(hashdata.encode("utf-8")).hexdigest()
def expression_flops(self):
@singledispatch
def _flops(expr):
raise AssertionError("Unhandled type %r" % type(expr))
@_flops.register(slate.AssembledVector)
@_flops.register(slate.Block)
@_flops.register(slate.Tensor)
@_flops.register(slate.Transpose)
@_flops.register(slate.Negative)
def _flops_none(expr):
return 0
@_flops.register(slate.Factorization)
def _flops_factorization(expr):
m, n = expr.shape
decomposition = expr.decomposition
# Extracted from Golub & Van Loan
# These all ignore lower-order terms...
if decomposition in {"PartialPivLU", "FullPivLU"}:
return 2/3 * n**3
elif decomposition in {"LLT", "LDLT"}:
return (1/3)*n**3
elif decomposition in {"HouseholderQR", "ColPivHouseholderQR", "FullPivHouseholderQR"}:
return 4/3 * n**3
elif decomposition in {"BDCSVD", "JacobiSVD"}:
return 12 * n**3
else:
# Don't know, but don't barf just because of it.
return 0
@_flops.register(slate.Inverse)
def _flops_inverse(expr):
m, n = expr.shape
assert m == n
# Assume LU factorisation
return (2/3)*n**3
@_flops.register(slate.Add)
def _flops_add(expr):
return int(np.prod(expr.shape))
@_flops.register(slate.Mul)
def _flops_mul(expr):
A, B = expr.operands
*rest_a, col = A.shape
_, *rest_b = B.shape
return 2*col*int(np.prod(rest_a))*int(np.prod(rest_b))
@_flops.register(slate.Solve)
def _flops_solve(expr):
Afac, B = expr.operands
_, *rest = B.shape
m, n = Afac.shape
# Forward elimination + back sub on factorised matrix
return (m*n + n**2)*int(np.prod(rest))
return int(sum(map(_flops, traverse_dags([self.expression]))))
def expression_flops(self):
@singledispatch
def _flops(expr):
raise AssertionError("Unhandled type %r" % type(expr))
@_flops.register(slate.AssembledVector)
@_flops.register(slate.Block)
@_flops.register(slate.Tensor)
@_flops.register(slate.Transpose)
@_flops.register(slate.Negative)
def _flops_none(expr):
return 0
@_flops.register(slate.Factorization)
def _flops_factorization(expr):
m, n = expr.shape
decomposition = expr.decomposition
# Extracted from Golub & Van Loan
# These all ignore lower-order terms...
if decomposition in {"PartialPivLU", "FullPivLU"}:
return 2/3 * n**3
elif decomposition in {"LLT", "LDLT"}:
return (1/3)*n**3
elif decomposition in {"HouseholderQR", "ColPivHouseholderQR", "FullPivHouseholderQR"}:
return 4/3 * n**3
elif decomposition in {"BDCSVD", "JacobiSVD"}:
return 12 * n**3
else:
# Don't know, but don't barf just because of it.
return 0
@_flops.register(slate.Inverse)
def _flops_inverse(expr):
m, n = expr.shape
assert m == n
# Assume LU factorisation
return (2/3)*n**3
@_flops.register(slate.Add)
def _flops_add(expr):
return int(np.prod(expr.shape))
@_flops.register(slate.Mul)
def _flops_mul(expr):
A, B = expr.operands
*rest_a, col = A.shape
_, *rest_b = B.shape
return 2*col*int(np.prod(rest_a))*int(np.prod(rest_b))
@_flops.register(slate.Solve)
def _flops_solve(expr):
Afac, B = expr.operands
_, *rest = B.shape
m, n = Afac.shape
# Forward elimination + back sub on factorised matrix
return (m*n + n**2)*int(np.prod(rest))
return int(sum(map(_flops, traverse_dags([self.expression]))))
def __init__(self, expression, tsfc_parameters=None):
"""Constructor for the LocalKernelBuilder class.
:arg expression: a :class:`TensorBase` object.
:arg tsfc_parameters: an optional `dict` of parameters to provide to
TSFC when constructing subkernels associated with the expression.
"""
assert isinstance(expression, slate.TensorBase)
# Collect terminals, expressions, and reference counts
temps = OrderedDict()
coeff_vecs = OrderedDict()
seen_coeff = set()
expression_dag = list(traverse_dags([expression]))
counter = Counter([expression])
for tensor in expression_dag:
counter.update(tensor.operands)
# Terminal tensors will always require a temporary.
if isinstance(tensor, slate.Tensor):
temps.setdefault(tensor, ast.Symbol("T%d" % len(temps)))
# 'AssembledVector's will always require a coefficient temporary.
if isinstance(tensor, slate.AssembledVector):
function = tensor._function
def dimension(e):
return create_element(e).space_dimension()
# Ensure coefficient temporaries aren't duplicated
if function not in seen_coeff:
if type(function.ufl_element()) == MixedElement:
shapes = [dimension(element) for element in function.ufl_element().sub_elements()]
else:
shapes = [dimension(function.ufl_element())]
# Local temporary
local_temp = ast.Symbol("VecTemp%d" % len(seen_coeff))
offset = 0
for i, shape in enumerate(shapes):
cinfo = CoefficientInfo(space_index=i,
offset_index=offset,
shape=(sum(shapes), ),
vector=tensor,
local_temp=local_temp)
coeff_vecs.setdefault(shape, []).append(cinfo)
offset += shape
seen_coeff.add(function)
self.expression = expression
self.tsfc_parameters = tsfc_parameters
self.temps = temps
self.ref_counter = counter
self.expression_dag = expression_dag
self.coefficient_vecs = coeff_vecs
self._setup()
def expression_hash(self):
from firedrake.slate.slac.utils import traverse_dags
hashdata = []
for op in traverse_dags([self]):
if isinstance(op, AssembledVector):
data = (type(op).__name__, op.arg_function_spaces[0].ufl_element()._ufl_signature_data_(), )
elif isinstance(op, Block):
data = (type(op).__name__, op._indices, )
elif isinstance(op, Factorization):
data = (type(op).__name__, op.decomposition, )
elif isinstance(op, Tensor):
data = (op.form.signature(), )
elif isinstance(op, (UnaryOp, BinaryOp)):
data = (type(op).__name__, )
else:
raise ValueError("Unhandled type %r" % type(op))
hashdata.append(data + (op.prec, ))
hashdata = "".join("%s" % (s, ) for s in hashdata)
return hashlib.sha512(hashdata.encode("utf-8")).hexdigest()
def generate_expr_data(expr):
"""This function generates a mapping of the form:
``temporaries = {node: symbol_name}``
where `node` objects are :class:`slate.TensorBase` nodes, and
`symbol_name` are :class:`coffee.base.Symbol` objects. In addition,
this function will return a list `aux_exprs` of any expressions that
require special handling in the compiler. This includes expressions
that require performing operations on already assembled data or
generating extra temporaries.
This mapping is used in the :class:`KernelBuilder` to provide direct
access to all temporaries associated with a particular slate expression.
:arg expression: a :class:`slate.TensorBase` object.
Returns: the terminal temporaries map and auxiliary temporaries.
"""
# Prepare temporaries map and auxiliary expressions list
# NOTE: Ordering here matters, especially when running
# Slate in parallel.
temps = OrderedDict()
aux_exprs = []
for tensor in traverse_dags([expr]):
if isinstance(tensor, Tensor):
temps.setdefault(tensor, ast.Symbol("T%d" % len(temps)))
elif isinstance(tensor, TensorOp):
# For Action, we need to declare a temporary later on for the
# acting coefficient. For inverses, we may declare additional
# temporaries (depending on reference count).
if isinstance(tensor, (Action, Inverse)):
aux_exprs.append(tensor)
# Aux expressions are visited pre-order. We want to declare as if we'd
# visited post-order (child temporaries first), so reverse.
aux_exprs = list(OrderedDict.fromkeys(reversed(aux_exprs)))
return temps, aux_exprs
def __init__(self, expression, tsfc_parameters=None):
"""Constructor for the LocalKernelBuilder class.
:arg expression: a :class:`TensorBase` object.
:arg tsfc_parameters: an optional `dict` of parameters to provide to
TSFC when constructing subkernels associated
with the expression.
"""
assert isinstance(expression, slate.TensorBase)
if expression.ufl_domain().variable_layers:
raise NotImplementedError(
"Variable layers not yet handled in Slate.")
# Collect terminals, expressions, and reference counts
temps = OrderedDict()
action_coeffs = OrderedDict()
seen_coeff = set()
expression_dag = list(traverse_dags([expression]))
counter = Counter([expression])
for tensor in expression_dag:
counter.update(tensor.operands)
# Terminal tensors will always require a temporary.
if isinstance(tensor, slate.Tensor):
temps.setdefault(tensor, ast.Symbol("T%d" % len(temps)))
# Actions will always require a coefficient temporary.
if isinstance(tensor, slate.Action):
actee, = tensor.actee
# Ensure coefficient temporaries aren't duplicated
if actee not in seen_coeff:
shapes = [(V.finat_element.space_dimension(), V.value_size)
for V in actee.function_space().split()]
c_shape = (sum(n * d for (n, d) in shapes), )
offset = 0
for fs_i, fs_shape in enumerate(shapes):
cinfo = CoefficientInfo(space_index=fs_i,
offset_index=offset,
shape=c_shape,
coefficient=actee)
action_coeffs.setdefault(fs_shape, []).append(cinfo)
offset += reduce(lambda x, y: x * y, fs_shape)
seen_coeff.add(actee)
self.expression = expression
self.tsfc_parameters = tsfc_parameters
self.temps = temps
# Terminal tensors do not need additional temps created for them
# and neither do Negative nodes.
self.aux_exprs = [
tensor for tensor in topological_sort(expression_dag)
if counter[tensor] > 1 and not isinstance(tensor, (slate.Tensor,
slate.Negative))
]
self.action_coefficients = action_coeffs
self._setup()
def __init__(self, expression, tsfc_parameters=None):
"""Constructor for the LocalKernelBuilder class.
:arg expression: a :class:`TensorBase` object.
:arg tsfc_parameters: an optional `dict` of parameters to provide to
TSFC when constructing subkernels associated with the expression.
"""
assert isinstance(expression, slate.TensorBase)
if expression.ufl_domain().variable_layers:
raise NotImplementedError("Variable layers not yet handled in Slate.")
# Collect terminals, expressions, and reference counts
temps = OrderedDict()
coeff_vecs = OrderedDict()
seen_coeff = set()
expression_dag = list(traverse_dags([expression]))
counter = Counter([expression])
for tensor in expression_dag:
counter.update(tensor.operands)
# Terminal tensors will always require a temporary.
if isinstance(tensor, slate.Tensor):
temps.setdefault(tensor, ast.Symbol("T%d" % len(temps)))
# 'AssembledVector's will always require a coefficient temporary.
if isinstance(tensor, slate.AssembledVector):
function = tensor._function
def dimension(e):
return create_element(e).space_dimension()
# Ensure coefficient temporaries aren't duplicated
if function not in seen_coeff:
if type(function.ufl_element()) == MixedElement:
shapes = [dimension(element) for element in function.ufl_element().sub_elements()]
else:
shapes = [dimension(function.ufl_element())]
# Local temporary
local_temp = ast.Symbol("VecTemp%d" % len(seen_coeff))
offset = 0
for i, shape in enumerate(shapes):
cinfo = CoefficientInfo(space_index=i,
offset_index=offset,
shape=(sum(shapes), ),
vector=tensor,
local_temp=local_temp)
coeff_vecs.setdefault(shape, []).append(cinfo)
offset += shape
seen_coeff.add(function)
self.expression = expression
self.tsfc_parameters = tsfc_parameters
self.temps = temps
self.ref_counter = counter
self.expression_dag = expression_dag
self.coefficient_vecs = coeff_vecs
self._setup()
