mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-11-08 13:33:42 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
97b5237e94de4cb56d8c7714801cbaec5d07db92
StarPilot/tinygrad_repo/test/runtime/test_kernel_opts.py
T
YassineYousfi 70f3c8897d modeld: tinygrad jit copies and update tinygrad (#39187)
2026-11-02 13:38:46 -06:00

410 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
from tinygrad.helpers import DEV
# 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))],
])
@unittest.skipIf(Device.DEFAULT == "AMD" and DEV.interface.startswith("MOCK"), "AMD emulator doesn't support direct pc read and write")
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_cmpeq_sum(self):
a, b = Tensor(list(range(17))), Tensor([i%3 for i in range(17)])
# The extra lanes compare 0 == 0, but must not contribute to the reduction.
helper_linearizer_opt((a == b).sum(), [[Opt(OptOps.PADTO, 0, 32)]])
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