[hexagon-mlir] Land the HMX engine path - #101
Closed
Alfred (Dino1844) wants to merge 1 commit into
Closed
Alfred (Dino1844) wants to merge 1 commit into
Alfred (Dino1844) wants to merge 1 commit into
Conversation
The whole HMX feature lived only in the worktree (54 modified tracked files plus a dozen untracked paths). One `git checkout` or `git clean` would have lost it, and the local patch cannot carry untracked files either, so the decision to keep this work on the `hmx` branch (STATE-OF-PLAY section 6) was never actually executed. This lands it, together with the structural and code-slop fixes from the repository audit. New paths: - `hmx` dialect: 11 ops, the `#hmx.crouton` tensor encoding and its `#hmx.crouton_memref_layout` counterpart, and `hmx.alloc_crouton` as the one crouton allocation point. - four dialect passes (`matmul-to-hmx`, `hmx-partition`, `weight-resident`, `hmx-workspace-resident`) and the `hmx-to-llvm` conversion. - runtime leaves under `bin/runtime/hmx`, compiled by the SDK clang with `-mhmx` into `libhmxapi.a` -- deliberately not a bitcode candidate, because only Qualcomm's clang lowers the HMX builtins. - the `hvx` dialect and `vector-row-reduce`, weight/workspace residency, the f32 activation ABI, and `hmx.stage`/`hmx.await` value edges feeding `scf::pipelineForLoop`. - tests for all of the above. Structural fixes: - `matmul-to-hmx` is gated on the same condition as the tile level below it: `weight-resident` / `hmx-partition` / `hmx-workspace-resident` all sit inside `if (enableBufferization)`, while attribution sat outside it. `hmx.matmul` has no other consumer and `HmxToLLVMPass` marks the dialect illegal without a MatmulOp pattern, so attributing with bufferization off produced IR that cannot be lowered (a compile failure reachable from `--no-enable-bufferization`). - attribution refusals now emit one bounded module-level warning on top of the per-op remark. Remarks have no visible output on the production path, which is how a single switch silently removed the entire HMX path: measured on device at 35 124 vs 17 931 us (1.96x), with zero `hmx_*` leaves in the kernel. - ensure/unlock pairing is decided from `hmx` dialect ops collected before the conversion instead of a callee-name prefix, with a post-conversion check that fails the build if a leaf caller was not recognised -- a missed pair hangs the next thread forever in `HAP_compute_res_hmx_lock`. - the implicit accumulator uses a dialect side-effect resource (`HmxEngineResource`, non-addressable) instead of a generic `MemWrite`, so CSE no longer serialises it against unrelated memory, while the ops that really write addressable memory keep their default effects. - the compiled-kernel metadata is a named dict behind one field contract (`backend/utils.py`) instead of a 14-entry positional tuple read with `len(x) > i` guards, which turned a stale cached JSON into "treat the feature as off" instead of an error. - `Passes.td` defaults for `enable-split-reduction`, `enable-conversion-to-fp16` and `enable-weight-resident` now match the Python backend, so lit exercises the production configuration instead of a second set of defaults. - `HexagonProfiler` no longer reboots the device from `__init__` (explicit `start()`), and a missing environment variable is named along with how to set it instead of pointing at a file that does not exist. Behaviour outside the backend: - `test_flash_attention.py` turns the HMX path on (`enableConvertToHexagonmem`) and checks an fp16 quantisation budget instead of `torch.allclose`. It was hardcoded off, so this correctness test never exercised an HMX leaf; on device 39 383 -> 18 098 us (2.18x, PASS, rel ~7e-05). - `.gitignore` keeps the vendored triton trees, the tarballs and the in-tree logs out of the repository. Verified: 63/63 ninja targets; manual lit 184 pass / 1 pre-existing fail (`return_alloc_from_loop.mlir`, untouched since 09-16) / 1 skip; standard 7-operator device batch green on the resulting build (`libtriton.so` f1beb6b0..., 2026-09-22 15:31) at 121 / 248 / 266 / 145 / 789 / 346 / 18098 us.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The whole HMX feature lived only in the worktree (54 modified tracked files plus a dozen untracked paths). One
git checkoutorgit cleanwould have lost it, and the local patch cannot carry untracked files either, so the decision to keep this work on thehmxbranch (STATE-OF-PLAY section 6) was never actually executed. This lands it, together with the structural and code-slop fixes from the repository audit.New paths:
hmxdialect: 11 ops, the#hmx.croutontensor encoding and its#hmx.crouton_memref_layoutcounterpart, andhmx.alloc_croutonas the one crouton allocation point.matmul-to-hmx,hmx-partition,weight-resident,hmx-workspace-resident) and thehmx-to-llvmconversion.bin/runtime/hmx, compiled by the SDK clang with-mhmxintolibhmxapi.a-- deliberately not a bitcode candidate, because only Qualcomm's clang lowers the HMX builtins.hvxdialect andvector-row-reduce, weight/workspace residency, the f32 activation ABI, andhmx.stage/hmx.awaitvalue edges feedingscf::pipelineForLoop.Structural fixes:
matmul-to-hmxis gated on the same condition as the tile level below it:weight-resident/hmx-partition/hmx-workspace-residentall sit insideif (enableBufferization), while attribution sat outside it.hmx.matmulhas no other consumer andHmxToLLVMPassmarks the dialect illegal without a MatmulOp pattern, so attributing with bufferization off produced IR that cannot be lowered (a compile failure reachable from--no-enable-bufferization).hmx_*leaves in the kernel.hmxdialect ops collected before the conversion instead of a callee-name prefix, with a post-conversion check that fails the build if a leaf caller was not recognised -- a missed pair hangs the next thread forever inHAP_compute_res_hmx_lock.HmxEngineResource, non-addressable) instead of a genericMemWrite, so CSE no longer serialises it against unrelated memory, while the ops that really write addressable memory keep their default effects.backend/utils.py) instead of a 14-entry positional tuple read withlen(x) > iguards, which turned a stale cached JSON into "treat the feature as off" instead of an error.Passes.tddefaults forenable-split-reduction,enable-conversion-to-fp16andenable-weight-residentnow match the Python backend, so lit exercises the production configuration instead of a second set of defaults.HexagonProfilerno longer reboots the device from__init__(explicitstart()), and a missing environment variable is named along with how to set it instead of pointing at a file that does not exist.Behaviour outside the backend:
test_flash_attention.pyturns the HMX path on (enableConvertToHexagonmem) and checks an fp16 quantisation budget instead oftorch.allclose. It was hardcoded off, so this correctness test never exercised an HMX leaf; on device 39 383 -> 18 098 us (2.18x, PASS, rel ~7e-05)..gitignorekeeps the vendored triton trees, the tarballs and the in-tree logs out of the repository.Verified: 63/63 ninja targets; manual lit 184 pass / 1 pre-existing fail (
return_alloc_from_loop.mlir, untouched since 09-16) / 1 skip; standard 7-operator device batch green on the resulting build (libtriton.sof1beb6b0..., 2026-09-22 15:31) at 121 / 248 / 266 / 145 / 789 / 346 / 18098 us.