From 37347528bf2c8190caa1be7346d0f65d4b0f490c Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Wed, 24 Jul 2024 19:33:04 +0800 Subject: [PATCH] shorter BufferOps.CONST creation (#5681) --- tinygrad/engine/schedule.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index b67053b070..71c75d9fc1 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -49,13 +49,9 @@ def _recursive_lazyop(buf:LazyBuffer, inputs:List[LazyBuffer], outputs:Tuple[Laz if buf.op is MetaOps.CONST: unbound_st, st_var_vals = st.simplify().unbind() var_vals.update(st_var_vals) - if isinstance(buf.arg, Variable): - val, var_val = buf.arg.unbind() - var_vals.__setitem__(val, var_val) - else: - assert isinstance(buf.arg, get_args(ConstType)), f"cannot create ConstBuffer with value {buf.arg}" - val = buf.arg - return LazyOp(BufferOps.CONST, (), ConstBuffer(val, buf.dtype, unbound_st)) + if isinstance(arg, Variable): var_vals.update([arg:=arg.unbind()]) + else: assert isinstance(arg, get_args(ConstType)), f"cannot create ConstBuffer with value {arg}" + return LazyOp(BufferOps.CONST, (), ConstBuffer(arg[0] if isinstance(arg, tuple) else arg, buf.dtype, unbound_st)) # if we aren't fusing it, it's a load and we add it to the inputs if buf.realized is not None or (buf in realizes and buf not in outputs):