Open lingzhi98 opened 2 months ago
yea I confirmed that we don't have a lowering for
Counter: aten::grid_sampler_2d
Value: 1
I guess the solution here would be to actually lower this op.
also found https://github.com/pytorch/xla/issues/6581
Lowering grid sampler 2d is not enough, we should lower grid sampler 3d also to fully support grid sampler op. That's why I suggest set autocast mode as fp32 firstly. Or there exists plan to support grid sampler 3d? I dont find it until now.
Hi, @wonjoolee95, is that ok to assign this ticket to you?
🐛 Bug
grid_sampler can not run with auto mixed precision mode.
Steps to reproduce the behavior:
RuntimeError: grid_sampler_2d_cpu not implemented for BFloat16
Expected behavior
Autocast the inputs of grid_sampler to fp32 datatype.
Environment
Additional context
torch.cuda set autocast mode of grid_sampler as promote, but torch xla can't. Due to torch_xla has no lowering of grid_sampler_2d/grid_sampler_3d, this op will fallback to torch.cpu implementation and no support of bfloat16. Maybe we should set autocast mode as fp32 firstly, and change it to promote until lowering is ready. By the way, grid sampler has no need of lowering, due to this op do nothing but only dispatch to gird_sampler_2d, grid_sampler_3d and cudnn_grid_sampler. https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/GridSampler.cpp#L1046