arrayfire / arrayfire/arrayfire
[Perf] Reduction over rows of a multi dimension array takes a while
- Ngôn ngữ chính
- C++
- Star
- 4.9k
- Fork
- 555
- Merge trung bình
- 1 giờ 24 phút
- Pull request đã merge (30 ngày)
- 1
Mô tả
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
-----------------
```cpp
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
---------
- [x] I have read [timing ArrayFire](http://arrayfire.org/docs/timing.htm) Yep
Hướng dẫn đóng góp
Đánh giá
Issue này chưa được đánh giá.