mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-11-07 21:13:41 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
63243d7e2d768514b61b5cd5d357200ecab7e208
StarPilot/tinygrad_repo/test/runtime/test_kernel_opts.py
T
Daniel Koepping 68984ff07b bump tinygrad (#39085)
2026-10-20 03:55:45 +00:00

403 lines
22 KiB
Python
Raw Blame History

import unittest
from tinygrad import Device, Tensor, dtypes
from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
from tinygrad.uop.ops import AxisType
from tinygrad.codegen.opt.postrange import Scheduler
# TODO: write a clean version of this
from test.runtime.test_linearizer import helper_linearizer_opt
class TestKernelOpts(unittest.TestCase):
def test_opt_without_axis(self):
ast = Tensor.empty(32, 32).sum(1).schedule_linear().src[-1].src[0]
for opt in [Opt(OptOps.SPLIT, None, (2, AxisType.UPCAST)), Opt(OptOps.PADTO, None, 32), Opt(OptOps.SWAP, None, 1)]:
with self.assertRaises(KernelOptError): Scheduler(ast, Device[Device.DEFAULT].renderer).apply_opt(opt)
def test_swap_invalid_arg(self):
ast = (Tensor.empty(32, 32) + 1).schedule_linear().src[-1].src[0]
for opt in [Opt(OptOps.SWAP, 0, None), Opt(OptOps.SWAP, 0, 99), Opt(OptOps.SWAP, 0, -1), Opt(OptOps.SWAP, 0, True)]:
with self.assertRaisesRegex(KernelOptError, "invalid swap axis"): Scheduler(ast, Device[Device.DEFAULT].renderer).apply_opt(opt)
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipIf(Device.DEFAULT == "AMD", "TODO: segfaults on MOCKKFD with AMD:LLVM")
def test_local_and_grouped_reduce(self):
N = 128
Tensor.manual_seed(1882)
a = Tensor.rand(4, 4, N, N)
b = Tensor.rand(4, 4, N)
r = (b.sqrt() + ((a+1).sum(axis=3).exp()))
helper_linearizer_opt(r, [
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))],
[Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL))], # Checking how it works with locals
[Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 1, (64, AxisType.LOCAL, True))], # Checking how it works with grouped reduce
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 0, (32, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL, True))],
# Checking how it works with locals + grouped reduce
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (64, AxisType.LOCAL, True))],
# Checking how it works with locals + grouped reduce + upcasts
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 4, (4, AxisType.UPCAST))],
# many local + many group
[Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL)),
Opt(OptOps.SPLIT, 3, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (2, AxisType.LOCAL))],
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))] * 4,
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL)),
Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (2, AxisType.LOCAL)),
Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 6, (2, AxisType.LOCAL)),
Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 8, (2, AxisType.LOCAL))],
])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_grouped_reduce_with_local_upcast_padto(self):
Tensor.manual_seed(7)
a = Tensor.rand(7, 11, 13)
helper_linearizer_opt(a.sum((1, 2)) + a.max((1, 2)), [
[Opt(OptOps.SPLIT, 0, (0, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (11, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 2, (0, AxisType.LOCAL, True)), Opt(OptOps.PADTO, 2, 32)],
])
b = Tensor.rand(17, 19)
helper_linearizer_opt(b.flip(0).pad(((2, 3), (0, 0))).sum(0), [
[Opt(OptOps.SPLIT, 1, (0, AxisType.LOCAL, True)), Opt(OptOps.PADTO, 0, 8),
Opt(OptOps.SPLIT, 0, (12, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (0, AxisType.LOCAL))],
])
x, w = Tensor.rand(1, 3, 15, 15), Tensor.rand(4, 3, 3, 3)
helper_linearizer_opt(x.conv2d(w, padding=1, stride=2), [
[Opt(OptOps.SPLIT, 5, (0, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 1, (0, AxisType.LOCAL))],
])
def test_unrolled_padded_cumsum(self):
Tensor.manual_seed(7)
a = Tensor.rand(13, 17)
helper_linearizer_opt(a.cumsum(1), [
[Opt(OptOps.SPLIT, 2, (0, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 0, 4)],
])
def test_upcasts(self):
N = 16
Tensor.manual_seed(1772)
a = Tensor.rand(N, N)
b = Tensor.rand(N, N)
r = (a+b).sqrt() * ((a+1).exp())
helper_linearizer_opt(r, [
[Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST))], # Checking how it works with upcasts
])
def test_full_upcast(self):
Tensor.manual_seed(1772)
a = Tensor.rand(4)
b = Tensor.rand(4)
r = (a+b).sqrt() * ((a+1).exp())
helper_linearizer_opt(r, [
[Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST))], # Checking how it works with upcasts
])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_matmul(self):
Tensor.manual_seed(1552)
# Keep the full reduction and enough output elements for every split, without a square output tile.
a = Tensor.rand(128, 128)
b = Tensor.rand(128, 32)
r = a@b
helper_linearizer_opt(r, [
[Opt(OptOps.SPLIT, 1, (32, AxisType.LOCAL))],
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL))],
[Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (8, AxisType.LOCAL))], # Checking how it works with locals
[Opt(OptOps.SPLIT, 2, (32, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 2, (32, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 2, (4, AxisType.UPCAST))], # Checking how it works with grouped_reduce
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (32, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (8, AxisType.LOCAL)),
Opt(OptOps.SPLIT, 4, (4, AxisType.LOCAL, True))], # Checking how it works with local+grouped_reduce
# Checking all together
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 4, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (4, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 1, (2, AxisType.UPCAST))],
# Full global upcast + local
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 4, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST))],
])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_matmul_upcast_group(self):
Tensor.manual_seed(1552)
helper_linearizer_opt(Tensor.rand(16, 128)@Tensor.rand(128, 16),
[[Opt(OptOps.SPLIT, 2, (32, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 2, (4, AxisType.UPCAST))]])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_double_reduce(self):
N = 128
Tensor.manual_seed(1552)
a = Tensor.rand(8, N, 8, N)
r = a.sum(axis=(1,3))
helper_linearizer_opt(r, [
# openCL / DEV=CL is 256 max threads
[Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL, True))], [Opt(OptOps.SPLIT, 2, (32, AxisType.LOCAL, True))],
# Checking how it works with 1 grouped_reduce.
[Opt(OptOps.SPLIT, 3, (2, AxisType.LOCAL, True))], [Opt(OptOps.SPLIT, 3, (32, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 4, (2, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 2, (16, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 4, (2, AxisType.LOCAL, True))],
[Opt(OptOps.SPLIT, 2, (4, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 4, (64, AxisType.LOCAL, True))], # Checking how it works with 2 grouped_reduces.
[Opt(OptOps.SPLIT, 2, (16, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 4, (2, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 2, (4, AxisType.UPCAST))],
# Checking how it works with 2 grouped_reduces + upcasts.
[Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 4, (32, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 4, (4, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (4, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 6, (4, AxisType.LOCAL, True))],
# Checking how it works with 2 grouped_reduces + upcasts + locals.
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (2, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 6, (32, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 5, (4, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 6, (4, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (2, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (8, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 6, (4, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 4, (4, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 5, (4, AxisType.UPCAST))], # Checking how it works with 2 grouped_reduces + upcasts + locals.
[Opt(OptOps.SPLIT, 0, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL)), Opt(OptOps.SPLIT, 4, (4, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 6, (4, AxisType.LOCAL, True)),
Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))], # No globals
])
def test_padto_matmul(self):
N = 17
Tensor.manual_seed(289)
a = Tensor.rand(N, N)
b = Tensor.rand(N, N)
helper_linearizer_opt(a@b, [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 1, 32)],
[Opt(OptOps.PADTO, 2, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.PADTO, 2, 32)],
# can optimize further post PADTO
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST)), Opt(OptOps.SPLIT, 1, (2, AxisType.UPCAST)),],
])
def test_padto_upcasted_not_ok(self):
N = 4
a = Tensor.rand(N, N)
b = Tensor.rand(N, N)
helper_linearizer_opt(a@b, [
[Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 1, (0, AxisType.UPCAST))],
[Opt(OptOps.SPLIT, 2, (0, AxisType.UPCAST))],
[Opt(OptOps.PADTO, 0, 8)],
[Opt(OptOps.PADTO, 1, 8)],
[Opt(OptOps.PADTO, 2, 8)],
])
with self.assertRaises(KernelOptError):
helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 0, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 1, 8)]])
with self.assertRaises(KernelOptError):
helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 1, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 1, 8)]])
with self.assertRaises(KernelOptError):
helper_linearizer_opt(a@b, [[Opt(OptOps.SPLIT, 2, (0, AxisType.UPCAST)), Opt(OptOps.PADTO, 2, 8)]])
def test_padto_sum_ok(self):
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)))
helper_linearizer_opt(a.sum(0), [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),],
])
helper_linearizer_opt(a.sum(1), [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),],
])
for axis in (0, 1):
helper_linearizer_opt(a.sum(), [[Opt(OptOps.PADTO, axis, 32)],])
helper_linearizer_opt(a.sum(0), [[Opt(OptOps.PADTO, axis, 32)],])
helper_linearizer_opt(b.sum(), [[Opt(OptOps.PADTO, axis, 32)],])
helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, axis, 32)],])
helper_linearizer_opt(b.sum(dtype=dtypes.bool), [[Opt(OptOps.PADTO, axis, 32)],])
# TODO: why?
if Device.DEFAULT != "WEBGPU":
helper_linearizer_opt(b.sum(0, dtype=dtypes.bool), [[Opt(OptOps.PADTO, axis, 32)],])
helper_linearizer_opt(b.sum(1, dtype=dtypes.bool), [[Opt(OptOps.PADTO, axis, 32)],])
# having unsafe ops after sum is fine
helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_padto_group_full_unroll_sum(self):
a = Tensor.ones(2, 28, 4096).realize()
out = ((a * 0.5).float().square()).sum(axis=(0, 2))
opts_to_apply = [Opt(OptOps.SPLIT, 2, (256, AxisType.LOCAL, True)), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.SPLIT, 3, (0, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 0, (7, AxisType.UPCAST))]
helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
def test_padto_unrolled_sum(self):
a = Tensor.arange(4*17, dtype=dtypes.float).reshape(4, 17).clone().realize()
for amt in (4, 0):
helper_linearizer_opt(a.sum(1), [[Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 1, (amt, AxisType.UPCAST))]])
def test_padto_unrolled_max(self):
a = (Tensor.arange(4*17, dtype=dtypes.float).reshape(4, 17) - 100).clone().realize()
for amt in (4, 0):
helper_linearizer_opt(a.max(1), [[Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 1, (amt, AxisType.UPCAST))]])
def test_padto_unrolled_upcast(self):
a = Tensor.arange(4*17, dtype=dtypes.float).reshape(4, 17).clone().realize()
helper_linearizer_opt(a.sum(1), [[Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 1, (0, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))]])
def test_padto_nested_reduce(self):
a = (Tensor.arange(2*3, dtype=dtypes.float).reshape(2, 3) + 1).clone().realize() # [[1, 2, 3], [4, 5, 6]]
# the pad gate has the outer reduce's range, the inner reduce must not resolve it with its own identity
pad_outer = [[Opt(OptOps.PADTO, 1, 4)]]
helper_linearizer_opt(a.max(1).sum(0), pad_outer, wanna_output=[[3+6]])
helper_linearizer_opt((-a).sum(1).max(0), pad_outer, wanna_output=[[-6]])
helper_linearizer_opt(a.prod(1).sum(0), pad_outer, wanna_output=[[6+120]])
# both reduce axes padded: the outer clause lifts out, the inner clause is the inner reduce's identity
helper_linearizer_opt(a.max(1).sum(0), [[Opt(OptOps.PADTO, 0, 4), Opt(OptOps.PADTO, 1, 4)]], wanna_output=[[3+6]])
def test_padto_unindexed_reduce(self):
# a repeat/cat along the reduced axis leaves a reduce range no buffer index uses
a = Tensor.arange(7*5, dtype=dtypes.float).reshape(7, 5).clone().realize()
helper_linearizer_opt(a.repeat((3, 16)).sum(1), [[Opt(OptOps.PADTO, 2, 32)]])
helper_linearizer_opt(a.repeat((1, 16)).sum(1), [[Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.PADTO, 1, 3)]])
helper_linearizer_opt(a.cat(a, dim=1).sum(1), [[Opt(OptOps.PADTO, 1, 4)]])
a = Tensor.full((7, 5), 2.0).clone().realize()
helper_linearizer_opt(a.repeat((1, 4)).prod(1), [[Opt(OptOps.PADTO, 1, 3)]])
def test_padto_reduce_identity(self):
a = Tensor.arange(7*5, dtype=dtypes.float).reshape(7, 5).clone().realize()
helper_linearizer_opt(a.sum(1), [[Opt(OptOps.PADTO, 1, 8), Opt(OptOps.PADTO, 1, 16)]])
helper_linearizer_opt((a > 0).all(1), [[Opt(OptOps.PADTO, 1, 8)]])
def test_padto_masked_reduce(self):
# a where with a defined false arm gives the padded iterations a value
a = Tensor.arange(7*17, dtype=dtypes.float).reshape(7, 17).clone().realize()
m = (Tensor.arange(7).reshape(7, 1) % 2 == 0).expand(7, 17)
helper_linearizer_opt(m.where(a, 1.0).sum(1), [[Opt(OptOps.PADTO, 1, 8)]])
helper_linearizer_opt(m.where(a, 1.0).sum(1), [[Opt(OptOps.PADTO, 1, 32)]])
helper_linearizer_opt((Tensor.arange(17).reshape(1, 17) < 5).expand(7, 17).where(a, 1.0).sum(1), [[Opt(OptOps.PADTO, 1, 8)]])
helper_linearizer_opt(m[:, :11].where(Tensor.ones(7, 11), 2.0).prod(1), [[Opt(OptOps.PADTO, 1, 8)]])
def test_padto_unrolled_prod(self):
a = (Tensor.arange(4*17, dtype=dtypes.float).reshape(4, 17) / 100 + 1).clone().realize()
helper_linearizer_opt(a.prod(1), [[Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 1, (0, AxisType.UPCAST)),
Opt(OptOps.SPLIT, 0, (2, AxisType.UPCAST))]])
def test_padto_arg(self):
a = Tensor.arange(4*17, dtype=dtypes.float).reshape(4, 17).clone().realize()
for arg in (-4, 0, 1, True):
with self.assertRaises(KernelOptError):
helper_linearizer_opt(a.sum(1), [[Opt(OptOps.PADTO, 1, arg)]])
def test_padto_sum(self):
N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp()
helper_linearizer_opt(a.exp().sum(), [[Opt(OptOps.PADTO, 0, 32)],])
helper_linearizer_opt(a.exp().sum(0), [[Opt(OptOps.PADTO, 1, 32)],])
b = a < 1
helper_linearizer_opt(b.sum(), [[Opt(OptOps.PADTO, 0, 32)],])
helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],])
def test_padto_max(self):
N = 18
# NOTE: this setup prevents 17 * 17 contiguous merged into one axis
a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100
helper_linearizer_opt(a.max(0), [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),],
])
helper_linearizer_opt(a.max(1), [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),],
])
helper_linearizer_opt(a.max(), [[Opt(OptOps.PADTO, 0, 32)],])
helper_linearizer_opt(a.max(0), [[Opt(OptOps.PADTO, 1, 32)],])
def test_padto_where(self):
Tensor.manual_seed(0)
N = 17
a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0).int()
helper_linearizer_opt(a.max(0), [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),],
])
def test_padto_where_multioutput(self):
Tensor.manual_seed(0)
N = 17
r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
a0 = r.where(1, 0).int()
a1 = r.where(2, 0).int()
helper_linearizer_opt([a0.max(0), a1.max(0)], [
[Opt(OptOps.PADTO, 0, 32)],
[Opt(OptOps.PADTO, 0, 32), Opt(OptOps.SPLIT, 0, (8, AxisType.UPCAST)),],
])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
def test_color_shapes_with_local(self):
N = 32
Tensor.manual_seed(1552)
a = Tensor.rand(N, N)
b = Tensor.rand(N, N)
r = a@b
opts_shapes = [
([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL))], [("blue",16),("blue",32),("cyan",2),("red",32)]),
([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (2, AxisType.LOCAL))],
[("blue",16),("blue",32),("cyan",2),("green",2),("red",16)]),
# check to ensure local_dims are stable for full upcast of the first reduce
([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.UPCAST))], [("blue",16),("blue",32),("cyan",2),("magenta",32)]),
# check behavior for full upcast on an existing group
([Opt(OptOps.SPLIT, 0, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (0, AxisType.LOCAL)),Opt(OptOps.SPLIT, 3, (2, AxisType.UPCAST))],
[("blue",16),("blue",32),("cyan",2),("green",16),("magenta",2)]),
([Opt(OptOps.SPLIT, 2, (2, AxisType.LOCAL)),Opt(OptOps.SPLIT, 2, (0, AxisType.UPCAST))],
[("blue",32),("blue",32),("red",16),("magenta",2)]),
]
helper_linearizer_opt(r, [x[0] for x in opts_shapes], color_sizes=[x[1] for x in opts_shapes])
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
def test_arange_opts(self):
a = Tensor.arange(128).clone()
# NOTE: arange no longer has reduce ops available for opt
helper_linearizer_opt(a, [
[Opt(op=OptOps.SPLIT, axis=0, arg=(8, AxisType.LOCAL))],
[Opt(op=OptOps.SPLIT, axis=0, arg=(8, AxisType.LOCAL)), Opt(op=OptOps.SPLIT, axis=0, arg=(0, AxisType.UPCAST))],
])
def test_top_split_non_reduce_axis(self):
# top puts the new range as the outer factor of any split
helper_linearizer_opt(Tensor.rand(64, 64).sum(1), [[Opt(OptOps.SPLIT, 0, (16, AxisType.UPCAST, True))]])
def test_double_sum_group(self):
a = Tensor.rand(4, 4, 4)
r = a.sum((1, 2)).sum()
with self.assertRaises(KernelOptError):
helper_linearizer_opt(r, [[Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL, True))],])
r = a.sum((1, 2)).sum()
with self.assertRaises(KernelOptError):
helper_linearizer_opt(r, [[Opt(OptOps.SPLIT, 1, (4, AxisType.UPCAST)), Opt(OptOps.SPLIT, 0, (16, AxisType.LOCAL, True))],])
r = a.sum((1, 2)).sum()
with self.assertRaises(KernelOptError):
helper_linearizer_opt(r, [[Opt(OptOps.SPLIT, 1, (4, AxisType.LOCAL, True)), Opt(OptOps.SPLIT, 1, (16, AxisType.LOCAL, True))],])
if __name__ == '__main__':
unittest.main()
Reference in New Issue View Git Blame Copy Permalink