diff --git a/meshmode/array_context.py b/meshmode/array_context.py index 440ae127..def123f1 100644 --- a/meshmode/array_context.py +++ b/meshmode/array_context.py @@ -1675,7 +1675,10 @@ def transform_loopy_program(self, t_unit): " strategy as some loops are uninferred which mesh entity" " they belong to.", stacklevel=2) - return super().transform_loopy_program(original_t_unit) + new_t_unit = super().transform_loopy_program(original_t_unit) + self.transform_loopy_cache.store_if_not_present( + original_t_unit, new_t_unit) + return new_t_unit for insn in knl.instructions: for assignee in insn.assignee_var_names(): @@ -1689,7 +1692,10 @@ def transform_loopy_program(self, t_unit): " strategy as some instructions couldn't be inferred as" " einsums", stacklevel=2) - return super().transform_loopy_program(original_t_unit) + new_t_unit = super().transform_loopy_program(original_t_unit) + self.transform_loopy_cache.store_if_not_present( + original_t_unit, new_t_unit) + return new_t_unit # }}} @@ -1797,7 +1803,10 @@ def transform_loopy_program(self, t_unit): "transform_loopy_program not broad enough (yet)." " Falling back to a possibly slower" " transformation strategy.") - return super().transform_loopy_program(original_t_unit) + new_t_unit = super().transform_loopy_program(original_t_unit) + self.transform_loopy_cache.store_if_not_present( + original_t_unit, new_t_unit) + return new_t_unit # }}}