[pypy-commit] pypy numpy-single-jitdriver: this branch is meant to reduce the number of jit merge points required for numpy

fijal noreply at buildbot.pypy.org
Thu Feb 2 16:08:08 CET 2012


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-single-jitdriver
Changeset: r52031:de5495b0f7fb
Date: 2012-02-02 17:07 +0200
http://bitbucket.org/pypy/pypy/changeset/de5495b0f7fb/

Log:	this branch is meant to reduce the number of jit merge points
	required for numpy to one. This should make it easier to do future
	optimizations

diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -86,8 +86,9 @@
 
     def apply_transformations(self, arr, transformations):
         v = self
-        for transform in transformations:
-            v = v.transform(arr, transform)
+        if transformations is not None:
+            for transform in transformations:
+                v = v.transform(arr, transform)
         return v
 
     def transform(self, arr, t):
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
@@ -12,7 +12,7 @@
 from pypy.rpython.lltypesystem import lltype, rffi
 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib.rstring import StringBuilder
-from pypy.module.micronumpy.interp_iter import (ArrayIterator, OneDimIterator,
+from pypy.module.micronumpy.interp_iter import (ArrayIterator,
     SkipLastAxisIterator, Chunk, ViewIterator)
 from pypy.module.micronumpy.appbridge import get_appbridge_cache
 
@@ -750,22 +750,10 @@
         raise NotImplementedError
 
     def compute(self):
-        result = W_NDimArray(self.size, self.shape, self.find_dtype())
-        shapelen = len(self.shape)
-        sig = self.find_sig()
-        frame = sig.create_frame(self)
-        ri = ArrayIterator(self.size)
-        while not ri.done():
-            numpy_driver.jit_merge_point(sig=sig,
-                                         shapelen=shapelen,
-                                         result_size=self.size,
-                                         frame=frame,
-                                         ri=ri,
-                                         self=self, result=result)
-            result.setitem(ri.offset, sig.eval(frame, self))
-            frame.next(shapelen)
-            ri = ri.next(shapelen)
-        return result
+        from pypy.module.micronumpy import loop
+        ra = ResultArray(self, self.size, self.shape, self.res_dtype)
+        loop.compute(ra)
+        return ra.left
 
     def force_if_needed(self):
         if self.forced_result is None:
@@ -864,6 +852,27 @@
         return signature.Call2(self.ufunc, self.name, self.calc_dtype,
                                self.left.create_sig(), self.right.create_sig())
 
+class ComputationArray(BaseArray):
+    """ A base class for all objects that describe operations for computation
+    """
+
+class ResultArray(Call2):
+    def __init__(self, child, size, shape, dtype, res=None, order='C'): 
+        if res is None:
+            res = W_NDimArray(size, shape, dtype, order)
+        Call2.__init__(self, None, 'assign', shape, dtype, dtype, res, child)
+
+    def create_sig(self):
+        return signature.ResultSignature(self.res_dtype, self.left.create_sig(),
+                                         self.right.create_sig())
+
+class Reduce(ComputationArray):
+    def __init__(self):
+        pass
+    
+    def create_sig(self):
+        return signature.ReduceSignature(self.func)
+
 class SliceArray(Call2):
     def __init__(self, shape, dtype, left, right, no_broadcast=False):
         self.no_broadcast = no_broadcast
@@ -883,10 +892,6 @@
                                             lsig, rsig)
 
 class AxisReduce(Call2):
-    """ NOTE: this is only used as a container, you should never
-    encounter such things in the wild. Remove this comment
-    when we'll make AxisReduce lazy
-    """
     _immutable_fields_ = ['left', 'right']
 
     def __init__(self, ufunc, name, shape, dtype, left, right, dim):
@@ -1039,9 +1044,9 @@
                            parent.order, parent)
         self.start = start
 
-    def create_iter(self):
+    def create_iter(self, transforms=None):
         return ViewIterator(self.start, self.strides, self.backstrides,
-                            self.shape)
+                            self.shape).apply_transformations(self, transforms)
 
     def setshape(self, space, new_shape):
         if len(self.shape) < 1:
@@ -1090,8 +1095,8 @@
         self.shape = new_shape
         self.calc_strides(new_shape)
 
-    def create_iter(self):
-        return ArrayIterator(self.size)
+    def create_iter(self, transforms=None):
+        return ArrayIterator(self.size).apply_transformations(self, transforms)
 
     def create_sig(self):
         return signature.ArraySignature(self.dtype)
@@ -1427,6 +1432,12 @@
     def create_sig(self):
         return signature.FlatSignature(self.base.dtype)
 
+    def create_iter(self, transforms=None):
+        return ViewIterator(self.base.start, self.base.strides, 
+                    self.base.backstrides,
+                    self.base.shape).apply_transformations(self.base,
+                                                           transforms)
+
     def descr_base(self, space):
         return space.wrap(self.base)
 
