Skip to content

[hexagon-mlir] Land the HMX engine path - #101

Closed
Alfred (Dino1844) wants to merge 1 commit into
qualcomm:mainfrom
Dino1844:hmx
Closed

Alfred (Dino1844) wants to merge 1 commit into
qualcomm:mainfrom
Dino1844:hmx

Conversation

@Dino1844

Copy link
Copy Markdown

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.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant