Gpu concat kernel improvements

[pfultz2]

Motivation

Concat kernels where the concatenation happens on the fast (last) axis with a small number of elements per input (< 64) perform poorly with the existing simple algorithm: each input writes a narrow slice of the output, so global writes are strided and uncoalesced, and the launch does little work per wavefront.

This PR adds a block_tile algorithm to the concat JIT kernel that stages the concatenated output tile in LDS. Each workgroup processes a group of output slices: the per-op inputs (including any fused pointwise pre-ops) are written into a {ngroups, nops, max_size} LDS tile, and after a barrier the tile is written out contiguously in blocked form with the fused post-op applied. This turns the scattered per-op writes into coalesced output writes.

Supporting this required generalizing reduce_dims to handle shape sets that differ along one axis (the concat axis), so the surrounding equal dimensions can still be collapsed, plus some kernel-infrastructure additions (multi-index subscripts on tensor_view, preprocessor enum machinery, and a blocked stride loop). The new unit-test coverage added along the way caught and fixed several bugs in the initial implementation.

Technical Details

Kernel (kernels/concat.hpp)

• concat::run is now parameterized on an algorithm. simple keeps the previous behavior (global-strided writes directly to the output slice). block_tile<NGroups> stages results into a {NGroups, nops, max_size} LDS tile (allocated via uninitialized_buffer) during run and writes the output in finish after __syncthreads(), using block_stride<per_block, 8> so each thread writes 8 contiguous elements.
• Slices are scheduled per workgroup with slice_schedule<per_block>(idx, slice_axes<-1>(), slice_group<NGroups>()); the group/element index of each slice element is derived with shape indexing (make_shape + multi) rather than ad-hoc arithmetic.

Host compiler (jit/concat.cpp, compile_gen)

• block_tile is selected when the concat axis is the fast axis, every op contributes the same number of elements (< 64), and the LDS tile (group * nops * max_size * type_size) fits in the 64KB workgroup limit; otherwise it falls back to simple.
• The group count comes from tile::compute_factor(lens[axis - 1], 16) — compute_tile_factor moved into the tile class, exported, and unit-tested (including its check-before-multiply overshoot behavior).

reduce_dims

• mask_shape no longer gives up when shapes disagree on an axis with neither length being 1; the differing axis is masked with stride 0, which blocks merging across it while letting the equal dimensions on either side collapse (e.g. {64,16,160,160} / {64,48,160,160} → {64,16,25600} / {64,48,25600}).
• Masking still bails out (returning the shapes unchanged) when an incompatible axis is adjacent to another masked axis, since a stride-0 pair would incorrectly allow those axes to merge — this preserves correctness for cases like pad ({1,3,224,224} vs {1,3,229,229}). TODO tests document the mergeable-but-unmerged dimensions this early exit leaves behind.

Kernel infrastructure

• index.hpp: ngroup() uses ceiling division; new block_stride<Group, Block> iterates a range in per-thread blocks of Block contiguous elements with a distributed tail.
• tensor_view.hpp: variadic index_to_offset constructor enables multi-index subscripts (output[{group, depth, k}]), including under MIGRAPHX_DEBUG source-location capture (debug.hpp generates the forwarding constructors with MIGRAPHX_PP_ENUM).
• pp.hpp: adds MIGRAPHX_PP_ENUM, MIGRAPHX_PP_GENERATE, MIGRAPHX_PP_BOOL/NOT, deferred-expansion helpers, and data-carrying variants of the argument-transform macros; MIGRAPHX_PP_REPEAT is now curried.
• functional.hpp: pack/pack_forward are declared before unpack_each (fixes a declaration-order failure for non-ADL types); arg_c special-cases N == 0 to reduce template instantiations.

Testing

• New GPU kernel unit-test suites for pp.hpp, functional.hpp, shape.hpp, and block_stride (~180 cases), with TEST_CASE_TEMPLATE/TEST_CASE_REGISTER support added to the kernel test harness. The block_stride coverage test (visit-each-element-exactly-once via a serial group emulator) is a regression test for a tail-loop bug it caught during development.
• tile::compute_factor unit tests in test_gpu_compile_gen; new reduce_dims tests for the differing-axis reductions, the incompatible-adjacency bail-out, and TODO cases.
• Verify tests: test_concat_axis_neg_1 is templated over type/sizes to exercise block_tile (including a float instantiation), and test_concat_lds_overflow (30 inputs × 60 elements × 16 groups = 115KB tile) is a regression test that the gate falls back to simple instead of generating a kernel that exceeds the LDS limit.

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.

