<s> import itertools <EOL> from . . import names <EOL> from . . builder import build_fn <EOL> from . . ndtypes import Int64 , repeat_tuple , NoneType , ScalarT , TupleT , ArrayT <EOL> from . . syntax import ( ParFor , IndexReduce , IndexScan , Index , Map , OuterMap , Var , Const , Expr ) <EOL> from . . syntax . helpers import get_types , none , zero_i64 <EOL> from . . syntax . adverb_helpers import max_rank_arg , max_rank <EOL> from transform import Transform <EOL> class IndexifyAdverbs ( Transform ) : <EOL> """<STR_LIT>""" <EOL> def fresh_input_name ( self , expr ) : <EOL> if expr is Var : <EOL> return names . refresh ( expr . name ) <EOL> else : <EOL> return names . fresh ( "<STR_LIT:input>" ) <EOL> def fresh_fn_name ( self , prefix , fn ) : <EOL> return names . fresh ( prefix + names . original ( self . get_fn ( fn ) . name ) ) <EOL> def fresh_index_names ( self , n ) : <EOL> name_supply = itertools . cycle ( [ "<STR_LIT:i>" , "<STR_LIT>" , "<STR_LIT:k>" , "<STR_LIT:l>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> return [ names . fresh ( name_supply . next ( ) ) for _ in xrange ( n ) ] <EOL> _indexed_fn_cache = { } <EOL> def indexify_fn ( self , fn , <EOL> axis , <EOL> array_args , <EOL> cartesian_product = False , <EOL> output = None , <EOL> index_offsets = None ) : <EOL> """<STR_LIT>""" <EOL> array_args = tuple ( array_args ) <EOL> array_arg_types = tuple ( get_types ( array_args ) ) <EOL> closure_args = self . closure_elts ( fn ) <EOL> closure_arg_types = tuple ( get_types ( closure_args ) ) <EOL> n_closure_args = len ( closure_args ) <EOL> fn = self . get_fn ( fn ) <EOL> axes = self . normalize_axes ( array_args , axis ) <EOL> key = ( fn . cache_key , <EOL> axes , <EOL> closure_arg_types , <EOL> array_arg_types , <EOL> output is None , <EOL> cartesian_product , <EOL> index_offsets , <EOL> ) <EOL> def mk_closure ( ) : <EOL> new_fn = self . _indexed_fn_cache [ key ] <EOL> if output is None : <EOL> return self . closure ( new_fn , closure_args + array_args ) <EOL> else : <EOL> return self . closure ( new_fn , ( output , ) + closure_args + array_args ) <EOL> if key in self . _indexed_fn_cache : <EOL> result = mk_closure ( ) <EOL> return result <EOL> if cartesian_product : <EOL> n_indices = <NUM_LIT:0> <EOL> for axis , arg_t in zip ( axes , array_arg_types ) : <EOL> if axis is None : <EOL> n_indices += self . rank ( arg_t ) <EOL> elif isinstance ( arg_t , ArrayT ) : <EOL> n_indices += <NUM_LIT:1> <EOL> else : <EOL> if any ( axis is None for axis in axes ) : <EOL> assert all ( axis is None for axis in axes ) , "<STR_LIT>" % axes <EOL> n_indices = max_rank ( array_arg_types ) <EOL> else : <EOL> assert all ( isinstance ( axis , ( int , long ) ) for axis in axes ) , "<STR_LIT>" % axes <EOL> n_indices = <NUM_LIT:1> <EOL> index_input_types = ( Int64 , ) * n_indices <EOL> if output is None : <EOL> inner_input_types = closure_arg_types + array_arg_types + index_input_types <EOL> new_return_type = fn . return_type <EOL> else : <EOL> inner_input_types = ( output . type , ) + closure_arg_types + array_arg_types + index_input_types <EOL> new_return_type = NoneType <EOL> input_names = [ ] <EOL> if output is not None : <EOL> if output is Var : <EOL> local_output_name = names . refresh ( output . name ) <EOL> else : <EOL> local_output_name = names . fresh ( "<STR_LIT>" ) <EOL> input_names . append ( local_output_name ) <EOL> for old_input_name in fn . arg_names : <EOL> input_names . append ( names . refresh ( old_input_name ) ) <EOL> input_names . extend ( self . fresh_index_names ( n_indices ) ) <EOL> new_fn_name = self . fresh_fn_name ( "<STR_LIT>" , fn ) <EOL> new_fn , builder , input_vars = build_fn ( inner_input_types , <EOL> new_return_type , <EOL> name = new_fn_name , <EOL> input_names = input_names ) <EOL> index_input_vars = input_vars [ - n_indices : ] <EOL> if output is None : <EOL> output_var = None <EOL> closure_arg_vars = input_vars [ : n_closure_args ] <EOL> array_arg_vars = input_vars [ n_closure_args : - n_indices ] <EOL> else : <EOL> output_var = input_vars [ <NUM_LIT:0> ] <EOL> closure_arg_vars = input_vars [ <NUM_LIT:1> : n_closure_args + <NUM_LIT:1> ] <EOL> array_arg_vars = input_vars [ n_closure_args + <NUM_LIT:1> : - n_indices ] <EOL> slice_values = self . get_slices ( builder , array_arg_vars , axes , index_input_vars , <EOL> cartesian_product , index_offsets ) <EOL> elt_result = builder . call ( fn , tuple ( closure_arg_vars ) + tuple ( slice_values ) ) <EOL> if output is None : <EOL> builder . return_ ( elt_result ) <EOL> else : <EOL> if len ( index_input_vars ) > <NUM_LIT:1> : <EOL> builder . setidx ( output_var , builder . tuple ( index_input_vars ) , elt_result ) <EOL> else : <EOL> builder . setidx ( output_var , index_input_vars [ <NUM_LIT:0> ] , elt_result ) <EOL> builder . return_ ( none ) <EOL> new_fn . created_by = self . fn . created_by <EOL> new_fn . transform_history = self . fn . transform_history <EOL> self . _indexed_fn_cache [ key ] = new_fn <EOL> return mk_closure ( ) <EOL> def get_slices ( self , builder , array_arg_vars , axes , index_input_vars , <EOL> cartesian_product = False , index_offsets = None ) : <EOL> slice_values = [ ] <EOL> axes = self . normalize_axes ( array_arg_vars , axes ) <EOL> if cartesian_product : <EOL> idx_counter = <NUM_LIT:0> <EOL> for i , curr_array in enumerate ( array_arg_vars ) : <EOL> axis = axes [ i ] <EOL> rank = curr_array . type . rank <EOL> if rank <= <NUM_LIT:1> and axis is None : <EOL> axis = <NUM_LIT:0> <EOL> if axis is None : <EOL> start = idx_counter <EOL> stop = idx_counter + rank <EOL> curr_indices = index_input_vars [ start : stop ] <EOL> idx_counter = stop <EOL> curr_slice = builder . index ( curr_array , curr_indices ) <EOL> elif rank > axis : <EOL> idx = index_input_vars [ idx_counter ] <EOL> idx_counter += <NUM_LIT:1> <EOL> curr_slice = builder . slice_along_axis ( curr_array , axis , idx ) <EOL> else : <EOL> curr_slice = curr_array <EOL> slice_values . append ( curr_slice ) <EOL> else : <EOL> for i , curr_array in enumerate ( array_arg_vars ) : <EOL> axis = axes [ i ] <EOL> rank = curr_array . type . rank <EOL> if rank <= <NUM_LIT:1> and axis is None : <EOL> axis = <NUM_LIT:0> <EOL> if axis is None : <EOL> assert len ( index_input_vars ) <= rank , "<STR_LIT>" % ( curr_array , curr_array . type ) <EOL> curr_indices = index_input_vars [ - rank : ] <EOL> curr_slice = builder . index ( curr_array , curr_indices ) <EOL> elif rank > axis : <EOL> curr_idx = index_input_vars [ <NUM_LIT:0> ] <EOL> if index_offsets is not None : <EOL> assert len ( index_offsets ) > i <EOL> curr_offset = index_offsets [ i ] <EOL> if not isinstance ( curr_offset , Expr ) : <EOL> curr_offset = builder . int ( curr_offset ) <EOL> curr_idx = builder . add ( curr_idx , curr_offset ) <EOL> curr_slice = builder . slice_along_axis ( curr_array , axis , curr_idx ) <EOL> else : <EOL> curr_slice = curr_array <EOL> slice_values . append ( curr_slice ) <EOL> return slice_values <EOL> def create_map_output_array ( self , <EOL> fn , array_args , axes , <EOL> cartesian_product = False , <EOL> name = "<STR_LIT>" ) : <EOL> """<STR_LIT>""" <EOL> assert self . is_fn ( fn ) , "<STR_LIT>" % ( fn , ) <EOL> assert isinstance ( array_args , ( list , tuple ) ) , "<STR_LIT>" % ( array_args , ) <EOL> axes = self . normalize_axes ( array_args , axes ) <EOL> n_indices = <NUM_LIT:0> <EOL> for arg , axis in zip ( array_args , axes ) : <EOL> r = self . rank ( arg ) <EOL> if r == <NUM_LIT:0> : <EOL> continue <EOL> if axis is None : <EOL> if cartesian_product : <EOL> n_indices += self . rank ( arg ) <EOL> else : <EOL> n_indices = max ( n_indices , self . rank ( arg ) ) <EOL> elif r <= axis : <EOL> continue <EOL> else : <EOL> if cartesian_product : <EOL> n_indices += <NUM_LIT:1> <EOL> else : <EOL> n_indices = max ( n_indices , <NUM_LIT:1> ) <EOL> inner_args = self . get_slices ( builder = self , <EOL> array_arg_vars = array_args , <EOL> axes = axes , <EOL> index_input_vars = [ zero_i64 ] * ( n_indices ) , <EOL> cartesian_product = cartesian_product ) <EOL> if cartesian_product : <EOL> extra_dims = [ ] <EOL> for array , axis in zip ( array_args , axes ) : <EOL> rank = self . rank ( array ) <EOL> if axis is None : <EOL> dim = <NUM_LIT:1> <EOL> elif rank > axis : <EOL> dim = self . shape ( array , axis ) <EOL> else : <EOL> dim = <NUM_LIT:1> <EOL> extra_dims . append ( dim ) <EOL> outer_shape_tuple = self . tuple ( extra_dims ) <EOL> else : <EOL> outer_shape_tuple = self . iter_bounds ( array_args , axes ) <EOL> if isinstance ( outer_shape_tuple . type , ScalarT ) : <EOL> outer_shape_tuple = self . tuple ( [ outer_shape_tuple ] ) <EOL> return self . create_output_array ( fn , inner_args , outer_shape_tuple , name ) <EOL> def insert_parfor ( self , index_fn , bounds , n_read_only , n_write_only = <NUM_LIT:1> ) : <EOL> read_only = ( [ False ] * n_write_only ) + ( [ True ] * n_read_only ) <EOL> write_only = ( [ True ] * n_write_only ) + ( [ False ] * n_read_only ) <EOL> stmt = ParFor ( index_fn , bounds , read_only = read_only , write_only = write_only ) <EOL> self . blocks += stmt <EOL> def transform_Map ( self , expr , output = None ) : <EOL> args = self . transform_expr_list ( expr . args ) <EOL> axes = self . normalize_axes ( args , expr . axis ) <EOL> old_fn = expr . fn <EOL> if output is None : <EOL> output = self . create_map_output_array ( old_fn , args , axes ) <EOL> bounds = self . iter_bounds ( args , axes ) <EOL> index_fn = self . indexify_fn ( expr . fn , axes , args , <EOL> cartesian_product = False , <EOL> output = output ) <EOL> self . insert_parfor ( index_fn , bounds , n_read_only = len ( args ) , n_write_only = <NUM_LIT:1> ) <EOL> return output <EOL> def transform_OuterMap ( self , expr ) : <EOL> args = self . transform_expr_list ( expr . args ) <EOL> axes = self . normalize_axes ( args , expr . axis ) <EOL> fn = expr . fn <EOL> outer_shape = self . iter_bounds ( args , axes , cartesian_product = True ) <EOL> zero = self . int ( <NUM_LIT:0> ) <EOL> first_values = [ self . slice_along_axis ( arg , axis , zero ) <EOL> for ( arg , axis ) in zip ( args , axes ) ] <EOL> output = self . create_output_array ( fn , first_values , outer_shape ) <EOL> loop_body = self . indexify_fn ( fn , axes , args , <EOL> cartesian_product = True , <EOL> output = output ) <EOL> self . insert_parfor ( loop_body , outer_shape , len ( args ) ) <EOL> return output <EOL> def transform_IndexMap ( self , expr , output = None ) : <EOL> shape = expr . shape <EOL> fn = expr . fn <EOL> dims = self . tuple_elts ( shape ) <EOL> n_dims = len ( dims ) <EOL> if n_dims == <NUM_LIT:1> : shape = dims [ <NUM_LIT:0> ] <EOL> if output is None : output = self . create_output_array ( fn , [ shape ] , shape ) <EOL> old_closure_args = self . closure_elts ( fn ) <EOL> old_closure_arg_types = get_types ( old_closure_args ) <EOL> fn = self . get_fn ( fn ) <EOL> closure_arg_names = [ self . fresh_input_name ( clos_arg ) for clos_arg in old_closure_args ] <EOL> new_closure_vars = [ Var ( name , type = t ) <EOL> for name , t in <EOL> zip ( closure_arg_names , old_closure_arg_types ) ] <EOL> old_input_types = fn . input_types <EOL> last_input_type = old_input_types [ - <NUM_LIT:1> ] <EOL> index_is_tuple = isinstance ( last_input_type , TupleT ) <EOL> if index_is_tuple : <EOL> index_types = last_input_type . elt_types <EOL> else : <EOL> index_types = old_input_types [ - n_dims : ] <EOL> idx_names = self . fresh_index_names ( n_dims ) <EOL> assert len ( index_types ) == n_dims , "<STR_LIT>" % ( dims , len ( index_types ) ) <EOL> output_name = names . refresh ( "<STR_LIT>" ) <EOL> new_input_names = [ output_name ] + closure_arg_names + idx_names <EOL> new_input_types = [ output . type ] + old_closure_arg_types + list ( index_types ) <EOL> new_fn_name = names . fresh ( "<STR_LIT>" + names . original ( fn . name ) ) <EOL> new_fn , builder , input_vars = build_fn ( new_input_types , NoneType , <EOL> name = new_fn_name , <EOL> input_names = new_input_names ) <EOL> new_fn . created_by = self . fn . created_by <EOL> output_var = input_vars [ <NUM_LIT:0> ] <EOL> idx_vars = input_vars [ - n_dims : ] <EOL> if not self . is_none ( expr . start_index ) : <EOL> if isinstance ( expr . start_index . type , ScalarT ) : <EOL> idx_vars = [ builder . add ( idx , expr . start_index , "<STR_LIT>" ) for idx in idx_vars ] <EOL> else : <EOL> start_indices = builder . tuple_elts ( expr . start_index ) <EOL> assert len ( start_indices ) == len ( idx_vars ) , "<STR_LIT>" % ( idx_vars , start_indices ) <EOL> idx_vars = [ builder . add ( idx , offset , "<STR_LIT>" % i ) <EOL> for i , ( idx , offset ) <EOL> in enumerate ( zip ( idx_vars , start_indices ) ) ] <EOL> if index_is_tuple : <EOL> elt_result = builder . call ( fn , new_closure_vars + [ builder . tuple ( idx_vars ) ] ) <EOL> else : <EOL> elt_result = builder . call ( fn , new_closure_vars + idx_vars ) <EOL> if len ( idx_vars ) == <NUM_LIT:1> : <EOL> builder . setidx ( output_var , idx_vars [ <NUM_LIT:0> ] , elt_result ) <EOL> else : <EOL> builder . setidx ( output_var , builder . tuple ( idx_vars ) , elt_result ) <EOL> builder . return_ ( none ) <EOL> new_closure = self . closure ( new_fn , ( output , ) + tuple ( old_closure_args ) ) <EOL> self . insert_parfor ( new_closure , shape , len ( old_closure_args ) ) <EOL> return output <EOL> def transform_Reduce ( self , expr ) : <EOL> fn = expr . fn <EOL> combine = expr . combine <EOL> init = expr . init <EOL> args = [ ] <EOL> axes = [ ] <EOL> raw_axes = self . normalize_axes ( expr . args , expr . axis ) <EOL> for axis , arg in zip ( raw_axes , expr . args ) : <EOL> if self . is_none ( axis ) : <EOL> args . append ( self . ravel ( arg ) ) <EOL> axes . append ( <NUM_LIT:0> ) <EOL> else : <EOL> args . append ( arg ) <EOL> axes . append ( axis ) <EOL> max_arg = max_rank_arg ( args ) <EOL> nelts = self . shape ( max_arg , axis ) <EOL> if self . is_none ( axis ) : <EOL> nelts = self . prod ( nelts ) <EOL> if init is None or self . is_none ( init ) : <EOL> init_args = [ self . index_along_axis ( arg , axis , self . int ( <NUM_LIT:0> ) ) for arg , axis in zip ( args , axes ) ] <EOL> init = self . call ( fn , init_args ) <EOL> index_offsets = ( <NUM_LIT:1> , ) * len ( axes ) <EOL> assert init . type == fn . return_type <EOL> nelts = self . sub ( nelts , self . int ( <NUM_LIT:1> ) , "<STR_LIT>" ) <EOL> else : <EOL> index_offsets = None <EOL> index_fn = self . indexify_fn ( fn , <EOL> axis , <EOL> args , <EOL> cartesian_product = False , <EOL> index_offsets = index_offsets ) <EOL> return IndexReduce ( fn = index_fn , <EOL> init = init , <EOL> combine = combine , <EOL> shape = nelts , <EOL> type = expr . type ) <EOL> def transform_Scan ( self , expr , output = None ) : <EOL> combine = expr . combine <EOL> init = expr . init <EOL> args = [ ] <EOL> axes = [ ] <EOL> for ( arg , axis ) in zip ( expr . args , self . normalize_axes ( expr . args , expr . axis ) ) : <EOL> if self . is_none ( axis ) : <EOL> args . append ( self . ravel ( arg ) ) <EOL> axis = <NUM_LIT:0> <EOL> else : <EOL> args . append ( arg ) <EOL> axes . append ( axis ) <EOL> bounds = self . iter_bounds ( args , axes ) <EOL> if self . is_none ( init ) : <EOL> assert len ( args ) == <NUM_LIT:1> <EOL> init_args = [ self . index_along_axis ( arg , axis , self . int ( <NUM_LIT:0> ) ) <EOL> for arg in args ] <EOL> init = self . call ( expr . fn , init_args ) <EOL> index_offsets = ( <NUM_LIT:1> , ) * len ( bounds ) <EOL> bounds = tuple ( self . sub ( bound , self . int ( <NUM_LIT:1> ) , "<STR_LIT>" ) for bound in bounds ) <EOL> else : <EOL> index_offsets = None <EOL> index_fn = self . indexify_fn ( expr . fn , <EOL> axis , <EOL> args , <EOL> cartesian_product = False , <EOL> index_offsets = index_offsets ) <EOL> return IndexScan ( fn = index_fn , <EOL> init = init , <EOL> combine = combine , <EOL> emit = expr . emit , <EOL> shape = bounds , <EOL> type = expr . type ) <EOL> def transform_Filter ( self , expr ) : <EOL> assert False , "<STR_LIT>" <EOL> def transform_FilterReduce ( self , expr ) : <EOL> raise NotImplemented <EOL> def transform_IndexFilter ( self , expr ) : <EOL> raise NotImplemented <EOL> def transform_IndexFilterReduce ( self , expr ) : <EOL> raise NotImplemented <EOL> def transform_Assign ( self , stmt ) : <EOL> """<STR_LIT>""" <EOL> if stmt . lhs . __class__ is Index : <EOL> rhs_class = stmt . rhs . __class__ <EOL> if rhs_class is Map : <EOL> self . transform_Map ( stmt . rhs , output = stmt . lhs ) <EOL> return None <EOL> elif rhs_class is OuterMap : <EOL> self . transform_OuterMap ( stmt . rhs , output = stmt . lhs ) <EOL> return Transform . transform_Assign ( self , stmt ) </s>
