<s> from dsltools import ScopedSet <EOL> from . . analysis . collect_vars import collect_var_names , collect_binding_names <EOL> from . . analysis . escape_analysis import may_alias <EOL> from . . analysis . syntax_visitor import SyntaxVisitor <EOL> from . . ndtypes import ImmutableT , ArrayT , PtrT <EOL> from . . syntax import Var , Assign , Return , While , If , Index , Alloc , AllocArray <EOL> from . . syntax import Array , ArrayView , Slice , Struct <EOL> from transform import Transform <EOL> class Find_LICM_Candidates ( SyntaxVisitor ) : <EOL> def __init__ ( self ) : <EOL> SyntaxVisitor . __init__ ( self ) <EOL> self . mutable_types = None <EOL> self . volatile_vars = ScopedSet ( ) <EOL> self . depends_on = { } <EOL> self . safe_to_move = set ( [ ] ) <EOL> self . curr_block_id = None <EOL> self . block_contains_return = set ( [ ] ) <EOL> self . may_alias = None <EOL> def visit_fn ( self , fn ) : <EOL> self . volatile_vars . push ( fn . arg_names ) <EOL> self . may_alias = may_alias ( fn ) <EOL> SyntaxVisitor . visit_fn ( self , fn ) <EOL> return self . safe_to_move <EOL> def mark_safe_assignments ( self , block , volatile_set ) : <EOL> for stmt in block : <EOL> klass = stmt . __class__ <EOL> if klass is Assign and stmt . lhs . __class__ is Var : <EOL> name = stmt . lhs . name <EOL> dependencies = self . depends_on . get ( name , set ( [ ] ) ) <EOL> volatile = name in volatile_set or any ( d in volatile_set for d in dependencies ) <EOL> if not volatile : <EOL> self . safe_to_move . add ( name ) <EOL> elif klass is If : <EOL> if id ( stmt . true ) in self . block_contains_return or id ( stmt . false ) in self . block_contains_return : <EOL> break <EOL> elif klass is While : <EOL> if id ( stmt . body ) in self . block_contains_return : <EOL> break <EOL> elif klass is Return : <EOL> break <EOL> def mark_curr_block_returns ( self ) : <EOL> self . block_contains_return . add ( self . curr_block_id ) <EOL> def does_block_return ( self , block ) : <EOL> return id ( block ) in self . block_contains_return <EOL> def visit_Return ( self , stmt ) : <EOL> self . mark_curr_block_returns ( ) <EOL> def visit_block ( self , stmts ) : <EOL> self . curr_block_id = id ( stmts ) <EOL> SyntaxVisitor . visit_block ( self , stmts ) <EOL> def visit_merge ( self , merge , both_branches = True ) : <EOL> pass <EOL> def visit_ForLoop ( self , stmt ) : <EOL> self . volatile_vars . push ( stmt . merge . keys ( ) ) <EOL> self . volatile_vars . add ( stmt . var . name ) <EOL> SyntaxVisitor . visit_ForLoop ( self , stmt ) <EOL> if self . does_block_return ( stmt . body ) : <EOL> self . block_contains_return ( ) <EOL> volatile_in_scope = self . volatile_vars . pop ( ) <EOL> self . mark_safe_assignments ( stmt . body , volatile_in_scope ) <EOL> def visit_While ( self , stmt ) : <EOL> self . volatile_vars . push ( stmt . merge . keys ( ) ) <EOL> SyntaxVisitor . visit_While ( self , stmt ) <EOL> if self . does_block_return ( stmt . body ) : <EOL> self . block_contains_return ( ) <EOL> volatile_in_scope = self . volatile_vars . pop ( ) <EOL> self . mark_safe_assignments ( stmt . body , volatile_in_scope ) <EOL> def visit_Var ( self , expr ) : <EOL> self . volatile_vars . add ( expr . name ) <EOL> def visit_If ( self , stmt ) : <EOL> self . volatile_vars . push ( stmt . merge . keys ( ) ) <EOL> self . visit_expr ( stmt . cond ) <EOL> SyntaxVisitor . visit_If ( self , stmt ) <EOL> if self . does_block_return ( stmt . true ) or self . does_block_return ( stmt . false ) : <EOL> self . mark_curr_block_returns ( ) <EOL> volatile_in_scope = self . volatile_vars . pop ( ) <EOL> self . mark_safe_assignments ( stmt . true , volatile_in_scope ) <EOL> self . mark_safe_assignments ( stmt . false , volatile_in_scope ) <EOL> def is_mutable_alloc ( self , expr ) : <EOL> return isinstance ( expr . type , ( PtrT , ArrayT ) ) <EOL> def visit_Assign ( self , stmt ) : <EOL> lhs_names = collect_binding_names ( stmt . lhs ) <EOL> rhs_names = collect_var_names ( stmt . rhs ) <EOL> for x in lhs_names : <EOL> dependencies = self . depends_on . get ( x , set ( [ ] ) ) <EOL> dependencies . update ( rhs_names ) <EOL> self . depends_on [ x ] = dependencies <EOL> if any ( x in self . volatile_vars for x in rhs_names ) : <EOL> self . volatile_vars . update ( lhs_names ) <EOL> elif self . is_mutable_alloc ( stmt . rhs ) : <EOL> if len ( lhs_names ) == <NUM_LIT:1> and len ( self . may_alias . get ( lhs_names [ <NUM_LIT:0> ] , [ ] ) ) <= <NUM_LIT:1> : <EOL> pass <EOL> else : <EOL> self . volatile_vars . update ( lhs_names ) <EOL> if stmt . lhs . __class__ is Index : <EOL> assert stmt . lhs . value . __class__ is Var , "<STR_LIT>" % stmt <EOL> self . volatile_vars . add ( stmt . lhs . value . name ) <EOL> class LoopInvariantCodeMotion ( Transform ) : <EOL> def __init__ ( self ) : <EOL> Transform . __init__ ( self ) <EOL> def pre_apply ( self , fn ) : <EOL> self . analysis = Find_LICM_Candidates ( ) <EOL> self . safe_to_move = self . analysis . visit_fn ( fn ) <EOL> self . binding_depth = { } <EOL> self . mark_binding_depths ( fn . arg_names , <NUM_LIT:0> ) <EOL> def mark_binding_depths ( self , names , depth_offset = <NUM_LIT:0> ) : <EOL> curr_depth = len ( self . blocks . _blocks ) - <NUM_LIT:1> + depth_offset <EOL> for name in names : <EOL> self . binding_depth [ name ] = curr_depth <EOL> def transform_ForLoop ( self , stmt ) : <EOL> self . mark_binding_depths ( stmt . merge . iterkeys ( ) , <NUM_LIT:1> ) <EOL> self . mark_binding_depths ( [ stmt . var . name ] , <NUM_LIT:1> ) <EOL> return Transform . transform_ForLoop ( self , stmt ) <EOL> def transform_While ( self , stmt ) : <EOL> self . mark_binding_depths ( stmt . merge . iterkeys ( ) , <NUM_LIT:1> ) <EOL> return Transform . transform_While ( self , stmt ) <EOL> def transform_If ( self , stmt ) : <EOL> self . mark_binding_depths ( stmt . merge . iterkeys ( ) , <NUM_LIT:1> ) <EOL> return Transform . transform_If ( self , stmt ) <EOL> def transform_Assign ( self , stmt ) : <EOL> if stmt . lhs . __class__ is Var and stmt . rhs . __class__ is not Index : <EOL> name = stmt . lhs . name <EOL> if name in self . safe_to_move : <EOL> deps = self . analysis . depends_on [ name ] <EOL> if all ( d in self . binding_depth for d in deps ) : <EOL> if len ( deps ) > <NUM_LIT:0> : <EOL> target_level = max ( self . binding_depth [ d ] for d in deps ) <EOL> else : <EOL> target_level = <NUM_LIT:0> <EOL> if target_level >= <NUM_LIT:0> and target_level < self . blocks . depth ( ) : <EOL> self . blocks . _blocks [ target_level ] . append ( stmt ) <EOL> self . binding_depth [ name ] = target_level <EOL> return None <EOL> self . mark_binding_depths ( collect_binding_names ( stmt . lhs ) ) <EOL> return stmt </s>
