Sum factorisation on simplices

tests/tsfc/test_codegen.py

@@ -1,6 +1,7 @@
+import numpy
 import pytest
 
-from gem import impero_utils
+from gem import gem, impero_utils
 from gem.gem import Index, Indexed, IndexSum, Product, Variable
 
 
@@ -24,6 +25,56 @@ def gencode(expr):
     assert len(gencode(e1).children) == len(gencode(e2).children)
 
 
+def test_jagged_index_codegen(monkeypatch):
+    import islpy as isl
+    import loopy as lp
+    import tsfc.loopy
+
+    # Execute the generated code so we check the numbers, not just the loop bounds
+    monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget())
+
+    n = 4
+    extent = n + 1
+    npts = 3
+    ndof = (n + 1) * (n + 2) // 2
+
+    rng = numpy.random.default_rng(7)
+    # Table zero-padded outside the simplex lattice p + q > n, and a
+    # clamped Morton index table for the coefficient gather
+    B = rng.random((extent, extent, npts))
+    morton = numpy.zeros((extent, extent), dtype=gem.uint_type)
+    for p_, q_ in numpy.ndindex(morton.shape):
+        if p_ + q_ > n:
+            B[p_, q_] = 0.0
+        else:
+            morton[p_, q_] = (p_ + q_) * (p_ + q_ + 1) // 2 + q_
+    c = rng.random(ndof)
+
+    i = Index(name="i", extent=npts)
+    p = Index(name="p", extent=extent)
+    q = gem.JaggedIndex(name="q", extent=extent, parents=(p,))
+
+    dof = gem.VariableIndex(Indexed(gem.Literal(morton, dtype=gem.uint_type), (p, q)))
+    integrand = Product(Indexed(Variable("c", (ndof,)), (dof,)),
+                        Indexed(gem.Literal(B), (p, q, i)))
+    expr = IndexSum(integrand, (p, q))
+
+    u = Variable("u", (npts,))
+    impero_c = impero_utils.compile_gem([(Indexed(u, (i,)), expr)], (i, p, q))
+    args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(npts,)),
+            lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))]
+    knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64)
+
+    # The jagged loop must have a domain parametrized by its parent iname
+    assert any(dom.get_var_names(isl.dim_type.param)
+               for dom in knl.default_entrypoint.domains)
+
+    u_out = numpy.zeros(npts)
+    knl(c=c, u=u_out)
+    u_ref = numpy.tensordot(c[morton], B, axes=((0, 1), (0, 1)))
+    assert numpy.allclose(u_out, u_ref, rtol=1e-14)
+
+
 if __name__ == "__main__":
     import os
     import sys

tests/tsfc/test_pickle_gem.py

@@ -17,6 +17,20 @@ def test_pickle_gem(protocol):
     assert repr(expr) == repr(unpickled)
 
 
+@pytest.mark.parametrize('protocol', range(3))
+def test_pickle_jagged_index(protocol):
+    p = gem.Index(name='p', extent=4)
+    q = gem.JaggedIndex(name='q', extent=4, parents=(p,))
+    expr = gem.IndexSum(gem.Indexed(gem.Variable('A', (4, 4)), (p, q)), (p, q))
+
+    unpickled = pickle.loads(pickle.dumps(expr, protocol))
+    assert repr(expr) == repr(unpickled)
+    up, uq = unpickled.multiindex
+    assert isinstance(uq, gem.JaggedIndex)
+    assert uq.extent == 4
+    assert uq.parents == (up,)
+
+
 @pytest.mark.parametrize('protocol', range(3))
 def test_listtensor(protocol):
     expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()])

tsfc/loopy.py

@@ -123,6 +123,7 @@ def __init__(self, target=None):
         self.indices = {}  # indices for declarations and referencing values, from ImperoC
         self.active_indices = {}  # gem index -> pymbolic variable
         self.index_extent = OrderedDict()  # pymbolic variable for indices -> extent
+        self.index_parents = {}  # iname -> parent inames bounding a jagged index
         self.gem_to_pymbolic = {}  # gem node -> pymbolic variable
         self.name_gen = UniqueNameGenerator()
         self.target = target
@@ -257,7 +258,7 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name
     instructions, event_name, preamble = profile_insns(kernel_name, instructions, log)
 
     # Create domains
-    domains = create_domains(ctx.index_extent.items())
+    domains = create_domains(ctx.index_extent.items(), ctx.index_parents)
 
     # Create loopy kernel
     knl = lp.make_kernel(
@@ -276,16 +277,23 @@ def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_name
     return knl, event_name
 
 
-def create_domains(indices):
+def create_domains(indices, index_parents=None):
     """ Create ISL domains from indices
 
     :arg indices: iterable of (index_name, extent) pairs
+    :arg index_parents: optional mapping from index_name to a tuple of parent
+        index names; the domain of a jagged index is parametrized by its
+        parents, with upper bound extent minus the sum of the parents.
     :returns: A list of ISL sets representing the iteration domain of the indices."""
 
     domains = []
     for idx, extent in indices:
-        inames = isl.make_zero_and_vars([idx])
-        domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(inames[0] + extent))))
+        parents = index_parents.get(idx, ()) if index_parents else ()
+        inames = isl.make_zero_and_vars([idx], parents)
+        bound = inames[0] + extent
+        for parent in parents:
+            bound = bound - inames[parent]
+        domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(bound))))
 
     if not domains:
         domains = [isl.BasicSet("[] -> {[]}")]
@@ -316,6 +324,13 @@ def statement_for(tree, ctx):
     assert extent
     idx = ctx.name_gen(ctx.index_names[tree.index])
     ctx.index_extent[idx] = extent
+    if isinstance(tree.index, gem.JaggedIndex) and \
+            all(parent in ctx.active_indices for parent in tree.index.parents):
+        # Tighten the loop bound of a jagged index nested inside its parents.
+        # If a parent loop is not in scope, the rectangular bound `extent`
+        # remains correct: jagged expressions are zero-padded.
+        ctx.index_parents[idx] = tuple(ctx.active_indices[parent].name
+                                       for parent in tree.index.parents)
     with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active:
         return statement(tree.children[0], ctx_active)