BUG: Decreased performance when assembling expressions with new constants

[jrmaddison]

Describe the bug
Assembly of expressions using constants can be slow if new Constant objects are instantiated.

Steps to Reproduce
Example code:

from firedrake import *

import time

mesh = UnitIntervalMesh(1)

# Same Constant, same value
t0 = time.perf_counter()
c = Constant(0.0)
for _ in range(100):
    assemble(c * dx(mesh))
t1 = time.perf_counter()
print(f"{t1 - t0}")

# Different Constant, same value
t0 = time.perf_counter()
for _ in range(100):
    assemble(Constant(0.0) * dx(mesh))
t1 = time.perf_counter()
print(f"{t1 - t0}")

# Same Constant, different value
t0 = time.perf_counter()
c = Constant(0.0)
for i in range(100):
    c.assign(i)
    assemble(c * dx(mesh))
t1 = time.perf_counter()
print(f"{t1 - t0}")

With revision c5cd37a and an empty cache:

0.5365729830009514
0.09846351200394565
0.08098157500353409

With revision a94b01c and an empty cache:

0.5236760989937466
39.5197532769962
1.8689304560029996

Environment:
Ubuntu 22.04

[connorjward]

Is this fixed by #2996?

[jrmaddison]

That raises an exception:

Traceback (most recent call last):
  File "/home/maddison/build/tlm_adjoint/tlm_adjoint/tmp/test.py", line 11, in <module>
    assemble(c * dx(mesh))
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/measure.py", line 423, in __rmul__
    return Form([integral])
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/form.py", line 294, in __init__
    self._constants = extract_constants(self)
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/algorithms/analysis.py", line 116, in extract_constants
    return sorted_by_count(extract_type(a, Constant))
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/algorithms/analysis.py", line 70, in extract_type
    return set(o for e in iter_expressions(a)
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/algorithms/analysis.py", line 70, in <genexpr>
    return set(o for e in iter_expressions(a)
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/corealg/traversal.py", line 134, in traverse_unique_terminals
    for op in unique_pre_traversal(expr, visited=visited):
  File "/home/maddison/firedrake/firedrake/src/ufl/ufl/corealg/traversal.py", line 70, in unique_pre_traversal
    visited.add(expr)
TypeError: unhashable type: 'Constant'

After adding:

    def __hash__(self):
        return self._ufl_compute_hash_()

to the Constant class, and with an empty cache:

0.6268445020032232
41.26168250699993
0.546623861999251

[wence-]

0.6268445020032232
41.26168250699993
0.546623861999251

The hash of a constant as implemented in #2996 depends on its count (which is unique for every new constant). With coefficients/arguments, ufl signature computation renumbers before computing the signature to determine the hash of a form. But it looks like it doesn't do so with these new constants. So every time you make a new constant you're seeing a "new" form with a new signature and paying the cost of compiling it.

[jrmaddison]

The previous performance is restored if Constant also inherits from ufl.Coefficient. But is there a reason not to inherit from both ufl.constantvalue.ConstantValue and ufl.Coefficient?

[jrmaddison]

A related issue:

from firedrake import *
import mpi4py.MPI as MPI

mesh = UnitIntervalMesh(100)
if MPI.COMM_WORLD.rank == 0:
    Constant(1.0)
print(assemble(Constant(1.0) * dx(mesh)))

leads to, with an empty cache,

pyop2.exceptions.CompilationError: Generated code differs across ranks (see output in /home/maddison/firedrake/firedrake/.cache/pyop2/mismatching-kernels)

EDIT: When run on more than one process.

[wence-]

EDIT: When run on more than one process.

Yeah, this is definitely a consequence of not relabeling constants before handing off to the compilation pipeline.

Constants are supposed to have no domain (and therefore not be collective over any communicator), however the implementation of labelling a constant with a count that is then used in code-gen/form caching makes them logically collective.

[wence-]

The previous performance is restored if Constant also inherits from ufl.Coefficient. But is there a reason not to inherit from both ufl.constantvalue.ConstantValue and ufl.Coefficient?

I suspect this means that UFL's ConstantValue needs to be relabelled in form preprocessing (or else Firedrake's use of it is wrong).

[jrmaddison]

Whether or not new compilation is triggered also depends on the order in which forms are assembled. I'm seeing this when running individual tests -- even if the whole test suite has been run, running a single test can lead to new compilation.

[wence-]

Whether or not new compilation is triggered also depends on the order in which forms are assembled. I'm seeing this when running individual tests -- even if the whole test suite has been run, running a single test can lead to new compilation.

This is a consequence of the same thing: if a Constants "count" ends up in the generated code, and the count is dependent on the order in which they are created...

[michalhabera]

I am not in the picture for the meshfree Constant you are preparing, but if you want more stable signature it must have rules for renumbering implemented here https://github.com/FEniCS/ufl/blob/main/ufl/form.py#L631 Upstream UFL only renumbers Domains, Coefficients and ufl.constant.Constant objects. See also FEniCS/ufl@07b61c3

[connorjward]

I have added this to the agenda for the next Firedrake meeting. Clearly a renumbering is required but we need to be careful about how we do this.

[connorjward]

I've implemented a fix in #3011 (with a corresponding UFL PR).

[wence-]

I think you are abusing the UFL type hierarchy here. ConstantValue in UFL parlance really is a "form-definition-time" constant (i.e. a compile-time literal: codegen for it just pastes in the .value slot). Whereas you're subclassing it to mean a runtime literal (that the assembler needs to provide a value for).

[connorjward]

We chose this approach because we thought that UFL shouldn't care whether or not values are provided at compile-time or runtime. The distinction is more of a Firedrake/TSFC implementation detail.

[wence-]

Ah, I see. I suppose it speaks to, in some senses, a bit of a misdesign in the way UFL handles domains. Why does a UFL Constant need to have a domain at all? I think in the past we have discussed (or at least thought about) requiring that the measure (dx say) be given a domain when constructing a form. This would be required to match any coefficients/arguments that are defined with a domain, but allows constants to not have domains at all. I think originally this was not the case because most UFL algorithms are designed to operate on the integrand expression (which therefore has thrown away information that might live in the measure).

[connorjward]

That seems like a reasonable approach.

[jrmaddison]

The following workaround seems to work:

from firedrake import *
from firedrake import Constant as _Constant

import time
import ufl


class Constant(_Constant, ufl.classes.Coefficient):
    def __init__(self, *args, **kwargs):
        _Constant.__init__(self, *args, **kwargs)
        
        domain = self.ufl_domain()
        if domain is None:
            cell = None
        else:
            cell = domain.ufl_cell()
        if len(self.ufl_shape) == 0:
            element = ufl.classes.FiniteElement("R", cell, 0)
        elif len(self.ufl_shape) == 1:
            element = ufl.classes.VectorElement("R", cell, 0,
                                                dim=self.ufl_shape[0])
        else:
            element = ufl.classes.TensorElement("R", cell, 0,
                                                shape=self.ufl_shape)
        space = ufl.classes.FunctionSpace(domain, element)
        ufl.classes.Coefficient.__init__(self, space)


mesh = UnitIntervalMesh(1)

# Same Constant, same value
t0 = time.perf_counter()
c = Constant(0.0)
for _ in range(100):
    assemble(c * dx(mesh))
t1 = time.perf_counter()
print(f"{t1 - t0}")

# Different Constant, same value
t0 = time.perf_counter()
for _ in range(100):
    assemble(Constant(0.0) * dx(mesh))
t1 = time.perf_counter()
print(f"{t1 - t0}")

# Same Constant, different value
t0 = time.perf_counter()
c = Constant(0.0)
for i in range(100):
    c.assign(i)
    assemble(c * dx(mesh))
t1 = time.perf_counter()
print(f"{t1 - t0}")

with timings (with an empty cache)

0.5934124190007424
0.16270772999996552
0.10279183799866587

Is this a bad idea?