Repository: pytorch
Issue: Program killed by SIGFPE in the backward pass of torch.sparse.sampled_addmm when input has specific shape
Labels: module: sparse,module: crash,module: cuda,triaged

### 🐛 Describe the bug

Recently I'm trying to implement graph attention using PyTorch, so I'm utilizing `torch.sparse.sampled_addmm`. Unfortunately, I find that the backward pass crashes if the operation is on GPU and the sparse `input` has specific shape, but not necessarily a lot of nonzero elements. Here is a simple script to reproduce:
```
N = 65536
mask = torch.sparse_coo_tensor(torch.stack([torch.arange(N), torch.arange(N)], dim=0), torch.ones(N)).cuda().to_sparse_csr()
mask.requires_grad = True
q = torch.randn(N, 128).cuda()
q.requires_grad = True
k = torch.randn(N, 128).cuda()
k.requires_grad = True
attn = torch.sparse.sampled_addmm(mask, q, k.t(), beta=0)
attn.backward(torch.ones_like(attn))
```
This script evaluates `q @ k.t()` only at the diagonal. Unfortunately, this script crashes under SIGFPE error signal. It seems that this error only happens when `N` is a multiple of 65536. The error is gone when `N` is 32768, 65537 or 98304 (32768\*3), or if the tensors are moved to CPU, but persists when `N` is changed to 131072 or 196608 (65536\*3). It also doesn't seem to depend on the number of nonzero elements. For example, if `mask` is changed to block diagonal with blocks of ones, the behavior is the same:
```
indices = torch.repeat_interleave(torch.arange(N), 128)
indices = torch.stack([indices, torch.arange(len(indices)) % 128 + indices // 128 * 128], dim=0)
mask = torch.sparse_coo_tensor(indices, torch.ones(indices.shape[1])).cuda().to_sparse_csr()
```

In the meantime, the workaround is to create a `torch.autograd.Function` which wraps `torch.sparse.sampled_addmm` and manually implement the gradients using `torch.mm`. By the way, this alternative approach is also faster, at least on my own GPU, but such topic may require a new issue.
```
class SampledMMFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, mat1, mat2):
        ctx.save_for_backward(mat1, mat2)
        return torch.sparse.sampled_addmm(input, mat1, mat2, beta=0)

    @staticmethod
    def backward(ctx, grad_output):
        mat1, mat2 = ctx.saved_tensors
        return torch.zeros_like(grad_output), torch.mm(grad_output, mat2.t()), torch.mm(mat1.t(), grad_output)
```

