tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 14:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

CPUGraph support for clang (#10014)

Browse Source 
Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
quortus
2025-04-24 07:52:35 -04:00
committed by GitHub
co-authored by George Hotz
parent c3ad7b2a84
commit 9e49721c47
3 changed files with 17 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_graph.py
+11 -7
View File
@@ -1,5 +1,5 @@
import numpy as np
import unittest, ctypes
import functools, unittest, ctypes
from tinygrad.device import Device, Buffer
from tinygrad.tensor import Tensor, _to_np_dtype
@@ -107,8 +107,12 @@ class TestGraph(unittest.TestCase):
helper_test_graphs(Device[d0].graph, graphs)
def skip_if_not_multigraph(self):
graph = g.func if isinstance(g:=Device[Device.DEFAULT].graph, functools.partial) else g
if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
def test_order_copy_writed(self):
if not issubclass(Device[Device.DEFAULT].graph, MultiGraphRunner): self.skipTest("graph does not supported (not MultiGraphRunner)")
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(4)]
@@ -120,7 +124,7 @@ class TestGraph(unittest.TestCase):
helper_test_graphs(Device[d0].graph, graphs)
def test_order_copy_then_read(self):
if not issubclass(Device[Device.DEFAULT].graph, MultiGraphRunner): self.skipTest("graph does not supported (not MultiGraphRunner)")
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(4)]
@@ -151,7 +155,7 @@ class TestGraph(unittest.TestCase):
helper_test_graphs(Device[d0].graph, graphs)
def test_copies_2_devs(self):
if not issubclass(Device[Device.DEFAULT].graph, MultiGraphRunner): self.skipTest("graph does not supported (not MultiGraphRunner)")
self.skip_if_not_multigraph()
d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
@@ -164,7 +168,7 @@ class TestGraph(unittest.TestCase):
helper_test_graphs(Device[d0].graph, graphs)
def test_copies_after_graph_global(self):
if not issubclass(Device[Device.DEFAULT].graph, MultiGraphRunner): self.skipTest("graph does not supported (not MultiGraphRunner)")
self.skip_if_not_multigraph()
d0, d1, d2, d3 = Device.DEFAULT, f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3"
b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
@@ -212,7 +216,7 @@ class TestGraph(unittest.TestCase):
helper_test_graphs(Device[d0].graph, graphs)
def test_graph_after_copies_devs(self):
if not issubclass(Device[Device.DEFAULT].graph, MultiGraphRunner): self.skipTest("graph does not supported (not MultiGraphRunner)")
self.skip_if_not_multigraph()
d0, d1, d2, d3 = Device.DEFAULT, f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3"
b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
@@ -240,7 +244,7 @@ class TestGraph(unittest.TestCase):
helper_test_graphs(Device[d0].graph, graphs)
def test_graph_offset_bufs(self):
if not issubclass(Device[Device.DEFAULT].graph, MultiGraphRunner): self.skipTest("graph does not supported (not MultiGraphRunner)")
self.skip_if_not_multigraph()
d0 = Device.DEFAULT
if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")
tinygrad/runtime/graph/cpu.py
+2 -1
View File
@@ -31,7 +31,8 @@ class CPUGraph(GraphRunner):
batched.append("}")
prep = [device.renderer._render(cast(CompiledRunner, ji.prg).p.uops) for i,ji in enumerate(jit_cache)]
funcs = dedup(device.renderer._render_body(prep[i][0], *prep[i][1:], cast(CompiledRunner, ji.prg).p.uops) for i,ji in enumerate(jit_cache))
funcs = dedup(device.renderer._render_body(prep[i][0], *prep[i][1:], cast(CompiledRunner, ji.prg).p.uops, ["static"])
for i,ji in enumerate(jit_cache))
defines = dedup(itertools.chain.from_iterable(device.renderer._render_defines(cast(CompiledRunner, ji.prg).p.uops) for ji in jit_cache))
entry = device.renderer._render_entry("batched", targs)
tinygrad/runtime/ops_cpu.py
+4 -2
View File
@@ -1,4 +1,4 @@
import platform, subprocess, sys
import functools, platform, subprocess, sys
from tinygrad.helpers import capstone_flatdump, getenv
from tinygrad.device import Compiled, Compiler, MallocAllocator, CPUProgram
from tinygrad.runtime.support.elf import jit_loader
@@ -19,6 +19,8 @@ class ClangJITCompiler(Compiler):
def disassemble(self, lib:bytes): return capstone_flatdump(lib)
class ClangDevice(Compiled):
def __init__(self, device:str): super().__init__(device, MallocAllocator, ClangRenderer(), ClangJITCompiler(), CPUProgram)
def __init__(self, device:str):
from tinygrad.runtime.graph.cpu import CPUGraph
super().__init__(device, MallocAllocator, ClangRenderer(), ClangJITCompiler(), CPUProgram, functools.partial(CPUGraph, self))
CPUDevice = ClangDevice