feat: add B10 math and fused ops - #2191
Conversation
📝 WalkthroughWalkthroughAdds ChangesFused tile math operations
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PythonAPI as pypto.language.op.tile_ops
participant IR as tile_ops.py
participant Backend as pto_ops_elementwise.cpp
participant RuntimeTests as test_math_fused.py
PythonAPI->>IR: construct axpy/add_relu/pow/pows
IR->>Backend: lower tile operation
Backend-->>RuntimeTests: emit PTO instruction
RuntimeTests->>RuntimeTests: execute and compare expected tensors
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 080aa62f2c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
6e46dce to
cdc163a
Compare
|
按 #2166 复核当前 head
其中第 1 项是固定 toolchain 下的硬 blocker，不能用跳过 PTOAS 的 UT 作为替代验证。 |
cdc163a to
dc090f8
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 13
🧹 Nitpick comments (1)
python/pypto/ir/op/tile_ops.py (1)
951-951: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRuff A001:
powshadows a builtin.The name intentionally mirrors the
tile.powop, so silence it in config rather than renaming or adding an inline# noqa. Based on learnings, prefer fixing this inruff.toml/pyproject.tomlper-file-ignoresforpython/pypto/**instead of inline suppressions.🔧 Suggested config entry
[lint.per-file-ignores] "python/pypto/ir/op/tile_ops.py" = ["A001"] # op names mirror tile.* IR opcodes🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/pypto/ir/op/tile_ops.py` at line 951, Suppress Ruff A001 for the intentional pow symbol in the project lint configuration, preferably by adding the per-file ignore for python/pypto/ir/op/tile_ops.py in ruff.toml or pyproject.toml. Keep the pow function name unchanged and do not add an inline noqa.Sources: Learnings, Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/en/dev/codegen/00-pto_codegen.md`:
- Around line 145-148: Update the mapping-table rows for tile.pow and tile.pows
in both docs/en/dev/codegen/00-pto_codegen.md (lines 145-148) and
docs/zh/dev/codegen/00-pto_codegen.md (lines 143-146) to state that tmp is
optional and only required for floating-point inputs, being rejected for integer
inputs. In the tile.add_relu rows at both sites, add the caveat that
pto.taddrelu is emission-only and rejected by the pinned PTOAS on A2/A3.
In `@docs/en/dev/ptoas-op-status.md`:
- Around line 129-130: Update the ST status cells for the pto.tpow and pto.tpows
rows to ❌, since current evidence is limited to a partial A2 floating/tmp
hardware path. Preserve the existing Notes text documenting that partial A2/A3
verification and UT-only coverage, and do not mark ST ✅ until the full dtype ×
precision × target × tail matrix has hardware execution evidence.
- Line 126: Update the pto.taddrelu matrix note to reference the pinned
PTOAS_VERSION v0.48, unless the row intentionally targets a newer upstream
release; in that case, explicitly clarify the reason for the v0.53 reference.
In `@docs/zh/dev/ptoas-op-status.md`:
- Around line 115-116: Update the ST status fields for pto.tpow and pto.tpows to
remain negative (❌), while preserving the existing A2/A3 floating-point
validation remarks and UT coverage notes. Do not mark system-test coverage as
complete until the full dtype, precision, target, and tail matrix has run.
In `@python/pypto/language/op/tile_ops.py`:
- Around line 2057-2060: Add documented per-file Ruff configuration ignores for
the intentional pow DSL API: ignore A001 in python/pypto/language/op/tile_ops.py
for pow, and ignore A004 in python/pypto/language/op/__init__.py and
python/pypto/language/__init__.py for its re-exports. Use linter configuration
rather than inline noqa suppressions.
In `@src/backend/common/pto_ops_elementwise.cpp`:
- Around line 533-540: Do not expose or emit tile.add_relu while PTOAS v0.48 is
pinned and lacks support for pto.taddrelu. Update the tile.add_relu registration
in src/backend/common/pto_ops_elementwise.cpp and its corresponding API in
python/pypto/language/op/tile_ops.py to gate or remove the unsupported path, and
update docs/zh/dev/ptoas-op-status.md to keep the documented status consistent;
skip_ptoas must not be treated as sufficient protection.
In `@src/ir/op/tile_ops/elementwise.cpp`:
- Around line 493-497: Update DeduceTileAxpyType and the tile.axpy registration
to admit BF16 source, scalar, and destination types only when targeting A5,
while preserving existing FP16/FP32 and dtype-matching rules and rejecting BF16
on unsupported targets. Add coverage for valid A5 BF16 tile.axpy programs.
- Around line 511-543: Update DeduceTilePowType in
src/ir/op/tile_ops/elementwise.cpp:511-543 to reject BF16 in default precision
and permit BF16 with high_precision only when the target architecture is A5,
preserving existing dtype and tmp validation. Update the tile.pow/tile.pows
registration in src/backend/common/pto_ops_elementwise.cpp:542-547 to require
row-major layout for every present base, exponent, output, and floating-point
tmp operand. Add negative coverage for BF16 on non-A5 targets, high-precision
BF16 on non-A5 targets, and non-row-major operands.
In `@tests/st/runtime/ops/test_math_fused.py`:
- Around line 418-443: Add BF16 coverage to the axpy cases assembled in _CASES,
including at least a BF16-to-BF16 case for platform A5 if TAXPY supports it. If
BF16 is unsupported, update the A5 IR type validation to reject it with a
target-specific message and add a negative test asserting that rejection.
- Around line 371-372: Update the axpy expected-value calculation in the test’s
op handling so multiplication and addition occur in FP32, then cast the final
result once to the source tensor dtype. Preserve the existing FP16 behavior
expectations while avoiding intermediate FP16 rounding.
- Around line 43-62: Update the test dtype mapping and unsigned integer tensor
generation around _TORCH_DT and _INTEGER_DTYPES to avoid torch.uint16 and
torch.uint32. Create unsigned test inputs using a wider supported signed dtype
such as torch.int64, and cast device outputs to the corresponding comparison
dtype before asserting results, while preserving coverage for DataType.UINT16
and DataType.UINT32.
In `@tests/ut/ir/operators/test_math_fused_ops.py`:
- Around line 89-98: Expand test_pow_rejects_wrong_tmp_contract to cover
target-aware BF16 default-precision rejection on A2/A3 and non-row-major layout
rejection for base, exp, dst, and caller-supplied tmp tiles. Add cases that also
validate incompatible tmp shape/layout is rejected before codegen, using the
existing tile.pow/tile.pows APIs and matching the established ValueError
behavior.
- Around line 124-127: Update the tile.add_relu coverage around AddReluProgram
to expose the current PTOAS limitation: either gate the public API and emission
based on the supported toolchain version, or add a negative compile test that
explicitly asserts compilation fails with the pinned PTOAS on A2/A3. Do not
leave the existing skip_ptoas=True emission-only assertion as the sole
validation.
---
Nitpick comments:
In `@python/pypto/ir/op/tile_ops.py`:
- Line 951: Suppress Ruff A001 for the intentional pow symbol in the project
lint configuration, preferably by adding the per-file ignore for
python/pypto/ir/op/tile_ops.py in ruff.toml or pyproject.toml. Keep the pow
function name unchanged and do not add an inline noqa.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 53ae84ff-09cf-493c-a603-edf3ea806689
📒 Files selected for processing (14)
docs/en/dev/codegen/00-pto_codegen.mddocs/en/dev/ptoas-op-status.mddocs/zh/dev/codegen/00-pto_codegen.mddocs/zh/dev/ptoas-op-status.mdpython/pypto/debug/torch_codegen.pypython/pypto/ir/op/tile_ops.pypython/pypto/language/__init__.pypython/pypto/language/op/__init__.pypython/pypto/language/op/tile_ops.pysrc/backend/common/pto_ops_elementwise.cppsrc/ir/op/tile_ops/elementwise.cpptests/st/runtime/ops/test_math_fused.pytests/ut/debug/test_torch_codegen.pytests/ut/ir/operators/test_math_fused_ops.py
| | `tile.axpy(src, scalar, dst)` | `pto.taxpy` (reuses `dst`) | | ||
| | `tile.add_relu(src0, src1)` | `pto.taddrelu` | | ||
| | `tile.pow(base, exp, tmp, high_precision=...)` | `pto.tpow` | | ||
| | `tile.pows(base, exp, tmp, high_precision=...)` | `pto.tpows` | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Mapping table rows replicated in both languages omit the same two contract details. The new rows present tmp as unconditional and pto.taddrelu as fully supported; in reality tmp is required only for floating-point inputs (and rejected for integer inputs), and pto.taddrelu is rejected by the pinned PTOAS on A2/A3.
docs/en/dev/codegen/00-pto_codegen.md#L145-L148: marktmpoptional/float-only in thetile.pow/tile.powsrows and add an "emission only; rejected by the pinned PTOAS on A2/A3" caveat to thetile.add_relurow.docs/zh/dev/codegen/00-pto_codegen.md#L143-L146: apply the same two wording changes so the Chinese table stays in sync.
📍 Affects 2 files
docs/en/dev/codegen/00-pto_codegen.md#L145-L148(this comment)docs/zh/dev/codegen/00-pto_codegen.md#L143-L146
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/en/dev/codegen/00-pto_codegen.md` around lines 145 - 148, Update the
mapping-table rows for tile.pow and tile.pows in both
docs/en/dev/codegen/00-pto_codegen.md (lines 145-148) and
docs/zh/dev/codegen/00-pto_codegen.md (lines 143-146) to state that tmp is
optional and only required for floating-point inputs, being rejected for integer
inputs. In the tile.add_relu rows at both sites, add the caveat that
pto.taddrelu is emission-only and rejected by the pinned PTOAS on A2/A3.
| | pto.trelu | TRELU | tile | ✅ | ✅ | ❌ | ✅ | — | | | ||
| | pto.tlrelu | TLRELU | tile | ✅ | ✅ | ❌ | ✅ | — | | | ||
| | pto.taddrelu | VADDRELU | tile | ✅ | ❌ | ❌ | ❌ | — | MISSING: lacks a complete frontend/codegen/ST path | | ||
| | pto.taddrelu | VADDRELU | tile | ✅ | ✅ | ❌ | ❌ | — | exact emission covered; PTOAS v0.53 A2/A3 legalization still marks the op illegal | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -n -i 'ptoas' --glob '!docs/**' -g '*.toml' -g '*.txt' -g '*.cmake' -g 'CMakeLists.txt' -g '*.json' -g '*.yaml' -g '*.yml' .Repository: hw-native-sys/pypto
Length of output: 208
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "tracked files containing ptoas (case insensitive):"
git ls-files -z | xargs -0 grep -n -i 'ptoas' || true
echo
echo "candidate files:"
fd -i 'ptoas|readme|install|setup|pyproject|poetry|requirements|mkdocs' . -t f | sed -n '1,120p'
echo
echo "docs line context:"
sed -n '100,145p' docs/en/dev/ptoas-op-status.md
echo
echo "repo root files:"
find . -maxdepth 2 -type f \( -iname 'readme*' -o -iname 'install*' -o -iname 'pyproject.toml' -o -name 'poetry.lock' -o -name 'requirements*.txt' -o -iname '*.cmake' \) | sort | sed -n '1,160p'Repository: hw-native-sys/pypto
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "toolchain versions.env:"
if [ -f toolchain/versions.env ]; then
sed -n '1,220p' toolchain/versions.env | rg -n 'PTOAS|PTO_ISA|VERSION|sha256' || true
else
echo "toolchain/versions.env not found"
fd -i 'versions.env' . -t f | sed -n '1,20p'
fi
echo
echo "docs matrix header/context:"
sed -n '1,70p' docs/en/dev/ptoas-op-status.md
echo
echo "all version/env-like files:"
fd -i 'versions|env|config|version|toolchain|readme|readme.*|install|setup|tool' . -t f -d 3 | sort | sed -n '1,200p'Repository: hw-native-sys/pypto
Length of output: 5421
Update the pinned PTOAS version in the matrix note.
toolchain/versions.env pins PTOAS_VERSION=v0.48, while this row still cites “PTOAS v0.53”; use the pinned version or clarify why this matrix targets a newer upstream version than CI uses.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/en/dev/ptoas-op-status.md` at line 126, Update the pto.taddrelu matrix
note to reference the pinned PTOAS_VERSION v0.48, unless the row intentionally
targets a newer upstream release; in that case, explicitly clarify the reason
for the v0.53 reference.
| | pto.tpow | TPOW | tile | ✅ | ✅ | ❌ | ✅ | — | floating/tmp path verified on A2/A3 hardware; integer and A5 high-precision contracts covered by UT | | ||
| | pto.tpows | TPOWS | tile | ✅ | ✅ | ❌ | ✅ | — | floating/tmp path verified on A2/A3 hardware; integer and A5 high-precision contracts covered by UT | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
ST ✅ overstates the evidence for tpow/tpows.
The notes on these same rows concede that integer dtypes and A5 high precision are only UT-covered, and the reported hardware run covers a limited A2 path. Under this file's own criteria (L17-24, ST✅ = same-name execution evidence on hardware), the row should stay ❌ until the full dtype × precision × target × tail matrix has executed, with the partial A2 result recorded in Notes.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/en/dev/ptoas-op-status.md` around lines 129 - 130, Update the ST status
cells for the pto.tpow and pto.tpows rows to ❌, since current evidence is
limited to a partial A2 floating/tmp hardware path. Preserve the existing Notes
text documenting that partial A2/A3 verification and UT-only coverage, and do
not mark ST ✅ until the full dtype × precision × target × tail matrix has
hardware execution evidence.
| | pto.tpow | TPOW | tile | ✅ | ✅ | ❌ | ✅ | — | 浮点/tmp 路径已通过 A2/A3 真机验证；整数与 A5 high-precision 契约由 UT 覆盖 | | ||
| | pto.tpows | TPOWS | tile | ✅ | ✅ | ❌ | ✅ | — | 浮点/tmp 路径已通过 A2/A3 真机验证；整数与 A5 high-precision 契约由 UT 覆盖 | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Keep TPOW/TPOWS system-test status negative.
The remarks only claim A2/A3 floating-point runtime coverage; integer and A5 high-precision cases are UT-only. Marking ST as ✅ overstates validation until the full dtype, precision, target, and tail matrix runs.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@docs/zh/dev/ptoas-op-status.md` around lines 115 - 116, Update the ST status
fields for pto.tpow and pto.tpows to remain negative (❌), while preserving the
existing A2/A3 floating-point validation remarks and UT coverage notes. Do not
mark system-test coverage as complete until the full dtype, precision, target,
and tail matrix has run.
| def pow(base: Tile, exp: Tile, tmp: Tile | None = None, *, high_precision: bool = False) -> Tile: | ||
| """Element-wise tile power; floating-point inputs require ``tmp``.""" | ||
| tmp_expr = None if tmp is None else tmp.unwrap() | ||
| return Tile(expr=_ir_ops.pow(base.unwrap(), exp.unwrap(), tmp_expr, high_precision=high_precision)) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Resolve the new Ruff builtin-shadowing errors via configuration. The public pl.pow API triggers A001/A004 in its definition and both re-exports.
python/pypto/language/op/tile_ops.py#L2057-L2060: add a documented per-file A001 ignore for the intentional DSL API name.python/pypto/language/op/__init__.py#L96-L97: add the corresponding per-file A004 ignore.python/pypto/language/__init__.py#L143-L144: add the corresponding per-file A004 ignore.
Based on learnings, prefer fixing root causes in linter configuration rather than using inline # noqa suppressions.
🧰 Tools
🪛 Ruff (0.16.0)
[error] 2057-2057: Variable pow is shadowing a Python builtin
(A001)
📍 Affects 3 files
python/pypto/language/op/tile_ops.py#L2057-L2060(this comment)python/pypto/language/op/__init__.py#L96-L97python/pypto/language/__init__.py#L143-L144
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/pypto/language/op/tile_ops.py` around lines 2057 - 2060, Add
documented per-file Ruff configuration ignores for the intentional pow DSL API:
ignore A001 in python/pypto/language/op/tile_ops.py for pow, and ignore A004 in
python/pypto/language/op/__init__.py and python/pypto/language/__init__.py for
its re-exports. Use linter configuration rather than inline noqa suppressions.
Sources: Learnings, Linters/SAST tools
| _TORCH_DT = { | ||
| DataType.INT8: torch.int8, | ||
| DataType.UINT8: torch.uint8, | ||
| DataType.INT16: torch.int16, | ||
| DataType.UINT16: torch.uint16, | ||
| DataType.INT32: torch.int32, | ||
| DataType.UINT32: torch.uint32, | ||
| DataType.FP16: torch.float16, | ||
| DataType.FP32: torch.float32, | ||
| DataType.BF16: torch.bfloat16, | ||
| } | ||
| _FLOAT_DTYPES = {DataType.FP16, DataType.FP32, DataType.BF16} | ||
| _INTEGER_DTYPES = [ | ||
| DataType.INT8, | ||
| DataType.UINT8, | ||
| DataType.INT16, | ||
| DataType.UINT16, | ||
| DataType.INT32, | ||
| DataType.UINT32, | ||
| ] |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
fd -H -t f -g 'pyproject.toml' -g 'requirements*.txt' . -x rg -n -i 'torch' {}Repository: hw-native-sys/pypto
Length of output: 270
🌐 Web query:
PyTorch torch.uint16 torch.uint32 which version added limited support
💡 Result:
PyTorch introduced limited support for torch.uint16, torch.uint32, and torch.uint64 in version 2.2 [1]. These unsigned integer dtypes were primarily added to improve interoperability and support usage with torch.compile rather than for general-purpose eager mode operations [2][3]. Because support remains limited, many operations—such as standard binary arithmetic (add, sub, mul)—may not be fully implemented across all backends (CPU, MPS, etc.), often resulting in NotImplementedError exceptions [4]. Key development context: - Initial introduction: The types were made available in the core codebase through pull request #116594, which was integrated for the PyTorch 2.2 release cycle [1]. - Status: Official documentation continues to label these types as having limited support [2][5]. They are designed to assist with specific workflows rather than replacing signed integer types for standard eager execution [2]. - Ongoing development: As of 2026, work continues to expand support, such as adding JIT CUDA kernel support [6] and addressing missing operator implementations [4].
Citations:
- 1: Add unsigned integer dtypes to PyTorch pytorch/pytorch#116594
- 2: https://docs.pytorch.org/docs/2.5/tensors.html
- 3: https://docs.pytorch.org/docs/2.11/tensor%5Fattributes.html
- 4: [CPU] Binary ops (add, sub, mul) not implemented for uint16/uint32/uint64 pytorch/pytorch#176298
- 5: https://docs.pytorch.org/docs/2.7/_sources/tensors.rst.txt
- 6: Add uint16, uint32, uint64 support to JIT CUDA kernels pytorch/pytorch#174303
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate test file =="
fd -a 'test_math_fused.py' . | sed 's#^\./##'
file="$(fd 'test_math_fused.py' . | head -n1)"
if [ -n "${file:-}" ]; then
wc -l "$file"
sed -n '1,140p' "$file" | cat -n
fi
echo
echo "== pytorch dependency files =="
git ls-files | rg '(^|/)(pyproject.toml|requirements.*\.txt|environment.*\.ya?ml|setup\.cfg|Pipfile)$' || true
for f in $(git ls-files | rg '(^|/)(pyproject.toml|requirements.*\.txt|environment.*\.ya?ml|setup\.cfg|Pipfile)$'); do
echo "--- $f ---"
rg -n -i 'torch' "$f" || true
done
echo
echo "== torch available? =="
python3 - <<'PY'
try:
import torch
print("torch.version", getattr(torch, "__version__", None))
print("torch uint16 attr", getattr(torch, "uint16", "<missing>"))
print("torch uint32 attr", getattr(torch, "uint32", "<missing>"))
print("torch uint64 attr", getattr(torch, "uint64", "<missing>"))
except Exception as e:
print("import_error", type(e).__name__, str(e))
PYRepository: hw-native-sys/pypto
Length of output: 6445
🌐 Web query:
PyTorch torch 2.2 torch.uint16 torch.uint32 available
💡 Result:
In PyTorch 2.2, torch.uint16 and torch.uint32 were not available as native data types [1][2]. These unsigned integer types (along with torch.uint64) were introduced in subsequent versions of PyTorch (starting with PyTorch 2.3) [3]. Even following their introduction, they are characterized as having limited support, primarily intended to assist with operations in torch.compile rather than broad eager-mode usage [4]. Users requiring unsigned integer support beyond torch.uint8 in older versions are typically advised to use signed integer variants (such as torch.int16 or torch.int32) [4].
Citations:
- 1: https://docs.pytorch.org/docs/2.2/tensor%5Fattributes.html
- 2: https://docs.pytorch.org/docs/2.2/tensors.html
- 3: torch.is_signed on new uint dtypes raises Unknown ScalarType pytorch/pytorch#125124
- 4: https://docs.pytorch.org/docs/2.5/tensors.html
🌐 Web query:
site:docs.pytorch.org torch.uint16 dtype PyTorch
💡 Result:
In PyTorch, torch.uint16 is a supported data type, though it currently has limited support [1]. Key details regarding torch.uint16 include: 1. Support Limitations: The dtype has limited support in eager mode. It primarily exists to assist with usage in torch.compile [1]. 2. Recommendations: Due to these limitations, if you require eager mode support and do not strictly need the extra range provided by the unsigned 16-bit integer, PyTorch documentation recommends using the signed variant, torch.int16, instead [1]. 3. Context of Use: While it can be found in certain contexts (for example, torchvision.io.decode_image may output a uint16 tensor when handling 16-bit PNGs), it is often advised to convert these tensors to other types like float or uint8 using functions such as torchvision.transforms.v2.functional.to_dtype for broader compatibility [2].
Citations:
- 1: https://docs.pytorch.org/docs/2.7/_sources/tensors.rst.txt
- 2: https://docs.pytorch.org/vision/stable/generated/torchvision.io.decode_image.html
Avoid relying on torch.uint16 / torch.uint32 in this ST path.
torch>=2.0.0 can resolve to a release where these dtypes are unavailable; even where available, PyTorch documents uint16/uint32 as limited-support dtypes, so eager cases like .to(torch.uint16)/.to(torch.uint32) are easier to trigger collection/runtime failures. Generate the unsigned test tensors via a wider signed dtype (for example int64) and compare after casting the device output.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/st/runtime/ops/test_math_fused.py` around lines 43 - 62, Update the
test dtype mapping and unsigned integer tensor generation around _TORCH_DT and
_INTEGER_DTYPES to avoid torch.uint16 and torch.uint32. Create unsigned test
inputs using a wider supported signed dtype such as torch.int64, and cast device
outputs to the corresponding comparison dtype before asserting results, while
preserving coverage for DataType.UINT16 and DataType.UINT32.
| if self.op_name == "axpy": | ||
| expected = src1 + src0 * self.scalar |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
FP16 reference is computed entirely in FP16.
src1 + src0 * self.scalar evaluates in the tensor dtype, so the FP16→FP16 axpy case double-rounds (multiply, then add) while the device may accumulate in FP32. If the harness tolerance is tight this will be flaky. Computing the reference in FP32 and casting once at the end is the safer form.
♻️ Suggested change
- expected = src1 + src0 * self.scalar
+ expected = src1.to(torch.float32) + src0.to(torch.float32) * self.scalar📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if self.op_name == "axpy": | |
| expected = src1 + src0 * self.scalar | |
| if self.op_name == "axpy": | |
| expected = src1.to(torch.float32) + src0.to(torch.float32) * self.scalar |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/st/runtime/ops/test_math_fused.py` around lines 371 - 372, Update the
axpy expected-value calculation in the test’s op handling so multiplication and
addition occur in FP32, then cast the final result once to the source tensor
dtype. Preserve the existing FP16 behavior expectations while avoiding
intermediate FP16 rounding.
| for _src_dtype, _dst_dtype, _scalar in ( | ||
| (DataType.FP16, DataType.FP16, -1.5), | ||
| (DataType.FP16, DataType.FP32, 0.0), | ||
| (DataType.FP32, DataType.FP32, 2.0), | ||
| ): | ||
| _CASES.append( | ||
| _case( | ||
| _platform, | ||
| "axpy", | ||
| _src_dtype, | ||
| _dst_dtype, | ||
| COMBINED_TAIL, | ||
| _scalar, | ||
| ) | ||
| ) | ||
| for _valid_shape in (FULL, ROW_TAIL, COL_TAIL): | ||
| _CASES.append( | ||
| _case( | ||
| _platform, | ||
| "axpy", | ||
| DataType.FP32, | ||
| DataType.FP32, | ||
| _valid_shape, | ||
| 2.0, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
No BF16 coverage for axpy.
The axpy matrix is FP16/FP32 only, on both platforms. If A5 supports BF16 for TAXPY, add at least a BF16→BF16 a5 case; if BF16 is intentionally excluded, the IR type check should reject it with a target-specific message and a negative test should assert that, so the exclusion is deliberate rather than incidental.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/st/runtime/ops/test_math_fused.py` around lines 418 - 443, Add BF16
coverage to the axpy cases assembled in _CASES, including at least a
BF16-to-BF16 case for platform A5 if TAXPY supports it. If BF16 is unsupported,
update the A5 IR type validation to reject it with a target-specific message and
add a negative test asserting that rejection.
| def test_pow_rejects_wrong_tmp_contract(): | ||
| base = _tile("base") | ||
| exp = _tile("exp") | ||
|
|
||
| with pytest.raises(ValueError, match="requires tmp"): | ||
| tile.pow(base, exp) | ||
| with pytest.raises(ValueError, match="forbids tmp"): | ||
| tile.pow(_tile("ibase", DataType.INT32), _tile("iexp", DataType.INT32), _tile("tmp", DataType.INT32)) | ||
| with pytest.raises(ValueError, match="high_precision"): | ||
| tile.pows(_tile("ibase", DataType.INT32), 2, high_precision=True) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Negative coverage is missing for dtype/precision/layout rejection.
The rejection tests only cover the tmp presence/absence contract and integer high_precision. There is no test that BF16 default precision is rejected on A2/A3, nor that non-row-major base/exp/dst/tmp tiles are rejected. Given pow/pows accept a caller-supplied scratch tile, a wrong-layout or wrong-shape tmp silently reaching codegen is the failure mode most worth pinning down here.
Want me to draft the additional negative cases once the target-aware validation lands in src/ir/op/tile_ops/elementwise.cpp?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/ut/ir/operators/test_math_fused_ops.py` around lines 89 - 98, Expand
test_pow_rejects_wrong_tmp_contract to cover target-aware BF16 default-precision
rejection on A2/A3 and non-row-major layout rejection for base, exp, dst, and
caller-supplied tmp tiles. Add cases that also validate incompatible tmp
shape/layout is rejected before codegen, using the existing tile.pow/tile.pows
APIs and matching the established ValueError behavior.
| ir.compile(AddReluProgram, output_dir=str(tmp_path), skip_ptoas=True, platform="a2a3") | ||
| pto_files = list(tmp_path.rglob("*.pto")) | ||
| assert pto_files | ||
| assert any("pto.taddrelu" in path.read_text() for path in pto_files) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | 🏗️ Heavy lift
skip_ptoas=True means taddrelu is never actually assembled.
This is the only validation for tile.add_relu, and it deliberately bypasses the assembler because the pinned PTOAS does not accept the op on A2/A3. As written, the API is publicly exported and emittable but cannot be assembled by the toolchain the repo pins — a user hitting pl.tile.add_relu gets a compile failure that no test covers. Either gate the API/emission behind the toolchain version until the PTOAS upgrade lands, or add an explicit negative test asserting the current failure mode so the gap is visible.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/ut/ir/operators/test_math_fused_ops.py` around lines 124 - 127, Update
the tile.add_relu coverage around AddReluProgram to expose the current PTOAS
limitation: either gate the public API and emission based on the supported
toolchain version, or add a negative compile test that explicitly asserts
compilation fails with the pinned PTOAS on A2/A3. Do not leave the existing
skip_ptoas=True emission-only assertion as the sole validation.
|
Closing this later PTOAS batch for now so work can proceed serially from B02 and B03. The branch is preserved for reopening when its turn arrives. |
2 participants
Summary
tile.axpywith destination reuse and exactpto.taxpyemissiontile.add_reluwith exactpto.taddreluemissiontile.pow/tile.powswith integer no-scratch and floating scratch/precision contractsValidation
/data/chenshenai/test2taxpy,tpow,tpowsall passed (3/3, one precompile worker)PTOAS limitation
pto.taddreluis emitted exactly and covered with askip_ptoas=Truecodegen assertion. Official PTOAS v0.53 parses the op but its A2/A3 conversion marks it explicitly illegal, including for a full tile, so same-name hardware execution remains blocked upstream and is recorded as such in the status matrix.