Skip to content

Remove concat->reshapes->slice - #5183

Merged
causten merged 15 commits into
developfrom
slice-reshapes-concat
Sep 2, 2026
Merged

Remove concat->reshapes->slice#5183
causten merged 15 commits into
developfrom
slice-reshapes-concat

Conversation

@pfultz2

@pfultz2 pfultz2 commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator

Motivation

Models that repack tensors (such as qkv weights) into a concat which an op decomposition then re-slices leave behind a concat->reshapes/transposes->slice chain that copies data at runtime for no reason. When a slice selects exactly one input of the concat, the slice can be rewritten directly on that input with the equivalent view ops. This also allows the kv-cache attention matcher in fuse_attention to drop its requirement that the first gemm input be a slice, since the slice is now removed before fusion.

Technical Details

  • Adds a find_slice_reshaped_concat matcher to simplify_reshapes that matches slice(view_ops...(concat)), where the view ops are reshape, reshape_lazy, squeeze, unsqueeze, flatten, or transpose. Direct slice(concat) is left to the existing find_concat_slice.
  • The element mapping of the view chain is tracked with shape_transform_descriptor rather than strides, so the transformation works with non-standard shapes (such as a layout-preserving concat output) and with non-unit dimensions before the concat axis.
  • Adds a shape_transform_descriptor::slice_axis(axis, slice_axes, starts, ends) member that restricts a source axis to the range selected by slicing the output dimensions. The subdimensions of the axis ordered by their split lineage form a mixed-radix decomposition, so the slice is accepted only when it selects one contiguous range [start, end) of that axis alone, which is returned to the caller. Unit subdimensions are renumbered to output order when possible to avoid generating a gratuitous transpose.
  • The pass then matches the selected range against the concat input offsets; when it covers exactly one segment, the slice is replaced by the ops from shape_transform_descriptor::generate applied to that input.
  • shape_transform_descriptor::apply now handles reshape_lazy the same as reshape.
  • Adds unit tests for the pass, the slice_axis member, and reshape_lazy descriptor tracking.

Changelog Category

Add a CHANGELOG.md entry for any option other than Not Applicable

    • Added: New functionality.
    • Changed: Changes to existing functionality.
    • Removed: Functionality or support that has been removed. (Compared to a previous release)
    • Optimized: Component performance that has been optimized or improved.
    • Resolved Issues: Known issues from a previous version that have been resolved.
    • Not Applicable: This PR is not to be included in the changelog.

Follow the LLVM AI Tool Use Policy for contributions using AI.

@pfultz2
pfultz2 requested review from CharlieL7, TedThemistokleous, bdevorem and kahmed10 and a balanced review from Copilot and removed request for CharlieL7 August 24, 2026 17:46
@pfultz2
pfultz2 marked this pull request as ready for review August 24, 2026 17:48
@pfultz2
pfultz2 requested a review from causten as a code owner August 24, 2026 17:48

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Optimizes concat/view/slice chains by forwarding aligned slices directly to concat inputs.

Changes:

  • Adds descriptor-based slice-axis tracking, including reshape_lazy.
  • Adds the new simplification matcher and broadens KV-cache attention matching.
  • Adds descriptor and reshape simplification tests.

Review performed as a single pass without agent fan-out; tests were not executed.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
src/simplify_reshapes.cpp Implements concat/view/slice elimination.
src/shape_transform_descriptor.cpp Adds slice-axis mapping and lazy reshape handling.
src/include/migraphx/shape_transform_descriptor.hpp Exposes the slice-axis API.
src/fuse_attention.cpp Relaxes the KV-cache query matcher.
test/simplify_reshapes_test.cpp Tests aligned, misaligned, leading-dimension, and nonstandard cases.
test/shape_transform_descriptor.cpp Tests slicing and reshape_lazy tracking.

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread src/fuse_attention.cpp Outdated
auto gemm1 = match::name("dot")(match::arg(0)(queries), match::arg(1)(k_transpose));
auto gemm1_maybe_cvt = match::skip(match::name("convert"))(gemm1);
auto scale = match::name("mul")(match::any_arg(0, 1)(gemm1_maybe_cvt));
auto gemm1 = match::name("dot")(match::arg(1)(k_transpose));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it would be a good idea to add a test here

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added the test.

Comment thread src/simplify_reshapes.cpp
Comment on lines +857 to +858
if(concat_ins->get_shape().dynamic())
return;
Comment thread test/simplify_reshapes_test.cpp Outdated
Comment on lines +5953 to +5955
// The slices do not align with the segment boundary, so the concat must remain
EXPECT(
std::any_of(m1.begin(), m1.end(), [](const auto& ins) { return ins.name() == "concat"; }));
@gh-app-migraphx-bot-pr-write

