#NEWLINE# def __init__( self , lib , dtype , N , C , K , D , H , W , T , R , S , M , P , Q , pad_d , pad_h , pad_w , str_d , str_h , str_w ) : #NEWLINE# #INDENT# super ( UpdateDirect , self ) . __init__ ( lib , dtype ) #NEWLINE# assert ( ( N % 32 ) == 0 ) , 'N dim must be multiple of 32' #NEWLINE# magic_RST = _magic32 ( ( ( ( C * R ) * S ) * T ) , ( ( R * S ) * T ) ) #NEWLINE# magic_RS = _magic32 ( ( ( ( R * S ) * T ) + 32 ) , ( R * S ) ) #NEWLINE# magic_S = _magic32 ( ( ( R * S ) + 32 ) , S ) #NEWLINE# grid_C = _ceil_div ( ( ( ( C * R ) * S ) * T ) , 128 ) #NEWLINE# sm_count = _get_sm_count ( ) #NEWLINE# if ( ( dtype . type is np . float32 ) and ( ( P * Q ) > ( 56 * 56 ) ) ) : #NEWLINE# #INDENT# K_tiles = ( 64 , ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# K_tiles = ( 128 , 64 ) #NEWLINE# #UNINDENT# if lib . deterministic : #NEWLINE# #INDENT# determ = 'D' #NEWLINE# if ( K <= 64 ) : #NEWLINE# #INDENT# K_tiles = ( 64 , ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# K_tiles = K_tiles [ 0 : 1 ] #NEWLINE# #UNINDENT# self . determ = ( ( ( ( C * T ) * R ) * S ) * K ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# determ = '' #NEWLINE# self . determ = 0 #NEWLINE# #UNINDENT# self . kernels = [ ] #NEWLINE# for ( tile_K , grid_K , offset_K ) in self . k_partitions ( K , K_tiles ) : #NEWLINE# #INDENT# kernel_name = ( '%s_updat%s_C128_K%d' % ( self . clss , determ , tile_K ) ) #NEWLINE# base_blocks = ( ( M * grid_C ) * grid_K ) #NEWLINE# ( grid_P , grid_Q , threads ) = self . update_grid ( kernel_name , base_blocks , P , Q , sm_count ) #NEWLINE# grid_PQ = ( grid_P * grid_Q ) #NEWLINE# magic_PQu = _magic64 ( grid_PQ ) #NEWLINE# magic_Qu = _magic64 ( grid_Q ) #NEWLINE# block = ( threads , 1 , 1 ) #NEWLINE# if ( ( ( R * S ) * T ) > 1 ) : #NEWLINE# #INDENT# grid = ( ( M * grid_PQ ) , grid_C , grid_K ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# grid = ( grid_C , grid_K , ( M * grid_PQ ) ) #NEWLINE# #UNINDENT# self . determ *= ( M * grid_PQ ) #NEWLINE# self . determ_shape = ( ( M * grid_PQ ) , ( ( ( ( C * T ) * R ) * S ) * K ) ) #NEWLINE# kernel = [ kernel_name , grid , block , None , None , None , None , None ] #NEWLINE# kernel . extend ( _flatten ( [ offset_K , N , K , D , H , W , ( W * N ) , ( ( H * W ) * N ) , ( ( ( D * H ) * W ) * N ) , C , ( ( ( C * R ) * S ) * T ) , ( ( R * S ) * T ) , magic_RST , ( R * S ) , magic_RS , S , magic_S , pad_d , pad_h , pad_w , str_d , str_h , str_w , P , Q , ( P * Q ) , ( Q * N ) , ( ( P * Q ) * N ) , ( ( ( M * P ) * Q ) * N ) , magic_Qu , magic_PQu , grid_P , grid_Q , grid_PQ , ( ( ( ( C * R ) * S ) * T ) * K ) ] ) ) #NEWLINE# self . kernels . append ( kernel ) #NEWLINE# #UNINDENT# lib . set_scratch_size ( ( ( self . determ or ( ( ( ( C * T ) * R ) * S ) * K ) ) * 4 ) )