DirectDiscretizationConnection: Fix kernel names, hash keys

inducer · Jun 17, 2021 · e179d35 · e179d35
1 parent 03d883a
commit e179d35
Showing 1 changed file with 17 additions and 17 deletions.
diff --git a/meshmode/discretization/connection/direct.py b/meshmode/discretization/connection/direct.py
@@ -370,7 +370,7 @@ def _apply_without_inplace_updates(self, ary):
         actx = ary.array_context
 
         @memoize_in(actx,
-                (DirectDiscretizationConnection, "resample_by_mat_batch_knl"))
+                (DirectDiscretizationConnection, "resample_by_mat_knl"))
         def batch_mat_knl():
             return make_loopy_program(
                 [
@@ -391,11 +391,11 @@ def batch_mat_knl():
                     lp.ValueArg("nelements_vec", np.int32),
                     "...",
                 ],
-                name="resample_by_mat_batch",
+                name="resample_by_mat",
             )
 
         @memoize_in(actx,
-                (DirectDiscretizationConnection, "resample_by_picking_batch_knl"))
+                (DirectDiscretizationConnection, "resample_by_picking_knl"))
         def batch_pick_knl():
             return make_loopy_program(
                 [
@@ -415,7 +415,7 @@ def batch_pick_knl():
                     lp.ValueArg("n_from_nodes", np.int32),
                     "...",
                 ],
-                name="resample_by_picking_batch",
+                name="resample_by_picking",
             )
 
         group_data = []
@@ -478,7 +478,8 @@ def batch_pick_knl():
     def _apply_with_inplace_updates(self, ary):
         actx = ary.array_context
 
-        @memoize_in(actx, (DirectDiscretizationConnection, "resample_by_mat_knl"))
+        @memoize_in(actx, (DirectDiscretizationConnection,
+            "resample_by_mat_knl_inplace"))
         def mat_knl():
             knl = make_loopy_program(
                 """{[iel, idof, j]:
@@ -499,32 +500,32 @@ def mat_knl():
                     lp.ValueArg("nelements_vec", np.int32),
                     "...",
                     ],
-                name="resample_by_mat")
+                name="resample_by_mat_inplace")
 
             return knl
 
         @memoize_in(actx,
-                (DirectDiscretizationConnection, "resample_by_picking_knl"))
+                (DirectDiscretizationConnection, "resample_by_picking_knl_inplace"))
         def pick_knl():
             knl = make_loopy_program(
                 """{[iel, idof]:
                     0<=iel<nelements and
                     0<=idof<n_to_nodes}""",
-                "result[to_element_indices[iel], idof] \
-                    = ary[from_element_indices[iel], pick_list[idof]]",
+                """
+                if from_element_indices[iel] != -1
+                    result[iel, idof] \
+                        = ary[from_element_indices[iel], pick_list[idof]]
+                end
+                """,
                 [
-                    lp.GlobalArg("result", None,
-                        shape="nelements_result, n_to_nodes",
-                        offset=lp.auto),
                     lp.GlobalArg("ary", None,
                         shape="nelements_vec, n_from_nodes",
                         offset=lp.auto),
-                    lp.ValueArg("nelements_result", np.int32),
                     lp.ValueArg("nelements_vec", np.int32),
                     lp.ValueArg("n_from_nodes", np.int32),
                     "...",
                     ],
-                name="resample_by_picking")
+                name="resample_by_picking_inplace")
 
             return knl
 
@@ -549,14 +550,13 @@ def pick_knl():
                             ary=ary[batch.from_group_index],
                             from_element_indices=batch.from_element_indices,
                             to_element_indices=batch.to_element_indices)
-
                 else:
                     actx.call_loopy(pick_knl(),
                             pick_list=point_pick_indices,
                             result=result[i_tgrp],
                             ary=ary[batch.from_group_index],
-                            from_element_indices=batch.from_element_indices,
-                            to_element_indices=batch.to_element_indices)
+                            from_element_indices=batch._global_from_element_indices(
+                                actx, self.to_discr.groups[i_tgrp]))
 
         return result