diff --git a/docs/env_vars.md b/docs/env_vars.md index 790eed06..4cbc1fbe 100644 --- a/docs/env_vars.md +++ b/docs/env_vars.md @@ -38,6 +38,9 @@ ENABLE_METHOD_CACHE | [1] | enable method cache (this is the default) EARLY_STOPPING | [# > 0] | stop after this many kernels DISALLOW_ASSIGN | [1] | disallow assignment of tensors NATIVE_EXPLOG | [1] | enable using native exp and log +CL_EXCLUDE | [name0,name1] | comma-separated list of device names to exclude when using OpenCL GPU backend (like `CL_EXCLUDE=gfx1036`) +CL_PLATFORM | [# >= 0] | index of the OpenCL [platform](https://documen.tician.de/pyopencl/runtime_platform.html#pyopencl.Platform) to run on. Defaults to 0. +RDNA | [1] | enable the specialized [RDNA 3](https://en.wikipedia.org/wiki/RDNA_3) assembler for AMD 7000-series GPUs. If not set, defaults to generic OpenCL codegen backend. ## File Specific Variables diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py index 46f11095..6f8fb52d 100644 --- a/tinygrad/runtime/ops_gpu.py +++ b/tinygrad/runtime/ops_gpu.py @@ -9,21 +9,21 @@ from tinygrad.runtime.lib import RawBufferCopyInOut from tinygrad.codegen.cstyle import CStyleCodegen, CStyleLanguage OSX_TIMING_RATIO = (125/3) if OSX else 1.0 # see test/external_osx_profiling.py to determine this ratio. it's in like GPU clocks or something -FLOAT16 = getenv("FLOAT16", 0) # TODO: if you fork and exit the child process after creating anything with cl on AMD, it hangs on e.wait() ROCM_LLVM_PATH = pathlib.Path("/opt/rocm/llvm/bin") #ROCM_LLVM_PATH = pathlib.Path(__file__).parent.parent.parent.parent / "extra/rocm/build/llvm-project/bin" if DEBUG >= 5: - from extra.helpers import enable_early_exec - early_exec = enable_early_exec() + early_exec = fromimport("extra.helpers", "enable_early_exec")() class _CL: def __init__(self): platforms: List[List[cl.Device]] = [y for y in ([x.get_devices(device_type=cl.device_type.GPU) for x in cl.get_platforms()] + [x.get_devices(device_type=cl.device_type.CPU) for x in cl.get_platforms()]) if len(y)] - if DEBUG >= 1: print(f"using {platforms[getenv('CL_PLATFORM', 0)]}") - self.cl_ctx: cl.Context = cl.Context(devices=[x for x in platforms[getenv('CL_PLATFORM', 0)] if x.name not in getenv('CL_EXCLUDE', "").split(",")]) + devices: List[cl.Device] = [x for x in platforms[getenv('CL_PLATFORM', 0)] if x.name not in getenv('CL_EXCLUDE', '').split(',')] + if DEBUG >= 1: print(f"using devices: {[d.hashable_model_and_version_identifier for d in devices]}") + self.cl_ctx: cl.Context = cl.Context(devices=devices) self.cl_queue: List[cl.CommandQueue] = [cl.CommandQueue(self.cl_ctx, device=device, properties=cl.command_queue_properties.PROFILING_ENABLE) for device in self.cl_ctx.devices] + def synchronize(self): for q in self.cl_queue: q.finish() CL = _CL() @@ -40,10 +40,12 @@ class CLBuffer(RawBufferCopyInOut): buf = cl.Buffer(CL.cl_ctx, cl.mem_flags.READ_WRITE, size * dtype.itemsize) setattr(buf, 'device', int(device)) # device is tracked on the underlying buffer super().__init__(size, dtype, buf) - def _copyin(self, x:np.ndarray): + + def _copyin(self, x: np.ndarray): assert not self.dtype.name.startswith("image"), f"can't copyin images {self.dtype}" cl.enqueue_copy(CL.cl_queue[self._buf.device], self._buf, np.require(x, requirements='C'), is_blocking=False) - def _copyout(self, x:np.ndarray): + + def _copyout(self, x: np.ndarray): assert not self.dtype.name.startswith("image"), f"can't copyout images {self.dtype}" cl.enqueue_copy(CL.cl_queue[self._buf.device], x, self._buf, is_blocking=True) @@ -58,8 +60,7 @@ class CLProgram: self.clprg = self._clprg.__getattr__(name) if DEBUG >= 5 and not OSX: if 'Adreno' in CL.cl_ctx.devices[0].name: - from disassemblers.adreno import disasm - disasm(self.binary()) + fromimport('disassemblers.adreno', 'disasm')(self.binary()) elif CL.cl_ctx.devices[0].name.startswith('gfx'): asm = early_exec(([ROCM_LLVM_PATH / "llvm-objdump", '-d', '-'], self.binary())) print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x])) @@ -91,7 +92,5 @@ class CLCodegen(CStyleCodegen): half_prekernel = "#pragma OPENCL EXTENSION cl_khr_fp16 : enable", barrier = "barrier(CLK_LOCAL_MEM_FENCE);", float4 = "(float4)", gid = [f'get_global_id({i})' for i in range(3)], lid = [f'get_local_id({i})' for i in range(3)], uses_vload=True) - supports_float4_alu = True - supports_float4 = True GPUBuffer = Compiled(CLBuffer, fromimport("tinygrad.codegen.assembly_rdna", "RDNACodegen") if getenv("RDNA") else CLCodegen, CLProgram, CL.synchronize)