Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[MPS] add support for aten::nextafter #109685

Closed
wants to merge 2 commits into from
Closed

[MPS] add support for aten::nextafter #109685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
da37c67
implemented nextafter for metal language 2.3
igm503 Sep 20, 2023
1392c39
added support for half; added correct behavior with NAN
igm503 Sep 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
46 changes: 46 additions & 0 deletions aten/src/ATen/native/mps/operations/BinaryKernel.mm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
#else
#include <ATen/ops/maximum.h>
#include <ATen/ops/minimum.h>
#include <ATen/ops/nextafter_native.h>
#include <ATen/ops/polar_native.h>
#include <ATen/ops/view_as_real.h>
#endif
Expand Down Expand Up @@ -181,6 +182,45 @@ kernel void complex_mul(constant void * input_ [[buffer(0)]],

REGISTER_COMPLEX_MUL_OP(float);
REGISTER_COMPLEX_MUL_OP(half);

template<typename T, typename U>
kernel void nextafter_kernel(constant void * input_ [[buffer(0)]],
constant void * other_ [[buffer(1)]],
device void * out_ [[buffer(2)]],
constant uint3 * offsets [[buffer(3)]],
uint tid [[thread_position_in_grid]]) {
device T* out = (device T*)((device uint8_t*)out_ + offsets[tid].x);
constant T* input = (constant T*)((constant uint8_t*)input_ + offsets[tid].y);
constant T* other = (constant T*)((constant uint8_t*)other_ + offsets[tid].z);

if (*input == *other)
{
*out = *other;
}
else if (isnan(*input) || isnan(*other))
{
*out = NAN;
}
else
{
U bits = as_type<U>(*input);
bits = bits + ((*other > *input) ? 1 : -1);
*out = as_type<T>(bits);
}
}

#define REGISTER_NEXTAFTER_OP(DTYPE, UTYPE) \
template \
[[host_name("nextafter_kernel_" #DTYPE)]] \
kernel void nextafter_kernel<DTYPE, UTYPE>( \
constant void * input, \
constant void * other, \
device void * out, \
constant uint3 * offsets, \
uint tid)

REGISTER_NEXTAFTER_OP(float, uint);
REGISTER_NEXTAFTER_OP(half, ushort);
)BINARY_METAL";

using namespace mps;
Expand Down Expand Up @@ -336,9 +376,15 @@ static void copysign_mps_kernel(TensorIteratorBase& iter) {
mps::binary_mps_impl(iter, "copysign");
}

static void nextafter_mps_kernel(TensorIteratorBase& iter) {
TORCH_CHECK_TYPE(isFloatingType(iter.common_dtype()), "nextafter_mps not implemented for non-floating types");
mps::binary_mps_impl(iter, "nextafter_kernel");
}

REGISTER_DISPATCH(fmax_stub, &fmax_mps_kernel);
REGISTER_DISPATCH(fmin_stub, &fmin_mps_kernel);
REGISTER_DISPATCH(copysign_stub, &copysign_mps_kernel);
REGISTER_DISPATCH(nextafter_stub, &nextafter_mps_kernel);

Tensor& polar_out_mps(const Tensor& abs, const Tensor& angle, Tensor& output) {
auto new_size = at::infer_size(abs.sizes(), angle.sizes());
Expand Down
2 changes: 1 addition & 1 deletion aten/src/ATen/native/native_functions.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9541,7 +9541,7 @@
structured: True
structured_inherits: TensorIteratorBase
dispatch:
CPU, CUDA: nextafter_out
CPU, CUDA, MPS: nextafter_out
tags: pointwise

- func: nextafter(Tensor self, Tensor other) -> Tensor
Expand Down
3 changes: 2 additions & 1 deletion test/test_mps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,8 @@ def mps_ops_grad_modifier(ops):
'exponential': [torch.float16, torch.float32],

# CPU errors
# derivative for aten::nextafter is not implemented on CPU
'nextafter': None,
# derivative for aten::floor_divide is not implemented on CPU
'floor_divide': [torch.float16, torch.float32],
# derivative for aten::narrow_copy is not implemented on CPU
Expand Down Expand Up @@ -556,7 +558,6 @@ def mps_ops_modifier(ops):
'nanquantile': None,
'nanmedian': None,
'native_dropout_backward': None,
'nextafter': None,
'normnuc': None,
'nn.functional.fractional_max_pool2d': None,
'nn.functional.fractional_max_pool3d': None,
Expand Down