diff --git a/pyopencl/tools.py b/pyopencl/tools.py index b16de3f64cd033cc0f4b3d2ce19b65a9e169d0ed..67302d8a970217a6473087eaf400f47ee8eaf141 100644 --- a/pyopencl/tools.py +++ b/pyopencl/tools.py @@ -32,7 +32,6 @@ from sys import intern # Do not add a pyopencl import here: This will add an import cycle. import numpy as np -from decorator import decorator from pytools import memoize, memoize_method from pyopencl._cl import bitlog2 # noqa: F401 from pytools.persistent_dict import KeyBuilder as KeyBuilderBase @@ -73,32 +72,41 @@ from pyopencl._cl import ( # noqa _first_arg_dependent_caches = [] -@decorator -def first_arg_dependent_memoize(func, cl_object, *args): - """Provides memoization for a function. Typically used to cache - things that get created inside a :class:`pyopencl.Context`, e.g. programs - and kernels. Assumes that the first argument of the decorated function is - an OpenCL object that might go away, such as a :class:`pyopencl.Context` or - a :class:`pyopencl.CommandQueue`, and based on which we might want to clear - the cache. +def first_arg_dependent_memoize(func): + def wrapper(cl_object, *args, **kwargs): + """Provides memoization for a function. Typically used to cache + things that get created inside a :class:`pyopencl.Context`, e.g. programs + and kernels. Assumes that the first argument of the decorated function is + an OpenCL object that might go away, such as a :class:`pyopencl.Context` or + a :class:`pyopencl.CommandQueue`, and based on which we might want to clear + the cache. - .. versionadded:: 2011.2 - """ - try: - ctx_dict = func._pyopencl_first_arg_dep_memoize_dic - except AttributeError: - # FIXME: This may keep contexts alive longer than desired. - # But I guess since the memory in them is freed, who cares. - ctx_dict = func._pyopencl_first_arg_dep_memoize_dic = {} - _first_arg_dependent_caches.append(ctx_dict) + .. versionadded:: 2011.2 + """ + if kwargs: + cache_key = (args, frozenset(kwargs.items())) + else: + cache_key = (args,) - try: - return ctx_dict[cl_object][args] - except KeyError: - arg_dict = ctx_dict.setdefault(cl_object, {}) - result = func(cl_object, *args) - arg_dict[args] = result - return result + try: + ctx_dict = func._pyopencl_first_arg_dep_memoize_dic + except AttributeError: + # FIXME: This may keep contexts alive longer than desired. + # But I guess since the memory in them is freed, who cares. + ctx_dict = func._pyopencl_first_arg_dep_memoize_dic = {} + _first_arg_dependent_caches.append(ctx_dict) + + try: + return ctx_dict[cl_object][cache_key] + except KeyError: + arg_dict = ctx_dict.setdefault(cl_object, {}) + result = func(cl_object, *args, **kwargs) + arg_dict[cache_key] = result + return result + + from functools import update_wrapper + update_wrapper(wrapper, func) + return wrapper context_dependent_memoize = first_arg_dependent_memoize diff --git a/setup.py b/setup.py index fa50aeed7409812c2d8751ece2f1ccb404eba7db..203ceb297d60f8b11a0e49cc5656381572dce0ea 100644 --- a/setup.py +++ b/setup.py @@ -245,7 +245,6 @@ def main(): install_requires=[ "numpy", "pytools>=2017.6", - "decorator>=3.2.0", "appdirs>=1.4.0", # "Mako>=0.3.6", ],