Qualcomm AI Engine Direct - Adding LPAI Custom Op Support and Example

[qti-horodnic]

Summary

Adding an end-to-end example of a custom PyTorch operator delegated to the LPAI backend running on the aDSP in direct mode.
Note: Custom ops are currently only supported in Direct Mode, not FastRPC Mode.

Changes required to enable custom ops:

• platform=HEXAGON - In direct mode the delegate is not __x86_64__ or __ANDROID__, so current_platform stayed UNKNOWN and every op package registration was silently skipped. Added a HEXAGON platform and an LPAI target to the compiler spec schema, and set current_platform from __hexagon__.
• registerOpPackage() - The 4th argument is backend specific. CPU/HTP take a processor target name; LPAI takes an optional target memory pool. The runtime now passes nullptr there and lets the backend pick its default.
• Deployment of the signed op package - the example signs the op package libraries via sign_library.sh and pushes them with SimpleADB's existing files= parameter. sign_library.sh gained an --op_package_dir option to sign them. On device an LPAI op package requires direct mode (registration over FastRPC is not supported), and both libQnnExampleLpaiOpPackage.so and libLpaiOpPackageIsland.so carry the kernel, so they must be rebuilt, signed and deployed together.

No partitioner change is needed: op package registration happens when the QnnManager for the compiler spec is created (InitBackend → BackendRegisterOpPackage), i.e. before partitioning, so validateOpConfig() already accepts an op package-backed node and it is validated like any other node.

Note: if the op were rejected it would fall back to the CPU implementation that the example registers for eager mode, which computes the same values, so the values would still match, and the test would pass even while the op package was never exercised. The example therefore inspects the lowered program and fails if the custom op is still present as a CPU operator, or if the graph contains no delegate at all. Both tests assert on that in addition to comparing the output.

This pr also fixes 2 existing bugs found, not specific to custom ops:

• QnnExecuTorchIdlWrapper's constructor could return early leaving method_ == nullptr, which execute_all() then dereferenced. The resulting DSP fault destroyed the real error message. It now reports Error::InvalidState.
• set_output_data_ptr() returning InvalidState for memory-planned outputs was treated as fatal even though it is documented as benign. Such outputs are now tracked and read back via method_->get_output().

LPAI op package conventions that are not currently documented in the SDK, and that the example kernel encodes with comments so the next author does not have to rediscover them:

• Scale may arrive as scale / 2^shift with shift > 31 (the eNPU reports shift=37). 1u << shift is undefined behavior for shifts >= 32 and evaluated to 0 here, turning the scale into +inf and every requantized value into NaN. The kernel uses ldexpf().
• getPerTensorQuantParams() reports biased storage: code = stored - offset, stored = code + offset, value = scale * code, with offset = -128 for 8-bit. Both directions must wrap modulo the storage width.
• getTensorDataType()'s signedness is unusable (it reports INT_8 for unsigned tensors); only the width can be trusted.
• Strides from getTensorLayout() are in bytes, and layoutOrder[0] is the slowest moving dimension (the last valid index is the fastest).
• getLayoutSupportFlag must be populated in the inference build as well, or on-device registration fails with AEE_EBADSTATE (0x8000040D).

Some more changes included in this pr:

• Added documentation for the flow in examples/qualcomm/custom_op/README.md.
• Added a custom_op_enablement.md agent skill covering op packages for both HTP and LPAI.
• Also corrected some stale flags in the QNN skill docs found along the way.
• sign_library.sh now checks the signer's exit status. Previously every invocation was yes | python $signer ... with no status check, so a failing elfsigner.py left the script exiting 0 and deployed stale or missing libraries.
• The DSP target the op package is built for is a separate --op_package_arch (default v79) rather than being derived from --htp_arch or --lpai_arch: on SM8850 the HTP is V81 and the LPAI hardware version is V6, while the op package builds for hexagon-V79.

Constraints:

