xt::argmax is about 3 times slower than np.argmax

[keeper]

Hi, I benchmark xt::argmax() recently, and found it takes around 110 ms (100 times average), while NumPy counter part only takes 35 ms.

The input is a (512, 768, 24) cube, and I argmax it on axis 0, which produce a (768, 24) matrix:

xt::xarray<double> cube = xt::random::randn<double>({512, 768, 24});
xt::argmax(cube, 0);

cube = np.random.randn(512, 768, 24)
np.argmax(cube, 0)

Compilable C++ source: https://gist.github.com/keeper/f065294442de562a72e41c47f2e18eff
Python scripts: https://gist.github.com/keeper/a664b6836939ee34b9d12eccdef087bd

Compiler version: g++ (Ubuntu 9.3.0-10ubuntu2) 9.3.0
Compile command: g++ -O3 -o main main.cpp

[emmenlau]

Is it possible that you did not activate the SIMD instructions in xtensor? See https://xtensor.readthedocs.io/en/latest/build-options.html#external-dependencies

[keeper]

@emmenlau
Apparently I did it in the wrong way, I used add_definitions(-DXTENSOR_ENABLE_XSIMD) (don't remember where I saw this) in my main project's CMakeList.txt and I thought SIMD was enabled.

However after adding set(XTENSOR_USE_XSIMD 1) properly my code doesn't compile, I'll figure it out it first then update the result. Thanks for the information!

[keeper]

After enable SIMD my main project doesn't output correct result at the moment, however the part using argmax() seems fine and the performance has no improve.

[tdegeus]

A small comment on you quest to compile with simd : If you use a modern version of CMake you can also use target_link_libraries(main PRIVATE xtensor xtensor::optimize xtensor::use_xsimd) for your executable.

[tdegeus]

Your result is strange but interesting. Could you try with xt::xtensor<double,3>?

[keeper]

xt::xtensor<double ,3> is slightly faster (~10 ms), but still slower than NumPy.

Update:
I've fix the problem with SIMD, and the main project passed the test with SIMD enabled, and argmax()'s performance remains the same.

[keeper]

I sneak some timers into xsort.hpp, and found the bottleneck might be transpose() called here: https://github.com/xtensor-stack/xtensor/blob/master/include/xtensor/xsort.hpp#L873

The timer reports total arg_func_impl() called in argmax() takes 92ms, and in arg_func_impl() :

transpose takes 80 ms
call_over_leading_axis (not leading_axis) takes 10 ms

[tdegeus]

Thats very useful info! Intuitively I would say that it might not be needed, as we could you the strides of the input and the output. But maybe it would be good to hear from the others @JohanMabille @SylvainCorlay @wolfv

[emmenlau]

Very interesting indeed! Thanks @keeper, nice work.

[JohanMabille]

Indeed, it should be able to operate on the transposed view directly (at least we should try and see what happens :D). Thanks @keeper for this detailed investigation, much appreciated!