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