feat(hipblaslt): select JIT solutions in heuristic queries by jaopaulolc · Pull Request #13202 · ROCm/rocm-libraries

jaopaulolc · 2026-10-06T20:43:25Z

JIRA ID

AIHPBLAS-4548
AIHPBLAS-4550
AIHPBLAS-4552
AIHPBLAS-4801
AIHPBLAS-5162

Motivation

Add JIT solution library and comgr builder can build a TensileLite kernel in the running process and store it on disk, but no query reads that store. This pull request connects those queries to the store. A kernel built for one query is reused by later queries and by later processes that can read the same directory. With the new setting left unset, the queries keep their current behavior.

Technical Details

HIPBLASLT_JIT is read once per process in hipblaslt-jit-mode.cpp. Unset, empty or 0 leaves the existing query path. 1 places results from the JIT library after the prebuilt library's exact-size matches and before its other matches. 2 returns only JIT library results. Any other value warns once and leaves the setting off. docs/reference/env-variables.rst documents this variable and HIPBLASLT_JIT_LIBRARY_PATH. When that path is unset, the library root is hipblaslt-jit-<uid> under the system temporary directory. TMPDIR is not used.

When the library has no matching kernel, Jit::generate builds one in a private directory, checks that the solution accepts the problem, links the code object, and publishes it. The published index is in [0x40000000, 0x60000000), the range the library pull request reserves. hipblasLtMatmul sends an index in that range to the JIT library and every other index to the prebuilt library. If the prebuilt library already contains an index in the JIT range, the JIT library is disabled for that process.

This pull request also declares the generation interfaces the library pull request leaves out: Backend, SolutionStore, Jit and makeJit in hipblaslt-jit-component.hpp and hipblaslt-jit-backend.hpp, the getJitAlgo, getLibraryAlgos and makeGemmRequest entry points in hipblaslt-jit.hpp, and DeviceTarget::make, makeLibraryStore and RocblasltToTensileContractionProblemGemm.

The backend in this pull request replays the committed test kernels from the library pull request. It publishes two copies. When K is a multiple of 512, the copy was built with workgroup mapping 8. The other copy, used for any K, was built with workgroup mapping 1. It does not call TensileLite's Python. Grouped GEMM is not filled from the JIT library. A build that was not given those kernels warns once: mode 1 is then unchanged, and mode 2 returns no algorithms.

flowchart TD
  Q["heuristic query"] --> M{"HIPBLASLT_JIT"}
  M -->|"unset or 0"| OLD["existing getBestSolutions path"]
  M -->|"1"| EQ["exact-size matches, then JIT library"]
  M -->|"2"| JL["JIT library only"]
  EQ --> JL
  JL -->|"no stored kernel"| GEN["build, check, publish"]
  GEN --> JL
  EQ --> OT["remaining prebuilt matches"]

These edits are outside the JIT build flag and apply to every build:

RocBlasLtStatusToHIPStatus maps rocblaslt_status_not_initialized and rocblaslt_status_type_mismatch to hipBLASLt statuses. It used to throw HIPBLAS_STATUS_INVALID_ENUM for those two values.
makeArgument records the algorithm only after solve succeeds, and it resets inputs.Synchronizer on each call.
getKernelNameFromAlgoIndex and getSolutionNameFromAlgoIndex return an empty string for an unknown index instead of using a null solution.
heuristicResult_copy copies the whole algorithm struct, not only its index bytes.
TensileLite's CachingLibrary includes the provider-row filter in its cache key. With no filter set, that key is the previous All filter.
gemmCreate clears prepared kernels on every build. Gemm::run with none returns HIPBLAS_STATUS_NOT_INITIALIZED instead of launching nothing and returning success. setProblem followed by run without initialize fails the same way with JIT on and with JIT off.

Add a HipKittens gfx950 JIT backend adds a second backend for gfx950. It is independent of replay, and a query asks each backend the build compiled.

Test Plan

On a host-only JIT build (-DHIPBLASLT_ENABLE_JIT=ON -DHIPBLASLT_ENABLE_DEVICE=OFF), ctest on clients/tests/jit covers the five tests from the library pull request plus:

