<s> """<STR_LIT>""" <EOL> import collections <EOL> from jinja2 import Environment , PackageLoader <EOL> import math <EOL> import numpy as np <EOL> from pycparser import c_parser , c_ast , c_generator <EOL> import pyopencl as cl <EOL> from pyopencl import array <EOL> import device <EOL> from device import * <EOL> from logger import warning <EOL> import plan <EOL> import petsc_base <EOL> from profiling import lineprof <EOL> from utils import verify_reshape , uniquify , maybe_setflags <EOL> class Kernel ( device . Kernel ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , code , name , opts = { } , include_dirs = [ ] ) : <EOL> device . Kernel . __init__ ( self , code , name , opts , include_dirs ) <EOL> class Instrument ( c_ast . NodeVisitor ) : <EOL> """<STR_LIT>""" <EOL> def instrument ( self , ast , kernel_name , instrument , constants ) : <EOL> self . _kernel_name = kernel_name <EOL> self . _instrument = instrument <EOL> self . _ast = ast <EOL> self . _constants = constants <EOL> self . generic_visit ( ast ) <EOL> ast . ext . insert ( <NUM_LIT:0> , self . _func_node . decl ) <EOL> def visit_FuncDef ( self , node ) : <EOL> if node . decl . name == self . _kernel_name : <EOL> self . _func_node = node <EOL> self . visit ( node . decl ) <EOL> def visit_IdentifierType ( self , node ) : <EOL> if node . names == [ '<STR_LIT>' , '<STR_LIT>' ] : <EOL> node . names = [ '<STR_LIT>' ] <EOL> def visit_ParamList ( self , node ) : <EOL> for i , p in enumerate ( node . params ) : <EOL> if self . _instrument [ i ] [ <NUM_LIT:0> ] : <EOL> p . storage . append ( self . _instrument [ i ] [ <NUM_LIT:0> ] ) <EOL> if self . _instrument [ i ] [ <NUM_LIT:1> ] : <EOL> p . type . quals . append ( self . _instrument [ i ] [ <NUM_LIT:1> ] ) <EOL> self . visit ( p ) <EOL> for cst in self . _constants : <EOL> if cst . _is_scalar : <EOL> t = c_ast . TypeDecl ( cst . _name , [ ] , c_ast . IdentifierType ( [ cst . _cl_type ] ) ) <EOL> else : <EOL> t = c_ast . PtrDecl ( [ ] , c_ast . TypeDecl ( cst . _name , [ "<STR_LIT>" ] , <EOL> c_ast . IdentifierType ( [ cst . _cl_type ] ) ) ) <EOL> decl = c_ast . Decl ( cst . _name , [ ] , [ ] , [ ] , t , None , <NUM_LIT:0> ) <EOL> node . params . append ( decl ) <EOL> def instrument ( self , instrument , constants ) : <EOL> ast = c_parser . CParser ( ) . parse ( self . _code ) <EOL> Kernel . Instrument ( ) . instrument ( ast , self . _name , instrument , constants ) <EOL> return c_generator . CGenerator ( ) . visit ( ast ) <EOL> class Arg ( device . Arg ) : <EOL> """<STR_LIT>""" <EOL> def _indirect_kernel_arg_name ( self , idx , subset ) : <EOL> if self . _is_global : <EOL> if self . _is_global_reduction : <EOL> return self . _reduction_local_name <EOL> else : <EOL> return self . name <EOL> if self . _is_direct : <EOL> if self . data . soa : <EOL> return "<STR_LIT>" % ( self . name , idx ) <EOL> return "<STR_LIT>" % ( self . name , idx , <EOL> self . data . cdim ) <EOL> if self . _is_indirect : <EOL> if self . _is_vec_map : <EOL> return self . _vec_name <EOL> if self . access is device . INC : <EOL> return self . _local_name ( ) <EOL> else : <EOL> return "<STR_LIT>" % ( self . _shared_name , self . _which_indirect , idx , <EOL> self . data . cdim ) <EOL> def _direct_kernel_arg_name ( self , idx = None , subset = False ) : <EOL> if self . _is_mat : <EOL> return self . _mat_entry_name <EOL> if self . _is_staged_direct : <EOL> return self . _local_name ( ) <EOL> elif self . _is_global_reduction : <EOL> return self . _reduction_local_name <EOL> elif self . _is_global : <EOL> return self . name <EOL> else : <EOL> if subset : <EOL> return "<STR_LIT>" % ( self . name , idx ) <EOL> return "<STR_LIT>" % ( self . name , idx ) <EOL> class Subset ( device . Subset ) : <EOL> def _allocate_device ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _device_data = array . to_device ( _queue , self . indices ) <EOL> class DeviceDataMixin ( device . DeviceDataMixin ) : <EOL> """<STR_LIT>""" <EOL> ClTypeInfo = collections . namedtuple ( '<STR_LIT>' , <EOL> [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] ) <EOL> CL_TYPES = { np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT:0>' , '<STR_LIT:0>' , '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT:0>' , '<STR_LIT>' , '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT:0>' , '<STR_LIT:0>' , '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT:0>' , '<STR_LIT>' , '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT:int>' , '<STR_LIT:0>' , '<STR_LIT>' , <EOL> '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT:float>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ) , <EOL> np . dtype ( '<STR_LIT>' ) : ClTypeInfo ( '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ) } <EOL> def _allocate_device ( self ) : <EOL> if self . state is DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> if self . soa : <EOL> shape = tuple ( reversed ( self . shape ) ) <EOL> else : <EOL> shape = self . shape <EOL> self . _device_data = array . zeros ( _queue , shape = shape , <EOL> dtype = self . dtype ) <EOL> self . state = DeviceDataMixin . HOST <EOL> def _to_device ( self ) : <EOL> self . _allocate_device ( ) <EOL> if self . state is DeviceDataMixin . HOST : <EOL> self . _device_data . set ( self . _maybe_to_soa ( self . _data ) , <EOL> queue = _queue ) <EOL> self . state = DeviceDataMixin . BOTH <EOL> def _from_device ( self ) : <EOL> flag = self . _data . flags [ '<STR_LIT>' ] <EOL> maybe_setflags ( self . _data , write = True ) <EOL> if self . state is DeviceDataMixin . DEVICE : <EOL> self . _device_data . get ( _queue , self . _data ) <EOL> self . _data = self . _maybe_to_aos ( self . _data ) <EOL> self . state = DeviceDataMixin . BOTH <EOL> maybe_setflags ( self . _data , write = flag ) <EOL> @ property <EOL> def _cl_type ( self ) : <EOL> return DeviceDataMixin . CL_TYPES [ self . dtype ] . clstring <EOL> @ property <EOL> def _cl_type_zero ( self ) : <EOL> return DeviceDataMixin . CL_TYPES [ self . dtype ] . zero <EOL> @ property <EOL> def _cl_type_min ( self ) : <EOL> return DeviceDataMixin . CL_TYPES [ self . dtype ] . min <EOL> @ property <EOL> def _cl_type_max ( self ) : <EOL> return DeviceDataMixin . CL_TYPES [ self . dtype ] . max <EOL> class Dat ( device . Dat , petsc_base . Dat , DeviceDataMixin ) : <EOL> pass <EOL> class Sparsity ( device . Sparsity ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . _block_sparse = False <EOL> super ( Sparsity , self ) . __init__ ( * args , ** kwargs ) <EOL> @ property <EOL> def colidx ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> setattr ( self , '<STR_LIT>' , <EOL> array . to_device ( _queue , <EOL> self . _colidx ) ) <EOL> return getattr ( self , '<STR_LIT>' ) <EOL> @ property <EOL> def rowptr ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> setattr ( self , '<STR_LIT>' , <EOL> array . to_device ( _queue , <EOL> self . _rowptr ) ) <EOL> return getattr ( self , '<STR_LIT>' ) <EOL> class Mat ( device . Mat , DeviceDataMixin ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> class Const ( device . Const , DeviceDataMixin ) : <EOL> """<STR_LIT>""" <EOL> @ property <EOL> def _array ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> setattr ( self , '<STR_LIT>' , array . to_device ( _queue , self . _data ) ) <EOL> return getattr ( self , '<STR_LIT>' ) <EOL> class Global ( device . Global , DeviceDataMixin ) : <EOL> """<STR_LIT>""" <EOL> @ property <EOL> def _array ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _device_data = array . to_device ( _queue , self . _data ) <EOL> return self . _device_data <EOL> def _allocate_reduction_array ( self , nelems ) : <EOL> self . _d_reduc_array = array . zeros ( _queue , nelems * self . cdim , dtype = self . dtype ) <EOL> @ property <EOL> def data ( self ) : <EOL> base . _trace . evaluate ( set ( [ self ] ) , set ( ) ) <EOL> if self . state is DeviceDataMixin . DEVICE : <EOL> self . _array . get ( _queue , ary = self . _data ) <EOL> if self . state is not DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> self . state = DeviceDataMixin . HOST <EOL> return self . _data <EOL> @ data . setter <EOL> def data ( self , value ) : <EOL> base . _trace . evaluate ( set ( ) , set ( [ self ] ) ) <EOL> self . _data = verify_reshape ( value , self . dtype , self . dim ) <EOL> if self . state is not DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> self . state = DeviceDataMixin . HOST <EOL> def _post_kernel_reduction_task ( self , nelems , reduction_operator ) : <EOL> assert reduction_operator in [ INC , MIN , MAX ] <EOL> def generate_code ( ) : <EOL> def headers ( ) : <EOL> if self . dtype == np . dtype ( '<STR_LIT>' ) : <EOL> return """<STR_LIT>""" <EOL> else : <EOL> return "<STR_LIT>" <EOL> op = { INC : '<STR_LIT>' , MIN : '<STR_LIT>' , MAX : '<STR_LIT>' } [ reduction_operator ] <EOL> return """<STR_LIT>""" % { '<STR_LIT>' : headers ( ) , '<STR_LIT>' : self . cdim , '<STR_LIT:type>' : self . _cl_type , '<STR_LIT>' : op } <EOL> src , kernel = _reduction_task_cache . get ( <EOL> ( self . dtype , self . cdim , reduction_operator ) , ( None , None ) ) <EOL> if src is None : <EOL> src = generate_code ( ) <EOL> prg = cl . Program ( _ctx , src ) . build ( options = "<STR_LIT>" ) <EOL> name = "<STR_LIT>" % ( self . _cl_type , self . cdim ) <EOL> kernel = prg . __getattr__ ( name ) <EOL> _reduction_task_cache [ <EOL> ( self . dtype , self . cdim , reduction_operator ) ] = ( src , kernel ) <EOL> kernel . set_arg ( <NUM_LIT:0> , self . _array . data ) <EOL> kernel . set_arg ( <NUM_LIT:1> , self . _d_reduc_array . data ) <EOL> kernel . set_arg ( <NUM_LIT:2> , np . int32 ( nelems ) ) <EOL> cl . enqueue_task ( _queue , kernel ) . wait ( ) <EOL> self . _array . get ( queue = _queue , ary = self . _data ) <EOL> self . state = DeviceDataMixin . BOTH <EOL> del self . _d_reduc_array <EOL> class Map ( device . Map ) : <EOL> """<STR_LIT>""" <EOL> def _to_device ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _device_values = array . to_device ( _queue , self . _values ) <EOL> class Plan ( plan . Plan ) : <EOL> @ property <EOL> def ind_map ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _ind_map_array = array . to_device ( _queue , super ( Plan , self ) . ind_map ) <EOL> return self . _ind_map_array <EOL> @ property <EOL> def ind_sizes ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _ind_sizes_array = array . to_device ( _queue , super ( Plan , self ) . ind_sizes ) <EOL> return self . _ind_sizes_array <EOL> @ property <EOL> def ind_offs ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _ind_offs_array = array . to_device ( _queue , super ( Plan , self ) . ind_offs ) <EOL> return self . _ind_offs_array <EOL> @ property <EOL> def loc_map ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _loc_map_array = array . to_device ( _queue , super ( Plan , self ) . loc_map ) <EOL> return self . _loc_map_array <EOL> @ property <EOL> def blkmap ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _blkmap_array = array . to_device ( _queue , super ( Plan , self ) . blkmap ) <EOL> return self . _blkmap_array <EOL> @ property <EOL> def offset ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _offset_array = array . to_device ( _queue , super ( Plan , self ) . offset ) <EOL> return self . _offset_array <EOL> @ property <EOL> def nelems ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _nelems_array = array . to_device ( _queue , super ( Plan , self ) . nelems ) <EOL> return self . _nelems_array <EOL> @ property <EOL> def nthrcol ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _nthrcol_array = array . to_device ( _queue , super ( Plan , self ) . nthrcol ) <EOL> return self . _nthrcol_array <EOL> @ property <EOL> def thrcol ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _thrcol_array = array . to_device ( _queue , super ( Plan , self ) . thrcol ) <EOL> return self . _thrcol_array <EOL> class Solver ( petsc_base . Solver ) : <EOL> def _solve ( self , A , x , b ) : <EOL> x . _from_device ( ) <EOL> b . _from_device ( ) <EOL> super ( Solver , self ) . _solve ( A , x , b ) <EOL> if x . state is not DeviceDataMixin . DEVICE_UNALLOCATED : <EOL> x . state = DeviceDataMixin . HOST <EOL> x . _to_device ( ) <EOL> class JITModule ( base . JITModule ) : <EOL> @ classmethod <EOL> def _cache_key ( cls , kernel , itspace , * args , ** kwargs ) : <EOL> parloop = kwargs . get ( '<STR_LIT>' ) <EOL> key = ( parloop . _is_direct , len ( parloop . _unique_dat_args ) , len ( parloop . _aliased_dat_args ) ) <EOL> return base . JITModule . _cache_key ( kernel , itspace , * args ) + key + ( kwargs [ '<STR_LIT>' ] [ '<STR_LIT>' ] , ) <EOL> def __init__ ( self , kernel , itspace_extents , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if self . _initialized : <EOL> return <EOL> self . _parloop = kwargs . get ( '<STR_LIT>' ) <EOL> self . _kernel = self . _parloop . _kernel <EOL> self . _conf = kwargs . get ( '<STR_LIT>' ) <EOL> self . _initialized = True <EOL> def compile ( self ) : <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> return self . _fun <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> def instrument_user_kernel ( ) : <EOL> inst = [ ] <EOL> for arg in self . _parloop . args : <EOL> i = None <EOL> if self . _parloop . _is_direct : <EOL> if ( arg . _is_direct and ( arg . data . _is_scalar or arg . data . soa ) ) or ( arg . _is_global and not arg . _is_global_reduction ) : <EOL> i = ( "<STR_LIT>" , None ) <EOL> else : <EOL> i = ( "<STR_LIT>" , None ) <EOL> else : <EOL> if arg . _is_direct or ( arg . _is_global and not arg . _is_global_reduction ) : <EOL> i = ( "<STR_LIT>" , None ) <EOL> elif ( arg . _is_indirect or arg . _is_vec_map ) and not arg . _is_indirect_reduction : <EOL> i = ( "<STR_LIT>" , None ) <EOL> else : <EOL> i = ( "<STR_LIT>" , None ) <EOL> inst . append ( i ) <EOL> for i in self . _parloop . _it_space . extents : <EOL> inst . append ( ( "<STR_LIT>" , None ) ) <EOL> return self . _parloop . _kernel . instrument ( inst , Const . _definitions ( ) ) <EOL> user_kernel = instrument_user_kernel ( ) <EOL> template = _jinja2_direct_loop if self . _parloop . _is_direct else _jinja2_indirect_loop <EOL> src = template . render ( { '<STR_LIT>' : self . _parloop , <EOL> '<STR_LIT>' : user_kernel , <EOL> '<STR_LIT>' : self . _conf , <EOL> '<STR_LIT>' : { '<STR_LIT>' : _AMD_fixes } , <EOL> '<STR_LIT>' : Const . _definitions ( ) <EOL> } ) . encode ( "<STR_LIT:ascii>" ) <EOL> self . _dump_generated_code ( src , ext = "<STR_LIT>" ) <EOL> prg = cl . Program ( _ctx , src ) . build ( ) <EOL> self . _fun = prg . __getattr__ ( self . _parloop . _stub_name ) <EOL> del self . _parloop <EOL> del self . _kernel <EOL> del self . _conf <EOL> return self . _fun <EOL> def __call__ ( self , thread_count , work_group_size , * args ) : <EOL> fun = self . compile ( ) <EOL> for i , arg in enumerate ( args ) : <EOL> fun . set_arg ( i , arg ) <EOL> with timed_region ( "<STR_LIT>" ) : <EOL> cl . enqueue_nd_range_kernel ( _queue , fun , ( thread_count , ) , <EOL> ( work_group_size , ) , g_times_l = False ) . wait ( ) <EOL> class ParLoop ( device . ParLoop ) : <EOL> @ property <EOL> def _matrix_args ( self ) : <EOL> return [ a for a in self . args if a . _is_mat ] <EOL> @ property <EOL> def _unique_matrix ( self ) : <EOL> return list ( uniquify ( a . data for a in self . _matrix_args ) ) <EOL> @ property <EOL> def _matrix_entry_maps ( self ) : <EOL> """<STR_LIT>""" <EOL> return list ( uniquify ( m for arg in self . args if arg . _is_mat for m in arg . map ) ) <EOL> @ property <EOL> def _requires_matrix_coloring ( self ) : <EOL> """<STR_LIT>""" <EOL> return not _supports_64b_atomics and not not self . _matrix_args <EOL> def _i_partition_size ( self ) : <EOL> warning ( '<STR_LIT>' ) <EOL> available_local_memory = _max_local_memory - <NUM_LIT> <EOL> available_local_memory -= <NUM_LIT:16> <EOL> available_local_memory -= ( _address_bits / <NUM_LIT:8> ) * ( len ( <EOL> self . _unique_dat_args ) + len ( self . _all_global_non_reduction_args ) ) <EOL> available_local_memory -= ( _address_bits / <NUM_LIT:8> ) * len ( self . _unique_indirect_dat_args ) <EOL> available_local_memory -= ( _address_bits / <NUM_LIT:8> ) * len ( self . _all_global_reduction_args ) <EOL> available_local_memory -= ( _address_bits / <NUM_LIT:8> ) * len ( self . _all_indirect_args ) <EOL> available_local_memory -= ( _address_bits / <NUM_LIT:8> ) * <NUM_LIT:7> <EOL> available_local_memory -= <NUM_LIT:4> <EOL> available_local_memory -= <NUM_LIT:7> <EOL> available_local_memory -= <NUM_LIT:12> + <NUM_LIT:3> <EOL> available_local_memory -= <NUM_LIT:4> + ( _address_bits / <NUM_LIT:8> ) * <NUM_LIT:2> * len ( self . _unique_indirect_dat_args ) <EOL> available_local_memory -= <NUM_LIT:2> * ( len ( self . _unique_indirect_dat_args ) - <NUM_LIT:1> ) <EOL> max_bytes = sum ( map ( lambda a : a . data . _bytes_per_elem , self . _all_indirect_args ) ) <EOL> return available_local_memory / ( <NUM_LIT:2> * _warpsize * max_bytes ) * ( <NUM_LIT:2> * _warpsize ) <EOL> def launch_configuration ( self ) : <EOL> if self . _is_direct : <EOL> per_elem_max_local_mem_req = self . _max_shared_memory_needed_per_set_element <EOL> shared_memory_offset = per_elem_max_local_mem_req * _warpsize <EOL> if per_elem_max_local_mem_req == <NUM_LIT:0> : <EOL> wgs = _max_work_group_size <EOL> else : <EOL> warning ( '<STR_LIT>' ) <EOL> available_local_memory = _max_local_memory - <NUM_LIT> <EOL> available_local_memory -= <NUM_LIT:16> <EOL> available_local_memory -= ( len ( self . _unique_dat_args ) + <EOL> len ( self . _all_global_non_reduction_args ) ) * ( _address_bits / <NUM_LIT:8> ) <EOL> available_local_memory -= len ( <EOL> self . _all_global_reduction_args ) * ( _address_bits / <NUM_LIT:8> ) <EOL> available_local_memory -= <NUM_LIT:7> <EOL> ps = available_local_memory / per_elem_max_local_mem_req <EOL> wgs = min ( _max_work_group_size , ( ps / _warpsize ) * _warpsize ) <EOL> nwg = min ( _pref_work_group_count , int ( <EOL> math . ceil ( self . _it_space . size / float ( wgs ) ) ) ) <EOL> ttc = wgs * nwg <EOL> local_memory_req = per_elem_max_local_mem_req * wgs <EOL> return { '<STR_LIT>' : ttc , <EOL> '<STR_LIT>' : wgs , <EOL> '<STR_LIT>' : nwg , <EOL> '<STR_LIT>' : local_memory_req , <EOL> '<STR_LIT>' : shared_memory_offset } <EOL> else : <EOL> return { '<STR_LIT>' : self . _i_partition_size ( ) } <EOL> @ collective <EOL> @ lineprof <EOL> def _compute ( self , part , fun , * arglist ) : <EOL> if part . size == <NUM_LIT:0> : <EOL> return <EOL> conf = self . launch_configuration ( ) <EOL> conf [ '<STR_LIT>' ] = isinstance ( part . set , Subset ) <EOL> if self . _is_indirect : <EOL> _plan = Plan ( part , <EOL> * self . _unwound_args , <EOL> partition_size = conf [ '<STR_LIT>' ] , <EOL> matrix_coloring = self . _requires_matrix_coloring ) <EOL> conf [ '<STR_LIT>' ] = _plan . nshared <EOL> conf [ '<STR_LIT>' ] = _plan . ninds <EOL> conf [ '<STR_LIT>' ] = min ( _max_work_group_size , <EOL> conf [ '<STR_LIT>' ] ) <EOL> conf [ '<STR_LIT>' ] = _plan . nblocks <EOL> conf [ '<STR_LIT>' ] = _warpsize <EOL> conf [ '<STR_LIT>' ] = self . _it_space . size <EOL> fun = JITModule ( self . kernel , self . it_space , * self . args , parloop = self , conf = conf ) <EOL> args = [ ] <EOL> for arg in self . _unique_args : <EOL> arg . data . _allocate_device ( ) <EOL> if arg . access is not device . WRITE : <EOL> arg . data . _to_device ( ) <EOL> for a in self . _unique_dat_args : <EOL> args . append ( a . data . array . data ) <EOL> for a in self . _all_global_non_reduction_args : <EOL> args . append ( a . data . _array . data ) <EOL> for a in self . _all_global_reduction_args : <EOL> a . data . _allocate_reduction_array ( conf [ '<STR_LIT>' ] ) <EOL> args . append ( a . data . _d_reduc_array . data ) <EOL> for cst in Const . _definitions ( ) : <EOL> args . append ( cst . _array . data ) <EOL> for m in self . _unique_matrix : <EOL> args . append ( m . _dev_array . data ) <EOL> m . _to_device ( ) <EOL> args . append ( m . _rowptr . data ) <EOL> args . append ( m . _colidx . data ) <EOL> for m in self . _matrix_entry_maps : <EOL> m . _to_device ( ) <EOL> args . append ( m . _device_values . data ) <EOL> if self . _is_direct : <EOL> args . append ( np . int32 ( part . size ) ) <EOL> args . append ( np . int32 ( part . offset ) ) <EOL> if conf [ '<STR_LIT>' ] : <EOL> part . set . _allocate_device ( ) <EOL> args . append ( part . set . _device_data . data ) <EOL> fun ( conf [ '<STR_LIT>' ] , conf [ '<STR_LIT>' ] , * args ) <EOL> else : <EOL> args . append ( np . int32 ( part . size ) ) <EOL> args . append ( np . int32 ( part . offset ) ) <EOL> if conf [ '<STR_LIT>' ] : <EOL> part . set . _allocate_device ( ) <EOL> args . append ( part . set . _device_data . data ) <EOL> args . append ( _plan . ind_map . data ) <EOL> args . append ( _plan . loc_map . data ) <EOL> args . append ( _plan . ind_sizes . data ) <EOL> args . append ( _plan . ind_offs . data ) <EOL> args . append ( _plan . blkmap . data ) <EOL> args . append ( _plan . offset . data ) <EOL> args . append ( _plan . nelems . data ) <EOL> args . append ( _plan . nthrcol . data ) <EOL> args . append ( _plan . thrcol . data ) <EOL> block_offset = <NUM_LIT:0> <EOL> args . append ( <NUM_LIT:0> ) <EOL> for i in range ( _plan . ncolors ) : <EOL> blocks_per_grid = int ( _plan . ncolblk [ i ] ) <EOL> threads_per_block = min ( _max_work_group_size , conf [ '<STR_LIT>' ] ) <EOL> thread_count = threads_per_block * blocks_per_grid <EOL> args [ - <NUM_LIT:1> ] = np . int32 ( block_offset ) <EOL> fun ( int ( thread_count ) , int ( threads_per_block ) , * args ) <EOL> block_offset += blocks_per_grid <EOL> for arg in self . args : <EOL> if arg . access is not READ : <EOL> arg . data . state = DeviceDataMixin . DEVICE <EOL> for a in self . _all_global_reduction_args : <EOL> a . data . _post_kernel_reduction_task ( conf [ '<STR_LIT>' ] , a . access ) <EOL> def _setup ( ) : <EOL> global _ctx <EOL> global _queue <EOL> global _pref_work_group_count <EOL> global _max_local_memory <EOL> global _address_bits <EOL> global _max_work_group_size <EOL> global _has_dpfloat <EOL> global _warpsize <EOL> global _AMD_fixes <EOL> global _reduction_task_cache <EOL> global _supports_64b_atomics <EOL> _ctx = cl . create_some_context ( ) <EOL> _queue = cl . CommandQueue ( _ctx , properties = cl . command_queue_properties . PROFILING_ENABLE ) <EOL> _pref_work_group_count = _queue . device . max_compute_units <EOL> _max_local_memory = _queue . device . local_mem_size <EOL> _address_bits = _queue . device . address_bits <EOL> _max_work_group_size = _queue . device . max_work_group_size <EOL> _has_dpfloat = '<STR_LIT>' in _queue . device . extensions or '<STR_LIT>' in _queue . device . extensions <EOL> if not _has_dpfloat : <EOL> warning ( '<STR_LIT>' ) <EOL> if '<STR_LIT>' in _queue . device . extensions : <EOL> _supports_64b_atomics = True <EOL> if _queue . device . type == cl . device_type . CPU : <EOL> _warpsize = <NUM_LIT:1> <EOL> elif _queue . device . type == cl . device_type . GPU : <EOL> _warpsize = <NUM_LIT:32> <EOL> _AMD_fixes = _queue . device . platform . vendor in [ '<STR_LIT>' ] <EOL> _reduction_task_cache = dict ( ) <EOL> _supports_64b_atomics = False <EOL> _debug = False <EOL> _ctx = None <EOL> _queue = None <EOL> _pref_work_group_count = <NUM_LIT:0> <EOL> _max_local_memory = <NUM_LIT:0> <EOL> _address_bits = <NUM_LIT:32> <EOL> _max_work_group_size = <NUM_LIT:0> <EOL> _has_dpfloat = False <EOL> _warpsize = <NUM_LIT:0> <EOL> _AMD_fixes = False <EOL> _reduction_task_cache = None <EOL> _jinja2_env = Environment ( loader = PackageLoader ( "<STR_LIT>" , "<STR_LIT>" ) ) <EOL> _jinja2_direct_loop = _jinja2_env . get_template ( "<STR_LIT>" ) <EOL> _jinja2_indirect_loop = _jinja2_env . get_template ( "<STR_LIT>" ) </s>
