hedge.backends.jit.compiler

1 """Just-in-time compiling backend.""" 2 3 from __future__ import division 4 5 __copyright__ = "Copyright (C) 2008 Andreas Kloeckner" 6 7 __license__ = """ 8 This program is free software: you can redistribute it and/or modify 9 it under the terms of the GNU General Public License as published by 10 the Free Software Foundation, either version 3 of the License, or 11 (at your option) any later version. 12 13 This program is distributed in the hope that it will be useful, 14 but WITHOUT ANY WARRANTY; without even the implied warranty of 15 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 16 GNU General Public License for more details. 17 18 You should have received a copy of the GNU General Public License 19 along with this program. If not, see U{http://www.gnu.org/licenses/}. 20 """ 21 22 23 24 25 import hedge.discretization 26 import hedge.optemplate 27 from pytools import memoize_method 28 from hedge.compiler import OperatorCompilerBase, FluxBatchAssign, \ 29 Assign

30 31 32 33 34 # compiler stuff -------------------------------------------------------------- 35 -class VectorExprAssign(Assign):

36 __slots__ = ["toolchain"] 37

38 - def get_executor_method(self, executor):

39 return executor.exec_vector_expr_assign

40 41 comment = "compiled" 42 43 @memoize_method

44 - def compiled(self, executor):

45 discr = executor.discr 46 47 if self.flop_count() > 500: 48 # reduce optimization level for complicated expressions 49 if "jit_dont_optimize_large_exprs" in discr.debug: 50 toolchain = discr.toolchain.with_optimization_level(0) 51 else: 52 toolchain = discr.toolchain.with_optimization_level(1) 53 else: 54 toolchain = discr.toolchain 55 56 from hedge.backends.vector_expr import \ 57 VectorExpressionInfo, simple_result_dtype_getter 58 from hedge.backends.jit.vector_expr import CompiledVectorExpression 59 return CompiledVectorExpression( 60 [VectorExpressionInfo( 61 name=name, 62 expr=expr, 63 do_not_return=dnr) 64 for name, expr, dnr in zip( 65 self.names, self.exprs, self.do_not_return)], 66 result_dtype_getter=simple_result_dtype_getter, 67 toolchain=toolchain, 68 wait_on_error="jit_wait_on_compile_error" in discr.debug)

69

70 71 72 -class CompiledFluxBatchAssign(FluxBatchAssign):

73 # members: compiled_func, arg_specs, is_boundary 74 75 @memoize_method

76 - def get_dependencies(self):

77 deps = set() 78 79 from hedge.tools import setify_field as setify 80 from hedge.optemplate import OperatorBinding, BoundaryPair 81 for f in self.fluxes: 82 assert isinstance(f, OperatorBinding) 83 if isinstance(f.field, BoundaryPair): 84 deps |= setify(f.field.field) | setify(f.field.bfield) 85 else: 86 deps |= setify(f.field) 87 88 dep_mapper = self.dep_mapper_factory() 89 90 from pytools import flatten 91 return set(flatten(dep_mapper(dep) for dep in deps))

92 93 @memoize_method

94 - def get_module(self, discr, dtype):

95 from hedge.backends.jit.flux import \ 96 get_interior_flux_mod, \ 97 get_boundary_flux_mod 98 99 if not self.is_boundary: 100 mod = get_interior_flux_mod( 101 self.fluxes, self.flux_var_info, 102 discr, dtype) 103 104 if discr.instrumented: 105 from hedge.tools import time_count_flop, gather_flops 106 mod.gather_flux = \ 107 time_count_flop( 108 mod.gather_flux, 109 discr.gather_timer, 110 discr.gather_counter, 111 discr.gather_flop_counter, 112 len(self.fluxes) 113 * gather_flops(discr) 114 * len(self.flux_var_info.arg_names)) 115 116 else: 117 mod = get_boundary_flux_mod( 118 self.fluxes, self.flux_var_info, discr, dtype) 119 120 if discr.instrumented: 121 from pytools.log import time_and_count_function 122 mod.gather_flux = time_and_count_function( 123 mod.gather_flux, discr.gather_timer) 124 125 return mod

