<s> import base <EOL> from base import * <EOL> from coffee . plan import ASTKernel <EOL> from mpi import collective <EOL> class Kernel ( base . Kernel ) : <EOL> def _ast_to_c ( self , ast , opts = { } ) : <EOL> """<STR_LIT>""" <EOL> ast_handler = ASTKernel ( ast ) <EOL> ast_handler . plan_gpu ( ) <EOL> return ast . gencode ( ) <EOL> def __init__ ( self , code , name , opts = { } , include_dirs = [ ] ) : <EOL> if self . _initialized : <EOL> return <EOL> self . _code = preprocess ( self . _ast_to_c ( code , opts ) , include_dirs ) <EOL> super ( Kernel , self ) . __init__ ( self . _code , name , opts = opts , include_dirs = include_dirs ) <EOL> class Arg ( base . Arg ) : <EOL> @ property <EOL> def name ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _is_indirect : <EOL> return "<STR_LIT>" % self . indirect_position <EOL> return "<STR_LIT>" % self . position <EOL> @ property <EOL> def _lmaoffset_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _shared_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> def _local_name ( self , idx = None ) : <EOL> if self . _is_direct : <EOL> return "<STR_LIT>" % self . name <EOL> else : <EOL> if self . _is_vec_map and idx is not None : <EOL> return "<STR_LIT>" % ( self . name , self . _which_indirect + idx ) <EOL> if self . _uses_itspace : <EOL> if idx is not None : <EOL> return "<STR_LIT>" % ( self . name , self . _which_indirect + idx ) <EOL> return "<STR_LIT>" % ( self . name , self . idx . index ) <EOL> return "<STR_LIT>" % ( self . name , self . idx ) <EOL> @ property <EOL> def _reduction_local_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _reduction_tmp_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _reduction_kernel_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _vec_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _map_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _size_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _mat_entry_name ( self ) : <EOL> return "<STR_LIT>" % self . name <EOL> @ property <EOL> def _is_staged_direct ( self ) : <EOL> return self . _is_direct and not ( self . data . _is_scalar or self . _is_soa ) <EOL> class DeviceDataMixin ( object ) : <EOL> DEVICE_UNALLOCATED = '<STR_LIT>' <EOL> HOST_UNALLOCATED = '<STR_LIT>' <EOL> DEVICE = '<STR_LIT>' <EOL> HOST = '<STR_LIT>' <EOL> BOTH = '<STR_LIT>' <EOL> @ property <EOL> def _bytes_per_elem ( self ) : <EOL> return self . dtype . itemsize * self . cdim <EOL> @ property <EOL> def _is_scalar ( self ) : <EOL> return self . cdim == <NUM_LIT:1> <EOL> @ property <EOL> def state ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _state <EOL> @ state . setter <EOL> def state ( self , value ) : <EOL> self . _state = value <EOL> @ property <EOL> @ collective <EOL> def data ( self ) : <EOL> """<STR_LIT>""" <EOL> base . _trace . evaluate ( self , self ) <EOL> if len ( self . _data ) is <NUM_LIT:0> and self . dataset . total_size > <NUM_LIT:0> : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> maybe_setflags ( self . _data , write = True ) <EOL> self . needs_halo_update = True <EOL> self . _from_device ( ) <EOL> if self . state is not DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> self . state = DeviceDataMixin . HOST <EOL> return self . _data [ : self . dataset . size ] <EOL> @ data . setter <EOL> @ collective <EOL> def data ( self , value ) : <EOL> base . _trace . evaluate ( set ( ) , set ( [ self ] ) ) <EOL> maybe_setflags ( self . _data , write = True ) <EOL> self . needs_halo_update = True <EOL> self . _data = verify_reshape ( value , self . dtype , self . shape ) <EOL> if self . state is not DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> self . state = DeviceDataMixin . HOST <EOL> @ property <EOL> def data_ro ( self ) : <EOL> """<STR_LIT>""" <EOL> base . _trace . evaluate ( reads = self ) <EOL> if len ( self . _data ) is <NUM_LIT:0> and self . dataset . total_size > <NUM_LIT:0> : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> maybe_setflags ( self . _data , write = True ) <EOL> self . _from_device ( ) <EOL> if self . state is not DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> self . state = DeviceDataMixin . BOTH <EOL> maybe_setflags ( self . _data , write = False ) <EOL> v = self . _data [ : self . dataset . size ] . view ( ) <EOL> v . setflags ( write = False ) <EOL> return v <EOL> def _maybe_to_soa ( self , data ) : <EOL> """<STR_LIT>""" <EOL> if self . soa : <EOL> shape = data . T . shape <EOL> return data . T . ravel ( ) . reshape ( shape ) <EOL> return data <EOL> def _maybe_to_aos ( self , data ) : <EOL> """<STR_LIT>""" <EOL> if self . soa : <EOL> tshape = data . T . shape <EOL> shape = data . shape <EOL> return data . reshape ( tshape ) . T . ravel ( ) . reshape ( shape ) <EOL> return data <EOL> def _allocate_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def _to_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def _from_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> class Dat ( DeviceDataMixin , base . Dat ) : <EOL> def __init__ ( self , dataset , data = None , dtype = None , name = None , <EOL> soa = None , uid = None ) : <EOL> self . state = DeviceDataMixin . DEVICE_UNALLOCATED <EOL> base . Dat . __init__ ( self , dataset , data , dtype , name , soa , uid ) <EOL> @ property <EOL> def array ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _device_data <EOL> @ array . setter <EOL> def array ( self , ary ) : <EOL> assert not getattr ( self , '<STR_LIT>' ) or self . shape == ary . shape <EOL> self . _device_data = ary <EOL> self . state = DeviceDataMixin . DEVICE <EOL> def _check_shape ( self , other ) : <EOL> """<STR_LIT>""" <EOL> if not self . shape == other . shape : <EOL> raise ValueError ( "<STR_LIT>" <EOL> % ( self . shape , other . shape ) ) <EOL> def halo_exchange_begin ( self ) : <EOL> if self . dataset . halo is None : <EOL> return <EOL> maybe_setflags ( self . _data , write = True ) <EOL> self . _from_device ( ) <EOL> super ( Dat , self ) . halo_exchange_begin ( ) <EOL> def halo_exchange_end ( self ) : <EOL> if self . dataset . halo is None : <EOL> return <EOL> maybe_setflags ( self . _data , write = True ) <EOL> super ( Dat , self ) . halo_exchange_end ( ) <EOL> if self . state in [ DeviceDataMixin . DEVICE , <EOL> DeviceDataMixin . BOTH ] : <EOL> self . _halo_to_device ( ) <EOL> self . state = DeviceDataMixin . DEVICE <EOL> def _halo_to_device ( self ) : <EOL> _lim = self . dataset . size * self . dataset . cdim <EOL> self . _device_data . ravel ( ) [ _lim : ] . set ( self . _data [ self . dataset . size : ] ) <EOL> class Const ( DeviceDataMixin , base . Const ) : <EOL> def __init__ ( self , dim , data , name , dtype = None ) : <EOL> base . Const . __init__ ( self , dim , data , name , dtype ) <EOL> self . state = DeviceDataMixin . HOST <EOL> @ property <EOL> def data ( self ) : <EOL> """<STR_LIT>""" <EOL> self . state = DeviceDataMixin . HOST <EOL> return self . _data <EOL> @ data . setter <EOL> def data ( self , value ) : <EOL> self . _data = verify_reshape ( value , self . dtype , self . dim ) <EOL> self . state = DeviceDataMixin . HOST <EOL> def _to_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def _from_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" % self ) <EOL> class Global ( DeviceDataMixin , base . Global ) : <EOL> def __init__ ( self , dim , data = None , dtype = None , name = None ) : <EOL> base . Global . __init__ ( self , dim , data , dtype , name ) <EOL> self . state = DeviceDataMixin . DEVICE_UNALLOCATED <EOL> @ property <EOL> def data_ro ( self ) : <EOL> return self . data <EOL> class Map ( base . Map ) : <EOL> def __init__ ( self , iterset , dataset , arity , values = None , name = None , <EOL> offset = None , parent = None , bt_masks = None ) : <EOL> base . Map . __init__ ( self , iterset , dataset , arity , values , name , offset , <EOL> parent , bt_masks ) <EOL> if len ( self . values_with_halo ) == <NUM_LIT:0> and self . iterset . total_size > <NUM_LIT:0> : <EOL> raise MapValueError ( "<STR_LIT>" ) <EOL> def _to_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def _from_device ( self ) : <EOL> """<STR_LIT>""" <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> class Mat ( base . Mat ) : <EOL> def __init__ ( self , datasets , dtype = None , name = None ) : <EOL> base . Mat . __init__ ( self , datasets , dtype , name ) <EOL> self . state = DeviceDataMixin . DEVICE_UNALLOCATED <EOL> class ParLoop ( base . ParLoop ) : <EOL> def __init__ ( self , kernel , itspace , * args , ** kwargs ) : <EOL> base . ParLoop . __init__ ( self , kernel , itspace , * args , ** kwargs ) <EOL> self . __unwound_args = [ ] <EOL> self . __unique_args = [ ] <EOL> self . _arg_dict = { } <EOL> seen = set ( ) <EOL> c = <NUM_LIT:0> <EOL> for arg in self . _actual_args : <EOL> if arg . _is_mat : <EOL> for a in arg : <EOL> self . __unwound_args . append ( a ) <EOL> elif arg . _is_vec_map or arg . _uses_itspace : <EOL> for d , m in zip ( arg . data , arg . map ) : <EOL> for i in range ( m . arity ) : <EOL> a = d ( arg . access , m [ i ] ) <EOL> a . position = arg . position <EOL> self . __unwound_args . append ( a ) <EOL> else : <EOL> for a in arg : <EOL> self . __unwound_args . append ( a ) <EOL> if arg . _is_dat : <EOL> key = ( arg . data , arg . map ) <EOL> if arg . _is_indirect : <EOL> arg . _which_indirect = c <EOL> if arg . _is_vec_map or arg . _flatten : <EOL> c += arg . map . arity <EOL> elif arg . _uses_itspace : <EOL> c += self . _it_space . extents [ arg . idx . index ] <EOL> else : <EOL> c += <NUM_LIT:1> <EOL> if key not in seen : <EOL> self . __unique_args . append ( arg ) <EOL> seen . add ( key ) <EOL> else : <EOL> self . __unique_args . append ( arg ) <EOL> def _get_arg_list ( self , propname , arglist_name , keep = lambda x : True ) : <EOL> attr = self . _arg_dict . get ( propname ) <EOL> if attr : <EOL> return attr <EOL> attr = filter ( keep , getattr ( self , arglist_name ) ) <EOL> self . _arg_dict [ propname ] = attr <EOL> return attr <EOL> @ property <EOL> def _is_direct ( self ) : <EOL> for arg in self . __unwound_args : <EOL> if arg . _is_indirect : <EOL> return False <EOL> return True <EOL> @ property <EOL> def _is_indirect ( self ) : <EOL> return not self . _is_direct <EOL> @ property <EOL> def _max_shared_memory_needed_per_set_element ( self ) : <EOL> staged = self . _all_staged_direct_args <EOL> reduction = self . _all_global_reduction_args <EOL> smax = <NUM_LIT:0> <EOL> rmax = <NUM_LIT:0> <EOL> if staged : <EOL> smax = max ( a . data . _bytes_per_elem for a in staged ) <EOL> if reduction : <EOL> rmax = max ( a . dtype . itemsize for a in reduction ) <EOL> return max ( smax , rmax ) <EOL> @ property <EOL> def _stub_name ( self ) : <EOL> return "<STR_LIT>" % self . kernel . name <EOL> @ property <EOL> def _has_itspace ( self ) : <EOL> return len ( self . _it_space . extents ) > <NUM_LIT:0> <EOL> @ property <EOL> def _needs_shared_memory ( self ) : <EOL> if self . _is_indirect : <EOL> return True <EOL> for arg in self . _actual_args : <EOL> if arg . _is_global_reduction : <EOL> return True <EOL> if arg . _is_staged_direct : <EOL> return True <EOL> return False <EOL> @ property <EOL> def _requires_coloring ( self ) : <EOL> """<STR_LIT>""" <EOL> return not not self . _all_inc_indirect_dat_args or self . _requires_matrix_coloring <EOL> @ property <EOL> def _requires_matrix_coloring ( self ) : <EOL> """<STR_LIT>""" <EOL> return False <EOL> @ property <EOL> def _unique_args ( self ) : <EOL> return self . __unique_args <EOL> @ property <EOL> def _unwound_args ( self ) : <EOL> return self . __unwound_args <EOL> @ property <EOL> def _unwound_indirect_args ( self ) : <EOL> keep = lambda x : x . _is_indirect <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _unique_dat_args ( self ) : <EOL> keep = lambda x : x . _is_dat <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _aliased_dat_args ( self ) : <EOL> keep = lambda x : x . _is_dat and all ( x is not y for y in self . _unique_dat_args ) <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _unique_vec_map_args ( self ) : <EOL> keep = lambda x : x . _is_vec_map <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _unique_indirect_dat_args ( self ) : <EOL> keep = lambda x : x . _is_indirect <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _unique_read_or_rw_indirect_dat_args ( self ) : <EOL> keep = lambda x : x . _is_indirect and x . access in [ READ , RW ] <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _unique_write_or_rw_indirect_dat_args ( self ) : <EOL> keep = lambda x : x . _is_indirect and x . access in [ WRITE , RW ] <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _unique_inc_indirect_dat_args ( self ) : <EOL> keep = lambda x : x . _is_indirect and x . access is INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_inc_indirect_dat_args ( self ) : <EOL> keep = lambda x : x . _is_indirect and x . access is INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_inc_non_vec_map_indirect_dat_args ( self ) : <EOL> keep = lambda x : x . _is_indirect and x . access is INC and not ( x . _is_vec_map or x . _uses_itspace ) <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_vec_map_args ( self ) : <EOL> keep = lambda x : x . _is_vec_map <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_itspace_dat_args ( self ) : <EOL> keep = lambda x : x . _is_dat and x . _uses_itspace <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_inc_itspace_dat_args ( self ) : <EOL> keep = lambda x : x . access is INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_non_inc_itspace_dat_args ( self ) : <EOL> keep = lambda x : x . access is not INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_inc_vec_map_args ( self ) : <EOL> keep = lambda x : x . _is_vec_map and x . access is INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_non_inc_vec_map_args ( self ) : <EOL> keep = lambda x : x . _is_vec_map and x . access is not INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_vec_like_args ( self ) : <EOL> keep = lambda x : x . _is_vec_map or ( x . _is_dat and x . _uses_itspace ) <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_inc_vec_like_args ( self ) : <EOL> keep = lambda x : x . access is INC <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_indirect_args ( self ) : <EOL> keep = lambda x : x . _is_indirect <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_direct_args ( self ) : <EOL> keep = lambda x : x . _is_direct <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_staged_direct_args ( self ) : <EOL> keep = lambda x : x . _is_staged_direct <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_staged_in_direct_args ( self ) : <EOL> keep = lambda x : x . access is not WRITE <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_staged_out_direct_args ( self ) : <EOL> keep = lambda x : x . access is not READ <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_global_reduction_args ( self ) : <EOL> keep = lambda x : x . _is_global_reduction <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _all_global_non_reduction_args ( self ) : <EOL> keep = lambda x : x . _is_global and not x . _is_global_reduction <EOL> return self . _get_arg_list ( '<STR_LIT>' , <EOL> '<STR_LIT>' , keep ) <EOL> @ property <EOL> def _has_matrix_arg ( self ) : <EOL> return any ( arg . _is_mat for arg in self . _unique_args ) </s>