• Qualcomm AI Engine Direct SDK >= 2.48 to build an op package at all (first release shipping QnnLpaiOpPackage.h and the LPAI op package makefiles), and >= 2.49 for the on-device path, which additionally requires direct mode.
• Only non-island mode is supported.
• The example op package declares 8-bit activations only (QNN_DATATYPE_UFIXED_POINT_8 / QNN_DATATYPE_UINT_8). A 16-bit activation arrives as QNN_DATATYPE_INT_32, and declaring INT_32 does pass validation and run, but only the first half of the output tensor comes back correct, reproducible with an identity requantization, which rules out the kernel's arithmetic and points at the backend's 16-bit tensor handover. That's an existing backend bug, unrelated to these changes.

Test plan

Two new tests in backends/qualcomm/tests/test_qnn_delegate.py:

• TestUtilsScript.test_custom_op_lpai builds the op package, signs it, deploys it, runs on device, and compares against the eager result.
• TestUtilsScript.test_custom_op_lpai_requant_edge_cases covers the two requantization paths the default run cannot reach, on the x86_64 simulator (the arithmetic is identical in both builds, so this avoids a DSP rebuild and re-sign): a small input whose code is biased into the upper half of the stored byte and has to be un-biased modulo the storage width, and an input above the calibrated range that has to saturate.

Verified on device:

python backends/qualcomm/tests/test_qnn_delegate.py \
  TestUtilsScript.test_custom_op_lpai \
  --executorch_root . --artifact_dir ./custom_op_lpai \
  --build_folder build-android --direct_build_folder build-direct \
  --backend lpai --soc_model SM8850 \
  --device 296753f6 \
  --host aisw-vm12-labsd

Passes: my_ops.mul3.default | True from the partitioner, _dom=adsp in the runner's domain URI, unique: [3.] with max abs err: 0.0, and 13.533 ms in qnn_executorch_execute_all.

Verified the requantization edge cases on the LPAI x86_64 simulator:

python backends/qualcomm/tests/test_qnn_delegate.py \
  TestUtilsScript.test_custom_op_lpai_requant_edge_cases \
  --executorch_root . --artifact_dir ./custom_op_lpai \
  --build_folder build-x86 --backend lpai --soc_model SM8850 \
  --enable_x86_64

