fix hevc (#17477)

* hevc tests

* x

.github/workflows/platform.yml

@@ -78,7 +78,8 @@ jobs:
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
+          test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 

test/testextra/test_hevc.py

@@ -1,7 +1,9 @@
 import unittest
 
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.helpers import fetch, round_up
+from tinygrad import Tensor, Device, Variable, dtypes
+from tinygrad.helpers import DEV, fetch, round_up
+from tinygrad.engine.realize import compile_linear
+from tinygrad.uop.ops import Ops
 from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu
 from extra.hevc.decode import hevc_decode
 
@@ -63,7 +65,7 @@ class TestHevc(unittest.TestCase):
     self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
     self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
 
-  @unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
+  @unittest.skipUnless(Device.DEFAULT == "NV" and not DEV.interface.startswith("MOCK"), "real NV only")
   def test_hevc_decode(self):
     url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
     dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
@@ -83,5 +85,22 @@ class TestHevc(unittest.TestCase):
       self.assertEqual(f.dtype, dtypes.uint8)
       self.assertEqual(f.device, "NV")
 
+  @unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
+  def test_hevc_decode_compile(self):
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
+
+    opaque, frame_info, _, _, luma_w, luma_h, _ = parse_hevc_file_headers(dat)
+    offset, sz, frame_pos, max_hist, _ = frame_info[1]
+    out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+    history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
+    decoded = Tensor(dat, device="NV")[offset:offset+sz].decode_hevc_frame(
+      Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
+
+    compiled = compile_linear(decoded.linear_with_vars()[0])
+    self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
+    encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
+    self.assertEqual(len(encdec_calls), 1)
+
 if __name__ == "__main__":
   unittest.main()

tinygrad/uop/spec.py

@@ -275,7 +275,7 @@ spec_kernel_graph = PatternMatcher([
   (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
   # all calls are on various sinks
-  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM)),), allow_any_len=True), lambda: True),
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
   # after on PARAM or AFTER
   (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),),
         allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),