jit-library and jit-library-concurrency check publication, the index range, and concurrent readers.
jit-end-to-end runs a replayed solution from getJitAlgo through hipblasLtMatmul and hipblaslt_ext::Gemm, and compares the output with a host reference.
jit-end-to-end-library does the same for indices from getLibraryAlgos.
jit-heuristic-off, jit-heuristic-fallback and jit-heuristic-forced check modes 0, 1 and 2. The fallback test checks the order (exact-size matches, then JIT, then the other prebuilt matches) only when the device library has an exact-size row. Otherwise it prints SKIP for that check.

clients/tests/src/provider_rows_gtest.cpp covers the provider-row filter. It is part of hipblaslt-test and was not run for this description.

Test Result

On gfx950, with HIP_VISIBLE_DEVICES=4, ctest on clients/tests/jit passed 12 of 12 at 2948ec5f6ae. The later merge d8e7db8c500 only brings the source-bundle test's link fix from the parent branch. That merge was not re-run locally. gfx90a and gfx942 were not run on this machine. This pull request's checks do not include the JIT workflow. That workflow runs on the library pull request, where run 38434784675 passed on gfx90a, gfx942 and gfx950 for eda3a798cd6. The workflow's path list now includes the rocBLASLt files this pull request edits, so a change to those files retriggers it on a pull request that already runs the workflow.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

Risk level

Medium. Queries are unchanged when JIT is not built, or when HIPBLASLT_JIT is unset or 0. Setting it to 1 or 2 changes which algorithms the query returns, and those kernels were not selected by a benchmark. Mode 1 places them ahead of every prebuilt match that is not an exact-size match.

The directory is not pruned, so it grows with each new problem. Deleting the root while no process is using it clears it.

The status mapping for rocblaslt_status_not_initialized and rocblaslt_status_type_mismatch changes in every build, including builds with JIT off. The other items in the list above also apply to every build.

Gemm::run with no prepared kernels also returns HIPBLAS_STATUS_NOT_INITIALIZED on every build. setProblem drops those kernels, so a later run without initialize does not reuse them when JIT is off.

Rollback is reverting this pull request, or leaving HIPBLASLT_JIT unset. The HipKittens pull request uses this branch as its base.

therock-pr-bot · 2026-10-06T22:23:28Z

✅ All Checks Passed — Ready for Review

Check	Status	Details
📝 PR Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—
🤖 therock-pr-bot	✅ Pass	—

🎉 All checks passed! This PR is ready for review.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

🙋 Wish to Override Policy?

therock-pr-bot · 2026-10-06T22:23:31Z

🎉 All checks passed! This PR is ready for review.

codecov · 2026-10-07T06:22:34Z

Codecov Report

❌ Patch coverage is 78.28947% with 33 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rary/src/amd_detail/rocblaslt/src/tensile_host.cpp	76.09%	12 Missing and 10 partials ⚠️
...cts/hipblaslt/library/src/amd_detail/hipblaslt.cpp	25.00%	3 Missing and 3 partials ⚠️
...c/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp	92.16%	1 Missing and 3 partials ⚠️
...ary/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@                                Coverage Diff                                @@
##           users/jolabega/jit-upstream-01-jit-interfaces   #13202      +/-   ##
=================================================================================
- Coverage                                          78.09%   78.08%   -0.00%     
=================================================================================
  Files                                               3005     3005              
  Lines                                             522055   522117      +62     
  Branches                                           81485    81495      +10     
=================================================================================
+ Hits                                              407656   407688      +32     
- Misses                                             85475    85497      +22     
- Partials                                           28924    28932       +8

Flag	Coverage Δ		*Carryforward flag
TensileLite-CPP	`53.88% <ø> (-<0.01%)`	⬇️
TensileLite-Unit	`85.08% <ø> (ø)`
hipBLAS	`90.62% <ø> (ø)`		Carriedforward from 4806b97
hipBLASLt	`70.02% <78.29%> (-0.05%)`	⬇️
hipCUB	`82.68% <ø> (ø)`		Carriedforward from 4806b97
hipDNN	`87.34% <ø> (ø)`		Carriedforward from 4806b97
hipFFT	`42.65% <ø> (ø)`		Carriedforward from 4806b97
hipRAND	`76.12% <ø> (ø)`		Carriedforward from 4806b97
hipSOLVER	`69.67% <ø> (ø)`		Carriedforward from 4806b97
hipSPARSE	`87.03% <ø> (ø)`		Carriedforward from 4806b97
rocALUTION	`82.89% <ø> (ø)`		Carriedforward from 4806b97
rocBLAS	`69.20% <ø> (ø)`		Carriedforward from 4806b97
rocFFT	`46.49% <ø> (ø)`		Carriedforward from 4806b97
rocRAND	`57.32% <ø> (ø)`		Carriedforward from 4806b97
rocSOLVER	`76.66% <ø> (ø)`		Carriedforward from 4806b97
rocSPARSE	`83.61% <ø> (ø)`		Carriedforward from 4806b97
rocThrust	`93.15% <ø> (ø)`		Carriedforward from 4806b97

