diff --git a/examples/python/call-external.py b/examples/python/call-external.py index ad5615c7b..a2c19b855 100644 --- a/examples/python/call-external.py +++ b/examples/python/call-external.py @@ -1,4 +1,5 @@ import numpy as np +from constantdict import constantdict import loopy as lp from loopy.diagnostic import LoopyError @@ -30,9 +31,10 @@ def with_types(self, arg_id_to_dtype, callables_table): "types") return (self.copy(name_in_target=name_in_target, - arg_id_to_dtype={0: vec_dtype, - 1: vec_dtype, - -1: vec_dtype}), + arg_id_to_dtype=constantdict({ + 0: vec_dtype, + 1: vec_dtype, + -1: vec_dtype})), callables_table) def with_descrs(self, arg_id_to_descr, callables_table): diff --git a/loopy/kernel/function_interface.py b/loopy/kernel/function_interface.py index 799e5d91b..ad5a8ffcd 100644 --- a/loopy/kernel/function_interface.py +++ b/loopy/kernel/function_interface.py @@ -565,9 +565,9 @@ def with_types(self, arg_id_to_dtype, callables_table): "the function %s." % (self.name)) def with_descrs(self, arg_id_to_descr, clbl_inf_ctx): - - arg_id_to_descr[-1] = ValueArgDescriptor() - return (self.copy(arg_id_to_descr=arg_id_to_descr), + new_arg_id_to_descr = dict(arg_id_to_descr) + new_arg_id_to_descr[-1] = ValueArgDescriptor() + return (self.copy(arg_id_to_descr=constantdict(new_arg_id_to_descr)), clbl_inf_ctx) def get_hw_axes_sizes(self, arg_id_to_arg, space, callables_table): @@ -782,6 +782,7 @@ def with_descrs(self, arg_id_to_descr, clbl_inf_ctx): # arg_id_to_descr expressions provided are from the caller's namespace, # need to register + new_arg_id_to_descr = dict(arg_id_to_descr) kw_to_pos, pos_to_kw = get_kw_pos_association(self.subkernel) kw_to_callee_idx = {arg.name: i @@ -789,7 +790,7 @@ def with_descrs(self, arg_id_to_descr, clbl_inf_ctx): new_args = self.subkernel.args[:] - for arg_id, descr in arg_id_to_descr.items(): + for arg_id, descr in new_arg_id_to_descr.items(): if isinstance(arg_id, int): arg_id = pos_to_kw[arg_id] @@ -837,15 +838,15 @@ def with_descrs(self, arg_id_to_descr, clbl_inf_ctx): for arg in subkernel.args: kw = arg.name if isinstance(arg, ArrayBase): - arg_id_to_descr[kw] = ( + new_arg_id_to_descr[kw] = ( ArrayArgDescriptor(shape=arg.shape, dim_tags=arg.dim_tags, address_space=arg.address_space)) else: assert isinstance(arg, ValueArg) - arg_id_to_descr[kw] = ValueArgDescriptor() + new_arg_id_to_descr[kw] = ValueArgDescriptor() - arg_id_to_descr[kw_to_pos[kw]] = arg_id_to_descr[kw] + new_arg_id_to_descr[kw_to_pos[kw]] = new_arg_id_to_descr[kw] # }}} @@ -879,8 +880,8 @@ def with_added_arg(self, arg_dtype, arg_descr): arg_id_to_descr[kw_to_pos[var_name]] = arg_descr return (self.copy(subkernel=subknl, - arg_id_to_dtype=arg_id_to_dtype, - arg_id_to_descr=arg_id_to_descr), + arg_id_to_dtype=constantdict(arg_id_to_dtype), + arg_id_to_descr=constantdict(arg_id_to_descr)), var_name) else: @@ -902,7 +903,7 @@ def with_packing_for_args(self): address_space=AddressSpace.GLOBAL) return self.copy(subkernel=self.subkernel, - arg_id_to_descr=arg_id_to_descr) + arg_id_to_descr=constantdict(arg_id_to_descr)) def get_used_hw_axes(self, callables_table): gsize, lsize = self.subkernel.get_grid_size_upper_bounds(callables_table, diff --git a/loopy/library/function.py b/loopy/library/function.py index 8b61ad41a..5e3eeefb8 100644 --- a/loopy/library/function.py +++ b/loopy/library/function.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING import numpy as np +from constantdict import constantdict from loopy.diagnostic import LoopyError from loopy.kernel.function_interface import ScalarCallable @@ -38,12 +39,12 @@ class MakeTupleCallable(ScalarCallable): def with_types(self, arg_id_to_dtype, callables_table): - new_arg_id_to_dtype = arg_id_to_dtype.copy() + new_arg_id_to_dtype = dict(arg_id_to_dtype) for i in range(len(arg_id_to_dtype)): if i in arg_id_to_dtype and arg_id_to_dtype[i] is not None: new_arg_id_to_dtype[-i-1] = new_arg_id_to_dtype[i] - return (self.copy(arg_id_to_dtype=new_arg_id_to_dtype, + return (self.copy(arg_id_to_dtype=constantdict(new_arg_id_to_dtype), name_in_target="loopy_make_tuple"), callables_table) def with_descrs(self, arg_id_to_descr, callables_table): @@ -52,7 +53,7 @@ def with_descrs(self, arg_id_to_descr, callables_table): (-id-1, ValueArgDescriptor()) for id in arg_id_to_descr.keys()} return ( - self.copy(arg_id_to_descr=new_arg_id_to_descr), + self.copy(arg_id_to_descr=constantdict(new_arg_id_to_descr)), callables_table) @@ -63,7 +64,7 @@ def with_types(self, arg_id_to_dtype, callables_table): if dtype is not None} new_arg_id_to_dtype[-1] = NumpyType(np.int32) - return (self.copy(arg_id_to_dtype=new_arg_id_to_dtype), + return (self.copy(arg_id_to_dtype=constantdict(new_arg_id_to_dtype)), callables_table) def emit_call(self, expression_to_code_mapper, expression, target): diff --git a/loopy/library/random123.py b/loopy/library/random123.py index f65fa7600..bd6113eb4 100644 --- a/loopy/library/random123.py +++ b/loopy/library/random123.py @@ -29,6 +29,7 @@ from typing import TYPE_CHECKING import numpy as np +from constantdict import constantdict from mako.template import Template from pymbolic.typing import not_none @@ -221,7 +222,7 @@ def with_types(self, arg_id_to_dtype, callables_table): new_arg_id_to_dtype = {-1: ctr_dtype, -2: ctr_dtype, 0: ctr_dtype, 1: key_dtype} return ( - self.copy(arg_id_to_dtype=new_arg_id_to_dtype, + self.copy(arg_id_to_dtype=constantdict(new_arg_id_to_dtype), name_in_target=fn+"_gen"), callables_table) @@ -230,7 +231,7 @@ def with_types(self, arg_id_to_dtype, callables_table): rng_variant.width), -2: ctr_dtype, 0: ctr_dtype, 1: key_dtype} - return self.copy(arg_id_to_dtype=new_arg_id_to_dtype, + return self.copy(arg_id_to_dtype=constantdict(new_arg_id_to_dtype), name_in_target=name), callables_table elif name == fn + "_f64": @@ -238,7 +239,7 @@ def with_types(self, arg_id_to_dtype, callables_table): rng_variant.width), -2: ctr_dtype, 0: ctr_dtype, 1: key_dtype} - return self.copy(arg_id_to_dtype=new_arg_id_to_dtype, + return self.copy(arg_id_to_dtype=constantdict(new_arg_id_to_dtype), name_in_target=name), callables_table return (self.copy(arg_id_to_dtype=arg_id_to_dtype), diff --git a/loopy/library/reduction.py b/loopy/library/reduction.py index 6ddc3fb86..81cc151cb 100644 --- a/loopy/library/reduction.py +++ b/loopy/library/reduction.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING import numpy as np +from constantdict import constantdict from pymbolic import var from pymbolic.primitives import expr_dataclass @@ -580,21 +581,21 @@ def with_types(self, arg_id_to_dtype, callables_table): index_dtype = arg_id_to_dtype[1] result_dtypes = self.name.reduction_op.result_dtypes(scalar_dtype, # pylint: disable=no-member index_dtype) - new_arg_id_to_dtype = arg_id_to_dtype.copy() + new_arg_id_to_dtype = dict(arg_id_to_dtype) new_arg_id_to_dtype[-1] = result_dtypes[0] new_arg_id_to_dtype[-2] = result_dtypes[1] name_in_target = self.name.reduction_op.prefix(scalar_dtype, # pylint: disable=no-member index_dtype) + "_op" - return self.copy(arg_id_to_dtype=new_arg_id_to_dtype, + return self.copy(arg_id_to_dtype=constantdict(new_arg_id_to_dtype), name_in_target=name_in_target), callables_table def with_descrs(self, arg_id_to_descr, callables_table): from loopy.kernel.function_interface import ValueArgDescriptor - new_arg_id_to_descr = arg_id_to_descr.copy() + new_arg_id_to_descr = dict(arg_id_to_descr) new_arg_id_to_descr[-1] = ValueArgDescriptor() return ( - self.copy(arg_id_to_descr=arg_id_to_descr), + self.copy(arg_id_to_descr=constantdict(arg_id_to_descr)), callables_table) diff --git a/loopy/preprocess.py b/loopy/preprocess.py index 7600c97bf..d14aeca33 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -588,7 +588,7 @@ def map_call(self, expr, expn_state, assignees=None): # }}} # specializing the function according to the parameter description - new_clbl, self.clbl_inf_ctx = clbl.with_descrs(arg_id_to_descr, + new_clbl, self.clbl_inf_ctx = clbl.with_descrs(constantdict(arg_id_to_descr), self.clbl_inf_ctx) self.clbl_inf_ctx, new_func_id = (self.clbl_inf_ctx diff --git a/loopy/target/c/__init__.py b/loopy/target/c/__init__.py index 06cc208ed..6cd60d36b 100644 --- a/loopy/target/c/__init__.py +++ b/loopy/target/c/__init__.py @@ -28,6 +28,7 @@ from typing import TYPE_CHECKING, Any, Sequence, cast import numpy as np +from constantdict import constantdict import pymbolic.primitives as p from cgen import ( @@ -563,9 +564,9 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=name, - arg_id_to_dtype={ + arg_id_to_dtype=constantdict({ 0: NumpyType(dtype), - -1: NumpyType(result_dtype)}), + -1: NumpyType(result_dtype)})), callables_table) # binary functions @@ -607,7 +608,7 @@ def with_types(self, arg_id_to_dtype, callables_table): dtype = NumpyType(dtype) return ( self.copy(name_in_target=name, - arg_id_to_dtype={-1: dtype, 0: dtype, 1: dtype}), + arg_id_to_dtype=constantdict({-1: dtype, 0: dtype, 1: dtype})), callables_table) elif name in ["max", "min"]: @@ -632,9 +633,10 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=f"lpy_{name}_{dtype.name}", - arg_id_to_dtype={-1: NumpyType(dtype), - 0: NumpyType(dtype), - 1: NumpyType(dtype)}), + arg_id_to_dtype=constantdict({ + -1: NumpyType(dtype), + 0: NumpyType(dtype), + 1: NumpyType(dtype)})), callables_table) elif name == "isnan": for id in arg_id_to_dtype: @@ -662,9 +664,9 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy( name_in_target=name, - arg_id_to_dtype={ + arg_id_to_dtype=constantdict({ 0: NumpyType(dtype), - -1: NumpyType(np.int32)}), + -1: NumpyType(np.int32)})), callables_table) def generate_preambles(self, target): @@ -738,9 +740,10 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=name_in_target, - arg_id_to_dtype={-1: arg_id_to_dtype[1], - 0: NumpyType(np.int32), - 1: arg_id_to_dtype[1]}), + arg_id_to_dtype=constantdict({ + -1: arg_id_to_dtype[1], + 0: NumpyType(np.int32), + 1: arg_id_to_dtype[1]})), callables_table) else: raise NotImplementedError(f"with_types for '{name}'") diff --git a/loopy/target/opencl.py b/loopy/target/opencl.py index d14dd9e30..8aa504b63 100644 --- a/loopy/target/opencl.py +++ b/loopy/target/opencl.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Literal, Sequence import numpy as np +from constantdict import constantdict from pymbolic import var from pytools import memoize_method @@ -217,9 +218,10 @@ def with_types(self, arg_id_to_dtype, callables_table): # OpenCL C 2.2, Section 6.13.3: abs returns *u*gentype from loopy.types import to_unsigned_dtype return (self.copy(name_in_target=name, - arg_id_to_dtype={ + arg_id_to_dtype=constantdict({ 0: NumpyType(dtype), - -1: NumpyType(to_unsigned_dtype(dtype))}), + -1: NumpyType(to_unsigned_dtype(dtype)) + })), callables_table) elif dtype.kind == "f": name = "fabs" @@ -251,8 +253,10 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=name, - arg_id_to_dtype={0: NumpyType(dtype), -1: - NumpyType(dtype)}), + arg_id_to_dtype=constantdict({ + 0: NumpyType(dtype), + -1: NumpyType(dtype) + })), callables_table) # }}} @@ -283,7 +287,9 @@ def with_types(self, arg_id_to_dtype, callables_table): dtype = NumpyType(dtype) return ( self.copy(name_in_target=name, - arg_id_to_dtype={-1: dtype, 0: dtype, 1: dtype}), + arg_id_to_dtype=constantdict({ + -1: dtype, 0: dtype, 1: dtype + })), callables_table) elif name in ["max", "min"]: @@ -305,7 +311,9 @@ def with_types(self, arg_id_to_dtype, callables_table): dtype = NumpyType(common_dtype) return ( self.copy(name_in_target=name, - arg_id_to_dtype={-1: dtype, 0: dtype, 1: dtype}), + arg_id_to_dtype=constantdict({ + -1: dtype, 0: dtype, 1: dtype + })), callables_table) else: # Unsupported type. @@ -328,8 +336,9 @@ def with_types(self, arg_id_to_dtype, callables_table): dtype = arg_id_to_dtype[0] scalar_dtype, _offset, _field_name = dtype.numpy_dtype.fields["s0"] return ( - self.copy(name_in_target=name, arg_id_to_dtype={-1: - NumpyType(scalar_dtype), 0: dtype, 1: dtype}), + self.copy(name_in_target=name, arg_id_to_dtype=constantdict({ + -1: NumpyType(scalar_dtype), 0: dtype, 1: dtype + })), callables_table) elif name == "pow": @@ -352,8 +361,11 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=name, - arg_id_to_dtype={-1: result_dtype, - 0: common_dtype, 1: common_dtype}), + arg_id_to_dtype=constantdict({ + -1: result_dtype, + 0: common_dtype, + 1: common_dtype + })), callables_table) elif name in _CL_SIMPLE_MULTI_ARG_FUNCTIONS: @@ -379,8 +391,9 @@ def with_types(self, arg_id_to_dtype, callables_table): raise LoopyError("%s does not support complex numbers" % name) - updated_arg_id_to_dtype = {id: NumpyType(dtype) for id in range(-1, - num_args)} + updated_arg_id_to_dtype = constantdict({ + id: NumpyType(dtype) for id in range(-1, num_args) + }) return ( self.copy(name_in_target=name, @@ -409,8 +422,9 @@ def with_types(self, arg_id_to_dtype, callables_table): NumpyType(dtype), count) return ( - self.copy(name_in_target="(%s%d) " % (base_tp_name, count), - arg_id_to_dtype=updated_arg_id_to_dtype), + self.copy( + name_in_target="(%s%d) " % (base_tp_name, count), + arg_id_to_dtype=constantdict(updated_arg_id_to_dtype)), callables_table) # does not satisfy any of the conditions needed for specialization. diff --git a/loopy/target/pyopencl.py b/loopy/target/pyopencl.py index 9add453d7..9352f5381 100644 --- a/loopy/target/pyopencl.py +++ b/loopy/target/pyopencl.py @@ -30,6 +30,7 @@ from warnings import warn import numpy as np +from constantdict import constantdict import pymbolic.primitives as p from cgen import ( @@ -114,8 +115,10 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=f"{tpname}_{name}", - arg_id_to_dtype={0: dtype, -1: NumpyType( - np.dtype(dtype.numpy_dtype.type(0).real))}), + arg_id_to_dtype=constantdict({ + 0: dtype, + -1: NumpyType(np.dtype(dtype.numpy_dtype.type(0).real)) + })), callables_table) if name in ["real", "imag", "conj"]: @@ -124,7 +127,7 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy( name_in_target=f"_lpy_{name}_{tpname}", - arg_id_to_dtype={0: dtype, -1: dtype}), + arg_id_to_dtype=constantdict({0: dtype, -1: dtype})), callables_table) if name in ["sqrt", "exp", "log", @@ -142,7 +145,7 @@ def with_types(self, arg_id_to_dtype, callables_table): return ( self.copy(name_in_target=f"{tpname}_{name}", - arg_id_to_dtype={0: dtype, -1: dtype}), + arg_id_to_dtype=constantdict({0: dtype, -1: dtype})), callables_table) # fall back to pure OpenCL for real-valued arguments