diff --git a/extra/assembly/amd/pcode.py b/extra/assembly/amd/pcode.py index 166f2c5e48..4151d9daba 100644 --- a/extra/assembly/amd/pcode.py +++ b/extra/assembly/amd/pcode.py @@ -440,7 +440,7 @@ class Parser: self.eat('COMMA') lo = self.parse() self.eat('RBRACE') - return (hi.cast(dtypes.uint64) << _u64(32)) | lo.cast(dtypes.uint64) + return (hi.cast(dt:=_BITS_DT.get((s:=lo.dtype.bitsize) * 2, dtypes.uint64)) << _const(dt, s)) | lo.cast(dt) if self.at('NUM'): num = self.eat('NUM').val if self.try_eat('QUOTE'): diff --git a/extra/assembly/amd/test/test_emu2_pcode.py b/extra/assembly/amd/test/test_emu2_pcode.py index 0be69661a7..a1b5450948 100644 --- a/extra/assembly/amd/test/test_emu2_pcode.py +++ b/extra/assembly/amd/test/test_emu2_pcode.py @@ -62,6 +62,16 @@ class TestWithSources(unittest.TestCase): dest, val = assigns[0] self.assertEqual(val.op, Ops.MUL) + def test_s_pack_ll_b32_b16(self): + """Test S_PACK_LL_B32_B16 packs two 16-bit values into 32-bit result.""" + s0 = UOp.const(dtypes.uint32, 0xDEADAAAA) + s1 = UOp.const(dtypes.uint32, 0xDEADBBBB) + _, assigns = parse_pcode(PCODE[SOP2Op.S_PACK_LL_B32_B16], {'S0': s0, 'S1': s1}) + self.assertEqual(len(assigns), 1) + dest, val = assigns[0] + self.assertTrue(dest.startswith('D0')) + self.assertEqual(val.simplify().arg, 0xBBBBAAAA) + class TestParseExpr(unittest.TestCase): """Test the parse_expr function directly.""" diff --git a/extra/gemm/asm/cdna/test_asm_gemm.py b/extra/gemm/asm/cdna/test_asm_gemm.py index 8281446353..53e7b6d5b5 100644 --- a/extra/gemm/asm/cdna/test_asm_gemm.py +++ b/extra/gemm/asm/cdna/test_asm_gemm.py @@ -48,6 +48,7 @@ class TestGemmLarge(unittest.TestCase): def test_gemm4(self): verify_asm_gemm(8, 4096, 14336, 4096, gpus=8) def test_gemm5(self): verify_asm_gemm(8, 4096, 4096, 14336, gpus=8) def test_gemm6(self): verify_asm_gemm(16, 4096, 4096, 14336, gpus=8) + def test_gemm7(self): verify_asm_gemm(1, 8192, 128256, 4096) if __name__ == "__main__": unittest.main() diff --git a/extra/models/llama.py b/extra/models/llama.py index 0448efd3c9..b1f5be15c7 100644 --- a/extra/models/llama.py +++ b/extra/models/llama.py @@ -88,9 +88,11 @@ class Attention: keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep) xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2) - attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2) + if Tensor.training: + attn = xq.scaled_dot_product_attention(keys, values, is_causal=True).transpose(1, 2) + else: + attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2) if getenv("STUB_ATTENTION"): - # TODO: do we need mask? from tinygrad.uop.ops import UOp, KernelInfo def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp: return UOp.sink(arg=KernelInfo(name="fa_custom_forward")) @@ -197,7 +199,9 @@ class Transformer: h = self.tok_embeddings(tokens) freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :] - mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None + if not Tensor.training and seqlen > 1: + mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) + else: mask = None for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask) logits = self.output(self.norm(h)) if math.isnan(temperature): return logits diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index bfc7594f45..98c04a5ea0 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -344,7 +344,7 @@ def scatter_add(self, dim, index, src, out): def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False): if src.is_tiny and dest.is_tiny: src_t, dest_t = unwrap(src), unwrap(dest) - if dest_t.uop.is_contiguous() or dest_t.uop.is_realized: src_t = src_t.contiguous() + if dest_t.uop.has_buffer_identity() or dest_t.uop.is_realized: src_t = src_t.contiguous() _apply_inplace(dest_t, src_t.cast(cast_dtype).to(to_device)) elif src.is_tiny and dest.is_cpu: dest.resize_(src.numel()).resize_(src.shape) diff --git a/test/external/external_fuzz_qcom_cpu_cache.py b/test/external/external_fuzz_qcom_cpu_cache.py new file mode 100644 index 0000000000..62832db27a --- /dev/null +++ b/test/external/external_fuzz_qcom_cpu_cache.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 +from tinygrad.tensor import Tensor +import numpy as np + +while True: + arr = np.ones(1000000, dtype=np.uint8) + print(f"numpy: {(arr + 1)[:10]}") + + ptr = arr.ctypes.data + tensor = Tensor.from_blob(ptr, arr.shape, dtype='uint8', device='QCOM').realize() + 1 + print(f"from_blob: {tensor.numpy()[:10]}") diff --git a/test/test_dtype.py b/test/test_dtype.py index 932a5553d1..7120f503a4 100644 --- a/test/test_dtype.py +++ b/test/test_dtype.py @@ -248,6 +248,18 @@ class TestBFloat16DTypeCast(unittest.TestCase): class TestHalfDType(TestDType): DTYPE = dtypes.half +@unittest.skipUnless(Ops.SHL in Device[Device.DEFAULT].renderer.code_for_op, "half decomp requires bitshift") +class TestEmulatedHalf(TestHalfDType): + @classmethod + def setUpClass(cls): + cls.stack = contextlib.ExitStack() + cls.stack.enter_context(Context(EMULATED_DTYPES="half")) + cls.DATA = rand_for_dtype(cls.DTYPE, 10) + + @classmethod + def tearDownClass(cls): cls.stack.close() + + class TestFloatDType(TestDType): DTYPE = dtypes.float diff --git a/test/test_dtype_alu.py b/test/test_dtype_alu.py index 9fdf54849f..b9d0b987fc 100644 --- a/test/test_dtype_alu.py +++ b/test/test_dtype_alu.py @@ -1,7 +1,7 @@ import unittest, operator, math from tinygrad import Context, Tensor, dtypes, Device from tinygrad.dtype import DType, truncate -from tinygrad.helpers import CI, getenv +from tinygrad.helpers import CI, EMULATED_DTYPES, getenv from tinygrad.tensor import _to_np_dtype from tinygrad.device import is_dtype_supported from tinygrad.runtime.ops_python import from_storage_scalar @@ -64,7 +64,10 @@ def universal_test(a, b, dtype, op): numpy_value = op[1](ta.numpy(), tb.numpy()) if dtype in dtypes.fp8s: numpy_value = truncate[dtype](numpy_value.item()) if dtype in dtypes.floats: - atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7)) + if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero + fe, fm = dtypes.finfo(dtype) + atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm) + else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7)) np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol) else: np.testing.assert_equal(tensor_value, numpy_value) @@ -117,6 +120,10 @@ class TestDTypeALU(unittest.TestCase): @given(ht.float16, ht.float16, strat.sampled_from(binary_operations)) def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op) + @given(ht.float16, ht.float16, strat.sampled_from(binary_operations)) + @Context(EMULATED_DTYPES="half") + def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op) + @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}") @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations)) def test_bfloat16(self, a, b, op): @@ -139,6 +146,10 @@ class TestDTypeALU(unittest.TestCase): @given(ht.float16, strat.sampled_from(unary_operations)) def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op) + @given(ht.float16, strat.sampled_from(unary_operations)) + @Context(EMULATED_DTYPES="half") + def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op) + @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}") @given(ht.bfloat16, strat.sampled_from(unary_operations)) def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op) diff --git a/test/test_schedule.py b/test/test_schedule.py index 772de58e62..5a28d03b70 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -175,6 +175,14 @@ class TestSchedule(unittest.TestCase): child.realize() assert a.uop.is_realized + def test_realize_view_of_realized_has_empty_schedule(self): + # views of realized buffers produce an empty schedule + t = Tensor.zeros((3, 3)).contiguous().realize() + v = t[1] # view - is_realized but not has_buffer_identity + assert v.uop.is_realized + sched, _ = Tensor.schedule_with_vars(v) + self.assertEqual(len(sched), 0) + # NOTE: because empty does not have a lowered ExecItem if realize is called on a childless empty, it never gets allocated. def test_childless_empty_never_allocates(self): a = Tensor.empty(10) diff --git a/test/test_setitem.py b/test/test_setitem.py index d207c3ec66..82f37f5324 100644 --- a/test/test_setitem.py +++ b/test/test_setitem.py @@ -54,7 +54,32 @@ class TestSetitem(unittest.TestCase): t = Tensor.ones(4) with self.assertRaises(RuntimeError): t[1] = 5 - @unittest.skip("TODO: flaky") + def test_setitem_chained_indexing(self): + # N[i][j] must work the same as N[i, j] + N1 = Tensor.zeros((3, 3)).contiguous().realize() + N1[1, 2] = 5 + N2 = Tensor.zeros((3, 3)).contiguous().realize() + N2[1][2] = 5 + np.testing.assert_equal(N1.numpy(), N2.numpy()) + + def test_setitem_detach(self): + # setitem on detached tensor should work + t = Tensor.zeros((3, 3)).contiguous().realize() + t.detach()[1, 2] = 5 + self.assertEqual(t[1, 2].item(), 5.0) + + def test_setitem_permute(self): + # setitem on permuted tensor should modify original + t = Tensor.zeros((2, 3)).contiguous().realize() + t.T[1, 0] = 5 # t.T is (3, 2), so [1, 0] maps to t[0, 1] + self.assertEqual(t[0, 1].item(), 5.0) + + def test_setitem_flip(self): + # setitem on flipped tensor should modify original + t = Tensor.zeros((3,)).contiguous().realize() + t[::-1][0] = 5 # flip, then set first element (which is last in original) + self.assertEqual(t[2].item(), 5.0) + def test_setitem_inplace_operator(self): t = Tensor.arange(4).reshape(2, 2).contiguous() t[1] += 2 diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 00d8d7d274..5ff129e414 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -11,7 +11,7 @@ from tinygrad.codegen.opt import Opt # import all pattern matchers here from tinygrad.codegen.gpudims import pm_add_gpudims from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load -from tinygrad.uop.decompositions import get_late_rewrite_patterns +from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_unsupported_dtypes_patterns, get_transcendental_patterns from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \ ReduceContext, correct_load_store, pm_render, pm_add_loads @@ -92,9 +92,12 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) - # decompositions supported_ops = tuple(ren.code_for_op.keys()) - pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, TRANSCENDENTAL>=2, bool(DISABLE_FAST_IDIV), - tuple(EMULATED_DTYPES.tolist(dtypes))) + pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV)) + pm_unsupported = get_unsupported_dtypes_patterns(ren.device, tuple(EMULATED_DTYPES.tolist(dtypes))) + pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions") + sink = graph_rewrite(sink, pm_unsupported, ctx=ren.device, name="unsupported dtypes", bottom_up=True) + sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental") # final rules for the renderer (without sym) extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index f9b78288c1..d98460aa6a 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -257,8 +257,7 @@ def _prepare_jit_inputs(args, kwargs): input_uops: list[UOp] = flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors]) if any(u.base.op is Ops.CONST for u in input_uops): raise JitError("JIT inputs cannot be const, create a buffer with .contiguous()") - input_buffers: list[Buffer] = flatten([b.bufs if isinstance(b:=u.base.realized, MultiBuffer) else [b] - for u in input_uops if u.base.realized is not None]) + input_buffers: list[Buffer] = flatten([b.bufs if isinstance(b, MultiBuffer) else [b] for u in input_uops if (b:=u.base.realized) is not None]) if len(set(input_buffers)) != len(input_buffers): raise JitError("duplicate inputs to JIT") inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops] _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))]) diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index 648dd455a8..f8c6c584d6 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -2,21 +2,33 @@ from __future__ import annotations import os, ctypes, functools, mmap, struct, array, math, sys, weakref, contextlib assert sys.platform != 'win32' from typing import Any -from tinygrad.device import BufferSpec, CompilerSet, CompilerPair +from tinygrad.device import BufferSpec, CompilerSet, CompilerPair, Device from tinygrad.runtime.support.hcq import HCQBuffer, HWQueue, HCQProgram, HCQCompiled, HCQAllocatorBase, HCQSignal, HCQArgsState, BumpAllocator from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface from tinygrad.runtime.autogen import kgsl, mesa from tinygrad.runtime.ops_cl import CLCompiler, CLDevice from tinygrad.renderer.cstyle import QCOMRenderer from tinygrad.renderer.nir import IR3Renderer -from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, prod, fromimport, cpu_profile, lo32, PROFILE, suppress_finalizing -from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC -from tinygrad.dtype import ImageDType +from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, fromimport, cpu_profile, lo32, suppress_finalizing +from tinygrad.helpers import next_power2, flatten, QCOM_IR3, QCOM_CC, PROFILE +from tinygrad.dtype import ImageDType, dtypes from tinygrad.runtime.support.system import System if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl # noqa: F401 # pylint: disable=unused-import BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2 +@functools.cache +def dcache_flush(): + from tinygrad.uop.ops import UOp, Ops, KernelInfo + from tinygrad.codegen import get_program + buf, n = UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=0), UOp(Ops.PARAM, dtypes.uint8.ptr(), arg=1) + i = UOp.range(n.cast(dtypes.int), 0, dtype=dtypes.int) + flush = UOp(Ops.CUSTOM, dtypes.void, (buf.cast(dtypes.ulong) + i.cast(dtypes.ulong) * UOp.const(dtypes.ulong, 64),), + arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");') + sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, dtypes.void, (), arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush")) + ps = get_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())))), Device["CPU"].renderer) + return Device["CPU"].runtime(ps.function_name, ps.lib) + #Parse C-style defines: ___SHIFT and ___MASK from the adreno module into the following format: # qreg.(=..., =..., ..., =...) def _qreg_exec(__reg, __val=0, **kwargs): @@ -391,6 +403,7 @@ class QCOMDevice(HCQCompiled): def _gpu_map(self, ptr:int, size:int, **kwargs) -> HCQBuffer: ptr_aligned, size_aligned = (ptr & ~0xfff), round_up(size + (ptr & 0xfff), 0x1000) + dcache_flush().fxn(ctypes.c_uint64(ptr_line_aligned:=ptr & ~63), ctypes.c_uint64(ceildiv(ptr + size - ptr_line_aligned, 64))) try: mi = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR) return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self, **kwargs) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index aa45bdfa6e..b50c710b9c 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -272,7 +272,7 @@ class Tensor(OpMixin): @disable_gc() def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor: """Triggers the computation needed to create these Tensor(s).""" - if len(to_realize:=[x for x in (self,)+lst if not x.uop.is_contiguous()]): + if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]): run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats) return self @@ -1286,7 +1286,7 @@ class Tensor(OpMixin): if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype) if self.requires_grad or v.requires_grad: raise NotImplementedError("setitem with requires_grad is not supported") self.realize() - if not self.uop.is_contiguous(): raise RuntimeError("setitem target needs to be contiguous") + if not self.uop.is_writable_view(): raise RuntimeError("setitem target must be a writable view backed by a buffer") res = self._getitem(indices, v) # if shapes match and data is not shared it's a copy and we assign to self if res.shape == self.shape and res.uop is not self.uop: diff --git a/tinygrad/uop/decompositions.py b/tinygrad/uop/decompositions.py index 0675f3934b..aa0677eaac 100644 --- a/tinygrad/uop/decompositions.py +++ b/tinygrad/uop/decompositions.py @@ -18,8 +18,8 @@ def exponent_bias(d:DType) -> int: return {dtypes.float64: 1023, dtypes.float32: def exponent_mask(d:DType) -> int: return {dtypes.float64: 2047, dtypes.float32: 255, dtypes.float16: 31}[d.scalar()] # **** utils **** -def shr(x:UOp, y:int) -> UOp: return x // (2**y) -def shl(x:UOp, y:int) -> UOp: return x * (2**y) +def shr(x:UOp|int, y:int) -> UOp: return x // (2**y) +def shl(x:UOp|int, y:int) -> UOp: return x * (2**y) def rintk(d:UOp) -> UOp: """round d:float to int away from 0""" @@ -319,7 +319,7 @@ def threefry2x32(x: UOp, key: UOp): l2i_dt = {dtypes.long: dtypes.int, dtypes.ulong: dtypes.uint} def unpack32(v:UOp) -> tuple[UOp, UOp]: return v.bitcast(dtypes.uint) & 0xFFFF, v.bitcast(dtypes.uint) >> 16 -def l2i_idx(idx:UOp, off:int) -> UOp: return idx.replace(src=(idx.src[0], idx.src[1]*2+off)) +def reindex(idx:UOp, off:int, mul=2) -> UOp: return idx.replace(src=(idx.src[0], idx.src[1]*mul+off)) # 4.3.1 is the relevant section in TAOCP def l2i(op: Ops, dt: DType, *uops:UOp): @@ -377,12 +377,39 @@ def l2i(op: Ops, dt: DType, *uops:UOp): case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1) case _: raise NotImplementedError(f"long decomposition of {op} unsupported") +# ***** floats ***** +f2f_dt = { dtypes.half: dtypes.ushort, dtypes.float: dtypes.uint } + +def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1))) + +def f2f(v, fr:DType, to:DType): + fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to) + # NB: denormals are zero! + if fe < te and fm < tm: + sign, nosign = shl((v & shl(1, fs-1)).cast(f2f_dt[to]), ts - fs), (v & (shl(1, fs-1) - 1)).cast(f2f_dt[to]) + exp, norm = shr(nosign, fm), shl(nosign, tm - fm) + shl(tb - fb, tm) + inf_or_nan = shl(nosign, tm - fm) | shl((shl(1, te) - 1), tm) + return (sign | exp.eq(0).where(0, exp.eq(shl(1, fe) - 1).where(inf_or_nan, norm))).bitcast(to) + elif fe > te and fm > tm: + sign, nosign, exp = shr(v, fs - ts) & shl(1, ts - 1), v & (shl(1, fs - 1) - 1), shr(v, fm) & (shl(1, fe) - 1) + norm = (rne(nosign, fm - tm) - shl(fb - tb, tm)).cast(f2f_dt[to]) + infnan = (sign | (shr(nosign, fm - tm) & (shl(1, tm) - 1)) | shl(shl(1, te) - 1, tm)).cast(f2f_dt[to]) + underflow, overflow = exp < (1 + fb - tb), exp > (shl(1, te) - 2 + (fb - tb)) + return exp.eq(shl(1, fe) - 1).where(infnan, sign.cast(f2f_dt[to]) | underflow.where(0, overflow.where(shl(shl(1, te) - 1, tm), norm))) + else: raise NotImplementedError(f"unsupported decomp {fr} -> {to}") + +def f2f_load(x: UOp) -> UOp: + if (n:=x.dtype.count) == 1: return f2f(x.replace(dtype=dtypes.ushort), dtypes.half, dtypes.float) + return UOp.vectorize(*(f2f(x.replace(dtype=dtypes.ushort, src=(reindex(x.src[0].src[0], i, 1),)), dtypes.half, dtypes.float) for i in range(n))) + +def f2f_store(st, idx, val): + if (n:=val.dtype.count) == 1: return st.replace(src=(idx, f2f(val.bitcast(dtypes.uint), dtypes.float, dtypes.half))) + return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.gep(i).bitcast(dtypes.uint), dtypes.float, dtypes.half))) for i in range(n))) + # ***** decomposition patterns ***** -powers_of_two: dict[int, int] = {2**i:i for i in range(64)} @functools.cache -def get_late_rewrite_patterns(ops:tuple[Ops, ...], device:str, force_transcendental:bool, disable_fast_idiv:bool, - emulated_dtypes:tuple[DType, ...]) -> PatternMatcher: +def get_transcendental_patterns(ops:tuple[Ops, ...], force_transcendental:bool) -> PatternMatcher: pat: list[tuple[UPat, Callable]] = [] for op,f in ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)): if op not in ops or force_transcendental: @@ -391,6 +418,12 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], device:str, force_transcenden lambda x,d: d.cast(dtypes.float32).alu(x.op).cast(x.dtype))] # rewrite SQRT to xpow 0.5 if Ops.SQRT not in ops or force_transcendental: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5)))) + return PatternMatcher(pat) + +powers_of_two: dict[int, int] = {2**i:i for i in range(64)} +@functools.cache +def get_late_rewrite_patterns(ops:tuple[Ops, ...], device:str, disable_fast_idiv:bool) -> PatternMatcher: + pat: list[tuple[UPat, Callable]] = [] # no real hardware supports THREEFRY, but NullRenderer does if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32)) # MAX can be rewritten as CMPLT + WHERE (max function is annoying on many cstyle backends) @@ -428,26 +461,42 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], device:str, force_transcenden if Ops.FDIV in ops: pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))] pat += [(UPat.var("a", dtypes.floats) * UPat.const(dtypes.floats, 1).alu(Ops.FDIV, UPat.var("b")), lambda a,b: a.alu(Ops.FDIV, b))] + return PatternMatcher(pat) + +@functools.cache +def get_unsupported_dtypes_patterns(device:str, emulated_dtypes:tuple[DType, ...]) -> PatternMatcher: + pat: list[tuple[UPat, Callable]] = [] if not is_dtype_supported(dtypes.long, device) or dtypes.long in emulated_dtypes: pat += [(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x: x.replace(dtype=l2i_dt[x.dtype.base].ptr(x.dtype.size * 2)) if hasattr(x.dtype, 'size') and x.dtype.base in l2i_dt else None)] - pat += [(UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: - None if x.tag is None else x.replace(dtype=l2i_dt[x.dtype], src=(x.src[0], x.src[1]*2+x.tag)))] + pat += [(UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype]))] pat += [(UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val: - st.replace(src=(l2i_idx(idx, 0), val.rtag(0))).group(st.replace(src=(l2i_idx(idx, 1), val.rtag(1)))) if val.tag is None else None)] + st.replace(src=(reindex(idx, 0), val.rtag(0))).group(st.replace(src=(reindex(idx, 1), val.rtag(1)))) if val.tag is None else None)] pat += [(UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x: l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt)))] pat += [(UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x: l2i(x.op, x.dtype, a)[x.tag] if x.tag is not None and a.dtype not in l2i_dt else None)] pat += [(UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x: - None if x.tag is None else (a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag])] + (a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag])] pat += [(UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x: l2i(x.op, x.dtype, a.rtag(0).cast(dt:=l2i_dt[a.dtype]), a.rtag(1).cast(dt)) if x.dtype not in l2i_dt and a.tag is None else None)] pat += [(UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x: - None if x.tag is None else l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt)) + l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt)) if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag])] pat += [(UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: - None if x.tag is None else x.replace(dtype=l2i_dt[x.dtype], src=(l2i_idx(idx, x.tag),)))] + x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag),)))] pat += [(UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x: - None if x.tag is None else UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))] + UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))] + if dtypes.half in emulated_dtypes: + pat += [(UPat((*GroupOp.Defines, Ops.INDEX), name="x"), lambda x: + x.replace(dtype=dtypes.uint16.ptr(x.dtype.size), tag=dtypes.half) if x.dtype.base == dtypes.half else None)] + pat += [(UPat(Ops.LOAD, dtypes.half, name="x"), f2f_load)] + pat += [(UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, dtypes.half, name="ld"),), name="bc"), lambda bc,ld: + ld.replace(dtype=dtypes.ushort).bitcast(bc.dtype))] + pat += [(UPat(Ops.BITCAST, (dtypes.ushort, dtypes.short, dtypes.bfloat16), src=(UPat.var("x", dtypes.float),), name="bc"), lambda bc,x: + bc.replace(src=(f2f(x.bitcast(dtypes.uint), dtypes.float, dtypes.half),)))] + pat += [(UPat(GroupOp.All, dtypes.half, name="x"), lambda x: + x.replace(dtype=dtypes.float.vec(x.dtype.count), src=tuple(s.cast(dtypes.float) if s.dtype == dtypes.half else s for s in x.src)))] + pat += [(UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.float)), name='st'), lambda st,idx,val: + f2f_store(st, idx, val) if (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == dtypes.half else None)] return PatternMatcher(pat) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 00635358aa..e7fe01ffc3 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -467,14 +467,15 @@ class UOp(OpMixin, Generic[OpT], metaclass=UOpMetaClass): return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(dtypes.bool, self.arg is not Invalid) def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs) - def is_contiguous(self): - # TODO: this is is_realized - if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].is_contiguous() + def is_writable_view(self) -> bool: + """Check if this UOp is a writable view backed by a buffer (injective mapping).""" + if self.op in {Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.DETACH}: return self.src[0].is_writable_view() + if self.op is Ops.MULTI: return all(x.is_writable_view() for x in self.src) return self.op is Ops.BUFFER def contiguous(self, *args, **kwargs): if self.op is Ops.CONTIGUOUS: return self - if self.is_contiguous(): return self + if self.has_buffer_identity(): return self return UOp(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs) def contiguous_backward(self): return self.alu(Ops.CONTIGUOUS_BACKWARD) def bufferize(self, *args, **kwargs): return UOp(Ops.BUFFERIZE, dtype=self.dtype, src=(self,)+args, **kwargs) @@ -640,6 +641,11 @@ class UOp(OpMixin, Generic[OpT], metaclass=UOpMetaClass): return self.src[0].buf_target() case _: raise RuntimeError(f"buf_target called on non load/index/store {self.op}") + def has_buffer_identity(self): + """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain).""" + if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity() + return self.op is Ops.BUFFER + @property def buffer(self) -> Buffer|MultiBuffer: from tinygrad.device import Buffer, MultiBuffer diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 002fd4390b..54d889f104 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -109,6 +109,7 @@ def uop_to_json(x:UOp) -> dict[int, dict]: if u.op is Ops.KERNEL: ast_str = f"SINK{tuple(s.op for s in u.arg.ast.src)}" if u.arg.ast.op is Ops.SINK else repr(u.arg.ast.op) argst = f"" + if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>" label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''))) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" for idx,x in enumerate(u.src[:1] if u.op in {Ops.BUFFERIZE, Ops.INDEX} else (u.src if u.op is not Ops.END else [])):