Improve _reduce_dims
@@ -138,39 +138,66 @@ cpdef str _get_kernel_params(tuple params, tuple args_info):
cpdef tuple _reduce_dims(list args, tuple params, tuple shape):
""" Remove contiguous stride to optimize CUDA kernel."""
+ cdef ndarray arr
+
+ if len(shape) <= 1 or len(args) == 0:
+ return shape
+
+ if len(args) == 1: # fast path for reduction
+ a = args[0]
+ if (<ParameterInfo>params[0]).raw or not isinstance(a, ndarray):
+ return shape
+ arr = a
+ arr = arr.reduced_view()
+ if arr is a:
+ return shape
+ else:
+ args[0] = arr
+ return arr.shape
+ return _reduced_view_core(args, params, shape)
+
+
+cdef tuple _reduced_view_core(list args, tuple params, tuple shape):
cdef int i, ax, last_ax, ndim
- cdef Py_ssize_t total_size
+ cdef Py_ssize_t x, total_size
cdef vector.vector[Py_ssize_t] vecshape, newshape, newstrides
cdef vector.vector[int] array_indexes, axes
- cdef vector.vector[vector.vector[Py_ssize_t]] args_strides
+ cdef vector.vector[int] strides_indexes
cdef ParameterInfo p
cdef ndarray arr
- cdef bint flag
ndim = len(shape)
- if ndim <= 1:
- return shape
-
+ array_indexes.reserve(len(args))
+ strides_indexes.reserve(len(args))
for i in range(len(args)):
p = params[i]
if not p.raw and isinstance(args[i], ndarray):
array_indexes.push_back(i)
arr = args[i]
if not arr._c_contiguous:
- args_strides.push_back(arr._strides)
+ strides_indexes.push_back(i)
- if args_strides.size() == 0:
+ if array_indexes.size() == 0:
+ return shape
+
+ if strides_indexes.size() == 0:
# The input arrays are all c_contiguous
- total_size = internal.prod(shape)
+ i = array_indexes[0]
+ arr = args[i]
+ total_size = arr.size
newshape.assign(<Py_ssize_t>1, total_size)
+ newstrides.resize(1)
for i in array_indexes:
arr = args[i]
- newstrides.assign(<Py_ssize_t>1, arr.dtype.itemsize)
+ newstrides[0] = arr.dtype.itemsize
# TODO(niboshi): Confirm update_x_contiguity flags
args[i] = arr._view(newshape, newstrides, False, True)
return total_size,
- vecshape = shape
+ axes.reserve(ndim)
+ vecshape.reserve(ndim)
+ for x in shape:
+ vecshape.push_back(x)
last_ax = -1
for ax in range(ndim):
if vecshape[ax] == 1:
@@ -178,8 +205,9 @@ cpdef tuple _reduce_dims(list args, tuple params, tuple shape):
if last_ax < 0:
last_ax = ax
continue
- for st in args_strides:
- if st[ax] * vecshape[ax] != st[last_ax]:
+ for i in strides_indexes:
+ arr = args[i]
+ if arr._strides[ax] * vecshape[ax] != arr._strides[last_ax]:
axes.push_back(last_ax)
break
else:
@@ -190,6 +218,8 @@ cpdef tuple _reduce_dims(list args, tuple params, tuple shape):
if <int>axes.size() == ndim:
return shape
+ newshape.reserve(axes.size())
+ newstrides.reserve(axes.size())
for ax in axes:
newshape.push_back(vecshape[ax])
for i in array_indexes:
