From ba7104614fc034682be9c9ab1d75dbc9465c80df Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Thu, 21 Nov 2024 14:52:49 -0600 Subject: [PATCH 1/4] replace pyrsistent.PMap, immutables.Map with immutabledict --- doc/conf.py | 6 +----- loopy/codegen/__init__.py | 8 +++---- loopy/frontend/fortran/translator.py | 4 ++-- loopy/kernel/__init__.py | 4 ++-- loopy/kernel/data.py | 5 ++--- loopy/kernel/tools.py | 4 ++-- loopy/preprocess.py | 12 +++++------ loopy/schedule/__init__.py | 4 ++-- loopy/schedule/tools.py | 4 ++-- loopy/schedule/tree.py | 18 ++++++++-------- loopy/symbolic.py | 12 +++++------ loopy/target/c/c_execution.py | 4 ++-- loopy/target/execution.py | 28 ++++++++++++------------- loopy/target/pyopencl_execution.py | 4 ++-- loopy/tools.py | 12 ++++------- loopy/transform/buffer.py | 4 ++-- loopy/transform/callable.py | 6 +++--- loopy/transform/data.py | 4 ++-- loopy/transform/fusion.py | 8 +++---- loopy/transform/pack_and_unpack_args.py | 4 ++-- loopy/transform/precompute.py | 4 ++-- loopy/transform/realize_reduction.py | 4 ++-- loopy/transform/save.py | 4 ++-- loopy/translation_unit.py | 18 ++++++++-------- pyproject.toml | 2 -- 25 files changed, 87 insertions(+), 100 deletions(-) diff --git a/doc/conf.py b/doc/conf.py index c4a13c445..01afefc5d 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -32,7 +32,7 @@ "pyopencl": ("https://documen.tician.de/pyopencl", None), "cgen": ("https://documen.tician.de/cgen", None), "pymbolic": ("https://documen.tician.de/pymbolic", None), - "pyrsistent": ("https://pyrsistent.readthedocs.io/en/latest/", None), + "immutabledict": ("https://immutabledict.corenting.fr/", None), } # Some modules need to import things just so that sphinx can resolve symbols in @@ -57,10 +57,6 @@ ["py:class", r"numpy\.float[0-9]+"], ["py:class", r"numpy\.complex[0-9]+"], - # As of 2022-06-22, it doesn't look like there's sphinx documentation - # available. - ["py:class", r"immutables\.(.+)"], - # Reference not found from ""? I'm not even sure where to look. ["py:class", r"Expression"], ] diff --git a/loopy/codegen/__init__.py b/loopy/codegen/__init__.py index 2e39d89bd..270a8310f 100644 --- a/loopy/codegen/__init__.py +++ b/loopy/codegen/__init__.py @@ -35,7 +35,7 @@ Union, ) -from immutables import Map +from immutabledict import immutabledict from loopy.codegen.result import CodeGenerationResult from loopy.library.reduction import ReductionOpFunction @@ -207,7 +207,7 @@ class CodeGenerationState: seen_functions: Set[SeenFunction] seen_atomic_dtypes: Set[LoopyType] - var_subst_map: Map[str, ExpressionT] + var_subst_map: immutabledict[str, ExpressionT] allow_complex: bool callables_table: CallablesTable is_entrypoint: bool @@ -418,7 +418,7 @@ def generate_code_for_a_single_kernel(kernel, callables_table, target, seen_dtypes=seen_dtypes, seen_functions=seen_functions, seen_atomic_dtypes=seen_atomic_dtypes, - var_subst_map=Map(), + var_subst_map=immutabledict(), allow_complex=allow_complex, var_name_generator=kernel.get_var_name_generator(), is_generating_device_code=False, @@ -519,7 +519,7 @@ def diverge_callee_entrypoints(program): new_callables[name] = clbl - return program.copy(callables_table=Map(new_callables)) + return program.copy(callables_table=immutabledict(new_callables)) @dataclass(frozen=True) diff --git a/loopy/frontend/fortran/translator.py b/loopy/frontend/fortran/translator.py index fc9eace87..283122dff 100644 --- a/loopy/frontend/fortran/translator.py +++ b/loopy/frontend/fortran/translator.py @@ -25,7 +25,7 @@ from warnings import warn import numpy as np -from immutables import Map +from immutabledict import immutabledict import islpy as isl from islpy import dim_type @@ -330,7 +330,7 @@ def specialize_fortran_division(t_unit): new_callables[name] = clbl - return t_unit.copy(callables_table=Map(new_callables)) + return t_unit.copy(callables_table=immutabledict(new_callables)) # }}} diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py index 967640260..40d8f1670 100644 --- a/loopy/kernel/__init__.py +++ b/loopy/kernel/__init__.py @@ -45,7 +45,7 @@ from warnings import warn import numpy as np -from immutables import Map +from immutabledict import immutabledict import islpy as isl from islpy import dim_type @@ -177,7 +177,7 @@ class LoopKernel(Taggable): Callable[["LoopKernel", str], Optional[Tuple[LoopyType, str]]]] = () linearization: Optional[Sequence[ScheduleItem]] = None iname_slab_increments: Mapping[InameStr, Tuple[int, int]] = field( - default_factory=Map) + default_factory=immutabledict) """ A mapping from inames to (lower_incr, upper_incr) tuples that will be separated out in the execution to generate diff --git a/loopy/kernel/data.py b/loopy/kernel/data.py index 5d1de0e5d..5a5cd4d25 100644 --- a/loopy/kernel/data.py +++ b/loopy/kernel/data.py @@ -25,7 +25,7 @@ THE SOFTWARE. """ - +from collections.abc import Mapping from dataclasses import dataclass, replace from enum import IntEnum from sys import intern @@ -43,7 +43,6 @@ import numpy # FIXME: imported as numpy to allow sphinx to resolve things import numpy as np -from immutables import Map from pymbolic import ArithmeticExpressionT, Variable from pytools import ImmutableRecord @@ -434,7 +433,7 @@ class _ArraySeparationInfo: should be used to realize this array. """ sep_axis_indices_set: FrozenSet[int] - subarray_names: Map[Tuple[int, ...], str] + subarray_names: Mapping[Tuple[int, ...], str] class ArrayArg(ArrayBase, KernelArgument): diff --git a/loopy/kernel/tools.py b/loopy/kernel/tools.py index 9a14aedd5..c9e58a659 100644 --- a/loopy/kernel/tools.py +++ b/loopy/kernel/tools.py @@ -2084,7 +2084,7 @@ def get_call_graph(t_unit, only_kernel_callables=False): :arg t_unit: An instance of :class:`TranslationUnit`. """ - from pyrsistent import pmap + from immutabledict import immutabledict from loopy.kernel import KernelState @@ -2111,7 +2111,7 @@ def get_call_graph(t_unit, only_kernel_callables=False): call_graph[name] = clbl.get_called_callables(t_unit.callables_table, recursive=False) - return pmap(call_graph) + return immutabledict(call_graph) # }}} diff --git a/loopy/preprocess.py b/loopy/preprocess.py index 3293e9a1e..b1b12efe3 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -29,7 +29,7 @@ from functools import partial import numpy as np -from immutables import Map +from immutabledict import immutabledict from pytools import ProcessLogger @@ -191,7 +191,7 @@ def make_arrays_for_sep_arrays(kernel: LoopKernel) -> LoopKernel: sep_info = _ArraySeparationInfo( sep_axis_indices_set=sep_axis_indices_set, - subarray_names=Map({ + subarray_names=immutabledict({ ind: vng(f"{arg.name}_s{'_'.join(str(i) for i in ind)}") for ind in np.ndindex(*cast(List[int], sep_shape))})) @@ -599,8 +599,6 @@ def map_call_with_kwargs(self, expr): raise NotImplementedError def __call__(self, expr, kernel, insn, assignees=None): - import immutables - from loopy.kernel.data import InstructionBase from loopy.symbolic import ExpansionState, UncachedIdentityMapper assert insn is None or isinstance(insn, InstructionBase) @@ -610,7 +608,7 @@ def __call__(self, expr, kernel, insn, assignees=None): kernel=kernel, instruction=insn, stack=(), - arg_context=immutables.Map()), assignees=assignees) + arg_context=immutabledict()), assignees=assignees) def map_kernel(self, kernel): @@ -744,7 +742,7 @@ def filter_reachable_callables(t_unit): t_unit.entrypoints) new_callables = {name: clbl for name, clbl in t_unit.callables_table.items() if name in (reachable_function_ids | t_unit.entrypoints)} - return t_unit.copy(callables_table=Map(new_callables)) + return t_unit.copy(callables_table=immutabledict(new_callables)) def _preprocess_single_kernel(kernel: LoopKernel, is_entrypoint: bool) -> LoopKernel: @@ -869,7 +867,7 @@ def preprocess_program(t_unit: TranslationUnit) -> TranslationUnit: new_callables[func_id] = in_knl_callable - t_unit = t_unit.copy(callables_table=Map(new_callables)) + t_unit = t_unit.copy(callables_table=immutabledict(new_callables)) # }}} diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py index 2460f5ed7..d3b1776e0 100644 --- a/loopy/schedule/__init__.py +++ b/loopy/schedule/__init__.py @@ -35,7 +35,7 @@ TypeVar, ) -from immutables import Map +from immutabledict import immutabledict import islpy as isl from pytools import ImmutableRecord, MinRecursionLimit, ProcessLogger @@ -2481,7 +2481,7 @@ def linearize(t_unit: TranslationUnit) -> TranslationUnit: else: raise NotImplementedError(type(clbl)) - return t_unit.copy(callables_table=Map(new_callables)) + return t_unit.copy(callables_table=immutabledict(new_callables)) # vim: foldmethod=marker diff --git a/loopy/schedule/tools.py b/loopy/schedule/tools.py index 3858462b1..a64d7f8f4 100644 --- a/loopy/schedule/tools.py +++ b/loopy/schedule/tools.py @@ -57,7 +57,7 @@ from functools import cached_property, reduce from typing import AbstractSet, Dict, FrozenSet, List, Sequence, Set, Tuple -from immutables import Map +from immutabledict import immutabledict from typing_extensions import TypeAlias import islpy as isl @@ -1048,7 +1048,7 @@ def _get_iname_to_tree_node_id_from_partial_loop_nest_tree( for iname in node: iname_to_tree_node_id[iname] = node - return Map(iname_to_tree_node_id) + return immutabledict(iname_to_tree_node_id) def get_loop_tree(kernel: LoopKernel) -> LoopTree: diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index e98724f83..eb27b55ad 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -39,7 +39,7 @@ from functools import cached_property from typing import Generic, TypeVar -from immutables import Map +from immutabledict import immutabledict from pytools import memoize_method @@ -71,13 +71,13 @@ class Tree(Generic[NodeT]): this allocates a new stack frame for each iteration of the operation. """ - _parent_to_children: Map[NodeT, tuple[NodeT, ...]] - _child_to_parent: Map[NodeT, NodeT | None] + _parent_to_children: immutabledict[NodeT, tuple[NodeT, ...]] + _child_to_parent: immutabledict[NodeT, NodeT | None] @staticmethod def from_root(root: NodeT) -> Tree[NodeT]: - return Tree(Map({root: ()}), - Map({root: None})) + return Tree(immutabledict({root: ()}), + immutabledict({root: None})) @cached_property def root(self) -> NodeT: @@ -182,7 +182,7 @@ def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: # {{{ update child to parent - child_to_parent_mut = self._child_to_parent.mutate() + child_to_parent_mut = dict(self._child_to_parent) del child_to_parent_mut[node] child_to_parent_mut[new_node] = parent @@ -193,7 +193,7 @@ def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: # {{{ update parent_to_children - parent_to_children_mut = self._parent_to_children.mutate() + parent_to_children_mut = dict(self._parent_to_children) del parent_to_children_mut[node] parent_to_children_mut[new_node] = children @@ -205,8 +205,8 @@ def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: # }}} - return Tree(parent_to_children_mut.finish(), - child_to_parent_mut.finish()) + return Tree(immutabledict(parent_to_children_mut), + immutabledict(child_to_parent_mut)) def move_node(self, node: NodeT, new_parent: NodeT | None) -> Tree[NodeT]: """ diff --git a/loopy/symbolic.py b/loopy/symbolic.py index ad502e1a5..678cdc01d 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -39,8 +39,8 @@ ) from warnings import warn -import immutables import numpy as np +from immutabledict import immutabledict import islpy as isl import pymbolic.primitives # FIXME: also import by full name to allow sphinx to resolve @@ -1026,12 +1026,12 @@ class ExpansionState(ImmutableRecord): a dict representing current argument values """ def __init__(self, kernel, instruction, stack, arg_context): - if not isinstance(arg_context, immutables.Map): + if not isinstance(arg_context, immutabledict): warn(f"Got a {type(arg_context)} for arg_context," - " expected `immutables.Map`. This is deprecated" + " expected `immutabledict`. This is deprecated" " and will result in an error from 2023.", DeprecationWarning, stacklevel=2) - arg_context = immutables.Map(arg_context) + arg_context = immutabledict(arg_context) super().__init__(kernel=kernel, instruction=instruction, stack=stack, @@ -1264,7 +1264,7 @@ def make_new_arg_context( from pymbolic.mapper.substitutor import make_subst_func arg_subst_map = SubstitutionMapper(make_subst_func(arg_context)) - return immutables.Map({ + return immutabledict({ formal_arg_name: arg_subst_map(arg_value) for formal_arg_name, arg_value in zip(arg_names, arguments)}) @@ -1307,7 +1307,7 @@ def __call__(self, expr, kernel, insn): kernel=kernel, instruction=insn, stack=(), - arg_context=immutables.Map())) + arg_context=immutabledict())) def map_instruction(self, kernel, insn): return insn diff --git a/loopy/target/c/c_execution.py b/loopy/target/c/c_execution.py index 9cde501a7..1d1f44fcf 100644 --- a/loopy/target/c/c_execution.py +++ b/loopy/target/c/c_execution.py @@ -24,13 +24,13 @@ import logging import os import tempfile +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Callable, Optional, Sequence, Tuple, Union import numpy as np from codepy.jit import compile_from_string from codepy.toolchain import GCCToolchain, ToolchainGuessError, guess_toolchain -from immutables import Map from pytools import memoize_method from pytools.codegen import CodeGenerator, Indentation @@ -493,7 +493,7 @@ def get_wrapper_generator(self): @memoize_method def translation_unit_info(self, - arg_to_dtype: Optional[Map[str, LoopyType]] = None) -> _KernelInfo: + arg_to_dtype: Optional[Mapping[str, LoopyType]] = None) -> _KernelInfo: t_unit = self.get_typed_and_scheduled_translation_unit(arg_to_dtype) from loopy.codegen import generate_code_v2 diff --git a/loopy/target/execution.py b/loopy/target/execution.py index 2443a1420..829791b61 100644 --- a/loopy/target/execution.py +++ b/loopy/target/execution.py @@ -39,7 +39,7 @@ cast, ) -from immutables import Map +from immutabledict import immutabledict from pymbolic import Variable, var from pytools.codegen import CodeGenerator, Indentation @@ -817,7 +817,7 @@ def check_for_required_array_arguments(self, input_args): "your argument.") def get_typed_and_scheduled_translation_unit_uncached( - self, arg_to_dtype: Optional[Map[str, LoopyType]] + self, arg_to_dtype: Optional[Mapping[str, LoopyType]] ) -> TranslationUnit: t_unit = self.t_unit @@ -827,15 +827,15 @@ def get_typed_and_scheduled_translation_unit_uncached( # FIXME: This is not so nice. This transfers types from the # subarrays of sep-tagged arrays to the 'main' array, because # type inference fails otherwise. - with arg_to_dtype.mutate() as mm: - for name, sep_info in self.sep_info.items(): - if entry_knl.arg_dict[name].dtype is None: - for sep_name in sep_info.subarray_names.values(): - if sep_name in arg_to_dtype: - mm.set(name, arg_to_dtype[sep_name]) - del mm[sep_name] + mm = dict(arg_to_dtype) + for name, sep_info in self.sep_info.items(): + if entry_knl.arg_dict[name].dtype is None: + for sep_name in sep_info.subarray_names.values(): + if sep_name in arg_to_dtype: + mm[name] = arg_to_dtype[sep_name] + del mm[sep_name] - arg_to_dtype = mm.finish() + arg_to_dtype = immutabledict(mm) from loopy.kernel.tools import add_dtypes t_unit = t_unit.with_kernel(add_dtypes(entry_knl, arg_to_dtype)) @@ -854,7 +854,7 @@ def get_typed_and_scheduled_translation_unit_uncached( return t_unit def get_typed_and_scheduled_translation_unit( - self, arg_to_dtype: Optional[Map[str, LoopyType]] + self, arg_to_dtype: Optional[Mapping[str, LoopyType]] ) -> TranslationUnit: from loopy import CACHING_ENABLED @@ -876,7 +876,7 @@ def get_typed_and_scheduled_translation_unit( return t_unit - def arg_to_dtype(self, kwargs) -> Optional[Map[str, LoopyType]]: + def arg_to_dtype(self, kwargs) -> Optional[immutabledict[str, LoopyType]]: if not self.has_runtime_typed_args: return None @@ -893,7 +893,7 @@ def arg_to_dtype(self, kwargs) -> Optional[Map[str, LoopyType]]: else: arg_to_dtype[arg_name] = NumpyType(dtype) - return Map(arg_to_dtype) + return immutabledict(arg_to_dtype) # {{{ debugging aids @@ -904,7 +904,7 @@ def get_highlighted_code(self, entrypoint, arg_to_dtype=None, code=None): def get_code( self, entrypoint: str, - arg_to_dtype: Optional[Map[str, LoopyType]] = None) -> str: + arg_to_dtype: Optional[Mapping[str, LoopyType]] = None) -> str: kernel = self.get_typed_and_scheduled_translation_unit(arg_to_dtype) from loopy.codegen import generate_code_v2 diff --git a/loopy/target/pyopencl_execution.py b/loopy/target/pyopencl_execution.py index be859ab70..b0879ea94 100644 --- a/loopy/target/pyopencl_execution.py +++ b/loopy/target/pyopencl_execution.py @@ -22,11 +22,11 @@ import logging +from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence import numpy as np -from immutables import Map from pytools import memoize_method from pytools.codegen import CodeGenerator, Indentation @@ -307,7 +307,7 @@ def get_wrapper_generator(self): @memoize_method def translation_unit_info( self, - arg_to_dtype: Optional[Map[str, LoopyType]] = None) -> _KernelInfo: + arg_to_dtype: Optional[Mapping[str, LoopyType]] = None) -> _KernelInfo: t_unit = self.get_typed_and_scheduled_translation_unit(arg_to_dtype) # FIXME: now just need to add the types to the arguments diff --git a/loopy/tools.py b/loopy/tools.py index bb4904bf2..6a43f9e9d 100644 --- a/loopy/tools.py +++ b/loopy/tools.py @@ -27,7 +27,7 @@ from typing import List import numpy as np -from immutables import Map +from immutabledict import immutabledict import islpy as isl from pytools import ProcessLogger, memoize_method @@ -78,15 +78,11 @@ def update_for_BasicSet(self, key_hash, key): # noqa key_hash.update(prn.get_str().encode("utf8")) def update_for_Map(self, key_hash, key): # noqa - if isinstance(key, Map): - self.update_for_dict(key_hash, key) - elif isinstance(key, isl.Map): + if isinstance(key, isl.Map): self.update_for_BasicSet(key_hash, key) else: raise AssertionError() - update_for_PMap = update_for_dict # noqa: N815 - # }}} @@ -811,7 +807,7 @@ def t_unit_to_python(t_unit, var_name="t_unit", .callables_table)) for name, clbl in t_unit.callables_table.items() if isinstance(clbl, CallableKernel)} - t_unit = t_unit.copy(callables_table=Map(new_callables)) + t_unit = t_unit.copy(callables_table=immutabledict(new_callables)) knl_python_code_srcs = [_kernel_to_python(clbl.subkernel, name in t_unit.entrypoints, @@ -826,7 +822,7 @@ def t_unit_to_python(t_unit, var_name="t_unit", "import loopy as lp", "import numpy as np", "from pymbolic.primitives import *", - "import immutables", + "from immutabledict import immutabledict", ]) body_str = "\n".join(knl_python_code_srcs + ["\n", merge_stmt]) diff --git a/loopy/transform/buffer.py b/loopy/transform/buffer.py index c8339f550..cd6b1fdc3 100644 --- a/loopy/transform/buffer.py +++ b/loopy/transform/buffer.py @@ -22,7 +22,7 @@ import logging -from immutables import Map +from immutabledict import immutabledict from pymbolic import var from pymbolic.mapper.substitutor import make_subst_func @@ -537,7 +537,7 @@ def buffer_array(program, *args, **kwargs): new_callables[func_id] = clbl - return program.copy(callables_table=Map(new_callables)) + return program.copy(callables_table=immutabledict(new_callables)) # vim: foldmethod=marker diff --git a/loopy/transform/callable.py b/loopy/transform/callable.py index 0210eaee2..8276aea6f 100644 --- a/loopy/transform/callable.py +++ b/loopy/transform/callable.py @@ -22,7 +22,7 @@ from collections.abc import Sequence -from immutables import Map +from immutabledict import immutabledict import islpy as isl from pytools import UniqueNameGenerator @@ -121,7 +121,7 @@ def merge(translation_units: Sequence[TranslationUnit]) -> TranslationUnit: return TranslationUnit( entrypoints=frozenset().union(*( t.entrypoints or frozenset() for t in translation_units)), - callables_table=Map(callables_table), + callables_table=immutabledict(callables_table), target=translation_units[0].target) @@ -597,7 +597,7 @@ def rename_callable( new_entrypoints = ((new_entrypoints | frozenset([new_name])) - frozenset([old_name])) - return t_unit.copy(callables_table=Map(new_callables_table), + return t_unit.copy(callables_table=immutabledict(new_callables_table), entrypoints=new_entrypoints) # }}} diff --git a/loopy/transform/data.py b/loopy/transform/data.py index c63604f8c..32b5026e9 100644 --- a/loopy/transform/data.py +++ b/loopy/transform/data.py @@ -25,7 +25,7 @@ from warnings import warn import numpy as np -from immutables import Map +from immutabledict import immutabledict from islpy import dim_type from pytools import MovedFunctionDeprecationWrapper @@ -425,7 +425,7 @@ def add_prefetch(t_unit, new_callables[func_id] = in_knl_callable - return t_unit.copy(callables_table=Map(new_callables)) + return t_unit.copy(callables_table=immutabledict(new_callables)) # }}} diff --git a/loopy/transform/fusion.py b/loopy/transform/fusion.py index 8e047c036..fe314e5a4 100644 --- a/loopy/transform/fusion.py +++ b/loopy/transform/fusion.py @@ -21,7 +21,7 @@ """ -from immutables import Map +from immutabledict import immutabledict import islpy as isl from islpy import dim_type @@ -117,8 +117,8 @@ def _merge_dicts(item_name, dict_a, dict_b): else: result[k] = v - if isinstance(dict_a, Map): - return Map(result) + if isinstance(dict_a, immutabledict): + return immutabledict(result) else: return result @@ -450,7 +450,7 @@ def fuse_kernels(kernels, suffixes=None, data_flow=None): new_callables[result.name] = CallableKernel(result) - return TranslationUnit(callables_table=Map(new_callables), + return TranslationUnit(callables_table=immutabledict(new_callables), target=result.target, entrypoints=frozenset([result.name])) diff --git a/loopy/transform/pack_and_unpack_args.py b/loopy/transform/pack_and_unpack_args.py index 2a82952c2..56f78eace 100644 --- a/loopy/transform/pack_and_unpack_args.py +++ b/loopy/transform/pack_and_unpack_args.py @@ -20,7 +20,7 @@ THE SOFTWARE. """ -from immutables import Map +from immutabledict import immutabledict from loopy.diagnostic import LoopyError from loopy.kernel import LoopKernel @@ -339,6 +339,6 @@ def pack_and_unpack_args_for_call(program, *args, **kwargs): new_callables[func_id] = in_knl_callable - return program.copy(callables_table=Map(new_callables)) + return program.copy(callables_table=immutabledict(new_callables)) # vim: foldmethod=marker diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index b0fbb5468..1d5994a05 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -25,7 +25,7 @@ from typing import FrozenSet, List, Optional, Sequence, Type, Union, cast import numpy as np -from immutables import Map +from immutabledict import immutabledict import islpy as isl from pymbolic import ArithmeticExpressionT, var @@ -1181,6 +1181,6 @@ def precompute(program, *args, **kwargs): new_callables[func_id] = clbl - return program.copy(callables_table=Map(new_callables)) + return program.copy(callables_table=immutabledict(new_callables)) # vim: foldmethod=marker diff --git a/loopy/transform/realize_reduction.py b/loopy/transform/realize_reduction.py index 7d1f3c870..88283360f 100644 --- a/loopy/transform/realize_reduction.py +++ b/loopy/transform/realize_reduction.py @@ -31,7 +31,7 @@ logger = logging.getLogger(__name__) -from immutables import Map +from immutabledict import immutabledict import islpy as isl from pymbolic.primitives import Expression @@ -2176,6 +2176,6 @@ def realize_reduction(t_unit, *args, **kwargs): subkernel=new_knl) callables_table[knl.name] = in_knl_callable - return t_unit.copy(callables_table=Map(callables_table)) + return t_unit.copy(callables_table=immutabledict(callables_table)) # vim: foldmethod=marker diff --git a/loopy/transform/save.py b/loopy/transform/save.py index bd25dec36..54147f9d2 100644 --- a/loopy/transform/save.py +++ b/loopy/transform/save.py @@ -23,7 +23,7 @@ import logging from functools import cached_property -from immutables import Map +from immutabledict import immutabledict from pytools import Record, memoize_method @@ -249,7 +249,7 @@ def __init__(self, kernel, callables_table): from collections import defaultdict self.insns_to_insert = [] self.insns_to_update = {} - self.updated_iname_objs = Map() + self.updated_iname_objs = immutabledict() self.updated_temporary_variables = {} # temporary name -> save or reload insn ids diff --git a/loopy/translation_unit.py b/loopy/translation_unit.py index 4afdfcef7..9ed692fa4 100644 --- a/loopy/translation_unit.py +++ b/loopy/translation_unit.py @@ -39,7 +39,7 @@ ) from warnings import warn -from immutables import Map +from immutabledict import immutabledict from typing_extensions import Concatenate, ParamSpec, Self from pymbolic.primitives import Call, Variable @@ -177,7 +177,7 @@ def map_call_with_kwargs(self, expr): # {{{ translation unit FunctionIdT = Union[str, ReductionOpFunction] -ConcreteCallablesTable = Map[FunctionIdT, InKernelCallable] +ConcreteCallablesTable = immutabledict[FunctionIdT, InKernelCallable] CallablesTable = Mapping[FunctionIdT, InKernelCallable] @@ -205,7 +205,7 @@ class TranslationUnit: .. attribute:: callables_table - An instance of :class:`pyrsistent.PMap` mapping the function + An instance of :class:`immutabledict.immutabledict` mapping the function identifiers in a kernel to their associated instances of :class:`~loopy.kernel.function_interface.InKernelCallable`. @@ -242,7 +242,7 @@ class TranslationUnit: def __post_init__(self): assert isinstance(self.entrypoints, abc_Set) - assert isinstance(self.callables_table, Map) + assert isinstance(self.callables_table, immutabledict) object.__setattr__(self, "_program_executor_cache", {}) @@ -272,7 +272,7 @@ def copy(self, **kwargs: Any) -> Self: new_callables[func_id] = clbl t_unit = replace( - self, callables_table=Map(new_callables), target=target) + self, callables_table=immutabledict(new_callables), target=target) return t_unit @@ -740,7 +740,7 @@ def finish_program(self, program): # }}} - return program.copy(callables_table=Map(new_callables)) + return program.copy(callables_table=immutabledict(new_callables)) def __getitem__(self, name): result = self.callables[name] @@ -761,7 +761,7 @@ def make_program(kernel: LoopKernel) -> TranslationUnit: """ return TranslationUnit( - callables_table=Map({ + callables_table=immutabledict({ kernel.name: CallableKernel(kernel)}), target=kernel.target, entrypoints=frozenset()) @@ -821,7 +821,7 @@ def _collective_transform( new_callables[func_id] = clbl - return t_unit.copy(callables_table=Map(new_callables)) + return t_unit.copy(callables_table=immutabledict(new_callables)) elif isinstance(t_unit_or_kernel, LoopKernel): kernel = t_unit_or_kernel return transform(kernel, *args, **kwargs) @@ -917,7 +917,7 @@ def resolve_callables(t_unit: TranslationUnit) -> TranslationUnit: else: raise NotImplementedError(f"{type(clbl)}") - t_unit = t_unit.copy(callables_table=Map(callables_table)) + t_unit = t_unit.copy(callables_table=immutabledict(callables_table)) validate_kernel_call_sites(t_unit) diff --git a/pyproject.toml b/pyproject.toml index 3204163f0..d44362928 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,8 +43,6 @@ dependencies = [ "codepy>=2017.1", "colorama", "Mako", - "pyrsistent", - "immutables", "immutabledict", "typing-extensions>=4", From 6dadbe7dc0cc942314a2acb253185060f2e288a6 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Dec 2024 12:57:25 -0600 Subject: [PATCH 2/4] go back to immutables.Map for Tree --- loopy/schedule/tree.py | 18 +++++++++--------- pyproject.toml | 1 + 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index 093bd6b2d..cc580713c 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -40,7 +40,7 @@ from functools import cached_property, reduce from typing import Generic, TypeVar -from immutabledict import immutabledict +from immutables import Map from pytools import memoize_method @@ -74,13 +74,13 @@ class Tree(Generic[NodeT]): this allocates a new stack frame for each iteration of the operation. """ - _parent_to_children: immutabledict[NodeT, tuple[NodeT, ...]] - _child_to_parent: immutabledict[NodeT, NodeT | None] + _parent_to_children: Map[NodeT, tuple[NodeT, ...]] + _child_to_parent: Map[NodeT, NodeT | None] @staticmethod def from_root(root: NodeT) -> Tree[NodeT]: - return Tree(immutabledict({root: ()}), - immutabledict({root: None})) + return Tree(Map({root: ()}), + Map({root: None})) @cached_property def root(self) -> NodeT: @@ -172,7 +172,7 @@ def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: # {{{ update child to parent - child_to_parent_mut = dict(self._child_to_parent) + child_to_parent_mut = self._child_to_parent.mutate() del child_to_parent_mut[node] child_to_parent_mut[new_node] = parent @@ -183,7 +183,7 @@ def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: # {{{ update parent_to_children - parent_to_children_mut = dict(self._parent_to_children) + parent_to_children_mut = self._parent_to_children.mutate() del parent_to_children_mut[node] parent_to_children_mut[new_node] = children @@ -195,8 +195,8 @@ def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: # }}} - return Tree(immutabledict(parent_to_children_mut), - immutabledict(child_to_parent_mut)) + return Tree(parent_to_children_mut.finish(), + child_to_parent_mut.finish()) def move_node(self, node: NodeT, new_parent: NodeT | None) -> Tree[NodeT]: """ diff --git a/pyproject.toml b/pyproject.toml index bf0b74f27..fa0348ae6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ "colorama", "Mako", "immutabledict", + "immutables", "typing-extensions>=4", ] From 91c14a1c910f0f089610992f3b79985122e0b64f Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Dec 2024 13:04:12 -0600 Subject: [PATCH 3/4] ruff fixes --- loopy/schedule/tree.py | 16 ++++++++-------- loopy/tools.py | 4 ++-- loopy/transform/precompute.py | 3 +-- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/loopy/schedule/tree.py b/loopy/schedule/tree.py index cc580713c..3861aa75c 100644 --- a/loopy/schedule/tree.py +++ b/loopy/schedule/tree.py @@ -149,11 +149,11 @@ def add_node(self, node: NodeT, parent: NodeT) -> Tree[NodeT]: siblings = self._parent_to_children[parent] - _parent_to_children_mut = self._parent_to_children.mutate() - _parent_to_children_mut[parent] = (*siblings, node) - _parent_to_children_mut[node] = () + parent_to_children_mut = self._parent_to_children.mutate() + parent_to_children_mut[parent] = (*siblings, node) + parent_to_children_mut[node] = () - return Tree(_parent_to_children_mut.finish(), + return Tree(parent_to_children_mut.finish(), self._child_to_parent.set(node, parent)) def replace_node(self, node: NodeT, new_node: NodeT) -> Tree[NodeT]: @@ -223,11 +223,11 @@ def move_node(self, node: NodeT, new_parent: NodeT | None) -> Tree[NodeT]: parents_new_children = tuple(frozenset(siblings) - frozenset([node])) new_parents_children = (*self.children(new_parent), node) - _parent_to_children_mut = self._parent_to_children.mutate() - _parent_to_children_mut[parent] = parents_new_children - _parent_to_children_mut[new_parent] = new_parents_children + parent_to_children_mut = self._parent_to_children.mutate() + parent_to_children_mut[parent] = parents_new_children + parent_to_children_mut[new_parent] = new_parents_children - return Tree(_parent_to_children_mut.finish(), + return Tree(parent_to_children_mut.finish(), self._child_to_parent.set(node, new_parent)) def __str__(self) -> str: diff --git a/loopy/tools.py b/loopy/tools.py index 49bf285f8..3a4b04e97 100644 --- a/loopy/tools.py +++ b/loopy/tools.py @@ -301,10 +301,10 @@ def empty_aligned(shape, dtype, order="C", n=64): # We now need to know how to offset base_ary # so it is correctly aligned - _array_aligned_offset = (n-address_from_numpy(base_ary)) % n + array_aligned_offset = (n-address_from_numpy(base_ary)) % n array = np.frombuffer( - base_ary[_array_aligned_offset:_array_aligned_offset-n].data, + base_ary[array_aligned_offset:array_aligned_offset-n].data, dtype=dtype).reshape(shape, order=order) return array diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index f41d0fb46..bfd61dee3 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -772,8 +772,7 @@ def precompute_for_single_kernel( if abm.non1_storage_axis_flags[i]: non1_storage_axis_names.append(saxis) else: - if saxis in new_iname_to_tag: - del new_iname_to_tag[saxis] + new_iname_to_tag.pop(saxis, None) if saxis in preexisting_precompute_inames: raise LoopyError("precompute axis %d (1-based) was " From 9f1132bfdf4ad0ec46fdcbc98b362bf9ee458883 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 13 Dec 2024 13:06:21 -0600 Subject: [PATCH 4/4] spelling --- doc/conf.py | 4 ++++ pyproject.toml | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/doc/conf.py b/doc/conf.py index d4371837f..271824c7e 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -57,6 +57,10 @@ ["py:class", r"numpy\.float[0-9]+"], ["py:class", r"numpy\.complex[0-9]+"], + # As of 2022-06-22, it doesn't look like there's sphinx documentation + # available. + ["py:class", r"immutables\.(.+)"], + # Reference not found from ""? I'm not even sure where to look. ["py:class", r"ExpressionNode"], ] diff --git a/pyproject.toml b/pyproject.toml index fa0348ae6..f1492df00 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,7 +110,7 @@ extend-ignore = [ # FIXME "NPY002", # numpy rng - "C408", # unnecssary dict() -> literal + "C408", # unnecessary dict() -> literal ] [tool.ruff.lint.per-file-ignores]