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

Revert "Green Uop unary check (#2792)" (#2799)

Browse Source 
This reverts commit d958777aed.
This commit is contained in:
chenyu
2023-12-16 12:49:28 -05:00
committed by GitHub
parent d958777aed
commit 0bb5d8f956
1 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/linearizer.py
+4 -8
View File
@@ -6,7 +6,7 @@ from enum import Enum, auto
from dataclasses import dataclass
from tinygrad.helpers import colored, ImageDType, DEBUG, dtypes, DType, prod, PtrDType, getenv, all_same, to_function_name, flatten
from tinygrad.ops import LazyOp, UnaryOps, BinaryOps, TernaryOps, ReduceOps, ConstBuffer, MemBuffer, BufferOps, vars_from_ast, get_lazyop_info
from tinygrad.ops import LazyOp, UnaryOps, BinaryOps, TernaryOps, ReduceOps, ConstBuffer, MemBuffer, BufferOps, vars_from_ast
from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.symbolic import Variable, NumNode, VariableOrNum, Node, SumNode, MulNode, DivNode, ModNode, LtNode, AndNode
from tinygrad.codegen.kernel import LocalBuffer, Kernel
@@ -64,10 +64,6 @@ class Linearizer(Kernel):
AndNode: lambda self,ops,ctx:
functools.reduce(lambda a,b: ctx.uop_alu_idx(a, b, ops, ctx, BinaryOps.MUL, dtype=dtypes.bool), self.nodes[1:], self.nodes[0].render(ops,ctx)) }
def get_uop_dtype(self, dtype: DType, amt=1):
if isinstance(dtype, ImageDType): dtype = dtypes.float
return dtype if amt == 1 else dtype.vec(amt)
def global_load(self, i:int, idxs:Sequence[Node], acc=None, barrier:Optional[UOp]=None) -> List[UOp]:
buf = self.bufs[i]
const = buf.val if isinstance(buf, ConstBuffer) else acc
@@ -87,7 +83,8 @@ class Linearizer(Kernel):
(g_idx, g_valid), amt, dim = self.sts[i].expr_idxs(fake_idxs), 1, None
else:
g_idx, g_valid = self.sts[i].expr_idxs(fake_idxs)
localtype = self.get_uop_dtype(buf.dtype, amt)
localtype = buf.dtype if amt == 1 else buf.dtype.vec(amt)
if isinstance(buf.dtype, ImageDType): localtype = dtypes.float if amt == 1 else dtypes.float.vec(amt)
e_idxs, e_valids = g_idx.expand(expand_vars), g_valid.expand(expand_vars)
@@ -479,7 +476,6 @@ class Linearizer(Kernel):
key = (uop, dtype, vin, arg)
if uop == UOps.ALU:
if arg in UnaryOps: assert dtype == vin[0].dtype, f"{arg} dtype mismatch {dtype=} != {vin[0].dtype=}"
if arg == BinaryOps.CMPLT: assert dtype == dtypes.bool, f"{arg} output dtype mismatch {dtype=} != {dtypes.bool}"
if arg == TernaryOps.WHERE: assert vin[0].dtype == dtypes.bool, f"{arg} selector dtype mismatch {vin[0].dtype=} != {dtypes.bool}"
@@ -540,5 +536,5 @@ class Linearizer(Kernel):
if input_acc[off] != acc[off]:
acc[off] = self.uop(UOps.PHI, input_acc[off].dtype, (input_acc[off], acc[off]) + tuple(loop_ctx))
else:
ret = [self.uop(UOps.ALU, dtype=dtypes.bool if x.op == BinaryOps.CMPLT else self.get_uop_dtype(get_lazyop_info(x).dtype) if x.op in UnaryOps else None, vin=val, arg=x.op) for val in zip(*values)] # noqa:E501 TODO use get_lazyop_info in all uops
ret = [self.uop(UOps.ALU, dtype=dtypes.bool if x.op == BinaryOps.CMPLT else None, vin=val, arg=x.op) for val in zip(*values)]
return ret