1
/*
2
 * This file implements assignment from an ndarray to another ndarray.
3
 *
4
 * Written by Mark Wiebe (mwwiebe@gmail.com)
5
 * Copyright (c) 2011 by Enthought, Inc.
6
 *
7
 * See LICENSE.txt for the license.
8
 */
9

10
#define PY_SSIZE_T_CLEAN
11
#include <Python.h>
12

13
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
14
#define _MULTIARRAYMODULE
15
#include <numpy/ndarraytypes.h>
16

17
#include "npy_config.h"
18
#include "npy_pycompat.h"
19

20
#include "convert_datatype.h"
21
#include "methods.h"
22
#include "shape.h"
23
#include "lowlevel_strided_loops.h"
24

25
#include "array_assign.h"
26

27
/*
28
 * Check that array data is both uint-aligned and true-aligned for all array
29
 * elements, as required by the copy/casting code in lowlevel_strided_loops.c
30
 */
31
NPY_NO_EXPORT int
32 1
copycast_isaligned(int ndim, npy_intp const *shape,
33
        PyArray_Descr *dtype, char *data, npy_intp const *strides)
34
{
35
    int aligned;
36
    int big_aln, small_aln;
37

38 1
    int uint_aln = npy_uint_alignment(dtype->elsize);
39 1
    int true_aln = dtype->alignment;
40

41
    /* uint alignment can be 0, meaning not uint alignable */
42 1
    if (uint_aln == 0) {
43
        return 0;
44
    }
45

46
    /*
47
     * As an optimization, it is unnecessary to check the alignment to the
48
     * smaller of (uint_aln, true_aln) if the data is aligned to the bigger of
49
     * the two and the big is a multiple of the small aln. We check the bigger
50
     * one first and only check the smaller if necessary.
51
     */
52 1
    if (true_aln >= uint_aln) {
53
        big_aln = true_aln;
54
        small_aln = uint_aln;
55
    }
56
    else {
57 1
        big_aln = uint_aln;
58 1
        small_aln = true_aln;
59
    }
60

61 1
    aligned = raw_array_is_aligned(ndim, shape, data, strides, big_aln);
62 1
    if (aligned && big_aln % small_aln != 0) {
63 0
        aligned = raw_array_is_aligned(ndim, shape, data, strides, small_aln);
64
    }
65
    return aligned;
66
}
67

68
/*
69
 * Assigns the array from 'src' to 'dst'. The strides must already have
70
 * been broadcast.
71
 *
72
 * Returns 0 on success, -1 on failure.
73
 */
74
NPY_NO_EXPORT int
75 1
raw_array_assign_array(int ndim, npy_intp const *shape,
76
        PyArray_Descr *dst_dtype, char *dst_data, npy_intp const *dst_strides,
77
        PyArray_Descr *src_dtype, char *src_data, npy_intp const *src_strides)
78
{
79
    int idim;
80
    npy_intp shape_it[NPY_MAXDIMS];
81
    npy_intp dst_strides_it[NPY_MAXDIMS];
82
    npy_intp src_strides_it[NPY_MAXDIMS];
83
    npy_intp coord[NPY_MAXDIMS];
84

85 1
    PyArray_StridedUnaryOp *stransfer = NULL;
86 1
    NpyAuxData *transferdata = NULL;
87 1
    int aligned, needs_api = 0;
88 1
    npy_intp src_itemsize = src_dtype->elsize;
89

90 1
    NPY_BEGIN_THREADS_DEF;
91

92 1
    aligned =
93 1
        copycast_isaligned(ndim, shape, dst_dtype, dst_data, dst_strides) &&
94 1
        copycast_isaligned(ndim, shape, src_dtype, src_data, src_strides);
95

96
    /* Use raw iteration with no heap allocation */
97 1
    if (PyArray_PrepareTwoRawArrayIter(
98
                    ndim, shape,
99
                    dst_data, dst_strides,
100
                    src_data, src_strides,
101
                    &ndim, shape_it,
102
                    &dst_data, dst_strides_it,
103
                    &src_data, src_strides_it) < 0) {
104
        return -1;
105
    }
106

107
    /*
108
     * Overlap check for the 1D case. Higher dimensional arrays and
109
     * opposite strides cause a temporary copy before getting here.
110
     */
111 1
    if (ndim == 1 && src_data < dst_data &&
112 1
                src_data + shape_it[0] * src_strides_it[0] > dst_data) {
113 1
        src_data += (shape_it[0] - 1) * src_strides_it[0];
114 1
        dst_data += (shape_it[0] - 1) * dst_strides_it[0];
115 1
        src_strides_it[0] = -src_strides_it[0];
116 1
        dst_strides_it[0] = -dst_strides_it[0];
117
    }
118

119
    /* Get the function to do the casting */
120 1
    if (PyArray_GetDTypeTransferFunction(aligned,
121
                        src_strides_it[0], dst_strides_it[0],
122
                        src_dtype, dst_dtype,
123
                        0,
124
                        &stransfer, &transferdata,
125
                        &needs_api) != NPY_SUCCEED) {
126
        return -1;
127
    }
128

129 1
    if (!needs_api) {
130 1
        NPY_BEGIN_THREADS;
131
    }
132

133 1
    NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
134
        /* Process the innermost dimension */
135 1
        if (stransfer(
136
                dst_data, dst_strides_it[0], src_data, src_strides_it[0],
137
                shape_it[0], src_itemsize, transferdata) < 0) {
138
            goto fail;
139
        }
140 1
    } NPY_RAW_ITER_TWO_NEXT(idim, ndim, coord, shape_it,
141
                            dst_data, dst_strides_it,
142
                            src_data, src_strides_it);
