#NEWLINE# @ context_dependent_memoize #NEWLINE# def _get_compensated_sum_kernel( dtype , rounding ) : #NEWLINE# #INDENT# _compensated_sum = '\n\n%(common)s\n\n__global__ void compensated_sum(unsigned* rand_state,\n          %(type)s* a_sum,\n          %(type)s* a_cmp,\n    const %(type)s* a_add,\n    float cmp_scale, float add_scale,\n    int row_strd, int col_strd, int n, int mantissa_bits)\n{\n    const int tid = threadIdx.x;\n    const int bid = blockIdx.x;\n\n    int offset = bid * row_strd + tid * col_strd;\n    int inc    = 32 * col_strd;\n\n    a_sum += offset;\n    a_cmp += offset;\n    a_add += offset;\n\n    %(inits)s\n\n    for (int i = tid; i < n; i += 32)\n    {\n        float s32 = %(cvt)s(__ldg((const %(type)s*)a_sum));\n        float c32 = %(cvt)s(__ldg((const %(type)s*)a_cmp));\n        float a32 = %(cvt)s(__ldg(a_add));\n\n        // Adjust amount to add by previous compensation\n        float y32 = a32 * add_scale - c32 * cmp_scale;\n\n        // Do the accumulation and truncate to the storage type\n        float rnd_sum = s32 + y32;\n        %(rnd_sum)s\n\n        // Convert accumulation back to fp32 so we can do more math on it\n        float t32 = %(cvt)s(t16);\n\n        // recover the low order bits that were lost in the truncation\n        float rnd_cmp = (t32 - s32) - y32;\n        %(rnd_cmp)s\n\n        *a_sum = t16;\n        *a_cmp = c16;\n\n        a_sum += inc;\n        a_cmp += inc;\n        a_add += inc;\n    }\n    %(finish)s\n}\n' #NEWLINE# template_vals = dict ( ) #NEWLINE# for key in ( 'common' , 'inits' , 'finish' ) : #NEWLINE# #INDENT# template_vals [ key ] = '' #NEWLINE# #UNINDENT# if ( dtype == 'f2' ) : #NEWLINE# #INDENT# template_vals [ 'common' ] += _common_fp16_to_fp32 #NEWLINE# #UNINDENT# if rounding : #NEWLINE# #INDENT# template_vals [ 'common' ] += _common_urand_gen #NEWLINE# template_vals [ 'common' ] += _common_round [ 'nearest' ] . get ( dtype , '' ) #NEWLINE# template_vals [ 'inits' ] += ( _init_rand_func + _init_rand_round_func ) #NEWLINE# template_vals [ 'finish' ] += _finish_rand_func #NEWLINE# mode = 'random' #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# mode = 'nearest' #NEWLINE# #UNINDENT# template_vals [ 'common' ] += _common_round [ mode ] . get ( dtype , '' ) #NEWLINE# template_vals [ 'type' ] = _ew_types [ dtype ] [ 'type' ] #NEWLINE# template_vals [ 'cvt' ] = _ew_types [ dtype ] [ 'cvt' ] #NEWLINE# no_op = 'float {0} = {1};' #NEWLINE# rnd_sum = _ew_strings [ 'round' ] [ mode ] . get ( dtype , no_op ) #NEWLINE# rnd_cmp = _ew_strings [ 'round' ] [ 'nearest' ] . get ( dtype , no_op ) #NEWLINE# template_vals [ 'rnd_sum' ] = rnd_sum . format ( 't16' , 'rnd_sum' ) #NEWLINE# template_vals [ 'rnd_cmp' ] = rnd_cmp . format ( 'c16' , 'rnd_cmp' ) #NEWLINE# code = ( _compensated_sum % template_vals ) #NEWLINE# module = SourceModule ( code ) #NEWLINE# kernel = module . get_function ( 'compensated_sum' ) #NEWLINE# kernel . prepare ( 'PPPPffiiii' ) #NEWLINE# return kernel