Use ArrayContext+assert instead of PyOpenCLActx in annotations

Author: inducer

doc/conf.py

@@ -55,6 +55,7 @@
     "CallInstruction": "class:loopy.kernel.instruction.CallInstruction",
     # arraycontext
     "Array": "obj:arraycontext.Array",
+    "ArrayContext": "class:arraycontext.ArrayContext",
     # boxtree
     "FMMTraversalInfo": "class:boxtree.traversal.FMMTraversalInfo",
     # sumpy

sumpy/array_context.py

@@ -40,6 +40,7 @@
 
     from numpy.typing import DTypeLike
 
+    from arraycontext import ArrayContext
     from loopy import TranslationUnit
     from loopy.codegen import PreambleInfo
     from pytools.tag import ToTagSetConvertible
@@ -112,7 +113,10 @@ def transform_loopy_program(self, t_unit: TranslationUnit):
         return t_unit
 
 
-def is_cl_cpu(actx: PyOpenCLArrayContext) -> bool:
+def is_cl_cpu(actx: ArrayContext) -> bool:
+    if not isinstance(actx, PyOpenCLArrayContext):
+        return False
+
     import pyopencl as cl
     return all(dev.type & cl.device_type.CPU for dev in actx.context.devices)
 

sumpy/distributed.py

@@ -33,13 +33,13 @@
 
 
 if TYPE_CHECKING:
-    from sumpy.array_context import PyOpenCLArrayContext
+    from arraycontext import ArrayContext
 
 
 class DistributedSumpyExpansionWrangler(
         DistributedExpansionWranglerMixin, SumpyExpansionWrangler):
     def __init__(
-            self, actx: PyOpenCLArrayContext,
+            self, actx: ArrayContext,
             comm, tree_indep, local_traversal, global_traversal,
             dtype, fmm_level_to_order, communicate_mpoles_via_allreduce=False,
             **kwargs):
@@ -53,7 +53,7 @@ def __init__(
         self.communicate_mpoles_via_allreduce = communicate_mpoles_via_allreduce
 
     def distribute_source_weights(self,
-            actx: PyOpenCLArrayContext, src_weight_vecs, src_idx_all_ranks):
+            actx: ArrayContext, src_weight_vecs, src_idx_all_ranks):
         src_weight_vecs_host = [
             actx.to_numpy(src_weight) for src_weight in src_weight_vecs
             ]
@@ -68,7 +68,7 @@ def distribute_source_weights(self,
         return local_src_weight_vecs_device
 
     def gather_potential_results(self,
-            actx: PyOpenCLArrayContext, potentials, tgt_idx_all_ranks):
+            actx: ArrayContext, potentials, tgt_idx_all_ranks):
         potentials_host_vec = [
             actx.to_numpy(potentials_dev) for potentials_dev in potentials
             ]
@@ -109,7 +109,7 @@ def reorder(x):
             return None
 
     def communicate_mpoles(self,
-            actx: PyOpenCLArrayContext, mpole_exps, return_stats=False):
+            actx: ArrayContext, mpole_exps, return_stats=False):
         mpole_exps_host = actx.to_numpy(mpole_exps)
         stats = super().communicate_mpoles(actx, mpole_exps_host, return_stats)
         mpole_exps[:] = mpole_exps_host

sumpy/e2e.py

@@ -25,6 +25,7 @@
 
 import logging
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
 import numpy as np
 from typing_extensions import override
@@ -34,10 +35,14 @@
 from pytools import memoize_method
 
 import sumpy.symbolic as sym
-from sumpy.array_context import PyOpenCLArrayContext, make_loopy_program
+from sumpy.array_context import make_loopy_program
 from sumpy.tools import KernelCacheMixin, to_complex_dtype
 
 
+if TYPE_CHECKING:
+    from arraycontext import ArrayContext
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -267,7 +272,7 @@ def get_optimized_kernel(self):
 
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg src_expansions:
         :arg src_box_starts:
