diff --git a/test/test_linearizer_dumb.py b/test/test_linearizer_dumb.py index 51ed289f56..4adca10df0 100644 --- a/test/test_linearizer_dumb.py +++ b/test/test_linearizer_dumb.py @@ -12,7 +12,6 @@ from tinygrad.engine.realize import get_program from tinygrad.renderer.ptx import PTXRenderer class TestLinearizerFailure(unittest.TestCase): - @unittest.expectedFailure @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL") def test_failure_beam_mnist(self): c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.uchar.ptr(4014080), arg=0, src=()) diff --git a/tinygrad/codegen/late/linearize.py b/tinygrad/codegen/late/linearize.py index a8ef31f3e4..d860125adf 100644 --- a/tinygrad/codegen/late/linearize.py +++ b/tinygrad/codegen/late/linearize.py @@ -222,6 +222,8 @@ def remove_blockend(x:UOp): if late_ops[i].op is Ops.BARRIER and late_ops[i+1].op is Ops.BARRIER: late_ops[i+1] = UOp(Ops.NOOP) arg = BasicBlock(parent_block.arg.lst+tuple(late_ops), tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt) return UOp(Ops.BLOCK, src=tuple(y for y in x.src if y is not parent_block)+parent_block.src, arg=arg) + # else the whole context ended by the blockend is already in this block and we can safely turn it into a block + return UOp(Ops.BLOCK, src=x.src, arg=BasicBlock(x.arg.lst, tuple([y for y in x.arg.ctx if y is not x.arg.end]), cnt=x.arg.cnt)) block_merge = PatternMatcher([ (UPat((Ops.BLOCK, Ops.BLOCKEND), name="x"), merge_block),