126

127 128 129 130 131 # flux kinds ------------------------------------------------------------------ 132 -class InteriorFluxKind(object):

133 - def __hash__(self):

134 return hash(self.__class__)

135

136 - def __str__(self):

137 return "interior"

138

139 - def __eq__(self, other):

140 return (other.__class__ == self.__class__)

141

142 -class BoundaryFluxKind(object):

143 - def __init__(self, tag):

144 self.tag = tag

145

146 - def __str__(self):

147 return "boundary(%s)" % self.tag

148

149 - def __hash__(self):

150 return hash((self.__class__, self.tag))

151

152 - def __eq__(self, other):

153 return (other.__class__ == self.__class__ 154 and other.tag == self.tag)

155

156 157 158 159 -class OperatorCompiler(OperatorCompilerBase):

160 - def __init__(self, discr):

161 OperatorCompilerBase.__init__(self, 162 max_vectors_in_batch_expr=100) 163 self.discr = discr

164

165 - def get_contained_fluxes(self, expr):

166 from hedge.optemplate import FluxCollector, BoundaryPair 167 from hedge.tools import is_obj_array 168 169 def get_deps(field): 170 if is_obj_array(field): 171 return set(field) 172 else: 173 return set([field])

174 175 def get_flux_deps(op_binding): 176 if isinstance(op_binding.field, BoundaryPair): 177 bpair = op_binding.field 178 return get_deps(bpair.field) | get_deps(bpair.bfield) 179 else: 180 return get_deps(op_binding.field)

181 182 def get_flux_kind(op_binding): 183 if isinstance(op_binding.field, BoundaryPair): 184 return BoundaryFluxKind(op_binding.field.tag) 185 else: 186 return InteriorFluxKind() 187 188 return [self.FluxRecord( 189 flux_expr=flux_binding, 190 kind=get_flux_kind(flux_binding), 191 dependencies=get_flux_deps(flux_binding)) 192 for flux_binding in FluxCollector()(expr)] 193

194 - def internal_map_flux(self, flux_bind):

195 from hedge.optemplate import IdentityMapper 196 return IdentityMapper.map_operator_binding(self, flux_bind)

197

198 - def map_operator_binding(self, expr):

199 from hedge.optemplate import FluxOperatorBase 200 if isinstance(expr.op, FluxOperatorBase): 201 return self.map_planned_flux(expr) 202 else: 203 return OperatorCompilerBase.map_operator_binding(self, expr)

204 205 # flux compilation --------------------------------------------------------

206 - def make_flux_batch_assign(self, names, fluxes, kind):

207 if isinstance(kind, BoundaryFluxKind): 208 return self.make_boundary_flux_batch_assign(names, fluxes, kind) 209 elif isinstance(kind, InteriorFluxKind): 210 return self.make_interior_flux_batch_assign(names, fluxes, kind) 211 else: 212 raise ValueError("invalid flux batch kind: %s" % kind)

213

214 - def make_interior_flux_batch_assign(self, names, fluxes, kind):

215 from hedge.backends.jit.flux import get_flux_var_info 216 return CompiledFluxBatchAssign(is_boundary=False, 217 names=names, fluxes=fluxes, kind=kind, 218 flux_var_info=get_flux_var_info(fluxes), 219 dep_mapper_factory=self.dep_mapper_factory)

220

221 - def make_boundary_flux_batch_assign(self, names, fluxes, kind):

222 from hedge.backends.jit.flux import get_flux_var_info 223 return CompiledFluxBatchAssign(is_boundary=True, 224 names=names, fluxes=fluxes, kind=kind, 225 flux_var_info=get_flux_var_info(fluxes), 226 dep_mapper_factory=self.dep_mapper_factory)

227 228 # vector math -------------------------------------------------------------

229 - def finalize_multi_assign(self, names, exprs, do_not_return, priority):

230 return VectorExprAssign(names=names, exprs=exprs, 231 do_not_return=do_not_return, 232 dep_mapper_factory=self.dep_mapper_factory, 233 priority=priority)

234

Source Code for Module hedge.backends.jit.compiler