diff --git a/loopy/__init__.py b/loopy/__init__.py index 2bea55f082eccebee940b3eb48fa175b0a509f63..db6953a6d2ce80e5bbba4d741c7123c9f390b45b 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -31,7 +31,8 @@ import islpy as isl from islpy import dim_type from loopy.symbolic import (RuleAwareIdentityMapper, RuleAwareSubstitutionMapper, - SubstitutionRuleMappingContext) + SubstitutionRuleMappingContext, + TaggedVariable, Reduction, LinearSubscript, ) from loopy.diagnostic import LoopyError @@ -62,12 +63,14 @@ from loopy.padding import (split_arg_axis, find_padding_multiple, from loopy.preprocess import (preprocess_kernel, realize_reduction, infer_unknown_types) from loopy.schedule import generate_loop_schedules, get_one_scheduled_kernel -from loopy.codegen import generate_code +from loopy.codegen import generate_code, generate_body from loopy.compiled import CompiledKernel from loopy.options import Options from loopy.auto_test import auto_test_vs_ref __all__ = [ + "TaggedVariable", "Reduction", "LinearSubscript", + "auto", "LoopKernel",