Skip to content

grid_sample: implement mode='bicubic' with 5-D input on XPU #5542

Closed

Description

@torchxpubot

Summary

torch.nn.functional.grid_sample gained support for mode='bicubic' with
5-D input on CPU and CUDA in pytorch/pytorch#194787 (upstream commit
20fb7bd04e). XPU does not implement it: grid_sampler_3d_kernel and
grid_sampler_3d_backward_kernel in
src/ATen/native/xpu/sycl/GridSampler.cpp call the mode-taking overload
check_grid_sampler_3d(input, grid, interpolation_mode)
(aten/src/ATen/native/GridSamplerUtils.h:88), which raises

NotImplementedError: grid_sampler(): bicubic interpolation with 5D input
                     is not implemented for xpu

Upstream kept that overload solely for torch-xpu-ops, with a TODO to drop
it once XPU implements the mode
(aten/src/ATen/native/GridSamplerUtils.h:83-86).

Why this is filed

The same upstream PR made sample_inputs_grid_sample yield bicubic for
5-D input on every backend, so every OpInfo-driven suite now hits the
unimplemented mode on XPU. Upstream marked the affected tests
unittest.expectedFailure for device_type="xpu" in the
nn.functional.grid_sample OpInfo skips tuple
(torch/testing/_internal/common_methods_invocations.py), each with a
TODO: drop these when mps / xpu implement 5-D bicubic.

One suite was missed: test/test_overrides.py
TestTorchFunctionRedispatchOpsDeviceXPU::test_redispatch_nn_functional_grid_sample_xpu_*,
which was consequently auto-disabled upstream as
pytorch/pytorch#197521, pytorch/pytorch#197522,
pytorch/pytorch#197523 and pytorch/pytorch#197524. That gap is being
closed by adding the matching xfail entry; this issue tracks the real
work.

What is needed

Port the CUDA implementation from aten/src/ATen/native/cuda/GridSampler.cu
and GridSampler.cuh to SYCL in src/ATen/native/xpu/sycl/GridSampler.cpp:

  • resolve_cubic_taps / compute_coordinates_sized /
    grid_sampler_unnormalize*_sized helpers (GridSampler.cuh:321-408)
  • a dedicated 5-D bicubic forward kernel (CUDA keeps it separate from
    grid_sampler_3d_kernel for register pressure reasons)
  • the matching 5-D bicubic backward kernel
  • switch both XPU 5-D entry points from
    check_grid_sampler_3d(input, grid, interpolation_mode) to the
    rank-only check_grid_sampler_3d(input, grid)

The double backward lives in device-agnostic
torch/csrc/autograd/FunctionsManual.cpp and already handles bicubic, so
it needs no XPU work.

Coordination note

The upstream expectedFailure entries are strict: once XPU implements the
mode they turn into "unexpected success" failures. A torch-xpu-ops PR
implementing this must land together with a pytorch/pytorch PR that
removes those entries and the test/test_overrides.py one.

Metadata

Metadata

Assignees

Labels

ai_generatedContent generated by an AI assistantenhancementNew feature or requestmodule: opsOther ATen/native ops: autograd, optimizer, fx, export

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions