From 0258c7fefc995c7b3a4e602c85b03fdb8f0ecb49 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 1 Aug 2026 22:37:06 -0400 Subject: [PATCH] minor argstr and alloc stack cleanup [PR] (#17361) --- tinygrad/codegen/late/regalloc.py | 2 +- tinygrad/uop/ops.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tinygrad/codegen/late/regalloc.py b/tinygrad/codegen/late/regalloc.py index 85cf7cfb1b..0d675e62ee 100644 --- a/tinygrad/codegen/late/regalloc.py +++ b/tinygrad/codegen/late/regalloc.py @@ -125,7 +125,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp): # alloc/dealloc stack if ctx.stack_size > 0: sp = ctx.ren.stack_pointer() - offset = UOp(Ops.CONST, sp.dtype, arg=ctx.stack_size) + offset = UOp.const(ctx.stack_size, sp.dtype) if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))] diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 4168089812..67129a675d 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from enum import Enum, auto from tinygrad.uop import Ops, GroupOp from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace, strong_dtype -from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar, weak_dtype +from tinygrad.dtype import PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar, weak_dtype from tinygrad.device import Buffer, MultiBuffer, canonicalize_device, TinyELF from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey @@ -272,7 +272,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return pretty_print(self) def argstr(self): if self.op is Ops.REDUCE: return f'({", ".join(map(str, self.arg))})' - return f"ConstFloat({float.__repr__(self.arg)})" if isinstance(self.arg, ConstFloat) else repr(self.arg) + return repr(self.arg) def tagstr(self): return f", tag={self.tag}" if self.tag is not None else "" def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)