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
38 from sfc.common import version
39 from sfc.common import default_parameters
40 from sfc.common import ParameterDict
41 from sfc.common import get_abs_ufc_h_path
42 from sfc.common import sfc_debug, sfc_info, sfc_warning, sfc_error, sfc_assert, write_file
43 from sfc.common.names import short_name, form_classname, dof_map_classname, finite_element_classname, base_element_classname
44 from sfc.codegeneration import generate_code, compiler_input
45
47 "A signature string uniquely identifying jit input."
48 h = hashlib.md5(repr(version))
49 h.update(prefix)
50 for e in ufl_elements:
51 h.update(repr(e))
52 for fd in formdatas:
53 h.update(fd.signature)
54 h.update(repr(parameters))
55 return h.hexdigest().lower()
56
57 -def compile_module(module_dir, hfilenames, cfilenames, signature, options):
58 """Assuming an existing directory module_dir with source files
59 hfilenames, cfilenames, create a python extension module and compile it."""
60
61 orig_dir = os.getcwd()
62 try:
63
64 system_headers = ["iostream", "stdexcept", "cmath", "ufc.h", "tr1/memory"]
65
66 cppargs = options.compilation.cppargs
67 if options.compilation.enable_debug_code:
68 cppargs.append("-DSFCDEBUG")
69
70 module_name = module_dir if signature is None else None
71
72 module = ufc_utils.build_ufc_module(h_files = hfilenames,
73 source_directory = module_dir,
74 sources = cfilenames,
75 system_headers = system_headers,
76 cppargs = cppargs,
77 signature = signature,
78 cache_dir = options.compilation.cache_dir,
79 modulename = module_name,
80 generate_interface = options.compilation.generate_interface,
81 generate_setup = options.compilation.generate_setup)
82 sfc_info("Successfully compiled module in '%s'." % module_dir)
83 return module
84
85 finally:
86 os.chdir(orig_dir)
87
96
98 module = None
99 if parameters.compilation.skip:
100
101 sfc_warning("Skipping module lookup.")
102
103 elif parameters.compilation.overwrite_cache:
104
105 sfc_warning("Complete implementation of overwrite_cache not done, but skipping cache lookup before code generation.")
106
107
108 else:
109
110 module = instant.import_module(signature)
111
112 if module:
113 sfc_info("Found module in cache.")
114
115 return module
116
118
119 module_dir_name = "_".join(["_".join(formclassnames),
120 "_".join(dmclassnames),
121 "_".join(feclassnames)])
122 module_dir_name = short_name(prefix, module_dir_name)
123
124
125 module_dir = os.path.abspath(module_dir_name)
126 if os.path.exists(module_dir):
127 msg = "Possibly overwriting files in existing module directory '%s'."
128 sfc_warning(msg % module_dir)
129 else:
130 os.mkdir(module_dir)
131 return module_dir
132
140
141 -def pack_jit_return_values(formdatas, ufl_elements, module,
142 formclassnames, feclassnames, dmclassnames,
143 list_input):
144
145 if formdatas:
146 prefix = "sfc dummy prefix for pydolfin which I have no idea what is used for"
147 ufc_forms = [getattr(module, name)() for name in formclassnames]
148 if list_input:
149 return ufc_forms, module, formdatas, [prefix]*len(ufc_forms)
150 else:
151 return ufc_forms[0], module, formdatas[0], prefix
152
153
154 elif ufl_elements:
155 ufc_elements = [(getattr(module, fe)(), getattr(module, dm)())
156 for (fe, dm) in zip(feclassnames, dmclassnames)]
157 if list_input:
158 return ufc_elements
159 else:
160 return ufc_elements[0]
161
162 msg = "Nothing to return, this shouldn't happen.\ninput =\n%r" % repr(input)
163 sfc_error(msg)
164
165 -def jit(input, parameters=None, common_cell=None, objects=None):
166
167 list_input = isinstance(input, list)
168 ufl_elements, formdatas = compiler_input(input, objects=objects,
169 common_cell=common_cell)
170 parameters = merge_with_default_parameters(parameters)
171
172
173 prefix = parameters.code.prefix or "sfc_jit"
174
175
176 signature = cache_signature(ufl_elements, formdatas, parameters, prefix)
177
178
179 formclassnames = [form_classname(fd.signature, parameters)
180 for fd in formdatas]
181 dmclassnames = [dof_map_classname(e) for e in ufl_elements]
182 feclassnames = [finite_element_classname(e) for e in ufl_elements]
183
184
185 module = lookup_module(parameters, signature)
186 if module is None:
187 module_dir = make_module_dir(prefix, formclassnames, dmclassnames, feclassnames)
188
189 write_parameters(module_dir, parameters, signature, ufl_elements)
190
191
192 orig_dir = os.getcwd()
193 try:
194 os.chdir(module_dir)
195 hfilenames, cfilenames = generate_code(ufl_elements + formdatas, objects,
196 parameters, common_cell=common_cell)
197 finally:
198
199 os.chdir(orig_dir)
200
201
202 if parameters.compilation.skip:
203 return None
204
205
206 module = compile_module(module_dir, hfilenames, cfilenames,
207 signature, parameters)
208
209
210 placed_in_cache = False
211 if placed_in_cache:
212 msg = "Module was placed in cache, deleting module directory '%s'."
213 sfc_info(msg % module_dir)
214 shutil.rmtree(module_dir)
215
216 if module is None:
217 sfc_error("Failed to load compiled module from cache.")
218
219
220
221
222
223 return pack_jit_return_values(formdatas, ufl_elements, module,
224 formclassnames, feclassnames, dmclassnames,
225 list_input)
226
227
228 -def jitf(input, objects, parameters=None, common_cell=None):
229 return jit(input, parameters=parameters, common_cell=common_cell, objects=objects)
230