extend nonzero to int64

[bhack]

Fixes #51871

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/125850

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures, 3 Unrelated Failures

As of commit a46722a with merge base 8f30f36 ():

NEW FAILURES - The following jobs have failed:

• pull / linux-focal-cuda12.1-py3.10-gcc9-bazel-test / build-and-test (default, 1, 1, linux.4xlarge.nvidia.gpu) (gh)
  /var/lib/jenkins/workspace/BUILD.bazel:430:11: Compiling aten/src/ATen/native/cuda/Nonzero.cu failed: (Exit 1): nvcc failed: error executing command (from target //:aten_cuda) external/local_cuda/cuda/bin/nvcc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wunused-but-set-parameter -Wno-free-nonheap-object -fno-omit-frame-pointer '-std=c++11' -MD -MF ... (remaining 332 arguments skipped)
• pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / build (gh)
  Process completed with exit code 1.
• pull / linux-focal-rocm6.1-py3.8 / build (gh)
  /var/lib/jenkins/workspace/aten/src/ATen/native/hip/Nonzero.hip:46:9: error: 'InputIteratorT' does not refer to a value

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / linux-focal-cuda11.8-py3.10-gcc9 / build (gh) (similar failure)
  Process completed with exit code 1.
• pull / linux-focal-cuda12.1-py3.10-gcc9 / build (gh) (similar failure)
  Process completed with exit code 1.
• pull / linux-jammy-cuda11.8-cudnn8-py3.8-clang12 / build (gh) (similar failure)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[bhack]

/cc @ezyang @eqy This is an explorative blackbox PR as I don't have free cuda resources right now and we don't have a quick way to setup the env to contribute sparse c++/cuda PR (see #125297).

But I made this editable on your side in the case you have the env ready and a quick fix is enough.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: bhack (0841181, bb7240c, a47573b, 637532c, 65dbb6c, 8cce2c7, 79ad3d3, 052bb61, a2554de, 679b9fd, dc29c0f, 5332381, a46722a, 2362079, 58eeb41, faa531f)

[eqy]

Should we add a (presumably large tensor) test for this?

[bhack]

Should we add a (presumably large tensor) test for this?

Do we had an INT_MAX test already somewhere that we could expand?

[eqy]

Unfortunately these are not really unified at the moment, but this should surface some examples: https://github.com/search?q=repo%3Apytorch%2Fpytorch+64bit+language%3APython+path%3A%2F%5Etest%5C%2F%2F&type=code

[bhack]

As we don't have a specific CUDA test do we want to find a workaround from python?

Can you suggest one from grep -R torch.nonzero test/?

[bhack]

I think I am going to close this as cub::DispatchSelectIf probably it will be slower then cub::DeviceSelect::Flagged we are currently using.

Probably we need to wait upstream for NVIDIA/cccl#1422

What do you think?

[bhack]

@ezyang Do you think we can we open a new ticket to lower this with Trition where and sum?
https://github.com/pytorch/pytorch/blob/a174c536f8f32b41da9efa647e364196423468a5/torch/_inductor/lowering.py#L2187C20-L2187C35

Edit:
The ticket is at #126003

[bhack]

Does this requries cub/cccl 2.4.0?

[ezyang]

yes, need a big tensor test. @eqy's link is good for examples

[bhack]

Ok thanks,
so I am going to close it as I don't have the env and currently spare GPU computing to write a brand new test and recompile it.
At least if we don't identify another python test that it is already indirectly using nonzero and that we could modify it with a big input.

[bhack]

Just to check if it could compile at least with the current CUB version.
Do you know what is this CI failure?

/usr/local/cuda/include/cub/agent/agent_select_if.cuh(264): error: function "at::native::<unnamed>::NonZeroOp<T>::operator() [with T=c10::complex<c10::Half>]" cannot be called with the given argument list
            argument types are: (int64_t)
            object type is: at::native::<unnamed>::NonZeroOp<c10::complex<c10::Half>>
                  selection_flags[ITEM] = select_op(items[ITEM]);

[bhack]

I think we need cub/cub/agent/agent_select_if.cuh changes introduced at NVIDIA/cccl#1379

So this mean that we need to wait for the next cuda 12.4 update and make it also conditional.

[ezyang]

This PR seems fine. I agree you may need to preprocessor your way to victory. CI will say.

[bhack]

@ezyang the new [CUDA] 12.5](https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html) delivers CUB 2.4.0 so it could be enough for this workaround.