Both subtests pass, each delegating the op: calibration=1.0, inference=0.25 matches eager, and calibration=1.0, inference=2.0 saturates to 3.0 (the graph's quantize node clamps the input to 1.0, so 3.0 is the correct answer rather than 6.0).

Also exercised the host-only path directly through the simulator:

python examples/qualcomm/custom_op/custom_ops_lpai.py \
  --build_folder build-x86 --backend lpai --soc_model SM8850 \
  --op_package_dir examples/qualcomm/custom_op/example_op_package_lpai/ExampleLpaiOpPackage \
  --build_op_package --enable_x86_64

Checked that the delegation assertion actually fails when the op is not delegated, by forcing a fallback with --skip_delegate_node_ops my_ops.mul3.default:

RuntimeError: ['my_ops::mul3'] was not delegated, it is executed by a CPU kernel in
method 'forward'. The op package was not exercised even if the output happens to match.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22659

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⚠️ 21 Awaiting Approval

As of commit 1f0dba0 with merge base db3390d ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux CUDA Wheels (gh)
• Build Aarch64 Linux Wheels (gh)
• Build documentation (gh)
• Build Linux CUDA Wheels (gh)
• Build Linux Wheels (gh)
• Build macOS Wheels (gh)
• Build Presets (gh)
• Build Windows Wheels (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[qti-horodnic]

@pytorchbot label "release notes: qualcomm"

[psiddh]

cc @mohankumarkumar @digantdesai @JacobSzwejbka

[psiddh]

@claude review this pr

[claude]

Claude finished @psiddh's task in 3m 58s —— View job

---

PR Review — LPAI Custom Op Support and Example

I read the runtime C++ changes, the export/schema plumbing, sign_library.sh, the example driver, the op package kernel, and the new tests. This is a genuinely high-quality, well-documented contribution — the rediscovered LPAI conventions captured as code comments (the ldexpf shift ≥ 32 UB, the biased-storage code = stored - offset, byte strides, getLayoutSupportFlag in the inference build) are the kind of thing that saves the next author days. The two incidental bug fixes are correct and cleanly scoped.

Below are the observations, none of which are blocking.

Correctness — looks solid

• QnnExecuTorchIdlWrapper::execute_all null guard (QnnExecuTorchIdlWrapper.cpp:174) — correct. The constructor returns early on several failure paths (method_meta fail :86, load_method fail :118, set_input fail :139) leaving method_ == nullptr, and execute_all dereferences it. Reporting Error::InvalidState instead of faulting is the right call.
• output_is_preallocated_ handling (QnnExecuTorchIdlWrapper.cpp:144-163, :268-281) — correct. Treating set_output_data_ptr returning non-Ok for memory-planned/constant outputs as benign and reading back via get_output(i) matches the documented contract, and the isTensor() check before toTensor() is appropriately defensive.
• BackendRegisterOpPackage (QnnBackendCommon.cpp) — the HEXAGON platform detection via __hexagon__ and the LPAI-specific nullptr 4th arg are correct, and the ordering (#elif before __ANDROID__) is fine since the DSP build defines neither __x86_64__ nor __ANDROID__.
• Kernel requantization (ExampleCustomOp_inference.c) — the ldexpf scale, the saturateCast NaN-safe ordering (!(value > min) catches NaN → returns min), clamping in the code domain rather than the biased-byte domain, and memcpy for the potentially-unaligned 16-bit path are all correct and the reasoning is well captured.

Minor suggestions

1. custom_ops_lpai.py — exception re-raise loses the traceback. At the very bottom:

   except Exception as e:
       if args.ip and args.port != -1:
           ...
       else:
           raise Exception(e)

   raise Exception(e) discards the original type and traceback, which hurts the standalone (no-IPC) debugging path this example is meant to support. A bare raise re-raises the original with its traceback intact. Fix this →

2. totalElements can overflow uint32_t (ExampleCustomOp_inference.c). It is accumulated as uint32_t across all dims. Fine for the example's (1,32,28,28), but a copy-paste into a kernel with a larger tensor could silently wrap. Since the comments explicitly invite reuse ("the interesting part to copy into a new kernel"), a one-word note or a uint64_t/size_t accumulator would make the template safer. Non-blocking.

3. get_lpai_device_lib_dir docstring drift. The suggested remediation string tells the user to run sign_library.sh --lpai_arch v<hw_ver> and "(add --direct_mode --htp_arch v<htp_arch> when deploying in direct mode)". Worth double-checking this matches the direct-mode invocation the rest of the PR standardizes on, so the error message doesn't point users at a subtly wrong command. Cosmetic.

4. Test IPC error surfacing. _run_custom_op_lpai launches the subprocess with stdout=subprocess.DEVNULL and relies on the conn.recv() JSON for the result. If the child dies before connecting to the Listener (e.g. an import error), listener.accept() / conn.recv() can hang or the failure mode is opaque. Not something to change here (it matches the existing TestUtilsScript pattern), just flagging that a crash-before-connect will present as a hang rather than a clear failure.

Things I checked and confirmed fine

• All imported symbols resolve: get_soc_to_lpai_hw_ver_map (utils/utils.py:1357), CustomOpsQuantAnnotator/IOQuantConfig (custom_op/annotator.py), QnnCustomOpPackageBuilder (custom_op/interface.py), and the test-class attributes host/device/direct_build_folder/enable_x86_64 (test_qnn_delegate.py:12008+).
• Schema .fbs and qc_schema.py enum additions (LPAI=3, HEXAGON=3) are kept in sync between the two files.
• sign_library.sh — the PIPESTATUS[1] check (rather than $?, which would be yes's SIGPIPE status) is correct, and the -e glob guard against a literal *.so is a nice touch.
• _assert_custom_op_delegated correctly closes the "eager fallback produces identical values" hole by inspecting the lowered program rather than trusting the numeric comparison — this is the crux of the test being meaningful, and it's done right.

Note for maintainers

The changed files include .claude/skills/qualcomm/ edits (SKILL.md, new custom_op_enablement.md, new_op_development.md) — worth a glance from whoever owns the agent-skill docs, as those are separate from the runtime review.

Overall: strong, careful work with excellent documentation of hard-won backend quirks. The suggestions above are all minor/optional.
· branch lpai_custom_ops_example

[qti-horodnic]

Addressed Claude's comments.