github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-07 22:23:55 -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
from tinygrad import Device, Tensor, dtypes
from tinygrad.helpers import CI
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
# TODO: write a clean version of this
@@ -177,9 +176,7 @@ class TestKernelOpts(unittest.TestCase):
], apply_tc=True, atol=atol, rtol=rtol)
def test_padto_matmul(self):
if (CI and Device.DEFAULT in ["AMD", "NV", "CUDA"]):
self.skipTest("super slow on CUDA and AMD because of the big grid dims")
N = 17 * 17
N = 17
Tensor.manual_seed(289)
a = 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)]])
def test_padto_sum_ok(self):
N = 18 * 18
N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
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)))
@@ -244,7 +241,7 @@ class TestKernelOpts(unittest.TestCase):
helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
def test_padto_sum_not_ok(self):
N = 18 * 18
N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp()
# 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)],])
def test_padto_max(self):
N = 18 * 18
N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one axis
a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100
@@ -282,7 +279,7 @@ class TestKernelOpts(unittest.TestCase):
def test_padto_where(self):
Tensor.manual_seed(0)
N = 17 * 17
N = 17
a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0)
helper_linearizer_opt(a.max(0), [
[Opt(OptOps.PADTO, 0, 32)],
@@ -291,7 +288,7 @@ class TestKernelOpts(unittest.TestCase):
def test_padto_where_multioutput(self):
Tensor.manual_seed(0)
N = 17 * 17
N = 17
r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
a0 = r.where(1, 0)
a1 = r.where(2, 0)