Compute floor_divide in float for Half and BFloat16 - #193766

Fixes #193754.

torch.floor_divide on bfloat16 and float16 gives a different answer depending on the length of the input:

>>> import torch
>>> a, b = torch.full((2,), 560.0, dtype=torch.bfloat16), torch.full((2,), 3.0, dtype=torch.bfloat16)
>>> torch.floor_divide(a, b)[0]
tensor(187., dtype=torch.bfloat16)
>>> torch.floor_divide(a[:1], b[:1])[0]
tensor(186., dtype=torch.bfloat16)

186 is right.

Cause

div_floor_kernel has two paths. When the divisor is a scalar it dispatches through AT_DISPATCH_REDUCED_FLOATING_TYPES and computes in opmath_t, so bfloat16 runs in float. The tensor-tensor path computes in scalar_t.

A size-1 tensor is squeezed to 0-dim, which makes iter.is_scalar(2) true, so it takes the widened path. That is the whole length dependence.

Staying in bfloat16 is what gets the wrong answer. From NOTE: [Floor Division in Python], the quotient is (a - fmod(a, b)) / b. For 560 and 3 that is 558 / 3, but 558 needs 10 significand bits and bfloat16 has 8, so it rounds to 560, and 560 / 3 rounds to 187.

std::fmod does not save this. c10::BFloat16 has its own overload returning c10::BFloat16, so every step in c10::div_floor_floating rounds.

The error is not bounded to one unit either. Over 4096 random values the largest difference from float math was 64.

Fix

Widen both loops in the tensor-tensor path to opmath_t, matching the scalar divisor path. Same one-line cast on CUDA, whose general path had the same split against its own is_cpu_scalar(2) branch.

In the vectorized loop this replaces the CPU_CAPABILITY_SVE256 && __ARM_FEATURE_BF16 specialization of div_floor_floating_vec, which already did exactly this widening for one backend and one dtype. Since it did not widen the scalar loop next to it, SVE256 has the same length dependence today, just in the other direction.

Written as if constexpr on is_reduced_floating_point_v rather than two specializations, so the float body is not instantiated for Half and BFloat16 at all. git diff -w on that hunk is small, the body is only reindented.

float and double are unaffected. opmath_type<float> is float, so the cast is a no-op and the if constexpr picks the existing body.

Correctness

4000 random pairs per dtype, comparing the tensor-tensor result against the same op elementwise on 0-dim tensors, and against float math narrowed once:

560 // 3 in bfloat16 is now 186 at lengths 1, 2, 3, 8, 16, 64 and 1000, and at 0-dim.

Performance

This is faster, not slower. Single threaded, arm64, best of 7.

At n = 4096: bfloat16 51.5 -> 11.8 us, float16 31.6 -> 9.1 us.

Vectorized<BFloat16>::fmod is map2(q, std::fmod) and ::floor is map(floor_impl), so the old path spilled each vector to memory and called scalar libm per lane, and std::fmod on BFloat16 converted to float and back anyway. Converting once up front gets Sleef_fmod and vrndmq_f32 on real float vectors. This is what the SVE comment predicted, that most of these instructions end up as f32 anyway.

Tests

test_div_floor_reduced_precision in test/test_binary_ufuncs.py. It compares against float math narrowed once with atol=0, rtol=0, since the default bfloat16 tolerance is wide enough to hide an off-by-one quotient. It also checks lengths 1, 2 and 3 so the scalar loop and the vectorized loop are both covered.

Before this change it fails on 240 / 4096 elements for bfloat16 and 18 / 4096 for float16.

test_binary_ufuncs.py passes, 12713 passed. test_ops.py -k "floor_divide or div_floor_rounding" passes, 169 passed.

I do not have a CUDA machine, so that one line is only reasoned about and left to CI.

cc @ptrblck @msaroufim @eqy @tinglvv @nWEIdia @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @aditew01