Blocks API: CUDA support

benchmarks/bench_blocks.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import argparse
 import queue
 import subprocess
 import threading
@@ -110,44 +111,41 @@ def _consume(frames):
         pass
 
 
-def _decode_sequential(path):
+def _make_blocks(path, device):
     demuxer = Demuxer(path)
-    decoder = PacketDecoder(demuxer)
-    converter = ColorConverter()
+    return demuxer, PacketDecoder(demuxer, device=device), ColorConverter(device=device)
+
+
+def _decode_sequential(path, device):
+    demuxer, decoder, converter = _make_blocks(path, device)
     _consume(_convert(converter, _decode(decoder, _demux(demuxer))))
 
 
-def _decode_prefetch_frames(path):
+def _decode_prefetch_frames(path, device):
     # [demux + decode] on one thread || [color-convert] on another.
-    demuxer = Demuxer(path)
-    decoder = PacketDecoder(demuxer)
-    converter = ColorConverter()
+    demuxer, decoder, converter = _make_blocks(path, device)
     frames = prefetch(_decode(decoder, _demux(demuxer)))
     _consume(_convert(converter, frames))
 
 
-def _decode_prefetch_packets(path):
+def _decode_prefetch_packets(path, device):
     # [demux] on one thread || [decode + color-convert] on another.
-    demuxer = Demuxer(path)
-    decoder = PacketDecoder(demuxer)
-    converter = ColorConverter()
+    demuxer, decoder, converter = _make_blocks(path, device)
     packets = prefetch(_demux(demuxer))
     _consume(_convert(converter, _decode(decoder, packets)))
 
 
-def _decode_prefetch_packets_and_frames(path):
+def _decode_prefetch_packets_and_frames(path, device):
     # [demux] || [decode] || [color-convert], each on its own thread.
-    demuxer = Demuxer(path)
-    decoder = PacketDecoder(demuxer)
-    converter = ColorConverter()
+    demuxer, decoder, converter = _make_blocks(path, device)
     packets = prefetch(_demux(demuxer))
     frames = prefetch(_decode(decoder, packets))
     _consume(_convert(converter, frames))
 
 
-def _decode_video_decoder(path):
+def _decode_video_decoder(path, device):
     # approximate seek mode to match the blocks
-    VideoDecoder(path, seek_mode="approximate").get_all_frames()
+    VideoDecoder(path, seek_mode="approximate", device=device).get_all_frames()
 
 
 def get_num_frames(path):
@@ -176,8 +174,15 @@ def bench(f, *args, num_exp=10, warmup=2, **kwargs):
 
 
 def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cpu", choices=("cpu", "cuda"))
+    args = parser.parse_args()
+
     path = make_video()
-    print(f"Video: {path}  ({_SOURCE} {_WIDTH}x{_HEIGHT} {_FPS}fps {_DURATION_S}s)\n")
+    print(
+        f"Video: {path}  ({_SOURCE} {_WIDTH}x{_HEIGHT} {_FPS}fps {_DURATION_S}s)  "
+        f"device={args.device}\n"
+    )
 
     methods = {
         "VideoDecoder": _decode_video_decoder,
@@ -189,7 +194,7 @@ def main():
 
     results = {}
     for name, fn in methods.items():
-        times_ns, cpu = bench(fn, path)
+        times_ns, cpu = bench(fn, path, args.device)
         results[name] = {
             "mean_ms": (times_ns / 1e6).mean().item(),
             "std_ms": (times_ns / 1e6).std().item(),