*This pull request uses carry forward flags. Click here to find out more.

Files with missing lines	Coverage Δ
.../amd_detail/rocblaslt/src/include/tensile_host.hpp	`28.57% <ø> (ø)`
...ary/src/amd_detail/rocblaslt/src/rocblaslt_mat.cpp	`92.43% <0.00%> (-0.08%)`	⬇️
...c/amd_detail/rocblaslt/src/rocblaslt_auxiliary.cpp	`84.60% <92.16%> (+0.05%)`	⬆️
...cts/hipblaslt/library/src/amd_detail/hipblaslt.cpp	`81.54% <25.00%> (-0.89%)`	⬇️
...rary/src/amd_detail/rocblaslt/src/tensile_host.cpp	`65.92% <76.09%> (+0.08%)`	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

jaopaulolc

Please address my comments.

jinchen62

Reviewed at 925ca44 against the base branch at f297932f, and built the host-only JIT configuration (HIPBLASLT_ENABLE_JIT=ON, HIPBLASLT_ENABLE_DEVICE=OFF) on gfx950.

Blocking

The head doesn't compile with JIT enabled (inline on hipblaslt-jit-backend.cpp:263). JIT run 40543334461 fails at the build step on gfx90a, gfx942 and gfx950. The description says this PR's checks don't include the JIT workflow, but it does run here, and it's red.
The JIT GPU tests fail independently of that. At 05a6166, before the wavesCount commit, 9 of 10 failed (run 40538632774) with "plain K=512", "K=1024 did not keep both plain-pair solutions", and _WGM1 chosen where _K512_WGM8 is expected. #13111 fails the same jit-publish and jit-launch checks at its head (run 40543243737), so this looks inherited from the base. The 12-of-12 result in the description is from 2948ec5f6ae, before both.
Every JIT index lookup reloads the code object and keeps it (inline on tensile_host.cpp:3490). I measured about 4 MiB of device memory per launch.

The other items are inline. Things I checked and ruled out: concurrent heuristic queries are safe, because the query state is thread_local; GemmRequest never copies alpha/beta from device memory, because there's no device-scalar pointer mode; and the override row can't overflow the results array, because requestedAlgoCount is decremented first. The provider-row filter and its cache key, the JIT-range collision check, the whole-struct heuristicResult_copy, and the null-solution guards in the name lookups look good.

I couldn't run the JIT tests end to end locally. Develop now needs a ROCm newer than 7.0 to build (hipDeviceAttributePciChipId), and ROCm 7.0's comgr aborts reading the committed kernels' metadata. The memory numbers come from calling SolutionAdapter::loadCodeObjectFile directly; the comments about generation cost come from reading the code.

CongMa13 · 2026-10-28T16:11:53Z

I reviewed this PR again. The following review threads are still unresolved, or comments have no author reply:

If you have already fixed these, please reply on each comment to say so and resolve the threads.

CongMa13 · 2026-10-28T18:40:08Z

I reviewed this PR again. The following review threads are still unresolved, or comments have no author reply:

If you have already fixed these, please reply on each comment to say so and resolve the threads.

CongMa13 · 2026-10-29T17:19:31Z

I reviewed this PR again. The following review threads are still unresolved, or comments have no author reply:

If you have already fixed these, please reply on each comment to say so and resolve the threads.

jaopaulolc · 2026-10-30T14:14:01Z

Those three threads already have a reply and are resolved:

The review at #13202 (review) is the summary of those inline comments. The compile failure, the code-object reload, and the other inline items on that review are answered on their own threads.

CongMa13

Two issues in the JIT selection and first-use generation paths are noted inline.

CongMa13

No new issues found. I agree with the following existing comments:

