diff --git a/loopy/__init__.py b/loopy/__init__.py index f21301d301f32a5d41f695f26fa158411e2b7828..049327d034bc48ff11d40abc92a955088ef6d5b5 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -11,7 +11,7 @@ register_mpz_with_pymbolic() -# TODO: Constant memory +# TODO: Constant memory (plus check for count) # TODO: Reuse of previously split dimensions for prefetch # (Or general merging) diff --git a/loopy/kernel.py b/loopy/kernel.py index 21bf43a4432a565603f8ab68f30a78392aaf28cd..6f14113a5dcafdba589182effef44ef626e466a1 100644 --- a/loopy/kernel.py +++ b/loopy/kernel.py @@ -83,16 +83,18 @@ class IndexTag(Record): def __hash__(self): raise RuntimeError("use .key to hash index tags") - @property - def key(self): - return type(self) class ParallelTag(IndexTag): pass -class ParallelTagWithAxis(ParallelTag): +class UniqueTag(IndexTag): + @property + def key(self): + return type(self) + +class ParallelTagWithAxis(ParallelTag, UniqueTag): __slots__ = ["axis", "forced_length"] def __init__(self, axis, forced_length=None): @@ -222,7 +224,8 @@ class LoopKernel(Record): def tag_key_to_iname(self): return dict( (tag.key, iname) - for iname, tag in self.iname_to_tag.iteritems()) + for iname, tag in self.iname_to_tag.iteritems() + if isinstance(tag, UniqueTag)) @property @memoize_method @@ -396,11 +399,12 @@ class LoopKernel(Record): new_tag_keys = set(tag.key for tag in [outer_tag, inner_tag] - if tag is not None) + if tag is not None + if isinstance(tag, UniqueTag)) repeated_tag_keys = new_tag_keys & set( - tag.key for tag in - self.iname_to_tag.itervalues()) + tag.key for tag in self.iname_to_tag.itervalues() + if isinstance(tag, UniqueTag)) if repeated_tag_keys: raise RuntimeError("repeated tag(s): %s" % repeated_tag_keys)