prims.div: equalize NVFuser prims with PyTorch for integer Tensors

[nikitaved]

NVFuser implementation of prims.div used reciprocal. That caused upcasting for integer inputs and, ultimately, rendered meta checks incorrect.

Fixes #808
Fixes #818

[nikitaved]

It is part one - the fix for integer types. We also have issues with floating types that need to get fixed to re-enable all the consistency tests.

[t-vi]

Thank you @nikitaved, awesome stuff!

[mruberry]

This is floor division, not true division

[nikitaved]

Good catch! Made a typo - will fix that. Thank you!

[nikitaved]

C division has always been the true division for me :)

[mruberry]

I don't think these results should be equal (I think we're exposing a bug in the nvFuser implementation of prims.div). prims.div should implement truncation division, which is not the same as Python's floor division. See, for example, the PyTorch executor's implementation of prims.div:

lightning-thunder/thunder/executors/torchex.py

Line 907 in 44d2f3c

def _div_prim_impl(a: Number | torch.Tensor, b: Number | torch.Tensor) -> torch.Tensor:

[mruberry]

fyi @kevinstephano @jjsjann123

[nikitaved]

div calls to prims.div, not ltorch.div, so they should be equivalent.

[mruberry]

Yes, but doesn't truediv call floor division?

[nikitaved]

Floor division has a decomp that is more than just prims.div, and this is what is being tested here.

[mruberry]

No, but it's a common misconception that they are. If the result is negative then they round differently. Floor division rounds down, truncation division rounds towards zero.

In Python

-9 // 5
-2

In C++

#include <iostream>

int main() {
    int a = (-9 / 5);
    std::cout << a; 
}

prints -1. This is because -9 / 5 is -1.8, and floor division takes the floor of -1.8, which is -2, while truncation division truncates the fractional part of -1.8, giving -1.

[mruberry]

I am confused -- I am not comparing these two operations head to head, I just make sure that the cpu version is consistent with the nvfuser version for both jitted prims.div and * // *.

You're correct, of course, and I just misread the code. You are comparing that the CPU and CUDA versions are consistent with each other. The fact that truncation division and floor division are distinct for integers is a separate issue, and doesn't impact this code. My mistake.

[mruberry]

fyi @kevinstephano @jjsjann123

Aha! @IvanYashchuk told me that nvFuser's div is performing truncation division for integers now. I mistakenly thought it wasn't. That's why this test passes

[t-vi]

while we are looking at torchex, I think the div prim impl has a typo (twice a instead of a and b) in the if condition:

lightning-thunder/thunder/executors/torchex.py

Line 908 in 44d2f3c

if dtypes.is_exact_dtype(to_dtype(a.dtype)) and dtypes.is_exact_dtype(to_dtype(a.dtype)):

[mruberry]

while we are looking at torchex, I think the div prim impl has a typo (twice a instead of a and b) in the if condition:

lightning-thunder/thunder/executors/torchex.py

Line 908 in 44d2f3c

if dtypes.is_exact_dtype(to_dtype(a.dtype)) and dtypes.is_exact_dtype(to_dtype(a.dtype)):

I think you're right!