Package sfc :: Module jit
[hide private]
[frames] | no frames]

Source Code for Module sfc.jit

  1  #!/usr/bin/env python 
  2  # -*- coding: utf-8 -*- 
  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  # Copyright (C) 2008-2009 Martin Sandve Alnes and Simula Resarch Laboratory 
  9  # 
 10  # This file is part of SyFi. 
 11  # 
 12  # SyFi is free software: you can redistribute it and/or modify 
 13  # it under the terms of the GNU General Public License as published by 
 14  # the Free Software Foundation, either version 2 of the License, or 
 15  # (at your option) any later version. 
 16  # 
 17  # SyFi is distributed in the hope that it will be useful, 
 18  # but WITHOUT ANY WARRANTY; without even the implied warranty of 
 19  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 
 20  # GNU General Public License for more details. 
 21  # 
 22  # You should have received a copy of the GNU General Public License 
 23  # along with SyFi. If not, see <http://www.gnu.org/licenses/>. 
 24  # 
 25  # Modified by Kent-Andre Mardal, 2010. 
 26  # 
 27  # First added:  2008-09-04 
 28  # Last changed: 2009-04-23 
 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   
46 -def cache_signature(ufl_elements, formdatas, parameters, prefix):
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 # headers to #include in swig file 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
88 -def merge_with_default_parameters(parameters):
89 if parameters is None: 90 parameters = default_parameters() 91 if isinstance(parameters, dict) and not isinstance(parameters, ParameterDict): 92 o = parameters 93 parameters = default_parameters() 94 parameters.update(o) 95 return parameters
96
97 -def lookup_module(parameters, signature):
98 module = None 99 if parameters.compilation.skip: 100 # Hacky hook for profiling without the C++ compilation 101 sfc_warning("Skipping module lookup.") 102 103 elif parameters.compilation.overwrite_cache: 104 # Possibly force recompilation 105 sfc_warning("Complete implementation of overwrite_cache not done, but skipping cache lookup before code generation.") 106 # TODO: delete existing module with signature or module_name 107 108 else: 109 # Check cache before generating code 110 module = instant.import_module(signature) 111 # Construct ufc::form object 112 if module: 113 sfc_info("Found module in cache.") 114 115 return module
116
117 -def make_module_dir(prefix, formclassnames, dmclassnames, feclassnames):
118 # Build module name, then shorten if it's too long 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 # Make sure the module directory is there and empty 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
133 -def write_parameters(module_dir, parameters, signature, ufl_elements):
134 # Write parameters, signature and form representations to files beside code 135 write_file(os.path.join(module_dir, "parameters.repr"), repr(parameters)) 136 write_file(os.path.join(module_dir, "signature"), signature) 137 for ue in ufl_elements: 138 sn = short_name("element_", base_element_classname(ue)) 139 write_file(os.path.join(module_dir,"%s.repr" % sn), repr(ue))
140
141 -def pack_jit_return_values(formdatas, ufl_elements, module, 142 formclassnames, feclassnames, dmclassnames, 143 list_input):
144 # Got forms, return forms 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 # Got elements, return elements 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 # To handle all input uniformly 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 # Default prefix shows where the module was compiled 173 prefix = parameters.code.prefix or "sfc_jit" 174 175 # Build module sig from input data on canonical form 176 signature = cache_signature(ufl_elements, formdatas, parameters, prefix) 177 178 # Build classnames (TODO: This is duplicate work, done in formrepresentation) 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 # Lookup module in cache before compiling 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 # Enter module directory and generate code! 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 # Restore original path 199 os.chdir(orig_dir) 200 201 # Hook for profiling without the C++ compilation 202 if parameters.compilation.skip: 203 return None 204 205 # Invoke Instant to compile generated files as a Python extension module 206 module = compile_module(module_dir, hfilenames, cfilenames, 207 signature, parameters) 208 209 # Remove temporary module directory if code was copied to cache 210 placed_in_cache = False # TODO 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 # If we get here, 'module' should be our compiled module. 220 221 # TODO: Handle mixed list of forms and elements? 222 223 return pack_jit_return_values(formdatas, ufl_elements, module, 224 formclassnames, feclassnames, dmclassnames, 225 list_input)
226 227 # TODO: Get rid of this one
228 -def jitf(input, objects, parameters=None, common_cell=None):
229 return jit(input, parameters=parameters, common_cell=common_cell, objects=objects)
230