<s> from . . import syntax <EOL> from . . ndtypes import ScalarT , TupleT <EOL> from . . syntax import Expr , While , ForLoop , Const <EOL> from . . syntax . helpers import zero , one , zero_i32 , zero_i64 , wrap_if_constant <EOL> from array_builder import ArrayBuilder <EOL> class LoopBuilder ( ArrayBuilder ) : <EOL> """<STR_LIT>""" <EOL> def loop_var ( self , name = "<STR_LIT:i>" , start_val = syntax . zero_i64 ) : <EOL> """<STR_LIT>""" <EOL> start_val = wrap_if_constant ( start_val ) <EOL> counter_type = start_val . type <EOL> counter_before = self . assign_name ( start_val , name + "<STR_LIT>" ) <EOL> counter = self . fresh_var ( counter_type , name ) <EOL> counter_after = self . fresh_var ( counter_type , name + "<STR_LIT>" ) <EOL> merge = { counter . name : ( counter_before , counter_after ) } <EOL> return counter , counter_after , merge <EOL> def loop ( self , <EOL> start , <EOL> niters , <EOL> loop_body , <EOL> return_stmt = False , <EOL> while_loop = False , <EOL> step = None , <EOL> merge = None ) : <EOL> if isinstance ( start , ( int , long ) ) : <EOL> start = self . int ( start ) <EOL> assert isinstance ( start , Expr ) <EOL> if isinstance ( niters , ( int , long ) ) : <EOL> niters = self . int ( niters ) <EOL> assert isinstance ( niters , Expr ) <EOL> if step is None : <EOL> step = one ( start . type ) <EOL> elif isinstance ( step , ( int , long ) ) : <EOL> step = self . int ( step ) <EOL> assert isinstance ( step , Expr ) <EOL> if merge is None : <EOL> merge = { } <EOL> if while_loop : <EOL> i , i_after , i_merge = self . loop_var ( "<STR_LIT:i>" , start ) <EOL> merge . update ( i_merge ) <EOL> cond = self . lt ( i , niters ) <EOL> self . blocks . push ( ) <EOL> loop_body ( i ) <EOL> self . assign ( i_after , self . add ( i , step ) ) <EOL> body = self . blocks . pop ( ) <EOL> loop_stmt = While ( cond , body , merge ) <EOL> else : <EOL> var = self . fresh_var ( start . type , "<STR_LIT:i>" ) <EOL> self . blocks . push ( ) <EOL> loop_body ( var ) <EOL> body = self . blocks . pop ( ) <EOL> loop_stmt = ForLoop ( var , start , niters , step , body , merge ) <EOL> if return_stmt : <EOL> return loop_stmt <EOL> else : <EOL> self . blocks += loop_stmt <EOL> class Accumulator : <EOL> def __init__ ( self , acc_type , fresh_var , assign ) : <EOL> self . acc_type = acc_type <EOL> self . fresh_var = fresh_var <EOL> self . assign = assign <EOL> self . start_var = fresh_var ( acc_type , "<STR_LIT>" ) <EOL> self . curr_var = self . start_var <EOL> def get ( self ) : <EOL> return self . curr_var <EOL> def update ( self , new_value ) : <EOL> new_var = self . fresh_var ( self . acc_type , "<STR_LIT>" ) <EOL> self . assign ( new_var , new_value ) <EOL> self . curr_var = new_var <EOL> def mk_acc ( self , init ) : <EOL> return self . Accumulator ( init . type , self . fresh_var , self . assign ) <EOL> def accumulate_loop ( self , start , stop , loop_body , init , return_stmt = False ) : <EOL> acc = self . mk_acc ( init ) <EOL> def loop_body_with_acc ( i ) : <EOL> loop_body ( acc , i ) <EOL> loop_stmt = self . loop ( start , stop , loop_body_with_acc , return_stmt = True ) <EOL> loop_stmt . merge [ acc . start_var . name ] = ( init , acc . curr_var ) <EOL> if return_stmt : <EOL> return loop_stmt , acc . start_var <EOL> else : <EOL> self . blocks += loop_stmt <EOL> return acc . start_var <EOL> def array_copy ( self , src , dest , return_stmt = False ) : <EOL> assert self . is_array ( dest ) <EOL> shape = self . shape ( dest ) <EOL> dims = self . tuple_elts ( shape ) <EOL> rank = len ( dims ) <EOL> index_vars = [ ] <EOL> def create_loops ( ) : <EOL> i = len ( index_vars ) <EOL> def loop_body ( index_var ) : <EOL> index_vars . append ( index_var ) <EOL> if i + <NUM_LIT:1> == rank : <EOL> index_tuple = self . tuple ( index_vars , "<STR_LIT>" ) <EOL> lhs = self . index ( dest , index_tuple , temp = False ) <EOL> rhs = self . index ( src , index_tuple , temp = True ) <EOL> self . assign ( lhs , rhs ) <EOL> else : <EOL> self . blocks += create_loops ( ) <EOL> start = syntax . zero_i64 <EOL> stop = dims [ i ] <EOL> if i > <NUM_LIT:0> or return_stmt : <EOL> return self . loop ( start , stop , loop_body , True ) <EOL> else : <EOL> return self . loop ( start , stop , loop_body , return_stmt ) <EOL> return create_loops ( ) <EOL> def _to_list ( self , bounds ) : <EOL> if isinstance ( bounds , Expr ) : <EOL> if isinstance ( bounds . type , ScalarT ) : <EOL> return [ bounds ] <EOL> else : <EOL> assert isinstance ( bounds . type , TupleT ) , "<STR_LIT>" % ( bounds , bounds . type ) <EOL> return self . tuple_elts ( bounds ) <EOL> elif isinstance ( bounds , ( int , long ) ) : <EOL> return [ bounds ] <EOL> else : <EOL> assert isinstance ( bounds , ( tuple , list ) ) <EOL> return bounds <EOL> def nested_loops ( self , <EOL> upper_bounds , <EOL> loop_body , <EOL> lower_bounds = None , <EOL> step_sizes = None , <EOL> index_vars_as_list = False ) : <EOL> upper_bounds = self . _to_list ( upper_bounds ) <EOL> n_loops = len ( upper_bounds ) <EOL> assert lower_bounds is None or len ( lower_bounds ) == n_loops <EOL> assert step_sizes is None or len ( step_sizes ) == n_loops <EOL> if lower_bounds is None : <EOL> lower_bounds = [ self . int ( <NUM_LIT:0> ) for _ in upper_bounds ] <EOL> else : <EOL> lower_bounds = self . _to_list ( lower_bounds ) <EOL> if step_sizes is None : <EOL> step_sizes = [ self . int ( <NUM_LIT:1> ) for _ in upper_bounds ] <EOL> else : <EOL> step_sizes = self . _to_list ( step_sizes ) <EOL> def build_loops ( index_vars = ( ) ) : <EOL> n_indices = len ( index_vars ) <EOL> if n_indices == n_loops : <EOL> if isinstance ( loop_body , Expr ) : <EOL> input_types = self . input_types ( loop_body ) <EOL> if len ( input_types ) == len ( index_vars ) : <EOL> result = self . call ( loop_body , index_vars ) <EOL> else : <EOL> result = self . call ( loop_body , [ self . tuple ( index_vars ) ] ) <EOL> else : <EOL> if index_vars_as_list : <EOL> idx_tuple = list ( index_vars ) <EOL> elif n_indices > <NUM_LIT:1> : <EOL> idx_tuple = self . tuple ( index_vars ) <EOL> else : <EOL> idx_tuple = index_vars [ <NUM_LIT:0> ] <EOL> assert hasattr ( loop_body , '<STR_LIT>' ) , "<STR_LIT>" % ( loop_body , ) <EOL> result = loop_body ( idx_tuple ) <EOL> assert self . is_none ( result ) , "<STR_LIT>" % ( result , ) <EOL> else : <EOL> def inner_loop_body ( idx ) : <EOL> build_loops ( index_vars + ( idx , ) ) <EOL> lower = lower_bounds [ n_indices ] <EOL> upper = upper_bounds [ n_indices ] <EOL> step = step_sizes [ n_indices ] <EOL> if lower . __class__ is Const and upper . __class__ is Const and step . __class__ is Const : <EOL> if step . value < <NUM_LIT:0> : <EOL> assert lower . value > upper . value , "<STR_LIT>" % ( lower , upper , step ) <EOL> elif step . value > <NUM_LIT:0> : <EOL> assert upper . value > lower . value , "<STR_LIT>" % ( lower , upper , step ) <EOL> self . loop ( lower , upper , inner_loop_body , step = step ) <EOL> build_loops ( ) </s>
