ops_python: add HIP tensor core mock and refactor METAL (#3354)

* ops_python: add HIP tensor core mock and refactor METAL

* Add tests to CI

* add DEBUG=2 to full tests

---------

Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
This commit is contained in:
Francis Lam
2024-02-09 03:46:06 -08:00
committed by GitHub
parent b385234961
commit ce21fdfb67
3 changed files with 45 additions and 25 deletions

View File

@@ -106,7 +106,7 @@ class TestLinearizer(unittest.TestCase):
def test_tensor_cores(self):
if not Device[Device.DEFAULT].compiler.linearizer_opts.has_tensor_cores:
self.skipTest("device doesn't have tensor cores")
for tc in tensor_cores[Device.DEFAULT]:
for tc in tensor_cores[Device[Device.DEFAULT].compiler.linearizer_opts.device]:
a, b = Tensor.rand(tc.dims[1], tc.dims[2], dtype=tc.dtype_in), Tensor.rand(tc.dims[2], tc.dims[0], dtype=tc.dtype_in)
np_a, np_b = a.numpy(), b.numpy()
r = a.matmul(b, acc_dtype=tc.dtype_out)
@@ -536,7 +536,7 @@ class TestLinearizerOpts(unittest.TestCase):
N = 128
Tensor.manual_seed(1552)
for tc in tensor_cores[Device.DEFAULT]:
for tc in tensor_cores[Device[Device.DEFAULT].compiler.linearizer_opts.device]:
a, b = Tensor.rand(N, N, dtype=tc.dtype_in), Tensor.rand(N, N, dtype=tc.dtype_in)
r = a.matmul(b, acc_dtype=tc.dtype_out)
(atol, rtol) = ((0.25, 0.01) if tc.dtype_out == dtypes.half else (3e-2, 1e-3)) if tc.dtype_in == dtypes.half else (1e-4, 1e-4)