`torch.topk` is two order of magnitude too slow on xpu (flex 170)

[fcharras]

Describe the bug

This snippet

out = torch.topk(X, 100)[0][0]
print(out)

with X being uniform randomly generated data of size 2**29 executes in about 4.25s on Flex 170 GPU, but it should be closer to 0.1s (100ms).

Versions

PyTorch version: 1.13.0a0+git6c9b55e
PyTorch CXX11 ABI: Yes
IPEX version: 1.13.120+xpu
IPEX commit: c2a37012e
Build type: Release

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: N/A
IGC version: 2023.1.0 (2023.1.0.20230320)
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.10 | packaged by conda-forge | (main, Mar 24 2023, 20:08:06) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.14.0-1047-oem-x86_64-with-glibc2.31
Is XPU available: True
DPCPP runtime version: N/A
MKL version: N/A
GPU models and configuration: 
[0] _DeviceProperties(name='Intel(R) Graphics [0x56c0]', platform_name='Intel(R) Level-Zero', dev_type='gpu, support_fp64=0, total_memory=13535MB, max_compute_units=512)
Intel OpenCL ICD version: 22.28.23726.1+i419~u20.04
Level Zero version: 1.3.23726.1+i419~u20.04

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 57 bits virtual
CPU(s):                          40
On-line CPU(s) list:             0-39
Thread(s) per core:              2
Core(s) per socket:              20
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           108
Model name:                      Intel(R) Xeon(R) D-2796NT CPU @ 2.00GHz
Stepping:                        1
Frequency boost:                 enabled
CPU MHz:                         800.000
CPU max MHz:                     2001.0000
CPU min MHz:                     800.0000
BogoMIPS:                        4000.00
Virtualization:                  VT-x
L1d cache:                       960 KiB
L1i cache:                       640 KiB
L2 cache:                        25 MiB
L3 cache:                        30 MiB
NUMA node0 CPU(s):               0-39
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT vulnerable
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 invpcid_single intel_ppin ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect wbnoinvd dtherm ida arat pln pts avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid fsrm md_clear pconfig flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] intel-extension-for-pytorch==1.13.120+xpu
[pip3] numpy==1.24.3
[pip3] torch==1.13.0a0+git6c9b55e
[conda] intel-extension-for-pytorch 1.13.120+xpu             pypi_0    pypi
[conda] mkl                       2023.1.0            intel_46342    intel
[conda] mkl-dpcpp                 2023.1.0            intel_46342    intel
[conda] numpy                     1.24.3          py310ha4c1d20_0    conda-forge
[conda] torch                     1.13.0a0+git6c9b55e          pypi_0    pypi

[fcharras]

From what I've gathered from oneapi libraries I suspect this is because the xpu topk is backed by an introselect-based search algorithm, but gpus prefer bucket, radix or bitonic-based search.

[gujinghui]

I recall we already have fix for big size with radix sort in topk.
You can try our master code base?

[fcharras]

Keyword search do not suggest that there is such an implementation in the master branch in this repo, but maybe there is in the intel extension for tensorflow. Will try both.

About the compile bundle, DPCPPROOT and MKLROOT are supposed to hint at the "latest" folder in the compiler and mkl folders of the oneapi basekit file tree ?

[fcharras]

Intel extension for tensorflow has the same issue.

I didn't manage to compile the pytorch extension from source, the compile bundle fails at compiling pytorch and the cloud filesystem I have access to is awfully slow, which makes any debug very painful. I'll wait for the next release, or would happily try nightly binaries if you have any to share.

[fcharras]

So I managed to install from compile_bundle.sh !

I see that indeed topk is about 10-20 times faster on xpu-master which is much in line with GPU performances. It would be interesting to compare with cuda-based implementation with GPUs with similar compute power, I think it's still behind by some factor.

Also, the optimized version on xpu-maste doesn't seem to work with an input of size 2**29 and k=100, I get PI_ERROR_OUT_OF_RESOURCES, it is somewhat surprising from radix sorting. But it works with size 2**28.

[fcharras]

@gujinghui Out of curiosity, would you have a link to the repository where the topk from xpu-master is implemented ?

