warp-lang 1.8.1__py3-none-manylinux_2_34_aarch64.whl → 1.9.1__py3-none-manylinux_2_34_aarch64.whl

warp/fem/field/virtual.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, ClassVar, Dict, Set
+from typing import Any, ClassVar, Dict, Optional, Set
 
 import warp as wp
 import warp.fem.operator as operator
@@ -22,7 +22,16 @@ from warp.fem.domain import GeometryDomain
 from warp.fem.linalg import basis_coefficient, generalized_inner, generalized_outer
 from warp.fem.quadrature import Quadrature
 from warp.fem.space import FunctionSpace, SpacePartition, SpaceRestriction
-from warp.fem.types import NULL_NODE_INDEX, DofIndex, Sample, get_node_coord, get_node_index_in_element
+from warp.fem.types import (
+    NULL_ELEMENT_INDEX,
+    NULL_NODE_INDEX,
+    DofIndex,
+    ElementIndex,
+    NodeElementIndex,
+    Sample,
+    get_node_coord,
+    get_node_index_in_element,
+)
 from warp.fem.utils import type_zero_element
 
 from .field import SpaceField
@@ -567,7 +576,13 @@ class LocalTrialField(LocalAdjointField):
         return s.trial_dof
 
 
-def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, accumulate_dtype: type):
+def make_linear_dispatch_kernel(
+    test: LocalTestField,
+    quadrature: Quadrature,
+    accumulate_dtype: type,
+    tile_size: int = 1,
+    kernel_options: Optional[Dict[str, Any]] = None,
+):
     global_test: TestField = test.global_field
     space_restriction = global_test.space_restriction
     domain = global_test.domain
@@ -583,8 +598,42 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
     TEST_OUTER_GRAD_BEGIN = test._TAYLOR_DOF_OFFSETS[LocalAdjointField.OUTER_GRAD_DOF]
 
     TEST_NODE_DOF_DIM = test.value_dof_count // test.node_dof_count
+    TEST_NODE_DOF_COUNT = test.node_dof_count
+
+    res_vec = cache.cached_vec_type(length=test.node_dof_count, dtype=accumulate_dtype)
+    qp_vec = cache.cached_vec_type(length=test.node_dof_count, dtype=float)
+
+    @cache.dynamic_func(f"{test.name}_{quadrature.name}")
+    def next_qp(
+        qp: int,
+        elem_offset: int,
+        qp_point_count: int,
+        element_index: ElementIndex,
+        test_element_index: NodeElementIndex,
+        element_end: int,
+        qp_arg: quadrature.Arg,
+        domain_arg: domain.ElementArg,
+        domain_index_arg: domain.ElementIndexArg,
+        test_arg: space_restriction.NodeArg,
+    ):
+        while qp >= qp_point_count and elem_offset < element_end:
+            # Next element
+            elem_offset += 1
+
+            if elem_offset < element_end:
+                qp -= qp_point_count
+                test_element_index = space_restriction.node_element_index(test_arg, elem_offset)
+                element_index = domain.element_index(domain_index_arg, test_element_index.domain_element_index)
+                qp_point_count = quadrature.point_count(
+                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index
+                )
 
-    @cache.dynamic_kernel(f"{test.name}_{quadrature.name}_{wp.types.get_type_code(accumulate_dtype)}")
+        return qp, elem_offset, qp_point_count, element_index, test_element_index
+
+    @cache.dynamic_kernel(
+        f"{test.name}_{quadrature.name}_{wp.types.get_type_code(accumulate_dtype)}_{tile_size}",
+        kernel_options=kernel_options,
+    )
     def dispatch_linear_kernel_fn(
         qp_arg: quadrature.Arg,
         domain_arg: domain.ElementArg,
@@ -594,33 +643,47 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
         local_result: wp.array3d(dtype=Any),
         result: wp.array2d(dtype=Any),
     ):
-        local_node_index, test_node_dof = wp.tid()
+        local_node_index, lane = wp.tid()
+
         node_index = space_restriction.node_partition_index(test_arg, local_node_index)
         element_beg, element_end = space_restriction.node_element_range(test_arg, node_index)
 
