11642291501

Committed 02 Nov 2024 11:16AM UTC coverage: 81.168% (+0.5%) from 80.657%

Build # 11642291501

Build Type

push

github

Committed by

web-flow

Commit Message

Upload to coveralls and docs from CI job running against python 3.12 (#726)

Run Details

3474 of 4280 relevant lines covered (81.17%)

0.81 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.79

/ffcx/codegeneration/integral_generator.py

# Copyright (C) 2015-2024 Martin Sandve Alnæs, Michal Habera, Igor Baratta, Chris Richardson
#
# Modified by Jørgen S. Dokken, 2024
#
# This file is part of FFCx. (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later
"""Integral generator."""

import collections
import logging
from numbers import Integral
from typing import Any

import ufl

import ffcx.codegeneration.lnodes as L
from ffcx.codegeneration import geometry
from ffcx.codegeneration.definitions import create_dof_index, create_quadrature_index
from ffcx.codegeneration.optimizer import optimize
from ffcx.ir.elementtables import piecewise_ttypes
from ffcx.ir.integral import BlockDataT
from ffcx.ir.representationutils import QuadratureRule

logger = logging.getLogger("ffcx")


def extract_dtype(v, vops: list[Any]):
    """Extract dtype from ufl expression v and its operands."""
    dtypes = []
    for op in vops:
        if hasattr(op, "dtype"):
            dtypes.append(op.dtype)
        elif hasattr(op, "symbol"):
            dtypes.append(op.symbol.dtype)
        elif isinstance(op, Integral):
            dtypes.append(L.DataType.INT)
        else:
            raise RuntimeError(f"Not expecting this type of operand {type(op)}")
    is_cond = isinstance(v, ufl.classes.Condition)
    return L.DataType.BOOL if is_cond else L.merge_dtypes(dtypes)


class IntegralGenerator:
    """Integral generator."""

    def __init__(self, ir, backend):
        """Initialise."""
        # Store ir
        self.ir = ir

        # Backend specific plugin with attributes
        # - symbols: for translating ufl operators to target language
        # - definitions: for defining backend specific variables
        # - access: for accessing backend specific variables
        self.backend = backend

        # Set of operator names code has been generated for, used in the
        # end for selecting necessary includes
        self._ufl_names = set()

        # Initialize lookup tables for variable scopes
        self.init_scopes()

        # Cache
        self.temp_symbols = {}

        # Set of counters used for assigning names to intermediate
        # variables
        self.symbol_counters = collections.defaultdict(int)

    def init_scopes(self):
        """Initialize variable scope dicts."""
        # Reset variables, separate sets for each quadrature rule
        self.scopes = {
            quadrature_rule: {} for quadrature_rule in self.ir.expression.integrand.keys()
        }
        self.scopes[None] = {}

    def set_var(self, quadrature_rule, v, vaccess):
        """Set a new variable in variable scope dicts.

        Scope is determined by quadrature_rule which identifies the
        quadrature loop scope or None if outside quadrature loops.

        Args:
            quadrature_rule: Quadrature rule
            v: the ufl expression
            vaccess: the LNodes expression to access the value in the code
        """
        self.scopes[quadrature_rule][v] = vaccess

    def get_var(self, quadrature_rule, v):
        """Lookup ufl expression v in variable scope dicts.

        Scope is determined by quadrature rule which identifies the
        quadrature loop scope or None if outside quadrature loops.

        If v is not found in quadrature loop scope, the piecewise
        scope (None) is checked.

        Returns the LNodes expression to access the value in the code.
        """
        if v._ufl_is_literal_:
            return L.ufl_to_lnodes(v)

        # quadrature loop scope
        f = self.scopes[quadrature_rule].get(v)

        # piecewise scope
        if f is None:
            f = self.scopes[None].get(v)
        return f

    def new_temp_symbol(self, basename):
        """Create a new code symbol named basename + running counter."""
        name = "%s%d" % (basename, self.symbol_counters[basename])
        self.symbol_counters[basename] += 1
        return L.Symbol(name, dtype=L.DataType.SCALAR)

    def get_temp_symbol(self, tempname, key):
        """Get a temporary symbol."""
        key = (tempname,) + key
        s = self.temp_symbols.get(key)
        defined = s is not None
        if not defined:
            s = self.new_temp_symbol(tempname)
            self.temp_symbols[key] = s
        return s, defined

    def generate(self):
        """Generate entire tabulate_tensor body.

        Assumes that the code returned from here will be wrapped in a
        context that matches a suitable version of the UFC
        tabulate_tensor signatures.
        """
        # Assert that scopes are empty: expecting this to be called only
        # once
        assert not any(d for d in self.scopes.values())

        parts = []

        # Generate the tables of quadrature points and weights
        parts += self.generate_quadrature_tables()

        # Generate the tables of basis function values and
        # pre-integrated blocks
        parts += self.generate_element_tables()

        # Generate the tables of geometry data that are needed
        parts += self.generate_geometry_tables()

        # Loop generation code will produce parts to go before
        # quadloops, to define the quadloops, and to go after the
        # quadloops
        all_preparts = []
        all_quadparts = []

        # Pre-definitions are collected across all quadrature loops to
        # improve re-use and avoid name clashes
        for rule in self.ir.expression.integrand.keys():
            # Generate code to compute piecewise constant scalar factors
            all_preparts += self.generate_piecewise_partition(rule)

            # Generate code to integrate reusable blocks of final
            # element tensor
            all_quadparts += self.generate_quadrature_loop(rule)

        # Collect parts before, during, and after quadrature loops
        parts += all_preparts
        parts += all_quadparts

        return L.StatementList(parts)

    def generate_quadrature_tables(self):
        """Generate static tables of quadrature points and weights."""
        parts = []

        # No quadrature tables for custom (given argument) or point
        # (evaluation in single vertex)
        skip = ufl.custom_integral_types + ufl.measure.point_integral_types
        if self.ir.expression.integral_type in skip:
            return parts

        # Loop over quadrature rules
        for quadrature_rule, _ in self.ir.expression.integrand.items():
            # Generate quadrature weights array
            wsym = self.backend.symbols.weights_table(quadrature_rule)
            parts += [L.ArrayDecl(wsym, values=quadrature_rule.weights, const=True)]

        # Add leading comment if there are any tables
        parts = L.commented_code_list(parts, "Quadrature rules")
        return parts

    def generate_geometry_tables(self):
        """Generate static tables of geometry data."""
        ufl_geometry = {
            ufl.geometry.FacetEdgeVectors: "facet_edge_vertices",
            ufl.geometry.CellFacetJacobian: "reference_facet_jacobian",
            ufl.geometry.ReferenceCellVolume: "reference_cell_volume",
            ufl.geometry.ReferenceFacetVolume: "reference_facet_volume",
            ufl.geometry.ReferenceCellEdgeVectors: "reference_edge_vectors",
            ufl.geometry.ReferenceFacetEdgeVectors: "facet_reference_edge_vectors",
            ufl.geometry.ReferenceNormal: "reference_facet_normals",
            ufl.geometry.FacetOrientation: "facet_orientation",
        }
        cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()}  # type: ignore

        for integrand in self.ir.expression.integrand.values():
            for attr in integrand["factorization"].nodes.values():
                mt = attr.get("mt")
                if mt is not None:
                    t = type(mt.terminal)
                    if t in ufl_geometry:
                        cells[t].add(
                            ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
                        )

        parts = []
        for i, cell_list in cells.items():
            for c in cell_list:
                parts.append(geometry.write_table(ufl_geometry[i], c))

        return parts

    def generate_element_tables(self):
        """Generate static tables.

        With precomputed element basis function values in quadrature points.
        """
        parts = []
        tables = self.ir.expression.unique_tables
        table_types = self.ir.expression.unique_table_types
        if self.ir.expression.integral_type in ufl.custom_integral_types:
            # Define only piecewise tables
            table_names = [name for name in sorted(tables) if table_types[name] in piecewise_ttypes]
        else:
            # Define all tables
            table_names = sorted(tables)

        for name in table_names:
            table = tables[name]
            parts += self.declare_table(name, table)

        # Add leading comment if there are any tables
        parts = L.commented_code_list(
            parts,
            [
                "Precomputed values of basis functions and precomputations",
                "FE* dimensions: [permutation][entities][points][dofs]",
            ],
        )
        return parts

    def declare_table(self, name, table):
        """Declare a table.

        If the dof dimensions of the table have dof rotations, apply
        these rotations.

        """
        table_symbol = L.Symbol(name, dtype=L.DataType.REAL)
        self.backend.symbols.element_tables[name] = table_symbol
        return [L.ArrayDecl(table_symbol, values=table, const=True)]

    def generate_quadrature_loop(self, quadrature_rule: QuadratureRule):
        """Generate quadrature loop with for this quadrature_rule."""
        # Generate varying partition
        definitions, intermediates_0 = self.generate_varying_partition(quadrature_rule)

        # Generate dofblock parts, some of this will be placed before or after quadloop
        tensor_comp, intermediates_fw = self.generate_dofblock_partition(quadrature_rule)
        assert all([isinstance(tc, L.Section) for tc in tensor_comp])

        # Check if we only have Section objects
        inputs = []
        for definition in definitions:
            assert isinstance(definition, L.Section)
            inputs += definition.output

        # Create intermediates section
        output = []
        declarations = []
        for fw in intermediates_fw:
            assert isinstance(fw, L.VariableDecl)
            output += [fw.symbol]
            declarations += [L.VariableDecl(fw.symbol, 0)]
            intermediates_0 += [L.Assign(fw.symbol, fw.value)]
        intermediates = [L.Section("Intermediates", intermediates_0, declarations, inputs, output)]

        iq_symbol = self.backend.symbols.quadrature_loop_index
        iq = create_quadrature_index(quadrature_rule, iq_symbol)

        code = definitions + intermediates + tensor_comp
        code = optimize(code, quadrature_rule)

        return [L.create_nested_for_loops([iq], code)]

    def generate_piecewise_partition(self, quadrature_rule):
        """Generate a piecewise partition."""
        # Get annotated graph of factorisation
        F = self.ir.expression.integrand[quadrature_rule]["factorization"]
        arraysymbol = L.Symbol(f"sp_{quadrature_rule.id()}", dtype=L.DataType.SCALAR)
        return self.generate_partition(arraysymbol, F, "piecewise", None)

    def generate_varying_partition(self, quadrature_rule):
        """Generate a varying partition."""
        # Get annotated graph of factorisation
        F = self.ir.expression.integrand[quadrature_rule]["factorization"]
        arraysymbol = L.Symbol(f"sv_{quadrature_rule.id()}", dtype=L.DataType.SCALAR)
        return self.generate_partition(arraysymbol, F, "varying", quadrature_rule)

    def generate_partition(self, symbol, F, mode, quadrature_rule):
        """Generate a partition."""
        definitions = []
        intermediates = []

        for i, attr in F.nodes.items():
            if attr["status"] != mode:
                continue
            v = attr["expression"]

            # Generate code only if the expression is not already in cache
            if not self.get_var(quadrature_rule, v):
                if v._ufl_is_literal_:
                    vaccess = L.ufl_to_lnodes(v)
                elif mt := attr.get("mt"):
                    tabledata = attr.get("tr")

                    # Backend specific modified terminal translation
                    vaccess = self.backend.access.get(mt, tabledata, quadrature_rule)
                    vdef = self.backend.definitions.get(mt, tabledata, quadrature_rule, vaccess)

                    if vdef:
                        assert isinstance(vdef, L.Section)
                    # Only add if definition is unique.
                    # This can happen when using sub-meshes
                    if vdef not in definitions:
                        definitions += [vdef]
                else:
                    # Get previously visited operands
                    vops = [self.get_var(quadrature_rule, op) for op in v.ufl_operands]
                    dtype = extract_dtype(v, vops)

                    # Mapping UFL operator to target language
                    self._ufl_names.add(v._ufl_handler_name_)
                    vexpr = L.ufl_to_lnodes(v, *vops)

                    j = len(intermediates)
                    vaccess = L.Symbol(f"{symbol.name}_{j}", dtype=dtype)
                    intermediates.append(L.VariableDecl(vaccess, vexpr))

                # Store access node for future reference
                self.set_var(quadrature_rule, v, vaccess)

        # Optimize definitions
        definitions = optimize(definitions, quadrature_rule)
        return definitions, intermediates

    def generate_dofblock_partition(self, quadrature_rule: QuadratureRule):
        """Generate a dofblock partition."""
        block_contributions = self.ir.expression.integrand[quadrature_rule]["block_contributions"]
        quadparts = []
        blocks = [
            (blockmap, blockdata)
            for blockmap, contributions in sorted(block_contributions.items())
            for blockdata in contributions
        ]

        block_groups = collections.defaultdict(list)

        # Group loops by blockmap, in Vector elements each component has
        # a different blockmap
        for blockmap, blockdata in blocks:
            scalar_blockmap = []
            assert len(blockdata.ma_data) == len(blockmap)
            for i, b in enumerate(blockmap):
                bs = blockdata.ma_data[i].tabledata.block_size
                offset = blockdata.ma_data[i].tabledata.offset
                b = tuple([(idx - offset) // bs for idx in b])
                scalar_blockmap.append(b)
            block_groups[tuple(scalar_blockmap)].append(blockdata)

        intermediates = []
        for blockmap in block_groups:
            block_quadparts, intermediate = self.generate_block_parts(
                quadrature_rule, blockmap, block_groups[blockmap]
            )
            intermediates += intermediate

            # Add computations
            quadparts.extend(block_quadparts)

        return quadparts, intermediates

    def get_arg_factors(self, blockdata, block_rank, quadrature_rule, iq, indices):
        """Get arg factors."""
        arg_factors = []
        tables = []
        for i in range(block_rank):
            mad = blockdata.ma_data[i]
            td = mad.tabledata
            scope = self.ir.expression.integrand[quadrature_rule]["modified_arguments"]
            mt = scope[mad.ma_index]
            arg_tables = []

            # Translate modified terminal to code
            # TODO: Move element table access out of backend?
            #       Not using self.backend.access.argument() here
            #       now because it assumes too much about indices.

            assert td.ttype != "zeros"

            if td.ttype == "ones":
                arg_factor = 1
            else:
                # Assuming B sparsity follows element table sparsity
                arg_factor, arg_tables = self.backend.access.table_access(
                    td, self.ir.expression.entity_type, mt.restriction, iq, indices[i]
                )

            tables += arg_tables
            arg_factors.append(arg_factor)

        return arg_factors, tables

    def generate_block_parts(
        self, quadrature_rule: QuadratureRule, blockmap: tuple, blocklist: list[BlockDataT]
    ):
        """Generate and return code parts for a given block.

        Returns parts occurring before, inside, and after the quadrature
        loop identified by the quadrature rule.

        Should be called with quadrature_rule=None for
        quadloop-independent blocks.
        """
        # The parts to return
        quadparts: list[L.LNode] = []
        intermediates: list[L.LNode] = []
        tables = []
        vars = []

        # RHS expressions grouped by LHS "dofmap"
        rhs_expressions = collections.defaultdict(list)

        block_rank = len(blockmap)
        iq_symbol = self.backend.symbols.quadrature_loop_index
        iq = create_quadrature_index(quadrature_rule, iq_symbol)

        for blockdata in blocklist:
            B_indices = []
            for i in range(block_rank):
                table_ref = blockdata.ma_data[i].tabledata
                symbol = self.backend.symbols.argument_loop_index(i)
                index = create_dof_index(table_ref, symbol)
                B_indices.append(index)

            ttypes = blockdata.ttypes
            if "zeros" in ttypes:
                raise RuntimeError(
                    "Not expecting zero arguments to be left in dofblock generation."
                )

            if len(blockdata.factor_indices_comp_indices) > 1:
                raise RuntimeError("Code generation for non-scalar integrals unsupported")

            # We have scalar integrand here, take just the factor index
            factor_index = blockdata.factor_indices_comp_indices[0][0]

            # Get factor expression
            F = self.ir.expression.integrand[quadrature_rule]["factorization"]

            v = F.nodes[factor_index]["expression"]
            f = self.get_var(quadrature_rule, v)

            # Quadrature weight was removed in representation, add it back now
            if self.ir.expression.integral_type in ufl.custom_integral_types:
                weights = self.backend.symbols.custom_weights_table
                weight = weights[iq.global_index]
            else:
                weights = self.backend.symbols.weights_table(quadrature_rule)
                weight = weights[iq.global_index]

            # Define fw = f * weight
            fw_rhs = L.float_product([f, weight])
            if not isinstance(fw_rhs, L.Product):
                fw = fw_rhs
            else:
                # Define and cache scalar temp variable
                key = (quadrature_rule, factor_index, blockdata.all_factors_piecewise)
                fw, defined = self.get_temp_symbol("fw", key)
                if not defined:
                    input = [f, weight]
                    # filter only L.Symbol in input
                    input = [i for i in input if isinstance(i, L.Symbol)]
                    output = [fw]

                    # assert input and output are Symbol objects
                    assert all(isinstance(i, L.Symbol) for i in input)
                    assert all(isinstance(o, L.Symbol) for o in output)

                    intermediates += [L.VariableDecl(fw, fw_rhs)]

            var = fw if isinstance(fw, L.Symbol) else fw.array
            vars += [var]
            assert not blockdata.transposed, "Not handled yet"

            # Fetch code to access modified arguments
            arg_factors, table = self.get_arg_factors(
                blockdata, block_rank, quadrature_rule, iq, B_indices
            )
            tables += table

            # Define B_rhs = fw * arg_factors
            B_rhs = L.float_product([fw] + arg_factors)

            A_indices = []
            for i in range(block_rank):
                index = B_indices[i]
                tabledata = blockdata.ma_data[i].tabledata
                offset = tabledata.offset
                if len(blockmap[i]) == 1:
                    A_indices.append(index.global_index + offset)
                else:
                    block_size = blockdata.ma_data[i].tabledata.block_size
                    A_indices.append(block_size * index.global_index + offset)
            rhs_expressions[tuple(A_indices)].append(B_rhs)

        # List of statements to keep in the inner loop
        keep = collections.defaultdict(list)

        for indices in rhs_expressions:
            keep[indices] = rhs_expressions[indices]

        body: list[L.LNode] = []

        A = self.backend.symbols.element_tensor
        A_shape = self.ir.expression.tensor_shape
        for indices in keep:
            multi_index = L.MultiIndex(list(indices), A_shape)
            for expression in keep[indices]:
                body.append(L.AssignAdd(A[multi_index], expression))

        # reverse B_indices
        B_indices = B_indices[::-1]
        body = [L.create_nested_for_loops(B_indices, body)]
        input = [*vars, *tables]
        output = [A]

        # Make sure we don't have repeated symbols in input
        input = list(set(input))

        # assert input and output are Symbol objects
        assert all(isinstance(i, L.Symbol) for i in input)
        assert all(isinstance(o, L.Symbol) for o in output)

        annotations = []
        if len(B_indices) > 1:
            annotations.append(L.Annotation.licm)

        quadparts += [L.Section("Tensor Computation", body, [], input, output, annotations)]

        return quadparts, intermediates

1	# Copyright (C) 2015-2024 Martin Sandve Alnæs, Michal Habera, Igor Baratta, Chris Richardson
2	#
3	# Modified by Jørgen S. Dokken, 2024
4	#
5	# This file is part of FFCx. (https://www.fenicsproject.org)
6	#
7	# SPDX-License-Identifier: LGPL-3.0-or-later
8	"""Integral generator."""
9
10	import collections	1✔
11	import logging	1✔
12	from numbers import Integral	1✔
13	from typing import Any	1✔
14
15	import ufl	1✔
16
17	import ffcx.codegeneration.lnodes as L	1✔
18	from ffcx.codegeneration import geometry	1✔
19	from ffcx.codegeneration.definitions import create_dof_index, create_quadrature_index	1✔
20	from ffcx.codegeneration.optimizer import optimize	1✔
21	from ffcx.ir.elementtables import piecewise_ttypes	1✔
22	from ffcx.ir.integral import BlockDataT	1✔
23	from ffcx.ir.representationutils import QuadratureRule	1✔
24
25	logger = logging.getLogger("ffcx")	1✔
26
27
28	def extract_dtype(v, vops: list[Any]):	1✔
29	"""Extract dtype from ufl expression v and its operands."""
30	dtypes = []	1✔
31	for op in vops:	1✔
32	if hasattr(op, "dtype"):	1✔
33	dtypes.append(op.dtype)	1✔
34	elif hasattr(op, "symbol"):	1✔
35	dtypes.append(op.symbol.dtype)	×
36	elif isinstance(op, Integral):	1✔
37	dtypes.append(L.DataType.INT)	1✔
38	else:
39	raise RuntimeError(f"Not expecting this type of operand {type(op)}")	×
40	is_cond = isinstance(v, ufl.classes.Condition)	1✔
41	return L.DataType.BOOL if is_cond else L.merge_dtypes(dtypes)	1✔
42
43
44	class IntegralGenerator:	1✔
45	"""Integral generator."""
46
47	def __init__(self, ir, backend):	1✔
48	"""Initialise."""
49	# Store ir
50	self.ir = ir	1✔
51
52	# Backend specific plugin with attributes
53	# - symbols: for translating ufl operators to target language
54	# - definitions: for defining backend specific variables
55	# - access: for accessing backend specific variables
56	self.backend = backend	1✔
57
58	# Set of operator names code has been generated for, used in the
59	# end for selecting necessary includes
60	self._ufl_names = set()	1✔
61
62	# Initialize lookup tables for variable scopes
63	self.init_scopes()	1✔
64
65	# Cache
66	self.temp_symbols = {}	1✔
67
68	# Set of counters used for assigning names to intermediate
69	# variables
70	self.symbol_counters = collections.defaultdict(int)	1✔
71
72	def init_scopes(self):	1✔
73	"""Initialize variable scope dicts."""
74	# Reset variables, separate sets for each quadrature rule
75	self.scopes = {	1✔
76	quadrature_rule: {} for quadrature_rule in self.ir.expression.integrand.keys()
77	}
78	self.scopes[None] = {}	1✔
79
80	def set_var(self, quadrature_rule, v, vaccess):	1✔
81	"""Set a new variable in variable scope dicts.
82
83	Scope is determined by quadrature_rule which identifies the
84	quadrature loop scope or None if outside quadrature loops.
85
86	Args:
87	quadrature_rule: Quadrature rule
88	v: the ufl expression
89	vaccess: the LNodes expression to access the value in the code
90	"""
91	self.scopes[quadrature_rule][v] = vaccess	1✔
92
93	def get_var(self, quadrature_rule, v):	1✔
94	"""Lookup ufl expression v in variable scope dicts.
95
96	Scope is determined by quadrature rule which identifies the
97	quadrature loop scope or None if outside quadrature loops.
98
99	If v is not found in quadrature loop scope, the piecewise
100	scope (None) is checked.
101
102	Returns the LNodes expression to access the value in the code.
103	"""
104	if v._ufl_is_literal_:	1✔
105	return L.ufl_to_lnodes(v)	1✔
106
107	# quadrature loop scope
108	f = self.scopes[quadrature_rule].get(v)	1✔
109
110	# piecewise scope
111	if f is None:	1✔
112	f = self.scopes[None].get(v)	1✔
113	return f	1✔
114
115	def new_temp_symbol(self, basename):	1✔
116	"""Create a new code symbol named basename + running counter."""
117	name = "%s%d" % (basename, self.symbol_counters[basename])	1✔
118	self.symbol_counters[basename] += 1	1✔
119	return L.Symbol(name, dtype=L.DataType.SCALAR)	1✔
120
121	def get_temp_symbol(self, tempname, key):	1✔
122	"""Get a temporary symbol."""
123	key = (tempname,) + key	1✔
124	s = self.temp_symbols.get(key)	1✔
125	defined = s is not None	1✔
126	if not defined:	1✔
127	s = self.new_temp_symbol(tempname)	1✔
128	self.temp_symbols[key] = s	1✔
129	return s, defined	1✔
130
131	def generate(self):	1✔
132	"""Generate entire tabulate_tensor body.
133
134	Assumes that the code returned from here will be wrapped in a
135	context that matches a suitable version of the UFC
136	tabulate_tensor signatures.
137	"""
138	# Assert that scopes are empty: expecting this to be called only
139	# once
140	assert not any(d for d in self.scopes.values())	1✔
141
142	parts = []	1✔
143
144	# Generate the tables of quadrature points and weights
145	parts += self.generate_quadrature_tables()	1✔
146
147	# Generate the tables of basis function values and
148	# pre-integrated blocks
149	parts += self.generate_element_tables()	1✔
150
151	# Generate the tables of geometry data that are needed
152	parts += self.generate_geometry_tables()	1✔
153
154	# Loop generation code will produce parts to go before
155	# quadloops, to define the quadloops, and to go after the
156	# quadloops
157	all_preparts = []	1✔
158	all_quadparts = []	1✔
159
160	# Pre-definitions are collected across all quadrature loops to
161	# improve re-use and avoid name clashes
162	for rule in self.ir.expression.integrand.keys():	1✔
163	# Generate code to compute piecewise constant scalar factors
164	all_preparts += self.generate_piecewise_partition(rule)	1✔
165
166	# Generate code to integrate reusable blocks of final
167	# element tensor
168	all_quadparts += self.generate_quadrature_loop(rule)	1✔
169
170	# Collect parts before, during, and after quadrature loops
171	parts += all_preparts	1✔
172	parts += all_quadparts	1✔
173
174	return L.StatementList(parts)	1✔
175
176	def generate_quadrature_tables(self):	1✔
177	"""Generate static tables of quadrature points and weights."""
178	parts = []	1✔
179
180	# No quadrature tables for custom (given argument) or point
181	# (evaluation in single vertex)
182	skip = ufl.custom_integral_types + ufl.measure.point_integral_types	1✔
183	if self.ir.expression.integral_type in skip:	1✔
184	return parts	×
185
186	# Loop over quadrature rules
187	for quadrature_rule, _ in self.ir.expression.integrand.items():	1✔
188	# Generate quadrature weights array
189	wsym = self.backend.symbols.weights_table(quadrature_rule)	1✔
190	parts += [L.ArrayDecl(wsym, values=quadrature_rule.weights, const=True)]	1✔
191
192	# Add leading comment if there are any tables
193	parts = L.commented_code_list(parts, "Quadrature rules")	1✔
194	return parts	1✔
195
196	def generate_geometry_tables(self):	1✔
197	"""Generate static tables of geometry data."""
198	ufl_geometry = {	1✔
199	ufl.geometry.FacetEdgeVectors: "facet_edge_vertices",
200	ufl.geometry.CellFacetJacobian: "reference_facet_jacobian",
201	ufl.geometry.ReferenceCellVolume: "reference_cell_volume",
202	ufl.geometry.ReferenceFacetVolume: "reference_facet_volume",
203	ufl.geometry.ReferenceCellEdgeVectors: "reference_edge_vectors",
204	ufl.geometry.ReferenceFacetEdgeVectors: "facet_reference_edge_vectors",
205	ufl.geometry.ReferenceNormal: "reference_facet_normals",
206	ufl.geometry.FacetOrientation: "facet_orientation",
207	}
208	cells: dict[Any, set[Any]] = {t: set() for t in ufl_geometry.keys()} # type: ignore	1✔
209
210	for integrand in self.ir.expression.integrand.values():	1✔
211	for attr in integrand["factorization"].nodes.values():	1✔
212	mt = attr.get("mt")	1✔
213	if mt is not None:	1✔
214	t = type(mt.terminal)	1✔
215	if t in ufl_geometry:	1✔
216	cells[t].add(	1✔
217	ufl.domain.extract_unique_domain(mt.terminal).ufl_cell().cellname()
218	)
219
220	parts = []	1✔
221	for i, cell_list in cells.items():	1✔
222	for c in cell_list:	1✔
223	parts.append(geometry.write_table(ufl_geometry[i], c))	1✔
224
225	return parts	1✔
226
227	def generate_element_tables(self):	1✔
228	"""Generate static tables.
229
230	With precomputed element basis function values in quadrature points.
231	"""
232	parts = []	1✔
233	tables = self.ir.expression.unique_tables	1✔
234	table_types = self.ir.expression.unique_table_types	1✔
235	if self.ir.expression.integral_type in ufl.custom_integral_types:	1✔
236	# Define only piecewise tables
237	table_names = [name for name in sorted(tables) if table_types[name] in piecewise_ttypes]	×
238	else:
239	# Define all tables
240	table_names = sorted(tables)	1✔
241
242	for name in table_names:	1✔
243	table = tables[name]	1✔
244	parts += self.declare_table(name, table)	1✔
245
246	# Add leading comment if there are any tables
247	parts = L.commented_code_list(	1✔
248	parts,
249	[
250	"Precomputed values of basis functions and precomputations",
251	"FE* dimensions: [permutation][entities][points][dofs]",
252	],
253	)
254	return parts	1✔
255
256	def declare_table(self, name, table):	1✔
257	"""Declare a table.
258
259	If the dof dimensions of the table have dof rotations, apply
260	these rotations.
261
262	"""
263	table_symbol = L.Symbol(name, dtype=L.DataType.REAL)	1✔
264	self.backend.symbols.element_tables[name] = table_symbol	1✔
265	return [L.ArrayDecl(table_symbol, values=table, const=True)]	1✔
266
267	def generate_quadrature_loop(self, quadrature_rule: QuadratureRule):	1✔
268	"""Generate quadrature loop with for this quadrature_rule."""
269	# Generate varying partition
270	definitions, intermediates_0 = self.generate_varying_partition(quadrature_rule)	1✔
271
272	# Generate dofblock parts, some of this will be placed before or after quadloop
273	tensor_comp, intermediates_fw = self.generate_dofblock_partition(quadrature_rule)	1✔
274	assert all([isinstance(tc, L.Section) for tc in tensor_comp])	1✔
275
276	# Check if we only have Section objects
277	inputs = []	1✔
278	for definition in definitions:	1✔
279	assert isinstance(definition, L.Section)	1✔
280	inputs += definition.output	1✔
281
282	# Create intermediates section
283	output = []	1✔
284	declarations = []	1✔
285	for fw in intermediates_fw:	1✔
286	assert isinstance(fw, L.VariableDecl)	1✔
287	output += [fw.symbol]	1✔
288	declarations += [L.VariableDecl(fw.symbol, 0)]	1✔
289	intermediates_0 += [L.Assign(fw.symbol, fw.value)]	1✔
290	intermediates = [L.Section("Intermediates", intermediates_0, declarations, inputs, output)]	1✔
291
292	iq_symbol = self.backend.symbols.quadrature_loop_index	1✔
293	iq = create_quadrature_index(quadrature_rule, iq_symbol)	1✔
294
295	code = definitions + intermediates + tensor_comp	1✔
296	code = optimize(code, quadrature_rule)	1✔
297
298	return [L.create_nested_for_loops([iq], code)]	1✔
299
300	def generate_piecewise_partition(self, quadrature_rule):	1✔
301	"""Generate a piecewise partition."""
302	# Get annotated graph of factorisation
303	F = self.ir.expression.integrand[quadrature_rule]["factorization"]	1✔
304	arraysymbol = L.Symbol(f"sp_{quadrature_rule.id()}", dtype=L.DataType.SCALAR)	1✔
305	return self.generate_partition(arraysymbol, F, "piecewise", None)	1✔
306
307	def generate_varying_partition(self, quadrature_rule):	1✔
308	"""Generate a varying partition."""
309	# Get annotated graph of factorisation
310	F = self.ir.expression.integrand[quadrature_rule]["factorization"]	1✔
311	arraysymbol = L.Symbol(f"sv_{quadrature_rule.id()}", dtype=L.DataType.SCALAR)	1✔
312	return self.generate_partition(arraysymbol, F, "varying", quadrature_rule)	1✔
313
314	def generate_partition(self, symbol, F, mode, quadrature_rule):	1✔
315	"""Generate a partition."""
316	definitions = []	1✔
317	intermediates = []	1✔
318
319	for i, attr in F.nodes.items():	1✔
320	if attr["status"] != mode:	1✔
321	continue	1✔
322	v = attr["expression"]	1✔
323
324	# Generate code only if the expression is not already in cache
325	if not self.get_var(quadrature_rule, v):	1✔
326	if v._ufl_is_literal_:	1✔
327	vaccess = L.ufl_to_lnodes(v)	×
328	elif mt := attr.get("mt"):	1✔
329	tabledata = attr.get("tr")	1✔
330
331	# Backend specific modified terminal translation
332	vaccess = self.backend.access.get(mt, tabledata, quadrature_rule)	1✔
333	vdef = self.backend.definitions.get(mt, tabledata, quadrature_rule, vaccess)	1✔
334
335	if vdef:	1✔
336	assert isinstance(vdef, L.Section)	1✔
337	# Only add if definition is unique.
338	# This can happen when using sub-meshes
339	if vdef not in definitions:	1✔
340	definitions += [vdef]	1✔
341	else:
342	# Get previously visited operands
343	vops = [self.get_var(quadrature_rule, op) for op in v.ufl_operands]	1✔
344	dtype = extract_dtype(v, vops)	1✔
345
346	# Mapping UFL operator to target language
347	self._ufl_names.add(v._ufl_handler_name_)	1✔
348	vexpr = L.ufl_to_lnodes(v, *vops)	1✔
349
350	j = len(intermediates)	1✔
351	vaccess = L.Symbol(f"{symbol.name}_{j}", dtype=dtype)	1✔
352	intermediates.append(L.VariableDecl(vaccess, vexpr))	1✔
353
354	# Store access node for future reference
355	self.set_var(quadrature_rule, v, vaccess)	1✔
356
357	# Optimize definitions
358	definitions = optimize(definitions, quadrature_rule)	1✔
359	return definitions, intermediates	1✔
360
361	def generate_dofblock_partition(self, quadrature_rule: QuadratureRule):	1✔
362	"""Generate a dofblock partition."""
363	block_contributions = self.ir.expression.integrand[quadrature_rule]["block_contributions"]	1✔
364	quadparts = []	1✔
365	blocks = [	1✔
366	(blockmap, blockdata)
367	for blockmap, contributions in sorted(block_contributions.items())
368	for blockdata in contributions
369	]
370
371	block_groups = collections.defaultdict(list)	1✔
372
373	# Group loops by blockmap, in Vector elements each component has
374	# a different blockmap
375	for blockmap, blockdata in blocks:	1✔
376	scalar_blockmap = []	1✔
377	assert len(blockdata.ma_data) == len(blockmap)	1✔
378	for i, b in enumerate(blockmap):	1✔
379	bs = blockdata.ma_data[i].tabledata.block_size	1✔
380	offset = blockdata.ma_data[i].tabledata.offset	1✔
381	b = tuple([(idx - offset) // bs for idx in b])	1✔
382	scalar_blockmap.append(b)	1✔
383	block_groups[tuple(scalar_blockmap)].append(blockdata)	1✔
384
385	intermediates = []	1✔
386	for blockmap in block_groups:	1✔
387	block_quadparts, intermediate = self.generate_block_parts(	1✔
388	quadrature_rule, blockmap, block_groups[blockmap]
389	)
390	intermediates += intermediate	1✔
391
392	# Add computations
393	quadparts.extend(block_quadparts)	1✔
394
395	return quadparts, intermediates	1✔
396
397	def get_arg_factors(self, blockdata, block_rank, quadrature_rule, iq, indices):	1✔
398	"""Get arg factors."""
399	arg_factors = []	1✔
400	tables = []	1✔
401	for i in range(block_rank):	1✔
402	mad = blockdata.ma_data[i]	1✔
403	td = mad.tabledata	1✔
404	scope = self.ir.expression.integrand[quadrature_rule]["modified_arguments"]	1✔
405	mt = scope[mad.ma_index]	1✔
406	arg_tables = []	1✔
407
408	# Translate modified terminal to code
409	# TODO: Move element table access out of backend?
410	# Not using self.backend.access.argument() here
411	# now because it assumes too much about indices.
412
413	assert td.ttype != "zeros"	1✔
414
415	if td.ttype == "ones":	1✔
416	arg_factor = 1	×
417	else:
418	# Assuming B sparsity follows element table sparsity
419	arg_factor, arg_tables = self.backend.access.table_access(	1✔
420	td, self.ir.expression.entity_type, mt.restriction, iq, indices[i]
421	)
422
423	tables += arg_tables	1✔
424	arg_factors.append(arg_factor)	1✔
425
426	return arg_factors, tables	1✔
427
428	def generate_block_parts(	1✔
429	self, quadrature_rule: QuadratureRule, blockmap: tuple, blocklist: list[BlockDataT]
430	):
431	"""Generate and return code parts for a given block.
432
433	Returns parts occurring before, inside, and after the quadrature
434	loop identified by the quadrature rule.
435
436	Should be called with quadrature_rule=None for
437	quadloop-independent blocks.
438	"""
439	# The parts to return
440	quadparts: list[L.LNode] = []	1✔
441	intermediates: list[L.LNode] = []	1✔
442	tables = []	1✔
443	vars = []	1✔
444
445	# RHS expressions grouped by LHS "dofmap"
446	rhs_expressions = collections.defaultdict(list)	1✔
447
448	block_rank = len(blockmap)	1✔
449	iq_symbol = self.backend.symbols.quadrature_loop_index	1✔
450	iq = create_quadrature_index(quadrature_rule, iq_symbol)	1✔
451
452	for blockdata in blocklist:	1✔
453	B_indices = []	1✔
454	for i in range(block_rank):	1✔
455	table_ref = blockdata.ma_data[i].tabledata	1✔
456	symbol = self.backend.symbols.argument_loop_index(i)	1✔
457	index = create_dof_index(table_ref, symbol)	1✔
458	B_indices.append(index)	1✔
459
460	ttypes = blockdata.ttypes	1✔
461	if "zeros" in ttypes:	1✔
462	raise RuntimeError(	×
463	"Not expecting zero arguments to be left in dofblock generation."
464	)
465
466	if len(blockdata.factor_indices_comp_indices) > 1:	1✔
467	raise RuntimeError("Code generation for non-scalar integrals unsupported")	×
468
469	# We have scalar integrand here, take just the factor index
470	factor_index = blockdata.factor_indices_comp_indices[0][0]	1✔
471
472	# Get factor expression
473	F = self.ir.expression.integrand[quadrature_rule]["factorization"]	1✔
474
475	v = F.nodes[factor_index]["expression"]	1✔
476	f = self.get_var(quadrature_rule, v)	1✔
477
478	# Quadrature weight was removed in representation, add it back now
479	if self.ir.expression.integral_type in ufl.custom_integral_types:	1✔
480	weights = self.backend.symbols.custom_weights_table	×
481	weight = weights[iq.global_index]	×
482	else:
483	weights = self.backend.symbols.weights_table(quadrature_rule)	1✔
484	weight = weights[iq.global_index]	1✔
485
486	# Define fw = f * weight
487	fw_rhs = L.float_product([f, weight])	1✔
488	if not isinstance(fw_rhs, L.Product):	1✔
489	fw = fw_rhs	×
490	else:
491	# Define and cache scalar temp variable
492	key = (quadrature_rule, factor_index, blockdata.all_factors_piecewise)	1✔
493	fw, defined = self.get_temp_symbol("fw", key)	1✔
494	if not defined:	1✔
495	input = [f, weight]	1✔
496	# filter only L.Symbol in input
497	input = [i for i in input if isinstance(i, L.Symbol)]	1✔
498	output = [fw]	1✔
499
500	# assert input and output are Symbol objects
501	assert all(isinstance(i, L.Symbol) for i in input)	1✔
502	assert all(isinstance(o, L.Symbol) for o in output)	1✔
503
504	intermediates += [L.VariableDecl(fw, fw_rhs)]	1✔
505
506	var = fw if isinstance(fw, L.Symbol) else fw.array	1✔
507	vars += [var]	1✔
508	assert not blockdata.transposed, "Not handled yet"	1✔
509
510	# Fetch code to access modified arguments
511	arg_factors, table = self.get_arg_factors(	1✔
512	blockdata, block_rank, quadrature_rule, iq, B_indices
513	)
514	tables += table	1✔
515
516	# Define B_rhs = fw * arg_factors
517	B_rhs = L.float_product([fw] + arg_factors)	1✔
518
519	A_indices = []	1✔
520	for i in range(block_rank):	1✔
521	index = B_indices[i]	1✔
522	tabledata = blockdata.ma_data[i].tabledata	1✔
523	offset = tabledata.offset	1✔
524	if len(blockmap[i]) == 1:	1✔
525	A_indices.append(index.global_index + offset)	×
526	else:
527	block_size = blockdata.ma_data[i].tabledata.block_size	1✔
528	A_indices.append(block_size * index.global_index + offset)	1✔
529	rhs_expressions[tuple(A_indices)].append(B_rhs)	1✔
530
531	# List of statements to keep in the inner loop
532	keep = collections.defaultdict(list)	1✔
533
534	for indices in rhs_expressions:	1✔
535	keep[indices] = rhs_expressions[indices]	1✔
536
537	body: list[L.LNode] = []	1✔
538
539	A = self.backend.symbols.element_tensor	1✔
540	A_shape = self.ir.expression.tensor_shape	1✔
541	for indices in keep:	1✔
542	multi_index = L.MultiIndex(list(indices), A_shape)	1✔
543	for expression in keep[indices]:	1✔
544	body.append(L.AssignAdd(A[multi_index], expression))	1✔
545
546	# reverse B_indices
547	B_indices = B_indices[::-1]	1✔
548	body = [L.create_nested_for_loops(B_indices, body)]	1✔
549	input = [vars, tables]	1✔
550	output = [A]	1✔
551
552	# Make sure we don't have repeated symbols in input
553	input = list(set(input))	1✔
554
555	# assert input and output are Symbol objects
556	assert all(isinstance(i, L.Symbol) for i in input)	1✔
557	assert all(isinstance(o, L.Symbol) for o in output)	1✔
558
559	annotations = []	1✔
560	if len(B_indices) > 1:	1✔
561	annotations.append(L.Annotation.licm)	1✔
562
563	quadparts += [L.Section("Tensor Computation", body, [], input, output, annotations)]	1✔
564
565	return quadparts, intermediates	1✔

FEniCS / ffcx / 11642291501

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous