advanced indexing, disable demoting to private temps to account for advanced indexing exprs

Author: kaushikcfd

meshmode/array_context.py

@@ -409,7 +409,7 @@ def _alias_global_temporaries(t_unit):
     from loopy.kernel import KernelState
     from loopy.schedule import (RunInstruction, EnterLoop, LeaveLoop,
                                 CallKernel, ReturnFromKernel, Barrier)
-    from loopy.schedule.tools import get_nearest_return_from_kernel
+    from loopy.schedule.tools import get_return_from_kernel_mapping
     from pytools import UniqueNameGenerator
     from collections import defaultdict
 
@@ -420,7 +420,7 @@ def _alias_global_temporaries(t_unit):
                           if tv.address_space == AddressSpace.GLOBAL)
     temp_to_live_interval_start = {}
     temp_to_live_interval_end = {}
-    return_from_kernel_idxs = get_nearest_return_from_kernel(kernel)
+    return_from_kernel_idxs = get_return_from_kernel_mapping(kernel)
 
     for sched_idx, sched_item in enumerate(kernel.linearization):
         if isinstance(sched_item, RunInstruction):
@@ -508,10 +508,13 @@ def _make_global_temporaries_private(t_unit):
                         for read_insn in read_insns)):
             if len({knl.insn_inames(read_insn) for read_insn in read_insns}) == 1:
                 knl = lp.assignment_to_subst(knl, tv.name)
-                knl = precompute_for_single_kernel(
-                    knl, t_unit.callables_table, f"{tv.name}_subst",
-                    sweep_inames=(),
-                    temporary_address_space=lp.AddressSpace.PRIVATE)
+                try:
+                    knl = precompute_for_single_kernel(
+                        knl, t_unit.callables_table, f"{tv.name}_subst",
+                        sweep_inames=(),
+                        temporary_address_space=lp.AddressSpace.PRIVATE)
+                except RuntimeError:
+                    pass
 
     return t_unit.with_kernel(knl)
 
@@ -578,7 +581,14 @@ class SingleGridWorkBalancingPytatoArrayContext(PytatoPyOpenCLArrayContextBase):
     def transform_loopy_program(self, t_unit):
         import loopy as lp
 
-        t_unit = _make_global_temporaries_private(t_unit)
+        # if len(t_unit.default_entrypoint.instructions) > 50:
+        #     import pudb; pu.db
+        #     1/0
+        #     with open("nozzle.knl", "w") as f:
+        #         f.write(str(t_unit))
+        #         1/0
+
+        # t_unit = _make_global_temporaries_private(t_unit)
         t_unit = _single_grid_work_group_transform(t_unit, self.queue.device)
         t_unit = lp.set_options(t_unit, "insert_gbarriers")
         t_unit = lp.linearize(lp.preprocess_kernel(t_unit))

meshmode/discretization/connection/direct.py

@@ -425,14 +425,28 @@ def batch_pick_knl():
                     )["result"]
 
                 else:
-                    batch_result = actx.call_loopy(
-                        batch_pick_knl(),
-                        pick_list=point_pick_indices,
-                        ary=ary[batch.from_group_index],
-                        from_element_indices=batch._global_from_element_indices(
-                            actx, self.to_discr.groups[i_tgrp]),
-                        n_to_nodes=self.to_discr.groups[i_tgrp].nunit_dofs
-                    )["result"]
+                    if actx.permits_advanced_indexing:
+                        from_vec = ary[batch.from_group_index]
+                        from_element_indices = actx.thaw(
+                            batch._global_from_element_indices(
+                                actx, self.to_discr.groups[i_tgrp])
+                            ).reshape((-1, 1))
+                        pick_list = actx.thaw(point_pick_indices)
+                        batch_result = actx.np.where(
+                            actx.np.not_equal(from_element_indices, -1),
+                            from_vec[from_element_indices, pick_list],
+                            0)
+                        assert batch_result.shape == (from_element_indices.size,
+                                                      pick_list.size)
+                    else:
+                        batch_result = actx.call_loopy(
+                             batch_pick_knl(),
+                             pick_list=point_pick_indices,
+                             ary=ary[batch.from_group_index],
+                             from_element_indices=batch._global_from_element_indices(
+                                 actx, self.to_discr.groups[i_tgrp]),
+                             n_to_nodes=self.to_discr.groups[i_tgrp].nunit_dofs
+                        )["result"]
 
                 batched_data.append(batch_result)