[ExecuTorch][WebGPU] Op-tests for argmax/argmin + int64-golden harness#21210
[ExecuTorch][WebGPU] Op-tests for argmax/argmin + int64-golden harness#21210JCNTH wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21210
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 46 New Failures, 3 Unrelated FailuresAs of commit a1092a2 with merge base 266e0dc ( NEW FAILURES - The following jobs have failed:
FLAKY - The following jobs failed but were likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
psiddh
left a comment
There was a problem hiding this comment.
Approving full WebGPU stack
Stack from ghstack (oldest at bottom):
Problem: The new argmax/argmin ops output int64 indices, but the op-test harness was float/int8-only (it cast any int output to fp32, so the driver would read index bytes as float — a garbage compare).
Solution: Extend the shared harness with an int64-golden path (mirroring the landed int8 path): the generator writes an int64 golden for int32/int64 raw outputs (
_write_int64, dtype "int64"), and the driver readsconst_data_ptr<int64_t>()and exact-compares to the int64 golden (load_int64_bin). The argmax/argmin suites usegolden_dtype="float32"so the golden's fp32 argmax matches the fp32 kernel exactly (an fp64 oracle could flip a near-tie index).Implementation:
generate_op_tests.pyadds_write_int64+ theis_int64branch.op_test_driver.cppadds the int64 exact-compare branch;driver_util.{h,cpp}addload_int64_bin.cases.pyregistersargmax/argminwith randn 2d/3d cases (max/min at an INTERIOR index — exercises the reduction walk) and a deliberate-tie case per op (argmax_tie_gen/argmin_tie_genplace a repeated extremum at idx 1 and 3 so the FIRST occurrence wins — this discriminates the strict->/<tie-break from a>=/<=bug).test_argmax.pyholds the modules + tie gens.Differential Revision: D112257677