143

144 1
    NPY_END_THREADS;
145 1
    NPY_AUXDATA_FREE(transferdata);
146
    return 0;
147 1
fail:
148 1
    NPY_END_THREADS;
149 1
    NPY_AUXDATA_FREE(transferdata);
150
    return -1;
151
}
152

153
/*
154
 * Assigns the array from 'src' to 'dst, wherever the 'wheremask'
155
 * value is True. The strides must already have been broadcast.
156
 *
157
 * Returns 0 on success, -1 on failure.
158
 */
159
NPY_NO_EXPORT int
160 1
raw_array_wheremasked_assign_array(int ndim, npy_intp const *shape,
161
        PyArray_Descr *dst_dtype, char *dst_data, npy_intp const *dst_strides,
162
        PyArray_Descr *src_dtype, char *src_data, npy_intp const *src_strides,
163
        PyArray_Descr *wheremask_dtype, char *wheremask_data,
164
        npy_intp const *wheremask_strides)
165
{
166
    int idim;
167
    npy_intp shape_it[NPY_MAXDIMS];
168
    npy_intp dst_strides_it[NPY_MAXDIMS];
169
    npy_intp src_strides_it[NPY_MAXDIMS];
170
    npy_intp wheremask_strides_it[NPY_MAXDIMS];
171
    npy_intp coord[NPY_MAXDIMS];
172

173 1
    PyArray_MaskedStridedUnaryOp *stransfer = NULL;
174 1
    NpyAuxData *transferdata = NULL;
175 1
    int aligned, needs_api = 0;
176 1
    npy_intp src_itemsize = src_dtype->elsize;
177

178 1
    NPY_BEGIN_THREADS_DEF;
179

180 1
    aligned =
181 1
        copycast_isaligned(ndim, shape, dst_dtype, dst_data, dst_strides) &&
182 1
        copycast_isaligned(ndim, shape, src_dtype, src_data, src_strides);
183

184
    /* Use raw iteration with no heap allocation */
185 1
    if (PyArray_PrepareThreeRawArrayIter(
186
                    ndim, shape,
187
                    dst_data, dst_strides,
188
                    src_data, src_strides,
189
                    wheremask_data, wheremask_strides,
190
                    &ndim, shape_it,
191
                    &dst_data, dst_strides_it,
192
                    &src_data, src_strides_it,
193
                    &wheremask_data, wheremask_strides_it) < 0) {
194
        return -1;
195
    }
196

197
    /*
198
     * Overlap check for the 1D case. Higher dimensional arrays cause
199
     * a temporary copy before getting here.
200
     */
201 1
    if (ndim == 1 && src_data < dst_data &&
202 1
                src_data + shape_it[0] * src_strides_it[0] > dst_data) {
203 0
        src_data += (shape_it[0] - 1) * src_strides_it[0];
204 0
        dst_data += (shape_it[0] - 1) * dst_strides_it[0];
205 0
        wheremask_data += (shape_it[0] - 1) * wheremask_strides_it[0];
206 0
        src_strides_it[0] = -src_strides_it[0];
207 0
        dst_strides_it[0] = -dst_strides_it[0];
208 0
        wheremask_strides_it[0] = -wheremask_strides_it[0];
209
    }
210

211
    /* Get the function to do the casting */
212 1
    if (PyArray_GetMaskedDTypeTransferFunction(aligned,
213
                        src_strides_it[0],
214
                        dst_strides_it[0],
215
                        wheremask_strides_it[0],
216
                        src_dtype, dst_dtype, wheremask_dtype,
217
                        0,
218
                        &stransfer, &transferdata,
219
                        &needs_api) != NPY_SUCCEED) {
220
        return -1;
221
    }
222

223 1
    if (!needs_api) {
224 1
        NPY_BEGIN_THREADS;
225
    }
226

227 1
    NPY_RAW_ITER_START(idim, ndim, coord, shape_it) {
228
        /* Process the innermost dimension */
229 1
        stransfer(dst_data, dst_strides_it[0], src_data, src_strides_it[0],
230
                    (npy_bool *)wheremask_data, wheremask_strides_it[0],
231
                    shape_it[0], src_itemsize, transferdata);
232 1
    } NPY_RAW_ITER_THREE_NEXT(idim, ndim, coord, shape_it,
233
                            dst_data, dst_strides_it,
234
                            src_data, src_strides_it,
235
                            wheremask_data, wheremask_strides_it);
236

237 1
    NPY_END_THREADS;
238

239 1
    NPY_AUXDATA_FREE(transferdata);
240

241 1
    return (needs_api && PyErr_Occurred()) ? -1 : 0;
242
}
243