Stack trace:
```
(gdb) bt
#0  0x00007fff4c5a54d2 in void at::native::(anonymous namespace)::_sparse_binary_op_intersection_kernel_impl<at::native::(anonymous namespace)::CUDAKernelLauncher, at::native::(anonymous namespace)::CUDAValueSelectionIntersectionKernel<at::native::(anonymous namespace)::MulOp>, long, 8l>(at::Tensor&, at::Tensor const&, at::Tensor const&, std::vector<long, std::allocator<long> >, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, bool, bool) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#1  0x00007fff4c5a9f27 in void at::native::(anonymous namespace)::_sparse_binary_op_intersection_kernel_out<at::native::(anonymous namespace)::CUDAKernelLauncher, at::native::(anonymous namespace)::CUDAValueSelectionIntersectionKernel<at::native::(anonymous namespace)::MulOp> >(at::Tensor&, at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, bool) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#2  0x00007fff4c56a96a in at::native::(anonymous namespace)::mul_sparse_sparse_out_cuda_kernel(at::Tensor&, at::Tensor const&, at::Tensor const&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#3  0x00007fffa5677703 in at::native::_mul_sparse_sparse_out(at::Tensor const&, at::Tensor const&, at::Tensor&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#4  0x00007fff4d1b6be9 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#5  0x00007fffa5a1c690 in at::_ops::mul_out::call(at::Tensor const&, at::Tensor const&, at::Tensor&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#6  0x00007fffa567c950 in at::native::mul_sparse(at::Tensor const&, at::Tensor const&) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#7  0x00007fff4d1b6a8c in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#8  0x00007fff4d1b6b07 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#9  0x00007fffa5a11feb in at::_ops::mul_Tensor::call(at::Tensor const&, at::Tensor const&) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#10 0x00007fffa5612eb9 in at::native::mul_out_sparse_csr(at::Tensor const&, at::Tensor const&, at::Tensor&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#11 0x00007fff4d1d9869 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#12 0x00007fffa5a1c690 in at::_ops::mul_out::call(at::Tensor const&, at::Tensor const&, at::Tensor&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#13 0x00007fffa561f010 in at::native::mul_sparse_csr(at::Tensor const&, at::Tensor const&) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#14 0x00007fff4d1d970c in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#15 0x00007fff4d1d9787 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#16 0x00007fffa5a11feb in at::_ops::mul_Tensor::call(at::Tensor const&, at::Tensor const&) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#17 0x00007fffa561fbb9 in at::native::sparse_mask_sparse_compressed(at::Tensor const&, at::Tensor const&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#18 0x00007fff4d1dd4cf in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#19 0x00007fff4d1dd557 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so
#20 0x00007fffa59b95be in at::_ops::sparse_mask::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#21 0x00007fffa85d9590 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#22 0x00007fffa85d9baa in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#23 0x00007fffa5a12ceb in at::_ops::sparse_mask::call(at::Tensor const&, at::Tensor const&) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#24 0x00007fffa9a197d0 in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#25 0x00007fffa8499a94 in torch::autograd::generated::SparseSampledAddmmBackward0::apply(std::vector<at::Tensor, std::allocator<at::Tensor> >&&) ()
--Type <RET> for more, q to quit, c to continue without paging--
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#26 0x00007fffa90d4bfb in  () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#27 0x00007fffa90ce0d5 in torch::autograd::Engine::evaluate_function(std::shared_ptr<torch::autograd::GraphTask>&, torch::autograd::Node*, torch::autograd::InputBuffer&, std::shared_ptr<torch::autograd::ReadyQueue> const&) () at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#28 0x00007fffa90cf662 in torch::autograd::Engine::thread_main(std::shared_ptr<torch::autograd::GraphTask> const&) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#29 0x00007fffa90c6268 in torch::autograd::Engine::thread_init(int, std::shared_ptr<torch::autograd::ReadyQueue> const&, bool) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_cpu.so
#30 0x00007fffb0d424cb in torch::autograd::python::PythonEngine::thread_init(int, std::shared_ptr<torch::autograd::ReadyQueue> const&, bool) ()
    at /usr/lib/python3.11/site-packages/torch/lib/libtorch_python.so
#31 0x00007fffb36e1943 in std::execute_native_thread_routine(void*) (__p=0x555564da71a0) at /usr/src/debug/gcc/gcc/libstdc++-v3/src/c++11/thread.cc:104
#32 0x00007ffff76aa9eb in start_thread (arg=<optimized out>) at pthread_create.c:444
#33 0x00007ffff772e7cc in clone3 () at ../sysdeps/unix/sysv/linux/x86_64/clone3.S:78
```

### Versions

PyTorch version: 2.1.1
Is debug build: False
CUDA used to build PyTorch: 12.3
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 13.2.1 20230801
Clang version: Could not collect
CMake version: version 3.27.8
Libc version: glibc-2.38

Python version: 3.11.6 (main, Nov 14 2023, 09:36:21) [GCC 13.2.1 20230801] (64-bit runtime)
Python platform: Linux-6.6.2-arch1-1-x86_64-with-glibc2.38
Is CUDA available: True
CUDA runtime version: 12.3.52
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Laptop GPU
Nvidia driver version: 545.29.06
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.9.6
/usr/lib/libcudnn_adv_infer.so.8.9.6
/usr/lib/libcudnn_adv_train.so.8.9.6
/usr/lib/libcudnn_cnn_infer.so.8.9.6
/usr/lib/libcudnn_cnn_train.so.8.9.6
/usr/lib/libcudnn_ops_infer.so.8.9.6
/usr/lib/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             16
On-line CPU(s) list:                0-15
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen 9 7940HS w/ Radeon 780M Graphics
CPU family:                         25
Model:                              116
Thread(s) per core:                 2
Core(s) per socket:                 8
Socket(s):                          1
Stepping:                           1
CPU(s) scaling MHz:                 32%
CPU max MHz:                        6228.0000
CPU min MHz:                        400.0000
BogoMIPS:                           7988.30
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good amd_lbr_v2 nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba perfmon_v2 ibrs ibpb stibp ibrs_enhanced vmmcall fsgsbase bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local user_shstk avx512_bf16 clzero irperf xsaveerptr rdpru wbnoinvd cppc arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif x2avic v_spec_ctrl vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid overflow_recov succor smca flush_l1d
Virtualization:                     AMD-V
L1d cache:                          256 KiB (8 instances)
L1i cache:                          256 KiB (8 instances)
L2 cache:                           8 MiB (8 instances)
L3 cache:                           16 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Vulnerable: Safe RET, no microcode
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced / Automatic IBRS, IBPB conditional, STIBP always-on, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.2
[pip3] torch==2.1.1
[conda] Could not collect

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer @jcaip @ptrblck