From cb6cfcc8f859d90556dc103e6a73cd9049ab35b2 Mon Sep 17 00:00:00 2001 From: taher <8665427+nullhook@users.noreply.github.com> Date: Wed, 15 Nov 2023 11:37:28 -0800 Subject: [PATCH] add icb support check for metal device (#2313) --- tinygrad/runtime/ops_metal.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index 18c007feb3..bca9972ba2 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -3,7 +3,7 @@ import os, subprocess, pathlib, ctypes, tempfile import Metal, Cocoa, libdispatch from typing import List, Any, Tuple, Dict, Union, Set, cast from tinygrad.codegen.kernel import LinearizerOptions -from tinygrad.helpers import prod, getenv, DEBUG, DType, dtypes, diskcache, dedup, CI +from tinygrad.helpers import prod, getenv, DEBUG, DType, dtypes, diskcache, dedup from tinygrad.ops import Compiled, BatchExecutor, JitItem, CompiledASTRunner from tinygrad.renderer.metal import MetalRenderer from tinygrad.runtime.lib import RawBufferMapped, RawBuffer, LRUAllocator @@ -18,6 +18,7 @@ class _METAL: def __init__(self): self.mtl_buffers_in_flight: List[Any] = [] self.device = Metal.MTLCreateSystemDefaultDevice() + self.supports_icb = (self.device.supportsFamily_(Metal.MTLGPUFamilyMac2) or self.device.supportsFamily_(Metal.MTLGPUFamilyApple3) or self.device.supportsFamily_(Metal.MTLGPUFamilyCommon2)) and self.device.argumentBuffersSupport() is Metal.MTLArgumentBuffersTier2 self.mtl_queue = self.device.newCommandQueueWithMaxCommandBufferCount_(1024) self.allocator = MetalAllocator(self.device.dedicatedMemorySize() or self.device.sharedMemorySize()) # TODO: is there a better way to do this? @@ -151,4 +152,4 @@ class MetalBatchExecutor(BatchExecutor): super().update_stats(var_vals, et) return et -MetalBuffer = Compiled(RawMetalBuffer, LinearizerOptions(device="METAL"), MetalRenderer, compile_metal, MetalProgram, METAL.synchronize, batch_executor=MetalBatchExecutor if not CI else BatchExecutor) +MetalBuffer = Compiled(RawMetalBuffer, LinearizerOptions(device="METAL"), MetalRenderer, compile_metal, MetalProgram, METAL.synchronize, batch_executor=MetalBatchExecutor if METAL.supports_icb else BatchExecutor)