244
/*
245
 * An array assignment function for copying arrays, broadcasting 'src' into
246
 * 'dst'. This function makes a temporary copy of 'src' if 'src' and
247
 * 'dst' overlap, to be able to handle views of the same data with
248
 * different strides.
249
 *
250
 * dst: The destination array.
251
 * src: The source array.
252
 * wheremask: If non-NULL, a boolean mask specifying where to copy.
253
 * casting: An exception is raised if the copy violates this
254
 *          casting rule.
255
 *
256
 * Returns 0 on success, -1 on failure.
257
 */
258
NPY_NO_EXPORT int
259 1
PyArray_AssignArray(PyArrayObject *dst, PyArrayObject *src,
260
                    PyArrayObject *wheremask,
261
                    NPY_CASTING casting)
262
{
263 1
    int copied_src = 0;
264

265
    npy_intp src_strides[NPY_MAXDIMS];
266

267
    /* Use array_assign_scalar if 'src' NDIM is 0 */
268 1
    if (PyArray_NDIM(src) == 0) {
269 1
        return PyArray_AssignRawScalar(
270 1
                            dst, PyArray_DESCR(src), PyArray_DATA(src),
271
                            wheremask, casting);
272
    }
273

274
    /*
275
     * Performance fix for expressions like "a[1000:6000] += x".  In this
276
     * case, first an in-place add is done, followed by an assignment,
277
     * equivalently expressed like this:
278
     *
279
     *   tmp = a[1000:6000]   # Calls array_subscript in mapping.c
280
     *   np.add(tmp, x, tmp)
281
     *   a[1000:6000] = tmp   # Calls array_assign_subscript in mapping.c
282
     *
283
     * In the assignment the underlying data type, shape, strides, and
284
     * data pointers are identical, but src != dst because they are separately
285
     * generated slices.  By detecting this and skipping the redundant
286
     * copy of values to themselves, we potentially give a big speed boost.
287
     *
288
     * Note that we don't call EquivTypes, because usually the exact same
289
     * dtype object will appear, and we don't want to slow things down
290
     * with a complicated comparison.  The comparisons are ordered to
291
     * try and reject this with as little work as possible.
292
     */
293 1
    if (PyArray_DATA(src) == PyArray_DATA(dst) &&
294 1
                        PyArray_DESCR(src) == PyArray_DESCR(dst) &&
295 1
                        PyArray_NDIM(src) == PyArray_NDIM(dst) &&
296 1
                        PyArray_CompareLists(PyArray_DIMS(src),
297 1
                                             PyArray_DIMS(dst),
298 1
                                             PyArray_NDIM(src)) &&
299 1
                        PyArray_CompareLists(PyArray_STRIDES(src),
300 1
                                             PyArray_STRIDES(dst),
301
                                             PyArray_NDIM(src))) {
302
        /*printf("Redundant copy operation detected\n");*/
303
        return 0;
304
    }
305

306 1
    if (PyArray_FailUnlessWriteable(dst, "assignment destination") < 0) {
307
        goto fail;
308
    }
309

310
    /* Check the casting rule */
311 1
    if (!PyArray_CanCastTypeTo(PyArray_DESCR(src),
312
                                PyArray_DESCR(dst), casting)) {
313 1
        npy_set_invalid_cast_error(
314
                PyArray_DESCR(src), PyArray_DESCR(dst), casting, NPY_FALSE);
315
        goto fail;
316
    }
317

318
    /*
319
     * When ndim is 1 and the strides point in the same direction,
320
     * the lower-level inner loop handles copying
321
     * of overlapping data. For bigger ndim and opposite-strided 1D
322
     * data, we make a temporary copy of 'src' if 'src' and 'dst' overlap.'
323
     */
324 1
    if (((PyArray_NDIM(dst) == 1 && PyArray_NDIM(src) >= 1 &&
325 1
                    PyArray_STRIDES(dst)[0] *
326 1
                            PyArray_STRIDES(src)[PyArray_NDIM(src) - 1] < 0) ||
327 1
                    PyArray_NDIM(dst) > 1 || PyArray_HASFIELDS(dst)) &&
328 1
                    arrays_overlap(src, dst)) {
329
        PyArrayObject *tmp;
330

331
        /*
332
         * Allocate a temporary copy array.
333
         */
334 1
        tmp = (PyArrayObject *)PyArray_NewLikeArray(dst,
335
                                        NPY_KEEPORDER, NULL, 0);
336 1
        if (tmp == NULL) {
337
            goto fail;
338
        }
339

340 1
        if (PyArray_AssignArray(tmp, src, NULL, NPY_UNSAFE_CASTING) < 0) {
341 0
            Py_DECREF(tmp);
342
            goto fail;
343
        }
344

345
        src = tmp;
346
        copied_src = 1;
347
    }
348

349
    /* Broadcast 'src' to 'dst' for raw iteration */
350 1
    if (PyArray_NDIM(src) > PyArray_NDIM(dst)) {
351 1
        int ndim_tmp = PyArray_NDIM(src);
352 1
        npy_intp *src_shape_tmp = PyArray_DIMS(src);
353 1
        npy_intp *src_strides_tmp = PyArray_STRIDES(src);
354
        /*
355
         * As a special case for backwards compatibility, strip
356
         * away unit dimensions from the left of 'src'
357
         */
358 1
        while (ndim_tmp > PyArray_NDIM(dst) && src_shape_tmp[0] == 1) {
359 1
            --ndim_tmp;
360 1
            ++src_shape_tmp;
361 1
            ++src_strides_tmp;
362
        }
363

364 1
        if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
365
                    ndim_tmp, src_shape_tmp,
366
                    src_strides_tmp, "input array",
367
                    src_strides) < 0) {
368
            goto fail;
369
        }
