[MPS] add support for aten::nextafter

[igm503]

Fixes #77764 (comment)

Adds support for aten::nextafter to the MPS backend. Supports float and half types.

Notes:

• I've added nextafter to the output_grad_check XFAILLIST since neither this nor the cpu implementations have grad functions
• Metal Shading Language 3.1 seems to have a native nextafter() function, so once that's available, this kernel can just call that.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/109685

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 1392c39 with merge base 0317626 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[qqaatw]

Small suggestion: As you mentioned, the MSL 3.1 supports nextafter natively. I think we could use __METAL_VERSION__ macro to branch off from the custom implementation used below Metal 3.1. Below is an example:

pytorch/aten/src/ATen/mps/IndexKernels.h

Lines 26 to 30 in d9fe171

	#if __METAL_VERSION__ >= 300
	constant IndexAB * indexAB [[buffer(0)]],
	#else
	constant IndexAB & indexAB [[buffer(0)]],
	#endif

[kulinseth]

Small suggestion: As you mentioned, the MSL 3.1 supports nextafter natively. I think we could use __METAL_VERSION__ macro to branch off from the custom implementation used below Metal 3.1. Below is an example:

pytorch/aten/src/ATen/mps/IndexKernels.h

Lines 26 to 30 in d9fe171

	#if __METAL_VERSION__ >= 300
	constant IndexAB * indexAB [[buffer(0)]],
	#else
	constant IndexAB & indexAB [[buffer(0)]],
	#endif

I think he meant when that’s available he can call that function directly . I didn’t see that native function being called in this pr . In metal 3.1 time we would still need this code path to have it working for older os

[qqaatw]

Small suggestion: As you mentioned, the MSL 3.1 supports nextafter natively. I think we could use __METAL_VERSION__ macro to branch off from the custom implementation used below Metal 3.1. Below is an example:

pytorch/aten/src/ATen/mps/IndexKernels.h

Lines 26 to 30 in d9fe171

	#if __METAL_VERSION__ >= 300
	constant IndexAB * indexAB [[buffer(0)]],
	#else
	constant IndexAB & indexAB [[buffer(0)]],
	#endif

I think he meant when that’s available he can call that function directly . I didn’t see that native function being called in this pr . In metal 3.1 time we would still need this code path to have it working for older os

I meant we can have both implementations, native one and custom one, in the kernel and use the macro to determine which one should be used. (it's ok to have it in a follow-up, don't want to block this PR).

[kulinseth]

Small suggestion: As you mentioned, the MSL 3.1 supports nextafter natively. I think we could use __METAL_VERSION__ macro to branch off from the custom implementation used below Metal 3.1. Below is an example:

pytorch/aten/src/ATen/mps/IndexKernels.h

Lines 26 to 30 in d9fe171

	#if __METAL_VERSION__ >= 300
	constant IndexAB * indexAB [[buffer(0)]],
	#else
	constant IndexAB & indexAB [[buffer(0)]],
	#endif

I think he meant when that’s available he can call that function directly . I didn’t see that native function being called in this pr . In metal 3.1 time we would still need this code path to have it working for older os

I meant we can have both implementations, native one and custom one, in the kernel and use the macro to determine which one should be used. (it's ok to have it in a follow-up, don't want to block this PR).

Agreed . Let’s do a follow up on this case , for simple cases unless it becomes perf critical it’s nice to have consistent code path from maintenance stand point

[igm503]

Small suggestion: As you mentioned, the MSL 3.1 supports nextafter natively. I think we could use __METAL_VERSION__ macro to branch off from the custom implementation used below Metal 3.1. Below is an example:

pytorch/aten/src/ATen/mps/IndexKernels.h

Lines 26 to 30 in d9fe171

	#if __METAL_VERSION__ >= 300
	constant IndexAB * indexAB [[buffer(0)]],
	#else
	constant IndexAB & indexAB [[buffer(0)]],
	#endif

I think he meant when that’s available he can call that function directly . I didn’t see that native function being called in this pr . In metal 3.1 time we would still need this code path to have it working for older os

I meant we can have both implementations, native one and custom one, in the kernel and use the macro to determine which one should be used. (it's ok to have it in a follow-up, don't want to block this PR).

Agreed . Let’s do a follow up on this case , for simple cases unless it becomes perf critical it’s nice to have consistent code path from maintenance stand point

So I'll merge this, but y'all think it would be worth following qqaatw's idea of branching based on MSL version? I'm happy to throw that in on a follow up pr if so.

[igm503]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here