ROCmSoftwarePlatform
Add bitonic topk
This implements a faster GPU topk.
- Update the ref version of topk to take a parameter for the indices, and also updated to handle any layout.
- Added a gpu bitonic topk version. This will do a bitonic sort per wavefront and then do a partial sort in shared memory to get the final topk values
- Added a
rewrite_topkpass that will split large topk's into 2 operators. This needs the indices to be passed along as they wont be the same for one batch.
Codecov Report
Attention: Patch coverage is 99.09091% with 1 line in your changes missing coverage. Please review.
| Files with missing lines | Patch % | Lines |
|---|---|---|
| src/include/migraphx/op/topk.hpp | 97.62% | 1 Missing :warning: |
Additional details and impacted files
@@ Coverage Diff @@
## develop #3862 +/- ##
===========================================
+ Coverage 92.11% 92.13% +0.02%
===========================================
Files 525 527 +2
Lines 24119 24179 +60
===========================================
+ Hits 22216 22277 +61
+ Misses 1903 1902 -1
| Files with missing lines | Coverage Δ | |
|---|---|---|
| src/include/migraphx/raw_data.hpp | 97.30% <100.00%> (+0.04%) |
:arrow_up: |
| src/include/migraphx/rewrite_topk.hpp | 100.00% <100.00%> (ø) |
|
| src/include/migraphx/shape.hpp | 93.02% <100.00%> (+1.85%) |
:arrow_up: |
| src/include/migraphx/tensor_view.hpp | 100.00% <100.00%> (ø) |
|
| src/rewrite_reduce.cpp | 100.00% <ø> (ø) |
|
| src/rewrite_topk.cpp | 100.00% <100.00%> (ø) |
|
| src/shape.cpp | 92.21% <100.00%> (+0.04%) |
:arrow_up: |
| src/include/migraphx/op/topk.hpp | 98.46% <97.62%> (-0.11%) |
:arrow_down: |
:rocket: New features to boost your workflow:
- :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
![codecov[bot] avatar](https://avatars.githubusercontent.com/in/254?v=4 "Published by a pub.dev verified publisher"){kind=small}
So the stable sorting is 2.5x slower on some config. We can recover some perf by lowering the split threshold.
Most of the perf cost comes from the wavefront sorting as the number of elements to sort gets larger, it seems to increase the register pressure. For now, I think this could be merged, and I can investigate the perf issues in the future.
| Test | Batch | Rate new dca39b |
Rate old 6f50ea |
Diff | Compare |
|---|---|---|---|---|---|
| torchvision-resnet50 | 64 | 3,257.96 | 3,256.95 | 0.03% | :white_check_mark: |
| torchvision-resnet50_fp16 | 64 | 6,937.19 | 6,934.08 | 0.04% | :white_check_mark: |
| torchvision-densenet121 | 32 | 2,456.15 | 2,455.08 | 0.04% | :white_check_mark: |
| torchvision-densenet121_fp16 | 32 | 4,231.74 | 4,220.67 | 0.26% | :white_check_mark: |
| torchvision-inceptionv3 | 32 | 1,627.50 | 1,628.68 | -0.07% | :white_check_mark: |
| torchvision-inceptionv3_fp16 | 32 | 2,717.09 | 2,720.07 | -0.11% | :white_check_mark: |
| cadene-inceptionv4 | 16 | 760.87 | 761.15 | -0.04% | :white_check_mark: |
| cadene-resnext64x4 | 16 | 818.67 | 818.95 | -0.03% | :white_check_mark: |
| slim-mobilenet | 64 | 7,479.06 | 7,474.40 | 0.06% | :white_check_mark: |
| slim-nasnetalarge | 64 | 217.92 | 217.97 | -0.02% | :white_check_mark: |
| slim-resnet50v2 | 64 | 3,459.44 | 3,459.20 | 0.01% | :white_check_mark: |
| bert-mrpc-onnx | 8 | 1,149.68 | 1,150.11 | -0.04% | :white_check_mark: |
| bert-mrpc-tf | 1 | 456.84 | 456.70 | 0.03% | :white_check_mark: |
| pytorch-examples-wlang-gru | 1 | 540.08 | 505.43 | 6.86% | :high_brightness: |
| pytorch-examples-wlang-lstm | 1 | 451.59 | 445.46 | 1.38% | :white_check_mark: |
| torchvision-resnet50_1 | 1 | 816.40 | 813.38 | 0.37% | :white_check_mark: |
| cadene-dpn92_1 | 1 | 430.08 | 427.66 | 0.57% | :white_check_mark: |
| cadene-resnext101_1 | 1 | 393.00 | 393.42 | -0.11% | :white_check_mark: |
| onnx-taau-downsample | 1 | 395.66 | 396.90 | -0.31% | :white_check_mark: |
| dlrm-criteoterabyte | 1 | 32.30 | 32.34 | -0.11% | :white_check_mark: |
| dlrm-criteoterabyte_fp16 | 1 | 51.32 | 51.26 | 0.12% | :white_check_mark: |
| agentmodel | 1 | 10,320.38 | 10,203.47 | 1.15% | :white_check_mark: |
| unet_fp16 | 2 | 58.65 | 58.72 | -0.11% | :white_check_mark: |
| resnet50v1_fp16 | 1 | 1,083.31 | 1,083.19 | 0.01% | :white_check_mark: |
| resnet50v1_int8 | 1 | 1,061.58 | 1,037.60 | 2.31% | :white_check_mark: |
| bert_base_cased_fp16 | 64 | 1,171.24 | 1,171.08 | 0.01% | :white_check_mark: |
| bert_large_uncased_fp16 | 32 | 356.29 | 356.26 | 0.01% | :white_check_mark: |
| bert_large_fp16 | 1 | 195.26 | 196.62 | -0.69% | :white_check_mark: |
| distilgpt2_fp16 | 16 | 2,232.40 | 2,232.90 | -0.02% | :white_check_mark: |
| yolov5s | 1 | 542.07 | 541.57 | 0.09% | :white_check_mark: |
| tinyllama | 1 | 43.98 | 43.85 | 0.29% | :white_check_mark: |
| vicuna-fastchat | 1 | 44.05 | 44.22 | -0.40% | :white_check_mark: |
| whisper-tiny-encoder | 1 | 422.17 | 420.96 | 0.29% | :white_check_mark: |
| whisper-tiny-decoder | 1 | 413.29 | 414.14 | -0.20% | :white_check_mark: |
| llama2_7b | 1 | nan | nan | nan% | :x: |
| qwen1.5-7b | 1 | 23.54 | 23.55 | -0.03% | :white_check_mark: |
| phi3-3.8b | 1 | nan | nan | nan% | :x: |
| mask-rcnn | 1 | 21.74 | 18.69 | 16.30% | :high_brightness: |
| llama3-8b | 1 | 21.74 | 21.73 | 0.06% | :white_check_mark: |
| whisper-large-encoder | 1 | 10.22 | 10.22 | 0.03% | :white_check_mark: |
| whisper-large-decoder | 1 | 100.19 | 96.12 | 4.23% | :high_brightness: |
| mistral-7b | 1 | 23.74 | 23.75 | -0.05% | :white_check_mark: |
| FLUX.1-schnell | 1 | 903.99 | 905.24 | -0.14% | :white_check_mark: |
| nan | nan | nan | nan | nan% | :x: |
This build is not recommended to merge :red_circle:
:x:bert-mrpc-tf: ERROR - check error output
2025-04-19 23:42:53.213069: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1745124178.620777 163647 gpu_device.cc:2022] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 62973 MB memory: -> device: 0, name: AMD Instinct MI250X/MI250, pci bus id: 0000:b3:00.0
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1745124179.443899 163647 mlir_graph_optimization_pass.cc:401] MLIR V1 optimization pass is not enabled
2025-04-19 23:43:07.685723: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.685806: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.686072: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.686126: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.686179: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.686235: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.686293: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
2025-04-19 23:43:07.686345: E external/local_xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc:250] bitcode module is required by this HLO module but was not found at ./opencl.bc
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
error: Failure when generating HSACO
2025-04-19 23:43:07.687596: E tensorflow/compiler/mlir/tools/kernel_gen/tf_framework_c_interface.cc:228] INTERNAL: Generating device code failed.
2025-04-19 23:43:07.688767: W tensorflow/core/framework/op_kernel.cc:1829] UNKNOWN: JIT compilation failed.
2025-04-19 23:43:07.688788: I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
2025-04-19 23:43:07.688799: I tensorflow/core/framework/local_rendezvous.cc:405] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
[[import/loss/output/_21]]
2025-04-19 23:43:07.688837: I tensorflow/core/framework/local_rendezvous.cc:424] Local rendezvous recv item cancelled. Key hash: 11217777527359497193
Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1407, in _do_call
return fn(*args)
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1390, in _run_fn
return self._call_tf_sessionrun(options, feed_dict, fetch_list,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1483, in _call_tf_sessionrun
return tf_session.TF_SessionRun_wrapper(self._session, options, feed_dict,
tensorflow.python.framework.errors_impl.UnknownError: 2 root error(s) found.
(0) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
[[import/loss/output/_21]]
(1) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
0 successful operations.
0 derived errors ignored.
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in
main()
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 324, in main
y_out = sess.run(y, feed_dict=tf_dict)
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 977, in run
result = self._run(None, fetches, feed_dict, options_ptr,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1220, in _run
results = self._do_run(handle, final_targets, final_fetches,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1400, in _do_run
return self._do_call(_run_fn, feeds, fetches, targets, options,
File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/client/session.py", line 1426, in _do_call
raise type(e)(node_def, op, message) # pylint: disable=no-value-for-parameter
tensorflow.python.framework.errors_impl.UnknownError: Graph execution error:
Detected at node 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' defined at (most recent call last):
Node: 'import/bert/embeddings/LayerNorm/moments/SquaredDifference'
Detected at node 'import/bert/embeddings/LayerNorm/moments/SquaredDifference' defined at (most recent call last):
Node: 'import/bert/embeddings/LayerNorm/moments/SquaredDifference'
2 root error(s) found.
(0) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
[[import/loss/output/_21]]
(1) UNKNOWN: JIT compilation failed.
[[{{node import/bert/embeddings/LayerNorm/moments/SquaredDifference}}]]
0 successful operations.
0 derived errors ignored.
Original stack trace for 'import/bert/embeddings/LayerNorm/moments/SquaredDifference'::red_circle:bert_large_uncased_fp16: FAILED: MIGraphX is not within tolerance - check verbose output
:x:llama2_7b: ERROR - check error output
Traceback (most recent call last):
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in
main()
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 205, in main
model = migraphx.parse_onnx(model_name, default_dim_value=batch)
RuntimeError: /src/AMDMIGraphX/src/onnx/onnx_parser.cpp:265: parse_from: PARSE_FROM: Failed reading onnx file: /new-saved-models/llama2_7b/decoder_model.onnx:x:phi3-3.8b: ERROR - check error output
Traceback (most recent call last):
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 340, in
main()
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 205, in main
model = migraphx.parse_onnx(model_name, default_dim_value=batch)
RuntimeError: /src/AMDMIGraphX/src/onnx/onnx_parser.cpp:265: parse_from: PARSE_FROM: Failed reading onnx file: /new-saved-models/phi3-3.8b/model.onnx:red_circle:mask-rcnn: FAILED: MIGraphX is not within tolerance - check verbose output