[pypy-commit] pypy numpy-back-to-applevel: write a jitdriver for take

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-back-to-applevel
Changeset: r51849:e656f86d3017
Date: 2012-01-27 13:41 +0200
http://bitbucket.org/pypy/pypy/changeset/e656f86d3017/

Log:	write a jitdriver for take

diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -55,6 +55,9 @@
         """NOT_RPYTHON"""
         self.fromcache = InternalSpaceCache(self).getorbuild
 
+    def _freeze_(self):
+        return True
+
     def issequence_w(self, w_obj):
         return isinstance(w_obj, ListObject) or isinstance(w_obj, W_NDimArray)
 
@@ -146,6 +149,9 @@
     def allocate_instance(self, klass, w_subtype):
         return instantiate(klass)
 
+    def newtuple(self, list_w):
+        raise ValueError
+
     def len_w(self, w_obj):
         if isinstance(w_obj, ListObject):
             return len(w_obj.items)
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
@@ -63,7 +63,7 @@
 )
 take_driver = jit.JitDriver(
     greens=['shapelen', 'sig'],
-    reds=['index_i', 'res_i', 'concr'],
+    reds=['index_i', 'res_i', 'concr', 'index', 'res'],
     name='numpy_take',
 )
 
@@ -614,21 +614,25 @@
             size *= elem
         res = W_NDimArray(size, res_shape[:], concr.dtype, concr.order)
         res_i = res.create_iter()
-        longdtype = interp_dtype.get_dtype_cache(space).w_longdtype
         shapelen = len(index.shape)
         sig = concr.find_sig()
         while not index_i.done():
-            take_driver.jit_merge_point(index_i=index_i,
+            take_driver.jit_merge_point(index_i=index_i, index=index,
                                         res_i=res_i, concr=concr,
+                                        res=res,
                                         shapelen=shapelen, sig=sig)
-            # XXX jitdriver + test_zjit
-            w_item = index.getitem(index_i.offset).convert_to(longdtype).item(
-                space)
+            w_item = index._getitem_long(space, index_i.offset)
             res.setitem(res_i.offset, concr.descr_getitem(space, w_item))
             index_i = index_i.next(shapelen)
             res_i = res_i.next(shapelen)
         return res
 
+    def _getitem_long(self, space, offset):
+        # an obscure hack to not have longdtype inside a jitted loop
+        longdtype = interp_dtype.get_dtype_cache(space).w_longdtype
+        return self.getitem(offset).convert_to(longdtype).item(
+            space)
+
     def descr_item(self, space, w_arg=None):
         if space.is_w(w_arg, space.w_None):
             if not isinstance(self, Scalar):
diff --git a/pypy/module/micronumpy/test/test_zjit.py b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -287,6 +287,27 @@
                                 'jump': 1,
                                 'arraylen_gc': 1})
 
+    def define_take():
+        return """
+        a = |10|
+        b = take(a, [1, 1, 3, 2])
+        b -> 2
+        """
+
+    def test_take(self):
+        result = self.run("take")
+        assert result == 3
+        self.check_simple_loop({'getinteriorfield_raw': 2,
+                                'cast_float_to_int': 1,
+                                'int_lt': 1,
+                                'int_ge': 2,
+                                'guard_false': 3,
+                                'setinteriorfield_raw': 1,
+                                'int_mul': 1,
+                                'int_add': 3,
+                                'jump': 1,
+                                'arraylen_gc': 2})
+
     def define_multidim():
         return """
         a = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]