forked from tinygrad/tinygrad
Merge branch 'master' into new_x86_backend
This commit is contained in:
ttomsa
@@ -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'):
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
+11
@@ -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]}")
|
||||
@@ -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
|
||||
|
||||
|
||||
+13
-2
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
+26
-1
@@ -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
|
||||
|
||||
@@ -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([])
|
||||
|
||||
@@ -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))])
|
||||
|
||||
@@ -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: <regname>_<field_x>__SHIFT and <regname>_<field_y>__MASK from the adreno module into the following format:
|
||||
# qreg.<regname>(<field_x>=..., <field_y>=..., ..., <field_n>=...)
|
||||
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)
|
||||
|
||||
+2
-2
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
+10
-4
@@ -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
|
||||
|
||||
@@ -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"<Kernel {len(list(u.arg.ast.toposort()))} {ast_str} {[str(m) for m in u.arg.metadata]}>"
|
||||
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 [])):
|
||||
|
||||
Reference in New Issue
Block a user