<s> """<STR_LIT>""" <EOL> from contextlib import contextmanager <EOL> from collections import OrderedDict <EOL> from copy import deepcopy as dcopy , copy as scopy <EOL> import os <EOL> from base import * <EOL> import base <EOL> import compilation <EOL> import host <EOL> from backends import _make_object <EOL> from caching import Cached <EOL> from profiling import lineprof , timed_region , profile <EOL> from logger import warning , info as log_info <EOL> from mpi import collective <EOL> from configuration import configuration <EOL> from utils import flatten , strip , as_tuple <EOL> import coffee <EOL> from coffee import base as ast <EOL> from coffee . utils import ast_make_alias <EOL> from coffee . visitors import FindInstances , SymbolReferences <EOL> try : <EOL> import slope_python as slope <EOL> except ImportError : <EOL> slope = None <EOL> class Arg ( host . Arg ) : <EOL> @ staticmethod <EOL> def specialize ( args , gtl_map , loop_id ) : <EOL> """<STR_LIT>""" <EOL> def convert ( arg , gtl_map , loop_id ) : <EOL> maps = as_tuple ( arg . map , Map ) <EOL> c_local_maps = [ None ] * len ( maps ) <EOL> for i , map in enumerate ( maps ) : <EOL> c_local_maps [ i ] = [ None ] * len ( map ) <EOL> for j , m in enumerate ( map ) : <EOL> c_local_maps [ i ] [ j ] = gtl_map [ "<STR_LIT>" % ( m . name , i , j ) ] <EOL> _arg = Arg ( arg . data , arg . map , arg . idx , arg . access , arg . _flatten ) <EOL> _arg . _loop_position = loop_id <EOL> _arg . position = arg . position <EOL> _arg . indirect_position = arg . indirect_position <EOL> _arg . _c_local_maps = c_local_maps <EOL> return _arg <EOL> try : <EOL> return [ convert ( arg , gtl_map , loop_id ) for arg in args ] <EOL> except TypeError : <EOL> return convert ( args , gtl_map , loop_id ) <EOL> @ staticmethod <EOL> def filter_args ( loop_args ) : <EOL> """<STR_LIT>""" <EOL> filtered_args = OrderedDict ( ) <EOL> for args in loop_args : <EOL> for a in args : <EOL> filtered_args [ a . data ] = filtered_args . get ( a . data , a ) <EOL> if a . access != filtered_args [ a . data ] . access : <EOL> if READ in [ a . access , filtered_args [ a . data ] . access ] : <EOL> filtered_args [ a . data ] . _access = RW <EOL> elif WRITE in [ a . access , filtered_args [ a . data ] . access ] : <EOL> filtered_args [ a . data ] . _access = WRITE <EOL> else : <EOL> filtered_args [ a . data ] . _access = RW <EOL> return filtered_args <EOL> def c_arg_bindto ( self , arg ) : <EOL> """<STR_LIT>""" <EOL> if self . ctype != arg . ctype : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> return "<STR_LIT>" % ( self . ctype , self . c_arg_name ( ) , arg . c_arg_name ( ) ) <EOL> def c_map_name ( self , i , j ) : <EOL> return self . _c_local_maps [ i ] [ j ] <EOL> @ property <EOL> def name ( self ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" % ( self . _loop_position , self . _position ) <EOL> class Kernel ( host . Kernel , tuple ) : <EOL> """<STR_LIT>""" <EOL> @ classmethod <EOL> def _cache_key ( cls , kernels , fused_ast = None , loop_chain_index = None ) : <EOL> keys = "<STR_LIT>" . join ( [ super ( Kernel , cls ) . _cache_key ( k . _code or k . _ast . gencode ( ) , <EOL> k . _name , k . _opts , k . _include_dirs , <EOL> k . _headers , k . _user_code ) <EOL> for k in kernels ] ) <EOL> return str ( loop_chain_index ) + keys <EOL> def _ast_to_c ( self , asts , opts ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( asts , ( ast . FunDecl , ast . Root ) ) : <EOL> asts = ast . Root ( asts ) <EOL> self . _ast = asts <EOL> self . _original_ast = dcopy ( asts ) <EOL> return super ( Kernel , self ) . _ast_to_c ( self . _ast , opts ) <EOL> def __init__ ( self , kernels , fused_ast = None , loop_chain_index = None ) : <EOL> """<STR_LIT>""" <EOL> if self . _initialized : <EOL> return <EOL> asts = fused_ast <EOL> if not asts : <EOL> kernels = dict ( zip ( [ k . cache_key [ <NUM_LIT:1> : ] for k in kernels ] , kernels ) ) . values ( ) <EOL> asts = [ k . _ast for k in kernels ] <EOL> kernels = as_tuple ( kernels , ( Kernel , host . Kernel , base . Kernel ) ) <EOL> Kernel . _globalcount += <NUM_LIT:1> <EOL> self . _kernels = kernels <EOL> self . _name = "<STR_LIT:_>" . join ( [ k . name for k in kernels ] ) <EOL> self . _opts = dict ( flatten ( [ k . _opts . items ( ) for k in kernels ] ) ) <EOL> self . _applied_blas = any ( k . _applied_blas for k in kernels ) <EOL> self . _include_dirs = list ( set ( flatten ( [ k . _include_dirs for k in kernels ] ) ) ) <EOL> self . _headers = list ( set ( flatten ( [ k . _headers for k in kernels ] ) ) ) <EOL> self . _user_code = "<STR_LIT:\n>" . join ( list ( set ( [ k . _user_code for k in kernels ] ) ) ) <EOL> self . _attached_info = False <EOL> self . _ast = asts <EOL> self . _code = None <EOL> self . _initialized = True <EOL> def __iter__ ( self ) : <EOL> for k in self . _kernels : <EOL> yield k <EOL> def __str__ ( self ) : <EOL> return "<STR_LIT>" % self . _name <EOL> class JITModule ( host . JITModule ) : <EOL> _cppargs = [ ] <EOL> _libraries = [ ] <EOL> _extension = '<STR_LIT>' <EOL> _wrapper = """<STR_LIT>""" <EOL> _kernel_wrapper = """<STR_LIT>""" <EOL> @ classmethod <EOL> def _cache_key ( cls , kernel , it_space , * args , ** kwargs ) : <EOL> key = ( hash ( kwargs [ '<STR_LIT>' ] ) , ) <EOL> all_args = kwargs [ '<STR_LIT>' ] <EOL> for kernel_i , it_space_i , args_i in zip ( kernel , it_space , all_args ) : <EOL> key += super ( JITModule , cls ) . _cache_key ( kernel_i , it_space_i , * args_i ) <EOL> return key <EOL> def __init__ ( self , kernel , it_space , * args , ** kwargs ) : <EOL> if self . _initialized : <EOL> return <EOL> self . _all_args = kwargs . pop ( '<STR_LIT>' ) <EOL> self . _executor = kwargs . pop ( '<STR_LIT>' ) <EOL> self . _it_space = it_space <EOL> super ( JITModule , self ) . __init__ ( kernel , it_space , * args , ** kwargs ) <EOL> def set_argtypes ( self , iterset , * args ) : <EOL> argtypes = [ slope . Executor . meta [ '<STR_LIT>' ] ] <EOL> for it_space in self . _it_space : <EOL> if isinstance ( it_space . iterset , Subset ) : <EOL> argtypes . append ( it_space . iterset . _argtype ) <EOL> for arg in args : <EOL> if arg . _is_mat : <EOL> argtypes . append ( arg . data . _argtype ) <EOL> else : <EOL> for d in arg . data : <EOL> argtypes . append ( d . _argtype ) <EOL> if arg . _is_indirect or arg . _is_mat : <EOL> maps = as_tuple ( arg . Map , Map ) <EOL> for map in maps : <EOL> for m in map : <EOL> argtypes . append ( m . _argtype ) <EOL> for c in Const . _definitions ( ) : <EOL> argtypes . append ( c . _argtype ) <EOL> return argtypes <EOL> def compile ( self ) : <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> self . _kernel = Kernel ( self . _kernel ) <EOL> slope_dir = os . environ [ '<STR_LIT>' ] <EOL> self . _kernel . _name = '<STR_LIT>' <EOL> self . _kernel . _headers . extend ( slope . Executor . meta [ '<STR_LIT>' ] ) <EOL> self . _kernel . _include_dirs . extend ( [ '<STR_LIT>' % ( slope_dir , <EOL> slope . get_include_dir ( ) ) ] ) <EOL> self . _libraries += [ '<STR_LIT>' % ( slope_dir , slope . get_lib_dir ( ) ) , <EOL> '<STR_LIT>' % slope . get_lib_name ( ) ] <EOL> compiler = coffee . plan . compiler . get ( '<STR_LIT:name>' ) <EOL> self . _cppargs += slope . get_compile_opts ( compiler ) <EOL> fun = super ( JITModule , self ) . compile ( ) <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> del self . _all_args <EOL> del self . _executor <EOL> return fun <EOL> def generate_code ( self ) : <EOL> indent = lambda t , i : ( '<STR_LIT:\n>' + '<STR_LIT:U+0020>' * i ) . join ( t . split ( '<STR_LIT:\n>' ) ) <EOL> code_dict = { } <EOL> code_dict [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> code_dict [ '<STR_LIT>' ] = "<STR_LIT>" % ( slope . Executor . meta [ '<STR_LIT>' ] , <EOL> slope . Executor . meta [ '<STR_LIT>' ] ) <EOL> _wrapper_args = '<STR_LIT:U+002CU+0020>' . join ( [ arg . c_wrapper_arg ( ) for arg in self . _args ] ) <EOL> _wrapper_decs = '<STR_LIT>' . join ( [ arg . c_wrapper_dec ( ) for arg in self . _args ] ) <EOL> if len ( Const . _defs ) > <NUM_LIT:0> : <EOL> _const_args = '<STR_LIT:U+002CU+0020>' <EOL> _const_args += '<STR_LIT:U+002CU+0020>' . join ( [ c_const_arg ( c ) for c in Const . _definitions ( ) ] ) <EOL> else : <EOL> _const_args = '<STR_LIT>' <EOL> _const_inits = '<STR_LIT>' . join ( [ c_const_init ( c ) for c in Const . _definitions ( ) ] ) <EOL> code_dict [ '<STR_LIT>' ] = _wrapper_args <EOL> code_dict [ '<STR_LIT>' ] = _const_args <EOL> code_dict [ '<STR_LIT>' ] = indent ( _wrapper_decs , <NUM_LIT:1> ) <EOL> code_dict [ '<STR_LIT>' ] = indent ( _const_inits , <NUM_LIT:1> ) <EOL> _loop_chain_body , _user_code , _ssinds_arg = [ ] , [ ] , [ ] <EOL> for i , loop in enumerate ( zip ( self . _kernel , self . _itspace , self . _all_args ) ) : <EOL> kernel , it_space , args = loop <EOL> loop_code_dict = host . JITModule ( kernel , it_space , * args ) . generate_code ( ) <EOL> args_dict = dict ( zip ( [ _a . data for _a in self . _args ] , self . _args ) ) <EOL> binding = OrderedDict ( zip ( args , [ args_dict [ a . data ] for a in args ] ) ) <EOL> if len ( binding ) != len ( args ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> binding = '<STR_LIT>' . join ( [ a0 . c_arg_bindto ( a1 ) for a0 , a1 in binding . items ( ) ] ) <EOL> loop_code_dict [ '<STR_LIT>' ] = binding <EOL> loop_code_dict [ '<STR_LIT>' ] = self . _executor . gtl_maps [ i ] [ '<STR_LIT>' ] <EOL> loop_code_dict [ '<STR_LIT>' ] = self . _executor . c_loop_init [ i ] <EOL> loop_code_dict [ '<STR_LIT>' ] = slope . Executor . meta [ '<STR_LIT>' ] <EOL> loop_code_dict [ '<STR_LIT>' ] = slope . Executor . meta [ '<STR_LIT>' ] <EOL> _loop_chain_body . append ( strip ( JITModule . _kernel_wrapper % loop_code_dict ) ) <EOL> _user_code . append ( kernel . _user_code ) <EOL> _ssinds_arg . append ( loop_code_dict [ '<STR_LIT>' ] ) <EOL> _loop_chain_body = "<STR_LIT>" . join ( _loop_chain_body ) <EOL> _user_code = "<STR_LIT:\n>" . join ( _user_code ) <EOL> _ssinds_arg = "<STR_LIT:U+002CU+0020>" . join ( [ s for s in _ssinds_arg if s ] ) <EOL> code_dict [ '<STR_LIT>' ] = indent ( _user_code , <NUM_LIT:1> ) <EOL> code_dict [ '<STR_LIT>' ] = _ssinds_arg <EOL> executor_code = indent ( self . _executor . c_code ( indent ( _loop_chain_body , <NUM_LIT:2> ) ) , <NUM_LIT:1> ) <EOL> code_dict [ '<STR_LIT>' ] = executor_code <EOL> return code_dict <EOL> class ParLoop ( host . ParLoop ) : <EOL> def __init__ ( self , kernel , it_space , * args , ** kwargs ) : <EOL> read_args = [ a . data for a in args if a . access in [ READ , RW ] ] <EOL> written_args = [ a . data for a in args if a . access in [ RW , WRITE , MIN , MAX , INC ] ] <EOL> inc_args = [ a . data for a in args if a . access in [ INC ] ] <EOL> LazyComputation . __init__ ( self , set ( read_args ) | Const . _defs , <EOL> set ( written_args ) , set ( inc_args ) ) <EOL> self . _kernel = kernel <EOL> self . _actual_args = args <EOL> self . _it_space = it_space <EOL> for i , arg in enumerate ( self . _actual_args ) : <EOL> arg . position = i <EOL> arg . indirect_position = i <EOL> for i , arg1 in enumerate ( self . _actual_args ) : <EOL> if arg1 . _is_dat and arg1 . _is_indirect : <EOL> for arg2 in self . _actual_args [ i : ] : <EOL> if arg2 . data is arg1 . data and arg2 . map is arg1 . map : <EOL> arg2 . indirect_position = arg1 . indirect_position <EOL> self . _all_args = kwargs . get ( '<STR_LIT>' , [ args ] ) <EOL> self . _inspection = kwargs . get ( '<STR_LIT>' ) <EOL> self . _executor = kwargs . get ( '<STR_LIT>' ) <EOL> @ collective <EOL> @ profile <EOL> def compute ( self ) : <EOL> """<STR_LIT>""" <EOL> arglist = self . prepare_arglist ( None , * self . args ) <EOL> with timed_region ( "<STR_LIT>" ) : <EOL> self . _compute ( * arglist ) <EOL> def prepare_arglist ( self , part , * args ) : <EOL> arglist = [ self . _inspection ] <EOL> for it_space in self . it_space : <EOL> if isinstance ( it_space . _iterset , Subset ) : <EOL> arglist . append ( it_space . _iterset . _indices . ctypes . data ) <EOL> for arg in args : <EOL> if arg . _is_mat : <EOL> arglist . append ( arg . data . handle . handle ) <EOL> else : <EOL> for d in arg . data : <EOL> arglist . append ( d . _data . ctypes . data ) <EOL> if arg . _is_indirect or arg . _is_mat : <EOL> maps = as_tuple ( arg . map , Map ) <EOL> for map in maps : <EOL> for m in map : <EOL> arglist . append ( m . _values . ctypes . data ) <EOL> for c in Const . _definitions ( ) : <EOL> arglist . append ( c . _data . ctypes . data ) <EOL> return arglist <EOL> @ collective <EOL> @ lineprof <EOL> def _compute ( self , * arglist ) : <EOL> kwargs = { <EOL> '<STR_LIT>' : self . _all_args , <EOL> '<STR_LIT>' : self . _executor , <EOL> } <EOL> fun = JITModule ( self . kernel , self . it_space , * self . args , ** kwargs ) <EOL> with timed_region ( "<STR_LIT>" ) : <EOL> fun ( * arglist ) <EOL> class Schedule ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , kernel ) : <EOL> self . _kernel = list ( kernel ) <EOL> def __call__ ( self , loop_chain ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> class PlainSchedule ( Schedule ) : <EOL> def __init__ ( self ) : <EOL> super ( PlainSchedule , self ) . __init__ ( [ ] ) <EOL> def __call__ ( self , loop_chain ) : <EOL> return loop_chain <EOL> class FusionSchedule ( Schedule ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , kernels , offsets ) : <EOL> super ( FusionSchedule , self ) . __init__ ( kernels ) <EOL> offsets = [ <NUM_LIT:0> ] + list ( offsets ) <EOL> loop_indices = [ range ( offsets [ i ] , o ) for i , o in enumerate ( offsets [ <NUM_LIT:1> : ] ) ] <EOL> self . _info = [ { '<STR_LIT>' : li } for li in loop_indices ] <EOL> def __call__ ( self , loop_chain ) : <EOL> fused_par_loops = [ ] <EOL> for kernel , info in zip ( self . _kernel , self . _info ) : <EOL> loop_indices = info [ '<STR_LIT>' ] <EOL> extra_args = info . get ( '<STR_LIT>' ) <EOL> iterregion = loop_chain [ loop_indices [ <NUM_LIT:0> ] ] . iteration_region <EOL> iterset = loop_chain [ loop_indices [ <NUM_LIT:0> ] ] . it_space . iterset <EOL> loops = [ loop_chain [ i ] for i in loop_indices ] <EOL> args = Arg . filter_args ( [ loop . args for loop in loops ] ) . values ( ) <EOL> extra_args = [ Dat ( * d ) ( * a ) for d , a in extra_args ] if extra_args else [ ] <EOL> args += extra_args <EOL> fused_par_loops . append ( _make_object ( '<STR_LIT>' , kernel , iterset , * args , <EOL> ** { '<STR_LIT>' : iterregion } ) ) <EOL> return fused_par_loops <EOL> class HardFusionSchedule ( FusionSchedule ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , schedule , fused ) : <EOL> self . _schedule = schedule <EOL> self . _fused = fused <EOL> self . _info = dcopy ( schedule . _info ) <EOL> for i , info in enumerate ( schedule . _info ) : <EOL> for k , v in info . items ( ) : <EOL> self . _info [ i ] [ k ] = [ i ] if k == '<STR_LIT>' else v <EOL> kernel = scopy ( schedule . _kernel ) <EOL> for fused_kernel , fused_map in fused : <EOL> base , fuse = fused_kernel . _kernels <EOL> base_idx , fuse_idx = kernel . index ( base ) , kernel . index ( fuse ) <EOL> pos = min ( base_idx , fuse_idx ) <EOL> kernel . insert ( pos , fused_kernel ) <EOL> self . _info [ pos ] [ '<STR_LIT>' ] = [ base_idx , fuse_idx ] <EOL> self . _info [ pos ] [ '<STR_LIT>' ] = [ ( ( fused_map . toset , None , np . int32 ) , <EOL> ( RW , fused_map ) ) ] <EOL> kernel . pop ( pos + <NUM_LIT:1> ) <EOL> pos = max ( base_idx , fuse_idx ) <EOL> self . _info . pop ( pos ) <EOL> kernel . pop ( pos ) <EOL> self . _kernel = kernel <EOL> def __call__ ( self , loop_chain , only_hard = False ) : <EOL> if not only_hard : <EOL> loop_chain = self . _schedule ( loop_chain ) <EOL> return super ( HardFusionSchedule , self ) . __call__ ( loop_chain ) <EOL> class TilingSchedule ( Schedule ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , schedule , inspection , executor ) : <EOL> self . _schedule = schedule <EOL> self . _inspection = inspection <EOL> self . _executor = executor <EOL> def __call__ ( self , loop_chain ) : <EOL> loop_chain = self . _schedule ( loop_chain ) <EOL> args = Arg . filter_args ( [ loop . args for loop in loop_chain ] ) . values ( ) <EOL> kernel = tuple ( ( loop . kernel for loop in loop_chain ) ) <EOL> all_args = tuple ( ( Arg . specialize ( loop . args , gtl_map , i ) for i , ( loop , gtl_map ) <EOL> in enumerate ( zip ( loop_chain , self . _executor . gtl_maps ) ) ) ) <EOL> it_space = tuple ( ( loop . it_space for loop in loop_chain ) ) <EOL> kwargs = { <EOL> '<STR_LIT>' : self . _inspection , <EOL> '<STR_LIT>' : all_args , <EOL> '<STR_LIT>' : self . _executor <EOL> } <EOL> return [ ParLoop ( kernel , it_space , * args , ** kwargs ) ] <EOL> class Inspector ( Cached ) : <EOL> """<STR_LIT>""" <EOL> _cache = { } <EOL> _modes = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> @ classmethod <EOL> def _cache_key ( cls , name , loop_chain , tile_size ) : <EOL> key = ( name , tile_size ) <EOL> for loop in loop_chain : <EOL> if isinstance ( loop , Mat . _Assembly ) : <EOL> continue <EOL> key += ( hash ( str ( loop . kernel . _ast ) ) , ) <EOL> for arg in loop . args : <EOL> if arg . _is_global : <EOL> key += ( arg . data . dim , arg . data . dtype , arg . access ) <EOL> elif arg . _is_dat : <EOL> if isinstance ( arg . idx , IterationIndex ) : <EOL> idx = ( arg . idx . __class__ , arg . idx . index ) <EOL> else : <EOL> idx = arg . idx <EOL> map_arity = arg . map . arity if arg . map else None <EOL> key += ( arg . data . dim , arg . data . dtype , map_arity , idx , arg . access ) <EOL> elif arg . _is_mat : <EOL> idxs = ( arg . idx [ <NUM_LIT:0> ] . __class__ , arg . idx [ <NUM_LIT:0> ] . index , <EOL> arg . idx [ <NUM_LIT:1> ] . index ) <EOL> map_arities = ( arg . map [ <NUM_LIT:0> ] . arity , arg . map [ <NUM_LIT:1> ] . arity ) <EOL> key += ( arg . data . dims , arg . data . dtype , idxs , map_arities , arg . access ) <EOL> return key <EOL> def __init__ ( self , name , loop_chain , tile_size ) : <EOL> if self . _initialized : <EOL> return <EOL> if not hasattr ( self , '<STR_LIT>' ) : <EOL> self . _inspected = <NUM_LIT:0> <EOL> self . _name = name <EOL> self . _tile_size = tile_size <EOL> self . _loop_chain = loop_chain <EOL> def inspect ( self , mode ) : <EOL> """<STR_LIT>""" <EOL> self . _inspected += <NUM_LIT:1> <EOL> if self . _heuristic_skip_inspection ( mode ) : <EOL> del self . _name <EOL> del self . _loop_chain <EOL> del self . _tile_size <EOL> return PlainSchedule ( ) <EOL> elif hasattr ( self , '<STR_LIT>' ) : <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> if self . mode != mode : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> return self . _schedule <EOL> elif not hasattr ( self , '<STR_LIT>' ) : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> if mode not in Inspector . _modes : <EOL> raise TypeError ( "<STR_LIT>" , <EOL> str ( Inspector . _modes ) ) <EOL> self . _mode = mode <EOL> mode = Inspector . _modes . index ( mode ) <EOL> with timed_region ( "<STR_LIT>" % self . _name ) : <EOL> self . _soft_fuse ( ) <EOL> if mode > <NUM_LIT:0> : <EOL> self . _hard_fuse ( ) <EOL> if mode > <NUM_LIT:1> : <EOL> self . _tile ( ) <EOL> self . _initialized = True <EOL> del self . _name <EOL> del self . _loop_chain <EOL> del self . _tile_size <EOL> return self . _schedule <EOL> def _heuristic_skip_inspection ( self , mode ) : <EOL> """<STR_LIT>""" <EOL> if mode == '<STR_LIT>' and self . _inspected < <NUM_LIT:2> : <EOL> return True <EOL> return False <EOL> def _filter_kernel_args ( self , loops , fundecl ) : <EOL> """<STR_LIT>""" <EOL> fused_loop_args = list ( flatten ( [ l . args for l in loops ] ) ) <EOL> unique_fused_loop_args = Arg . filter_args ( [ l . args for l in loops ] ) <EOL> fused_kernel_args = fundecl . args <EOL> binding = OrderedDict ( zip ( fused_loop_args , fused_kernel_args ) ) <EOL> new_fused_kernel_args , args_maps = [ ] , [ ] <EOL> for fused_loop_arg , fused_kernel_arg in binding . items ( ) : <EOL> unique_fused_loop_arg = unique_fused_loop_args [ fused_loop_arg . data ] <EOL> if fused_loop_arg is unique_fused_loop_arg : <EOL> new_fused_kernel_args . append ( fused_kernel_arg ) <EOL> continue <EOL> tobind_fused_kernel_arg = binding [ unique_fused_loop_arg ] <EOL> if tobind_fused_kernel_arg . is_const : <EOL> if fused_loop_arg . _is_written : <EOL> tobind_fused_kernel_arg . qual . remove ( '<STR_LIT>' ) <EOL> elif '<STR_LIT>' not in fused_kernel_arg . qual : <EOL> fused_kernel_arg . qual . append ( '<STR_LIT>' ) <EOL> binding [ fused_loop_arg ] = tobind_fused_kernel_arg <EOL> if fused_kernel_arg . sym . symbol == tobind_fused_kernel_arg . sym . symbol : <EOL> continue <EOL> alias = ast_make_alias ( dcopy ( fused_kernel_arg ) , <EOL> dcopy ( tobind_fused_kernel_arg ) ) <EOL> args_maps . append ( alias ) <EOL> fundecl . children [ <NUM_LIT:0> ] . children = args_maps + fundecl . children [ <NUM_LIT:0> ] . children <EOL> fundecl . args = new_fused_kernel_args <EOL> return binding <EOL> def _soft_fuse ( self ) : <EOL> """<STR_LIT>""" <EOL> def fuse ( self , loops , loop_chain_index ) : <EOL> kernels = [ l . kernel for l in loops ] <EOL> fuse_asts = [ k . _original_ast if k . _code else k . _ast for k in kernels ] <EOL> base_ast = dcopy ( fuse_asts [ <NUM_LIT:0> ] ) <EOL> retval = FindInstances . default_retval ( ) <EOL> base_fundecl = FindInstances ( ast . FunDecl ) . visit ( base_ast , ret = retval ) [ ast . FunDecl ] <EOL> if len ( base_fundecl ) != <NUM_LIT:1> : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> base_fundecl = base_fundecl [ <NUM_LIT:0> ] <EOL> for unique_id , _fuse_ast in enumerate ( fuse_asts [ <NUM_LIT:1> : ] , <NUM_LIT:1> ) : <EOL> fuse_ast = dcopy ( _fuse_ast ) <EOL> retval = FindInstances . default_retval ( ) <EOL> fuse_fundecl = FindInstances ( ast . FunDecl ) . visit ( fuse_ast , ret = retval ) [ ast . FunDecl ] <EOL> if len ( fuse_fundecl ) != <NUM_LIT:1> : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> fuse_fundecl = fuse_fundecl [ <NUM_LIT:0> ] <EOL> base_fundecl . name = "<STR_LIT>" % ( base_fundecl . name , fuse_fundecl . name ) <EOL> base_fundecl . args . extend ( fuse_fundecl . args ) <EOL> retval = SymbolReferences . default_retval ( ) <EOL> fuse_symbols = SymbolReferences ( ) . visit ( fuse_ast , ret = retval ) <EOL> for decl in fuse_fundecl . args : <EOL> for symbol , _ in fuse_symbols [ decl . sym . symbol ] : <EOL> symbol . symbol = "<STR_LIT>" % ( symbol . symbol , unique_id ) <EOL> base_fundecl . children [ <NUM_LIT:0> ] = ast . Block ( <EOL> [ ast . Block ( base_fundecl . children [ <NUM_LIT:0> ] . children , open_scope = True ) , <EOL> ast . FlatBlock ( "<STR_LIT>" ) , <EOL> ast . Block ( fuse_fundecl . children [ <NUM_LIT:0> ] . children , open_scope = True ) ] ) <EOL> self . _filter_kernel_args ( loops , base_fundecl ) <EOL> fused_ast = base_ast <EOL> return Kernel ( kernels , fused_ast , loop_chain_index ) <EOL> fused , fusing = [ ] , [ self . _loop_chain [ <NUM_LIT:0> ] ] <EOL> for i , loop in enumerate ( self . _loop_chain [ <NUM_LIT:1> : ] ) : <EOL> base_loop = fusing [ - <NUM_LIT:1> ] <EOL> if base_loop . it_space != loop . it_space or ( base_loop . is_indirect and loop . is_indirect ) : <EOL> fused . append ( ( fuse ( self , fusing , len ( fused ) ) , i + <NUM_LIT:1> ) ) <EOL> fusing = [ loop ] <EOL> elif ( base_loop . is_direct and loop . is_direct ) or ( base_loop . is_direct and loop . is_indirect ) or ( base_loop . is_indirect and loop . is_direct ) : <EOL> fusing . append ( loop ) <EOL> else : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> if fusing : <EOL> fused . append ( ( fuse ( self , fusing , len ( fused ) ) , len ( self . _loop_chain ) ) ) <EOL> fused_kernels , offsets = zip ( * fused ) <EOL> self . _schedule = FusionSchedule ( fused_kernels , offsets ) <EOL> self . _loop_chain = self . _schedule ( self . _loop_chain ) <EOL> def _hard_fuse ( self ) : <EOL> """<STR_LIT>""" <EOL> def has_raw_or_war ( loop1 , loop2 ) : <EOL> return loop2 . reads & loop1 . writes or loop2 . writes & loop1 . reads or loop1 . incs & ( loop2 . writes - loop2 . incs ) or loop2 . incs & ( loop1 . writes - loop1 . incs ) <EOL> def has_iai ( loop1 , loop2 ) : <EOL> return loop1 . incs & loop2 . incs <EOL> def fuse ( base_loop , loop_chain , fused ) : <EOL> """<STR_LIT>""" <EOL> for loop in loop_chain : <EOL> if has_raw_or_war ( loop , base_loop ) : <EOL> return <EOL> if loop . it_space == base_loop . it_space : <EOL> warning ( "<STR_LIT>" ) <EOL> continue <EOL> common_inc_data = has_iai ( base_loop , loop ) <EOL> if not common_inc_data : <EOL> continue <EOL> common_incs = [ a for a in base_loop . args + loop . args <EOL> if a . data in common_inc_data ] <EOL> maps = list ( set ( flatten ( [ a . map for a in common_incs ] ) ) ) <EOL> maps += [ m . factors for m in maps if hasattr ( m , '<STR_LIT>' ) ] <EOL> maps = list ( flatten ( maps ) ) <EOL> set1 , set2 = base_loop . it_space . iterset , loop . it_space . iterset <EOL> fused_map = [ m for m in maps if set1 == m . iterset and set2 == m . toset ] <EOL> if fused_map : <EOL> fused . append ( ( base_loop , loop , fused_map [ <NUM_LIT:0> ] , common_incs [ <NUM_LIT:1> ] ) ) <EOL> return <EOL> fused_map = [ m for m in maps if set1 == m . toset and set2 == m . iterset ] <EOL> if fused_map : <EOL> fused . append ( ( loop , base_loop , fused_map [ <NUM_LIT:0> ] , common_incs [ <NUM_LIT:0> ] ) ) <EOL> return <EOL> fused = [ ] <EOL> for i , l in enumerate ( self . _loop_chain , <NUM_LIT:1> ) : <EOL> fuse ( l , self . _loop_chain [ i : ] , fused ) <EOL> if not fused : <EOL> return <EOL> _fused = [ ] <EOL> for base_loop , fuse_loop , fused_map , fused_inc_arg in fused : <EOL> base , fuse = base_loop . kernel , fuse_loop . kernel <EOL> base_ast = dcopy ( base . _original_ast ) if base . _code else dcopy ( base . _ast ) <EOL> retval = FindInstances . default_retval ( ) <EOL> base_info = FindInstances ( ( ast . FunDecl , ast . PreprocessNode ) ) . visit ( base_ast , ret = retval ) <EOL> base_headers = base_info [ ast . PreprocessNode ] <EOL> base_fundecl = base_info [ ast . FunDecl ] <EOL> fuse_ast = dcopy ( fuse . _original_ast ) if fuse . _code else dcopy ( fuse . _ast ) <EOL> retval = FindInstances . default_retval ( ) <EOL> fuse_info = FindInstances ( ( ast . FunDecl , ast . PreprocessNode ) ) . visit ( fuse_ast , ret = retval ) <EOL> fuse_headers = fuse_info [ ast . PreprocessNode ] <EOL> fuse_fundecl = fuse_info [ ast . FunDecl ] <EOL> retval = SymbolReferences . default_retval ( ) <EOL> fuse_symbol_refs = SymbolReferences ( ) . visit ( fuse_ast , ret = retval ) <EOL> if len ( base_fundecl ) != <NUM_LIT:1> or len ( fuse_fundecl ) != <NUM_LIT:1> : <EOL> raise RuntimeError ( "<STR_LIT>" ) <EOL> base_fundecl = base_fundecl [ <NUM_LIT:0> ] <EOL> fuse_fundecl = fuse_fundecl [ <NUM_LIT:0> ] <EOL> body = ast . Block ( [ ] ) <EOL> fusion_name = '<STR_LIT>' % ( base_fundecl . name , fuse_fundecl . name ) <EOL> fusion_args = base_fundecl . args + fuse_fundecl . args <EOL> fusion_fundecl = ast . FunDecl ( base_fundecl . ret , fusion_name , <EOL> fusion_args , body ) <EOL> binding = self . _filter_kernel_args ( [ base_loop , fuse_loop ] , fusion_fundecl ) <EOL> fusion_fundecl . args += [ ast . Decl ( '<STR_LIT>' , ast . Symbol ( '<STR_LIT>' ) ) ] <EOL> base_funcall_syms = [ ast . Symbol ( d . sym . symbol ) <EOL> for d in base_fundecl . args ] <EOL> base_funcall = ast . FunCall ( base_fundecl . name , * base_funcall_syms ) <EOL> fuse_funcall_syms = [ ast . Symbol ( binding [ arg ] . sym . symbol ) <EOL> for arg in fuse_loop . args ] <EOL> fuse_funcall = ast . FunCall ( fuse_fundecl . name , * fuse_funcall_syms ) <EOL> if_cond = ast . Not ( ast . Symbol ( '<STR_LIT>' , ( '<STR_LIT:i>' , <NUM_LIT:0> ) ) ) <EOL> if_update = ast . Assign ( ast . Symbol ( '<STR_LIT>' , ( '<STR_LIT:i>' , <NUM_LIT:0> ) ) , ast . Symbol ( '<STR_LIT:1>' ) ) <EOL> if_exec = ast . If ( if_cond , [ ast . Block ( [ fuse_funcall , if_update ] , <EOL> open_scope = True ) ] ) <EOL> fuse_body = ast . Block ( [ if_exec ] , open_scope = True ) <EOL> fuse_for = ast . c_for ( '<STR_LIT:i>' , fused_map . arity , fuse_body , pragma = None ) <EOL> body . children . extend ( [ base_funcall , fuse_for . children [ <NUM_LIT:0> ] ] ) <EOL> ofs_syms , ofs_decls , ofs_vals = [ ] , [ ] , [ ] <EOL> init = lambda v : '<STR_LIT>' % '<STR_LIT:U+002CU+0020>' . join ( [ str ( j ) for j in v ] ) <EOL> for i , fuse_args in enumerate ( zip ( fuse_loop . args , fuse_fundecl . args ) ) : <EOL> fuse_loop_arg , fuse_kernel_arg = fuse_args <EOL> sym_id = fuse_kernel_arg . sym . symbol <EOL> if fuse_loop_arg == fused_inc_arg : <EOL> fuse_inc_refs = fuse_symbol_refs [ sym_id ] <EOL> fuse_inc_refs = [ sym for sym , parent in fuse_inc_refs <EOL> if not isinstance ( parent , ast . Decl ) ] <EOL> fuse_kernel_arg . sym . rank = binding [ fused_inc_arg ] . sym . rank <EOL> for b in fused_inc_arg . _block_shape : <EOL> for rc in b : <EOL> _ofs_vals = [ [ <NUM_LIT:0> ] for j in range ( len ( rc ) ) ] <EOL> for j , ofs in enumerate ( rc ) : <EOL> ofs_sym_id = '<STR_LIT>' % ( i , j ) <EOL> ofs_syms . append ( ofs_sym_id ) <EOL> ofs_decls . append ( ast . Decl ( '<STR_LIT:int>' , ast . Symbol ( ofs_sym_id ) ) ) <EOL> _ofs_vals [ j ] . append ( ofs ) <EOL> for s in fuse_inc_refs : <EOL> s . offset = tuple ( ( <NUM_LIT:1> , o ) for o in ofs_syms ) <EOL> ofs_vals . extend ( [ init ( o ) for o in _ofs_vals ] ) <EOL> fuse_kernel_arg . pragma = set ( [ ast . INC ] ) <EOL> elif fuse_loop_arg . _is_indirect : <EOL> fuse_arity = fuse_loop_arg . map . arity <EOL> base_arity = fuse_arity * fused_map . arity <EOL> cdim = fuse_loop_arg . data . dataset . cdim <EOL> size = fuse_arity * cdim <EOL> if fuse_loop_arg . _flatten and cdim > <NUM_LIT:1> : <EOL> ofs_tmp = '<STR_LIT>' % fuse_kernel_arg . sym . symbol <EOL> ofs_tmp_sym = ast . Symbol ( ofs_tmp , ( size , ) ) <EOL> ofs_tmp_decl = ast . Decl ( '<STR_LIT>' % fuse_loop_arg . ctype , ofs_tmp_sym ) <EOL> _ofs_vals = [ [ base_arity * j + k for k in range ( fuse_arity ) ] <EOL> for j in range ( cdim ) ] <EOL> _ofs_vals = [ [ fuse_arity * j + k for k in flatten ( _ofs_vals ) ] <EOL> for j in range ( fused_map . arity ) ] <EOL> _ofs_vals = list ( flatten ( _ofs_vals ) ) <EOL> ofs_idx_sym = '<STR_LIT>' % i <EOL> body . children . insert ( <NUM_LIT:0> , ast . Decl ( <EOL> '<STR_LIT:int>' , ast . Symbol ( ofs_idx_sym , ( len ( _ofs_vals ) , ) ) , <EOL> ast . ArrayInit ( init ( _ofs_vals ) ) , [ '<STR_LIT>' , '<STR_LIT>' ] ) ) <EOL> ofs_idx_syms = [ ast . Symbol ( ofs_idx_sym , ( '<STR_LIT:i>' , ) , ( ( size , j ) , ) ) <EOL> for j in range ( size ) ] <EOL> ofs_assigns = [ ofs_tmp_decl ] <EOL> ofs_assigns += [ ast . Assign ( ast . Symbol ( ofs_tmp , ( j , ) ) , <EOL> ast . Symbol ( sym_id , ( k , ) ) ) <EOL> for j , k in enumerate ( ofs_idx_syms ) ] <EOL> fuse_funcall . children [ fuse_loop . args . index ( fuse_loop_arg ) ] = ast . Symbol ( ofs_tmp ) <EOL> else : <EOL> ofs_sym = ast . Symbol ( '<STR_LIT>' , ( len ( ofs_vals ) , '<STR_LIT:i>' ) ) <EOL> ofs_assigns = [ ast . Assign ( sym_id , ast . Sum ( sym_id , ofs_sym ) ) ] <EOL> ofs_vals . append ( init ( [ j * size for j in range ( fused_map . arity ) ] ) ) <EOL> if_exec . children [ <NUM_LIT:0> ] . children [ <NUM_LIT:0> : <NUM_LIT:0> ] = ofs_assigns <EOL> body . children . insert ( <NUM_LIT:0> , ast . Decl ( <EOL> '<STR_LIT:int>' , ast . Symbol ( '<STR_LIT>' , ( len ( ofs_vals ) , fused_map . arity ) ) , <EOL> ast . ArrayInit ( init ( ofs_vals ) ) , [ '<STR_LIT>' , '<STR_LIT>' ] ) ) <EOL> if_exec . children [ <NUM_LIT:0> ] . children [ <NUM_LIT:0> : <NUM_LIT:0> ] = [ ast . Decl ( '<STR_LIT:int>' , ast . Symbol ( s ) , ast . Symbol ( '<STR_LIT>' , ( i , '<STR_LIT:i>' ) ) ) <EOL> for i , s in enumerate ( ofs_syms ) ] <EOL> fuse_funcall . children . extend ( [ ast . Symbol ( s ) for s in ofs_syms ] ) <EOL> fuse_fundecl . args . extend ( ofs_decls ) <EOL> fused_headers = set ( [ str ( h ) for h in base_headers + fuse_headers ] ) <EOL> fused_ast = ast . Root ( [ ast . PreprocessNode ( h ) for h in fused_headers ] + <EOL> [ base_fundecl , fuse_fundecl , fusion_fundecl ] ) <EOL> kernels = [ base , fuse ] <EOL> loop_chain_index = ( self . _loop_chain . index ( base_loop ) , <EOL> self . _loop_chain . index ( fuse_loop ) ) <EOL> _fused . append ( ( Kernel ( kernels , fused_ast , loop_chain_index ) , fused_map ) ) <EOL> self . _schedule = HardFusionSchedule ( self . _schedule , _fused ) <EOL> self . _loop_chain = self . _schedule ( self . _loop_chain , only_hard = True ) <EOL> def _tile ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> backend_map = { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT>' : '<STR_LIT>' } <EOL> slope_backend = backend_map [ configuration [ '<STR_LIT>' ] ] <EOL> slope . set_exec_mode ( slope_backend ) <EOL> log_info ( "<STR_LIT>" % slope_backend ) <EOL> except KeyError : <EOL> warning ( "<STR_LIT>" % configuration [ '<STR_LIT>' ] ) <EOL> inspector = slope . Inspector ( ) <EOL> arguments = [ ] <EOL> insp_sets , insp_maps , insp_loops = set ( ) , { } , [ ] <EOL> for loop in self . _loop_chain : <EOL> slope_desc = set ( ) <EOL> insp_sets . add ( ( loop . it_space . name , loop . it_space . core_size ) ) <EOL> for a in loop . args : <EOL> maps = as_tuple ( a . map , Map ) <EOL> if not maps : <EOL> slope_desc . add ( ( '<STR_LIT>' , a . access . _mode ) ) <EOL> continue <EOL> for i , map in enumerate ( maps ) : <EOL> for j , m in enumerate ( map ) : <EOL> map_name = "<STR_LIT>" % ( m . name , i , j ) <EOL> insp_maps [ m . name ] = ( map_name , m . iterset . name , <EOL> m . toset . name , m . values ) <EOL> slope_desc . add ( ( map_name , a . access . _mode ) ) <EOL> insp_loops . append ( ( loop . kernel . name , loop . it_space . name , list ( slope_desc ) ) ) <EOL> arguments . extend ( [ inspector . add_sets ( insp_sets ) ] ) <EOL> arguments . extend ( [ inspector . add_maps ( insp_maps . values ( ) ) ] ) <EOL> inspector . add_loops ( insp_loops ) <EOL> arguments . extend ( [ inspector . set_external_dats ( ) ] ) <EOL> arguments . extend ( [ inspector . set_tile_size ( self . _tile_size ) ] ) <EOL> argtypes , argvalues = zip ( * arguments ) <EOL> src = inspector . generate_code ( ) <EOL> rettype = slope . Executor . meta [ '<STR_LIT>' ] <EOL> slope_dir = os . environ [ '<STR_LIT>' ] <EOL> compiler = coffee . plan . compiler . get ( '<STR_LIT:name>' ) <EOL> cppargs = slope . get_compile_opts ( compiler ) <EOL> cppargs += [ '<STR_LIT>' % ( slope_dir , slope . get_include_dir ( ) ) ] <EOL> ldargs = [ '<STR_LIT>' % ( slope_dir , slope . get_lib_dir ( ) ) , <EOL> '<STR_LIT>' % slope . get_lib_name ( ) ] <EOL> fun = compilation . load ( src , "<STR_LIT>" , "<STR_LIT>" , cppargs , ldargs , <EOL> argtypes , rettype , compiler ) <EOL> inspection = fun ( * argvalues ) <EOL> executor = slope . Executor ( inspector ) <EOL> self . _schedule = TilingSchedule ( self . _schedule , inspection , executor ) <EOL> @ property <EOL> def mode ( self ) : <EOL> return self . _mode <EOL> def fuse ( name , loop_chain , tile_size ) : <EOL> """<STR_LIT>""" <EOL> if len ( loop_chain ) in [ <NUM_LIT:0> , <NUM_LIT:1> ] : <EOL> return loop_chain <EOL> remainder = [ ] <EOL> synch_points = [ l for l in loop_chain if isinstance ( l , Mat . _Assembly ) ] <EOL> if synch_points : <EOL> if len ( synch_points ) > <NUM_LIT:1> : <EOL> warning ( "<STR_LIT>" ) <EOL> synch_point = loop_chain . index ( synch_points [ <NUM_LIT:0> ] ) <EOL> remainder , loop_chain = loop_chain [ synch_point : ] , loop_chain [ : synch_point ] <EOL> if len ( loop_chain ) in [ <NUM_LIT:0> , <NUM_LIT:1> ] : <EOL> return loop_chain + remainder <EOL> if any ( [ isinstance ( l , ParLoop ) for l in loop_chain ] ) or any ( [ l . _reduced_globals for l in loop_chain ] ) : <EOL> return loop_chain + remainder <EOL> if any ( [ not hasattr ( l . kernel , '<STR_LIT>' ) or not l . kernel . _ast for l in loop_chain ] ) : <EOL> return loop_chain + remainder <EOL> if any ( [ isinstance ( l . kernel . _ast , ast . FlatBlock ) for l in loop_chain ] ) : <EOL> return loop_chain + remainder <EOL> if any ( a . _is_mixed for a in flatten ( [ l . args for l in loop_chain ] ) ) : <EOL> return loop_chain + remainder <EOL> if any ( [ l . is_layered for l in loop_chain ] ) : <EOL> return loop_chain + remainder <EOL> mode = '<STR_LIT>' <EOL> if tile_size > <NUM_LIT:0> : <EOL> mode = '<STR_LIT>' <EOL> if slope is None : <EOL> warning ( "<STR_LIT>" ) <EOL> return loop_chain + remainder <EOL> try : <EOL> os . environ [ '<STR_LIT>' ] <EOL> except KeyError : <EOL> warning ( "<STR_LIT>" ) <EOL> warning ( "<STR_LIT>" ) <EOL> return loop_chain + remainder <EOL> inspector = Inspector ( name , loop_chain , tile_size ) <EOL> schedule = inspector . inspect ( mode ) <EOL> return schedule ( loop_chain ) + remainder <EOL> @ contextmanager <EOL> def loop_chain ( name , num_unroll = <NUM_LIT:1> , tile_size = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> from base import _trace <EOL> trace = _trace . _trace <EOL> stamp = trace [ - <NUM_LIT:1> : ] <EOL> yield <EOL> if num_unroll < <NUM_LIT:1> : <EOL> return <EOL> start_point = trace . index ( stamp [ <NUM_LIT:0> ] ) + <NUM_LIT:1> if stamp else <NUM_LIT:0> <EOL> extracted_loop_chain = trace [ start_point : ] <EOL> total_loop_chain = loop_chain . unrolled_loop_chain + extracted_loop_chain <EOL> if len ( total_loop_chain ) / len ( extracted_loop_chain ) == num_unroll : <EOL> start_point = trace . index ( total_loop_chain [ <NUM_LIT:0> ] ) <EOL> trace [ start_point : ] = fuse ( name , total_loop_chain , tile_size ) <EOL> loop_chain . unrolled_loop_chain = [ ] <EOL> else : <EOL> loop_chain . unrolled_loop_chain . extend ( extracted_loop_chain ) <EOL> loop_chain . unrolled_loop_chain = [ ] </s>