[fengyuan14]

We have concurrency algorithm optimization in Topk.

1. Re a small shape case, we are using grouped-radix algorithm.
2. Re a large shape case, we are leveraging sort kernel.

And your case hits the path 2. Route to sort kernel. Our sort kernel has heuristic to choose radix and merge algorithm. As for Bitonic, It assumes additional memory consumption and operators for padding (exponent of 2). It is good for specific shape, but still meet pretty unfriendly case like 2049. Evaluate complex heuristic, we don't adopt Bitonic.

Basing on existing concurrency algorithm implementation, I suppose the performance should not be that bad. BTW, we still have some performance gap compared against CUDA impl (PyTorch CUDA is using CUB). We are optimizing.

So I are reproducing your data to find out what happened. And so far, I can get ~200ms with a 2 * * 28 tensor.

Can you try 2 * * 28?

As for 2 * * 29, I got some runtime error still. I am checking. And keep you in loop if any progress.

[fcharras]

My original post in this issue uses 1.13.120+xpu binaries. It works with 2**30, 2**29 but is about 10-100x slower than it should.

I then used the compile_bundle.sh to install from latest commit on xpu-master branch. With this version it doesn't work with 2**29, 2**30, but I have comparable performance from what you describe with 2**28 which is reasonably fast.

It's good to know you're working on it. It's very useful for gpu-accelerated search algorithms. I'm currently trying to implement a KNN algorithm that can run on all GPU brands, the performance of the topk is the main bottleneck, so I'll look closely at performance improvements.

Agree that bitonic is hard to use in practice. Regarding radix-based, are sycl-based implementations in intel extension able to use similar intrinsics such as ballot / popc as it does for cuda/rocm in pytorch for histogram compute ?

I think it would be interesting to compare to performance on nividia GPU with similar compute power, I think the gap is not yet 100% closed but with current state of xpu-master it seems to be on the way.

edit: do you have information about the ETA of the next release ?

[gujinghui]

My original post in this issue uses 1.13.120+xpu binaries. It works with 2**30, 2**29 but is about 10-100x slower than it should.

I then used the compile_bundle.sh to install from latest commit on xpu-master branch. With this version it doesn't work with 2**29, 2**30, but I have comparable performance from what you describe with 2**28 which is reasonably fast.

It's good to know you're working on it. It's very useful for gpu-accelerated search algorithms. I'm currently trying to implement a KNN algorithm that can run on all GPU brands, the performance of the topk is the main bottleneck, so I'll look closely at performance improvements.

Agree that bitonic is hard to use in practice. Regarding radix-based, are sycl-based implementations in intel extension able to use similar intrinsics such as ballot / popc as it does for cuda/rocm in pytorch for histogram compute ?

I think it would be interesting to compare to performance on nividia GPU with similar compute power, I think the gap is not yet 100% closed but with current state of xpu-master it seems to be on the way.

edit: do you have information about the ETA of the next release ?

Thanks for your attention.
The next release should be happened in this 3rd quarter.

[fengyuan14]

Regarding failure with 2 * * 29 or more large size, it is out of memory on Intel Arc, due to underlying sort kernel. Sort kernel needs two inputs, val, key, and two outputs, sorted_val, sorted_key. As for the case 2 * * 29, input val is 2GB, input key is 4GB (int64), and corresponding sorted_val and sorted_key are in same size. So the total memory consumption is 12GB.
Same implementation is in stock PyTorch CUDA backend, we may get OOM on same level NV card.

And the hard-coded heuristic comes from pytorch/pytorch#68632

BTW, CUDA will deprecate the implementation, when workable CUB supports scan by key,

We will follow up counterpart implementation (Including SYCL Intel extension forward progress support).

Regarding CUDA warp ballot like, SYCL Intel oneAPI extension has similar intrinsic support,

The scope could be work group level or sub group level (warp like).

[fcharras]

Thanks a lot for this comprehensive answer. The progress is very impressive and exciting.

[fcharras]

Closing the issue since the question is answered (fix will be shipped in the next release, expect 3rd quarter).