diff --git a/pypy/module/micronumpy/loop.py b/pypy/module/micronumpy/loop.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/micronumpy/loop.py
@@ -0,0 +1,30 @@
+
+""" This file is the main run loop as well as evaluation loops for various
+signatures
+"""
+
+from pypy.rlib.jit import JitDriver
+from pypy.module.micronumpy import signature
+
+def get_printable_location(shapelen, sig):
+    return 'numpy ' + sig.debug_repr() + ' [%d dims]' % (shapelen,)
+
+numpy_driver = JitDriver(
+    greens=['shapelen', 'sig'],
+    virtualizables=['frame'],
+    reds=['frame', 'arr'],
+    get_printable_location=signature.new_printable_location('numpy'),
+    name='numpy',
+)
+
+def compute(arr):
+    sig = arr.find_sig()
+    shapelen = len(arr.shape)
+    frame = sig.create_frame(arr)
+    while not frame.done():
+        numpy_driver.jit_merge_point(sig=sig,
+                                     shapelen=shapelen,
+                                     frame=frame, arr=arr)
+        sig.eval(frame, arr)
+        frame.next(shapelen)
+
diff --git a/pypy/module/micronumpy/signature.py b/pypy/module/micronumpy/signature.py
--- a/pypy/module/micronumpy/signature.py
+++ b/pypy/module/micronumpy/signature.py
@@ -4,6 +4,7 @@
      ConstantIterator, AxisIterator, ViewTransform,\
      BroadcastTransform
 from pypy.rlib.jit import hint, unroll_safe, promote
+from pypy.tool.pairtype import extendabletype
 
 """ Signature specifies both the numpy expression that has been constructed
 and the assembler to be compiled. This is a very important observation -
@@ -113,6 +114,8 @@
     return r_dict(sigeq_no_numbering, sighash)
 
 class Signature(object):
+    __metaclass_ = extendabletype
+    
     _attrs_ = ['iter_no', 'array_no']
     _immutable_fields_ = ['iter_no', 'array_no']
 
@@ -143,7 +146,6 @@
         self._create_iter(iterlist, arraylist, arr, [])
         return NumpyEvalFrame(iterlist, arraylist)
 
-
 class ConcreteSignature(Signature):
     _immutable_fields_ = ['dtype']
 
@@ -182,13 +184,10 @@
         assert isinstance(concr, ConcreteArray)
         storage = concr.storage
         if self.iter_no >= len(iterlist):
-            iterlist.append(self.allocate_iter(concr, transforms))
+            iterlist.append(concr.create_iter(transforms))
         if self.array_no >= len(arraylist):
             arraylist.append(storage)
 
-    def allocate_iter(self, arr, transforms):
-        return ArrayIterator(arr.size).apply_transformations(arr, transforms)
-
     def eval(self, frame, arr):
         iter = frame.iterators[self.iter_no]
         return self.dtype.getitem(frame.arrays[self.array_no], iter.offset)
@@ -220,22 +219,10 @@
         allnumbers.append(no)
         self.iter_no = no
 
-    def allocate_iter(self, arr, transforms):
-        return ViewIterator(arr.start, arr.strides, arr.backstrides,
-                            arr.shape).apply_transformations(arr, transforms)
-
 class FlatSignature(ViewSignature):
     def debug_repr(self):
         return 'Flat'
 
-    def allocate_iter(self, arr, transforms):
-        from pypy.module.micronumpy.interp_numarray import W_FlatIterator
-        assert isinstance(arr, W_FlatIterator)
-        return ViewIterator(arr.base.start, arr.base.strides, 
-                    arr.base.backstrides,
-                    arr.base.shape).apply_transformations(arr.base,
-                                                         transforms)
-
 class VirtualSliceSignature(Signature):
     def __init__(self, child):
         self.child = child
@@ -359,6 +346,17 @@
         return 'Call2(%s, %s, %s)' % (self.name, self.left.debug_repr(),
                                       self.right.debug_repr())
 
+class ResultSignature(Call2):
+    def __init__(self, dtype, left, right):
+        Call2.__init__(self, None, 'assign', dtype, left, right)
+
+    def eval(self, frame, arr):
+        from pypy.module.micronumpy.interp_numarray import ResultArray
+
+        assert isinstance(arr, ResultArray)
+        offset = frame.get_final_iter().offset
+        arr.left.setitem(offset, self.right.eval(frame, arr.right))
+
 class BroadcastLeft(Call2):
     def _invent_numbering(self, cache, allnumbers):
         self.left._invent_numbering(new_cache(), allnumbers)


More information about the pypy-commit mailing list