tinygrad icon indicating copy to clipboard operation
tinygrad copied to clipboard

Enable JIT test for METAL. Add METAL_NO_FAST_MATH flag.

Open chenyuxyz opened this issue 1 year ago • 1 comments

Metal by default enables fast-math, which may violate IEEE 754 standard (https://developer.apple.com/documentation/metal/mtlcompileoptions/1515914-fastmathenabled). With fast-math, rand * 2 * pi (used in randn implementation) will be different from numpy about 30% of the time, which breaks the exact number tests in test_jit.py. test_ops.py was not affected because it's all testing allclose

Add a flag METAL_NO_FAST_MATH to disable fast-math.

chenyuxyz avatar Jul 12 '23 18:07 chenyuxyz

Changes made in tinygrad/:

------------------------------------------------------------
files                             insertions       deletions
------------------------------------------------------------
tinygrad/jit.py                            3               1
tinygrad/runtime/ops_metal.py              3               0
------------------------------------------------------------
total                                      6               1
------------------------------------------------------------
lines added in the tinygrad folder: 5

tinyb0t avatar Jul 12 '23 18:07 tinyb0t

Tests failing, but overall I like the change

geohot avatar Jul 17 '23 15:07 geohot

METAL test behaves differently on my local M1 vs CI. Might be related to other Metal issue. WebGPU fails because it's using METAL backend. I marked this as draft now unless we want to merge one more skipped test for METAL?

chenyuxyz avatar Jul 17 '23 19:07 chenyuxyz

I have a better way to improve this and I want to separate (1) enable JIT test in CI for all backends, (2) improve JIT tests (currently some tests still pass if we remove TinyJIT decorator), and (3) flags to control metal compile.

I will close this one and open new ones separately to make review easier.

chenyuxyz avatar Jul 18 '23 14:07 chenyuxyz