370
    }
371
    else {
372 1
        if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
373 1
                    PyArray_NDIM(src), PyArray_DIMS(src),
374 1
                    PyArray_STRIDES(src), "input array",
375
                    src_strides) < 0) {
376
            goto fail;
377
        }
378
    }
379

380
    /* optimization: scalar boolean mask */
381 1
    if (wheremask != NULL &&
382 1
            PyArray_NDIM(wheremask) == 0 &&
383 1
            PyArray_DESCR(wheremask)->type_num == NPY_BOOL) {
384 1
        npy_bool value = *(npy_bool *)PyArray_DATA(wheremask);
385 1
        if (value) {
386
            /* where=True is the same as no where at all */
387
            wheremask = NULL;
388
        }
389
        else {
390
            /* where=False copies nothing */
391
            return 0;
392
        }
393
    }
394

395 1
    if (wheremask == NULL) {
396
        /* A straightforward value assignment */
397
        /* Do the assignment with raw array iteration */
398 1
        if (raw_array_assign_array(PyArray_NDIM(dst), PyArray_DIMS(dst),
399 1
                PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
400 1
                PyArray_DESCR(src), PyArray_DATA(src), src_strides) < 0) {
401
            goto fail;
402
        }
403
    }
404
    else {
405
        npy_intp wheremask_strides[NPY_MAXDIMS];
406

407
        /* Broadcast the wheremask to 'dst' for raw iteration */
408 1
        if (broadcast_strides(PyArray_NDIM(dst), PyArray_DIMS(dst),
409 1
                    PyArray_NDIM(wheremask), PyArray_DIMS(wheremask),
410 1
                    PyArray_STRIDES(wheremask), "where mask",
411
                    wheremask_strides) < 0) {
412
            goto fail;
413
        }
414

415
        /* A straightforward where-masked assignment */
416
         /* Do the masked assignment with raw array iteration */
417 1
        if (raw_array_wheremasked_assign_array(
418 1
                PyArray_NDIM(dst), PyArray_DIMS(dst),
419 1
                PyArray_DESCR(dst), PyArray_DATA(dst), PyArray_STRIDES(dst),
420 1
                PyArray_DESCR(src), PyArray_DATA(src), src_strides,
421 1
                PyArray_DESCR(wheremask), PyArray_DATA(wheremask),
422
                        wheremask_strides) < 0) {
423
            goto fail;
424
        }
425
    }
426

427 1
    if (copied_src) {
428 1
        Py_DECREF(src);
429
    }
430
    return 0;
431

432 1
fail:
433 1
    if (copied_src) {
434 0
        Py_DECREF(src);
435
    }
436
    return -1;
437
}

Read our documentation on viewing source code .

Loading