Chemgen/generate_chemistry.py at main · harikrishna1410/Chemgen · 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
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
import argparse
import os
from chemgen.parser import ckparser
from chemgen.chemistry import chemistry, chemistry_expressions
from chemgen.write_constants import write_coef_module, write_constants_header

def parse_arguments():
    parser = argparse.ArgumentParser(description='Generate GPU coefficients and getrates functions')

    # Main operation choice
    parser.add_argument('--mode', choices=['gpu_coef', 'getrates', "mini_app"], required=True,
                      help='Mode of operation: generate GPU coefficients or getrates function')

    # Input files
    parser.add_argument('--mech', help='Path to mechanism input file')
    parser.add_argument('--therm', help='Path to thermodynamic data file')
    parser.add_argument('--yaml_file', help='Path to YAML mechanism file')
    parser.add_argument('--output', required=True, help='Output file or directory')

    # GPU coefficient options
    parser.add_argument('--parallel-level', type=int, choices=[1,2,3,4], default=1,
                      help='Parallelization level for GPU coefficients')
    parser.add_argument('--veclen', type=int, default=16,
                      help='Vector length for GPU coefficients')
    parser.add_argument('--nreact-per-block', type=int, default=8,
                      help='Number of reactions per block')

    # Getrates options
    parser.add_argument('--language', choices=['python', 'fortran'], default='fortran',
                      help='Output language for getrates')
    parser.add_argument('--omp', action='store_true',
                      help='Enable OpenMP support')
    parser.add_argument('--vector', action='store_true',
                      help='Enable vectorized version')
    parser.add_argument('--module', action='store_true',
                      help='Generate as module (for Fortran)')
    parser.add_argument('--input-mw', action='store_true',
                      help='Take molecular weights as input')
    parser.add_argument('--rtypes-together', action='store_true',
                      help='Write reaction types together')

    #mini app options
    parser.add_argument('--ng', type=int, default=1,
                      help='Number of grid points')
    parser.add_argument('--ncpu', type=int, default=64,
                        help='Number of CPUs per node')
    parser.add_argument('--ngpu', type=int, default=8,
                        help='Number of GPUs per node')
    parser.add_argument('--nt', type=int, default=1,
                        help='Number of time steps')
    parser.add_argument('--time-cpu', action='store_true',
                        help='Time the CPU code')

    return parser.parse_args()

def generate_gpu_coefficients(args):
    # Create parser and chemistry objects
    ckp = ckparser()
    chem = chemistry(args.mech, ckp, therm_file=args.therm, yaml_file=args.yaml_file)

    # Generate coefficient module
    write_coef_module(args.output, chem,
                     parallel_level=args.parallel_level,
                     nreact_per_block=args.nreact_per_block)

    # Generate constants header
    write_constants_header(args.output, chem,
                         parallel_level=args.parallel_level,
                         veclen=args.veclen,
                         nreact_per_block=args.nreact_per_block)

def generate_getrates(args):
    # Create parser and chemistry objects
    ckp = ckparser()
    chem = chemistry(args.mech, ckp, therm_file=args.therm, yaml_file=args.yaml_file)

    # Create chemistry expressions object with specified options
    chem_expr = chemistry_expressions(chem,
                                   vec=args.vector,
                                   omp=args.omp,
                                   mod=args.module,
                                   language=args.language)

    # Write expressions to file
    chem_expr.write_expressions_to_file(args.output,
                                      write_rtypes_together=args.rtypes_together,
                                      input_MW=args.input_mw)

def generate_mini_app(args):
    # Create parser and chemistry objects
    ckp = ckparser()
    chem = chemistry(args.mech, ckp, therm_file=args.therm, yaml_file=args.yaml_file)


    src_dir = os.path.join(args.output,f"src")
    os.makedirs(src_dir,exist_ok=True)
    #write scalar f90
    # Create chemistry expressions object with specified options
    filename = os.path.join(src_dir,f"getrates.f90")
    chem_expr = chemistry_expressions(chem,
                                   vec=False,
                                   omp=False,
                                   mod=False,
                                   language=args.language)

    # Write expressions to file
    chem_expr.write_expressions_to_file(filename,
                                      write_rtypes_together=args.rtypes_together,
                                      input_MW=args.input_mw)

    #write vector f90
    # Create chemistry expressions object with specified options
    filename = os.path.join(src_dir,f"getrates_i.f90")
    chem_expr = chemistry_expressions(chem,
                                   vec=True,
                                   omp=False,
                                   mod=False,
                                   language=args.language)

    # Write expressions to file
    chem_expr.write_expressions_to_file(filename,
                                      write_rtypes_together=args.rtypes_together,
                                      input_MW=args.input_mw)

    #write chemgen f90
    # Create chemistry expressions object with specified options
    filename = os.path.join(src_dir,f"chemgen_m.f90")
    chem_expr = chemistry_expressions(chem,
                                   vec=False,
                                   omp=True,
                                   mod=True,
                                   language=args.language)

    # Write expressions to file
    chem_expr.write_expressions_to_file(filename,
                                      write_rtypes_together=args.rtypes_together,
                                      input_MW=args.input_mw)

    inc_dir = os.path.join(src_dir,"include")
    mod_dir = os.path.join(src_dir,"modules")
    os.makedirs(inc_dir, exist_ok=True)
    os.makedirs(mod_dir, exist_ok=True)
    for parallel_level in [1,2,3,4]:
        # Generate coefficient module
        write_coef_module(mod_dir, chem,
                         parallel_level=parallel_level,
                         nreact_per_block=args.nreact_per_block)
        # Generate constants header
        write_constants_header(inc_dir, chem,
                             parallel_level=parallel_level,
                             veclen=args.veclen,
                             nreact_per_block=args.nreact_per_block)

    #write hip kernels
    chem_expr.write_hip_kernels(src_dir)
    #write mini app
    chem_expr.write_chemistry_mini_app(src_dir,
                                       args.ng,
                                       ncpu=args.ncpu,
                                       ngpu=args.ngpu,
                                       nt=args.nt,
                                       input_MW=args.input_mw,
                                       time_cpu=args.time_cpu)

def main():
    args = parse_arguments()

    if args.mode == 'gpu_coef':
        generate_gpu_coefficients(args)
    elif args.mode == "mini_app":
        generate_mini_app(args)
    else:  # getrates
        generate_getrates(args)

if __name__ == "__main__":
    main()