def write_straprim_caller(myfile,
index,
a_coeffs,
b_coeffs,
c_coeffs,
dims,
num_multiplies,
level=1):
comment = '// M%d = (' % (index)
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 0, i, dims, level ) \
for i, c in enumerate(a_coeffs) if is_nonzero(c)])
comment += ') * ('
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 1, i, dims, level ) \
for i, c in enumerate(b_coeffs) if is_nonzero(c)])
comment += '); '
comment += '; '.join([
' %s += %s * M%d' % (getBlockName(2, i, dims, level), c, index)
for i, c in enumerate(c_coeffs) if is_nonzero(c)
])
comment += ';'
write_line(myfile, 1, comment)
write_stra_mat(myfile, 0, a_coeffs, index, 'ms, ks', dims, level)
write_stra_mat(myfile, 1, b_coeffs, index, 'ks, ns', dims, level)
write_stra_mat(myfile, 2, c_coeffs, index, 'ms, ns', dims, level)
#add = 'stra_gemm(comm, cfg, alpha, Av{0}, Bv{0}, beta, Cv{0});'.format( index )
#add = 'straprim_naive(comm, cfg, alpha, Av{0}, Bv{0}, beta, Cv{0});'.format( index )
add = 'straprim_ab(comm, cfg, alpha, Av{0}, Bv{0}, beta, Cv{0});'.format(
index)
write_line(myfile, 1, add)
write_line(myfile, 1, 'comm.barrier();')
write_break(myfile)
def write_straprim_caller(myfile,
index,
a_coeffs,
b_coeffs,
c_coeffs,
dims,
num_multiplies,
level=1):
comment = '// M%d = (' % (index)
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 0, i, dims, level ) \
for i, c in enumerate(a_coeffs) if is_nonzero(c)])
comment += ') * ('
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 1, i, dims, level ) \
for i, c in enumerate(b_coeffs) if is_nonzero(c)])
comment += '); '
comment += '; '.join([
' %s += %s * M%d' % (getBlockName(2, i, dims, level), c, index)
for i, c in enumerate(c_coeffs) if is_nonzero(c)
])
comment += ';'
write_line(myfile, 1, comment)
write_stra_mat(myfile, 0, a_coeffs, index, ['AC', 'AB'], dims, level)
write_stra_mat(myfile, 1, b_coeffs, index, ['AB', 'BC'], dims, level)
write_stra_mat(myfile, 2, c_coeffs, index, ['AC', 'BC'], dims, level)
myfile.write( \
'''\
if (Cv{0}.stride(!row_major) == 1)
{{
Av{0}.transpose();
Bv{0}.transpose();
Cv{0}.transpose();
stra_gemm(comm, cfg, alpha, Bv{0}, Av{0}, beta, Cv{0});
}} else {{
stra_gemm(comm, cfg, alpha, Av{0}, Bv{0}, beta, Cv{0});
}}
'''.format( index ) )
#Av{0}.swap(Bv{0});
#add = 'stra_gemm(comm, cfg, alpha, Av{0}, Bv{0}, beta, Cv{0});'.format( index )
#write_line( myfile, 1, add )
write_line(myfile, 1, 'comm.barrier();')
write_line(
myfile, 1,
'//std::cout << "stra_internal/stra_mult_M{0}:" << std::endl;'.format(
index))
write_line(myfile, 1, '//print_tensor_matrix( ct );')
write_break(myfile)
def write_straprim_caller(myfile, index, a_coeffs, b_coeffs, c_coeffs, dims, num_multiplies, level=1):
comment = '// M%d = (' % (index)
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 0, i, dims, level ) \
for i, c in enumerate(a_coeffs) if is_nonzero(c)])
comment += ') * ('
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 1, i, dims, level ) \
for i, c in enumerate(b_coeffs) if is_nonzero(c)])
comment += '); '
comment += '; '.join([' %s += %s * M%d' % ( getBlockName( 2, i, dims, level ), c, index ) for i, c in enumerate(c_coeffs) if is_nonzero(c)])
comment += ';'
write_line(myfile, 1, comment)
write_stra_mat( myfile, 0, a_coeffs, index, ['AC', 'AB'], dims, level )
write_stra_mat( myfile, 1, b_coeffs, index, ['AB', 'BC'], dims, level )
write_stra_mat( myfile, 2, c_coeffs, index, ['AC', 'BC'], dims, level )
myfile.write( \
'''\
//if (ct.stride(!row_major) == 1)
//{{
// Av{0}.transpose();
// Bv{0}.transpose();
// Cv{0}.transpose();
// straprim_naive<T,{1},{2},{3}>(comm, cfg, my_sub_len_AB, my_sub_len_AC, my_sub_len_BC,
// alpha,
// B{0}_list, B{0}_coeff_list, my_stride_B_AB, my_stride_B_BC,
// A{0}_list, A{0}_coeff_list, my_stride_A_AB, my_stride_A_AC,
// beta,
// C{0}_list, C{0}_coeff_list, my_stride_C_AC, my_stride_C_BC);
//}} else {{
straprim_naive<T,{1},{2},{3}>(comm, cfg, my_sub_len_AB, my_sub_len_AC, my_sub_len_BC,
alpha,
A{0}_list, A{0}_coeff_list, my_stride_A_AB, my_stride_A_AC,
B{0}_list, B{0}_coeff_list, my_stride_B_AB, my_stride_B_BC,
beta,
C{0}_list, C{0}_coeff_list, my_stride_C_AC, my_stride_C_BC);
//}}
'''.format( index, getNNZ(a_coeffs), getNNZ(b_coeffs), getNNZ(c_coeffs) ) )
#Av{0}.swap(Bv{0});
#add = 'stra_gemm(comm, cfg, alpha, Av{0}, Bv{0}, beta, Cv{0});'.format( index )
#write_line( myfile, 1, add )
write_line( myfile, 1, 'comm.barrier();' )
write_line( myfile, 1, '//std::cout << "stra_internal/stra_mult_M{0}:" << std::endl;'.format( index ) )
write_line( myfile, 1, '//print_tensor_matrix( ct );' )
write_break( myfile )
def write_stra_matC(myfile, coeff_idx, coeffs, idx, dim_name, dims, level):
mat_name = (getName(coeff_idx))[0]
nonzero_coeffs = [coeff for coeff in coeffs if is_nonzero(coeff)]
#add = 'std::array<unsigned, ' + str(len(nonzero_coeffs)) + '> ' + mat_name + str(idx) + '_subid = {'
#add += ', '.join(['%s' % getActualBlockIndex( coeff_idx, i, dims, level ) \
# for i, c in enumerate(coeffs) if is_nonzero(c)])
#add += '};'
#write_line(myfile, 1, add)
add = 'std::array<T*,' + str(
len(nonzero_coeffs)) + '> ' + mat_name + str(idx) + '_list = {'
#add += ', '.join( [ str(c) for i, c in enumerate(coeffs) if is_nonzero(c) ] )
add += ', '.join([
'const_cast<T*>(%s)' % (getBlockName(coeff_idx, i, dims, level))
for i, c in enumerate(coeffs) if is_nonzero(c)
])
add += '};'
write_line(myfile, 1, add)
add = 'std::array<T,' + str(
len(nonzero_coeffs)) + '> ' + mat_name + str(idx) + '_coeff_list = {'
add += ', '.join([str(c) for i, c in enumerate(coeffs) if is_nonzero(c)])
add += '};'
write_line(myfile, 1, add)
def write_straprim_caller(myfile,
index,
a_coeffs,
b_coeffs,
c_coeffs,
dims,
num_multiplies,
level=1):
comment = '// M%d = (' % (index)
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 0, i, dims, level ) \
for i, c in enumerate(a_coeffs) if is_nonzero(c)])
comment += ') * ('
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 1, i, dims, level ) \
for i, c in enumerate(b_coeffs) if is_nonzero(c)])
comment += '); '
comment += '; '.join([
' %s += %s * M%d' % (getBlockName(2, i, dims, level), c, index)
for i, c in enumerate(c_coeffs) if is_nonzero(c)
])
comment += ';'
write_line(myfile, 1, comment)
add = 'bl_dgemm_straprim_abc%d( ms, ns, ks, ' % index
add += ', '.join(['%s' % getBlockName( 0, i, dims, level ) \
for i, c in enumerate(a_coeffs) if is_nonzero(c)])
add += ', lda, '
add += ', '.join(['%s' % getBlockName( 1, i, dims, level ) \
for i, c in enumerate(b_coeffs) if is_nonzero(c)])
add += ', ldb, '
add += ', '.join(['%s' % getBlockName( 2, i, dims, level ) \
for i, c in enumerate(c_coeffs) if is_nonzero(c)])
add += ', ldc, packA, packB, bl_ic_nt );'
write_line(myfile, 1, add)
def write_stra_mat(myfile, coeff_idx, coeffs, idx, dim_name, dims, level):
mat_name = (getName(coeff_idx))[0]
nonzero_coeffs = [coeff for coeff in coeffs if is_nonzero(coeff)]
add = 'stra_matrix_view<T,' + str(len(nonzero_coeffs)) + '> '
add += mat_name + 'v' + str(idx) + '({' + dim_name + '}, {'
#add += ', '.join( ['const_cast<T*>(%s)' % (getSubMatName(coeff_idx, i, dims, level) ) for i, c in enumerate(coeffs) if is_nonzero(c)] )
add += ', '.join([
'const_cast<T*>(%s)' % (getBlockName(coeff_idx, i, dims, level))
for i, c in enumerate(coeffs) if is_nonzero(c)
])
add += '}, {'
add += ', '.join([str(c) for i, c in enumerate(coeffs) if is_nonzero(c)])
add += '}, {rs_' + mat_name + ', cs_' + mat_name + '});'
write_line(myfile, 1, add)
def write_straprim_abc_function(myfile, index, a_coeffs, b_coeffs, c_coeffs,
dims, level):
comment = '// M%d = (' % (index)
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 0, i, dims, level ) \
for i, c in enumerate(a_coeffs) if is_nonzero(c)])
comment += ') * ('
comment += ' + '.join([str(c) + ' * %s' % getBlockName( 1, i, dims, level ) \
for i, c in enumerate(b_coeffs) if is_nonzero(c)])
comment += '); '
comment += '; '.join([
' %s += %s * M%d' % (getBlockName(2, i, dims, level), c, index)
for i, c in enumerate(c_coeffs) if is_nonzero(c)
])
comment += ';'
write_line(myfile, 0, comment)
add = 'void bl_dgemm_straprim_abc%d( int m, int n, int k, ' % index
add += ', '.join(
['double* %s%d' % ('a', i) for i in range(getNNZ(a_coeffs))])
add += ', int lda, '
add += ', '.join(
['double* %s%d' % ('b', i) for i in range(getNNZ(b_coeffs))])
add += ', int ldb, '
add += ', '.join(
['double* %s%d' % ('c', i) for i in range(getNNZ(c_coeffs))])
add += ', int ldc, double *packA, double *packB, int bl_ic_nt ) {'
write_line(myfile, 0, add)
write_line(myfile, 1, 'int i, j, p, ic, ib, jc, jb, pc, pb;')
write_line(myfile, 1, 'for ( jc = 0; jc < n; jc += DGEMM_NC ) {')
write_line(myfile, 2, 'jb = min( n - jc, DGEMM_NC );')
write_line(myfile, 2, 'for ( pc = 0; pc < k; pc += DGEMM_KC ) {')
write_line(myfile, 3, 'pb = min( k - pc, DGEMM_KC );')
#write_line( myfile, 0, '#ifdef _PARALLEL_')
#write_line( myfile, 3, '#pragma omp parallel for num_threads( bl_ic_nt ) private( j )' )
#write_line( myfile, 0, '#endif')
write_line(myfile, 3, '{')
write_line(myfile, 4, 'int tid = omp_get_thread_num();')
write_line(myfile, 4, 'int my_start;')
write_line(myfile, 4, 'int my_end;')
write_line(myfile, 4, 'bl_get_range( jb, DGEMM_NR, &my_start, &my_end );')
write_line(myfile, 4, 'for ( j = my_start; j < my_end; j += DGEMM_NR ) {')
add = 'packB_add_stra_abc%d( min( jb - j, DGEMM_NR ), pb, ' % index
add += ', '.join([
'&%s%d[ pc + (jc+j)*ldb ]' % ('b', i) for i in range(getNNZ(b_coeffs))
])
add += ', ldb, &packB[ j * pb ] );'
write_line(myfile, 5, add)
write_line(myfile, 4, '}')
write_line(myfile, 3, '}')
write_line(myfile, 0, '#ifdef _PARALLEL_')
write_line(myfile, 0, '#pragma omp barrier')
write_line(myfile, 0, '#endif')
#write_line( myfile, 0, '#ifdef _PARALLEL_')
#write_line( myfile, 3, '#pragma omp parallel num_threads( bl_ic_nt ) private( ic, ib, i )' )
#write_line( myfile, 0, '#endif')
write_line(myfile, 3, '{')
#write_line( myfile, 0, '#ifdef _PARALLEL_')
write_line(myfile, 4, 'int tid = omp_get_thread_num();')
write_line(myfile, 4, 'int my_start;')
write_line(myfile, 4, 'int my_end;')
write_line(myfile, 4, 'bl_get_range( m, DGEMM_MR, &my_start, &my_end );')
#write_line( myfile, 0, '#else')
#write_line( myfile, 4, 'int tid = 0;' )
#write_line( myfile, 4, 'int my_start = 0;' )
#write_line( myfile, 4, 'int my_end = m;' )
#write_line( myfile, 0, '#endif')
write_line(myfile, 4,
'for ( ic = my_start; ic < my_end; ic += DGEMM_MC ) {')
write_line(myfile, 5, 'ib = min( my_end - ic, DGEMM_MC );')
write_line(myfile, 5, 'for ( i = 0; i < ib; i += DGEMM_MR ) {')
add = 'packA_add_stra_abc%d( min( ib - i, DGEMM_MR ), pb, ' % index
add += ', '.join([
'&%s%d[ pc*lda + (ic+i) ]' % ('a', i) for i in range(getNNZ(a_coeffs))
])
add += ', lda, &packA[ tid * DGEMM_MC * pb + i * pb ] );'
write_line(myfile, 6, add)
write_line(myfile, 5, '}')
add = 'bl_macro_kernel_stra_abc%d( ib, jb, pb, packA + tid * DGEMM_MC * pb, packB, ' % index
add += ', '.join(
['&%s%d[ jc * ldc + ic ]' % ('c', i) for i in range(getNNZ(c_coeffs))])
add += ', ldc );'
write_line(myfile, 5, add)
write_line(myfile, 4, '}')
write_line(myfile, 3, '}')
write_line(myfile, 0, '#ifdef _PARALLEL_')
write_line(myfile, 0, '#pragma omp barrier')
write_line(myfile, 0, '#endif')
write_line(myfile, 2, '}')
write_line(myfile, 1, '}')
write_line(myfile, 0, '#ifdef _PARALLEL_')
write_line(myfile, 0, '#pragma omp barrier')
write_line(myfile, 0, '#endif')
write_line(myfile, 0, '}')
write_break(myfile)