[Perf] Reduction over rows of a multi dimension array takes a while

Start with the findDistances reproducer and the af::sqrt(af::sum(af::pow(YBatch, 2), 1)) line, then run it with the stated CUDA setup while using ArrayFire timing guidance. Compare the reduction timing with the reported CuPy baseline and investigate the multi-dimensional reduction path. Done means the slowdown is reproduced and the CUDA reduction performance is improved or its cause is documented.

Basically I'm trying to find the norm of 3D array, over the 2nd and 3rd dimensions, and it's taking much longer than expected.

Description

I'm trying to rewrite code that was previously written in CuPy ton Arrayfire in a super speedy manner.

I'm using CUDA.

Doing this with CuPy takes just 0.9 seconds for the whole function to complete.

I used official installers.

And yes it can be produced reliably - meaning it happens every time.

Reproducible Code

inline af::array findDistances(af::array &X, af::array &A, af::array &B, float alpha = 1.2) {
        int k = A.dims(1) / 2;
        int m = B.dims(1);

        int n = X.dims(0);
        int d = X.dims(1);
        int D = B.dims(0) / 2;

        int batchSize = findDistanceBatchSize(alpha, n, d, k, m); // Comes out to 20

        af::array distances(n, 2 * k * m, af::dtype::f32);
        af::array ABatch(batchSize, 2 * k, A.type());
        af::array BBatch(batchSize, m, B.type());
        af::array XBatch(batchSize, 2 * k, m, d, X.type());
        af::array XBatchAdj(batchSize, 2 * k * m, d,
                            X.type()); // This is very large, around 7gb. Possible to do this without explicitly allocating the memory?
        af::array XSubset(batchSize, d, X.type());
        af::array XSubsetReshaped = af::constant(0, XBatchAdj.dims(), XBatchAdj.type());
        af::array YBatch = af::constant(0, XBatchAdj.dims(), XBatchAdj.type());

        for (int i = 0; i < n; i += batchSize) {
            int maxBatchIdx = i + batchSize - 1;
            ABatch = A(af::seq(i, maxBatchIdx), af::span);

            BBatch = B(ABatch, af::span);

            BBatch = af::moddims(BBatch, batchSize, 2 * k, m);

            XBatch = X(BBatch, af::span);

            XBatchAdj = af::moddims(XBatch, batchSize, 2 * k * m, d);

            XSubset = X(af::seq(i, maxBatchIdx), af::span);

            XSubsetReshaped = moddims(XSubset, batchSize, 1, d); // Insert new dim

            YBatch = XBatchAdj - XSubsetReshaped;

//            distances(af::seq(i, maxBatchIdx), af::span) =
            af::sqrt(af::sum(af::pow(YBatch, 2), 1)); // It gets hung up on this line. The assignment above breaks the code, so just to get an idea of runtime, I just put it on a new line
        }

        return distances;
    }

System Information

ArrayFire Version: 2.9.0
Device: RTX 3090. Running CUDA 12.6
Operating System: Ubuntu 20.04
Driver version: (nvidia driver): 560.28.03

Checklist

• I have read timing ArrayFire Yep