[Copilot]

Pull request overview

Optimizes GPU concat for small, equal-width fast-axis inputs using an LDS-backed tiled kernel, with supporting kernel utilities and tests.

Changes:

• Adds tiled concat selection and execution.
• Extends kernel indexing, shape, debug, and preprocessor utilities.
• Expands GPU kernel and concat verification coverage.

Review performed as a single pass without agent fan-out.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/reduce_dims.cpp	Enables reduction across differing fast dimensions.
src/targets/gpu/compile_gen.cpp	Exposes tile-factor computation.
src/targets/gpu/include/migraphx/gpu/compile_gen.hpp	Declares the tile-factor helper.
src/targets/gpu/jit/concat.cpp	Selects and launches tiled concat.
src/targets/gpu/kernels/include/migraphx/kernels/concat.hpp	Implements tiled and simple concat algorithms.
src/targets/gpu/kernels/include/migraphx/kernels/debug.hpp	Supports multi-argument source-location capture.
src/targets/gpu/kernels/include/migraphx/kernels/functional.hpp	Reorders helpers and simplifies arg_c<0>.
src/targets/gpu/kernels/include/migraphx/kernels/index.hpp	Adds ceiling group counts and block-stride traversal.
src/targets/gpu/kernels/include/migraphx/kernels/pp.hpp	Extends recursive preprocessor utilities.
src/targets/gpu/kernels/include/migraphx/kernels/reduce.hpp	Adapts reduction to the revised repeat macro.
src/targets/gpu/kernels/include/migraphx/kernels/shape.hpp	Asserts nonempty kernel shapes.
src/targets/gpu/kernels/include/migraphx/kernels/tensor_view.hpp	Adds multidimensional index construction.
src/targets/gpu/kernels/include/migraphx/kernels/test.hpp	Adds templated kernel-test macros.
test/gpu/compile_gen.cpp	Tests tile-factor behavior.
test/gpu/kernels/functional.cpp	Adds device functional-utility tests.
test/gpu/kernels/index.cpp	Tests block-stride traversal.
test/gpu/kernels/main.cpp	Discovers templated kernel tests.
test/gpu/kernels/pp.cpp	Adds device preprocessor tests.
test/gpu/kernels/shape.cpp	Adds device shape tests.
test/reduce_dims.cpp	Covers differing fast dimensions.
test/verify/test_concat_axis_neg_1.cpp	Exercises varied concat sizes and types.

---

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

[pfultz2]

This has been fixed and additional unit tests added.

[pfultz2]

Added a test for this and a check.

[pfultz2]

Right now the test harness doesnt support launching with different launch params so its not possible to fully unit test the index class yet.

[github-actions]

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

format.py

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 520 in c5df9f2

auto s = migraphx::make_shape(migraphx::index_ints<2, 3>{}, migraphx::index_ints<1, 2>{});

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 527 in c5df9f2

auto s = migraphx::make_shape(migraphx::index_ints<2, 3, 4>{}, migraphx::index_ints<1, 8, 2>{});

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 539 in c5df9f2

auto s = migraphx::make_shape(migraphx::index_ints<5, 4, 2, 3>{}, migraphx::index_ints<5, 1, 20, 60>{});

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 546 in c5df9f2

auto out_lens = migraphx::index_ints<2, 3>{};

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Lines 548 to 549 in c5df9f2

	auto out_shape =
	migraphx::make_shape_from_permutation(out_lens, permutation);

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 556 in c5df9f2

auto out_lens = migraphx::index_ints<2, 3, 4>{};

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Lines 558 to 559 in c5df9f2

	auto out_shape =
	migraphx::make_shape_from_permutation(out_lens, permutation);

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 566 in c5df9f2

auto out_lens = migraphx::index_ints<5, 4, 2, 3>{};

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Lines 568 to 569 in c5df9f2

	auto out_shape =
	migraphx::make_shape_from_permutation(out_lens, permutation);

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/gpu/kernels/shape.cpp

Line 573 in c5df9f2

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/reduce_dims.cpp

Line 195 in c5df9f2

std::vector<migraphx::shape> eshapes = {make_shape({64, 16, 25600}), make_shape({64, 48, 25600})};

---

[format.py] _{reported by reviewdog 🐶}

AMDMIGraphX/test/reduce_dims.cpp

Line 205 in c5df9f2

std::vector<migraphx::shape> eshapes = {make_shape({64, 8, 2, 25600}), make_shape({64, 4, 12, 25600})};