From 5bdf0c79515aeccbc1646eeb5fad9730f263e4d6 Mon Sep 17 00:00:00 2001 From: quortus <156855065+quortus@users.noreply.github.com> Date: Mon, 17 Feb 2025 11:08:20 +0100 Subject: [PATCH] Bitcast constant folding 2.0 (#9089) * Prevent const folding in test_payne_hanek_reduction * Do not use list as a default parameter * Bitcast constant folding --------- Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com> --- test/test_const_folding.py | 40 ++++++++++++++++++++++++++++++++++++-- test/test_uops.py | 9 +-------- tinygrad/ops.py | 8 +++++++- 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/test/test_const_folding.py b/test/test_const_folding.py index aaf1eb8e63..ebe4fc9226 100644 --- a/test/test_const_folding.py +++ b/test/test_const_folding.py @@ -1,7 +1,10 @@ -import unittest, math +import unittest, itertools, math +from typing import Any from tinygrad import Tensor, Device, dtypes -from tinygrad.ops import Ops +from tinygrad.dtype import DType +from tinygrad.ops import Ops, UOp from tinygrad.helpers import CI +from tinygrad.codegen.rewriter import full_graph_rewrite import numpy as np from tinygrad.device import is_dtype_supported @@ -97,6 +100,39 @@ class TestBinaryOpsConstFolding(unittest.TestCase): def test_tensor_one_pow(self): _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4])) +class TestBitcastConstFolding(unittest.TestCase): + def test_scalar_bitcast(self): + def t(cases: dict[DType, Any]): + for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()): + if not math.isnan(from_v): + r = full_graph_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0] + self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})") + self.assertEqual(r.dtype, to_dt, msg) + np.testing.assert_equal(r.arg, to_v, msg) + + t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False}) + t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True}) + + t({dtypes.int8: -1, dtypes.uint8: 2**8-1}) + t({dtypes.int16: -1, dtypes.uint16: 2**16-1, dtypes.float16: float('nan')}) + t({dtypes.int32: -1, dtypes.uint32: 2**32-1, dtypes.float32: float('nan')}) + t({dtypes.int64: -1, dtypes.uint64: 2**64-1, dtypes.float64: float('nan')}) + + t({dtypes.int8: -2**7, dtypes.uint8: 2**7}) + t({dtypes.int16: -2**15, dtypes.uint16: 2**15}) + t({dtypes.int32: -2**31, dtypes.uint32: 2**31}) + t({dtypes.int64: -2**63, dtypes.uint64: 2**63}) + + t({dtypes.int16: 13496, dtypes.uint16: 13496, dtypes.float16: 0.294921875}) + t({dtypes.int32: 1050081145, dtypes.uint32: 1050081145, dtypes.float32: 0.29485681653022766}) + t({dtypes.int64: 4598983288165178391, dtypes.uint64: 4598983288165178391, dtypes.float64: 0.29485681936461233}) + + def test_vec_bitcast(self): + r = full_graph_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0] + self.assertEqual(r.op, Ops.VECTORIZE) + self.assertEqual(r.dtype, dtypes.uint32.vec(3)) + self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75)) + # folds advance indexing into basic indexing class TestIndexingConstFolding(unittest.TestCase): def test_scalar_index(self): diff --git a/test/test_uops.py b/test/test_uops.py index 9934c3820d..7ea178789a 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -11,7 +11,7 @@ from tinygrad.ops import Ops, UOp, UPat, KernelInfo, exec_alu # noqa F401 from tinygrad.spec import spec from tinygrad.renderer import ProgramSpec from tinygrad.engine.schedule import fix_kernel_ops -from tinygrad.engine.realize import CompiledRunner, lower_schedule_item, get_kernel +from tinygrad.engine.realize import CompiledRunner, get_kernel from tinygrad.codegen.linearize import linearize_uop from tinygrad.codegen.rewriter import full_graph_rewrite, sym from tinygrad.device import is_dtype_supported @@ -242,13 +242,6 @@ class TestConstantFolding(unittest.TestCase): si = t.schedule() assert len(si) == 0 - def test_bitcast_const(self): - t = Tensor(1, dtype=dtypes.float).bitcast(dtypes.int) - si = t.schedule() - assert len(si) == 1 - ji = lower_schedule_item(si[-1]) - assert any(uop.op is Ops.BITCAST for uop in ji.prg.p.uops), f"{[uop.op for uop in ji.prg.p.uops]} does not contain bitcast" - class TestGatedStoreRewrite(unittest.TestCase): def test_tiny_gate_store(self): gmem = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(), (), 0) diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 772634a11b..dda0f11e63 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -1,6 +1,6 @@ from __future__ import annotations from typing import Any, Optional, Union, Callable, cast, TYPE_CHECKING, Type, Literal, get_args -import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref +import sys, time, functools, itertools, math, operator, hashlib, os, types, pickle, pathlib, inspect, weakref, struct from enum import auto, IntEnum, Enum from dataclasses import dataclass, field from collections import defaultdict @@ -1114,6 +1114,11 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None: if int(c.arg) == c.arg: return (y := x.pow(c.const_like(c.arg//2))) * y * (x if c.arg%2 == 1 else 1) return None +def fold_bitcast(root:UOp, c:UOp) -> UOp|None: + if (from_fmt:=c.dtype.scalar().fmt) is None or (to_fmt:=root.dtype.scalar().fmt) is None: return None + def convert(v:Any): return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0] + return root.const_like(convert(c.arg) if root.dtype.count == 1 else tuple(map(convert, c.arg))) + # def max_var_const(x:UOp, c1:UOp, c2:UOp): # if x.vmin >= 0: return x*c1 if c1.arg >= c2.arg else x*c2 # if x.vmax <= 0: return x*c2 if c1.arg >= c2.arg else x*c1 @@ -1157,6 +1162,7 @@ symbolic_simple = PatternMatcher([ # *** cast/bitcast *** (UPat(Ops.CAST, name="root", src=UPat.cvar("c")), lambda root, c: root.const_like(c.arg)), (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None), + (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # ** pow ** (UPat.var("x").alu(Ops.POW, UPat.cvar("c", vec=False)), simplify_pow), # positive const ** x