Commits

Konstantin Lopuhin committed f347a1a

make indices virtualizable

Comments (0)

Files changed (2)

pypy/module/micronumpy/iter.py

         return self.index
 
 class MultiDimViewIterator(ConcreteArrayIterator):
+    _virtualizable_ = ['indexes[*]'] # 'shape[*]', 'strides[*]', 'backstrides[*]']
+
     def __init__(self, array, start, strides, backstrides, shape):
+        self = jit.hint(self, access_directly=True, fresh_virtualizable=True)
         self.indexes = [0] * len(shape)
         self.array = array
         self.shape = shape

pypy/module/micronumpy/loop.py

 argmax = _new_argmin_argmax('max')
 
 # note that shapelen == 2 always
-dot_driver = jit.JitDriver(name = 'numpy_dot',
-                           greens = ['dtype'],
-                           reds = 'auto')
+dot_driver = jit.JitDriver(
+        name = 'numpy_dot',
+        greens = ['dtype'],
+        virtualizables = ['outi', #'lefti', 'righti' # TODO
+            ],
+        reds = ['outi', 'lefti', 'righti', 'result'])
 
 def multidim_dot(space, left, right, result, dtype, right_critical_dim):
     ''' assumes left, right are concrete arrays
     lefti = left.create_dot_iter(broadcast_shape, left_skip)
     righti = right.create_dot_iter(broadcast_shape, right_skip)
     while not outi.done():
-        dot_driver.jit_merge_point(dtype=dtype)
+        dot_driver.jit_merge_point(dtype=dtype,
+                outi=outi, lefti=lefti, righti=righti, result=result)
         lval = lefti.getitem().convert_to(space, dtype)
         rval = righti.getitem().convert_to(space, dtype)
         outval = outi.getitem().convert_to(space, dtype)