【KernelGen】Add nonzero_numpy operator

[Schopenhauer-loves-Hegel]

PR Category

Operator

Type of Change

New Feature

Description

Add nonzero_numpy operator implementation with Triton kernel.

• Implementation mode: manual_kernel
• Accuracy test: 24/24 passed

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

torch.bfloat16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0604	0.2958	0.204
[256, 256]	0.0634	0.3826	0.166
[1024, 1024]	0.0667	0.3940	0.169
[4096, 4096]	0.5237	0.8263	0.634
[1024, 65536]	1.9437	2.3741	0.819
[10000, 1]	0.0641	0.3143	0.204
[10000, 256]	0.1064	0.4239	0.251
[10000, 65536]	18.6106	21.4398	0.868

torch.bool

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0595	0.2004	0.297
[256, 256]	0.0638	0.2939	0.217
[1024, 1024]	0.0643	0.2924	0.220
[4096, 4096]	0.3146	0.5887	0.534
[1024, 65536]	1.1322	1.6725	0.677
[10000, 1]	0.0635	0.2020	0.315
[10000, 256]	0.0772	0.3083	0.250
[10000, 65536]	10.6116	14.5354	0.730

torch.float16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0622	0.2987	0.208
[256, 256]	0.0631	0.3833	0.165
[1024, 1024]	0.0659	0.3975	0.166
[4096, 4096]	0.5262	0.8275	0.636
[1024, 65536]	1.9399	2.3785	0.816
[10000, 1]	0.0633	0.2943	0.215
[10000, 256]	0.1063	0.4253	0.250
[10000, 65536]	18.5689	21.4910	0.864

torch.float32

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0608	0.2934	0.207
[256, 256]	0.0633	0.3826	0.166
[1024, 1024]	0.0750	0.3994	0.188
[4096, 4096]	0.5718	0.8310	0.688
[1024, 65536]	2.1408	2.4800	0.863
[10000, 1]	0.0637	0.2936	0.217
[10000, 256]	0.1134	0.4293	0.264
[10000, 65536]	20.6683	22.3866	0.923

torch.int16

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0622	0.2935	0.212
[256, 256]	0.0625	0.3829	0.163
[1024, 1024]	0.0651	0.3921	0.166
[4096, 4096]	0.5244	0.8313	0.631
[1024, 65536]	1.9430	2.3822	0.816
[10000, 1]	0.0642	0.2934	0.219
[10000, 256]	0.1073	0.4237	0.253
[10000, 65536]	18.5533	21.4803	0.864

torch.int32

Shape	Torch Latency (ms)	Gems Latency (ms)	Speedup
[64, 64]	0.0596	0.2933	0.203
[256, 256]	0.0629	0.3884	0.162
[1024, 1024]	0.0665	0.3933	0.169
[4096, 4096]	0.5714	0.8205	0.696
[1024, 65536]	2.1585	2.4437	0.883
[10000, 1]	0.0637	0.2925	0.218
[10000, 256]	0.1135	0.4233	0.268
[10000, 65536]	20.5569	22.0224	0.933

Overall: median speedup = 0.251x, mean speedup = 0.422x (48 data points)

---

Generated by auto_gen tool with Claude Code

[Schopenhauer-loves-Hegel]

Closing in favor of grouped PRs to reduce review overhead.

[tengqm]

@Schopenhauer-loves-Hegel You may want to rebase this onto 'master' for picking up the latest CI jobs.

[CLAassistant]

All committers have signed the CLA.

[factnn]

UT Results

Branch rebased to upstream/master and tests passed:

python -m pytest tests/ -m nonzero --ref cpu -x -q
18 passed

[w1120029931-bit]

please resolve comments

[factnn]

Thanks for the review! All addressed:

1. Renamed benchmark/test_nonzero_numpy_perf.py → benchmark/test_nonzero_numpy.py ✅
2. Renamed test_perf_nonzero_numpy → test_nonzero_numpy ✅
3. Already using BOOL_TYPES (not [torch.bool]) ✅
4. Re: gems_assert_equal — nonzero_numpy outputs integer index tensors (dtype=int64), not float values, so assert_equal is correct here. The float input only affects which indices are returned, not the dtype of the output.

[w1120029931-bit]

LGTM