From 0d7d4dd7311a1a2ebfb127630cdeb97c18b1d042 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 26 Jul 2024 18:38:10 -0400 Subject: [PATCH] UOp._min_max for MUL and MOD (#5741) --- tinygrad/codegen/uops.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index 322d13a426..ceb5aa51d4 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -110,6 +110,9 @@ class UOp: return self.const(-self.src[0].vmax.arg), self.const(-self.src[0].vmin.arg) if self.arg is BinaryOps.ADD: return self.const(self.src[0].vmin.arg+self.src[1].vmin.arg), self.const(self.src[0].vmax.arg+self.src[1].vmax.arg) + if self.arg is BinaryOps.MUL and self.src[0].vmin.arg >= 0 and self.src[1].vmin.arg >= 0: + return self.const(self.src[0].vmin.arg*self.src[1].vmin.arg), self.const(self.src[0].vmax.arg*self.src[1].vmax.arg) + if self.arg is BinaryOps.MOD and self.src[1].op is UOps.CONST: return self.const(0), self.const(self.src[1].arg-1) return None, None @dataclass(frozen=True, repr=False) # reuse repr from UOp