Copy link
Copy Markdown
Test Batch New Rate (1d9c82) Old Rate (3a503c)* Diff Status
torchvision-resnet50 64 3,299.99 3,264.92 1.07%
torchvision-resnet50_fp16 64 7,882.12 7,548.67 4.42%
torchvision-densenet121 32 2,492.00 2,483.99 0.32%
torchvision-densenet121_fp16 32 5,015.52 5,004.24 0.23%
torchvision-inceptionv3 32 2,059.07 2,058.51 0.03%
torchvision-inceptionv3_fp16 32 4,443.80 4,416.99 0.61%
cadene-inceptionv4 16 817.29 820.61 -0.40%
cadene-resnext64x4 16 785.01 782.78 0.28%
slim-mobilenet 64 8,373.47 8,386.36 -0.15%
slim-nasnetalarge 64 229.34 228.86 0.21%
slim-resnet50v2 64 3,230.74 3,180.91 1.57%
bert-mrpc-onnx 8 1,170.30 1,168.84 0.13%
bert-mrpc-tf 1 499.07 498.63 0.09%
pytorch-examples-wlang-gru 1 447.16 473.35 -5.53% 🔴
pytorch-examples-wlang-lstm 1 446.39 384.83 16.00% 🔆
torchvision-resnet50_1 1 1,052.11 1,046.63 0.52%
cadene-dpn92_1 1 443.23 437.32 1.35%
cadene-resnext101_1 1 363.66 365.89 -0.61%
onnx-taau-downsample 1 846.64 844.09 0.30%
dlrm-criteoterabyte 1 32.26 32.42 -0.50%
dlrm-criteoterabyte_fp16 1 51.60 51.80 -0.38%
agentmodel 1 14,514.86 9,209.12 57.61% 🔆
unet_fp16 2 58.36 58.80 -0.74%
resnet50v1_fp16 1 1,440.90 1,366.11 5.47% 🔆
resnet50v1_int8 1 1,768.49 1,883.96 -6.13% 🔴
bert_base_cased_fp16 64 1,098.71 1,098.16 0.05%
bert_large_uncased_fp16 32 347.23 345.59 0.47%
bert_large_fp16 1 206.33 206.59 -0.13%
distilgpt2_fp16 16 2,102.45 2,092.89 0.46%
yolov5s 1 564.02 558.33 1.02%
tinyllama 1 45.78 45.83 -0.10%
vicuna-fastchat 1 44.23 44.20 0.06%
whisper-tiny-encoder 1 412.78 411.87 0.22%
whisper-tiny-decoder 1 409.37 408.48 0.22%
llama2_7b 1 20.88 20.84 0.19%
qwen1.5-7b 1 23.67 23.58 0.37%
phi3-3.8b 1 28.37 26.72 6.19% 🔆
llama3-8b 1 22.79 21.80 4.51%
whisper-large-encoder 1 10.17 10.18 -0.09%
whisper-large-decoder 1 107.66 105.30 2.24%
mistral-7b 1 23.75 23.78 -0.10%
FLUX.1-schnell 1 813.68 755.22 7.74% 🔆

Regressions detected 🔴

* No develop baseline was found for this PR's branch point; compared against the latest available develop run instead.

@gh-app-migraphx-bot-pr-write

Copy link
Copy Markdown
Test Status Result
bert-mrpc-onnx PASSED: MIGraphX meets tolerance
bert-mrpc-tf ERROR - check error output
traceback
Traceback (most recent call last):
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 377, in
main()
File "/src/AMDMIGraphX/tools/accuracy/accuracy_checker.py", line 313, in main
import tensorflow as tf
File "/usr/local/lib/python3.12/dist-packages/tensorflow/init.py", line 40, in
from tensorflow.python import pywrap_tensorflow as _pywrap_tensorflow # pylint: disable=unused-import
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/tensorflow/python/pywrap_tensorflow.py", line 37, in
self_check.preload_check()
File "/usr/local/lib/python3.12/dist-packages/tensorflow/python/platform/self_check.py", line 63, in preload_check
from tensorflow.python.platform import _pywrap_cpu_feature_guard
ImportError: libnuma.so.1: cannot open shared object file: No such file or directory
pytorch-examples-wlang-gru 🔴 FAILED: MIGraphX is not within tolerance - check verbose output
pytorch-examples-wlang-lstm 🔴 FAILED: MIGraphX is not within tolerance - check verbose output
dlrm-criteoterabyte PASSED: MIGraphX meets tolerance
agentmodel PASSED: MIGraphX meets tolerance
unet PASSED: MIGraphX meets tolerance
resnet50v1 PASSED: MIGraphX meets tolerance
bert_base_cased_fp16 PASSED: MIGraphX meets tolerance
bert_large_uncased_fp16 🔴 FAILED: MIGraphX is not within tolerance - check verbose output
bert_large PASSED: MIGraphX meets tolerance
yolov5s PASSED: MIGraphX meets tolerance
tinyllama PASSED: MIGraphX meets tolerance
vicuna-fastchat PASSED: MIGraphX meets tolerance
whisper-tiny-encoder PASSED: MIGraphX meets tolerance
whisper-tiny-decoder PASSED: MIGraphX meets tolerance
distilgpt2_fp16 🔴 FAILED: MIGraphX is not within tolerance - check verbose output
llama2_7b PASSED: MIGraphX meets tolerance
qwen1.5-7b PASSED: MIGraphX meets tolerance
phi3-3.8b PASSED: MIGraphX meets tolerance
llama3-8b PASSED: MIGraphX meets tolerance
whisper-large-decoder PASSED: MIGraphX meets tolerance
mistral-7b PASSED: MIGraphX meets tolerance
FLUX.1-schnell PASSED: MIGraphX meets tolerance

@bdevorem bdevorem left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

needs a tidy and a changelog, but the code changes lgtm

@pfultz2
pfultz2 requested a review from a team as a code owner August 29, 2026 21:35
@causten causten added the high priority A PR with high priority for review and merging. label Aug 31, 2026

@CharlieL7 CharlieL7 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

# Conflicts:
#	CHANGELOG.md
#	src/fuse_attention.cpp
pfultz2 and others added 2 commits September 1, 2026 18:13
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@causten
causten merged commit 941b59e into develop Sep 2, 2026
12 of 13 checks passed
@causten
causten deleted the slice-reshapes-concat branch September 2, 2026 13:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

high priority A PR with high priority for review and merging.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants