feat: broaden Ascend AMP and dtype support

[lvyufeng]

AI Agent Information

• Agent/Tool: Claude Code CLI
• Model: Claude Opus 5
• Human Reviewer: @aoyulong
• Session Summary: After PR feat: add PrivateUse1 autocast and dtype support on NVIDIA H800/sm90 #127 added generic PrivateUse1 autocast policies, this change validates and extends AMP and dtype behavior on Ascend hardware. The implementation was developed from measured CANN 9.0 failures and verified against CPU reference semantics.

Summary

This change broadens Ascend AMP support and makes dtype behavior follow PyTorch semantics at CANN capability boundaries. It adds generated AMP unscale and foreach list-add routes, preserves float64 copies and casts, applies promotion-aware binary output allocation, and falls back to CPU for unsupported ACLNN unary and matmul dtype combinations. Comprehensive integration tests and backend dtype documentation are included.

Change Type

• Bug Fix
• New Feature
• Documentation
• Testing

Platforms Affected

• Ascend

Problem Analysis

What was broken/missing?

• GradScaler failed because _amp_foreach_non_finite_check_and_unscale_ had no Ascend PrivateUse1 implementation.
• Optimizer steps could fail because _foreach_add_.List was not routed on Ascend.
• Mixed tensor arithmetic such as float16 + float32 returned the first input dtype instead of PyTorch's promoted dtype.
• Ascend _to_copy clamped requested float64 storage to float32.
• CANN rejected specific dtype/operator combinations such as int16 negation and float64 matmul with error 161002.

Why did it happen?

The generic Ascend codegen assumed that ACL dtype-enum availability implied per-operator support and allocated binary outputs from self.options(). CANN 9.0 has narrower operator contracts than its dtype enum, while PrivateUse1 bypassed the normal native promotion and AMP auxiliary kernel paths for several operations.

Investigation process:

1. Read the PrivateUse1 autocast registration, Ascend codegen templates, copy path, matmul wrapper, existing CUDA/MUSA implementations, and generated dispatcher declarations.
2. Rebuilt current flagos/main at 81fd16c and reproduced the ACLNN int16 negation, float64 matmul, mixed promotion, and float64 copy failures on Ascend 910.
3. Measured a dtype-by-operator matrix and compared native results/errors with CPU PyTorch reference behavior.
4. Verified CANN boundaries with exact float64 round trips, integer unary sweeps, fallback probes, and generated-code idempotency.

Solution Design

Implementation approach:

• Add dtype_support.h with shared Ascend capability and fallback helpers.
• Extend scripts/codegen_ascend.py categories and templates so all Ascend generated output remains codegen-owned.
• Use at::result_type plus true-division promotion to allocate and marshal binary inputs correctly.
• Keep native ACLNN for supported floating-point paths; route unsupported integral/bool unary and matmul inputs through CPU reference operations and copy results back to Ascend.
• Remove the obsolete Ascend float64 clamp in _to_copy.
• Add AMP and dtype regression coverage and document measured Ascend boundaries.

Key design decisions:

• Do not advertise unsupported dtypes as AMP targets; AMP remains float16/bfloat16.
• Do not force unsupported dtypes through ACLNN. CPU fallback preserves correctness and gives a clear capability boundary.
• Use generator templates and regenerate ascend_kernels.cc; do not hand-edit generated kernels.
• Keep complex and quantized dtype support out of scope because the Ascend ACL tensor wrapper does not map them.

Code changes by file:

• scripts/codegen_ascend.py: add dtype-aware unary/matmul fallback, promotion-aware binary templates, AMP unscale, and foreach list-add categories.
• csrc/aten/backends/ascend/dtype_support.h: shared Ascend dtype capability and CPU-result movement helpers.
• csrc/aten/backends/ascend/generated/ascend_kernels.cc: regenerated kernels from the updated generator.
• csrc/aten/backends/ascend/matmul.cc: fallback for unsupported matmul dtype combinations.
• csrc/aten/copy_ops.cc: preserve requested float64 dtype on Ascend copies/casts.
• torch_fl/configs/backends_ascend.conf: generated registrations for AMP and foreach routes.
• tests/integration/test_amp.py: AMP out-overload regression coverage.
• tests/integration/test_dtype_coverage.py: factory, unary, binary, reduction, indexing, comparison, copy, promotion, and fallback coverage.
• docs/reference/dtype-support.md: per-backend dtype and AMP contract.
• docs/reference/operator-support.md: measured Ascend route evidence and update history.

