MSELoss #435
MSELoss #435
Conversation
| if reduction == Reduction.NONE.value: | ||
| return func(inp, target) | ||
|
|
||
| M = inp.numel() |
There was a problem hiding this comment.
suggest making inp and target contiguous
src/flag_gems/ops/mse_loss.py
Outdated
| M = inp.numel() | ||
| dtype = inp.dtype | ||
| if dtype is torch.bool: | ||
| inp = inp.to(torch.int64) |
There was a problem hiding this comment.
what about target? does it need to be upcasted unto torch.int64?
There was a problem hiding this comment.
I fould that torch mse do not support bool and int.
RuntimeError: "mse_cuda" not implemented for 'Bool'
RuntimeError: "mse_cuda" not implemented for 'Bool'
src/flag_gems/ops/mse_loss.py
Outdated
| @triton.jit | ||
| def kernel_1(inp, target, mid, M, BLOCK_SIZE: tl.constexpr, reduction: tl.constexpr): | ||
| if tl.constexpr(inp.dtype.element_ty == tl.float16) or tl.constexpr( | ||
| inp.dtype.element_ty == tl.bfloat16 |
There was a problem hiding this comment.
dtype of input tensor will be fixed as constant when compiling. it's not necessary to specify it explicitly.
src/flag_gems/ops/mse_loss.py
Outdated
| ): | ||
| cdtype = tl.float32 | ||
| else: | ||
| cdtype = inp.dtype.element_ty |
There was a problem hiding this comment.
assigning cdtype as tl.float32 is more simple.
src/flag_gems/ops/mse_loss.py
Outdated
| ): | ||
| cdtype = tl.float32 | ||
| else: | ||
| cdtype = mid.dtype.element_ty |
src/flag_gems/ops/mse_loss.py
Outdated
| mid_size = triton.cdiv(M, block_size) | ||
| block_mid = triton.next_power_of_2(mid_size) | ||
|
|
||
| mid = torch.empty((mid_size,), dtype=dtype, device=inp.device) |
There was a problem hiding this comment.
initializing mid as torch.float32 might improve the precision.
tests/test_reduction_ops.py
Outdated
| def test_accuracy_mse_loss(shape, dtype, reduction): | ||
| dim = 1 | ||
| inp = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) | ||
| target = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) |
There was a problem hiding this comment.
May requires_grad be set to false?
| inp = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) | ||
| target = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) | ||
|
|
||
| ref_inp = to_reference(inp, True) |
There was a problem hiding this comment.
Since the dtype is float and the operation involves reduction, setting upcast=True is necessary to obtain a higher precision for reference ?
|
@StrongSpoon @0x45f Hi, I don not understand why the converage CI failed. |
PR Category
Type of Change
Description
Issue
Closes: #396
Progress
Performance
On NVIDIA A10
