types (#720)

* types

* cleanups

* don't use None, use LocalBuffer

* eh

tinygrad/codegen/cstyle.py

@@ -1,6 +1,6 @@
-from typing import Final, Dict, Callable, ClassVar, List, Optional, NamedTuple, DefaultDict, Tuple, Set, Any
+from typing import Final, Dict, Callable, ClassVar, List, Optional, NamedTuple, DefaultDict, Tuple, Set
 import math, collections
-from tinygrad.codegen.linearizer import Linearizer, UOps
+from tinygrad.codegen.linearizer import Linearizer, UOps, UOp, LocalBuffer
 from tinygrad.ops import ASTRunner, Op, UnaryOps, BinaryOps, FusedOps
 from tinygrad.helpers import getenv, all_same, partition, ImageDType, DEBUG, dtypes
 from tinygrad.runtime.lib import RawConst
@@ -56,7 +56,7 @@ code_for_op: Final[Dict[Op, Callable]] = {
   BinaryOps.CMPEQ: lambda a,b: f"({a}=={b})", FusedOps.MULACC: lambda a,b,c: f"(({b}*{c})+{a})"
 }
 
-def uops_to_cstyle(uops:List[Tuple[UOps, Optional[str], Any]], bufs:List[LazyBuffer], bufnames:List[str], lang:CStyleLanguage) -> Tuple[str, List[int], List[int]]:
+def uops_to_cstyle(uops:List[UOp], bufs:List[LazyBuffer], bufnames:List[str], lang:CStyleLanguage) -> Tuple[str, List[int], List[int]]:
   def group_float4(grp:List[str]) -> str:
     if all(g.endswith(e) for g,e in zip(grp, [".x", ".y", ".z", ".w"])) and all_same([g.split(".")[0] for g in grp]): return grp[0].split(".")[0]
     else: return f"{lang.float4}({','.join(g for g in grp)})"
@@ -70,7 +70,7 @@ def uops_to_cstyle(uops:List[Tuple[UOps, Optional[str], Any]], bufs:List[LazyBuf
   depth = 0
   def kk(s): kernel.append("  "*depth+s)
 
-  for uop,newvar,args in uops:
+  for uop,newvar,vin,args in uops:
     if uop == UOps.LOOP:
       root = None
       for i,var in enumerate(args[0]):
@@ -115,58 +115,58 @@ def uops_to_cstyle(uops:List[Tuple[UOps, Optional[str], Any]], bufs:List[LazyBuf
         depth -= 1
         kk("}"*len(args[0]) + f" /* {args[1]} */")
     if uop == UOps.CONST:
-      if args[0] == -math.inf:
+      if args == -math.inf:
         kk(f"float {newvar} = -INFINITY;")
       else:
-        kk(f"float {newvar} = {args[0]}f;")
+        kk(f"float {newvar} = {args}f;")
     if uop == UOps.ALU:
-      if newvar is None:
-        kk(f"{args[2]} = {code_for_op[args[0]](*args[1])};")
+      if newvar in vin:
+        kk(f"{newvar} = {code_for_op[args](*vin)};")
       else:
-        kk(f"float {newvar} = {code_for_op[args[0]](*args[1])};")
+        kk(f"float {newvar} = {code_for_op[args](*vin)};")
     # TODO: refactor the next 14 lines
     if uop == UOps.LOAD:
       # TODO: merge with CONST?
-      if bufs[args[0]] is not None and isinstance(bufs[args[0]].realized, RawConst):
+      if bufs[args.i] is not None and isinstance(bufs[args.i].realized, RawConst):
         # nan? inf?
-        val = f"{bufs[args[0]].realized._buf}f"
+        val = f"{bufs[args.i].realized._buf}f"
       else:
-        if lang.uses_vload and bufs[args[0]] is not None and bufs[args[0]].dtype == dtypes.float16:
-          val = f"vload_half({args[1].render(render_cl)}, {bufnames[args[0]]})"
+        if lang.uses_vload and bufs[args.i].dtype == dtypes.float16:
+          val = f"vload_half({args.idx.render(render_cl)}, {bufnames[args.i]})"
         else:
-          val = f"{bufnames[args[0]]}[{args[1].render(render_cl)}]"
+          val = f"{bufnames[args.i]}[{args.idx.render(render_cl)}]"
       # NOTE: if min and max are both 0, it should be a CONST in the Linearizer
-      if args[2].min == 1: kk(f"float {newvar} = {val};")
-      else: kk(f"float {newvar} = ({args[2].render(render_cl)}) ? ({val}) : 0.0f;")
+      if args.valid.min == 1: kk(f"float {newvar} = {val};")
+      else: kk(f"float {newvar} = ({args.valid.render(render_cl)}) ? ({val}) : 0.0f;")
     if uop == UOps.LOAD4:
-      if bufs[args[0]] is not None and isinstance(bufs[args[0]].dtype, ImageDType):
+      if isinstance(bufs[args.i].dtype, ImageDType):
         prekernel.add("const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n")
-        idx, idy = to_image_idx(bufs[args[0]].dtype.shape, args[1], args[2])
-        val = f"read_imagef({bufnames[args[0]]}, smp, (int2)({idx.render(render_cl)}, {idy.render(render_cl)}))"
+        idx, idy = to_image_idx(bufs[args.i].dtype.shape, args.idx, args.valid)
+        val = f"read_imagef({bufnames[args.i]}, smp, (int2)({idx.render(render_cl)}, {idy.render(render_cl)}))"
       else:
-        val = f"(({lang.buffer_prefix if bufs[args[0]] is not None else lang.smem_prefix}float4*){bufnames[args[0]]})[{(args[1]//4).render(render_cl)}]"
+        val = f"(({lang.smem_prefix if isinstance(bufs[args.i], LocalBuffer) else lang.buffer_prefix}float4*){bufnames[args.i]})[{(args.idx//4).render(render_cl)}]"
       # NOTE: if min and max are both 0, it should be a CONST in the Linearizer
       if args[2].min == 1: kk(f"float4 {newvar} = {val};")
-      else: kk(f"float4 {newvar} = ({args[2].render(render_cl)}) ? ({val}) : {group_float4(['0.0f']*4)};")
+      else: kk(f"float4 {newvar} = ({args.valid.render(render_cl)}) ? ({val}) : {group_float4(['0.0f']*4)};")
     if uop == UOps.STORE:
-      assert args[2].min == 1, "store must be valid"
-      if lang.uses_vload and bufs[args[0]] is not None and bufs[args[0]].dtype == dtypes.float16:
-        kk(f"vstore_half({args[3]}, {args[1].render(render_cl)}, {bufnames[args[0]]});")
+      assert args.valid.min == 1, "store must be valid"
+      if lang.uses_vload and bufs[args.i].dtype == dtypes.float16:
+        kk(f"vstore_half({vin[0]}, {args.idx.render(render_cl)}, {bufnames[args.i]});")
       else:
-        kk(f"{bufnames[args[0]]}[{args[1].render(render_cl)}] = {args[3]};")
+        kk(f"{bufnames[args.i]}[{args.idx.render(render_cl)}] = {vin[0]};")
     if uop == UOps.STORE4:
-      assert args[2].min == 1, "store must be valid"
-      if bufs[args[0]] is not None and isinstance(bufs[args[0]].dtype, ImageDType):
-        idx, idy = to_image_idx(bufs[args[0]].dtype.shape, args[1], args[2])
-        kk(f"write_imagef({bufnames[args[0]]}, (int2)({idx.render(render_cl)}, {idy.render(render_cl)}), {group_float4(args[3])});")
+      assert args.valid.min == 1, "store must be valid"
+      if isinstance(bufs[args[0]].dtype, ImageDType):
+        idx, idy = to_image_idx(bufs[args.i].dtype.shape, args[1], args[2])
+        kk(f"write_imagef({bufnames[args.i]}, (int2)({idx.render(render_cl)}, {idy.render(render_cl)}), {group_float4(vin)});")
       else:
-        kk(f"(({lang.buffer_prefix if bufs[args[0]] is not None else lang.smem_prefix}float4*){bufnames[args[0]]})[{(args[1]//4).render(render_cl)}] = {group_float4(args[3])};")
+        kk(f"(({lang.smem_prefix if isinstance(bufs[args.i], LocalBuffer) else lang.buffer_prefix}float4*){bufnames[args.i]})[{(args.idx//4).render(render_cl)}] = {group_float4(vin)};")
     if uop == UOps.DEFINE_LOCAL:
       kk(lang.smem_prefix + f"float {args[0]}[{args[1]}];")
 
   buftypes = [(i,f"{'read_only' if i > 0 else 'write_only'} image2d_t" if x.dtype.name.startswith('image') else
                ("const " if i > 0 else "")+lang.buffer_prefix+x.dtype.name+"*"+lang.buffer_suffix) for i,x in enumerate(bufs)
-               if x is not None and not isinstance(x.realized, RawConst)]
+               if not isinstance(x, LocalBuffer) and not isinstance(x.realized, RawConst)]
   prg = ''.join([f"{lang.kernel_prefix} void KERNEL_NAME_PLACEHOLDER(",] +
     [', '.join([f'{t} {bufnames[i]}' for i,t in buftypes] + lang.extra_args)] +
     [") {\n"] + list(prekernel) + ['\n'.join(kernel), "\n}"])

tinygrad/codegen/linearizer.py

@@ -1,9 +1,9 @@
-from typing import List, Tuple, Any, Optional, cast, Dict, DefaultDict
+from typing import List, Tuple, Any, Optional, cast, Dict, DefaultDict, NamedTuple
 import itertools, math
 from collections import defaultdict
 from enum import Enum, auto
 
-from tinygrad.helpers import dedup, colored, all_same, ImageDType, DEBUG, prod, dtypes, mnum
+from tinygrad.helpers import dedup, colored, all_same, ImageDType, DEBUG, prod, dtypes, mnum, DType
 from tinygrad.ops import LazyOp, get_lazyops, get_buffers, FlopCounter, get_lazyop_info, map_buffers, UnaryOps
 from tinygrad.lazy import LazyBuffer
 from tinygrad.ops import MovementOps, ReduceOps, BinaryOps, FusedOps
@@ -12,6 +12,22 @@ from tinygrad.shape.symbolic import Variable, SumNode, ModNode
 
 class UOps(Enum): LOOP = auto(); DEFINE_LOCAL = auto(); LOAD = auto(); ALU = auto(); CONST = auto(); ENDLOOP = auto(); STORE = auto(); LOAD4 = auto(); STORE4 = auto() # noqa: E702
 
+class LocalBuffer(NamedTuple):
+  dtype: DType = dtypes.float32
+  realized: None = None
+
+class MemOp(NamedTuple):
+  i: int
+  idx: Variable
+  valid: Variable
+
+class UOp(NamedTuple):
+  uop: UOps
+  out: Optional[str]
+  vin: List[str]
+  arg: Any
+  def __repr__(self): return f"{str(self.uop):20s}: {self.out if self.out is not None else '':10s} {str(self.vin):32s} {self.arg}"
+
 def get_first_reduce(shapes):
   for i in range(len(shapes[0])):
     if not all_same([x[i] for x in shapes]): return i
@@ -89,13 +105,21 @@ class Linearizer:
     self.registers = [Register(f"data{i}") for i in range(len(self.bufs))]
     self.group_for_reduce: List[int] = []
 
+  def can_merge_float4(self, i:int, idxs:List[Variable], offset:int) -> bool:
+    if offset%4 != 0: return False
+    float4_index = Variable("FLOAT4_INDEX", 0, 3)
+    idxy_test, valid_test = self.sts[i].expr_idxs(float4_index+offset, idxs)
+    if DEBUG >= 4: print(f"attempting to fuse buf {i} :", check_no_mul(idxy_test, float4_index), idxy_test//4, valid_test//4)
+    # float4_index must not be in after divide or in valid. NOTE: this forces it to always be aligned too, maybe not required?
+    return check_no_mul(idxy_test, float4_index) and "FLOAT4_INDEX" not in (idxy_test//4).render() and "FLOAT4_INDEX" not in (valid_test//4).render()
+
   def linearize(self):
     # uops
-    self.uops: List[Tuple[UOps, Optional[str], Any]] = []
+    self.uops: List[UOp] = []
 
     # add a local buffer for multistage reduce
     if len(self.group_for_reduce):
-      self.bufs.append(None)
+      self.bufs.append(LocalBuffer())
       # TODO: the strides of this can be controlled
       st = ShapeTracker(tuple([1] * self.first_reduce + self.group_for_reduce + [1] * (self.shape_len - len(self.group_for_reduce) - self.first_reduce) + [x[0] for x in self.registers[0].axis]))
       buftoken = Register("temp")
@@ -105,35 +129,27 @@ class Linearizer:
         st.views[-1] = View(st.shape[0:-1], st.views[-1].strides[0:-1], st.views[-1].offset)
       self.sts.append(st)
       self.registers.append(buftoken)
-      self.uop(UOps.DEFINE_LOCAL, (self.registers[-1].name, self.sts[-1].size()*self.registers[-1].size()))
+      self.uop(UOps.DEFINE_LOCAL, None, [], (self.registers[-1].name, self.sts[-1].size()*self.registers[-1].size()))
 
-    # TODO: add upcasting to float4 here
-    def global_buf(i, idxs, store=None):
-      should_upcast = self.supports_float4 and self.registers[i].can_float4() and (self.bufs[i] is None or self.bufs[i].dtype != dtypes.float16 or isinstance(self.bufs[i].dtype, ImageDType))
+    # print
+    if DEBUG >= 3: self.printbufs()
+
+    def global_buf(i, idxs:List[Variable], store=None):
+      should_upcast = self.supports_float4 and self.registers[i].can_float4() and self.bufs[i].dtype != dtypes.float16
       cache: Dict[int, str] = {}
       store_offset: Dict[int, int] = {y:x for x,y in enumerate(self.registers[i].offsets())}  # NOTE: for stores, these should be unique
       def op(offset):
         if offset in cache: return cache[offset]
-        will_merge = False
-        if should_upcast and offset%4 == 0:
-          float4_index = Variable("FLOAT4_INDEX", 0, 3)
-          idxy_test, valid_test = self.sts[i].expr_idxs(float4_index+offset, idxs)
-          if DEBUG >= 4: print(f"attempting to fuse buf {i} :", check_no_mul(idxy_test, float4_index), idxy_test//4, valid_test//4)
-          # float4_index must not be in after divide or in valid. NOTE: this forces it to always be aligned too, maybe not required?
-          will_merge = check_no_mul(idxy_test, float4_index) and "FLOAT4_INDEX" not in (idxy_test//4).render() and "FLOAT4_INDEX" not in (valid_test//4).render()
+        will_merge = should_upcast and self.can_merge_float4(i, idxs, offset)
         if store is not None:
           if offset in store_offset:
-            if will_merge:
-              offsets = []
-              for j in range(0, 4):
-                offsets.append(store[store_offset[offset+j]])
-                del store_offset[offset+j]
-              self.uop(UOps.STORE4, (i, *self.sts[i].expr_idxs(offset, idxs), offsets))
-            else:
-              self.uop(UOps.STORE, (i, *self.sts[i].expr_idxs(offset, idxs), store[store_offset[offset]]))
-              del store_offset[offset]
+            offsets = []
+            for j in range(0, 4 if will_merge else 1):
+              offsets.append(store[store_offset[offset+j]])
+              del store_offset[offset+j]
+            self.uop(UOps.STORE4 if will_merge else UOps.STORE, None, offsets, MemOp(i, *self.sts[i].expr_idxs(offset, idxs)))
         else:
-          reg = self.uop(UOps.LOAD4 if will_merge else UOps.LOAD, (i, *self.sts[i].expr_idxs(offset, idxs)), self.registers[i].name+"_"+mnum(offset))
+          reg = self.uop(UOps.LOAD4 if will_merge else UOps.LOAD, self.registers[i].name+"_"+mnum(offset), [], MemOp(i, *self.sts[i].expr_idxs(offset, idxs)))
           if will_merge:
             for j in range(0, 4): cache[offset+j] = reg+"."+"xyzw"[j]
           else:
@@ -153,13 +169,13 @@ class Linearizer:
 
     # global loop
     global_idxs = [Variable(f"gidx{i}", 0, self.full_shape[i]-1 if i < self.first_reduce else 0) for i in range(0, self.first_reduce+len(self.group_for_reduce))]
-    self.uop(UOps.LOOP, (global_idxs, "global"))
+    self.uop(UOps.LOOP, None, [], (global_idxs, "global"))
 
     # local loop
     if self.group_for_reduce:
       # NOTE: this is assuming the global size = the local size in these dims. in general, this doesn't have to be true
       local_idxs = [Variable(f"lidx{i}", 0, self.full_shape[i]-1 if i >= self.first_reduce else 0) for i in range(0, self.first_reduce+len(self.group_for_reduce))]
-      self.uop(UOps.LOOP, (local_idxs, "local"))
+      self.uop(UOps.LOOP, None, [], (local_idxs, "local"))
       gl_idxs = [x*(y.max+1)+y for x,y in zip(global_idxs, local_idxs)]
     else:
       # without local idxs, it's just the global idxs
@@ -168,11 +184,11 @@ class Linearizer:
     # reduce op
     if self.reduceop is not None:
       # define accumulator
-      acc = [self.uop(UOps.CONST, ({ReduceOps.SUM: 0.0, ReduceOps.MAX: -math.inf}[cast(ReduceOps, self.reduceop.op)],), ssa('acc')) for _ in self.registers[0].offsets()]
+      acc = [self.uop(UOps.CONST, ssa('acc'), [], {ReduceOps.SUM: 0.0, ReduceOps.MAX: -math.inf}[cast(ReduceOps, self.reduceop.op)]) for _ in self.registers[0].offsets()]
 
       # reduce loop
       reduce_idxs = [Variable(f"ridx{i}", 0, self.full_shape[i]-1) for i in range(self.first_reduce+len(self.group_for_reduce), self.shape_len)]
-      self.uop(UOps.LOOP, (reduce_idxs, "reduce"))
+      self.uop(UOps.LOOP, None, [], (reduce_idxs, "reduce"))
 
       # load earlybufs
       loaded_buffers.update({b:global_buf(i, gl_idxs+reduce_idxs) for i,b in enumerate(self.bufs) if b in self.earlybufs and i != 0})
@@ -181,12 +197,12 @@ class Linearizer:
       self.ast_parse(self.reduceop, [acc[off] for off in self.registers[self.full_buf_index].acc_offsets()], loaded_buffers, ssa, do_reduce=True)
 
       # end the reduce loop
-      self.uop(UOps.ENDLOOP, (reduce_idxs, "reduce"))
+      self.uop(UOps.ENDLOOP, None, [], (reduce_idxs, "reduce"))
 
       # end the local loop, do the local reduce
       if self.group_for_reduce:
         global_buf(-1, local_idxs, acc)  # store accumulators
-        self.uop(UOps.ENDLOOP, (local_idxs, "local"))   # this is a barrier on GPUs
+        self.uop(UOps.ENDLOOP, None, [], (local_idxs, "local"))   # this is a barrier on GPUs
 
         # if any group_for_reduce items aren't reduces, upcast them here
         for j in self.upcast_in_mid_reduce_axes:
@@ -197,11 +213,11 @@ class Linearizer:
         # NOTE: this structure is the same as the reduce op above
 
         # define late accumulator
-        acc = [self.uop(UOps.CONST, ({ReduceOps.SUM: 0.0, ReduceOps.MAX: -math.inf}[cast(ReduceOps, self.reduceop.op)],), ssa('lacc')) for _ in self.registers[-1].offsets()]
+        acc = [self.uop(UOps.CONST, ssa('lacc'), [], {ReduceOps.SUM: 0.0, ReduceOps.MAX: -math.inf}[cast(ReduceOps, self.reduceop.op)]) for _ in self.registers[-1].offsets()]
 
         # late reduce loop
         end_local_idxs = [Variable(f"tidx{i}", 0, self.full_shape[i]-1 if i >= self.first_reduce else 0) for i in range(0, self.first_reduce+len(self.group_for_reduce))]
-        self.uop(UOps.LOOP, (end_local_idxs, "late_reduce"))
+        self.uop(UOps.LOOP, None, [], (end_local_idxs, "late_reduce"))
 
         # load localbufs
         loaded_buffers["LOCAL_BUFFER"] = global_buf(-1, end_local_idxs)
@@ -210,10 +226,10 @@ class Linearizer:
         self.ast_parse(LazyOp(self.reduceop.op, ("LOCAL_BUFFER",)), [acc[off] for off in self.registers[-1].acc_offsets()], loaded_buffers, ssa, do_reduce=True)
 
         # end the late reduce loop
-        self.uop(UOps.ENDLOOP, (end_local_idxs, "late_reduce"))
+        self.uop(UOps.ENDLOOP, None, [], (end_local_idxs, "late_reduce"))
 
     # load latebufs
-    loaded_buffers.update({b:global_buf(i, global_idxs) for i,b in enumerate(self.bufs) if b not in self.earlybufs and i != 0 and b is not None})
+    loaded_buffers.update({b:global_buf(i, global_idxs) for i,b in enumerate(self.bufs) if b not in self.earlybufs and i != 0 and not isinstance(b, LocalBuffer)})
 
     # run late AST
     val = self.ast_parse(self.ast, acc, loaded_buffers, ssa)
@@ -222,20 +238,16 @@ class Linearizer:
     global_buf(0, global_idxs, val)
 
     # end the global loop
-    self.uop(UOps.ENDLOOP, (global_idxs, "global"))
+    self.uop(UOps.ENDLOOP, None, [], (global_idxs, "global"))
 
     # kernel function definition
     self.function_name = ("r_" if self.reduceop else "E_") + '_'.join([str(x) for x in self.full_shape])
 
-    # print
-    if DEBUG >= 3:
-      self.printbufs()
-      for x in self.uops:
-        print(x)
 
-  def uop(self, uop:UOps, arg:Any, name:Optional[str]=None):
-    self.uops.append((uop, name, arg))
-    return name
+  def uop(self, uop:UOps, out:Optional[str], vin:List[str], arg:Any):
+    self.uops.append(UOp(uop, out, vin, arg))
+    if DEBUG >= 3: print(self.uops[-1])
+    return out
 
   def ast_parse(self, x, acc, loaded_buffers, ssa, do_reduce=False) -> List[str]:
     if not isinstance(x, LazyOp): return loaded_buffers[x]
@@ -246,12 +258,12 @@ class Linearizer:
       x = LazyOp(FusedOps.MULACC, x.src[0].src, x.arg)
     values = [self.ast_parse(v, acc, loaded_buffers, ssa) for v in x.src]
     if isinstance(x.op, (ReduceOps, FusedOps)):
-      return [self.uop(UOps.ALU, ({ReduceOps.SUM:BinaryOps.ADD, ReduceOps.MAX:BinaryOps.MAX, FusedOps.MULACC:FusedOps.MULACC}[x.op], val, val[0]), None) for val in zip(acc, *values)]
+      return [self.uop(UOps.ALU, val[0], list(val), {ReduceOps.SUM:BinaryOps.ADD, ReduceOps.MAX:BinaryOps.MAX, FusedOps.MULACC:FusedOps.MULACC}[x.op]) for val in zip(acc, *values)]
     else:
-      return [self.uop(UOps.ALU, (x.op, val), ssa('alu')) for val in zip(*values)]
+      return [self.uop(UOps.ALU, ssa('alu'), list(val), x.op) for val in zip(*values)]
 
   @property
-  def first_reduce(self) -> int: return get_first_reduce([x.shape for i,x in enumerate(self.sts) if self.bufs[i] is not None])
+  def first_reduce(self) -> int: return get_first_reduce([x.shape for i,x in enumerate(self.sts) if not isinstance(self.bufs[i], LocalBuffer)])
 
   @property
   def full_shape(self) -> Tuple[int, ...]: return self.sts[self.full_buf_index].shape

tinygrad/codegen/llvmir.py

@@ -1,7 +1,7 @@
-from typing import Final, Dict, Callable, Any, List, Optional, Tuple
+from typing import Final, Dict, Callable, Any, List, Optional
 import functools
 from llvmlite import ir  # type: ignore
-from tinygrad.codegen.linearizer import Linearizer, UOps
+from tinygrad.codegen.linearizer import Linearizer, UOps, UOp
 from tinygrad.helpers import dtypes
 from tinygrad.ops import Op, ASTRunner, UnaryOps, BinaryOps, FusedOps
 from tinygrad.lazy import LazyBuffer
@@ -32,7 +32,7 @@ code_for_op: Final[Dict[Op, Callable]] = {
   FusedOps.MULACC: lambda builder,x,y,z: builder.fadd(builder.fmul(y,z, flags=('fast',)), x, flags=('fast',)),
 }
 
-def uops_to_llvm_ir(uops:List[Tuple[UOps, Optional[str], Any]], bufs:List[LazyBuffer]) -> str:
+def uops_to_llvm_ir(uops:List[UOp], bufs:List[LazyBuffer]) -> str:
   # all llvm stuff goes into a module
   module = ir.Module(name=__file__)
 
@@ -51,9 +51,9 @@ def uops_to_llvm_ir(uops:List[Tuple[UOps, Optional[str], Any]], bufs:List[LazyBu
   lvars: Dict[Optional[str], Any] = {}  # this Any is an llvm type
   render_llvm[Variable] = lambda self,ops,ctx: lvars[self.expr]
 
-  for uop,newvar,args in uops:
+  for uop,newvar,vin,args in uops:
     if uop == UOps.CONST:
-      lvars[newvar] = ir.Constant(ir.FloatType(), args[0])
+      lvars[newvar] = ir.Constant(ir.FloatType(), args)
       reduce_phis.append(newvar)
     if uop == UOps.LOOP:
       for var in args[0]:
@@ -81,22 +81,22 @@ def uops_to_llvm_ir(uops:List[Tuple[UOps, Optional[str], Any]], bufs:List[LazyBu
         bb.append(ir.IRBuilder(func.append_basic_block(f"loop_exit_{var.expr}")))
         bb[-2].cbranch(bb[-2].icmp_unsigned("==", idx_p1, int_const(var.max+1)), bb[-1]._block, block._block)
     if uop == UOps.LOAD:
-      idx, valid = args[1].render(render_llvm, bb[-1]), args[2].render(render_llvm, bb[-1])
-      if args[2].min == 0:
+      idx, valid = args.idx.render(render_llvm, bb[-1]), args.valid.render(render_llvm, bb[-1])
+      if args.valid.min == 0:
         aug_idx = bb[-1].select(valid, idx, int_const(0))
-        val= bb[-1].select(valid, bb[-1].load(bb[-1].gep(func.args[args[0]], [aug_idx], inbounds=True)), ir.Constant(func_dtypes[args[0]], 0))
+        val = bb[-1].select(valid, bb[-1].load(bb[-1].gep(func.args[args.i], [aug_idx], inbounds=True)), ir.Constant(func_dtypes[args[0]], 0))
       else:
-        val = bb[-1].load(bb[-1].gep(func.args[args[0]], [idx], inbounds=True))
-      if func_dtypes[args[0]] != ir.FloatType(): val = bb[-1].fpext(val, ir.FloatType())
+        val = bb[-1].load(bb[-1].gep(func.args[args.i], [idx], inbounds=True))
+      if func_dtypes[args.i] != ir.FloatType(): val = bb[-1].fpext(val, ir.FloatType())
       lvars[newvar] = val
     if uop == UOps.STORE:
-      assert args[2].min == 1, "store must be valid"
-      idx = args[1].render(render_llvm, bb[-1])
-      element = lvars[args[3]]
+      assert args.valid.min == 1, "store must be valid"
+      idx = args.idx.render(render_llvm, bb[-1])
+      element = lvars[vin[0]]
       if func_dtypes[0] != ir.FloatType(): element = bb[-1].fptrunc(element, func_dtypes[0])
-      bb[-1].store(element, bb[-1].gep(func.args[args[0]], [idx], inbounds=True))
+      bb[-1].store(element, bb[-1].gep(func.args[args.i], [idx], inbounds=True))
     if uop == UOps.ALU:
-      lvars[newvar if newvar is not None else args[2]] = code_for_op[args[0]](bb[-1], *[lvars[x] for x in args[1]])
+      lvars[newvar] = code_for_op[args](bb[-1], *[lvars[x] for x in vin])
 
   bb[-1].ret_void()
   return str(module)

tinygrad/ops.py

@@ -85,7 +85,7 @@ class ASTRunner:
     return self
 
   def exec(self, bufs) -> Optional[float]:
-    rawbufs = [x.realized for x in bufs if x is not None and not isinstance(x.realized, RawConst)]
+    rawbufs = [x.realized for x in bufs if x.realized is not None and not isinstance(x.realized, RawConst)]
     if GlobalCounters.cache is not None: GlobalCounters.cache.append((self, rawbufs))
     return self(rawbufs)