Changes by commit:

1. 55cb6a2 - feat: broaden Ascend AMP and dtype support: implements generated Ascend AMP/dtype behavior, tests, and documentation.

Verification

Pre-submission Checklist

• Linting passed (ruff check, ruff format --check)
• Type checking passed (not applicable)
• All tests pass (the full unit suite has one unrelated environment failure described below)
• Manual testing completed (Ascend 910 / CANN 9.0 reproduction probes)
• No debug/temporary code
• Documentation updated
• Commit messages follow conventions
• All text in English

Linting Results

$ ruff check .
All checks passed!

$ ruff format --check .
177 files already formatted

Test Results

$ ASCEND_RT_VISIBLE_DEVICES=2 pytest tests/integration/test_amp.py tests/integration/test_dtype_coverage.py -q
199 passed in 2.08s

$ pytest tests/unit/test_profiler_privateuse1.py tests/unit/test_vendor_routing.py -q
26 passed, 3 skipped, 1 warning

$ pytest tests/unit/ -q
1 failed, 95 passed, 29 skipped, 1 warning

The single full-unit failure is the pre-existing tests/unit/bpu/test_x86_env.py::test_env_carries_library_path_and_stubs assertion. The current shell exports Ascend paths in PYTHONPATH; the test expects the isolated temporary HBDK stubs path. The affected BPU test passes in its intended clean environment and no BPU files are changed by this PR.

Manual Verification

$ ASCEND_RT_VISIBLE_DEVICES=2 python dtype_probe.py
neg int16       -> int16 CPU-reference parity
uint8 neg       -> uint8 wraparound parity
fp16 + fp32     -> torch.float32
int16 / int64   -> torch.float32
fp64 copy       -> exact round trip, including 1e300
fp64 matmul     -> CPU fallback, exact result parity

$ python scripts/codegen_ascend.py --no-conf
second generation diff: empty (idempotent)

Code Quality Verification

Style Consistency

• Matched existing code style in modified files
• Followed naming conventions and existing dispatcher patterns
• Comment density matches surrounding code
• Reused existing copy, dispatch, and CPU fallback utilities

Edge Cases Considered

1. Integer and bool unary operations, including uint8 negation wraparound and bool negation errors.
2. Mixed float and integer promotion, including integral true division.
3. Float64 storage, device-to-device cast, CPU/device round trip, and matmul fallback.
4. AMP finite and non-finite GradScaler paths, including the .out overload.
5. Generator idempotency and CANN unsupported dtype combinations.

Potential Risks

1. CPU fallback adds device-to-host and host-to-device transfer overhead for unsupported dtype/operator combinations.
2. CANN capability boundaries can vary by toolkit release; the documented matrix is measured against CANN 9.0 on Ascend 910.

Rollback Plan

Revert commit 55cb6a2. This removes the new generated routes and restores the prior Ascend dtype behavior without changing CUDA, MetaX, DCU, or MUSA code.

Related Work

• Related to feat: add PrivateUse1 autocast and dtype support on NVIDIA H800/sm90 #127 (PrivateUse1 autocast and dtype support)
• Builds on feat: add Ascend MSPTI profiler integration #126 (Ascend MSPTI profiler integration)

Explicitly Not Included

• Complex dtype support.
• Quantized dtype support.
• Claiming uniform dtype parity across CUDA, Ascend, MetaX, DCU, and MUSA.
• Native ACLNN implementations for dtype/operator combinations rejected by CANN.

Human Review Notes

Areas needing special attention:

1. Review BinaryResultType and generated binary templates for promotion behavior across all routed binary categories.
2. Review the CPU fallback boundaries in dtype_support.h against future CANN toolkit releases.
3. Review generated AMP unscale and foreach list-add registration signatures against the PyTorch 2.10 dispatcher schemas.

Questions for reviewer:

1. Should the Ascend dtype capability table be expanded with additional measured CANN operator-specific entries in a follow-up?
2. Should a dedicated clean-environment BPU CI job be added to prevent host PYTHONPATH contamination from affecting the existing x86 environment test?

---

🤖 Generated with Claude Code