Rudimentary loops (range and xrange). Really good speedups

so far

2 files changed, 115 insertions, 35 deletions

diff --git a/weave/accelerate_tools.py b/weave/accelerate_tools.py
index c62ee6f97..8b45b16b4 100644
--- a/weave/accelerate_tools.py
+++ b/weave/accelerate_tools.py
@@ -9,11 +9,11 @@ C++ equivalents to Python functions.
 """
 #**************************************************************************#
 
-from types import FunctionType,IntType,FloatType,StringType,TypeType
+from types import FunctionType,IntType,FloatType,StringType,TypeType,XRangeType
 import inspect
 import md5
 import weave
-import bytecodecompiler
+from bytecodecompiler import CXXCoder,Type_Descriptor,Function_Descriptor
 
 def CStr(s):
     "Hacky way to get legal C string from Python string"
@@ -22,7 +22,7 @@ def CStr(s):
     r = repr('"'+s) # Better for embedded quotes
     return '"'+r[2:-1]+'"'
 
-class Basic(bytecodecompiler.Type_Descriptor):
+class Basic(Type_Descriptor):
     def check(self,s):
         return "%s(%s)"%(self.checker,s)
     def inbound(self,s):
@@ -60,10 +60,12 @@ class String(Basic):
 
 import Numeric
 
-class Vector(bytecodecompiler.Type_Descriptor):
+class Vector(Type_Descriptor):
     cxxtype = 'PyArrayObject*'
     refcount = 1
-    prerequisites = bytecodecompiler.Type_Descriptor.prerequisites+\
+    module_init_code = 'import_array'
+
+    prerequisites = Type_Descriptor.prerequisites+\
                    ['#include "Numeric/arrayobject.h"']
     dims = 1
     def check(self,s):
@@ -74,19 +76,53 @@ class Vector(bytecodecompiler.Type_Descriptor):
     def outbound(self,s):
         return "(PyObject*)(%s)"%s
     def binopMixed(self,symbol,a,b):
-        return '*((%s*)(%s->data+%s*%s->strides[0]))'%(self.cxxbase,a,b,a),Integer()
+        return '*((%s*)(%s->data+%s*%s->strides[0]))'%(self.cxxbase,a,b,a),Integer
 
 class IntegerVector(Vector):
     typecode = 'PyArray_INT'
     cxxbase = 'int'
 
+class XRange(Type_Descriptor):
+    cxxtype = 'XRange'
+    prerequisites = ['''
+    class XRange {
+    public:
+    XRange(long aLow, long aHigh, long aStep=1)
+    : low(aLow),high(aHigh),step(aStep)
+    {
+    }
+    long low;
+    long high;
+    long step;
+    };''']
+
+# -----------------------------------------------
+# Singletonize the type names
+# -----------------------------------------------
+Integer = Integer()
+Double = Double()
+String = String()
+IntegerVector = IntegerVector()
+XRange = XRange()
+
 typedefs = {
-    IntType: Integer(),
-    FloatType: Double(),
-    StringType: String(),
-    (Numeric.ArrayType,1,'l'): IntegerVector(),
+    IntType: Integer,
+    FloatType: Double,
+    StringType: String,
+    (Numeric.ArrayType,1,'l'): IntegerVector,
+    XRangeType : XRange,
     }
 
+import math
+functiondefs = {
+    (range,(Integer,Integer)):
+    Function_Descriptor(code='XRange(%s)',return_type=XRange),
+
+    (math.sin,(Double,)):
+    Function_Descriptor(code='sin(%s)',return_type=Double),
+    }
+    
+
 
 ##################################################################
 #                      FUNCTION LOOKUP_TYPE                      #
@@ -189,30 +225,41 @@ class accelerate:
 ##################################################################
 #                        CLASS PYTHON2CXX                        #
 ##################################################################
-class Python2CXX(bytecodecompiler.CXXCoder):
+class Python2CXX(CXXCoder):
     def typedef_by_value(self,v):
         T = lookup_type(v)
-        if T not in self.used: self.used.append(T)
+        if T not in self.used:
+            self.used.append(T)
         return T
 
+    def function_by_signature(self,signature):
+        descriptor = functiondefs[signature]
+        if descriptor.return_type not in self.used:
+            self.used.append(descriptor.return_type)
+        return descriptor
+
     def __init__(self,f,signature,name=None):
         # Make sure function is a function
         import types
         assert type(f) == FunctionType
         # and check the input type signature
         assert reduce(lambda x,y: x and y,
-                      map(lambda x: isinstance(x,bytecodecompiler.Type_Descriptor),
+                      map(lambda x: isinstance(x,Type_Descriptor),
                           signature),
                       1),'%s not all type objects'%signature
         self.arg_specs = []
         self.customize = weave.base_info.custom_info()
-        self.customize.add_module_init_code('import_array()\n')
 
-        bytecodecompiler.CXXCoder.__init__(self,f,signature,name)
+        CXXCoder.__init__(self,f,signature,name)
+
         return
 
     def function_code(self):
-        return self.wrapped_code()
+        code = self.wrapped_code()
+        for T in self.used:
+            if T != None and T.module_init_code:
+                self.customize.add_module_init_code(T.module_init_code)
+        return code
 
     def python_function_definition_code(self):
         return '{ "%s", wrapper_%s, METH_VARARGS, %s },\n'%(
diff --git a/weave/bytecodecompiler.py b/weave/bytecodecompiler.py
index 5ee5290d8..17df27e67 100644
--- a/weave/bytecodecompiler.py
+++ b/weave/bytecodecompiler.py
@@ -10,16 +10,33 @@ from types import *
 import string
 import inspect
 
-
 ##################################################################
-#                    CLASS CXXTYPEDESCRIPTION                    #
+#                       CLASS __DESCRIPTOR                       #
 ##################################################################
-class Type_Descriptor:
+class __Descriptor:
     prerequisites = []
     refcount = 0
     def __repr__(self):
         return self.__module__+'.'+self.__class__.__name__
 
+##################################################################
+#                     CLASS TYPE_DESCRIPTOR                      #
+##################################################################
+class Type_Descriptor(__Descriptor):
+    module_init_code = ''
+
+##################################################################
+#                   CLASS FUNCTION_DESCRIPTOR                    #
+##################################################################
+class Function_Descriptor(__Descriptor):
+    def __init__(self,code,return_type,support=''):
+	self.code	= code
+	self.return_type	= return_type
+        self.support = support
+	return
+
+        
+            
 
 haveArgument = 90 # Opcodes greater-equal to this have argument
 byName = {
@@ -1050,11 +1067,10 @@ class CXXCoder(ByteCodeMeaning):
         try:
             F = self.function.func_globals[self.codeobject.co_names[var_num]]
         except:
-            F = getattr(__builtins__,self.codeobject.co_names[var_num])
+            F = __builtins__[self.codeobject.co_names[var_num]]
 
         # For functions, we see if we know about this function
         if callable(F):
-            assert functiondefs.has_key(F),"Function %s is known"%F
             self.push(F,type(F))
             return
 
@@ -1083,6 +1099,32 @@ class CXXCoder(ByteCodeMeaning):
         self.push(native,t)
         return
 
+    def SETUP_LOOP(self,pc,delta):
+        "Pushes a block for a loop onto the block stack. The block spans from the current instruction with a size of delta bytes."
+        return
+
+    def FOR_LOOP(self,pc,delta):
+        "Iterate over a sequence. TOS is the current index, TOS1 the sequence. First, the next element is computed. If the sequence is exhausted, increment byte code counter by delta. Otherwise, push the sequence, the incremented counter, and the current item onto the stack."
+        # Pull off control variable and range info
+        v2,t2 = self.pop()
+        v1,t1 = self.pop()
+        self.emit('for(%s=%s.low; %s<%s.high; %s += %s.step) {'%(
+            v2,v1,v2,v1,v2,v1))
+
+        # Put range back on for assignment
+        self.push(v2,t2)
+        return
+
+    def JUMP_ABSOLUTE(self,pc,target):
+        "Set byte code counter to target."
+        self.emit('}')
+        return
+
+    def POP_BLOCK(self,pc):
+        "Removes one block from the block stack. Per frame, there is a stack of blocks, denoting nested loops, try statements, and such."
+        return
+
+
     ##################################################################
     #                       MEMBER STORE_FAST                        #
     ##################################################################
@@ -1145,30 +1187,21 @@ class CXXCoder(ByteCodeMeaning):
  
         # Pull function object off stack and get descriptor
         f,t = self.pop()
-        descriptor = functiondefs[f]
+        signature = (f,tuple(types))
+        descriptor = self.function_by_signature(signature)
         #self.prerequisites += descriptor['prerequisite']+'\n'
         
-        # Look through descriptors for a match
-        for inputs,outputs,format in descriptor:
-            if inputs == types:
-                break
-        else:
-            raise TypeError,f
-
         # Build a rhs
-        rhs = format%string.join(args,',')
+        rhs = descriptor.code%string.join(args,',')
 
         # Build a statement
-        assert len(outputs) == 1,"Single valued return"
-        assert typedefs.has_key(outputs[0]),"Know about type %s"%outputs[0]
-        description = typedefs[outputs[0]]
         temp = self.unique()
         self.emit('%s %s = %s;\n'%(
-            description.cxxtype,
+            descriptor.return_type.cxxtype,
             temp,
             rhs))
 
-        self.push(temp,outputs[0])
+        self.push(temp,descriptor.return_type)
         return