From a5fc63dd31b553073df3e5adae1e56f3f07f322c Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Thu, 1 Aug 2019 20:55:55 +0800 Subject: [PATCH 1/3] [TOPI] Memoize winograd matrix --- python/tvm/contrib/pickle_memoize.py | 16 +++++++++++----- topi/python/topi/nn/winograd_util.py | 3 +++ 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/python/tvm/contrib/pickle_memoize.py b/python/tvm/contrib/pickle_memoize.py index c65797106fc5..b5abf9b9b7a6 100644 --- a/python/tvm/contrib/pickle_memoize.py +++ b/python/tvm/contrib/pickle_memoize.py @@ -34,9 +34,11 @@ class Cache(object): ---------- key: str The file key to the function + save_at_exit: bool + Whether save the cache to file when the program exits """ cache_by_key = {} - def __init__(self, key): + def __init__(self, key, save_at_exit): cache_dir = ".pkl_memoize_py{0}".format(sys.version_info[0]) if not os.path.exists(cache_dir): os.mkdir(cache_dir) @@ -49,6 +51,7 @@ def __init__(self, key): else: self.cache = {} self.dirty = False + self.save_at_exit = save_at_exit def save(self): if self.dirty: @@ -60,16 +63,19 @@ def save(self): def _atexit(): """Save handler.""" for value in Cache.cache_by_key.values(): - value.save() + if value.save_at_exit: + value.save() -def memoize(key): +def memoize(key, save_at_exit=False): """Memoize the result of function and reuse multiple times. Parameters ---------- key: str The unique key to the file + save_at_exit: bool + Whether save the cache to file when the program exits Returns ------- @@ -81,9 +87,9 @@ def _register(f): allow_types = (string_types, int, float) fkey = key + "." + f.__name__ + ".pkl" if fkey not in Cache.cache_by_key: - Cache.cache_by_key[fkey] = Cache(fkey) + Cache.cache_by_key[fkey] = Cache(fkey, save_at_exit) cache = Cache.cache_by_key[fkey] - cargs = tuple(x.cell_contents for x in f.__closure__) + cargs = tuple(x.cell_contents for x in f.__closure__) if f.__closure__ else () cargs = (len(cargs),) + cargs def _memoized_f(func, *args, **kwargs): diff --git a/topi/python/topi/nn/winograd_util.py b/topi/python/topi/nn/winograd_util.py index db57f7671618..48a2a40e526c 100644 --- a/topi/python/topi/nn/winograd_util.py +++ b/topi/python/topi/nn/winograd_util.py @@ -26,6 +26,7 @@ from functools import reduce import numpy as np from ..util import const_matrix +from tvm.contrib.pickle_memoize import memoize # pylint: disable=invalid-name @@ -131,6 +132,8 @@ def _interpolation_points(degree): return np.array(in_pts[degree-1], dtype=np.float64) + +@memoize("topi.nn.winograd_util", save_at_exit=False) def winograd_transform_matrices(tile_size, kernel_size, out_dtype): """Compute the A, B, and G transform matrices for `tile_size` as a `tvm.Expr`. """ From bc694eef4528b6c32f2cad76dc12ea21f7fe9d53 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Thu, 1 Aug 2019 21:08:08 +0800 Subject: [PATCH 2/3] lint --- topi/python/topi/nn/winograd_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/topi/python/topi/nn/winograd_util.py b/topi/python/topi/nn/winograd_util.py index 48a2a40e526c..e936d8ac639c 100644 --- a/topi/python/topi/nn/winograd_util.py +++ b/topi/python/topi/nn/winograd_util.py @@ -25,8 +25,8 @@ from operator import mul from functools import reduce import numpy as np -from ..util import const_matrix from tvm.contrib.pickle_memoize import memoize +from ..util import const_matrix # pylint: disable=invalid-name From a47d168422b90dc1167681c159f9b3aa2ad9ae34 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Thu, 1 Aug 2019 21:46:31 +0800 Subject: [PATCH 3/3] Fix name --- topi/python/topi/nn/winograd_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/topi/python/topi/nn/winograd_util.py b/topi/python/topi/nn/winograd_util.py index e936d8ac639c..464b63301b40 100644 --- a/topi/python/topi/nn/winograd_util.py +++ b/topi/python/topi/nn/winograd_util.py @@ -133,7 +133,7 @@ def _interpolation_points(degree): return np.array(in_pts[degree-1], dtype=np.float64) -@memoize("topi.nn.winograd_util", save_at_exit=False) +@memoize("topi.nn.winograd_matrices", save_at_exit=False) def winograd_transform_matrices(tile_size, kernel_size, out_dtype): """Compute the A, B, and G transform matrices for `tile_size` as a `tvm.Expr`. """