Package hedge :: Package backends :: Package jit :: Module diff
[hide private]
[frames] | no frames]

Source Code for Module hedge.backends.jit.diff

  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  import numpy 
 25  from pytools import memoize_method 
26 27 28 29 30 -class JitDifferentiator:
31 - def __init__(self, discr):
32 self.discr = discr
33 34 @memoize_method
35 - def make_diff(self, elgroup, dtype):
36 from hedge._internal import UniformElementRanges 37 assert isinstance(elgroup.ranges, UniformElementRanges) 38 39 ldis = elgroup.local_discretization 40 discr = self.discr 41 from codepy.cgen import \ 42 FunctionDeclaration, FunctionBody, Typedef, \ 43 Const, Reference, Value, POD, \ 44 Statement, Include, Line, Block, Initializer, Assign, \ 45 For, \ 46 Define 47 48 from pytools import to_uncomplex_dtype 49 50 from codepy.bpl import BoostPythonModule 51 mod = BoostPythonModule() 52 53 S = Statement 54 mod.add_to_preamble([ 55 Include("hedge/volume_operators.hpp"), 56 Include("boost/foreach.hpp"), 57 ]) 58 59 mod.add_to_module([ 60 S("namespace ublas = boost::numeric::ublas"), 61 S("using namespace hedge"), 62 S("using namespace pyublas"), 63 Line(), 64 Define("DOFS_PER_EL", ldis.node_count()), 65 Define("DIMENSIONS", discr.dimensions), 66 Line(), 67 Typedef(POD(dtype, "value_type")), 68 Typedef(POD(to_uncomplex_dtype(dtype), "uncomplex_type")), 69 ]) 70 71 fdecl = FunctionDeclaration( 72 Value("void", "diff"), 73 [ 74 Const(Reference(Value("uniform_element_ranges", "ers"))), 75 Value("numpy_array<value_type>", "field") 76 ]+[ 77 Value("ublas::matrix<uncomplex_type>", "diffmat_rst%d" % rst) 78 for rst in range(discr.dimensions) 79 ]+[ 80 Value("numpy_array<value_type>", "result%d" % i) 81 for i in range(discr.dimensions) 82 ]+[ 83 Value("numpy_array<double>", "coeffs"), 84 Value("numpy_array<npy_uint32>", "el_nrs"), 85 POD(numpy.uint32, "total_el_count"), 86 ] 87 ) 88 89 def make_it(name, is_const=True, tpname="value_type"): 90 if is_const: 91 const = "const_" 92 else: 93 const = "" 94 95 return Initializer( 96 Value("numpy_array<%s>::%siterator" % (tpname, const), name+"_it"), 97 "%s.begin()" % name)
98 99 fbody = Block([ 100 S("assert(DOFS_PER_EL == ers.el_size())"), 101 Line(), 102 make_it("field"), 103 ]+[ 104 make_it("result%d" % i, is_const=False) 105 for i in range(discr.dimensions) 106 ]+[ 107 Line(), 108 make_it("coeffs", tpname="double"), 109 make_it("el_nrs", tpname="npy_uint32"), 110 Line(), 111 For("element_number_t eg_el_nr = 0", 112 "eg_el_nr < ers.size()", 113 "++eg_el_nr", 114 Block([ 115 Initializer( 116 Value("node_number_t", "el_base"), 117 "ers.start() + eg_el_nr*DOFS_PER_EL"), 118 Initializer( 119 Value("element_number_t", "el_nr"), 120 "el_nrs_it[eg_el_nr]"), 121 Line(), 122 For("unsigned i = 0", 123 "i < DOFS_PER_EL", 124 "++i", 125 Block([ 126 Initializer(Value("value_type", "drst_%d" % rst), 0) 127 for rst in range(discr.dimensions) 128 ]+[ 129 Line(), 130 ]+[ 131 For("unsigned j = 0", 132 "j < DOFS_PER_EL", 133 "++j", 134 Block([ 135 S("drst_%(rst)d += diffmat_rst%(rst)d(i, j)*field_it[el_base+j]" 136 % {"rst":rst}) 137 for rst in range(discr.dimensions) 138 ]) 139 ), 140 Line(), 141 ]+[ 142 Assign("result%d_it[el_base+i]" % xyz, 143 " + ".join( 144 "uncomplex_type(coeffs_it[total_el_count*(" 145 "DIMENSIONS*%(xyz)d + %(rst)d) + el_nr])" 146 " * drst_%(rst)d" 147 % {"xyz":xyz, "rst":rst} 148 for rst in range(discr.dimensions) 149 ) 150 ) 151 for xyz in range(discr.dimensions) 152 153 ]) 154 ) 155 ]) 156 ) 157 ]) 158 159 mod.add_function(FunctionBody(fdecl, fbody)) 160 161 #print "----------------------------------------------------------------" 162 #print mod.generate() 163 #raw_input() 164 165 compiled_func = mod.compile(self.discr.toolchain, 166 wait_on_error="jit_wait_on_compile_error" in discr.debug).diff 167 168 if self.discr.instrumented: 169 from hedge.tools import time_count_flop 170 171 compiled_func = time_count_flop(compiled_func, 172 discr.diff_timer, discr.diff_counter, 173 discr.diff_flop_counter, 174 flops=discr.dimensions*( 175 2 # mul+add 176 * ldis.node_count() * len(elgroup.members) 177 * ldis.node_count() 178 + 179 2 * discr.dimensions 180 * len(elgroup.members) * ldis.node_count()), 181 increment=discr.dimensions) 182 183 return compiled_func
184
185 - def __call__(self, op_class, field, xyz_needed):
186 result = [self.discr.volume_zeros(dtype=field.dtype) 187 for i in range(self.discr.dimensions)] 188 from hedge.tools import is_zero 189 if not is_zero(field): 190 for eg in self.discr.element_groups: 191 coeffs = op_class.coefficients(eg) 192 193 from pytools import to_uncomplex_dtype 194 uncomplex_dtype = to_uncomplex_dtype(field.dtype) 195 args = ([eg.ranges, field] 196 + [m.astype(uncomplex_dtype) for m in op_class.matrices(eg)] 197 + result 198 + [coeffs, eg.member_nrs, coeffs.shape[2]]) 199 200 diff_routine = self.make_diff(eg, field.dtype) 201 diff_routine(*args) 202 203 return [result[i] for i in xyz_needed]
204