Optimizer concatenate for c_contiguous arrays
@@ -1984,27 +1984,28 @@ cpdef ndarray _repeat(ndarray a, repeats, axis=None):
cpdef ndarray concatenate(tup, axis, shape, dtype):
- cdef ndarray ret
- cdef int base
- cdef int cum
+ cdef ndarray a, x, ret
+ cdef int i, base, cum
+ cdef bint all_same_type, all_one_and_contiguous
ret = ndarray(shape, dtype=dtype)
all_same_type = True
- all_one = True
+ all_one_and_contiguous = True
+ dtype = tup[0].dtype
for a in tup:
- all_same_type &= a.dtype == tup[0].dtype
- all_one &= a._shape[axis] == 1
+ all_same_type &= a.dtype == dtype
+ all_one_and_contiguous &= (a._c_contiguous and a._shape[axis] == 1)
if all_same_type:
base = internal.prod_ssize_t(shape[axis+1:])
x = array([a.data.ptr for a in tup])
- x_strides = array([a.strides for a in tup], 'i')
- if all_one:
+ if all_one_and_contiguous:
_concatenate_kernel_one(
- x, axis, x_strides, array(shape, 'i'), base, ret, size=ret.size)
+ x, axis, len(shape), base, ret, size=ret.size)
else:
+ x_strides = array([a.strides for a in tup], 'i')
cum = 0
cum_sizes = numpy.empty(len(tup), 'i')
for i, a in enumerate(tup):
@@ -2026,29 +2027,28 @@ cpdef ndarray concatenate(tup, axis, shape, dtype):
cdef _concatenate_kernel_one = ElementwiseKernel(
- '''raw P x, int32 axis, raw int32 x_strides, raw int32 shape,
- int32 base''',
+ '''raw P x, int32 axis, int32 ndim, int32 base''',
'raw T y',
'''
+ const int* shape = y.shape();
int n = shape[axis];
int array_ind = i / base % n;
- char* ptr = reinterpret_cast<char*>(x[array_ind]);
+ T* ptr = reinterpret_cast<T*>(x[array_ind]);
int ind_rest = i;
- for (int j = shape.size() - 1; j >= 0; --j) {
- int ind[] = {array_ind, j};
- int next_ind = ind_rest / shape[j];
- int offset;
- if (j == axis) {
- offset = 0;
- } else {
- offset = ind_rest - next_ind * shape[j];
+ int stride = 1;
+ for (int j = ndim - 1; j >= 0; --j) {
+ int dim_size = shape[j];
+ int next_ind = ind_rest / dim_size;
+ if (j != axis) {
+ int offset = ind_rest - next_ind * dim_size;
+ ptr += stride * offset;
+ stride *= dim_size;
}
- ptr += x_strides[ind] * offset;
ind_rest = next_ind;
}
- y[i] = *reinterpret_cast<T*>(ptr);
+ y[i] = *ptr;
''',
'cupy_concatenate_one'
)
