1
2
3 """
4 This module contains the Just-In-Time compiler tools.
5 The main function is jit(form, options), see its documentation.
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 import shutil
31 import os
32 import hashlib
33 import instant
34 import ufc_utils
35 import ufl
36 import ufl.algorithms
37 from ufl.algorithms import preprocess
38
39 from sfc.common import version
40 from sfc.common import default_parameters
41 from sfc.common import ParameterDict
42 from sfc.common import get_abs_ufc_h_path
43 from sfc.common import sfc_debug, sfc_info, sfc_warning, sfc_error, sfc_assert, write_file
44 from sfc.common.names import short_name, form_classname, dof_map_classname, finite_element_classname, base_element_classname
45 from sfc.codegeneration import generate_code, compiler_input
46
48 "A signature string uniquely identifying jit input."
49
50
51 if not isinstance(input, list):
52 input = [input]
53 newinput = []
54 for i in input:
55 if isinstance(i, ufl.form.Form):
56 i = i.signature()
57
58 newinput.append(i)
59 return "sfc.jit(%r, %r) # SFC version %r" % (newinput, options, version)
60
61 -def compile_module(module_dir, hfilenames, cfilenames, signature, options):
62 """Assuming an existing directory module_dir with source files
63 hfilenames, cfilenames, create a python extension module and compile it."""
64
65 orig_dir = os.getcwd()
66 try:
67
68 system_headers = ["iostream", "stdexcept", "cmath", "ufc.h", "tr1/memory"]
69
70 cppargs = options.compilation.cppargs
71 if options.compilation.enable_debug_code:
72 cppargs.append("-DSFCDEBUG")
73
74 module_name = module_dir if signature is None else None
75
76 module = ufc_utils.build_ufc_module(h_files = hfilenames,
77 source_directory = module_dir,
78 sources = cfilenames,
79 system_headers = system_headers,
80 cppargs = cppargs,
81 signature = signature,
82 cache_dir = options.compilation.cache_dir,
83 modulename = module_name,
84 generate_interface = options.compilation.generate_interface,
85 generate_setup = options.compilation.generate_setup)
86 sfc_info("Successfully compiled module in '%s'." % module_dir)
87 return module
88
89 finally:
90 os.chdir(orig_dir)
91
92 -def jit(input, parameters=None, common_cell=None):
93 ufl_elements, formdatas = compiler_input(input, common_cell=common_cell)
94
95
96 if parameters is None:
97 parameters = default_parameters()
98 if isinstance(parameters, dict) and not isinstance(parameters, ParameterDict):
99 o = parameters
100 parameters = default_parameters()
101 parameters.update(o)
102
103
104 prefix = parameters.code.prefix
105 if not prefix:
106 prefix = "sfc_jit"
107
108
109 signature = cache_signature(input, parameters)
110 formclassnames = [form_classname(fd.preprocessed_form, parameters) for fd in formdatas]
111
112 dmclassnames = [dof_map_classname(e) for e in ufl_elements]
113 feclassnames = [finite_element_classname(e) for e in ufl_elements]
114
115
116 module = None
117
118
119
120 if parameters.compilation.skip:
121 pass
122 else:
123
124 if parameters.compilation.overwrite_cache:
125 sfc_warning("Complete implementation of overwrite_cache not done, but skipping cache lookup before code generation.")
126
127 else:
128
129 module = instant.import_module(signature)
130
131 if module:
132 sfc_info("Found module in cache.")
133
134 if module is None:
135
136 orig_dir = os.getcwd()
137
138
139 module_dir_name = "_".join(["_".join(formclassnames), "_".join(dmclassnames), "_".join(feclassnames)])
140 module_dir_name = short_name(prefix, module_dir_name)
141
142
143 module_dir = os.path.abspath(module_dir_name)
144 if os.path.exists(module_dir):
145 sfc_warning("Possibly overwriting files in existing module directory '%s'." % module_dir)
146 else:
147 os.mkdir(module_dir)
148
149
150 try:
151 os.chdir(module_dir)
152
153
154 write_file("parameters.repr", repr(parameters))
155 write_file("signature", signature)
156 for fd in formdatas:
157 write_file("%s.repr" % short_name("form_", fd.name), repr(fd.preprocessed_form))
158 for ue in ufl_elements:
159 write_file("%s.repr" % short_name("element_", base_element_classname(ue)), repr(ue))
160
161
162 hfilenames, cfilenames = generate_code(ufl_elements + formdatas, None,
163 parameters, common_cell=common_cell)
164
165
166 if parameters.compilation.skip:
167 return None
168
169 finally:
170
171 os.chdir(orig_dir)
172
173
174 module = compile_module(module_dir, hfilenames, cfilenames, signature, parameters)
175
176
177 placed_in_cache = False
178 if placed_in_cache:
179 sfc_info("Module was placed in cache, deleting module directory '%s'." % module_dir)
180 shutil.rmtree(module_dir)
181
182 if module is None:
183 sfc_error("Failed to load compiled module from cache.")
184
185
186
187
188
189
190 if formdatas:
191 prefix = "sfc dummy prefix for pydolfin which I have no idea what is used for"
192 ufc_forms = [getattr(module, name)() for name in formclassnames]
193 return ufc_forms[0], module, formdatas[0], prefix
194
195
196 elif ufl_elements:
197 ufc_elements = [(getattr(module, fe)(), getattr(module, dm)()) for (fe, dm) in zip(feclassnames, dmclassnames)]
198 return ufc_elements[0]
199
200 msg = "Nothing to return, this shouldn't happen.\ninput =\n%r" % repr(input)
201 sfc_error(msg)
202
203
204 -def jitf(input, objects, parameters = None):
205 """Generate code from input and parameters.
206
207 @param input:
208 TODO
209 @param parameters:
210 TODO
211 @return:
212 FIXME: Agree on return values with FFC! Currently FFC returns (ufc_form, module, formdata), and dolfin assumes this
213 """
214
215
216 list_input = isinstance(input, list)
217 ufl_elements, formdatas = compiler_input(input, objects=objects)
218
219
220 if parameters is None:
221 parameters = default_parameters()
222 if isinstance(parameters, dict) and not isinstance(parameters, ParameterDict):
223 o = parameters
224 parameters = default_parameters()
225 parameters.update(o)
226
227
228 prefix = parameters.code.prefix
229 if not prefix:
230 prefix = "sfc_jit"
231
232
233 signature = cache_signature(input, parameters)
234 formclassnames = [form_classname(fd.preprocessed_form, parameters) for fd in formdatas]
235 dmclassnames = [dof_map_classname(e) for e in ufl_elements]
236 feclassnames = [finite_element_classname(e) for e in ufl_elements]
237
238
239 module = None
240
241
242 if parameters.compilation.skip:
243 pass
244 else:
245
246 if parameters.compilation.overwrite_cache:
247 sfc_warning("Complete implementation of overwrite_cache not done, but skipping cache lookup before code generation.")
248
249 else:
250
251 module = instant.import_module(signature)
252
253 if module:
254 sfc_info("Found module in cache.")
255
256 if module is None:
257
258 orig_dir = os.getcwd()
259
260
261 module_dir_name = "_".join(["_".join(formclassnames), "_".join(dmclassnames), "_".join(feclassnames)])
262 module_dir_name = short_name(prefix, module_dir_name)
263
264
265 module_dir = os.path.abspath(module_dir_name)
266 if os.path.exists(module_dir):
267 sfc_warning("Possibly overwriting files in existing module directory '%s'." % module_dir)
268 else:
269 os.mkdir(module_dir)
270
271
272 try:
273 os.chdir(module_dir)
274
275
276 write_file("parameters.repr", repr(parameters))
277 write_file("signature", signature)
278 for fd in formdatas:
279 write_file("%s.repr" % short_name("form_", fd.name), repr(fd.preprocessed_form))
280 for ue in ufl_elements:
281 write_file("%s.repr" % short_name("element_", base_element_classname(ue)), repr(ue))
282
283
284 hfilenames, cfilenames = generate_code(ufl_elements + formdatas, objects,
285 parameters)
286
287
288 if parameters.compilation.skip:
289 return None
290
291 finally:
292
293 os.chdir(orig_dir)
294
295
296 module = compile_module(module_dir, hfilenames, cfilenames, signature, parameters)
297
298
299 placed_in_cache = False
300 if placed_in_cache:
301 sfc_info("Module was placed in cache, deleting module directory '%s'." % module_dir)
302 shutil.rmtree(module_dir)
303
304 if module is None:
305 sfc_error("Failed to load compiled module from cache.")
306
307
308
309
310
311
312 if formdatas:
313 prefix = "sfc dummy prefix for pydolfin which I have no idea what is used for"
314 ufc_forms = [getattr(module, name) for name in formclassnames]
315 if list_input:
316 return ufc_forms, module, formdatas, [prefix]*len(ufc_forms)
317 else:
318 return ufc_forms[0], module, formdatas[0], prefix
319
320
321 elif ufl_elements:
322 ufc_elements = [(getattr(module, fe)(), getattr(module, dm)()) for (fe, dm) in zip(feclassnames, dmclassnames)]
323 if list_input:
324 return ufc_elements
325 else:
326 return ufc_elements[0]
327
328 msg = "Nothing to return, this shouldn't happen.\ninput =\n%r" % repr(input)
329 sfc_error(msg)
330