mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-11 23:46:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix tests

Browse Source
This commit is contained in:
George Hotz
2026-06-01 16:22:16 -07:00
parent f532e1b2d0
commit 08fe658f74
3 changed files with 2 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tinygrad/renderer/llvmir.py
View File

@@ -176,7 +176,6 @@ class LLVMRenderer(Renderer):
args.append((r[u], u.dtype))
elif u.op in (Ops.DEFINE_LOCAL, Ops.DEFINE_REG):
r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}"
#assert isinstance(u.dtype, PtrDType)
size = u.max_numel()
if u.op is Ops.DEFINE_REG:
kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}]")

4
tinygrad/uop/ops.py
View File

@@ -769,7 +769,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.GEP}:
return self.src[0].addrspace
if self.op in GroupOp.Movement: return self.src[0].addrspace
if self.op is Ops.STACK or self.op in GroupOp.Elementwise:
if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise:
ad = [x.addrspace for x in self.src if x.addrspace is not None]
if not len(ad) or not all_same(ad): return None
return ad[0]
@@ -1116,7 +1116,7 @@ class ProgramInfo:
if u.op is Ops.DEFINE_VAR: _vars.append(u)
if u.op is Ops.PARAM: _globals.append(u.arg.slot)
if u.op in (Ops.STORE, Ops.LOAD):
if (idx:=u.src[0]).op is Ops.INDEX or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX):
if (idx:=u.src[0]).op in (Ops.INDEX, Ops.SHRINK):
if (buf:=idx.src[0]).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg.slot)
if u.op is Ops.SPECIAL:
if u.arg[0] == 'i': local_size = None

2
tinygrad/uop/spec.py
View File

@@ -81,8 +81,6 @@ spec_shared = PatternMatcher([
(UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.UNROLL, Ops.INS))), lambda: True),
# TOOD: these should be buffer with different addrspace
#(UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
#(UPat(Ops.DEFINE_REG, src=()), lambda: True),
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG)), lambda: True),
# AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, or another AFTER