diff --git a/src/pystencils/__init__.py b/src/pystencils/__init__.py index 4f8b2660777dbe7a1fa535cfbbc2e64458a0d692..6cb375b61e63904c0cd3c2e6e9d6a3be86be6b29 100644 --- a/src/pystencils/__init__.py +++ b/src/pystencils/__init__.py @@ -18,7 +18,7 @@ from .field import Field, FieldType, fields from .types import create_type, create_numeric_type from .cache import clear_cache from .kernel_decorator import kernel, kernel_config -from .kernelcreation import create_kernel, create_staggered_kernel +from .codegen.driver import create_kernel, create_staggered_kernel from .codegen import Kernel from .jit import no_jit from .backend.exceptions import KernelConstraintsError diff --git a/src/pystencils/backend/transformations/canonicalize_symbols.py b/src/pystencils/backend/transformations/canonicalize_symbols.py index c0406c25d820df0a1c3821074395b8709b482113..17875caa17d470484d6d509fdd023a1d7867ba1f 100644 --- a/src/pystencils/backend/transformations/canonicalize_symbols.py +++ b/src/pystencils/backend/transformations/canonicalize_symbols.py @@ -102,10 +102,11 @@ class CanonicalizeSymbols: cc.mark_as_updated(lhs.symbol) case PsLoop(ctr, _, _, _, _): + decl_symb = ctr.symbol for c in node.children[::-1]: self.visit(c, cc) cc.mark_as_updated(ctr.symbol) - cc.end_lifespan(ctr.symbol) + cc.end_lifespan(decl_symb) case PsConditional(cond, then, els): if els is not None: diff --git a/tests/nbackend/transformations/test_canonicalize_symbols.py b/tests/nbackend/transformations/test_canonicalize_symbols.py index a11e9bd1353ef98c5ae23e0d86ddce3e5c9d579c..2758d123417eb8e3015ed1d6b4d8cf0ba7c14611 100644 --- a/tests/nbackend/transformations/test_canonicalize_symbols.py +++ b/tests/nbackend/transformations/test_canonicalize_symbols.py @@ -104,13 +104,16 @@ def test_loop_counters(): loops = factory.loops_from_ispace(ispace, body) - loops_copy = loops.clone() + loops_clone = loops.clone() + loops_clone2 = loops.clone() - ast = PsBlock([loops, loops_copy]) + ast = PsBlock([loops, loops_clone, loops_clone2]) ast = canonicalize(ast) - assert loops_copy.counter.symbol.name == "ctr_0" - assert not loops_copy.counter.symbol.get_dtype().const - assert loops.counter.symbol.name == "ctr_0__0" + assert loops_clone2.counter.symbol.name == "ctr_0" + assert not loops_clone2.counter.symbol.get_dtype().const + assert loops_clone.counter.symbol.name == "ctr_0__0" + assert not loops_clone.counter.symbol.get_dtype().const + assert loops.counter.symbol.name == "ctr_0__1" assert not loops.counter.symbol.get_dtype().const