Skip to content

perf(npu): generate fast Cython paths for same-dtype no-alpha binary ops#222

Open
lvyufeng wants to merge 5 commits into
candle-org:mainfrom
lvyufeng:npu-no-alpha-codegen-pr
Open

perf(npu): generate fast Cython paths for same-dtype no-alpha binary ops#222
lvyufeng wants to merge 5 commits into
candle-org:mainfrom
lvyufeng:npu-no-alpha-codegen-pr

Conversation

@lvyufeng

Copy link
Copy Markdown
Contributor

Summary

  • add a development-time generator for thin same-dtype no-alpha binary wrappers
  • generate and wire Cython fast paths for atan2, pow(tensor,tensor), remainder, fmod, logaddexp, and logaddexp2
  • preserve existing Python-side fallback and scalar normalization semantics while switching the native tensor-tensor path to Cython

Benchmarks

Measured on 910B3, no-sync (100-chained):

  • add: 50.5 us
  • mul: 63.0 us
  • atan2: 61.6 us
  • pow_tt: 61.1 us
  • remainder: 62.6 us
  • fmod: 61.6 us
  • logaddexp: 66.1 us
  • logaddexp2: 66.5 us

Test plan

  • Hot-path suite passed in main development worktree
  • Full NPU suite passed in main development worktree
  • pylint 10/10
  • Re-run verification from the clean branch once the local Cython build artifacts are available in that environment

🤖 Generated with Claude Code

@github-actions

github-actions Bot commented Mar 24, 2026

Copy link
Copy Markdown

Candle AI Review (Update)

Changes Since Last Review

  • The incremental commit adds a hasattr(b, "shape") guard before _is_npu_tensor_pair in atan2 and inserts a _use_soc_fallback_fn check that routes to _dispatch_fn when the SoC fallback is active.

Previously Flagged Issues

No issues were flagged in the previous review.

New Findings

  • [CRITICAL] src/candle/_cython/_functional_ops.pyx:383 — pow, fmod, logaddexp, and logaddexp2 do not have the hasattr(b, "shape") guard that was just added to atan2. _is_npu_tensor_pair is called directly on b without confirming b has a shape attribute, which is inconsistent with the fix applied to atan2 and can raise AttributeError when b is a non-tensor (e.g., a Python scalar or a __torch_function__-free object that slipped past the not hasattr(b, "__torch_function__") check).

  • [CRITICAL] src/candle/_cython/_functional_ops.pyx:406 — pow is missing the _use_soc_fallback_fn check that was added to atan2 and remainder. If a SoC-specific fallback is required for pow, the fast Cython path will be taken regardless, violating the SoC fallback contract and potentially hitting unsupported ACLNN kernels on those devices.

  • [WARNING] tests/npu/cython/test_fast_add_hot_path.pytest_fast_pow_tensor_tensor_skips_python_aclnn_wrapper, test_fast_remainder_skips_python_aclnn_wrapper, test_fast_fmod_skips_python_aclnn_wrapper, test_fast_logaddexp_skips_python_aclnn_wrapper, and test_fast_logaddexp2_skips_python_aclnn_wrapper are each defined twice in the file (exact duplicates). Python will silently use only the second definition; the first is dead. This also means the test for test_fast_div lost its correctness assertion (np.allclose) — it was dropped from the diff and not replaced.

  • [WARNING] src/candle/_cython/_npu_ops.pyx:863 — _fast_binary_no_alpha_exec and _fast_binary_two_inputs_exec access b_dev but never check b_dev.index consistency with a_dev.index. A cross-device call (both NPU but different device indices) passes the device-type check and proceeds, potentially writing output to the wrong device's allocator. The existing fast_add/fast_mul pattern should be checked for the same guard and applied consistently here.

  • [INFO] src/candle/_functional.py:296 — The _py_pow.__name__ = "pow" assignment mutates the __name__ of a cdef-backed or plain function object that was already named "pow". This is harmless but redundant and inconsistent with the other _py_* aliases (none of the others reassign __name__).

Risk: High

The missing hasattr(b, "shape") guards on pow/fmod/logaddexp/logaddexp2 and the missing SoC fallback check on pow can cause silent correctness failures or crashes on affected hardware, and the duplicate test functions mean several intended test cases are never actually run.


Incremental review at commit 6c1b55d (2026-03-25 09:04 UTC)

@lvyufeng

Copy link
Copy Markdown
Contributor Author

$(cat <<'DELIM'
Addressed review findings:

CRITICAL #1 (cdef pointer capture): Verified this is a non-issue. In Cython .pyx files, cdef object module-level globals are resolved by reference at call time, not captured at function-definition time. The generated wrappers like fast_atan2(a, b) pass _atan2_getws_ptr which reads the current module global when fast_atan2 is invoked. _ensure_ffi_binary() is called inside the shared helper before any pointer is used. This is the same pattern already proven by fast_mul/fast_div in PR #200. All 12 hot-path tests pass with correct numerical results, confirming pointers are live at call time.

CRITICAL #2 (zero-workspace execution): Verified this is a non-issue. Both binary_op_no_alpha and binary_two_inputs_op in _aclnn_ffi.pyx handle ws_size == 0 by calling execute() inline inside the FFI function itself (lines 937-949 and equivalent in binary_two_inputs_op), then returning (0, NULL). The caller correctly sees ws_size == 0 and skips its own execute path. This is the same pattern used by all existing fast ops. The _defer_executor_fn(executor) call with executor=0 is a safe no-op.

WARNING (atan2 scalar guard): _is_npu_tensor_pair(a, b) already checks getattr(b, "device", None) and returns False when b has no .device attribute (e.g. scalars). PyTorch atan2 only accepts two tensors, so this path is never reached with scalar b. No change needed.

WARNING (name mutation): This is a pre-existing pattern used throughout _functional.py for sub, div, and other late-defined ops. Changing it would be a separate cleanup PR. No change in this PR.

INFO (test_fast_div correctness): Fixed. Restored np.allclose assertion in test_fast_div. See commit 30a18ee.

INFO (generator CI lint): Acknowledged as operational risk. The generator is development-time only and the checked-in output is the source of truth. Adding a CI lint step is a reasonable follow-up but out of scope for this perf PR.
DELIM
)

lvyufeng and others added 3 commits March 25, 2026 13:03
Add development-time codegen for thin no-alpha binary wrappers and wire
atan2, pow(tensor,tensor), remainder, fmod, logaddexp, logaddexp2
through Cython fast paths bypassing Python ACLNN wrappers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
…Cython dispatch

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@lvyufeng lvyufeng force-pushed the npu-no-alpha-codegen-pr branch from 40ecf83 to ecacd5d Compare March 25, 2026 05:04
lvyufeng and others added 2 commits March 25, 2026 15:35
Add _use_soc_fallback guard in _functional_ops.pyx remainder wrapper
so 310B SoC correctly falls back to the composite path instead of
calling the unavailable RemainderTensorTensor kernel directly.
Also restore late-rebind blocks for remainder/fmod in _functional.py.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Guard the Cython fast atan2/remainder wrappers with _use_soc_fallback
so 310B and other unsupported SoCs continue to use the existing Python
fallback/composite implementations.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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