PERF: avoid an unnecessary copy in np.take with out=#31940
Conversation
|
@mattip since you mentioned this I thought you might be interested :) |
There was a problem hiding this comment.
If the consensus is to just let out be modified, that is like basic guarantee in terms of c++ exception safety. I'd be tempted to add an additional message to the raised error in that case. Like if the error is raised and out= is set, you also say in the error message that out may be modified. cc @seberg @mattip what do you think? It can be done using by getting the exception that would be raised which is IndexError: index 4 is out of bounds and adding a note to it using the add_note method they have (numpy is already 3.12+). Something like
if (npy_fasttake(
dest, src, indices_data, n, m, max_item, nelem, chunk,
clipmode, itemsize, needs_refcounting, src_descr, dst_descr,
axis) < 0) {
if (out != NULL && obj == out) {
PyObject *exc = PyErr_GetRaisedException();
if (exc != NULL) {
PyObject *res = PyObject_CallMethod(exc, "add_note", "s",
"the output array `out` may have been partially modified");
if (res == NULL) {
PyErr_Clear();
}
Py_XDECREF(res);
PyErr_SetRaisedException(exc);
}
}
goto fail;
}There was a problem hiding this comment.
That might be a little bit of an excessive testing for this kind of tiny change IMO.
|
I like the add-note idea, although it applies to many operations (i.e. if we do that, I think it should probably be a central helper and maybe added to some other places). FWIW, it seems OK to me, although I don't care too much either way ( |
e87e7ba to
d27656a
Compare
|
thanks both! I added a compatibility release note documenting that out can now be partially modified if an out of bounds index raises (the relaxed guarantee) then dropped the refcount-asserting test and trimmed the rest down to the essentials so there's no getrefcount asserts left. leak coverage can just be left to the sanitized CI @ikrommyd @seberg I left add_note out for now since i agree it makes more sense as a central helper across the ops that can partially write out :) |
I'd also be happy to open a pr for this too if you think it makes sense |
PR summary
Fixes #28636.
np.take(a, indices, out=out)in the defaultmode="raise"was slower than the same call withoutout.PyArray_TakeFromforced a defensiveENSURECOPYofoutand wrote through a temporary that was copied back on success and only sooutstayed unchanged if an out-of-bounds index raisedthis removes that copy so in
raisemode the result is written directly intooutjust likemode="wrap"/"clip"already do. The overlap-with-selfcopy (arrays_overlap) was kept so aliasing stays correct.there IS a behavior change on an out-of-bounds index since
outmay now be partially written beforeIndexErroris raised (previously it was left unchanged).takestill raises. This proooobably warrants a release note and I'll write it after reviewBefore/after
out=7.66us -> 4.90us(1.56x)99.5us -> 47.3us(2.11x)(200, 64) -> (4000, 64):214us -> 64.6us(3.31x)The no-
outpath is unchanged and the speedup grows with output size (since the removed copy isO(n))AI Disclosure
Used AI to double check my code + find any edge cases I might've missed :)