CoreNeuron/coreneuron/kinderiv.py at e291683db9ff656badd9929d5fe86f0fbc33ad8b · BlueBrain/CoreNeuron · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
#!/usr/bin/python

# =============================================================================
# Copyright (c) 2016 - 2021 Blue Brain Project/EPFL
#
# See top-level LICENSE file for details.
# =============================================================================

# ----------------------------------------------------------------------------
# read the translated mod files and construct _kinderiv.h
# this is implemented to avoid function callbacks using
# function pointers while using methods like euler, newton
# or kinetic schemes. This is required for GPU implementation.
#
# if _kinderiv.h already exists and is the same as what this script generates
# then do not update it (to avoid re-compiling all the mech c files if
# only one of them changes)
# ----------------------------------------------------------------------------
from __future__ import print_function
import os
import sys
import filecmp

kftmp = '_tmp_kinderiv.h'
kf = '_kinderiv.h'

HEADER_GUARDS = '''
#ifndef _kinderiv_h
#define _kinderiv_h
'''

euler = []
deriv = []
kin = []


def process_modc(f):
    for line in f:
        word = line.split()
        if len(word) > 3 and word[0] == '/*':
            if word[1] == '_derivimplicit_':
                deriv.append([word[2], word[3], fname, word[1]])
            elif word[1] == '_kinetic_':
                kin.append([word[2], word[3], fname, word[1]])
            elif word[1] == '_euler_':
                euler.append([word[2], word[3], fname, word[1]])


def write_out_kinderiv(fout):
    fout.write(HEADER_GUARDS)
    fout.write("\n/* data used to construct this file */\n")
    for l in [deriv, kin,  euler]:
        for item in l:
            fout.write('/*')
            for word in item:
                fout.write(' %s' % word)
            fout.write(' */\n')

    fout.write("\n/* declarations */\n")
    fout.write("\nnamespace coreneuron {\n")

    if deriv or kin or euler:
        fout.write('nrn_pragma_omp(declare target)\n')

    for item in deriv:
        fout.write('#pragma acc routine seq\n')
        fout.write('extern int %s%s(_threadargsproto_);\n' % (item[0], item[1]))
        fout.write('#pragma acc routine seq\n')
        fout.write('extern int _newton_%s%s(_threadargsproto_);\n' % (item[0], item[1]))

    for item in kin:
        fout.write('#pragma acc routine seq\n')
        fout.write('extern int %s%s(void*, double*, _threadargsproto_);\n' % (item[0], item[1]))

    for item in euler:
        fout.write('#pragma acc routine seq\n')
        fout.write('extern int %s%s(_threadargsproto_);\n' % (item[0], item[1]))

    if deriv or kin or euler:
        fout.write('nrn_pragma_omp(end declare target)\n')

    fout.write("\n/* callback indices */\n")
    derivoffset = 1
    kinoffset = 1

    for i, item in enumerate(deriv):
        fout.write('#define _derivimplicit_%s%s %d\n' % (item[0], item[1], i + derivoffset))
    for i, item in enumerate(kin):
        fout.write('#define _kinetic_%s%s %d\n' % (item[0], item[1], i + kinoffset))
    for i, item in enumerate(euler):
        fout.write('#define _euler_%s%s %d\n' % (item[0], item[1], i + derivoffset))

    fout.write("\n/* switch cases */\n")
    fout.write("\n#define _NRN_DERIVIMPLICIT_CASES \\\n")
    for item in (deriv):
        fout.write("  case _derivimplicit_%s%s: %s%s(_threadargs_); break; \\\n" % (
            item[0], item[1], item[0], item[1]))
    fout.write("\n")

    fout.write("\n#define _NRN_DERIVIMPLICIT_NEWTON_CASES \\\n")
    for item in deriv:
        fout.write("  case _derivimplicit_%s%s: _newton_%s%s(_threadargs_); break; \\\n" % (
            item[0], item[1], item[0], item[1]))
    fout.write("\n")

    fout.write("\n#define _NRN_KINETIC_CASES \\\n")
    for item in kin:
        fout.write("  case _kinetic_%s%s: %s%s(so, rhs, _threadargs_); break; \\\n" % (
            item[0], item[1], item[0], item[1]))
    fout.write("\n")

    fout.write("\n#define _NRN_EULER_CASES \\\n")
    for item in euler:
        fout.write("  case _euler_%s%s: %s%s(_threadargs_); break; \\\n" % (
            item[0], item[1], item[0], item[1]))
    fout.write("\n")

    fout.write('\n#endif\n')
    fout.write("\n} //namespace coreneuron\n")


if __name__ == '__main__':
    use_dir = "."
    if len(sys.argv) == 2:
        use_dir = sys.argv[1]
        assert os.path.isdir(use_dir), "Dir {} doesnt exist".format(use_dir)

    fnames = sorted(f for f in os.listdir(use_dir) if f.endswith(".cpp"))

    for fname in fnames:
        with open(os.path.join(use_dir, fname), "r") as f:
            process_modc(f)

    with open(kftmp, "w") as fout:
        write_out_kinderiv(fout)

    # if kf exists and is same as kftmp, just remove kftmp. Otherwise rename kftmp to kf
    if os.path.isfile(kf) and filecmp.cmp(kftmp, kf):
        os.remove(kftmp)
    else:
        os.rename(kftmp, kf)