Mixed wavesCount semantics for JIT and prebuilt results — @jinchen62 — comment #4392314406
Mode 1 provider-order coverage and test process isolation — @jinchen62 — comment #4392314507
Confirm the JIT workflow's GPU runner cost with CI owners — @jinchen62 — comment #4392314518
Propagate code-object load failures — @jinchen62 — comment #4392314531 (the fork-safety part was addressed; the load result is still discarded)

CongMa13

I found two P2 issues in the JIT CI integration: the provider-row implementation is missing from the workflow path filter, and the JIT lane is not documented in the component testing strategy.

I agree with the following existing comments:

Mixed wavesCount semantics for JIT and prebuilt results — @jinchen62 — comment #4392314406
Mode 1 provider-order coverage is skipped in CI — @jinchen62 — comment #4392314507
Confirm the JIT workflow's GPU runner cost — @jinchen62 — comment #4392314518
Propagate code-object load failures — @jinchen62 — comment #4392314531
Correct the PR description's JIT-check statement — @jinchen62 — review #5671968669

CongMa13 · 2026-11-06T15:47:20Z

+      - "projects/hipblaslt/library/src/amd_detail/rocblaslt/src/handle.cpp"
      - "projects/hipblaslt/tensilelite/include/Tensile/CachingLibrary.hpp"
      - "projects/hipblaslt/tensilelite/include/Tensile/ExactLogicLibrary.hpp"
      - "projects/hipblaslt/tensilelite/include/Tensile/ProviderRows.hpp"


This filter includes ProviderRows.hpp but omits projects/hipblaslt/tensilelite/src/ProviderRows.cpp, which implements the provider-row behavior used by heuristic selection. A later change to the implementation alone will skip the JIT workflow. Please add this path or a pattern that covers the implementation sources.

