From 903f61f89c728096f9cc2a60b20edae7d20cbda7 Mon Sep 17 00:00:00 2001 From: lvyufeng Date: Sun, 8 Mar 2026 15:05:55 +0800 Subject: [PATCH] mindtorch_v2/mps: real Metal GPU compute backend for 30+ ops Replace NumPy CPU fallbacks with actual Metal GPU compute kernels for element-wise, reduction, matmul, activation, and in-place ops. - Add metal_shaders.py: 105 MSL kernels (f32/f16 variants) - Add metal_compute.py: GPU dispatch engine with lazy compilation and pipeline caching (PyObjC + ctypes dual-path) - Extend runtime.py: compile_library, make_compute_pipeline, get_compute_encoder with ctypes fallbacks - Update mps_kernels.py: MPSMatrixMultiplication GPU matmul (4.3x faster than Accelerate BLAS at 1024x1024) - Update ops.py: GPU dispatch for 9 binary, 16 unary, 4 reduction, 4 NN fused, and 8 in-place ops with automatic f64/non-contiguous CPU fallback Co-Authored-By: Claude Opus 4.6 --- .../_backends/mps/metal_compute.py | 673 ++++++++++++++++++ .../_backends/mps/metal_shaders.py | 439 ++++++++++++ src/mindtorch_v2/_backends/mps/mps_kernels.py | 189 ++++- src/mindtorch_v2/_backends/mps/ops.py | 476 ++++++++++++- src/mindtorch_v2/_backends/mps/runtime.py | 99 +++ 5 files changed, 1866 insertions(+), 10 deletions(-) create mode 100644 src/mindtorch_v2/_backends/mps/metal_compute.py create mode 100644 src/mindtorch_v2/_backends/mps/metal_shaders.py diff --git a/src/mindtorch_v2/_backends/mps/metal_compute.py b/src/mindtorch_v2/_backends/mps/metal_compute.py new file mode 100644 index 000000000..a686c68db --- /dev/null +++ b/src/mindtorch_v2/_backends/mps/metal_compute.py @@ -0,0 +1,673 @@ +"""Metal GPU compute dispatch engine. + +Lazily compiles MSL kernels, caches compute pipeline states, and dispatches +element-wise, reduction, and in-place kernels on Metal GPU. +""" +import ctypes +import struct +import threading + +from .runtime import get_runtime, buffer_contents, _HAS_PYOBJC + +_MTLSize = None +if _HAS_PYOBJC: + from Metal import MTLSizeMake as _MTLSizeMake # pylint: disable=import-error,no-name-in-module + + def _MTLSize(w, h, d): + return _MTLSizeMake(w, h, d) +else: + def _MTLSize(w, h, d): # noqa: E302 + return None + +# --------------------------------------------------------------------------- +# Singleton dispatcher +# --------------------------------------------------------------------------- +_dispatcher = None +_dispatcher_lock = threading.Lock() + + +def get_dispatcher(): + """Return the singleton MetalKernelDispatcher (lazy init).""" + global _dispatcher + if _dispatcher is not None: + return _dispatcher + with _dispatcher_lock: + if _dispatcher is None: + _dispatcher = MetalKernelDispatcher() + return _dispatcher + + +class MetalKernelDispatcher: + """Compiles MSL kernels lazily, caches pipelines, dispatches compute work.""" + + def __init__(self): + self._library = None + self._pipeline_cache = {} + self._lock = threading.Lock() + + # ------------------------------------------------------------------ + # Compilation + # ------------------------------------------------------------------ + + def ensure_compiled(self): + """Compile all MSL source on first use.""" + if self._library is not None: + return + with self._lock: + if self._library is not None: + return + from .metal_shaders import MSL_SOURCE + rt = get_runtime() + self._library = rt.compile_library(MSL_SOURCE) + + def _get_pipeline(self, kernel_name): + """Get or create a cached compute pipeline for *kernel_name*.""" + if kernel_name in self._pipeline_cache: + return self._pipeline_cache[kernel_name] + self.ensure_compiled() + rt = get_runtime() + if _HAS_PYOBJC: + fn = self._library.newFunctionWithName_(kernel_name) + else: + fn = _library_get_function_ctypes(self._library, kernel_name) + if fn is None or (isinstance(fn, int) and fn == 0): + raise RuntimeError(f"Metal kernel '{kernel_name}' not found in library") + pipeline = rt.make_compute_pipeline(fn) + self._pipeline_cache[kernel_name] = pipeline + return pipeline + + # ------------------------------------------------------------------ + # Thread dispatch helpers + # ------------------------------------------------------------------ + + @staticmethod + def _threads_per_group(pipeline): + if _HAS_PYOBJC: + return min(256, int(pipeline.maxTotalThreadsPerThreadgroup())) + return 256 # safe default for Apple Silicon + + # ------------------------------------------------------------------ + # Dispatch: unary (a → out) + # ------------------------------------------------------------------ + + def dispatch_unary(self, kernel_name, a_metal_buf, out_metal_buf, numel): + """Encode and execute a unary element-wise kernel.""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_metal_buf, 0, 0) + enc.setBuffer_offset_atIndex_(out_metal_buf, 0, 1) + n_bytes = struct.pack("I", numel) + enc.setBytes_length_atIndex_(n_bytes, 4, 2) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_unary_ctypes(enc, pipeline, a_metal_buf, out_metal_buf, + numel, groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: binary (a, b → out) + # ------------------------------------------------------------------ + + def dispatch_binary(self, kernel_name, a_buf, b_buf, out_buf, numel): + """Encode and execute a binary element-wise kernel.""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBuffer_offset_atIndex_(b_buf, 0, 1) + enc.setBuffer_offset_atIndex_(out_buf, 0, 2) + n_bytes = struct.pack("I", numel) + enc.setBytes_length_atIndex_(n_bytes, 4, 3) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_binary_ctypes(enc, pipeline, a_buf, b_buf, out_buf, + numel, groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: binary with scalar (a, scalar → out) + # ------------------------------------------------------------------ + + def dispatch_binary_scalar(self, kernel_name, a_buf, scalar, out_buf, + numel, scalar_fmt="f"): + """Encode and execute a binary-scalar kernel (scalar embedded via setBytes).""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + scalar_bytes = struct.pack(scalar_fmt, scalar) + scalar_size = len(scalar_bytes) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBytes_length_atIndex_(scalar_bytes, scalar_size, 1) + enc.setBuffer_offset_atIndex_(out_buf, 0, 2) + n_bytes = struct.pack("I", numel) + enc.setBytes_length_atIndex_(n_bytes, 4, 3) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_binary_scalar_ctypes(enc, pipeline, a_buf, + scalar_bytes, scalar_size, + out_buf, numel, groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: reduction (input → scalar output, two-pass) + # ------------------------------------------------------------------ + + def dispatch_reduction(self, partial_kernel, final_kernel, a_buf, out_buf, + numel): + """Two-pass parallel reduction: per-threadgroup partials → final.""" + rt = get_runtime() + pipeline_p = self._get_pipeline(partial_kernel) + pipeline_f = self._get_pipeline(final_kernel) + tpg = self._threads_per_group(pipeline_p) + num_groups = (numel + tpg - 1) // tpg + + # Allocate partial-results buffer (one float per group) + partials_buf = rt.create_buffer(num_groups * 4) + + # --- Pass 1: per-threadgroup partial --- + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline_p) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBuffer_offset_atIndex_(partials_buf, 0, 1) + enc.setBytes_length_atIndex_(struct.pack("I", numel), 4, 2) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(num_groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_unary_ctypes(enc, pipeline_p, a_buf, partials_buf, + numel, num_groups, tpg) + rt.commit_and_wait(cmd) + + # --- Pass 2: reduce partials → single output --- + final_tpg = self._threads_per_group(pipeline_f) + # Final pass uses a single threadgroup + final_tpg = max(final_tpg, num_groups) + # Round up to next power-of-2 for proper tree reduction + final_tpg = 1 + while final_tpg < num_groups: + final_tpg *= 2 + final_tpg = min(final_tpg, 256) + + cmd2 = rt.create_command_buffer() + enc2 = rt.get_compute_encoder(cmd2) + if _HAS_PYOBJC: + enc2.setComputePipelineState_(pipeline_f) + enc2.setBuffer_offset_atIndex_(partials_buf, 0, 0) + enc2.setBuffer_offset_atIndex_(out_buf, 0, 1) + enc2.setBytes_length_atIndex_(struct.pack("I", num_groups), 4, 2) + enc2.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(1, 1, 1), _MTLSize(final_tpg, 1, 1)) + enc2.endEncoding() + else: + _encode_unary_ctypes(enc2, pipeline_f, partials_buf, out_buf, + num_groups, 1, final_tpg) + rt.commit_and_wait(cmd2) + + # ------------------------------------------------------------------ + # Dispatch: argmax/argmin reduction (two-pass with value+index) + # ------------------------------------------------------------------ + + def dispatch_arg_reduction(self, partial_kernel, final_kernel, + a_buf, out_buf, numel): + """Two-pass argmax/argmin: partials carry (value, index) pairs.""" + rt = get_runtime() + pipeline_p = self._get_pipeline(partial_kernel) + pipeline_f = self._get_pipeline(final_kernel) + tpg = self._threads_per_group(pipeline_p) + num_groups = (numel + tpg - 1) // tpg + + partial_vals_buf = rt.create_buffer(num_groups * 4) # float per group + partial_idxs_buf = rt.create_buffer(num_groups * 4) # uint per group + + # --- Pass 1 --- + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline_p) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBuffer_offset_atIndex_(partial_vals_buf, 0, 1) + enc.setBuffer_offset_atIndex_(partial_idxs_buf, 0, 2) + enc.setBytes_length_atIndex_(struct.pack("I", numel), 4, 3) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(num_groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_arg_partial_ctypes(enc, pipeline_p, a_buf, + partial_vals_buf, partial_idxs_buf, + numel, num_groups, tpg) + rt.commit_and_wait(cmd) + + # --- Pass 2 --- + final_tpg = 1 + while final_tpg < num_groups: + final_tpg *= 2 + final_tpg = min(final_tpg, 256) + + cmd2 = rt.create_command_buffer() + enc2 = rt.get_compute_encoder(cmd2) + if _HAS_PYOBJC: + enc2.setComputePipelineState_(pipeline_f) + enc2.setBuffer_offset_atIndex_(partial_vals_buf, 0, 0) + enc2.setBuffer_offset_atIndex_(partial_idxs_buf, 0, 1) + enc2.setBuffer_offset_atIndex_(out_buf, 0, 2) + enc2.setBytes_length_atIndex_(struct.pack("I", num_groups), 4, 3) + enc2.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(1, 1, 1), _MTLSize(final_tpg, 1, 1)) + enc2.endEncoding() + else: + _encode_arg_final_ctypes(enc2, pipeline_f, partial_vals_buf, + partial_idxs_buf, out_buf, + num_groups, final_tpg) + rt.commit_and_wait(cmd2) + + # ------------------------------------------------------------------ + # Dispatch: in-place unary (a → a) + # ------------------------------------------------------------------ + + def dispatch_inplace_unary(self, kernel_name, a_buf, numel): + """In-place unary: writes output back to input buffer.""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBytes_length_atIndex_(struct.pack("I", numel), 4, 1) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_inplace_unary_ctypes(enc, pipeline, a_buf, numel, + groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: in-place binary (a, b → a) + # ------------------------------------------------------------------ + + def dispatch_inplace_binary(self, kernel_name, a_buf, b_buf, numel): + """In-place binary: a[i] op= b[i].""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBuffer_offset_atIndex_(b_buf, 0, 1) + enc.setBytes_length_atIndex_(struct.pack("I", numel), 4, 2) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_unary_ctypes(enc, pipeline, a_buf, b_buf, numel, + groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: in-place binary scalar (a, scalar → a) + # ------------------------------------------------------------------ + + def dispatch_inplace_binary_scalar(self, kernel_name, a_buf, scalar, + numel, scalar_fmt="f"): + """In-place binary-scalar: a[i] op= scalar.""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + scalar_bytes = struct.pack(scalar_fmt, scalar) + scalar_size = len(scalar_bytes) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBytes_length_atIndex_(scalar_bytes, scalar_size, 1) + enc.setBytes_length_atIndex_(struct.pack("I", numel), 4, 2) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_inplace_scalar_ctypes(enc, pipeline, a_buf, + scalar_bytes, scalar_size, + numel, groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: fill (scalar → buffer) + # ------------------------------------------------------------------ + + def dispatch_fill(self, kernel_name, out_buf, scalar, numel, + scalar_fmt="f"): + """Fill buffer with a scalar value.""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + tpg = self._threads_per_group(pipeline) + groups = (numel + tpg - 1) // tpg + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + scalar_bytes = struct.pack(scalar_fmt, scalar) + scalar_size = len(scalar_bytes) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(out_buf, 0, 0) + enc.setBytes_length_atIndex_(scalar_bytes, scalar_size, 1) + enc.setBytes_length_atIndex_(struct.pack("I", numel), 4, 2) + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups, 1, 1), _MTLSize(tpg, 1, 1)) + enc.endEncoding() + else: + _encode_inplace_scalar_ctypes(enc, pipeline, out_buf, + scalar_bytes, scalar_size, + numel, groups, tpg) + + rt.commit_and_wait(cmd) + + # ------------------------------------------------------------------ + # Dispatch: copy (src → dst) + # ------------------------------------------------------------------ + + def dispatch_copy(self, kernel_name, src_buf, dst_buf, numel): + """Copy src buffer to dst buffer.""" + self.dispatch_binary(kernel_name, src_buf, dst_buf, numel) + + # ------------------------------------------------------------------ + # Dispatch: softmax 2D (input → output, with rows/cols) + # ------------------------------------------------------------------ + + def dispatch_softmax_2d(self, kernel_name, a_buf, out_buf, rows, cols): + """Dispatch softmax over last dim of a 2D tensor.""" + rt = get_runtime() + pipeline = self._get_pipeline(kernel_name) + + cmd = rt.create_command_buffer() + enc = rt.get_compute_encoder(cmd) + + if _HAS_PYOBJC: + enc.setComputePipelineState_(pipeline) + enc.setBuffer_offset_atIndex_(a_buf, 0, 0) + enc.setBuffer_offset_atIndex_(out_buf, 0, 1) + enc.setBytes_length_atIndex_(struct.pack("I", rows), 4, 2) + enc.setBytes_length_atIndex_(struct.pack("I", cols), 4, 3) + tpg_x = min(32, cols) + tpg_y = min(8, rows) + groups_x = (cols + tpg_x - 1) // tpg_x + groups_y = (rows + tpg_y - 1) // tpg_y + enc.dispatchThreadgroups_threadsPerThreadgroup_( + _MTLSize(groups_x, groups_y, 1), _MTLSize(tpg_x, tpg_y, 1)) + enc.endEncoding() + else: + _encode_softmax_ctypes(enc, pipeline, a_buf, out_buf, + rows, cols) + + rt.commit_and_wait(cmd) + + +# --------------------------------------------------------------------------- +# ctypes encoding helpers (fallback when no PyObjC) +# --------------------------------------------------------------------------- + +def _library_get_function_ctypes(library, name): + """Get a MTLFunction from a MTLLibrary by name (ctypes path).""" + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + # Create NSString for function name + ns_string_class = _libobjc.objc_getClass(b"NSString") + sel_alloc = _libobjc.sel_registerName(b"alloc") + sel_init = _libobjc.sel_registerName(b"initWithUTF8String:") + ns_str = _libobjc.objc_msgSend(ns_string_class, sel_alloc) + name_bytes = name.encode("utf-8") + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_char_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + ns_str = _libobjc.objc_msgSend(ns_str, sel_init, name_bytes) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + + sel = _libobjc.sel_registerName(b"newFunctionWithName:") + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + fn = _libobjc.objc_msgSend(library, sel, ns_str) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + return fn + + +def _ctypes_set_buffer(enc, buf, offset, index): + """setBuffer:offset:atIndex: via ctypes.""" + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + sel = _libobjc.sel_registerName(b"setBuffer:offset:atIndex:") + _libobjc.objc_msgSend.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint64, + ] + _libobjc.objc_msgSend.restype = None + _libobjc.objc_msgSend(enc, sel, buf, ctypes.c_uint64(offset), + ctypes.c_uint64(index)) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + + +def _ctypes_set_bytes(enc, data_bytes, length, index): + """setBytes:length:atIndex: via ctypes.""" + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + sel = _libobjc.sel_registerName(b"setBytes:length:atIndex:") + buf = ctypes.create_string_buffer(data_bytes) + _libobjc.objc_msgSend.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p, ctypes.c_uint64, ctypes.c_uint64, + ] + _libobjc.objc_msgSend.restype = None + _libobjc.objc_msgSend(enc, sel, buf, ctypes.c_uint64(length), + ctypes.c_uint64(index)) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + + +def _ctypes_set_pipeline(enc, pipeline): + """setComputePipelineState: via ctypes.""" + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + sel = _libobjc.sel_registerName(b"setComputePipelineState:") + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = None + _libobjc.objc_msgSend(enc, sel, pipeline) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + + +def _ctypes_dispatch_threadgroups(enc, groups, tpg): + """dispatchThreadgroups:threadsPerThreadgroup: via ctypes. + + MTLSize is a struct {uint64, uint64, uint64}. + """ + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + + class MTLSize(ctypes.Structure): + _fields_ = [("width", ctypes.c_uint64), + ("height", ctypes.c_uint64), + ("depth", ctypes.c_uint64)] + + sel = _libobjc.sel_registerName(b"dispatchThreadgroups:threadsPerThreadgroup:") + grid = MTLSize(groups, 1, 1) + tpg_s = MTLSize(tpg, 1, 1) + _libobjc.objc_msgSend.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, MTLSize, MTLSize, + ] + _libobjc.objc_msgSend.restype = None + _libobjc.objc_msgSend(enc, sel, grid, tpg_s) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + + +def _ctypes_end_encoding(enc): + """endEncoding via ctypes.""" + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + sel = _libobjc.sel_registerName(b"endEncoding") + _libobjc.objc_msgSend.restype = None + _libobjc.objc_msgSend(enc, sel) + _libobjc.objc_msgSend.restype = ctypes.c_void_p + + +def _encode_unary_ctypes(enc, pipeline, a_buf, out_buf, numel, groups, tpg): + """Encode a unary kernel dispatch via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_buffer(enc, out_buf, 0, 1) + _ctypes_set_bytes(enc, struct.pack("I", numel), 4, 2) + _ctypes_dispatch_threadgroups(enc, groups, tpg) + _ctypes_end_encoding(enc) + + +def _encode_binary_ctypes(enc, pipeline, a_buf, b_buf, out_buf, + numel, groups, tpg): + """Encode a binary kernel dispatch via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_buffer(enc, b_buf, 0, 1) + _ctypes_set_buffer(enc, out_buf, 0, 2) + _ctypes_set_bytes(enc, struct.pack("I", numel), 4, 3) + _ctypes_dispatch_threadgroups(enc, groups, tpg) + _ctypes_end_encoding(enc) + + +def _encode_binary_scalar_ctypes(enc, pipeline, a_buf, scalar_bytes, + scalar_size, out_buf, numel, groups, tpg): + """Encode a binary-scalar kernel dispatch via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_bytes(enc, scalar_bytes, scalar_size, 1) + _ctypes_set_buffer(enc, out_buf, 0, 2) + _ctypes_set_bytes(enc, struct.pack("I", numel), 4, 3) + _ctypes_dispatch_threadgroups(enc, groups, tpg) + _ctypes_end_encoding(enc) + + +def _encode_arg_partial_ctypes(enc, pipeline, a_buf, vals_buf, idxs_buf, + numel, groups, tpg): + """Encode argmax/argmin partial pass via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_buffer(enc, vals_buf, 0, 1) + _ctypes_set_buffer(enc, idxs_buf, 0, 2) + _ctypes_set_bytes(enc, struct.pack("I", numel), 4, 3) + _ctypes_dispatch_threadgroups(enc, groups, tpg) + _ctypes_end_encoding(enc) + + +def _encode_arg_final_ctypes(enc, pipeline, vals_buf, idxs_buf, out_buf, + num_groups, tpg): + """Encode argmax/argmin final pass via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, vals_buf, 0, 0) + _ctypes_set_buffer(enc, idxs_buf, 0, 1) + _ctypes_set_buffer(enc, out_buf, 0, 2) + _ctypes_set_bytes(enc, struct.pack("I", num_groups), 4, 3) + _ctypes_dispatch_threadgroups(enc, 1, tpg) + _ctypes_end_encoding(enc) + + +def _encode_inplace_unary_ctypes(enc, pipeline, a_buf, numel, groups, tpg): + """Encode an in-place unary kernel via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_bytes(enc, struct.pack("I", numel), 4, 1) + _ctypes_dispatch_threadgroups(enc, groups, tpg) + _ctypes_end_encoding(enc) + + +def _encode_inplace_scalar_ctypes(enc, pipeline, a_buf, scalar_bytes, + scalar_size, numel, groups, tpg): + """Encode an in-place binary-scalar kernel via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_bytes(enc, scalar_bytes, scalar_size, 1) + _ctypes_set_bytes(enc, struct.pack("I", numel), 4, 2) + _ctypes_dispatch_threadgroups(enc, groups, tpg) + _ctypes_end_encoding(enc) + + +def _encode_softmax_ctypes(enc, pipeline, a_buf, out_buf, rows, cols): + """Encode softmax 2D kernel via ctypes.""" + _ctypes_set_pipeline(enc, pipeline) + _ctypes_set_buffer(enc, a_buf, 0, 0) + _ctypes_set_buffer(enc, out_buf, 0, 1) + _ctypes_set_bytes(enc, struct.pack("I", rows), 4, 2) + _ctypes_set_bytes(enc, struct.pack("I", cols), 4, 3) + tpg_x = min(32, cols) + tpg_y = min(8, rows) + groups_x = (cols + tpg_x - 1) // tpg_x + groups_y = (rows + tpg_y - 1) // tpg_y + + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + + class MTLSize(ctypes.Structure): + _fields_ = [("width", ctypes.c_uint64), + ("height", ctypes.c_uint64), + ("depth", ctypes.c_uint64)] + + sel = _libobjc.sel_registerName(b"dispatchThreadgroups:threadsPerThreadgroup:") + grid = MTLSize(groups_x, groups_y, 1) + tpg_s = MTLSize(tpg_x, tpg_y, 1) + _libobjc.objc_msgSend.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, MTLSize, MTLSize, + ] + _libobjc.objc_msgSend.restype = None + _libobjc.objc_msgSend(enc, sel, grid, tpg_s) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + _ctypes_end_encoding(enc) diff --git a/src/mindtorch_v2/_backends/mps/metal_shaders.py b/src/mindtorch_v2/_backends/mps/metal_shaders.py new file mode 100644 index 000000000..132dce3c7 --- /dev/null +++ b/src/mindtorch_v2/_backends/mps/metal_shaders.py @@ -0,0 +1,439 @@ +"""Metal Shading Language (MSL) kernel source strings for GPU compute. + +All kernels are compiled once at first use via MetalKernelDispatcher. +Naming convention: _f32 / _f16 for each dtype variant. +""" + +# --------------------------------------------------------------------------- +# Helper: generate element-wise kernels for both f32 and f16 +# --------------------------------------------------------------------------- + +_UNARY_TEMPLATE = """ +kernel void {name}_{suffix}(device const {type}* a [[buffer(0)]], + device {type}* out [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) {{ + if (id < N) {{ + {type} x = a[id]; + out[id] = {expr}; + }} +}} +""" + +_BINARY_TEMPLATE = """ +kernel void {name}_{suffix}(device const {type}* a [[buffer(0)]], + device const {type}* b [[buffer(1)]], + device {type}* out [[buffer(2)]], + constant uint& N [[buffer(3)]], + uint id [[thread_position_in_grid]]) {{ + if (id < N) out[id] = {expr}; +}} +""" + +_BINARY_SCALAR_TEMPLATE = """ +kernel void {name}_scalar_{suffix}(device const {type}* a [[buffer(0)]], + constant {type}& scalar [[buffer(1)]], + device {type}* out [[buffer(2)]], + constant uint& N [[buffer(3)]], + uint id [[thread_position_in_grid]]) {{ + if (id < N) out[id] = {expr}; +}} +""" + + +def _gen_unary(name, expr, types=("float", "half")): + """Generate unary kernel source for given types.""" + parts = [] + for t in types: + suffix = "f32" if t == "float" else "f16" + parts.append(_UNARY_TEMPLATE.format(name=name, suffix=suffix, type=t, expr=expr)) + return "".join(parts) + + +def _gen_binary(name, expr, types=("float", "half")): + """Generate binary kernel source for given types.""" + parts = [] + for t in types: + suffix = "f32" if t == "float" else "f16" + parts.append(_BINARY_TEMPLATE.format(name=name, suffix=suffix, type=t, expr=expr)) + scalar_expr = expr.replace("b[id]", "scalar") + parts.append(_BINARY_SCALAR_TEMPLATE.format(name=name, suffix=suffix, type=t, expr=scalar_expr)) + return "".join(parts) + + +# --------------------------------------------------------------------------- +# Build the full MSL source +# --------------------------------------------------------------------------- + +_HEADER = """ +#include +using namespace metal; +""" + +# --- Binary element-wise --- +_BINARY_OPS = ( + ("add", "a[id] + b[id]"), + ("sub", "a[id] - b[id]"), + ("mul", "a[id] * b[id]"), + ("div", "a[id] / b[id]"), + ("maximum", "max(a[id], b[id])"), + ("minimum", "min(a[id], b[id])"), + ("pow", "pow(a[id], b[id])"), + ("fmod", "fmod(a[id], b[id])"), + ("remainder", "a[id] - b[id] * floor(a[id] / b[id])"), +) + +# --- Unary element-wise --- +_UNARY_OPS = ( + ("neg", "-x"), + ("abs", "abs(x)"), + ("sqrt", "sqrt(x)"), + ("rsqrt", "rsqrt(x)"), + ("exp", "exp(x)"), + ("log", "log(x)"), + ("log2", "log2(x)"), + ("sin", "sin(x)"), + ("cos", "cos(x)"), + ("tanh", "tanh(x)"), + ("sigmoid", "1.0 / (1.0 + exp(-x))"), + ("sign", "sign(x)"), + ("floor", "floor(x)"), + ("ceil", "ceil(x)"), + ("round", "rint(x)"), + ("relu", "max(x, ({type})0)"), + ("gelu", "0.5f * x * (1.0f + tanh(0.7978845608f * (x + 0.044715f * x * x * x)))"), + ("silu", "x / (1.0f + exp(-x))"), +) + +# --- Reduction kernels (two-pass parallel) --- +_REDUCTION_SOURCE = """ +// ---- sum reduction ---- +kernel void sum_partial_f32(device const float* input [[buffer(0)]], + device float* partials [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint gid [[thread_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint group_id [[threadgroup_position_in_grid]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float shared[256]; + float val = (gid < N) ? input[gid] : 0.0f; + shared[lid] = val; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s) shared[lid] += shared[lid + s]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) partials[group_id] = shared[0]; +} + +kernel void sum_final_f32(device const float* partials [[buffer(0)]], + device float* output [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint lid [[thread_position_in_threadgroup]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float shared[256]; + float val = (lid < N) ? partials[lid] : 0.0f; + shared[lid] = val; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s) shared[lid] += shared[lid + s]; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) output[0] = shared[0]; +} + +// ---- max reduction ---- +kernel void max_partial_f32(device const float* input [[buffer(0)]], + device float* partials [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint gid [[thread_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint group_id [[threadgroup_position_in_grid]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float shared[256]; + float val = (gid < N) ? input[gid] : -INFINITY; + shared[lid] = val; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s) shared[lid] = max(shared[lid], shared[lid + s]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) partials[group_id] = shared[0]; +} + +kernel void max_final_f32(device const float* partials [[buffer(0)]], + device float* output [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint lid [[thread_position_in_threadgroup]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float shared[256]; + float val = (lid < N) ? partials[lid] : -INFINITY; + shared[lid] = val; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s) shared[lid] = max(shared[lid], shared[lid + s]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) output[0] = shared[0]; +} + +// ---- min reduction ---- +kernel void min_partial_f32(device const float* input [[buffer(0)]], + device float* partials [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint gid [[thread_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint group_id [[threadgroup_position_in_grid]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float shared[256]; + float val = (gid < N) ? input[gid] : INFINITY; + shared[lid] = val; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s) shared[lid] = min(shared[lid], shared[lid + s]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) partials[group_id] = shared[0]; +} + +kernel void min_final_f32(device const float* partials [[buffer(0)]], + device float* output [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint lid [[thread_position_in_threadgroup]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float shared[256]; + float val = (lid < N) ? partials[lid] : INFINITY; + shared[lid] = val; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s) shared[lid] = min(shared[lid], shared[lid + s]); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) output[0] = shared[0]; +} + +// ---- argmax reduction ---- +kernel void argmax_partial_f32(device const float* input [[buffer(0)]], + device float* partial_vals [[buffer(1)]], + device uint* partial_idxs [[buffer(2)]], + constant uint& N [[buffer(3)]], + uint gid [[thread_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint group_id [[threadgroup_position_in_grid]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float s_val[256]; + threadgroup uint s_idx[256]; + float val = (gid < N) ? input[gid] : -INFINITY; + uint idx = gid; + s_val[lid] = val; + s_idx[lid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s && s_val[lid + s] > s_val[lid]) { + s_val[lid] = s_val[lid + s]; + s_idx[lid] = s_idx[lid + s]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) { + partial_vals[group_id] = s_val[0]; + partial_idxs[group_id] = s_idx[0]; + } +} + +kernel void argmax_final_f32(device const float* partial_vals [[buffer(0)]], + device const uint* partial_idxs [[buffer(1)]], + device uint* output [[buffer(2)]], + constant uint& N [[buffer(3)]], + uint lid [[thread_position_in_threadgroup]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float s_val[256]; + threadgroup uint s_idx[256]; + float val = (lid < N) ? partial_vals[lid] : -INFINITY; + uint idx = (lid < N) ? partial_idxs[lid] : 0; + s_val[lid] = val; + s_idx[lid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s && s_val[lid + s] > s_val[lid]) { + s_val[lid] = s_val[lid + s]; + s_idx[lid] = s_idx[lid + s]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) output[0] = s_idx[0]; +} + +// ---- argmin reduction ---- +kernel void argmin_partial_f32(device const float* input [[buffer(0)]], + device float* partial_vals [[buffer(1)]], + device uint* partial_idxs [[buffer(2)]], + constant uint& N [[buffer(3)]], + uint gid [[thread_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint group_id [[threadgroup_position_in_grid]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float s_val[256]; + threadgroup uint s_idx[256]; + float val = (gid < N) ? input[gid] : INFINITY; + uint idx = gid; + s_val[lid] = val; + s_idx[lid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s && s_val[lid + s] < s_val[lid]) { + s_val[lid] = s_val[lid + s]; + s_idx[lid] = s_idx[lid + s]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) { + partial_vals[group_id] = s_val[0]; + partial_idxs[group_id] = s_idx[0]; + } +} + +kernel void argmin_final_f32(device const float* partial_vals [[buffer(0)]], + device const uint* partial_idxs [[buffer(1)]], + device uint* output [[buffer(2)]], + constant uint& N [[buffer(3)]], + uint lid [[thread_position_in_threadgroup]], + uint group_size [[threads_per_threadgroup]]) { + threadgroup float s_val[256]; + threadgroup uint s_idx[256]; + float val = (lid < N) ? partial_vals[lid] : INFINITY; + uint idx = (lid < N) ? partial_idxs[lid] : 0; + s_val[lid] = val; + s_idx[lid] = idx; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint s = group_size / 2; s > 0; s >>= 1) { + if (lid < s && s_val[lid + s] < s_val[lid]) { + s_val[lid] = s_val[lid + s]; + s_idx[lid] = s_idx[lid + s]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + if (lid == 0) output[0] = s_idx[0]; +} + +// ---- fill kernel ---- +kernel void fill_f32(device float* out [[buffer(0)]], + constant float& val [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) { + if (id < N) out[id] = val; +} + +kernel void fill_f16(device half* out [[buffer(0)]], + constant half& val [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) { + if (id < N) out[id] = val; +} + +// ---- copy kernel ---- +kernel void copy_f32(device const float* src [[buffer(0)]], + device float* dst [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) { + if (id < N) dst[id] = src[id]; +} + +kernel void copy_f16(device const half* src [[buffer(0)]], + device half* dst [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) { + if (id < N) dst[id] = src[id]; +} + +// ---- softmax (per-row, last dim) ---- +kernel void softmax_f32(device const float* input [[buffer(0)]], + device float* output [[buffer(1)]], + constant uint& rows [[buffer(2)]], + constant uint& cols [[buffer(3)]], + uint2 pos [[thread_position_in_grid]]) { + uint row = pos.y; + uint col = pos.x; + if (row >= rows || col >= cols) return; + + // Find max in this row + float max_val = -INFINITY; + for (uint c = 0; c < cols; c++) { + max_val = max(max_val, input[row * cols + c]); + } + + // Compute exp(x - max) + float exp_val = exp(input[row * cols + col] - max_val); + + // Sum of exp in this row + float sum_exp = 0.0f; + for (uint c = 0; c < cols; c++) { + sum_exp += exp(input[row * cols + c] - max_val); + } + + output[row * cols + col] = exp_val / sum_exp; +} +""" + +# --- In-place binary kernels --- +_INPLACE_BINARY_TEMPLATE = """ +kernel void {name}_inplace_{suffix}(device {type}* a [[buffer(0)]], + device const {type}* b [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) {{ + if (id < N) a[id] {op}= b[id]; +}} + +kernel void {name}_inplace_scalar_{suffix}(device {type}* a [[buffer(0)]], + constant {type}& scalar [[buffer(1)]], + constant uint& N [[buffer(2)]], + uint id [[thread_position_in_grid]]) {{ + if (id < N) a[id] {op}= scalar; +}} +""" + +_INPLACE_UNARY_SOURCE = """ +kernel void relu_inplace_f32(device float* a [[buffer(0)]], + constant uint& N [[buffer(1)]], + uint id [[thread_position_in_grid]]) { + if (id < N && a[id] < 0.0f) a[id] = 0.0f; +} + +kernel void relu_inplace_f16(device half* a [[buffer(0)]], + constant uint& N [[buffer(1)]], + uint id [[thread_position_in_grid]]) { + if (id < N && a[id] < (half)0.0f) a[id] = (half)0.0f; +} +""" + + +def _build_msl_source(): + """Assemble the complete MSL source string.""" + parts = [_HEADER] + + # Binary ops + for name, expr in _BINARY_OPS: + parts.append(_gen_binary(name, expr)) + + # Unary ops + for name, expr in _UNARY_OPS: + # Handle {type} placeholder in expressions like relu + for t, suffix in (("float", "f32"), ("half", "f16")): + e = expr.replace("{type}", t) + parts.append(_UNARY_TEMPLATE.format(name=name, suffix=suffix, type=t, expr=e)) + + # Reductions and special kernels + parts.append(_REDUCTION_SOURCE) + + # In-place binary ops + for name, op in (("add", "+"), ("sub", "-"), ("mul", "*"), ("div", "/")): + for t, suffix in (("float", "f32"), ("half", "f16")): + parts.append(_INPLACE_BINARY_TEMPLATE.format( + name=name, op=op, suffix=suffix, type=t)) + + parts.append(_INPLACE_UNARY_SOURCE) + + return "".join(parts) + + +MSL_SOURCE = _build_msl_source() diff --git a/src/mindtorch_v2/_backends/mps/mps_kernels.py b/src/mindtorch_v2/_backends/mps/mps_kernels.py index 1be033292..b527b2595 100644 --- a/src/mindtorch_v2/_backends/mps/mps_kernels.py +++ b/src/mindtorch_v2/_backends/mps/mps_kernels.py @@ -8,15 +8,20 @@ pip install pyobjc-framework-Metal pyobjc-framework-MetalPerformanceShaders """ +import ctypes import numpy as np +from .runtime import get_runtime, buffer_contents, _HAS_PYOBJC + # --------------------------------------------------------------------------- # Try to import pyobjc MPS bindings # --------------------------------------------------------------------------- _HAS_MPS_PYOBJC = False +_MPS = None +_Metal = None try: - import Metal # noqa: F401 - import MetalPerformanceShaders as MPS # noqa: F401 + import Metal as _Metal # noqa: F401 + import MetalPerformanceShaders as _MPS # noqa: F401 _HAS_MPS_PYOBJC = True except ImportError: pass @@ -44,9 +49,6 @@ def mps_matmul(a_np, b_np): ------- np.ndarray, shape (M, N) """ - # For now, always use Accelerate BLAS which is already highly optimised - # on Apple Silicon. When pyobjc MPS is available, large matrices could - # be routed through MPSMatrixMultiplication for additional GPU parallelism. M, K = a_np.shape K2, N = b_np.shape assert K == K2, f"inner dimensions must match: {K} vs {K2}" @@ -64,7 +66,182 @@ def mps_matmul(a_np, b_np): b_np.ctypes.data, N, 0.0, out.ctypes.data, N) else: - # Non-float types: fall through to numpy np.matmul(a_np, b_np, out=out) return out + + +# --------------------------------------------------------------------------- +# GPU MatMul via MPSMatrixMultiplication (PyObjC path) +# --------------------------------------------------------------------------- + +# MPSDataType constants +_MPSDataTypeFloat32 = 0x10000000 | 32 # 268435488 +_MPSDataTypeFloat16 = 0x10000000 | 16 # 268435472 + + +def _mps_dtype_code(np_dtype): + """Map numpy dtype to MPSDataType constant.""" + if np_dtype == np.float32: + return _MPSDataTypeFloat32 + if np_dtype == np.float16: + return _MPSDataTypeFloat16 + return None + + +def mps_matmul_gpu(a_metal_buf, b_metal_buf, M, K, N, dtype_code, + itemsize): + """GPU matrix multiply via MPSMatrixMultiplication. + + Parameters + ---------- + a_metal_buf : Metal buffer, shape (M, K) row-major + b_metal_buf : Metal buffer, shape (K, N) row-major + M, K, N : int — matrix dimensions + dtype_code : int — MPSDataType constant + itemsize : int — bytes per element (4 for f32, 2 for f16) + + Returns + ------- + out_metal_buf : Metal buffer containing result (M, N) + """ + rt = get_runtime() + out_nbytes = M * N * itemsize + out_buf = rt.create_buffer(max(out_nbytes, 1)) + + if _HAS_MPS_PYOBJC: + return _mps_matmul_pyobjc(rt, a_metal_buf, b_metal_buf, + out_buf, M, K, N, dtype_code, itemsize) + return _mps_matmul_ctypes(rt, a_metal_buf, b_metal_buf, + out_buf, M, K, N, dtype_code, itemsize) + + +def _mps_matmul_pyobjc(rt, a_buf, b_buf, out_buf, M, K, N, + dtype_code, itemsize): + """MPSMatrixMultiplication via PyObjC.""" + # Create matrix descriptors via class factory method + a_desc = _MPS.MPSMatrixDescriptor.matrixDescriptorWithRows_columns_rowBytes_dataType_( + M, K, K * itemsize, dtype_code) + b_desc = _MPS.MPSMatrixDescriptor.matrixDescriptorWithRows_columns_rowBytes_dataType_( + K, N, N * itemsize, dtype_code) + c_desc = _MPS.MPSMatrixDescriptor.matrixDescriptorWithRows_columns_rowBytes_dataType_( + M, N, N * itemsize, dtype_code) + + # Wrap Metal buffers as MPSMatrix + a_mat = _MPS.MPSMatrix.alloc().initWithBuffer_descriptor_(a_buf, a_desc) + b_mat = _MPS.MPSMatrix.alloc().initWithBuffer_descriptor_(b_buf, b_desc) + c_mat = _MPS.MPSMatrix.alloc().initWithBuffer_descriptor_(out_buf, c_desc) + + # Create and encode MPSMatrixMultiplication (C = alpha*A*B + beta*C) + mm_kernel = _MPS.MPSMatrixMultiplication.alloc().initWithDevice_transposeLeft_transposeRight_resultRows_resultColumns_interiorColumns_alpha_beta_( + rt.device, False, False, M, N, K, 1.0, 0.0) + + cmd = rt.create_command_buffer() + mm_kernel.encodeToCommandBuffer_leftMatrix_rightMatrix_resultMatrix_( + cmd, a_mat, b_mat, c_mat) + rt.commit_and_wait(cmd) + + return out_buf + + +def _mps_matmul_ctypes(rt, a_buf, b_buf, out_buf, M, K, N, + dtype_code, itemsize): + """MPSMatrixMultiplication via ctypes/libobjc fallback. + + This is complex due to the number of ObjC messages required. + Falls back to CPU BLAS if any step fails. + """ + try: + from .runtime import _libobjc, _load_objc_libs + _load_objc_libs() + + def _cls(name): + return _libobjc.objc_getClass(name.encode()) + + def _sel(name): + return _libobjc.sel_registerName(name.encode()) + + def _msg(obj, sel_name, *args, argtypes=None, restype=ctypes.c_void_p): + sel = _sel(sel_name) + if argtypes: + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + argtypes + _libobjc.objc_msgSend.restype = restype + result = _libobjc.objc_msgSend(obj, sel, *args) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + return result + + # Create MPSMatrixDescriptor via class factory method + desc_cls = _cls("MPSMatrixDescriptor") + a_desc = _msg( + desc_cls, + "matrixDescriptorWithRows:columns:rowBytes:dataType:", + ctypes.c_uint64(M), ctypes.c_uint64(K), + ctypes.c_uint64(K * itemsize), ctypes.c_uint32(dtype_code), + argtypes=[ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint32], + ) + b_desc = _msg( + desc_cls, + "matrixDescriptorWithRows:columns:rowBytes:dataType:", + ctypes.c_uint64(K), ctypes.c_uint64(N), + ctypes.c_uint64(N * itemsize), ctypes.c_uint32(dtype_code), + argtypes=[ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint32], + ) + c_desc = _msg( + desc_cls, + "matrixDescriptorWithRows:columns:rowBytes:dataType:", + ctypes.c_uint64(M), ctypes.c_uint64(N), + ctypes.c_uint64(N * itemsize), ctypes.c_uint32(dtype_code), + argtypes=[ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint32], + ) + + # Create MPSMatrix objects + mat_cls = _cls("MPSMatrix") + a_mat = _msg( + _msg(mat_cls, "alloc"), + "initWithBuffer:descriptor:", + a_buf, a_desc, + argtypes=[ctypes.c_void_p, ctypes.c_void_p], + ) + b_mat = _msg( + _msg(mat_cls, "alloc"), + "initWithBuffer:descriptor:", + b_buf, b_desc, + argtypes=[ctypes.c_void_p, ctypes.c_void_p], + ) + c_mat = _msg( + _msg(mat_cls, "alloc"), + "initWithBuffer:descriptor:", + out_buf, c_desc, + argtypes=[ctypes.c_void_p, ctypes.c_void_p], + ) + + # Create MPSMatrixMultiplication kernel + mm_cls = _cls("MPSMatrixMultiplication") + mm_kernel = _msg( + _msg(mm_cls, "alloc"), + "initWithDevice:transposeLeft:transposeRight:resultRows:resultColumns:interiorColumns:alpha:beta:", + rt.device, + ctypes.c_bool(False), ctypes.c_bool(False), + ctypes.c_uint64(M), ctypes.c_uint64(N), ctypes.c_uint64(K), + ctypes.c_double(1.0), ctypes.c_double(0.0), + argtypes=[ctypes.c_void_p, ctypes.c_bool, ctypes.c_bool, + ctypes.c_uint64, ctypes.c_uint64, ctypes.c_uint64, + ctypes.c_double, ctypes.c_double], + ) + + # Encode to command buffer + cmd = rt.create_command_buffer() + _msg(mm_kernel, + "encodeToCommandBuffer:leftMatrix:rightMatrix:resultMatrix:", + cmd, a_mat, b_mat, c_mat, + argtypes=[ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p, ctypes.c_void_p], + restype=None) + rt.commit_and_wait(cmd) + + return out_buf + + except Exception: + # If ctypes path fails, fall back to CPU BLAS + return None diff --git a/src/mindtorch_v2/_backends/mps/ops.py b/src/mindtorch_v2/_backends/mps/ops.py index 2a1bc9468..dbf5c6d01 100644 --- a/src/mindtorch_v2/_backends/mps/ops.py +++ b/src/mindtorch_v2/_backends/mps/ops.py @@ -1,15 +1,89 @@ import math +import ctypes +import struct import numpy as np from ..._dtype import bool as bool_dtype from ..._dtype import int64 as int64_dtype +from ..._dtype import float16 as float16_dtype from ..._dtype import float32 as float32_dtype from ..._dtype import float64 as float64_dtype from ..._dtype import to_numpy_dtype -from ..._storage import mps_typed_storage_from_numpy +from ..._storage import mps_typed_storage_from_numpy, _MPSUntypedStorage, TypedStorage from ..._tensor import Tensor from . import accelerate as _accel +# --------------------------------------------------------------------------- +# GPU dispatch helpers +# --------------------------------------------------------------------------- +_GPU_FLOAT_DTYPES = frozenset({float32_dtype, float16_dtype}) + + +def _can_use_gpu(t): + """Check if tensor can use Metal GPU kernels.""" + return (t.dtype in _GPU_FLOAT_DTYPES + and t.is_contiguous() + and t.numel() > 0 + and hasattr(t._storage, '_untyped') + and hasattr(t._storage._untyped, '_metal_buffer') + and t._storage._untyped._metal_buffer is not None) + + +def _metal_buf(t): + """Get the raw Metal buffer from a tensor.""" + return t._storage._untyped._metal_buffer + + +def _kernel_suffix(dtype): + """Return MSL kernel suffix for dtype.""" + return "f16" if dtype == float16_dtype else "f32" + + +def _scalar_fmt(dtype): + """Return struct format char for scalar encoding.""" + return "e" if dtype == float16_dtype else "f" + + +def _itemsize(dtype): + """Return byte size per element.""" + return 2 if dtype == float16_dtype else 4 + + +def _alloc_output_buf(numel, dtype): + """Allocate a Metal buffer for output.""" + from .runtime import get_runtime + rt = get_runtime() + return rt.create_buffer(numel * _itemsize(dtype)) + + +def _from_metal_buffer(metal_buf, shape, stride, dtype, device): + """Wrap an existing Metal buffer into a Tensor without copying data.""" + from .runtime import buffer_contents + nbytes = 1 + for s in shape: + nbytes *= s + nbytes *= _itemsize(dtype) + nbytes = max(nbytes, 1) + untyped = _MPSUntypedStorage(metal_buf, nbytes, device=device) + ptr = buffer_contents(metal_buf) + numel = 1 + for s in shape: + numel *= s + data = np.frombuffer( + (ctypes.c_uint8 * nbytes).from_address(ptr), + dtype=to_numpy_dtype(dtype), + count=max(numel, 1), + ) + size = max(numel, 1) + storage = TypedStorage(untyped, dtype, size, data=data) + return Tensor(storage, shape, stride) + + +def _get_dispatcher(): + """Lazy import of the Metal kernel dispatcher singleton.""" + from .metal_compute import get_dispatcher + return get_dispatcher() + def _to_numpy(t): return t._numpy_view() @@ -28,18 +102,60 @@ def _from_numpy(arr, dtype, device): def add(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"add_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"add_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b return _from_numpy(a_np + b_np, a.dtype, a.device) def mul(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"mul_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"mul_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b return _from_numpy(a_np * b_np, a.dtype, a.device) def div(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"div_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"div_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b out = np.true_divide(a_np, b_np) @@ -74,15 +190,33 @@ def _blas_gemm(a_np, b_np, dtype): def matmul(a, b): + # GPU path: MPSMatrixMultiplication for 2D contiguous float32/float16 + if (_can_use_gpu(a) and _can_use_gpu(b) + and len(a.shape) == 2 and len(b.shape) == 2): + from .mps_kernels import mps_matmul_gpu, _mps_dtype_code + np_dt = to_numpy_dtype(a.dtype) + dtype_code = _mps_dtype_code(np_dt) + if dtype_code is not None: + M, K = a.shape + K2, N = b.shape + if K == K2: + out_buf = mps_matmul_gpu( + _metal_buf(a), _metal_buf(b), + M, K, N, dtype_code, _itemsize(a.dtype)) + if out_buf is not None: + from ..._tensor import _compute_strides + out_shape = (M, N) + out_stride = _compute_strides(out_shape) + return _from_metal_buffer(out_buf, out_shape, out_stride, + a.dtype, a.device) + # CPU path (Accelerate BLAS or numpy) a_np = _to_numpy(a) b_np = _to_numpy(b) - # 2D x 2D: use BLAS if possible if a_np.ndim == 2 and b_np.ndim == 2: a_c = np.ascontiguousarray(a_np) b_c = np.ascontiguousarray(b_np) if _can_use_blas(a_c) and _can_use_blas(b_c): return _from_numpy(_blas_gemm(a_c, b_c, a.dtype), a.dtype, a.device) - # Batched: try per-batch BLAS for 3D if a_np.ndim == 3 and b_np.ndim == 3: a_c = np.ascontiguousarray(a_np) b_c = np.ascontiguousarray(b_np) @@ -94,7 +228,6 @@ def matmul(a, b): for i in range(batch): out[i] = _blas_gemm(a_c[i], b_c[i], a.dtype) return _from_numpy(out, a.dtype, a.device) - # Fallback to numpy return _from_numpy(a_np @ b_np, a.dtype, a.device) @@ -107,10 +240,24 @@ def bmm(a, b): def relu(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"relu_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.maximum(_to_numpy(a), 0), a.dtype, a.device) def gelu(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"gelu_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) arr = _to_numpy(a) out = 0.5 * arr * (1.0 + np.vectorize(math.erf)(arr / math.sqrt(2.0))) return _from_numpy(np.ascontiguousarray(out), a.dtype, a.device) @@ -138,10 +285,30 @@ def _check_dim_range(d): for d in dim: _check_dim_range(d) + # GPU path: full-tensor reduction (dim=None) + if dim is None and not keepdim and _can_use_gpu(a) and a.dtype == float32_dtype: + d = _get_dispatcher() + out_buf = _alloc_output_buf(1, a.dtype) + d.dispatch_reduction("sum_partial_f32", "sum_final_f32", + _metal_buf(a), out_buf, a.numel()) + from ..._tensor import _compute_strides + return _from_metal_buffer(out_buf, (), (), a.dtype, a.device) + return _from_numpy(_to_numpy(a).sum(axis=dim, keepdims=keepdim), a.dtype, a.device) def mean_(a, dim=None, keepdim=False): + # GPU path: full-tensor mean (dim=None) + if dim is None and not keepdim and _can_use_gpu(a) and a.dtype == float32_dtype: + d = _get_dispatcher() + # mean = sum / N + sum_buf = _alloc_output_buf(1, a.dtype) + d.dispatch_reduction("sum_partial_f32", "sum_final_f32", + _metal_buf(a), sum_buf, a.numel()) + out_buf = _alloc_output_buf(1, a.dtype) + n = float(a.numel()) + d.dispatch_binary_scalar("div_scalar_f32", sum_buf, n, out_buf, 1) + return _from_metal_buffer(out_buf, (), (), a.dtype, a.device) return _from_numpy(_to_numpy(a).mean(axis=dim, keepdims=keepdim), a.dtype, a.device) @@ -162,6 +329,18 @@ def any_(a, dim=None, keepdim=False): def argmax(a, dim=None, keepdim=False): + # GPU path: full-tensor argmax (dim=None) + if dim is None and _can_use_gpu(a) and a.dtype == float32_dtype: + d = _get_dispatcher() + out_buf = _alloc_output_buf(1, int64_dtype) # uint output + d.dispatch_arg_reduction("argmax_partial_f32", "argmax_final_f32", + _metal_buf(a), out_buf, a.numel()) + # Read uint32 result, convert to int64 + from .runtime import buffer_contents + ptr = buffer_contents(out_buf) + idx_val = struct.unpack("I", (ctypes.c_char * 4).from_address(ptr))[0] + out = np.array(int(idx_val), dtype=np.int64) + return _from_numpy(out, int64_dtype, a.device) arr = _to_numpy(a) if dim is None: out = np.array(np.argmax(arr), dtype=np.int64) @@ -174,6 +353,17 @@ def argmax(a, dim=None, keepdim=False): def argmin(a, dim=None, keepdim=False): + # GPU path: full-tensor argmin (dim=None) + if dim is None and _can_use_gpu(a) and a.dtype == float32_dtype: + d = _get_dispatcher() + out_buf = _alloc_output_buf(1, int64_dtype) + d.dispatch_arg_reduction("argmin_partial_f32", "argmin_final_f32", + _metal_buf(a), out_buf, a.numel()) + from .runtime import buffer_contents + ptr = buffer_contents(out_buf) + idx_val = struct.unpack("I", (ctypes.c_char * 4).from_address(ptr))[0] + out = np.array(int(idx_val), dtype=np.int64) + return _from_numpy(out, int64_dtype, a.device) arr = _to_numpy(a) if dim is None: out = np.array(np.argmin(arr), dtype=np.int64) @@ -717,24 +907,61 @@ def ge(a, b): def add_(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_inplace_binary(f"add_inplace_{sfx}", _metal_buf(a), + _metal_buf(b), numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_inplace_binary_scalar(f"add_inplace_scalar_{sfx}", + _metal_buf(a), scalar, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return a arr = _to_numpy(a) arr += _to_numpy(b) if isinstance(b, Tensor) else b return a def mul_(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_inplace_binary(f"mul_inplace_{sfx}", _metal_buf(a), + _metal_buf(b), numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_inplace_binary_scalar(f"mul_inplace_scalar_{sfx}", + _metal_buf(a), scalar, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return a arr = _to_numpy(a) arr *= _to_numpy(b) if isinstance(b, Tensor) else b return a def relu_(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + d.dispatch_inplace_unary(f"relu_inplace_{sfx}", _metal_buf(a), a.numel()) + return a arr = _to_numpy(a) np.maximum(arr, 0, out=arr) return a def zero_(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + d.dispatch_fill(f"fill_{sfx}", _metal_buf(a), 0.0, a.numel(), + scalar_fmt=_scalar_fmt(a.dtype)) + return a arr = _to_numpy(a) arr.fill(0) return a @@ -803,6 +1030,12 @@ def geometric_(a, p, generator=None): def fill_(a, value): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + d.dispatch_fill(f"fill_{sfx}", _metal_buf(a), float(value), a.numel(), + scalar_fmt=_scalar_fmt(a.dtype)) + return a arr = _to_numpy(a) arr.fill(value) return a @@ -815,6 +1048,12 @@ def clamp_(a, min_val=None, max_val=None): def copy_(a, src): + if _can_use_gpu(a) and _can_use_gpu(src): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = min(a.numel(), src.numel()) + d.dispatch_copy(f"copy_{sfx}", _metal_buf(src), _metal_buf(a), numel) + return a arr = _to_numpy(a) src_arr = _to_numpy(src) np.copyto(arr, src_arr) @@ -874,12 +1113,38 @@ def _ndtr_inv(p): def sub_(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_inplace_binary(f"sub_inplace_{sfx}", _metal_buf(a), + _metal_buf(b), numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_inplace_binary_scalar(f"sub_inplace_scalar_{sfx}", + _metal_buf(a), scalar, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return a arr = _to_numpy(a) arr -= _to_numpy(b) if isinstance(b, Tensor) else b return a def div_(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_inplace_binary(f"div_inplace_{sfx}", _metal_buf(a), + _metal_buf(b), numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_inplace_binary_scalar(f"div_inplace_scalar_{sfx}", + _metal_buf(a), scalar, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return a arr = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b arr /= b_np @@ -894,30 +1159,79 @@ def contiguous(a): def abs(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"abs_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.abs(_to_numpy(a)), a.dtype, a.device) def neg(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"neg_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.negative(_to_numpy(a)), a.dtype, a.device) def exp(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"exp_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.exp(_to_numpy(a)), a.dtype, a.device) def log(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"log_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.log(_to_numpy(a)), a.dtype, a.device) def sqrt(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"sqrt_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.sqrt(_to_numpy(a)), a.dtype, a.device) def sin(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"sin_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.sin(_to_numpy(a)), a.dtype, a.device) def cos(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"cos_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.cos(_to_numpy(a)), a.dtype, a.device) @@ -926,24 +1240,59 @@ def tan(a): def tanh(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"tanh_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.tanh(_to_numpy(a)), a.dtype, a.device) def sigmoid(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"sigmoid_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) arr = _to_numpy(a) out = 1.0 / (1.0 + np.exp(-arr)) return _from_numpy(out, a.dtype, a.device) def floor(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"floor_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.floor(_to_numpy(a)), a.dtype, a.device) def ceil(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"ceil_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.ceil(_to_numpy(a)), a.dtype, a.device) def round(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"round_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.round(_to_numpy(a)), a.dtype, a.device) @@ -958,6 +1307,20 @@ def frac(a): def pow(a, b): + if isinstance(a, Tensor) and _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"pow_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"pow_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) if isinstance(a, Tensor): arr_a = _to_numpy(a) ref = a @@ -972,6 +1335,13 @@ def pow(a, b): def log2(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"log2_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.log2(_to_numpy(a)), a.dtype, a.device) @@ -984,12 +1354,26 @@ def exp2(a): def rsqrt(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"rsqrt_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) arr = _to_numpy(a) out = 1.0 / np.sqrt(arr) return _from_numpy(out, a.dtype, a.device) def sign(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"sign_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.sign(_to_numpy(a)), a.dtype, a.device) @@ -1058,6 +1442,13 @@ def softplus(a): def silu(a): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_unary(f"silu_{sfx}", _metal_buf(a), out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) arr = _to_numpy(a) out = arr / (1.0 + np.exp(-arr)) return _from_numpy(out, a.dtype, a.device) @@ -1283,6 +1674,20 @@ def hypot(a, b): def remainder(a, b): + if isinstance(a, Tensor) and _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"remainder_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"remainder_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) if isinstance(a, Tensor) else a b_np = _to_numpy(b) if isinstance(b, Tensor) else b ref = a if isinstance(a, Tensor) else b @@ -1290,6 +1695,14 @@ def remainder(a, b): def fmod(a, b): + if _can_use_gpu(a) and isinstance(b, Tensor) and _can_use_gpu(b): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_binary(f"fmod_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) return _from_numpy(np.fmod(_to_numpy(a), _to_numpy(b)), a.dtype, a.device) @@ -1535,6 +1948,19 @@ def pad(a, pad_widths, mode='constant', value=0): def softmax(a, dim): + # GPU path: softmax over last dim for 2D contiguous float32 + ndim = len(a.shape) + actual_dim = dim if dim >= 0 else dim + ndim + if (_can_use_gpu(a) and a.dtype == float32_dtype + and actual_dim == ndim - 1 and ndim >= 1): + d = _get_dispatcher() + numel = a.numel() + cols = a.shape[-1] + rows = numel // cols + out_buf = _alloc_output_buf(numel, a.dtype) + d.dispatch_softmax_2d("softmax_f32", _metal_buf(a), out_buf, + rows, cols) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) arr = _to_numpy(a) exp_arr = np.exp(arr - np.max(arr, axis=dim, keepdims=True)) result = exp_arr / np.sum(exp_arr, axis=dim, keepdims=True) @@ -2095,6 +2521,20 @@ def adaptive_avg_pool2d(input, output_size): # --------------------------------------------------------------------------- def sub(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"sub_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"sub_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b return _from_numpy(a_np - b_np, a.dtype, a.device) @@ -2113,12 +2553,40 @@ def reciprocal(a): def maximum(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"maximum_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"maximum_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b return _from_numpy(np.maximum(a_np, b_np), a.dtype, a.device) def minimum(a, b): + if _can_use_gpu(a): + d = _get_dispatcher() + sfx = _kernel_suffix(a.dtype) + numel = a.numel() + out_buf = _alloc_output_buf(numel, a.dtype) + if isinstance(b, Tensor) and _can_use_gpu(b): + d.dispatch_binary(f"minimum_{sfx}", _metal_buf(a), _metal_buf(b), + out_buf, numel) + else: + scalar = float(b) if not isinstance(b, Tensor) else float(_to_numpy(b).ravel()[0]) + d.dispatch_binary_scalar(f"minimum_scalar_{sfx}", _metal_buf(a), + scalar, out_buf, numel, + scalar_fmt=_scalar_fmt(a.dtype)) + return _from_metal_buffer(out_buf, a.shape, a.stride, a.dtype, a.device) a_np = _to_numpy(a) b_np = _to_numpy(b) if isinstance(b, Tensor) else b return _from_numpy(np.minimum(a_np, b_np), a.dtype, a.device) diff --git a/src/mindtorch_v2/_backends/mps/runtime.py b/src/mindtorch_v2/_backends/mps/runtime.py index 3e60d3e16..4dd7bd18a 100644 --- a/src/mindtorch_v2/_backends/mps/runtime.py +++ b/src/mindtorch_v2/_backends/mps/runtime.py @@ -97,6 +97,40 @@ def synchronize(self): cb = self.create_command_buffer() self.commit_and_wait(cb) + def compile_library(self, source_code): + """Compile MSL source string into a MTLLibrary.""" + self._ensure_init() + if _HAS_PYOBJC: + lib, err = self._device.newLibraryWithSource_options_error_( + source_code, None, None) + if err is not None: + raise RuntimeError(f"Metal shader compilation failed: {err}") + return lib + return _compile_library_ctypes(self._device, source_code) + + def make_compute_pipeline(self, function): + """Create MTLComputePipelineState from a MTLFunction.""" + self._ensure_init() + if _HAS_PYOBJC: + pipeline, err = self._device.newComputePipelineStateWithFunction_error_( + function, None) + if err is not None: + raise RuntimeError(f"Failed to create compute pipeline: {err}") + return pipeline + return _make_compute_pipeline_ctypes(self._device, function) + + def get_compute_encoder(self, cmd_buffer): + """Create a MTLComputeCommandEncoder from a command buffer.""" + if _HAS_PYOBJC: + return cmd_buffer.computeCommandEncoder() + return _get_compute_encoder_ctypes(cmd_buffer) + + def buffer_length(self, metal_buffer): + """Return the length in bytes of a Metal buffer.""" + if _HAS_PYOBJC: + return int(metal_buffer.length()) + return _buffer_length_ctypes(metal_buffer) + def device_name(self): """Return the Metal device name.""" self._ensure_init() @@ -228,3 +262,68 @@ def buffer_contents(metal_buffer): if _HAS_PYOBJC: return int(metal_buffer.contents()) return int(_get_buffer_contents_ctypes(metal_buffer)) + + +# --------------------------------------------------------------------------- +# ctypes fallback: compute pipeline helpers +# --------------------------------------------------------------------------- + +def _compile_library_ctypes(device, source_code): + """Compile MSL source via objc_msgSend(device, newLibraryWithSource:options:error:).""" + _load_objc_libs() + # Create NSString from source_code + ns_string_class = _libobjc.objc_getClass(b"NSString") + sel_alloc = _libobjc.sel_registerName(b"alloc") + sel_init = _libobjc.sel_registerName(b"initWithUTF8String:") + ns_str = _libobjc.objc_msgSend(ns_string_class, sel_alloc) + src_bytes = source_code.encode("utf-8") + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_char_p] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + ns_str = _libobjc.objc_msgSend(ns_str, sel_init, src_bytes) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + + # Call newLibraryWithSource:options:error: + sel = _libobjc.sel_registerName(b"newLibraryWithSource:options:error:") + err = ctypes.c_void_p(0) + _libobjc.objc_msgSend.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), + ] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + lib = _libobjc.objc_msgSend(device, sel, ns_str, None, ctypes.byref(err)) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + if lib is None or lib == 0: + raise RuntimeError("Metal shader compilation failed (ctypes path)") + return lib + + +def _make_compute_pipeline_ctypes(device, function): + """Create compute pipeline state from a MTLFunction via ctypes.""" + _load_objc_libs() + sel = _libobjc.sel_registerName(b"newComputePipelineStateWithFunction:error:") + err = ctypes.c_void_p(0) + _libobjc.objc_msgSend.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), + ] + _libobjc.objc_msgSend.restype = ctypes.c_void_p + pipeline = _libobjc.objc_msgSend(device, sel, function, ctypes.byref(err)) + _libobjc.objc_msgSend.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + if pipeline is None or pipeline == 0: + raise RuntimeError("Failed to create compute pipeline (ctypes path)") + return pipeline + + +def _get_compute_encoder_ctypes(cmd_buffer): + """Create compute command encoder via ctypes.""" + return _objc_msg(cmd_buffer, "computeCommandEncoder") + + +def _buffer_length_ctypes(metal_buffer): + """Get Metal buffer length via ctypes.""" + _load_objc_libs() + _libobjc.objc_msgSend.restype = ctypes.c_uint64 + length = _libobjc.objc_msgSend( + metal_buffer, _libobjc.sel_registerName(b"length")) + _libobjc.objc_msgSend.restype = ctypes.c_void_p + return int(length)