Skip to content

Compile MLIR prefills with tuning candidates - #5268

Closed
umangyadav wants to merge 4 commits into
developfrom
fix/compile-benchmark-device-ops-parallel
Closed

umangyadav wants to merge 4 commits into
developfrom
fix/compile-benchmark-device-ops-parallel

Conversation

@umangyadav

Copy link
Copy Markdown
Member

Problem

The MLIR tuning path compiled inserted hip::fill operations from time_solution(). Since candidate benchmarking is serial, this moved fill compilation out of the parallel candidate-compilation phase.

Follow-up to #5252 and the maintainer feedback in #5252 (comment).

Changes

  • Compile MLIR prefill operations alongside each tuning candidate.
  • Carry the compiled fill code objects into the replacement callback and insert them before their associated MLIR kernels.
  • Flatten tuple output allocation shapes when resolving prefill buffer shapes.
  • Remove device-op compilation from the serial benchmark path.

Validation

  • python3 tools/format.py origin/develop --clang-format-path /opt/rocm/llvm/bin/clang-format
  • cmake --build build --target analyze --parallel 32
  • git diff --check origin/develop...HEAD

Cppcheck was not available in the local build environment.

Made with Cursor

Co-authored-by: Cursor <cursoragent@cursor.com>
@umangyadav
umangyadav requested a review from causten as a code owner September 15, 2026 22:24
Copilot AI lite review requested due to automatic review settings September 15, 2026 22:24
@umangyadav umangyadav self-assigned this Sep 15, 2026
@umangyadav
umangyadav requested a review from pfultz2 September 15, 2026 22:25

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.

🟡 Changes recommended

Add enabled regression coverage for the new prefill integration path.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

This PR moves MLIR prefill compilation into parallel tuning-candidate compilation and removes serial benchmark-time compilation.

Changes:

  • Carries compiled prefill code into replacement callbacks.
  • Handles flattened tuple-output allocation shapes.
  • Inserts fills before associated MLIR kernels.
File summaries
File Description
src/targets/gpu/jit/mlir.cpp Compiles and inserts MLIR prefill operations.
src/targets/gpu/compile_ops.cpp Removes serial device-op compilation during benchmarking.

The new path lacks an enabled GPU MLIR regression covering prefilled candidates and tuple-output allocation.

Review details
  • Files reviewed: 2/2 changed files
  • Comments generated: 1
  • Review effort level: Lite

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

Comment thread src/targets/gpu/jit/mlir.cpp

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.

🟡 Changes recommended

Tuple-output prefills can be eliminated before benchmarking, and the integration lacks enabled regression coverage.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

src/targets/gpu/jit/mlir.cpp:301

  • The new compiled-prefill integration has no enabled regression test. The only MLIR prefill tests in test/gpu/mlir.cpp are skipped, while the existing hip::fill tests exercise the compiler only in isolation; therefore incorrect prefill-index mapping, tuple-output flattening, or code-object offsets can pass CI. Add an enabled GPU MLIR test that compiles and executes a prefilled candidate with a tuple output allocation.
                return migraphx::gpu::compile_op(
                    fill.name(), ctx, {inputs.at(mco.prefill_indices.at(i))}, fill.to_value());
  • Files reviewed: 2/2 changed files
  • Comments generated: 1
  • Review effort level: Balanced

Comment thread src/targets/gpu/jit/mlir.cpp
umangyadav and others added 2 commits September 15, 2026 23:04
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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.

🔵 Needs a closer look

The modified split-candidate prefill path lacks deterministic regression coverage.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

src/targets/gpu/jit/mlir.cpp:393

  • The split fallback's new prefill compilation is not covered by the added regression. prefill_tuple_output requires exactly one MLIR code object plus two fills (code_objects.size() == 3), while this branch starts with two MLIR code objects, so it cannot validate this path or the mcos.size() + i replacement offset. Add a case that deterministically forces is_module_fusible to return false and verifies the split candidate's fill is compiled, retained, and executed before the first MLIR kernel.
  • Files reviewed: 3/3 changed files
  • Comments generated: 0 new
  • Review effort level: Balanced

result.reserve(indices.size());
std::transform(
indices.begin(), indices.end(), std::back_inserter(result), [&](const auto i) {
auto fill = migraphx::make_op("hip::fill", {{"value", mco.prefill_values.at(i)}});

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.

is this cached? I guess fill kernels (with the same shapes) will be repeated over and over. I'd be nice to have a dynamic shape fill kernel if the dynamic functionality is ready, so it gets compiled only once.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

It is not cached. I think #5270 addresses it right ?

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.

yes, I got the idea while reviewing the code here. The results on that PR are on top of this PR.

@umangyadav umangyadav added the high priority A PR with high priority for review and merging. label Sep 16, 2026

@pfultz2 pfultz2 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.

We should use the precompile fill instead which is what is done with #5275.

@umangyadav

Copy link
Copy Markdown
Member Author

Closing this one. #5275

@umangyadav umangyadav closed this Sep 17, 2026
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.

4 participants