-        val_sum = accumulate_dtype(0.0)
-
-        for n in range(element_beg, element_end):
-            test_element_index = space_restriction.node_element_index(test_arg, n)
-            element_index = domain.element_index(domain_index_arg, test_element_index.domain_element_index)
-
-            qp_point_count = quadrature.point_count(
-                domain_arg, qp_arg, test_element_index.domain_element_index, element_index
+        val_sum = res_vec()
+
+        elem_offset = element_beg - 1
+        qp_point_count = int(0)
+        qp = lane
+        test_element_index = NodeElementIndex()
+        element_index = ElementIndex(NULL_ELEMENT_INDEX)
+
+        while elem_offset < element_end:
+            qp, elem_offset, qp_point_count, element_index, test_element_index = next_qp(
+                qp,
+                elem_offset,
+                qp_point_count,
+                element_index,
+                test_element_index,
+                element_end,
+                qp_arg,
+                domain_arg,
+                domain_index_arg,
+                test_arg,
             )
-            for k in range(qp_point_count):
+
+            if qp < qp_point_count:
                 qp_index = quadrature.point_index(
-                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, qp
                 )
                 qp_eval_index = quadrature.point_evaluation_index(
-                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, qp
                 )
                 coords = quadrature.point_coords(
-                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, qp
                 )
 
                 qp_result = local_result[qp_eval_index]
 
-                qp_sum = float(0.0)
+                qp_sum = qp_vec()
 
                 if wp.static(0 != TEST_INNER_COUNT):
                     w = test.space.element_inner_weight(
@@ -631,9 +694,12 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
                         test_element_index.node_index_in_element,
                         qp_index,
                     )
-                    for val_dof in range(TEST_NODE_DOF_DIM):
-                        test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
-                        qp_sum += basis_coefficient(w, val_dof) * qp_result[TEST_INNER_BEGIN, test_dof]
+                    for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                        for val_dof in range(TEST_NODE_DOF_DIM):
+                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
+                            qp_sum[test_node_dof] += (
+                                basis_coefficient(w, val_dof) * qp_result[TEST_INNER_BEGIN, test_dof]
+                            )
 
                 if wp.static(0 != TEST_OUTER_COUNT):
                     w = test.space.element_outer_weight(
@@ -644,9 +710,12 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
                         test_element_index.node_index_in_element,
                         qp_index,
                     )
-                    for val_dof in range(TEST_NODE_DOF_DIM):
-                        test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
-                        qp_sum += basis_coefficient(w, val_dof) * qp_result[TEST_OUTER_BEGIN, test_dof]
+                    for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                        for val_dof in range(TEST_NODE_DOF_DIM):
+                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
+                            qp_sum[test_node_dof] += (
+                                basis_coefficient(w, val_dof) * qp_result[TEST_OUTER_BEGIN, test_dof]
+                            )
 
                 if wp.static(0 != TEST_INNER_GRAD_COUNT):
                     w_grad = test.space.element_inner_weight_gradient(
@@ -657,13 +726,14 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
                         test_element_index.node_index_in_element,
                         qp_index,
                     )
-                    for val_dof in range(TEST_NODE_DOF_DIM):
-                        test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
-                        for grad_dof in range(TEST_INNER_GRAD_COUNT):
-                            qp_sum += (
-                                basis_coefficient(w_grad, val_dof, grad_dof)
-                                * qp_result[grad_dof + TEST_INNER_GRAD_BEGIN, test_dof]
-                            )
+                    for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                        for val_dof in range(TEST_NODE_DOF_DIM):
+                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
+                            for grad_dof in range(TEST_INNER_GRAD_COUNT):
+                                qp_sum[test_node_dof] += (
+                                    basis_coefficient(w_grad, val_dof, grad_dof)
+                                    * qp_result[grad_dof + TEST_INNER_GRAD_BEGIN, test_dof]
+                                )
 
                 if wp.static(0 != TEST_OUTER_GRAD_COUNT):
                     w_grad = test.space.element_outer_weight_gradient(
@@ -674,23 +744,36 @@ def make_linear_dispatch_kernel(test: LocalTestField, quadrature: Quadrature, ac
                         test_element_index.node_index_in_element,
                         qp_index,
                     )
-                    for val_dof in range(TEST_NODE_DOF_DIM):
-                        test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
-                        for grad_dof in range(TEST_OUTER_GRAD_COUNT):
-                            qp_sum += (
-                                basis_coefficient(w_grad, val_dof, grad_dof)
-                                * qp_result[grad_dof + TEST_OUTER_GRAD_BEGIN, test_dof]
-                            )
+                    for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                        for val_dof in range(TEST_NODE_DOF_DIM):
+                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + val_dof
+                            for grad_dof in range(TEST_OUTER_GRAD_COUNT):
+                                qp_sum[test_node_dof] += (
+                                    basis_coefficient(w_grad, val_dof, grad_dof)
+                                    * qp_result[grad_dof + TEST_OUTER_GRAD_BEGIN, test_dof]
+                                )
 
-                val_sum += accumulate_dtype(qp_sum)
+                val_sum += res_vec(qp_sum)
+                qp += wp.static(tile_size)
 
-        result[node_index, test_node_dof] += result.dtype(val_sum)
+        if wp.static(tile_size == 1):
+            for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                result[node_index, test_node_dof] += result.dtype(val_sum[test_node_dof])
+        else:
+            t_sum = wp.tile_sum(wp.tile(val_sum, preserve_type=True))[0]
+            for test_node_dof in range(lane, TEST_NODE_DOF_COUNT, wp.static(tile_size)):
+                result[node_index, test_node_dof] += result.dtype(t_sum[test_node_dof])
 
     return dispatch_linear_kernel_fn
 
 
 def make_bilinear_dispatch_kernel(
-    test: LocalTestField, trial: LocalTrialField, quadrature: Quadrature, accumulate_dtype: type
+    test: LocalTestField,
+    trial: LocalTrialField,
+    quadrature: Quadrature,
+    accumulate_dtype: type,
+    tile_size: int = 1,
+    kernel_options: Optional[Dict[str, Any]] = None,
 ):
     global_test: TestField = test.global_field
     space_restriction = global_test.space_restriction
@@ -718,12 +801,24 @@ def make_bilinear_dispatch_kernel(
 
     TEST_NODE_DOF_DIM = test.value_dof_count // test.node_dof_count
     TRIAL_NODE_DOF_DIM = trial.value_dof_count // trial.node_dof_count
+    TEST_TRIAL_NODE_DOF_DIM = TEST_NODE_DOF_DIM * TRIAL_NODE_DOF_DIM
+
+    TEST_NODE_DOF_COUNT = test.node_dof_count
+    TRIAL_NODE_DOF_COUNT = trial.node_dof_count
+    TEST_TAYLOR_DOF_COUNT = test.TAYLOR_DOF_COUNT
+    TRIAL_TAYLOR_DOF_COUNT = trial.TAYLOR_DOF_COUNT
 
     MAX_NODES_PER_ELEMENT = trial.space.topology.MAX_NODES_PER_ELEMENT
 
     trial_dof_vec = cache.cached_vec_type(length=trial.TAYLOR_DOF_COUNT, dtype=float)
+    test_dof_vec = cache.cached_vec_type(length=test.TAYLOR_DOF_COUNT, dtype=float)
+
+    val_t = cache.cached_mat_type(shape=(test.node_dof_count, trial.node_dof_count), dtype=accumulate_dtype)
 
-    @cache.dynamic_kernel(f"{trial.name}_{test.name}_{quadrature.name}{wp.types.get_type_code(accumulate_dtype)}")
+    @cache.dynamic_kernel(
+        f"{trial.name}_{test.name}_{quadrature.name}{wp.types.get_type_code(accumulate_dtype)}_{tile_size}",
+        kernel_options=kernel_options,
+    )
     def dispatch_bilinear_kernel_fn(
         qp_arg: quadrature.Arg,
         domain_arg: domain.ElementArg,
@@ -733,163 +828,166 @@ def make_bilinear_dispatch_kernel(
         trial_partition_arg: trial.space_partition.PartitionArg,
         trial_topology_arg: trial.space_partition.space_topology.TopologyArg,
         trial_space_arg: trial.space.SpaceArg,
-        local_result: wp.array4d(dtype=trial_dof_vec),
+        local_result: wp.array4d(dtype=float),
         triplet_rows: wp.array(dtype=int),
         triplet_cols: wp.array(dtype=int),
         triplet_values: wp.array3d(dtype=Any),
     ):
-        test_local_node_index, test_node_dof, trial_node_dof, trial_node = wp.tid()
+        test_node_offset, trial_node, lane = wp.tid()
 
-        test_node_index = space_restriction.node_partition_index(test_arg, test_local_node_index)
-        element_beg, element_end = space_restriction.node_element_range(test_arg, test_node_index)
+        test_node_index = space_restriction.node_partition_index_from_element_offset(test_arg, test_node_offset)
 
-        for element in range(element_beg, element_end):
-            test_element_index = space_restriction.node_element_index(test_arg, element)
-            element_index = domain.element_index(domain_index_arg, test_element_index.domain_element_index)
-            test_node = test_element_index.node_index_in_element
+        test_element_index = space_restriction.node_element_index(test_arg, test_node_offset)
+        element_index = domain.element_index(domain_index_arg, test_element_index.domain_element_index)
+        test_node = test_element_index.node_index_in_element
 
-            element_trial_node_count = trial.space.topology.element_node_count(
-                domain_arg, trial_topology_arg, element_index
-            )
+        element_trial_node_count = trial.space.topology.element_node_count(
+            domain_arg, trial_topology_arg, element_index
+        )
 
-            qp_point_count = wp.where(
-                trial_node < element_trial_node_count,
-                quadrature.point_count(domain_arg, qp_arg, test_element_index.domain_element_index, element_index),
-                0,
+        if trial_node >= element_trial_node_count:
+            block_offset = test_node_offset * MAX_NODES_PER_ELEMENT + trial_node
+            triplet_rows[block_offset] = NULL_NODE_INDEX
+            triplet_cols[block_offset] = NULL_NODE_INDEX
+            return
+
+        qp_point_count = quadrature.point_count(
+            domain_arg, qp_arg, test_element_index.domain_element_index, element_index
+        )
+        qp_dof_count = qp_point_count * TEST_TRIAL_NODE_DOF_DIM
+
+        val_sum = val_t()
+
+        for dof in range(lane, qp_dof_count, wp.static(tile_size)):
+            k = dof // TEST_TRIAL_NODE_DOF_DIM
+            test_trial_val_dof = dof - k * TEST_TRIAL_NODE_DOF_DIM
+            test_val_dof = test_trial_val_dof // TRIAL_NODE_DOF_DIM
+            trial_val_dof = test_trial_val_dof - test_val_dof * TRIAL_NODE_DOF_DIM
+
+            qp_index = quadrature.point_index(
+                domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+            )
+            qp_eval_index = quadrature.point_evaluation_index(
+                domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+            )
+            coords = quadrature.point_coords(
+                domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
             )
 
-            val_sum = accumulate_dtype(0.0)
+            # test shape functions
+            w_test = test_dof_vec()
 
-            for k in range(qp_point_count):
-                qp_index = quadrature.point_index(
-                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
-                )
-                qp_eval_index = quadrature.point_evaluation_index(
-                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
-                )
-                coords = quadrature.point_coords(
-                    domain_arg, qp_arg, test_element_index.domain_element_index, element_index, k
+            if wp.static(0 != TEST_INNER_COUNT):
+                w_test_inner = test.space.element_inner_weight(
+                    domain_arg, test_space_arg, element_index, coords, test_node, qp_index
                 )
+                w_test[TEST_INNER_BEGIN] = basis_coefficient(w_test_inner, test_val_dof)
 
-                qp_result = local_result[qp_eval_index]
-                trial_result = float(0.0)
-
-                if wp.static(0 != TEST_INNER_COUNT):
-                    w_test_inner = test.space.element_inner_weight(
-                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
-                    )
+            if wp.static(0 != TEST_OUTER_COUNT):
+                w_test_outer = test.space.element_outer_weight(
+                    domain_arg, test_space_arg, element_index, coords, test_node, qp_index
+                )
+                w_test[TEST_OUTER_BEGIN] = basis_coefficient(w_test_outer, test_val_dof)
 
-                if wp.static(0 != TEST_OUTER_COUNT):
-                    w_test_outer = test.space.element_outer_weight(
-                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
+            if wp.static(0 != TEST_INNER_GRAD_COUNT):
+                w_test_grad_inner = test.space.element_inner_weight_gradient(
+                    domain_arg, test_space_arg, element_index, coords, test_node, qp_index
+                )
+                for grad_dof in range(TEST_INNER_GRAD_COUNT):
+                    w_test[TEST_INNER_GRAD_BEGIN + grad_dof] = basis_coefficient(
+                        w_test_grad_inner, test_val_dof, grad_dof
                     )
 
-                if wp.static(0 != TEST_INNER_GRAD_COUNT):
-                    w_test_grad_inner = test.space.element_inner_weight_gradient(
-                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
+            if wp.static(0 != TEST_OUTER_GRAD_COUNT):
+                w_test_grad_outer = test.space.element_outer_weight_gradient(
+                    domain_arg, test_space_arg, element_index, coords, test_node, qp_index
+                )
+                for grad_dof in range(TEST_OUTER_GRAD_COUNT):
+                    w_test[TEST_OUTER_GRAD_BEGIN + grad_dof] = basis_coefficient(
+                        w_test_grad_outer, test_val_dof, grad_dof
                     )
 
-                if wp.static(0 != TEST_OUTER_GRAD_COUNT):
-                    w_test_grad_outer = test.space.element_outer_weight_gradient(
-                        domain_arg, test_space_arg, element_index, coords, test_node, qp_index
-                    )
+            # trial shape functions
+            w_trial = trial_dof_vec()
 
-                if wp.static(0 != TRIAL_INNER_COUNT):
-                    w_trial_inner = trial.space.element_inner_weight(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
-                    )
+            if wp.static(0 != TRIAL_INNER_COUNT):
+                w_trial_inner = trial.space.element_inner_weight(
+                    domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
+                )
+                w_trial[TRIAL_INNER_BEGIN] = basis_coefficient(w_trial_inner, trial_val_dof)
 
-                if wp.static(0 != TRIAL_OUTER_COUNT):
-                    w_trial_outer = trial.space.element_outer_weight(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
-                    )
+            if wp.static(0 != TRIAL_OUTER_COUNT):
+                w_trial_outer = trial.space.element_outer_weight(
+                    domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
+                )
+                w_trial[TRIAL_OUTER_BEGIN] = basis_coefficient(w_trial_outer, trial_val_dof)
 
-                if wp.static(0 != TRIAL_INNER_GRAD_COUNT):
-                    w_trial_grad_inner = trial.space.element_inner_weight_gradient(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
+            if wp.static(0 != TRIAL_INNER_GRAD_COUNT):
+                w_trial_grad_inner = trial.space.element_inner_weight_gradient(
+                    domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
+                )
+                for grad_dof in range(TRIAL_INNER_GRAD_COUNT):
+                    w_trial[TRIAL_INNER_GRAD_BEGIN + grad_dof] = basis_coefficient(
+                        w_trial_grad_inner, trial_val_dof, grad_dof
                     )
 
-                if wp.static(0 != TRIAL_OUTER_GRAD_COUNT):
-                    w_trial_grad_outer = trial.space.element_outer_weight_gradient(
-                        domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
+            if wp.static(0 != TRIAL_OUTER_GRAD_COUNT):
+                w_trial_grad_outer = trial.space.element_outer_weight_gradient(
+                    domain_arg, trial_space_arg, element_index, coords, trial_node, qp_index
+                )
+                for grad_dof in range(TRIAL_OUTER_GRAD_COUNT):
+                    w_trial[TRIAL_OUTER_GRAD_BEGIN + grad_dof] = basis_coefficient(
+                        w_trial_grad_outer, trial_val_dof, grad_dof
                     )
 
-                for trial_val_dof in range(TRIAL_NODE_DOF_DIM):
+            # triple product test @ qp @ trial
+            for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                test_dof = test_node_dof * TEST_NODE_DOF_DIM + test_val_dof
+                for trial_node_dof in range(TRIAL_NODE_DOF_COUNT):
+                    dof_res = float(0.0)
                     trial_dof = trial_node_dof * TRIAL_NODE_DOF_DIM + trial_val_dof
-                    test_result = trial_dof_vec(0.0)
-
-                    if wp.static(0 != TEST_INNER_COUNT):
-                        for test_val_dof in range(TEST_NODE_DOF_DIM):
-                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + test_val_dof
-                            test_result += (
-                                basis_coefficient(w_test_inner, test_val_dof)
-                                * qp_result[test_dof, trial_dof, TEST_INNER_BEGIN]
-                            )
 
-                    if wp.static(0 != TEST_OUTER_COUNT):
-                        for test_val_dof in range(TEST_NODE_DOF_DIM):
-                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + test_val_dof
-                            test_result += (
-                                basis_coefficient(w_test_outer, test_val_dof)
-                                * qp_result[test_dof, trial_dof, TEST_OUTER_BEGIN]
+                    for test_taylor_dof in range(TEST_TAYLOR_DOF_COUNT):
+                        test_res = float(0.0)
+                        for trial_taylor_dof in range(TRIAL_TAYLOR_DOF_COUNT):
+                            taylor_dof = test_taylor_dof * TRIAL_TAYLOR_DOF_COUNT + trial_taylor_dof
+                            test_res += (
+                                local_result[test_dof, trial_dof, qp_eval_index, taylor_dof] * w_trial[trial_taylor_dof]
                             )
+                        dof_res += w_test[test_taylor_dof] * test_res
 
-                    if wp.static(0 != TEST_INNER_GRAD_COUNT):
-                        for test_val_dof in range(TEST_NODE_DOF_DIM):
-                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + test_val_dof
-                            for grad_dof in range(TEST_INNER_GRAD_COUNT):
-                                test_result += (
-                                    basis_coefficient(w_test_grad_inner, test_val_dof, grad_dof)
-                                    * qp_result[test_dof, trial_dof, grad_dof + TEST_INNER_GRAD_BEGIN]
-                                )
+                    val_sum[test_node_dof, trial_node_dof] += accumulate_dtype(dof_res)
 
-                    if wp.static(0 != TEST_OUTER_GRAD_COUNT):
-                        for test_val_dof in range(TEST_NODE_DOF_DIM):
-                            test_dof = test_node_dof * TEST_NODE_DOF_DIM + test_val_dof
-                            for grad_dof in range(TEST_OUTER_GRAD_COUNT):
-                                test_result += (
-                                    basis_coefficient(w_test_grad_outer, test_val_dof, grad_dof)
-                                    * qp_result[test_dof, trial_dof, grad_dof + TEST_OUTER_GRAD_BEGIN]
-                                )
+        # write block value
+        block_offset = test_node_offset * MAX_NODES_PER_ELEMENT + trial_node
+        if wp.static(tile_size) > 1:
+            val_sum = wp.tile_sum(wp.tile(val_sum, preserve_type=True))[0]
 
-                    if wp.static(0 != TRIAL_INNER_COUNT):
-                        trial_result += basis_coefficient(w_trial_inner, trial_val_dof) * test_result[TRIAL_INNER_BEGIN]
+            for dof in range(lane, wp.static(TEST_NODE_DOF_COUNT * TRIAL_NODE_DOF_COUNT), wp.static(tile_size)):
+                test_node_dof = dof // TRIAL_NODE_DOF_COUNT
+                trial_node_dof = dof - TRIAL_NODE_DOF_COUNT * test_node_dof
 
-                    if wp.static(0 != TRIAL_OUTER_COUNT):
-                        trial_result += basis_coefficient(w_trial_outer, trial_val_dof) * test_result[TRIAL_OUTER_BEGIN]
-
-                    if wp.static(0 != TRIAL_INNER_GRAD_COUNT):
-                        for grad_dof in range(TRIAL_INNER_GRAD_COUNT):
-                            trial_result += (
-                                basis_coefficient(w_trial_grad_inner, trial_val_dof, grad_dof)
-                                * test_result[grad_dof + TRIAL_INNER_GRAD_BEGIN]
-                            )
-
-                    if wp.static(0 != TRIAL_OUTER_GRAD_COUNT):
-                        for grad_dof in range(TRIAL_OUTER_GRAD_COUNT):
-                            trial_result += (
-                                basis_coefficient(w_trial_grad_outer, trial_val_dof, grad_dof)
-                                * test_result[grad_dof + TRIAL_OUTER_GRAD_BEGIN]
-                            )
-
-                val_sum += accumulate_dtype(trial_result)
-
-            block_offset = element * MAX_NODES_PER_ELEMENT + trial_node
-            triplet_values[block_offset, test_node_dof, trial_node_dof] = triplet_values.dtype(val_sum)
-
-            # Set row and column indices
-            if test_node_dof == 0 and trial_node_dof == 0:
-                if trial_node < element_trial_node_count:
-                    trial_node_index = trial.space_partition.partition_node_index(
-                        trial_partition_arg,
-                        trial.space.topology.element_node_index(
-                            domain_arg, trial_topology_arg, element_index, trial_node
-                        ),
+                triplet_values[block_offset, test_node_dof, trial_node_dof] = triplet_values.dtype(
+                    val_sum[test_node_dof, trial_node_dof]
+                )
+        else:
+            for test_node_dof in range(TEST_NODE_DOF_COUNT):
+                for trial_node_dof in range(TRIAL_NODE_DOF_COUNT):
+                    triplet_values[block_offset, test_node_dof, trial_node_dof] = triplet_values.dtype(
+                        val_sum[test_node_dof, trial_node_dof]
                     )
-                else:
-                    trial_node_index = NULL_NODE_INDEX  # will get ignored when converting to bsr
 
-                triplet_rows[block_offset] = test_node_index
-                triplet_cols[block_offset] = trial_node_index
+        # Set row and column indices
+        if lane == 0:
+            if trial_node < element_trial_node_count:
+                trial_node_index = trial.space_partition.partition_node_index(
+                    trial_partition_arg,
+                    trial.space.topology.element_node_index(domain_arg, trial_topology_arg, element_index, trial_node),
+                )
+            else:
+                trial_node_index = NULL_NODE_INDEX  # will get ignored when converting to bsr
+
+            triplet_rows[block_offset] = test_node_index
+            triplet_cols[block_offset] = trial_node_index
 
     return dispatch_bilinear_kernel_fn

warp/fem/geometry/geometry.py

@@ -542,17 +542,17 @@ class Geometry:
 
         pos_type = cache.cached_vec_type(self.dimension, dtype=float)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_lookup(args: self.CellArg, pos: pos_type, max_dist: float):
             return unfiltered_cell_lookup(args, pos, max_dist, null_filter_data, null_filter_target)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_lookup(args: self.CellArg, pos: pos_type, guess: Sample):
             guess_pos = self.cell_position(args, guess)
             max_dist = wp.length(guess_pos - pos)
             return unfiltered_cell_lookup(args, pos, max_dist, null_filter_data, null_filter_target)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_lookup(args: self.CellArg, pos: pos_type):
             max_dist = 0.0
             return unfiltered_cell_lookup(args, pos, max_dist, null_filter_data, null_filter_target)
@@ -561,13 +561,13 @@ class Geometry:
         filtered_cell_lookup = self.make_filtered_cell_lookup(filter_func=_array_load)
         pos_type = cache.cached_vec_type(self.dimension, dtype=float)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_lookup(
             args: self.CellArg, pos: pos_type, max_dist: float, filter_array: wp.array(dtype=Any), filter_target: Any
         ):
             return filtered_cell_lookup(args, pos, max_dist, filter_array, filter_target)
 
-        @cache.dynamic_func(suffix=self.name)
+        @cache.dynamic_func(suffix=self.name, allow_overloads=True)
         def cell_lookup(args: self.CellArg, pos: pos_type, filter_array: wp.array(dtype=Any), filter_target: Any):
             max_dist = 0.0
             return filtered_cell_lookup(args, pos, max_dist, filter_array, filter_target)