From 73c8dae60d0eaa3a27a252f7967d2dab8378109e Mon Sep 17 00:00:00 2001
From: Sieds Lykles <93992551+S-Lykles@users.noreply.github.com>
Date: Sat, 20 Sep 2025 06:29:19 +0200
Subject: [PATCH] add missing remove_blockend case (#12251)

* add missing remove_blockend case

* remove expectedFailure

* better comment
---
 test/test_linearizer_dumb.py       | 1 -
 tinygrad/codegen/late/linearize.py | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

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),