<s> from . . import syntax <EOL> from . . builder import Builder <EOL> from . . ndtypes import SliceT , make_array_type , ArrayT <EOL> from . . ndtypes import NoneT , ScalarT , Int64 , PtrT , IntT <EOL> from . . syntax import Const , Index , Tuple , Var , ArrayView , Assign , Slice , Struct <EOL> from . . syntax . helpers import zero_i64 , one_i64 <EOL> from transform import Transform <EOL> class LowerIndexing ( Transform ) : <EOL> def pre_apply ( self , fn ) : <EOL> self . bindings = { } <EOL> def tuple_proj ( self , tup , idx , explicit_struct = False ) : <EOL> if tup . __class__ is Var and tup . name in self . bindings : <EOL> stored = self . bindings [ tup . name ] <EOL> if stored . __class__ is Tuple : <EOL> return stored . elts [ idx ] <EOL> else : <EOL> return stored . args [ idx ] <EOL> else : <EOL> return Builder . tuple_proj ( self , tup , idx , <EOL> explicit_struct = explicit_struct ) <EOL> def attr ( self , obj , field ) : <EOL> if obj . __class__ is Var and obj . name in self . bindings : <EOL> stored = self . bindings [ obj . name ] <EOL> stored_class = stored . __class__ <EOL> if stored_class is Struct : <EOL> pos = stored . type . field_pos ( field ) <EOL> return stored . args [ pos ] <EOL> elif stored_class is Slice or stored_class is ArrayView : <EOL> return getattr ( stored , field ) <EOL> return Builder . attr ( self , obj , field ) <EOL> def array_slice ( self , arr , indices ) : <EOL> data_ptr = self . attr ( arr , "<STR_LIT:data>" ) <EOL> shape = self . attr ( arr , "<STR_LIT>" ) <EOL> strides = self . attr ( arr , "<STR_LIT>" ) <EOL> elt_offset = self . attr ( arr , "<STR_LIT>" ) <EOL> size = self . attr ( arr , "<STR_LIT:size>" ) <EOL> new_strides = [ ] <EOL> new_shape = [ ] <EOL> for ( i , idx ) in enumerate ( indices ) : <EOL> stride_i = self . tuple_proj ( strides , i ) <EOL> shape_i = self . tuple_proj ( shape , i ) <EOL> idx_t = idx . type <EOL> if isinstance ( idx_t , ScalarT ) : <EOL> offset_i = self . mul ( idx , stride_i , "<STR_LIT>" % i ) <EOL> elt_offset = self . add ( elt_offset , offset_i ) <EOL> elif idx_t . __class__ is NoneT : <EOL> new_strides . append ( stride_i ) <EOL> new_shape . append ( shape_i ) <EOL> elif idx_t . __class__ is SliceT : <EOL> if isinstance ( idx_t . start_type , NoneT ) : <EOL> start = zero_i64 <EOL> else : <EOL> start = self . attr ( idx , "<STR_LIT:start>" ) <EOL> if isinstance ( idx_t . step_type , NoneT ) : <EOL> step = one_i64 <EOL> else : <EOL> step = self . attr ( idx , "<STR_LIT>" ) <EOL> if isinstance ( idx_t . stop_type , NoneT ) : <EOL> stop = shape_i <EOL> else : <EOL> stop = self . attr ( idx , "<STR_LIT>" ) <EOL> offset_i = self . mul ( start , stride_i , "<STR_LIT>" % i ) <EOL> elt_offset = self . add ( elt_offset , offset_i ) <EOL> dim_i = self . sub ( stop , start , "<STR_LIT>" % i ) <EOL> new_shape . append ( self . cast ( dim_i , Int64 ) ) <EOL> new_strides . append ( self . mul ( stride_i , step ) ) <EOL> else : <EOL> raise RuntimeError ( "<STR_LIT>" % idx_t ) <EOL> elt_t = arr . type . elt_type <EOL> new_rank = len ( new_strides ) <EOL> new_array_t = make_array_type ( elt_t , new_rank ) <EOL> new_strides = self . tuple ( new_strides , "<STR_LIT>" ) <EOL> new_shape = self . tuple ( new_shape , "<STR_LIT>" ) <EOL> return ArrayView ( data_ptr , new_shape , new_strides , <EOL> elt_offset , size , <EOL> type = new_array_t ) <EOL> def transform_Index ( self , expr ) : <EOL> arr = self . transform_expr ( expr . value ) <EOL> idx = self . transform_expr ( expr . index ) <EOL> idx = self . assign_name ( idx , "<STR_LIT>" ) <EOL> arr_t = arr . type <EOL> if arr_t . __class__ is PtrT : <EOL> assert isinstance ( idx . type , IntT ) <EOL> return expr <EOL> assert arr_t . __class__ is ArrayT , "<STR_LIT>" % ( arr , arr . type ) <EOL> if self . is_tuple ( idx ) : <EOL> indices = self . tuple_elts ( idx ) <EOL> else : <EOL> indices = [ idx ] <EOL> n_given = len ( indices ) <EOL> n_required = arr_t . rank <EOL> if n_given < n_required : <EOL> extra_indices = [ syntax . helpers . slice_none ] * ( n_required - n_given ) <EOL> indices . extend ( extra_indices ) <EOL> if syntax . helpers . all_scalars ( indices ) : <EOL> data_ptr = self . attr ( arr , "<STR_LIT:data>" ) <EOL> strides = self . attr ( arr , "<STR_LIT>" ) <EOL> offset_elts = self . attr ( arr , "<STR_LIT>" ) <EOL> for ( i , idx_i ) in enumerate ( indices ) : <EOL> stride_i = self . tuple_proj ( strides , i ) <EOL> elts_i = self . mul ( stride_i , idx_i , "<STR_LIT>" % i ) <EOL> offset_elts = self . add ( offset_elts , elts_i , "<STR_LIT>" ) <EOL> return self . index ( data_ptr , offset_elts , temp = False ) <EOL> else : <EOL> return self . array_slice ( arr , indices ) <EOL> def transform_Assign ( self , stmt ) : <EOL> lhs = stmt . lhs <EOL> lhs_class = lhs . __class__ <EOL> rhs = self . transform_expr ( stmt . rhs ) <EOL> if lhs_class is Tuple : <EOL> for ( i , _ ) in enumerate ( lhs . type . elt_types ) : <EOL> lhs_i = self . tuple_proj ( lhs , i ) <EOL> rhs_i = self . tuple_proj ( rhs , i ) <EOL> assert lhs_i . __class__ not in ( ArrayView , Tuple ) <EOL> self . assign ( lhs_i , rhs_i ) <EOL> return None <EOL> elif lhs_class is Index : <EOL> lhs = self . transform_Index ( lhs ) <EOL> if lhs . __class__ is ArrayView : <EOL> copy_loop = self . array_copy ( src = rhs , dest = lhs , return_stmt = True ) <EOL> copy_loop = self . transform_stmt ( copy_loop ) <EOL> return copy_loop <EOL> elif lhs_class is Var and stmt . rhs . __class__ in ( Slice , Struct , ArrayView , Tuple ) : <EOL> self . bindings [ lhs . name ] = rhs <EOL> return Assign ( lhs , rhs ) </s>
