[pypy-commit] pypy default: fixed numpy's setslice to play nicely with the jit.

[justinpeel]

Author: Justin Peel <notmuchtotell at gmail.com>
Branch: 
Changeset: r45854:631e083bd996
Date: 2011-07-21 19:20 -0600
http://bitbucket.org/pypy/pypy/changeset/631e083bd996/

Log:	fixed numpy's setslice to play nicely with the jit.

diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -17,8 +17,8 @@
                              reds = ['result_size', 'i', 'self', 'result'])
 all_driver = jit.JitDriver(greens=['signature'], reds=['i', 'size', 'self'])
 any_driver = jit.JitDriver(greens=['signature'], reds=['i', 'size', 'self'])
-slice_driver1 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'storage', 'arr'])
-slice_driver2 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'storage', 'arr'])
+slice_driver1 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'source', 'dest'])
+slice_driver2 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'source', 'dest'])
 
 def add(v1, v2):
     return v1 + v2
@@ -247,23 +247,25 @@
     def descr_mean(self, space):
         return space.wrap(space.float_w(self.descr_sum(space))/self.find_size())
 
-    def _sliceloop1(self, start, stop, step, arr, storage):
+    def _sliceloop1(self, start, stop, step, source, dest):
         i = start
         j = 0
         while i < stop:
-            #slice_driver1.jit_merge_point(signature=arr.signature,
-            #        step=step, stop=stop, i=i, j=j, arr=arr, storage=storage)
-            storage[i] = arr.eval(j)
+            slice_driver1.jit_merge_point(signature=source.signature,
+                    step=step, stop=stop, i=i, j=j, source=source,
+                    dest=dest)
+            dest.storage[i] = source.eval(j)
             j += 1
             i += step
 
-    def _sliceloop2(self, start, stop, step, arr, storage):
+    def _sliceloop2(self, start, stop, step, source, dest):
         i = start
         j = 0
         while i > stop:
-            #slice_driver2.jit_merge_point(signature=arr.signature,
-            #        step=step, stop=stop, i=i, j=j, arr=arr, storage=storage)
-            storage[i] = arr.eval(j)
+            slice_driver2.jit_merge_point(signature=source.signature,
+                    step=step, stop=stop, i=i, j=j, source=source,
+                    dest=dest)
+            dest.storage[i] = source.eval(j)
             j += 1
             i += step
 
@@ -448,9 +450,9 @@
             stop = self.calc_index(stop)
         step = self.step * step
         if step > 0:
-            self._sliceloop1(start, stop, step, arr, self.parent.storage)
+            self._sliceloop1(start, stop, step, arr, self.parent)
         else:
-            self._sliceloop2(start, stop, step, arr, self.parent.storage)
+            self._sliceloop2(start, stop, step, arr, self.parent)
 
     def calc_index(self, item):
         return (self.start + item * self.step)
@@ -489,9 +491,9 @@
         if not isinstance(arr, BaseArray):
             arr = convert_to_array(space, arr)
         if step > 0:
-            self._sliceloop1(start, stop, step, arr, self.storage)
+            self._sliceloop1(start, stop, step, arr, self)
         else:
-            self._sliceloop2(start, stop, step, arr, self.storage)
+            self._sliceloop2(start, stop, step, arr, self)
 
     def __del__(self):
         lltype.free(self.storage, flavor='raw')