src/gmxlib/nonbonded/nb_kernel_avx_256_double/make_nb_kernel_avx_256_double.py

   1 #!/usr/bin/python
   2
   3 import sys
   4 import os
   5 sys.path.append ( "../preprocessor" )
   6 from gmxpreprocess import gmxpreprocess
   7
   8 # "The happiest programs are programs that write other programs."
   9 #
  10 #
  11 # This script controls the generation of Gromacs nonbonded kernels.
  12 #
  13 # We no longer generate kernels on-the-fly, so this file is not run
  14 # during a Gromacs compile - only when we need to update the kernels (=rarely).
  15 #
  16 # To maximize performance, each combination of interactions in Gromacs
  17 # has a separate nonbonded kernel without conditionals in the code.
  18 # To avoid writing hundreds of different routines for each architecture,
  19 # we instead use a custom preprocessor so we can encode the conditionals
  20 # and expand for-loops (e.g, for water-water interactions)
  21 # from a general kernel template. While that file will contain quite a
  22 # few preprocessor directives, it is still an order of magnitude easier
  23 # to maintain than ~200 different kernels (not to mention it avoids bugs).
  24 #
  25 # To actually generate the kernels, this program iteratively calls the
  26 # preprocessor with different define settings corresponding to all
  27 # combinations of coulomb/van-der-Waals/geometry options.
  28 #
  29 # A main goal in the design was to make this new generator _general_. For
  30 # this reason we have used a lot of different fields to identify a particular
  31 # kernel and interaction. Basically, each kernel will have a name like
  32 #
  33 # nbkernel_ElecXX_VdwYY_GeomZZ_VF_QQ()
  34 #
  35 # Where XX/YY/ZZ/VF are strings to identify what the kernel computes.
  36 #
  37 # Elec/Vdw describe the type of interaction for electrostatics and van der Waals.
  38 # The geometry settings correspond e.g. to water-water or water-particle kernels,
  39 # and finally the VF setting is V,F,or VF depending on whether we calculate
  40 # only the potential, only the force, or both of them. The final string (QQ)
  41 # is the architecture/language/optimization of the kernel.
  42 #
  43 Arch       = 'avx_256_double'
  44
  45 # Explanation of the 'properties':
  46 #
  47 # It is cheap to compute r^2, and the kernels require various other functions of r for
  48 # different kinds of interaction. Depending on the needs of the kernel and the available
  49 # processor instructions, this will be done in different ways.
  50 #
  51 # 'rinv' means we need 1/r, which is calculated as 1/sqrt(r^2).
  52 # 'rinvsq' means we need 1/(r*r). This is calculated as rinv*rinv if we already did rinv, otherwise 1/r^2.
  53 # 'r' is similarly calculated as r^2*rinv when needed
  54 # 'table' means the interaction is tabulated, in which case we will calculate a table index before the interaction
  55 # 'shift' means the interaction will be modified by a constant to make it zero at the cutoff.
  56 # 'cutoff' means the interaction is set to 0.0 outside the cutoff
  57 #
  58
  59 FileHeader = \
  60 '/*\n' \
  61 ' * Note: this file was generated by the Gromacs '+Arch+' kernel generator.\n' \
  62 ' *\n' \
  63 ' *                This source code is part of\n' \
  64 ' *\n' \
  65 ' *                 G   R   O   M   A   C   S\n' \
  66 ' *\n' \
  67 ' * Copyright (c) 2001-2012, The GROMACS Development Team\n' \
  68 ' *\n' \
  69 ' * Gromacs is a library for molecular simulation and trajectory analysis,\n' \
  70 ' * written by Erik Lindahl, David van der Spoel, Berk Hess, and others - for\n' \
  71 ' * a full list of developers and information, check out http://www.gromacs.org\n' \
  72 ' *\n' \
  73 ' * This program is free software; you can redistribute it and/or modify it under\n' \
  74 ' * the terms of the GNU Lesser General Public License as published by the Free\n' \
  75 ' * Software Foundation; either version 2 of the License, or (at your option) any\n' \
  76 ' * later version.\n' \
  77 ' *\n' \
  78 ' * To help fund GROMACS development, we humbly ask that you cite\n' \
  79 ' * the papers people have written on it - you can find them on the website.\n' \
  80 ' */\n'
  81
  82 ###############################################
  83 # ELECTROSTATICS
  84 # Interactions and flags for them
  85 ###############################################
  86 ElectrostaticsList = {
  87     'None'                    : [],
  88     'Coulomb'                 : ['rinv','rinvsq'],
  89     'ReactionField'           : ['rinv','rinvsq'],
  90     'GeneralizedBorn'         : ['rinv','r'],
  91     'CubicSplineTable'        : ['rinv','r','table'],
  92     'Ewald'                   : ['rinv','rinvsq','r'],
  93 }
  94
  95
  96 ###############################################
  97 # VAN DER WAALS
  98 # Interactions and flags for them
  99 ###############################################
 100 VdwList = {
 101     'None'                    : [],
 102     'LennardJones'            : ['rinvsq'],
 103 #    'Buckingham'              : ['rinv','rinvsq','r'], # Disabled for AVX_256 to reduce number of kernels and simply the template
 104     'CubicSplineTable'        : ['rinv','r','table'],
 105 }
 106
 107
 108 ###############################################
 109 # MODIFIERS
 110 # Different ways to adjust/modify interactions to conserve energy
 111 ###############################################
 112 ModifierList = {
 113     'None'                    : [],
 114     'ExactCutoff'             : ['exactcutoff'],        # Zero the interaction outside the cutoff, used for reaction-field-zero
 115     'PotentialShift'          : ['shift','exactcutoff'],
 116     'PotentialSwitch'         : ['rinv','r','switch','exactcutoff']
 117 }
 118
 119
 120 ###############################################
 121 # GEOMETRY COMBINATIONS
 122 ###############################################
 123 GeometryNameList = [
 124     [ 'Particle' , 'Particle' ],
 125     [ 'Water3'   , 'Particle' ],
 126     [ 'Water3'   , 'Water3'   ],
 127     [ 'Water4'   , 'Particle' ],
 128     [ 'Water4'   , 'Water4'   ]
 129 ]
 130
 131
 132 ###############################################
 133 # POTENTIAL / FORCE
 134 ###############################################
 135 VFList = [
 136     'PotentialAndForce',
 137 # 'Potential',   # Not used yet
 138     'Force'
 139 ]
 140
 141
 142 ###############################################
 143 # GEOMETRY PROPERTIES
 144 ###############################################
 145 # Dictionaries with lists telling which interactions are present
 146 # 1,2,3 means particles 1,2,3 (but not 0) have electrostatics!
 147 GeometryElectrostatics = {
 148     'Particle'  : [ 0 ],
 149     'Particle2' : [ 0 , 1 ],
 150     'Particle3' : [ 0 , 1 , 2 ],
 151     'Particle4' : [ 0 , 1 , 2 , 3 ],
 152     'Water3'    : [ 0 , 1 , 2 ],
 153     'Water4'    : [ 1 , 2 , 3 ]
 154 }
 155
 156 GeometryVdw = {
 157     'Particle'  : [ 0 ],
 158     'Particle2' : [ 0 , 1 ],
 159     'Particle3' : [ 0 , 1 , 2 ],
 160     'Particle4' : [ 0 , 1 , 2 , 3 ],
 161     'Water3'    : [ 0 ],
 162     'Water4'    : [ 0 ]
 163 }
 164
 165
 166
 167
 168 # Dictionary to abbreviate all strings (mixed from all the lists)
 169 Abbreviation = {
 170     'None'                    : 'None',
 171     'Coulomb'                 : 'Coul',
 172     'Ewald'                   : 'Ew',
 173     'ReactionField'           : 'RF',
 174     'GeneralizedBorn'         : 'GB',
 175     'CubicSplineTable'        : 'CSTab',
 176     'LennardJones'            : 'LJ',
 177     'Buckingham'              : 'Bham',
 178     'PotentialShift'          : 'Sh',
 179     'PotentialSwitch'         : 'Sw',
 180     'ExactCutoff'             : 'Cut',
 181     'PotentialAndForce'       : 'VF',
 182     'Potential'               : 'V',
 183     'Force'                   : 'F',
 184     'Water3'                  : 'W3',
 185     'Water4'                  : 'W4',
 186     'Particle'                : 'P1',
 187     'Particle2'               : 'P2',
 188     'Particle3'               : 'P3',
 189     'Particle4'               : 'P4'
 190 }
 191
 192
 193 ###############################################
 194 # Functions
 195 ###############################################
 196
 197 # Return a string with the kernel name from current settings
 198 def MakeKernelFileName(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom):
 199     ElecStr = 'Elec' + Abbreviation[KernelElec]
 200     if(KernelElecMod!='None'):
 201         ElecStr = ElecStr + Abbreviation[KernelElecMod]
 202     VdwStr  = 'Vdw'  + Abbreviation[KernelVdw]
 203     if(KernelVdwMod!='None'):
 204         VdwStr = VdwStr + Abbreviation[KernelVdwMod]
 205     GeomStr = 'Geom' + Abbreviation[KernelGeom[0]] + Abbreviation[KernelGeom[1]]
 206     return 'nb_kernel_' + ElecStr + '_' + VdwStr + '_' + GeomStr + '_' + Arch
 207
 208 def MakeKernelName(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelVF):
 209     ElecStr = 'Elec' + Abbreviation[KernelElec]
 210     if(KernelElecMod!='None'):
 211         ElecStr = ElecStr + Abbreviation[KernelElecMod]
 212     VdwStr  = 'Vdw'  + Abbreviation[KernelVdw]
 213     if(KernelVdwMod!='None'):
 214         VdwStr = VdwStr + Abbreviation[KernelVdwMod]
 215     GeomStr = 'Geom' + Abbreviation[KernelGeom[0]] + Abbreviation[KernelGeom[1]]
 216     VFStr   = Abbreviation[KernelVF]
 217     return 'nb_kernel_' + ElecStr + '_' + VdwStr + '_' + GeomStr + '_' + VFStr + '_' + Arch
 218
 219 # Return a string with a declaration to use for the kernel;
 220 # this will be a sequence of string combinations as well as the actual function name
 221 # Dont worry about field widths - that is just pretty-printing for the header!
 222 def MakeKernelDecl(KernelName,KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelOther,KernelVF):
 223     KernelStr   = '\"'+KernelName+'\"'
 224     ArchStr     = '\"'+Arch+'\"'
 225     ElecStr     = '\"'+KernelElec+'\"'
 226     ElecModStr  = '\"'+KernelElecMod+'\"'
 227     VdwStr      = '\"'+KernelVdw+'\"'
 228     VdwModStr   = '\"'+KernelVdwMod+'\"'
 229     GeomStr     = '\"'+KernelGeom[0]+KernelGeom[1]+'\"'
 230     OtherStr    = '\"'+KernelOther+'\"'
 231     VFStr       = '\"'+KernelVF+'\"'
 232
 233     ThisSpec = ArchStr+', '+ElecStr+', '+ElecModStr+', '+VdwStr+', '+VdwModStr+', '+GeomStr+', '+OtherStr+', '+VFStr
 234     ThisDecl = '    { '+KernelName+', '+KernelStr+', '+ThisSpec+' }'
 235     return ThisDecl
 236
 237
 238 # Returns 1 if this kernel should be created, 0 if we should skip it
 239 # This routine is not critical - it is not the end of the world if we create more kernels,
 240 # but since the number is pretty large we save both space and compile-time by reducing it a bit.
 241 def KeepKernel(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelVF):
 242
 243     # No need for kernels without interactions
 244     if(KernelElec=='None' and KernelVdw=='None'):
 245         return 0
 246
 247     # No need for modifiers without interactions
 248     if((KernelElec=='None' and KernelElecMod!='None') or (KernelVdw=='None' and KernelVdwMod!='None')):
 249         return 0
 250
 251     # No need for LJ-only water optimization, or water optimization with implicit solvent.
 252     if('Water' in KernelGeom[0] and (KernelElec=='None' or 'GeneralizedBorn' in KernelElec)):
 253         return 0
 254
 255     # Non-matching table settings are pointless
 256     if( ('Table' in KernelElec) and ('Table' in KernelVdw) and KernelElec!=KernelVdw ):
 257         return 0
 258
 259     # Try to reduce the number of different switch/shift options to get a reasonable number of kernels
 260     # For electrostatics, reaction-field can use 'exactcutoff', and ewald can use switch or shift.
 261     if(KernelElecMod=='ExactCutoff' and KernelElec!='ReactionField'):
 262         return 0
 263     if(KernelElecMod in ['PotentialShift','PotentialSwitch'] and KernelElec!='Ewald'):
 264         return 0
 265     # For Vdw, we support switch and shift for Lennard-Jones/Buckingham
 266     if((KernelVdwMod=='ExactCutoff') or
 267        (KernelVdwMod in ['PotentialShift','PotentialSwitch'] and KernelVdw not in ['LennardJones','Buckingham'])):
 268         return 0
 269
 270     # Choose either switch or shift and don't mix them...
 271     if((KernelElecMod=='PotentialShift' and KernelVdwMod=='PotentialSwitch') or
 272        (KernelElecMod=='PotentialSwitch' and KernelVdwMod=='PotentialShift')):
 273         return 0
 274
 275     # Don't use a Vdw kernel with a modifier if the electrostatics one does not have one
 276     if(KernelElec!='None' and KernelElecMod=='None' and KernelVdwMod!='None'):
 277         return 0
 278
 279     # Don't use an electrostatics kernel with a modifier if the vdw one does not have one,
 280     # unless the electrostatics one is reaction-field with exact cutoff.
 281     if(KernelVdw!='None' and KernelVdwMod=='None' and KernelElecMod!='None'):
 282         if(KernelElec=='ReactionField' and KernelVdw!='CubicSplineTable'):
 283             return 0
 284         elif(KernelElec!='ReactionField'):
 285             return 0
 286
 287     return 1
 288
 289
 290
 291 #
 292 # The preprocessor will automatically expand the interactions for water and other
 293 # geometries inside the kernel, but to get this right we need to setup a couple
 294 # of defines - we do them in a separate routine to keep the main loop clean.
 295 #
 296 # While this routine might look a bit complex it is actually quite straightforward,
 297 # and the best news is that you wont have to modify _anything_ for a new geometry
 298 # as long as you correctly define its Electrostatics/Vdw geometry in the lists above!
 299 #
 300 def SetDefines(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelVF,defines):
 301     # What is the _name_ for the i/j group geometry?
 302     igeometry            = KernelGeom[0]
 303     jgeometry            = KernelGeom[1]
 304     # define so we can access it in the source when the preprocessor runs
 305     defines['GEOMETRY_I'] = igeometry
 306     defines['GEOMETRY_J'] = jgeometry
 307
 308     # For the i/j groups, extract a python list of which sites have electrostatics
 309     # For SPC/TIP3p this will be [1,1,1], while TIP4p (no elec on first site) will be [0,1,1,1]
 310     ielec                = GeometryElectrostatics[igeometry]
 311     jelec                = GeometryElectrostatics[jgeometry]
 312     # Zero out the corresponding lists in case we dont do Elec
 313     if(KernelElec=='None'):
 314         ielec = []
 315         jelec = []
 316
 317     # Extract similar interaction lists for Vdw interactions (example for SPC: [1,0,0])
 318     iVdw                 = GeometryVdw[igeometry]
 319     jVdw                 = GeometryVdw[jgeometry]
 320
 321     # Zero out the corresponding lists in case we dont do Vdw
 322     if(KernelVdw=='None'):
 323         iVdw = []
 324         jVdw = []
 325
 326     # iany[] and jany[] contains lists of the particles actually used (for interactions) in this kernel
 327     iany = list(set(ielec+iVdw))  # convert to+from set to make elements unique
 328     jany = list(set(jelec+jVdw))
 329
 330     defines['PARTICLES_ELEC_I'] = ielec
 331     defines['PARTICLES_ELEC_J'] = jelec
 332     defines['PARTICLES_VDW_I']  = iVdw
 333     defines['PARTICLES_VDW_J']  = jVdw
 334     defines['PARTICLES_I']      = iany
 335     defines['PARTICLES_J']      = jany
 336
 337     # elecij,Vdwij are sets with pairs of particles for which the corresponding interaction is done
 338     # (and anyij again corresponds to either electrostatics or Vdw)
 339     elecij = []
 340     Vdwij  = []
 341     anyij  = []
 342
 343     for i in ielec:
 344         for j in jelec:
 345             elecij.append([i,j])
 346
 347     for i in iVdw:
 348         for j in jVdw:
 349             Vdwij.append([i,j])
 350
 351     for i in iany:
 352         for j in jany:
 353             if [i,j] in elecij or [i,j] in Vdwij:
 354                 anyij.append([i,j])
 355
 356     defines['PAIRS_IJ']     = anyij
 357
 358     # Make an 2d list-of-distance-properties-to-calculate for i,j
 359     ni = max(iany)+1
 360     nj = max(jany)+1
 361     # Each element properties[i][j] is an empty list
 362     properties = [ [ [] for j in range(0,nj) ] for i in range (0,ni) ]
 363     # Add properties to each set
 364     for i in range(0,ni):
 365         for j in range(0,nj):
 366             if [i,j] in elecij:
 367                 properties[i][j] = properties[i][j] + ['electrostatics'] + ElectrostaticsList[KernelElec] + ModifierList[KernelElecMod]
 368             if [i,j] in Vdwij:
 369                 properties[i][j] = properties[i][j] + ['vdw'] + VdwList[KernelVdw] + ModifierList[KernelVdwMod]
 370             # Add rinv if we need r
 371             if 'r' in properties[i][j]:
 372                 properties[i][j] = properties[i][j] + ['rinv']
 373             # Add rsq if we need rinv or rinsq
 374             if 'rinv' in properties[i][j] or 'rinvsq' in properties[i][j]:
 375                 properties[i][j] = properties[i][j] + ['rsq']
 376
 377     defines['INTERACTION_FLAGS']    = properties
 378
 379
 380
 381 def PrintStatistics(ratio):
 382     ratio = 100.0*ratio
 383     print '\rGenerating %s nonbonded kernels... %5.1f%%' % (Arch,ratio),
 384     sys.stdout.flush()
 385
 386
 387
 388 defines = {}
 389 kerneldecl = []
 390
 391 cnt     = 0.0
 392 nelec   = len(ElectrostaticsList)
 393 nVdw    = len(VdwList)
 394 nmod    = len(ModifierList)
 395 ngeom   = len(GeometryNameList)
 396
 397 ntot    = nelec*nmod*nVdw*nmod*ngeom
 398
 399 numKernels = 0
 400
 401 fpdecl = open('nb_kernel_' + Arch + '.c','w')
 402 fpdecl.write( FileHeader )
 403 fpdecl.write( '#ifndef nb_kernel_' + Arch + '_h\n' )
 404 fpdecl.write( '#define nb_kernel_' + Arch + '_h\n\n' )
 405 fpdecl.write( '#include "../nb_kernel.h"\n\n' )
 406
 407 for KernelElec in ElectrostaticsList:
 408     defines['KERNEL_ELEC'] = KernelElec
 409
 410     for KernelElecMod in ModifierList:
 411         defines['KERNEL_MOD_ELEC'] = KernelElecMod
 412
 413         for KernelVdw in VdwList:
 414             defines['KERNEL_VDW'] = KernelVdw
 415
 416             for KernelVdwMod in ModifierList:
 417                 defines['KERNEL_MOD_VDW'] = KernelVdwMod
 418
 419                 for KernelGeom in GeometryNameList:
 420
 421                     cnt += 1
 422                     KernelFilename = MakeKernelFileName(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom) + '.c'
 423                     fpkernel = open(KernelFilename,'w')
 424                     defines['INCLUDE_HEADER'] = 1  # Include header first time in new file
 425                     DoHeader = 1
 426
 427                     for KernelVF in VFList:
 428
 429                         KernelName = MakeKernelName(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelVF)
 430
 431                         defines['KERNEL_NAME'] = KernelName
 432                         defines['KERNEL_VF']   = KernelVF
 433
 434                         # Check if this is a valid/sane/usable combination
 435                         if not KeepKernel(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelVF):
 436                             continue;
 437
 438                         # The overall kernel settings determine what the _kernel_ calculates, but for the water
 439                         # kernels this does not mean that every pairwise interaction has e.g. Vdw interactions.
 440                         # This routine sets defines of what to calculate for each pair of particles in those cases.
 441                         SetDefines(KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelVF,defines)
 442
 443                         if(DoHeader==1):
 444                             fpkernel.write( FileHeader )
 445
 446                         gmxpreprocess('nb_kernel_template_' + Arch + '.pre', KernelName+'.tmp' , defines, force=1,contentType='C')
 447                         numKernels = numKernels + 1
 448
 449                         defines['INCLUDE_HEADER'] = 0   # Header has been included once now
 450                         DoHeader=0
 451
 452                         # Append temp file contents to the common kernelfile
 453                         fptmp = open(KernelName+'.tmp','r')
 454                         fpkernel.writelines(fptmp.readlines())
 455                         fptmp.close()
 456                         os.remove(KernelName+'.tmp')
 457
 458                         # Add a declaration for this kernel
 459                         fpdecl.write('nb_kernel_t ' + KernelName + ';\n');
 460
 461                         # Add declaration to the buffer
 462                         KernelOther=''
 463                         kerneldecl.append(MakeKernelDecl(KernelName,KernelElec,KernelElecMod,KernelVdw,KernelVdwMod,KernelGeom,KernelOther,KernelVF))
 464
 465                     filesize = fpkernel.tell()
 466                     fpkernel.close()
 467                     if(filesize==0):
 468                         os.remove(KernelFilename)
 469
 470                     PrintStatistics(cnt/ntot)
 471                 pass
 472             pass
 473         pass
 474     pass
 475 pass
 476
 477 # Write out the list of settings and corresponding kernels to the declaration file
 478 fpdecl.write( '\n\n' )
 479 fpdecl.write( 'nb_kernel_info_t\n' )
 480 fpdecl.write( 'kernellist_'+Arch+'[] =\n' )
 481 fpdecl.write( '{\n' )
 482 for decl in kerneldecl[0:-1]:
 483     fpdecl.write( decl + ',\n' )
 484 fpdecl.write( kerneldecl[-1] + '\n' )
 485 fpdecl.write( '};\n\n' )
 486 fpdecl.write( 'int\n' )
 487 fpdecl.write( 'kernellist_'+Arch+'_size = sizeof(kernellist_'+Arch+')/sizeof(kernellist_'+Arch+'[0]);\n\n')
 488 fpdecl.write( '#endif\n')
 489 fpdecl.close()