Segmentation fault on kernel.grad (passing in torch.Tensor on wrong device)

[oliver-batchelor]

Passing in a torch.Tensor with the wrong device for the taichi backend into the kernel.grad causes a segmentation fault shown below. Forward pass handles it OK.

Example shows passing a cuda tensor when using arch=ti.cpu, if the tensor is also on cpu the program runs fine.

python crash.py

[Taichi] version 1.7.0, llvm 15.0.4, commit 59156b3c, linux, python 3.10.0
[Taichi] Starting on arch=x64
True
[E 12/21/23 19:33:09.236 3705923] Received signal 11 (Segmentation fault)


***********************************
* Taichi Compiler Stack Traceback *
***********************************
/local/mambaforge/envs/torch/lib/python3.10/site-packages/taichi/_lib/core/taichi_python.cpython-310-x86_64-linux-gnu.so(+0x4805a34) [0x7fc9486b0a34]
/local/mambaforge/envs/torch/lib/python3.10/site-packages/taichi/_lib/core/taichi_python.cpython-310-x86_64-linux-gnu.so(+0x1b3618e) [0x7fc9459e118e]
/lib/x86_64-linux-gnu/libc.so.6(+0x42520) [0x7fc958bd0520]
[0x7fc9436c5372]

import taichi as ti
import torch

@ti.kernel
def double_kernel(x:ti.types.ndarray(dtype=ti.f32, ndim=1), 
           out:ti.types.ndarray(dtype=ti.f32, ndim=1)):
  for i in x:
    out[i] = 2*x[i]



if __name__ == '__main__':
    ti.init(log_level=ti.DEBUG, debug=True)

    device = torch.device('cuda:0')
    params = torch.rand(100, dtype=torch.float32, device=device)
    out = torch.ones_like(params)

    double_kernel(params, out)

    out.requires_grad_(True)
    out.grad = torch.ones_like(out)

    double_kernel.grad(params, out)