@@ -511,7 +516,7 @@ def get_optimized_kernel(self, result_dtype):
 
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg src_expansions:
         :arg src_box_starts:
@@ -624,7 +629,7 @@ def get_optimized_kernel(self, result_dtype):
 
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg src_rscale:
         :arg translation_classes_level_start:
@@ -729,7 +734,7 @@ def get_optimized_kernel(self, result_dtype):
             knl = optimization(knl)
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg src_expansions
         :arg preprocessed_src_expansions
@@ -830,7 +835,7 @@ def get_optimized_kernel(self, result_dtype):
         knl = lp.add_inames_for_unused_hw_axes(knl)
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg tgt_expansions
         :arg tgt_expansions_before_postprocessing
@@ -943,7 +948,7 @@ def get_kernel(self):
 
         return loopy_knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg src_expansions:
         :arg src_box_starts:
@@ -1050,7 +1055,7 @@ def get_kernel(self):
 
         return loopy_knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg src_expansions:
         :arg src_box_starts:

sumpy/e2p.py

@@ -24,17 +24,22 @@
 """
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 import loopy as lp
 import pytools.obj_array as obj_array
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION  # noqa: F401
 
-from sumpy.array_context import PyOpenCLArrayContext, make_loopy_program
+from sumpy.array_context import make_loopy_program
 from sumpy.tools import KernelCacheMixin, gather_loopy_arguments
 
 
+if TYPE_CHECKING:
+    from arraycontext import ArrayContext
+
+
 __doc__ = """
 
 Expansion-to-particle
@@ -203,7 +208,7 @@ def get_optimized_kernel(self):
 
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         """
         :arg expansions:
         :arg target_boxes:
@@ -331,7 +336,7 @@ def get_optimized_kernel(self):
 
         return knl
 
-    def __call__(self, actx: PyOpenCLArrayContext, **kwargs):
+    def __call__(self, actx: ArrayContext, **kwargs):
         centers = kwargs.pop("centers")
         # "1" may be passed for rscale, which won't have its type
         # meaningfully inferred. Make the type of rscale explicit.

sumpy/fmm.py

@@ -77,9 +77,8 @@
     from numpy.typing import DTypeLike
 
     import pyopencl
-    from arraycontext import Array
+    from arraycontext import Array, ArrayContext
 
-    from sumpy.array_context import PyOpenCLArrayContext
     from sumpy.expansion.local import LocalExpansionBase
     from sumpy.expansion.multipole import MultipoleExpansionBase
 
@@ -114,7 +113,7 @@ class SumpyTreeIndependentDataForWrangler(TreeIndependentDataForWrangler):
     strength_usage: Sequence[int] | None
 
     def __init__(self,
-                 array_context: PyOpenCLArrayContext,
+                 array_context: ArrayContext,
                  multipole_expansion_factory: MultipoleExpansionFromOrderFactory,
                  local_expansion_factory: LocalExpansionFromOrderFactory,
                  target_kernels: Sequence[Kernel],
@@ -134,7 +133,7 @@ def __init__(self,
         """
         super().__init__()
 
-        self._setup_actx: PyOpenCLArrayContext = array_context
+        self._setup_actx: ArrayContext = array_context
 
         self.multipole_expansion_factory = multipole_expansion_factory
         self.local_expansion_factory = local_expansion_factory
@@ -422,7 +421,7 @@ def order_to_size(order: int):
         return build_csr_level_starts(self.level_orders, order_to_size,
                 level_starts=self.m2l_translation_class_level_start_box_nrs())
 
