| 85 | class UOpMetaClass(type): |
| 86 | ucache:dict[tuple, weakref.ReferenceType[UOp]] = {} |
| 87 | def __call__(cls, op:Ops, dtype:DType=dtypes.void, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None, |
| 88 | metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None): |
| 89 | if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret |
| 90 | UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key)) |
| 91 | if metadata is not None: all_metadata[created] = metadata |
| 92 | # NOTE: this value is set by pickle when pickling a realized tensor |
| 93 | if _buffer is not None: |
| 94 | assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}" |
| 95 | buffers[created] = _buffer |
| 96 | if SPEC > 1: |
| 97 | from tinygrad.uop.spec import spec_full, test_pyrender |
| 98 | if SPEC > 2: |
| 99 | # SPEC=3 checks the shape |
| 100 | _ = created._shape |
| 101 | if SPEC > 3: |
| 102 | test_pyrender(created) |
| 103 | with Context(CHECK_OOB=0): fret = cast(bool|None, spec_full.rewrite(created)) |
| 104 | if fret is not True: raise RuntimeError(f"SPEC ISSUE {fret}: {created}") |
| 105 | return created |
| 106 | |
| 107 | # some uops map to other stuff |
| 108 | buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/BUFFER_VIEW uops to their device Buffers |