[Maintain][Manifest] flip All/Any/CountNonzero reduction ops to implemented#1431
[Maintain][Manifest] flip All/Any/CountNonzero reduction ops to implemented#1431lcy-seso wants to merge 1 commit into
Conversation
lcy-seso
added
maintain
|
/gemini review |
There was a problem hiding this comment.
Pull request overview
Promotes the boolean/count reduction ops (AllFwdOp, AnyFwdOp, CountNonzeroFwdOp) to status: implemented in the reduction manifest, making them discoverable/dispatchable via the manifest now that implementation + conformance coverage exists.
Changes:
- Flip
AllFwdOpmanifest status fromspec-onlytoimplemented. - Flip
AnyFwdOpmanifest status fromspec-onlytoimplemented. - Flip
CountNonzeroFwdOpmanifest status fromspec-onlytoimplemented.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the status of AllFwdOp, AnyFwdOp, and CountNonzeroFwdOp to 'implemented' in the reduction manifest. Feedback highlights several inconsistencies between the manifest definitions and the actual Python implementations, specifically regarding default dimension values, shape rules for empty dimension lists, and missing supported data types in the input signatures.
| ref_api: "torch.all" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
While flipping the status to implemented is correct given the implementation provided, there is a discrepancy between the manifest and the Python implementation in tileops/ops/reduction/all_op.py. The manifest defines the default for dim as null (None), which matches torch.all behavior (full reduction). However, the implementation in all_op.py (line 65) defaults to -1. This inconsistency should be resolved to ensure the op behaves as specified when used via the manifest-driven API or when instantiated without arguments.
There was a problem hiding this comment.
Verified — AllFwdOp.__init__ (tileops/ops/reduction/all_op.py:65) defaults dim=-1, manifest says null. The code-only prereq PR will flip the default to None (same change planned for _ReduceOpBase, _WelfordReduceOp, AnyFwdOp, CountNonzeroFwdOp).
| ref_api: "torch.any" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
Similar to AllFwdOp, there is a discrepancy between the manifest default for dim (null) and the implementation in tileops/ops/reduction/any_op.py (line 65), which defaults to -1. The implementation should be updated to default to None to align with the manifest and torch.any semantics.
There was a problem hiding this comment.
Confirmed — same root cause as AllFwdOp. AnyFwdOp.__init__ (tileops/ops/reduction/any_op.py:65) defaults dim=-1. Code-only prereq PR fixes it.
| ref_api: "torch.count_nonzero" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
There are a few inconsistencies in the manifest for CountNonzeroFwdOp that should be addressed now that it is marked as implemented:
- Default Value Discrepancy: The manifest defaults
dimtonull(line 739), but the implementation intileops/ops/reduction/count_nonzero.py(line 72) defaults to-1. - Shape Rule Inconsistency: The
shape_rules(line 742) useand len(dim) > 0for the list/tuple check, which causesdim=[]to be treated as a full reduction (returning a scalar). In contrast,AllFwdOpandAnyFwdOptreatdim=[]as a no-op (returning the input shape). Given thattorch.count_nonzerodoes not supportdim=[], it should either be forbidden or follow the no-op convention of the other logical reductions. - Missing Dtype: The
inputssignature (line 735) is missingbool, which is supported by both PyTorch and the provided implementation incount_nonzero.py.
There was a problem hiding this comment.
Item-by-item verdict:
- Default
dimmismatch (real) —CountNonzeroFwdOp.__init__(tileops/ops/reduction/count_nonzero.py:72) defaultsdim=-1. Will be fixed in the code-only prereq PR. shape_rulesempty-list inconsistency (real, manifest-only) — CountNonzero uses_empty_dim_policy="full", while All/Any inherit"reject". The shape_rules text for CountNonzero / All / Any does need to be unified with the runtime policy and aligned to PyTorch (torch.count_nonzerorejectsdim=[]). Outside the status-flip carve-out — separate manifest-only PR.- Missing
booldtype (real, manifest-only) — signature dtype list will be expanded to includebool(andint32/int64/complexfor All/Any/CountNonzero where the impl supports them via_pre_kernelfp32 promotion) in the same manifest-only follow-up.
There was a problem hiding this comment.
Code Review
This pull request updates the status of AllFwdOp, AnyFwdOp, and CountNonzeroFwdOp from spec-only to implemented in the reduction manifest. However, the reviewer identified several critical issues that make this change premature. Specifically, a kernel bug in logical_reduce.py could lead to out-of-bounds reads due to missing row masking. There are also inconsistencies between the manifest and the Python implementation concerning default dimension values and the list of supported data types.
| ref_api: "torch.all" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
Marking AllFwdOp as implemented is premature due to several issues that should be addressed to ensure consistency and correctness:
- Kernel Bug: In
tileops/kernels/reduction/logical_reduce.py, the non-tiled kernel_logical_reduce_kernel(lines 114-123) only performs row masking if_needs_pad(column padding) is true. If the hidden dimensionNis aligned (e.g., 256) but the number of rowsMis not a multiple ofblock_m, theT.copyat line 123 will perform an out-of-bounds read. - Default Value Discrepancy: The manifest defines the default
dimasnull(full reduction) at line 647, whereas the Python implementation inall_op.pydefaults to-1(last-dim reduction) at line 65. This inconsistency can lead to unexpected behavior in manifest-driven tools or benchmarks. - Supported Dtypes: The manifest signature (line 643) only lists
float16 | bfloat16 | float32 | bool, but the implementation inall_op.py(lines 47-48) and the kernel logic explicitly supportint32,int64, andcomplextypes via pre-conversion tofloat32.
There was a problem hiding this comment.
Item-by-item verdict:
-
Kernel out-of-bounds on aligned-N path — verified false positive. Ran
compute-sanitizer --tool memcheckonAllFwdOpwithM=129, N=256, block_m=8(the smallest tile config; would expose any tail-row OOB) plus 15 other shapes (M ∈ {129, 257, 513, 1025} × N ∈ {128, 256, 512, 1024}). Zero errors, output canary intact, numerical results matchtorch.all. TileLang'sT.copy(x[a, b], shared)lowering inserts implicit row-bounds guards from the source tensor's static extentT.Tensor[(M, N), dtype], so the manualif pid_m*block_m + i < Mis redundant. The "else" branch is safe. No kernel fix needed. -
Default
dim=-1↔ manifestnull(real) — will be fixed in the code-only prereq PR. -
Missing
int32/int64/complexin signature dtype list (real, manifest-only) — the impl does promote those via_pre_kernel. The signature dtype list will be expanded in the follow-up manifest-only PR (outside the status-flip carve-out).
| ref_api: "torch.any" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
Marking AnyFwdOp as implemented is premature for the same reasons as AllFwdOp:
- Kernel Bug: The non-tiled kernel in
logical_reduce.pylacks row masking whenNis aligned butMis not a multiple ofblock_m, leading to potential out-of-bounds reads inT.copy. - Default Value Discrepancy: The manifest default for
dimisnull(line 693), while the Python default is-1(line 65 ofany_op.py). - Supported Dtypes: The manifest signature (line 689) should be updated to include
int32,int64, andcomplextypes supported by the implementation.
There was a problem hiding this comment.
Same verdicts as the AllFwdOp block:
- Kernel OOB — verified false positive (
compute-sanitizermemcheck clean across 16 shape configs). - Default
dim— real; code-only prereq PR. - Dtype list — real, manifest-only follow-up.
| ref_api: "torch.count_nonzero" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
Marking CountNonzeroFwdOp as implemented is premature due to the following issues:
- Kernel Bug: Similar to the other logical reduction ops, the non-tiled kernel in
logical_reduce.pyis missing row masking for alignedN, causing unsafeT.copyoperations whenM % block_m != 0. - Default Value Discrepancy: The manifest default for
dimisnull(line 739), while the Python default is-1(line 72 ofcount_nonzero.py). - Supported Dtypes: The manifest signature (line 735) is missing
bool,int32,int64, andcomplextypes, all of which are supported by the implementation and are common use cases forcount_nonzero.
There was a problem hiding this comment.
Same verdicts as the AllFwdOp / AnyFwdOp blocks for CountNonzeroFwdOp:
- Kernel OOB — verified false positive (
compute-sanitizermemcheck clean). - Default
dim— real; code-only prereq PR. - Dtype list (
bool,int32,int64,complex) — real, manifest-only follow-up.
Ibuki-wind
left a comment
Ibuki-wind
left a comment
There was a problem hiding this comment.
Overall
The logical-reduction status flip is not yet contract-true; fix the op defaults and kernel tail handling before promoting these entries.
| ref_api: "torch.all" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
LogicalReduceKernel still has an unmasked tail-row read/write path in _logical_reduce_kernel: when N is aligned, the else branch does T.copy(x[pid_m * block_m, 0], shared_buf) without checking pid_m * block_m + i < M, so the last block can read past the end of x and write past the end of out for shapes like M=129, N=512 -> keep these entries spec-only or mask every load/store before promoting them.
There was a problem hiding this comment.
Verified false positive after empirical check. Ran compute-sanitizer --tool memcheck on the exact shape you flagged (M=129, N=256, tile config block_m=8, threads=128) plus 15 sibling shapes (M ∈ {129, 257, 513, 1025} × N ∈ {128, 256, 512, 1024}):
- 0 memcheck errors across all configs.
- 8 KiB input canary (zero-initialized region adjacent to
x) intact after every run. - Numerical results match
torch.all/torch.any/torch.count_nonzeroon each shape.
Reason: TileLang lowers T.copy(x[pid_m * block_m, 0], shared_buf) against the source tensor's statically-declared extent T.Tensor[(M, N), dtype]. The lowered loads carry implicit row-bounds predicates derived from M, so the explicit if pid_m*block_m + i < M is unnecessary — the codegen already inserts equivalent guards. The DSL reads as "raw CUDA" but the safety wrapper sits one level lower.
If you have a counter-example shape that triggers the OOB I'd happily re-run sanitizer on it, but on the configs the kernel actually compiles to for these ops, there's no real-world OOB to fix.
| ref_api: "torch.all" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
AllFwdOp still defaults dim to -1 in __init__ and inherits _empty_dim_policy = "reject", but the manifest entry advertises dim: null and treats dim=[] as a no-op, so manifest-driven callers that omit dim or pass an empty sequence do not get the promoted semantics -> change the op default and empty-dim handling to match the manifest before marking this implemented.
There was a problem hiding this comment.
Confirmed on the default-dim side — AllFwdOp.__init__ (tileops/ops/reduction/all_op.py:65) defaults dim=-1, manifest says null. Code-only prereq PR will fix it on All/Any/CountNonzero plus the two base classes.
On _empty_dim_policy = "reject" vs the manifest's "treat dim=[] as no-op" shape_rules — PyTorch's torch.all(x, dim=[]) actually errors, so "reject" is the right semantics. The manifest shape_rules text is the side that's wrong (matches CountNonzero's "full" policy rather than All/Any's "reject"). That's a manifest-only fix and will be bundled into the manifest follow-up (outside the status-flip carve-out).
| ref_api: "torch.any" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
AnyFwdOp still defaults dim to -1 in __init__ and inherits _empty_dim_policy = "reject", but the manifest entry advertises dim: null and treats dim=[] as a no-op, so manifest-driven callers that omit dim or pass an empty sequence do not get the promoted semantics -> change the op default and empty-dim handling to match the manifest before marking this implemented.
There was a problem hiding this comment.
Confirmed same as the AllFwdOp thread — AnyFwdOp.__init__ (tileops/ops/reduction/any_op.py:65) defaults dim=-1. Code-only prereq PR will flip it. The _empty_dim_policy = "reject" ↔ manifest shape_rules disagreement also matches AllFwdOp's analysis — PyTorch errors on dim=[] so "reject" is correct; manifest shape_rules need a follow-up fix.
| ref_api: "torch.count_nonzero" | ||
| family: reduction | ||
| status: spec-only # dim=None not yet implemented; update when impl lands | ||
| status: implemented |
There was a problem hiding this comment.
CountNonzeroFwdOp still defaults dim to -1 in __init__, while the manifest entry advertises dim: null; workload-driven and benchmark code that omit dim will reduce the last axis instead of a full-tensor reduction -> switch the op default to None or leave the entry spec-only.
There was a problem hiding this comment.
Confirmed — CountNonzeroFwdOp.__init__ (tileops/ops/reduction/count_nonzero.py:72) defaults dim=-1, manifest says null. Code-only prereq PR will fix it (along with All/Any and the two base classes).
|
Reviewer-feedback verification summary After re-reading every comment on PRs #1428 / #1430 / #1431 and running targeted checks, the verified picture is:
Plan:
Detailed per-comment verdicts are posted as inline replies on each thread. |
…mented AllFwdOp, AnyFwdOp, CountNonzeroFwdOp now satisfy their spec via logical_reduce kernel + dedicated op files; drop the stale "dim=None not yet implemented" comment along with the status flip. Co-Authored-By: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
lcy-seso
force-pushed
the
fix/manifest/issue-1401
branch
from
b1056ae to
bde01bd
Compare
Closes #1401
Summary
AllFwdOp,AnyFwdOp,CountNonzeroFwdOpfromstatus: spec-onlytostatus: implementedintileops/manifest/reduction.yaml.Context
Third of a sibling set:
After all three merge, the remaining 7 spec-only reduction ops flip to
implemented, bringing the reduction family to 19/19status: implemented.Note on AC-5
AC-5 in the linked issue says "All 13 reduction.yaml ops are status: implemented after merge." The literal count is stale —
reduction.yamlcurrently holds 19 ops (12 alreadyimplementedonupstream/main, 7spec-onlybefore this sibling set). The substantive intent — "reduction family fully implemented after merge" — is satisfied once #1428, #1430, and this PR all land.Test plan
tileops/manifest/reduction.yamlshowsstatus: implementedfor AllFwdOp, AnyFwdOp, CountNonzeroFwdOp.python scripts/validate_manifest.pyexits 0.pytest tests/test_validate_manifest.pypasses.tileops/manifest/reduction.yaml(3 lines changed).implementedafter the three sibling PRs merge (literal "13" is a stale count; see note above).