diff --git a/loopy/preprocess.py b/loopy/preprocess.py index 0b19ff416b2ca15e9d626a1644b9223367da2f69..3e8c70e536e6852cea5f603cda5265a676062237 100644 --- a/loopy/preprocess.py +++ b/loopy/preprocess.py @@ -1946,13 +1946,13 @@ def realize_reduction(kernel, insn_id_filter=None, unknown_types_ok=True, def realize_ilp(kernel): logger.debug("%s: add axes to temporaries for ilp" % kernel.name) - from loopy.kernel.data import IlpBaseTag, VectorizeTag + from loopy.kernel.data import (IlpBaseTag, VectorizeTag, + filter_iname_tags_by_type) - privatizing_inames = frozenset(iname - for iname in kernel.all_inames() - if isinstance( - kernel.iname_to_tag.get(iname), - (IlpBaseTag, VectorizeTag))) + privatizing_inames = frozenset( + iname for iname, tags in six.iteritems(kernel.iname_to_tags) + if filter_iname_tags_by_type(tags, (IlpBaseTag, VectorizeTag)) + ) from loopy.transform.privatize import privatize_temporaries_with_inames return privatize_temporaries_with_inames(kernel, privatizing_inames) diff --git a/loopy/transform/precompute.py b/loopy/transform/precompute.py index 4755ca1774a15480a2c6b255380dd724e47f9042..1af16f92f8f18b4c08742b0059267350b6904835 100644 --- a/loopy/transform/precompute.py +++ b/loopy/transform/precompute.py @@ -993,12 +993,9 @@ def precompute(kernel, subst_use, sweep_inames=[], within=None, from loopy import tag_inames kernel = tag_inames(kernel, new_iname_to_tag) - from loopy.kernel.data import AutoFitLocalIndexTag - has_automatic_axes = any( - isinstance(tag, AutoFitLocalIndexTag) - for tag in new_iname_to_tag.values()) + from loopy.kernel.data import AutoFitLocalIndexTag, filter_iname_tags_by_type - if has_automatic_axes: + if filter_iname_tags_by_type(new_iname_to_tag.values(), AutoFitLocalIndexTag): from loopy.kernel.tools import assign_automatic_axes kernel = assign_automatic_axes(kernel)