github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-25 14:58:46 -05:00
Code Issues Packages Projects Releases Wiki Activity

use dict for UOp parents (#6036)

Browse Source
This commit is contained in:
qazal
2024-08-12 20:06:45 +08:00
committed by GitHub
parent 7d1f118731
commit bb653fa0a5
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tinygrad/codegen/uops.py
View File

@@ -7,7 +7,7 @@ from dataclasses import dataclass
from tinygrad.dtype import ConstType, dtypes, DType
from tinygrad.shape.symbolic import sint, Variable
from tinygrad.ops import UnaryOps, BinaryOps, TernaryOps, exec_alu
from tinygrad.helpers import prod, pretty_print
from tinygrad.helpers import merge_dicts, prod, pretty_print
# the order of these UOps controls the order of the toposort
class UOps(Enum):
@@ -87,9 +87,9 @@ class UOp:
@staticmethod
def store(*src:UOp, **kwargs): return type((src:=(*src, *kwargs.values()))[0])(UOps.STORE, None, src)
@functools.cached_property
def parents(self) -> Set[UOp]: return set.union(set(self.src), *[x.parents for x in self.src])
def parents(self) -> Dict[UOp, None]: return merge_dicts([{x:None for x in self.src}]+[x.parents for x in self.src])
@property # parents with self
def sparents(self) -> Set[UOp]: return set([self]).union(self.parents)
def sparents(self) -> Dict[UOp, None]: return {**self.parents, self:None}
def vars(self) -> Set[UOp]: return set([x for x in self.sparents if x.op is UOps.DEFINE_VAR])
def const_factor(self) -> int:
"""largest known int that divides self"""