-    def multipole_expansion_zeros(self, actx: PyOpenCLArrayContext) -> Array:
+    def multipole_expansion_zeros(self, actx: ArrayContext) -> Array:
         """Return an expansions array (which must support addition)
         capable of holding one multipole or local expansion for every
         box in the tree.
@@ -441,7 +440,7 @@ def local_expansion_zeros(self, actx) -> Array:
                 dtype=self.dtype)
 
     def m2l_translation_classes_dependent_data_zeros(
-            self, actx: PyOpenCLArrayContext):
+            self, actx: ArrayContext):
         data_level_starts = (
             self.m2l_translation_classes_dependent_data_level_starts())
         level_start_box_nrs = (
@@ -497,7 +496,7 @@ def order_to_size(order):
                 level_starts=self.tree_level_start_box_nrs)
 
     def m2l_preproc_mpole_expansion_zeros(
-            self, actx: PyOpenCLArrayContext, template_ary):
+            self, actx: ArrayContext, template_ary):
         level_starts = self.m2l_preproc_mpole_expansions_level_starts()
 
         result = []
@@ -522,7 +521,7 @@ def m2l_preproc_mpole_expansions_view(self, mpole_exps, level):
     m2l_work_array_level_starts = m2l_preproc_mpole_expansions_level_starts
 
     def output_zeros(self,
-                actx: PyOpenCLArrayContext
+                actx: ArrayContext
             ) -> obj_array.ObjectArray1D[Array]:
         """Return a potentials array (which must support addition) capable of
         holding a potential value for each target in the tree. Note that
@@ -587,7 +586,7 @@ def box_target_list_kwargs(self):
 
     # }}}
 
-    def run_opencl_fft(self, actx: PyOpenCLArrayContext,
+    def run_opencl_fft(self, actx: ArrayContext,
             input_vec, inverse, wait_for):
         app = self.tree_indep.opencl_fft_app(input_vec.shape, input_vec.dtype,
             inverse)
@@ -601,7 +600,7 @@ def run_opencl_fft(self, actx: PyOpenCLArrayContext,
         return result
 
     def form_multipoles(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             level_start_source_box_nrs, source_boxes,
             src_weight_vecs):
         mpoles = self.multipole_expansion_zeros(actx)
@@ -635,7 +634,7 @@ def form_multipoles(self,
         return mpoles
 
     def coarsen_multipoles(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             level_start_source_parent_box_nrs,
             source_parent_boxes,
             mpoles):
@@ -689,7 +688,7 @@ def coarsen_multipoles(self,
         return mpoles
 
     def eval_direct(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             target_boxes, source_box_starts,
             source_box_lists, src_weight_vecs):
         pot = self.output_zeros(actx)
@@ -791,7 +790,7 @@ def _add_m2l_precompute_kwargs(self, kwargs_for_m2l,
             self.translation_classes_data.from_sep_siblings_translation_classes
 
     def multipole_to_local(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             level_start_target_box_nrs,
             target_boxes, src_box_starts, src_box_lists,
             mpole_exps):
@@ -915,7 +914,7 @@ def multipole_to_local(self,
         return local_exps
 
     def eval_multipoles(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             target_boxes_by_source_level, source_boxes_by_level, mpole_exps):
         pot = self.output_zeros(actx)
 
@@ -956,7 +955,7 @@ def eval_multipoles(self,
         return pot
 
     def form_locals(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             level_start_target_or_target_parent_box_nrs,
             target_or_target_parent_boxes, starts, lists, src_weight_vecs):
         local_exps = self.local_expansion_zeros(actx)
@@ -997,7 +996,7 @@ def form_locals(self,
         return local_exps
 
     def refine_locals(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             level_start_target_or_target_parent_box_nrs,
             target_or_target_parent_boxes,
             local_exps):
@@ -1040,7 +1039,7 @@ def refine_locals(self,
         return local_exps
 
     def eval_locals(self,
-            actx: PyOpenCLArrayContext,
+            actx: ArrayContext,
             level_start_target_box_nrs, target_boxes, local_exps):
         pot = self.output_zeros(actx)
         level_start_target_box_nrs = actx.to_numpy(level_start_target_box_nrs)
@@ -1077,7 +1076,7 @@ def eval_locals(self,
 
         return pot
 
-    def finalize_potentials(self, actx: PyOpenCLArrayContext, potentials):
+    def finalize_potentials(self, actx: ArrayContext, potentials):
         return potentials
 
 # }}}