geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

remove logic to vectorize reduces (#6536)

Browse Source 
* remove logic to vectorize reduces

* fix tests
This commit is contained in:
George Hotz
2024-09-16 14:04:48 +08:00
committed by GitHub
parent 607113fcdf
commit 42ba887daa
2 changed files with 8 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_linearizer.py
+7 -6
View File
@@ -227,6 +227,7 @@ class TestLinearizer(unittest.TestCase):
@unittest.skipIf(CI and Device.DEFAULT in {"AMD"}, "AMD CI doesn't support multiple sync threads yet")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@unittest.skip("this is not supported, it worked by luck")
def test_double_reduce_multireduce(self):
Tensor.manual_seed(0)
x = Tensor.randn(8, 32, 8, 16, dtype=dtypes.float).realize()
@@ -984,11 +985,10 @@ class TestLinearizer(unittest.TestCase):
k.hand_coded_optimizations()
k.linearize()
accs = [u for u in k.uops if u.op is UOps.DEFINE_ACC]
stores = [u for u in k.uops if u.op is UOps.STORE]
# the first store is to lds and can be upcasted
assert accs[0].dtype == stores[0].src[-1].dtype == dtypes.float.vec(4)
assert stores[0].src[-1].dtype == dtypes.float.vec(4)
assert stores[0].src[0].op is UOps.DEFINE_LOCAL
# the second store is to gds with no upcasts
assert stores[1].src[2].dtype == dtypes.float
@@ -1301,7 +1301,7 @@ class TestLinearizer(unittest.TestCase):
# check that the float4 cast collapses
store_vals = [u.src[-1] for u in k.uops if u.op is UOps.STORE]
for val in store_vals:
assert val.dtype == dtypes.float.vec(4) and val.op is not UOps.VECTORIZE
assert val.dtype == dtypes.float.vec(4) # and val.op is not UOps.VECTORIZE
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -1340,7 +1340,7 @@ class TestLinearizer(unittest.TestCase):
barrier = [u for u in k.uops if u.op is UOps.BARRIER][0]
# check that the float4 cast collapses for all stores
for store in local_stores+global_stores:
assert store.src[2].dtype.count > 1 and store.src[2].op is not UOps.VECTORIZE
assert store.src[2].dtype.count > 1 # and store.src[2].op is not UOps.VECTORIZE
# # check the children's vins
# TODO: src ALU are not the same, should it?
# assert barrier.src == tuple(local_stores)
@@ -1357,7 +1357,7 @@ class TestLinearizer(unittest.TestCase):
# the float4 value stores directly in lds and we skip upcast
assert stores[0].src[-1].dtype == dtypes.float.vec(4)
assert stores[0].src[-1].op is not UOps.VECTORIZE
#assert stores[0].src[-1].op is not UOps.VECTORIZE
# the global store doesn't change
assert stores[1].src[2].dtype == dtypes.float
@@ -1627,6 +1627,7 @@ class TestFloat4(unittest.TestCase):
count = TestFloat4.count_half4(k)
assert count == expected, f"{count=}, {expected=}"
@unittest.skip("this doesn't happen anymore")
def test_float4_acc(self):
# from float32 stable diffusion red tinybox
ast = UOp(UOps.SINK, src=(
@@ -1656,7 +1657,7 @@ class TestFloat4(unittest.TestCase):
count = len([uop for uop in k.uops if uop.op is UOps.DEFINE_ACC and uop.dtype == dtypes.float.vec(4)])
assert count == expected, f"{count=}, {expected=}"
@unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@unittest.skip("this doesn't happen anymore")
def test_float2_acc(self):
# from resnet
ast = UOp(UOps.SINK, src=(
tinygrad/codegen/uopgraph.py
+1 -14
View File
@@ -5,7 +5,7 @@ from collections import defaultdict
from tinygrad.dtype import dtypes, PtrDType, ImageDType
from tinygrad.ops import UnaryOps, BinaryOps, exec_alu, UOp, UOps, END_FOR_UOP, type_verify, print_uops, identity_element
from tinygrad.ops import UPat, PatternMatcher, graph_rewrite
from tinygrad.helpers import DEBUG, getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, CI, all_same, partition
from tinygrad.helpers import DEBUG, getenv, flatten, dedup, TRANSCENDENTAL, AMX, prod, CI, partition
from tinygrad.codegen.transcendental import xexp2, xlog2, xsin, TRANSCENDENTAL_SUPPORTED_DTYPES
if TYPE_CHECKING: from tinygrad.renderer import Renderer
@@ -59,17 +59,6 @@ def fold_expanded(ex, buf):
# remove Nones for STORE
return UOp(ex.op, ex.dtype, tuple(x for x in new_srcs if x is not None), ex.arg) if len(used) else None
def vectorize_reduce(vec:UOp):
if all_same(vec.src): return None # don't REDUCE the same thing multiple times
if not all_same([(x.src[1:], x.arg) for x in vec.src]): return None # must have the same reduce ranges
if not vec.dtype or vec.dtype.scalar() not in {dtypes.float, dtypes.half}: return None # only fold float/half like this
return UOp(UOps.REDUCE, vec.dtype, (UOp(UOps.VECTORIZE, vec.dtype, tuple(x.src[0] for x in vec.src)),) + vec.src[0].src[1:], vec.src[0].arg)
def vectorize_alu(vec:UOp):
if not all_same([x.arg for x in vec.src]): return None
return UOp(vec.src[0].op, vec.dtype, tuple(UOp(UOps.VECTORIZE, vec.src[0].src[i].dtype.vec(vec.dtype.count),
tuple(x.src[i] for x in vec.src)) for i in range(len(vec.src[0].src))), vec.src[0].arg)
def fix_unfoldable_image_load(load:UOp, buf:UOp):
if not isinstance(buf.dtype, ImageDType) or load.src[1].dtype.count == 2: return None
id4 = load.src[1] % 4
@@ -84,8 +73,6 @@ def fix_unfoldable_image_load(load:UOp, buf:UOp):
float4_folding = PatternMatcher([
(UPat(UOps.EXPAND, src=UPat(UOps.LOAD, src=(UPat.var("buf"), UPat()), allow_any_len=True), name="ex"), fold_expanded),
(UPat((UOps.BARRIER, UOps.SINK), src=UPat(UOps.STORE, src=(UPat.var("buf"), UPat(), UPat()), allow_any_len=True), name="ex"), fold_expanded),
(UPat(UOps.VECTORIZE, src=UPat(UOps.REDUCE), name="vec"), vectorize_reduce),
(UPat(UOps.VECTORIZE, src=UPat((UOps.ALU, UOps.CAST, UOps.BITCAST)), name="vec"), vectorize_alu),
])
# ***** mod *****