Skip to content

Conversation

@drbh
Copy link
Collaborator

@drbh drbh commented Nov 6, 2025

This PR adds support for a new metal-cpp kernel dependency. This is a follow up to the metal-cpp support in hf nix: huggingface/hf-nix#128 and enables kernels to use the cpp headers to drive metal kernels.

Changes:

  • adds dep to build2cmake
  • currently prefers the cpp relu over objc impl in examples/relu/build.toml. since they both produce metal output we can only build one at a time..
  • update xcrunHost command to respect the cryptex location of the latest metal toolchain on the host. This may need to be improved in line with Metal 4 Update to macOS SDK 26 and Metal 4 #290

example usage

# /// script
# requires-python = ">=3.10"
# dependencies = ["kernels", "torch"]
# ///
from kernels import get_local_kernel
import torch
from pathlib import Path

relu = get_local_kernel(Path("examples/relu/result"), "relu").relu

input = torch.tensor([-1.0, -1.5, 0.0, 2.0, 3.5], device="mps", dtype=torch.float16)
out = relu(input)
ref = torch.relu(input)

assert torch.allclose(out, ref), f"Float16 failed: {out} != {ref}"

print(out.cpu().numpy())
print(ref.cpu().numpy())

print("PASS")
# [0.  0.  0.  2.  3.5]
# [0.  0.  0.  2.  3.5]
# PASS

Comment on lines -19 to +31
"relu_metal/relu.mm",
"relu_metal/relu.metal",
"relu_metal/common.h",
"relu_metal_cpp/relu.cpp",
"relu_metal_cpp/metallib_loader.mm",
"relu_metal_cpp/relu_cpp.metal",
"relu_metal_cpp/common.h",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be best to keep relu as is and add an extra relu-metal-cpp example.


# Build inputs
apple-sdk_15,
metal-cpp,
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's keep them lexicographically sorted.

if [[ "$*" =~ "metallib" ]]; then
# If metallib is requested, find the air-lld from the Metal toolchain
METALLIB_BIN=$(ls /var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain*/Metal.xctoolchain/usr/bin/air-lld 2>/dev/null | head -n 1)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

metallib is a binary too :)

Also, this path is not stable, e.g. it doesn't exist on my macOS 26 installation.

Comment on lines +82 to +103
# Remove the '-sdk macosx metallib' and other unsupported flags from the command arguments
ARGS=$(echo "$@" | sed 's/-sdk macosx metallib //' | sed 's/-mmacosx-version-min=[^ ]* //')
# Add platform version for macOS 15+ to support Metal 3.2 / AIR 2.7
$METALLIB_BIN -platform_version macos 15.0 15.0 $ARGS
elif [[ "$*" =~ "metal" ]]; then
# If metal is requested, find the metal compiler from the Metal toolchain
METAL_BIN=$(ls /var/run/com.apple.security.cryptexd/mnt/com.apple.MobileAsset.MetalToolchain*/Metal.xctoolchain/usr/bin/metal 2>/dev/null | head -n 1)
if [ -z "$METAL_BIN" ]; then
echo "Error: Metal compiler not found" >&2
exit 1
fi
# Remove the '-sdk macosx metal' from the command arguments
ARGS=$(echo "$@" | sed 's/-sdk macosx metal //')
$METAL_BIN $ARGS
else
# In all other cases, just use the host xcrun
unset DEVELOPER_DIR
/usr/bin/xcrun $@
fi
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think most of this is not needed with #290 merged, we do not call xcrun anymore for metal and metallib.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants