github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

update padto opts tests (#12679)

Browse Source
This commit is contained in:
chenyu
2025-10-14 17:00:42 -04:00
committed by GitHub
parent e8380968f2
commit d25ceffe8d
1 changed files with 6 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
test/opt/test_kernel_opts.py
View File

@@ -1,6 +1,5 @@
import unittest import unittest
from tinygrad import Device, Tensor, dtypes from tinygrad import Device, Tensor, dtypes
from tinygrad.helpers import CI
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
# TODO: write a clean version of this # TODO: write a clean version of this
@@ -177,9 +176,7 @@ class TestKernelOpts(unittest.TestCase):
], apply_tc=True, atol=atol, rtol=rtol) ], apply_tc=True, atol=atol, rtol=rtol)
def test_padto_matmul(self): def test_padto_matmul(self):
if (CI and Device.DEFAULT in ["AMD", "NV", "CUDA"]): N = 17
self.skipTest("super slow on CUDA and AMD because of the big grid dims")
N = 17 * 17
Tensor.manual_seed(289) Tensor.manual_seed(289)
a = Tensor.rand(N, N) a = Tensor.rand(N, N)
b = Tensor.rand(N, N) b = Tensor.rand(N, N)
@@ -213,7 +210,7 @@ class TestKernelOpts(unittest.TestCase):
helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]]) helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])
def test_padto_sum_ok(self): def test_padto_sum_ok(self):
N = 18 * 18 N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one dimension # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
a = Tensor.rand(N, N).realize().shrink(((0, 17), (0, 17))) * 100 a = Tensor.rand(N, N).realize().shrink(((0, 17), (0, 17))) * 100
b = (Tensor.rand(N, N) < 0.5).realize().shrink(((0, 17), (0, 17))) b = (Tensor.rand(N, N) < 0.5).realize().shrink(((0, 17), (0, 17)))
@@ -244,7 +241,7 @@ class TestKernelOpts(unittest.TestCase):
helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],]) helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
def test_padto_sum_not_ok(self): def test_padto_sum_not_ok(self):
N = 18 * 18 N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one dimension # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp() a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp()
# exp is not safe to pad # exp is not safe to pad
@@ -261,7 +258,7 @@ class TestKernelOpts(unittest.TestCase):
helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],]) helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],])
def test_padto_max(self): def test_padto_max(self):
N = 18 * 18 N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one axis # NOTE: this setup prevents 17 * 17 contiguous merged into one axis
a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100 a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100
@@ -282,7 +279,7 @@ class TestKernelOpts(unittest.TestCase):
def test_padto_where(self): def test_padto_where(self):
Tensor.manual_seed(0) Tensor.manual_seed(0)
N = 17 * 17 N = 17
a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0) a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0)
helper_linearizer_opt(a.max(0), [ helper_linearizer_opt(a.max(0), [
[Opt(OptOps.PADTO, 0, 32)], [Opt(OptOps.PADTO, 0, 32)],
@@ -291,7 +288,7 @@ class TestKernelOpts(unittest.TestCase):
def test_padto_where_multioutput(self): def test_padto_where_multioutput(self):
Tensor.manual_seed(0) Tensor.manual_seed(0)
N = 17 * 17 N = 17
r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1 r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
a0 = r.where(1, 0) a0 = r.where(1, 0)
a1 = r.where(2, 0) a1 = r.where(2, 0)