Int64 support for UpsampleNearest3d

[kdrozd-dev]

Fixes: #2510.

Align with changes from: pytorch/pytorch#144865. Improve error messages to resemble the cuda ones and change the relevant int32 variables to int64.

[pbielak]

Looks good, but similar to my changes with the MaxPool3D kernel (see PR: #3558 and #3632 ), we should check the impact of the int64_t "hardcoding". Maybe a dispatch over int32/int64 would give better perf results.

[kdrozd-dev]

Looks good, but similar to my changes with the MaxPool3D kernel (see PR: #3558 and #3632 ), we should check the impact of the int64_t "hardcoding". Maybe a dispatch over int32/int64 would give better perf results.

Seems like a good idea will switch to dispatch approach after a quick benchmark

[kdrozd-dev]

3 runs with 20 execs per test case each

Case	Shape	dtype	Output numel	Speedup (med)	Speedup (avg)	Stable
small-f32	(2, 64, 16, 16, 16)	float32	4M	0.869x	0.899x	no
small-bf16	(2, 64, 16, 16, 16)	bfloat16	4M	1.044x	1.062x	no
small-f16	(2, 64, 16, 16, 16)	float16	4M	1.482x	1.287x	no
med-f32	(4, 128, 32, 32, 32)	float32	134M	1.034x	0.869x	no
med-bf16	(4, 128, 32, 32, 32)	bfloat16	134M	0.986x	1.089x	no
med-f16	(4, 128, 32, 32, 32)	float16	134M	0.925x	0.848x	no
large-f32	(1, 32, 64, 128, 128)	float32	268M	1.053x	1.053x	YES
large-bf16	(1, 32, 64, 128, 128)	bfloat16	268M	1.058x	1.058x	YES
xl-bf16	(1, 64, 64, 128, 256)	bfloat16	1.07B	1.054x	1.054x	YES
xl-f32	(1, 64, 64, 128, 256)	float32	1.07B	1.055x	1.055x	YES
med-exact-f32	(4, 128, 32, 32, 32)	float32	134M	0.847x	0.845x	no
med-scale3-f32	(2, 64, 16, 16, 16)	float32	14M	1.302x	1.063x	no

The benchmark shows that for stable test cases dispatch based approach is faster by around 5%. For small kernels results were mostly noise and varied greatly.

[EikanWang]

@copilot please review the changes in depth

[Copilot]

@copilot please review the changes in depth

I've reviewed the changes in depth. Here's my assessment:

Overall approach: The PR correctly implements an int32/int64 index dispatch pattern using AT_DISPATCH_INDEX_TYPES with canUse32BitIndexMath, which gives optimal performance for small tensors (int32 path) while supporting large tensors (int64 path). This aligns with the CUDA implementation from pytorch/pytorch#144865 and the benchmark results show ~5% speedup for large tensors.

Key changes reviewed:

1. Forward kernel (UpsampleNearest3dKernelFunctor) — index_t is correctly applied to dst_idx, dst_c_stride, src_c_stride, and src_idx which are the variables that can overflow with large tensors. The per-dimension variables (c, dst_z, src_z, etc.) remain int which is correct since individual dimensions won't exceed int32 limits.

2. Backward kernel (UpsampleNearest3dBackwardFunctor) — Same pattern applied correctly. index_t used for stride computations and src_idx inside the triple-nested loop.

3. Template function signatures — Changed from unsigned int n to int64_t n, and ceil_div calls updated to use static_cast<int64_t>(work_group_size).

4. TORCH_CHECK removal — The old TORCH_CHECK(output.numel() <= std::numeric_limits<int32_t>::max()) is correctly removed since the dispatch now handles both cases.

One minor observation: The UpsampleNearest3dBackwardFunctor constructor (line 282-308) doesn't initialize index_bw_op_ in its member initializer list, though it's declared as a member (line 325). This is a pre-existing issue (it exists in the original code), but it means index_bw_op_ will be default-constructed rather than using the passed index_bw_op parameter. This doesn't affect correctness since NearestBwIndexOp and NearestExactBwIndexOp are stateless functors, but it's worth noting.

No blocking issues found — the implementation is correct and well-aligned with the upstream CUDA approach.

[kdrozd-dev]

/merge -f "unrelated CI failure"

[chuanqi129]

✅ PR has been successfully merged by @kdrozd-dev (force merge).