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

clean up is_address

Browse Source
This commit is contained in:
geohot
2026-07-10 20:11:23 -07:00
parent 5e869b4a08
commit 6e3f7e6a84
2 changed files with 8 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/renderer/isa/x86.py
+2 -8
View File
@@ -133,14 +133,8 @@ class X86GroupOp:
All = set(X86Ops)
def is_address(x:UOp) -> bool:
if x.op is Ops.PARAM: return x.arg.addrspace is AddrSpace.GLOBAL
if x.op is Ops.BUFFER: return True
if x.op is Ops.INS:
if x.arg == X86Ops.LEA or (x.arg == X86Ops.DEFINE and x.tag == (RSP,)): return True
return x.dtype is dtypes.uint64 and x.arg in {X86Ops.MOV, X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE} and \
(x.shape == () or any(is_address(s) for s in x.src[:2]))
if x.op in {Ops.INDEX, Ops.SHRINK, Ops.AFTER, Ops.NOOP} and x.src: return is_address(x.src[0])
return x.op is Ops.WHERE and is_address(x.src[1])
# addresses are GLOBAL/LOCAL/REG addrspace values, or uint64 INS values (LEA/MOV/CMOV/DEFINE-RSP from isel)
return x.addrspace not in (None, AddrSpace.ALU) or (x.op is Ops.INS and x.dtype is dtypes.uint64)
# ***** X86 legalization *****
tinygrad/uop/ops.py
+6 -2
View File
@@ -831,10 +831,14 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if self.op is Ops.BUFFER: return self.arg.addrspace
if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU
if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT}:
return self.src[0].addrspace
if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.NOOP}:
return self.src[0].addrspace if self.src else None
if self.op in GroupOp.Movement: return self.src[0].addrspace
if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise:
# WHERE's condition (src[0]) is never an address, either branch being a pointer makes the result a pointer
if self.op is Ops.WHERE:
ad = [x.addrspace for x in self.src[1:] if x.addrspace not in (None, AddrSpace.ALU)]
return ad[0] if ad else None
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]