Added in d1df9741553: projects/hipblaslt/tensilelite/src/ProviderRows.cpp is now in the paths: filter next to ProviderRows.hpp. I checked the filter with a script that applies GitHub path-filter matching to the tracked files: all entries match at least one file, including the new one (the top of the stack, #13209, adds the HipKittens entries and every one of its 24 entries matches).

CongMa13 · 2026-11-06T15:47:20Z

+      - "projects/hipblaslt/library/src/amd_detail/hipblaslt-ext.cpp"
      - "projects/hipblaslt/clients/tests/jit/**"
      - "projects/hipblaslt/scripts/jit/**"
+      - "projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp"


Please add this JIT lane to projects/hipblaslt/TESTING.md, including its trigger paths, tested architectures, and whether it gates merges. That file is the component's current CI and gating reference, but it has no JIT lane entry.

Done in d1df9741553. projects/hipblaslt/TESTING.md now lists the JIT lane in three places: the short version (not required), a new row in the Validation Gates table, and the Informational list under PR Test Classification. The row gives the workflow (hipBLASLt JIT CI, checks JIT (gfx90a), JIT (gfx942), JIT (gfx950)), what it runs (host-only build with HIPBLASLT_ENABLE_JIT=ON, the jit-cpu tests with every GPU hidden, then the jit-gpu tests, one GPU runner per architecture), the trigger (pull requests that change a path in the paths: filter, and manual dispatch) and the paths by area. On gating: it is not a required check. gh pr checks --required on the top PR lists only Component CI Summary, pre-commit and Math CI Summary, and the develop ruleset list in this file does not include a JIT job, so a red JIT job does not block a merge. The HipKittens kernel directory and its CMake file are added to that row in #13209 (9b3b5998ceb).

The attempt key omitted K and the operand types, so a second shape was never built, and a publish reloaded every code object already in use. Library-test children now start with posix_spawn, and setProblem drops prepared kernels on every build.

A matmul with no caller algorithm bound the prebuilt library before selection, so a JIT index from mode 1 or 2 could not be launched. Selection now follows the heuristic ranking and the launch resolves that index. Concurrent generation of one problem waits for the first attempt instead of building it twice.

…tic layer

jaopaulolc requested review from a team as code owners October 6, 2026 20:43

github-actions Bot added documentation project: hipblaslt project: hipsparselt ci:hipsparselt-fast labels Oct 6, 2026

jaopaulolc added this pull request to stack #13185 October 6, 2026 20:49

jaopaulolc added the rocm:gemm algos Label to identify PRs of the GEMM Algorithms team. label Oct 6, 2026

assistant-librarian Bot added the organization: ROCm label Oct 6, 2026

jaopaulolc mentioned this pull request Oct 6, 2026

feat(hipblaslt): publish JIT solutions to a persistent library #13208

Closed

1 task

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from b1900cf to 53196e2 Compare October 6, 2026 23:13

jaopaulolc removed this pull request from stack #13185 October 6, 2026 23:16

jaopaulolc added this pull request to stack #13210 October 6, 2026 23:16

jaopaulolc requested a review from a team as a code owner October 6, 2026 23:17

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from 53196e2 to fda6a1c Compare October 6, 2026 23:17

jaopaulolc mentioned this pull request Oct 7, 2026

feat(hipblaslt): run JIT GEMM solutions through the normal launch path #13131

Closed

1 task

jaopaulolc removed this pull request from stack #13210 October 7, 2026 14:50

jaopaulolc added this pull request to stack #13245 October 7, 2026 14:50

jaopaulolc removed this pull request from stack #13245 October 7, 2026 14:51

jaopaulolc changed the base branch from users/jolabega/jit-upstream-01b-jit-gemm-run to users/jolabega/jit-upstream-01-jit-interfaces October 7, 2026 15:27

jaopaulolc added this pull request to stack #13248 October 7, 2026 15:47

jaopaulolc mentioned this pull request Oct 7, 2026

feat(hipblaslt): add JIT solution library and comgr builder #13111

Open

1 task

github-actions Bot added the github actions label Oct 7, 2026

jaopaulolc commented Oct 7, 2026

View reviewed changes

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from 2b879a2 to d42ce30 Compare October 8, 2026 15:57

jaopaulolc requested a review from talumbau October 27, 2026 16:33

github-actions Bot added the ci:hipsparselt-fast label Oct 27, 2026

jinchen62 reviewed Oct 27, 2026

View reviewed changes

ROCm deleted a comment from cgmillette Oct 27, 2026

jaopaulolc requested a review from CongMa13 October 27, 2026 21:56

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from a1d2150 to 4162910 Compare October 28, 2026 14:38

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from 4162910 to 3cfa9ff Compare October 28, 2026 14:46

jaopaulolc requested a review from jinchen62 October 28, 2026 15:00

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from 1ee5961 to d1db3d9 Compare October 28, 2026 16:09

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from d1db3d9 to 33a6160 Compare October 28, 2026 16:37

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from 33a6160 to ed1b5ff Compare October 30, 2026 14:13

CongMa13 reviewed Nov 2, 2026

View reviewed changes

Comment thread projects/hipblaslt/library/src/amd_detail/rocblaslt/src/tensile_host.cpp Resolved

Comment thread projects/hipblaslt/library/src/amd_detail/hipblaslt-jit-library.cpp Resolved

jaopaulolc force-pushed the users/jolabega/jit-upstream-01c-jit-heuristic branch from ed1b5ff to dc855b7 Compare November 2, 2026 20:37

CongMa13 reviewed Nov 3, 2026

View reviewed changes

caio96 mentioned this pull request Nov 4, 2026

perf(hipblaslt): remove per-call locks from the matmul launch path #15091

Open

1 task

CongMa13 reviewed Nov 6, 2026

View reviewed changes

jaopaulolc added 4 commits November 6, 2026 19:19

feat(hipblaslt): select JIT solutions in heuristic queries

e855801

refactor(hipblaslt): move generation interface declarations to heuris…

fc3e6e3

…tic layer

Conversation

jaopaulolc commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Risk level

Uh oh!

therock-pr-bot Bot commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Checks Passed — Ready for Review

Uh oh!

therock-pr-bot Bot commented Oct 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Oct 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

jaopaulolc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jinchen62 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

CongMa13 commented Oct 28, 2026

Uh oh!

CongMa13 commented Oct 28, 2026

Uh oh!

CongMa13 commented Oct 29, 2026

Uh oh!

jaopaulolc commented Oct 30, 2026

Uh oh!

CongMa13 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

CongMa13 left a comment

Choose a reason for hiding this comment

Uh oh!

CongMa13 left a comment

Choose a reason for hiding this comment

Uh oh!

CongMa13 Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

jaopaulolc Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

CongMa13 Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

jaopaulolc Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

jaopaulolc commented Oct 6, 2026 •

edited

Loading

therock-pr-bot Bot commented Oct 6, 2026 •

edited

Loading

therock-pr-bot Bot commented Oct 6, 2026 •

edited

Loading

codecov Bot commented Oct 7, 2026 •

edited

Loading