Performance issue with `np_compat.clip`

[Lasloruhberg]

Hello,

While benchmarking a port of some NumPy code to this library, I observed a noticeable performance drop in my simple tests, with runtimes up to roughly 2x slower than the original NumPy implementation.

Before investigating further, I'd like to check whether this is expected, whether my usage is suboptimal, or whether it might be an issue worth looking into. I've attached a minimal reproducible example below.

Thanks!

Python 3.11
Numpy version: 2.4.6
array_api_compat version: 1.15.0

Output:

4000 250 3 Numpy inplace 0.50s
4000 250 3 Numpy standard 0.66s
4000 250 3 Array_compat 1.44s
4000 250 3 Array_compat inplace 1.32s

import numpy as np
import time
from array_api_compat import array_namespace
from copy import deepcopy

rng=np.random.default_rng(0)

N=4000
s=250
c=3
base_imgs=[]
for i in range(N):
    base_imgs.append(rng.random((s,s,c),dtype=np.float32))
           
#### Test numpy speed vs arraycompat speed
     
  
minimum = 0.001
maximum = 0.997
print("Made img list")

# copy images to allow inplace modification
imgs=deepcopy(base_imgs)
start_time=time.perf_counter()
for data in imgs:
    np.clip(data, minimum, maximum, out=data) 
    np.subtract(data, minimum, out=data) 
    np.true_divide(data, maximum - minimum, out=data)
    np.log(data,out=data)

end_time= time.perf_counter()              
print(N,s,c, f"Numpy inplace {end_time-start_time:.2f}s")

del imgs
imgs=deepcopy(base_imgs)
start_time=time.perf_counter()
 
for data in imgs:
    #xp=array_namespace(data)
    data = np.clip(data, minimum, maximum)
    data = data - minimum
    data = data / (maximum - minimum)
    data = np.log(data)
end_time= time.perf_counter()              
print(N,s,c, f"Numpy standard {end_time-start_time:.2f}s")


del imgs
imgs=deepcopy(base_imgs)
xp=array_namespace(imgs[0])

data=deepcopy(imgs[0])
data = xp.clip(data, minimum, maximum)
data = data - minimum
data = data / (maximum - minimum)
data = xp.log(data)
del data

start_time=time.perf_counter()

for data in imgs:
    
    data = xp.clip(data, minimum, maximum)
    data = data - minimum
    data = data / (maximum - minimum)
    data = xp.log(data)
end_time= time.perf_counter()              
print(N,s,c, f"Array_compat {end_time-start_time:.2f}s")

start_time=time.perf_counter()

for data in imgs:
    
    xp.clip(data, minimum, maximum,out=data)
    xp.subtract(data, minimum, out=data) 
    xp.true_divide(data, maximum - minimum, out=data)
    xp.log(data,out=data)
end_time= time.perf_counter()              
print(N,s,c, f"Array_compat inplace {end_time-start_time:.2f}s")

[ev-br]

A regular way to check performance would be to use line_profiler, or more conveniently via %lprun magic in ipython.

It might be that xp.clip does a bit of work for compatibility with the Array API spec. Don't take my word for it though, do check: any second-guessing performance bottlenecks is almost always off.

[Lasloruhberg]

Yes that seems to be happening, checking with lprun i find

That arraycopmpat does inded spend 70% of the time on the clip, almost 10x as much as numpy.
Is this to be expected/something that is planned to be improved?
Do you have a suggestion around this?

(Numpy: 26 872.1 time units per hit vs
ArrayCompat: 266 455.0 )

Interestingly is the "inplace" clipping is almost as time consuming at the 'normal' one for arraycompat

--------------------  xp transform
Timer unit: 1e-09 s

Total time: 1.56056 s
File: /tmp/ipykernel_4331/2816937264.py
Function: transform at line 13

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    13                                           def transform(imgs, xp, minimum, maximum):
    14      4001     965188.0    241.2      0.1      for data in imgs:
    15      4000 1060335365.0 265083.8     67.9          data = xp.clip(data, minimum, maximum)
    16      4000   73424451.0  18356.1      4.7          data = data - minimum
    17      4000   81551902.0  20388.0      5.2          data = data / (maximum - minimum)
    18      4000   70111366.0  17527.8      4.5          data = data + 1.0
    19      4000  274176368.0  68544.1     17.6          data = xp.log(data)

--------------------  xp transform inplace
Timer unit: 1e-09 s

Total time: 1.51467 s
File: /tmp/ipykernel_4331/2816937264.py
Function: transform_inplace at line 23

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    23                                           def transform_inplace(imgs, xp, minimum, maximum):
    24      4001     785175.0    196.2      0.1      for data in imgs:
    25      4000 1065820129.0 266455.0     70.4          xp.clip(data, minimum, maximum, out=data)
    26      4000   48011336.0  12002.8      3.2          xp.subtract(data, minimum, out=data)
    27      4000   81333614.0  20333.4      5.4          xp.true_divide(data, maximum - minimum, out=data)
    28      4000   41930211.0  10482.6      2.8          xp.add(data, 1., out=data)
    29      4000  276793648.0  69198.4     18.3          xp.log(data, out=data)


--------------------  np transform inplace
Timer unit: 1e-09 s

Total time: 0.546395 s
File: /tmp/ipykernel_4331/2816937264.py
Function: numpy_transform_inplace at line 31

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    31                                           def numpy_transform_inplace(imgs, minimum, maximum):
    32      4001     472546.0    118.1      0.1      for data in imgs:
    33      4000  107488209.0  26872.1     19.7          np.clip(data, minimum, maximum, out=data)
    34      4000   44802958.0  11200.7      8.2          np.subtract(data, minimum, out=data)
    35      4000   79853928.0  19963.5     14.6          np.true_divide(data, maximum - minimum, out=data)
    36      4000   43401206.0  10850.3      7.9          np.add(data, 1., out=data)
    37      4000  270376023.0  67594.0     49.5          np.log(data, out=data)
imgs= deepcopy(base_imgs)



--------------------  np transform
Timer unit: 1e-09 s

Total time: 0.711996 s
File: /tmp/ipykernel_4331/2816937264.py
Function: numpy_transform at line 4

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     4                                           def numpy_transform(imgs,  minimum, maximum):
     5      4001     990212.0    247.5      0.1      for data in imgs:
     6      4000  205662280.0  51415.6     28.9          data = np.clip(data, minimum, maximum)
     7      4000   73283941.0  18321.0     10.3          data = data - minimum
     8      4000   82819565.0  20704.9     11.6          data = data / (maximum - minimum)
     9      4000   73059950.0  18265.0     10.3          data = data + 1.0
    10      4000  276179834.0  69045.0     38.8          data = np.log(data)

[ev-br]

np_compat.clip goes through the "generic" code path,

array-api-compat/src/array_api_compat/common/_aliases.py

Line 375 in 2865123

def clip(

, which does have to do a bit of work to simulate the spec-dictated behavior. It's very much possible that there's a better way to reuse np.clip for numpy-specific version, similar to how torch_compat.clip goes through torch.clamp for some inputs,

array-api-compat/src/array_api_compat/torch/_aliases.py

Line 861 in 2865123

def clip(

I don't think addressing the overhead is on anybody's radar at the moment, so it's up to a motivated contributor to investigate, improve and send a PR (hint hint).

What is the optimization budget, given the constraint of being spec-compliant, I do not know off-hand either. Likely no way to tell without trying.