diff --git a/benchmarks/decoders/gpu_diagnostic.py b/benchmarks/decoders/gpu_diagnostic.py new file mode 100644 index 000000000..80dfb3b36 --- /dev/null +++ b/benchmarks/decoders/gpu_diagnostic.py @@ -0,0 +1,394 @@ +#!/usr/bin/env python3 +""" +GPU decode diagnostic - isolates each layer of the ROCm video decode stack. + +Tests each component independently to identify exactly where the failure occurs: + 1. HIP runtime + device info + 2. rocDecode library availability + 3. torchcodec CPU-only decode (baseline) + 4. torchcodec core API GPU decode (low-level) + 5. torchcodec VideoDecoder GPU decode (high-level) + +Run: + python3 gpu_diagnostic.py [video_path] +""" + +import ctypes +import os +import subprocess +import sys +import traceback +from pathlib import Path + + +def section(title): + print(f"\n{'='*70}") + print(f" {title}") + print(f"{'='*70}") + + +def test_hip_runtime(): + """Test 1: HIP runtime and device enumeration.""" + section("1. HIP Runtime & Device Info") + + import torch + print(f" torch.version.hip = {torch.version.hip}") + print(f" torch.cuda.is_available() = {torch.cuda.is_available()}") + print(f" torch.cuda.device_count() = {torch.cuda.device_count()}") + + if not torch.cuda.is_available(): + print(" FAIL: No GPU available") + return False + + for i in range(torch.cuda.device_count()): + props = torch.cuda.get_device_properties(i) + vram = getattr(props, 'total_memory', getattr(props, 'total_mem', 0)) + print(f" Device {i}: {props.name} ({props.gcnArchName}) " + f"VRAM={vram // (1024**2)}MB") + + # Test basic HIP allocation on device 0 + print("\n Testing HIP memory ops on cuda:0...", end=" ", flush=True) + try: + t = torch.zeros(1024, device="cuda:0") + t += 1 + torch.cuda.synchronize() + del t + print("OK") + except Exception as e: + print(f"FAIL: {e}") + return False + + return True + + +def test_rocdecode_library(): + """Test 2: Check if librocdecode.so is loadable and has expected symbols.""" + section("2. rocDecode Library") + + lib_names = ["librocdecode.so", "librocdecode.so.0", "librocdecode.so.1"] + lib = None + loaded_name = None + + for name in lib_names: + try: + lib = ctypes.CDLL(name) + loaded_name = name + break + except OSError: + continue + + if lib is None: + print(" FAIL: Could not load librocdecode.so") + print(" Checked:", ", ".join(lib_names)) + + # Check if the file exists anywhere + try: + result = subprocess.run( + ["find", "/", "-name", "librocdecode*", "-type", "f"], + capture_output=True, text=True, timeout=10 + ) + if result.stdout.strip(): + print(" Found files:") + for line in result.stdout.strip().split("\n"): + print(f" {line}") + else: + print(" No librocdecode files found on system") + except Exception: + pass + + # Check LD_LIBRARY_PATH + print(f" LD_LIBRARY_PATH = {os.environ.get('LD_LIBRARY_PATH', '(not set)')}") + return False + + print(f" Loaded: {loaded_name}") + + # Check for expected symbols + expected_symbols = [ + "rocDecCreateDecoder", + "rocDecDestroyDecoder", + "rocDecGetDecoderCaps", + "rocDecDecodeFrame", + "rocDecGetVideoFrame", + "rocDecGetErrorName", + "rocDecCreateVideoParser", + "rocDecParseVideoData", + "rocDecDestroyVideoParser", + ] + + missing = [] + for sym in expected_symbols: + try: + getattr(lib, sym) + except AttributeError: + missing.append(sym) + + if missing: + print(f" FAIL: Missing symbols: {missing}") + return False + + print(f" All {len(expected_symbols)} expected symbols found") + + # Try to check version + try: + result = subprocess.run( + ["dpkg", "-l", "rocdecode*"], + capture_output=True, text=True, timeout=5 + ) + for line in result.stdout.strip().split("\n"): + if "rocdecode" in line.lower(): + print(f" Package: {line.strip()}") + except Exception: + pass + + return True + + +def test_cpu_decode(video_path): + """Test 3: CPU-only decode to establish baseline.""" + section("3. CPU Decode (baseline)") + + from torchcodec.decoders import VideoDecoder + + print(f" Video: {video_path}") + print(" Creating CPU decoder...", end=" ", flush=True) + try: + decoder = VideoDecoder(str(video_path), device="cpu") + print("OK") + except Exception as e: + print(f"FAIL: {e}") + return False + + meta = decoder.metadata + print(f" Codec: {meta.codec} Resolution: {meta.width}x{meta.height}") + print(f" Frames: {meta.num_frames} FPS: {meta.average_fps:.1f}") + + print(" Decoding first frame on CPU...", end=" ", flush=True) + try: + frame = next(iter(decoder)) + print(f"OK (shape={frame.data.shape}, dtype={frame.data.dtype})") + except Exception as e: + print(f"FAIL: {e}") + return False + + print(" Decoding 10 frames on CPU...", end=" ", flush=True) + try: + decoder2 = VideoDecoder(str(video_path), device="cpu") + count = 0 + for f in decoder2: + count += 1 + if count >= 10: + break + print(f"OK ({count} frames)") + except Exception as e: + print(f"FAIL: {e}") + return False + + return True + + +def test_core_api_gpu(video_path): + """Test 4: Core API GPU decode (low-level, matches existing gpu_benchmark.py).""" + section("4. Core API GPU Decode") + + import torch + import torchcodec._core + + if not torch.cuda.is_available(): + print(" SKIP: No GPU") + return False + + print(" Creating decoder from file...", end=" ", flush=True) + try: + decoder = torchcodec._core.create_from_file(str(video_path)) + print("OK") + except Exception as e: + print(f"FAIL: {e}") + traceback.print_exc() + return False + + print(" Adding video stream on cuda:0...", end=" ", flush=True) + try: + torchcodec._core._add_video_stream( + decoder, + stream_index=-1, + device="cuda:0", + num_threads=1, + ) + print("OK") + except Exception as e: + print(f"FAIL: {e}") + traceback.print_exc() + return False + + print(" Decoding first frame via core API...", end=" ", flush=True) + try: + frame, *_ = torchcodec._core.get_next_frame(decoder) + torch.cuda.synchronize() + print(f"OK (shape={frame.shape}, device={frame.device})") + except Exception as e: + print(f"FAIL: {e}") + traceback.print_exc() + return False + + print(" Decoding 10 more frames...", end=" ", flush=True) + try: + count = 1 + for _ in range(10): + frame, *_ = torchcodec._core.get_next_frame(decoder) + count += 1 + torch.cuda.synchronize() + print(f"OK ({count} total frames)") + except Exception as e: + print(f"FAIL after {count} frames: {e}") + traceback.print_exc() + return False + + return True + + +def test_videodecoder_gpu(video_path): + """Test 5: High-level VideoDecoder GPU decode.""" + section("5. VideoDecoder GPU Decode") + + import torch + from torchcodec.decoders import VideoDecoder + + if not torch.cuda.is_available(): + print(" SKIP: No GPU") + return False + + print(" Creating VideoDecoder on cuda:0...", end=" ", flush=True) + try: + decoder = VideoDecoder( + str(video_path), + device="cuda:0", + num_ffmpeg_threads=1, + ) + print("OK") + except Exception as e: + print(f"FAIL: {e}") + traceback.print_exc() + return False + + print(" Decoding first frame...", end=" ", flush=True) + try: + frame = next(iter(decoder)) + torch.cuda.synchronize() + print(f"OK (shape={frame.data.shape}, device={frame.data.device})") + except Exception as e: + print(f"FAIL: {e}") + traceback.print_exc() + return False + + print(" Decoding 10 frames...", end=" ", flush=True) + try: + decoder2 = VideoDecoder( + str(video_path), + device="cuda:0", + num_ffmpeg_threads=1, + ) + count = 0 + for f in decoder2: + count += 1 + if count >= 10: + break + torch.cuda.synchronize() + print(f"OK ({count} frames)") + except Exception as e: + print(f"FAIL after {count} frames: {e}") + traceback.print_exc() + return False + + return True + + +def test_drm_info(): + """Bonus: DRM render node info for multi-GPU debugging.""" + section("DRM / GPU Topology Info") + + # List DRI render nodes + dri_path = Path("/dev/dri") + if dri_path.exists(): + nodes = sorted(dri_path.glob("renderD*")) + print(f" DRI render nodes: {len(nodes)}") + for n in nodes: + print(f" {n}") + else: + print(" /dev/dri not found") + + # Check KFD + kfd = Path("/dev/kfd") + print(f" /dev/kfd exists: {kfd.exists()}") + + # rocminfo GPU topology + try: + result = subprocess.run( + ["rocminfo"], capture_output=True, text=True, timeout=10 + ) + agents = [] + current_agent = {} + for line in result.stdout.splitlines(): + if "Agent " in line and "Agent " in line: + if current_agent: + agents.append(current_agent) + current_agent = {"header": line.strip()} + if "Name:" in line: + current_agent["name"] = line.split(":")[-1].strip() + if "Device Type:" in line: + current_agent["type"] = line.split(":")[-1].strip() + if current_agent: + agents.append(current_agent) + + gpu_agents = [a for a in agents if a.get("type", "").strip() == "GPU"] + print(f"\n rocminfo: {len(gpu_agents)} GPU agents") + for i, a in enumerate(gpu_agents): + print(f" GPU {i}: {a.get('name', 'unknown')}") + except Exception as e: + print(f" rocminfo: {e}") + + # HIP_VISIBLE_DEVICES + print(f"\n HIP_VISIBLE_DEVICES = {os.environ.get('HIP_VISIBLE_DEVICES', '(not set)')}") + print(f" ROCR_VISIBLE_DEVICES = {os.environ.get('ROCR_VISIBLE_DEVICES', '(not set)')}") + print(f" CUDA_VISIBLE_DEVICES = {os.environ.get('CUDA_VISIBLE_DEVICES', '(not set)')}") + + +def main(): + # Find test video + test_resources = Path("/workspace/test_resources") + repo_resources = Path(__file__).resolve().parent.parent.parent / "test" / "resources" + + if len(sys.argv) > 1: + video_path = Path(sys.argv[1]) + elif (test_resources / "nasa_13013.mp4").exists(): + video_path = test_resources / "nasa_13013.mp4" + elif (repo_resources / "nasa_13013.mp4").exists(): + video_path = repo_resources / "nasa_13013.mp4" + else: + print("ERROR: No test video found. Pass a video path as argument.") + sys.exit(1) + + print(f"Using video: {video_path}") + + results = {} + + # Run diagnostics in order + test_drm_info() + results["hip_runtime"] = test_hip_runtime() + results["rocdecode_lib"] = test_rocdecode_library() + results["cpu_decode"] = test_cpu_decode(video_path) + results["core_api_gpu"] = test_core_api_gpu(video_path) + results["videodecoder_gpu"] = test_videodecoder_gpu(video_path) + + # Summary + section("SUMMARY") + for test, passed in results.items(): + status = "PASS" if passed else "FAIL" + print(f" {test:<25} {status}") + + all_passed = all(results.values()) + print(f"\n Overall: {'ALL PASSED' if all_passed else 'FAILURES DETECTED'}") + sys.exit(0 if all_passed else 1) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/decoders/rocm_vs_cuda_benchmark.py b/benchmarks/decoders/rocm_vs_cuda_benchmark.py new file mode 100644 index 000000000..8e4c29f25 --- /dev/null +++ b/benchmarks/decoders/rocm_vs_cuda_benchmark.py @@ -0,0 +1,798 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +""" +Cross-platform GPU benchmark for TorchCodec: ROCm vs CUDA + +Runs identical workloads on both AMD (rocDecode/VCN) and NVIDIA (NVDEC) +hardware decoders, producing comparable JSON output for head-to-head +analysis. + +Usage: + # Auto-detect GPU and run all benchmarks + python rocm_vs_cuda_benchmark.py + + # Specify output file + python rocm_vs_cuda_benchmark.py --output results.json + + # Quick mode (fewer iterations, faster turnaround) + python rocm_vs_cuda_benchmark.py --quick + + # Test specific codecs only + python rocm_vs_cuda_benchmark.py --codecs h264,h265 + + # Include multi-threaded decode benchmark + python rocm_vs_cuda_benchmark.py --threads 4 +""" + +import argparse +import json +import os +import platform +import subprocess +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass, field +from pathlib import Path + +import torch +import torch.utils.benchmark as benchmark + +import torchcodec +from torchcodec.decoders import VideoDecoder + + +# --------------------------------------------------------------------------- +# Test video management +# --------------------------------------------------------------------------- +TEST_RESOURCES = Path(__file__).resolve().parent.parent.parent / "test" / "resources" +# Docker containers copy test resources to /workspace/test_resources +DOCKER_TEST_RESOURCES = Path("/workspace/test_resources") + + +def _find_test_resources() -> Path: + """Find test resources directory (works both in-repo and in Docker).""" + if TEST_RESOURCES.exists(): + return TEST_RESOURCES + if DOCKER_TEST_RESOURCES.exists(): + return DOCKER_TEST_RESOURCES + return TEST_RESOURCES # fallback, will fail gracefully later + + +# Videos bundled with the repo, keyed by codec +_res = _find_test_resources() +BUNDLED_VIDEOS = { + "h264": _res / "nasa_13013.mp4", + "h265": _res / "h265_video.mp4", + "av1": _res / "av1_video.mkv", +} + +# Synthetic video specs for generation when bundled videos are insufficient +SYNTHETIC_SPECS = [ + { + "name": "synthetic_1080p_h264_30s", + "resolution": "1920x1080", + "codec": "libx264", + "duration": 30, + "fps": 30, + "gop": 30, + "pix_fmt": "yuv420p", + }, + { + "name": "synthetic_1080p_h264_120s", + "resolution": "1920x1080", + "codec": "libx264", + "duration": 120, + "fps": 60, + "gop": 60, + "pix_fmt": "yuv420p", + }, + { + "name": "synthetic_4k_h264_10s", + "resolution": "3840x2160", + "codec": "libx264", + "duration": 10, + "fps": 30, + "gop": 30, + "pix_fmt": "yuv420p", + }, +] + + +def generate_synthetic_video(spec: dict, output_dir: Path) -> Path: + """Generate a synthetic test video using ffmpeg.""" + outfile = output_dir / f"{spec['name']}.mp4" + if outfile.exists(): + return outfile + + cmd = [ + "ffmpeg", "-y", + "-f", "lavfi", + "-i", f"mandelbrot=s={spec['resolution']}", + "-t", str(spec["duration"]), + "-c:v", spec["codec"], + "-r", str(spec["fps"]), + "-g", str(spec["gop"]), + "-pix_fmt", spec["pix_fmt"], + str(outfile), + ] + try: + subprocess.check_call(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + except (subprocess.CalledProcessError, FileNotFoundError): + print(f"Warning: could not generate {outfile.name} (ffmpeg not available?)") + return None + return outfile + + +# --------------------------------------------------------------------------- +# GPU detection +# --------------------------------------------------------------------------- +@dataclass +class GPUInfo: + vendor: str # "AMD" or "NVIDIA" + name: str + architecture: str + vram_mb: int + driver_version: str + runtime_version: str + device_count: int + + +def detect_gpu() -> GPUInfo: + """Detect GPU vendor and capabilities.""" + if not torch.cuda.is_available(): + raise RuntimeError("No GPU available. torch.cuda.is_available() returned False.") + + props = torch.cuda.get_device_properties(0) + device_count = torch.cuda.device_count() + + # Check if this is ROCm (HIP) + if hasattr(torch.version, "hip") and torch.version.hip is not None: + vendor = "AMD" + runtime_version = torch.version.hip + + # Try to get architecture from rocminfo + arch = "unknown" + try: + out = subprocess.check_output( + ["rocminfo"], stderr=subprocess.DEVNULL, text=True + ) + for line in out.splitlines(): + if "gfx" in line.lower() and "name:" in line.lower(): + arch = line.split(":")[-1].strip() + break + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + # Try to get driver version + driver = "unknown" + try: + out = subprocess.check_output( + ["rocm-smi", "--showdriverversion"], + stderr=subprocess.DEVNULL, text=True + ) + for line in out.splitlines(): + if "driver" in line.lower() and "version" in line.lower(): + driver = line.split(":")[-1].strip() + break + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + else: + vendor = "NVIDIA" + runtime_version = torch.version.cuda + arch = f"sm_{props.major}{props.minor}" + + driver = "unknown" + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], + stderr=subprocess.DEVNULL, text=True + ) + driver = out.strip().split("\n")[0] + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + return GPUInfo( + vendor=vendor, + name=props.name, + architecture=arch, + vram_mb=getattr(props, 'total_memory', getattr(props, 'total_mem', 0)) // (1024 * 1024), + driver_version=driver, + runtime_version=runtime_version, + device_count=device_count, + ) + + +# --------------------------------------------------------------------------- +# Benchmark data structures +# --------------------------------------------------------------------------- +@dataclass +class BenchmarkResult: + test_name: str + codec: str + video_file: str + video_resolution: str + video_duration_s: float + video_fps: float + num_frames_decoded: int + device: str # "gpu" or "cpu" + gpu_vendor: str + gpu_name: str + decode_fps_median: float + decode_fps_p25: float + decode_fps_p75: float + time_median_s: float + time_iqr_s: float + seek_mode: str = "exact" + resize: str = "none" + num_threads: int = 1 + extra: dict = field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# Benchmark workloads +# --------------------------------------------------------------------------- +def get_video_info(video_path: str) -> dict: + """Extract metadata from a video file.""" + decoder = VideoDecoder(video_path) + meta = decoder.metadata + return { + "codec": meta.codec, + "width": meta.width, + "height": meta.height, + "duration_s": meta.duration_seconds, + "fps": meta.average_fps, + "num_frames": meta.num_frames, + } + + +def _gpu_sync(): + """Synchronize GPU after decode operations to avoid async errors.""" + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def bench_sequential_decode( + video_path: str, + device: str, + num_frames: int, + seek_mode: str = "approximate", + min_run_time: float = 10.0, +) -> benchmark.Measurement: + """Benchmark sequential frame decoding from the start.""" + def _decode(): + decoder = VideoDecoder( + video_path, + device=device, + seek_mode=seek_mode, + num_ffmpeg_threads=1 if device != "cpu" else 0, + ) + count = 0 + for frame in decoder: + count += 1 + if count >= num_frames: + break + _gpu_sync() + return count + + t = benchmark.Timer( + stmt="_decode()", + globals={"_decode": _decode}, + label="sequential_decode", + sub_label=f"device={device}", + description=f"first {num_frames} frames", + ) + return t.blocked_autorange(min_run_time=min_run_time) + + +def bench_full_video_decode( + video_path: str, + device: str, + seek_mode: str = "approximate", + min_run_time: float = 10.0, +) -> tuple[benchmark.Measurement, int]: + """Benchmark decoding an entire video.""" + frame_count_holder = [0] + + def _decode(): + decoder = VideoDecoder( + video_path, + device=device, + seek_mode=seek_mode, + num_ffmpeg_threads=1 if device != "cpu" else 0, + ) + count = 0 + for frame in decoder: + count += 1 + frame_count_holder[0] = count + _gpu_sync() + return count + + # Warm up once to get frame count + _decode() + frame_count = frame_count_holder[0] + + t = benchmark.Timer( + stmt="_decode()", + globals={"_decode": _decode}, + label="full_video_decode", + sub_label=f"device={device}", + description=f"all {frame_count} frames", + ) + return t.blocked_autorange(min_run_time=min_run_time), frame_count + + +def bench_random_seek_decode( + video_path: str, + device: str, + num_seeks: int = 50, + min_run_time: float = 10.0, +) -> benchmark.Measurement: + """Benchmark random seek + decode operations.""" + meta = VideoDecoder(video_path).metadata + duration = meta.duration_seconds + + # Generate deterministic random seek points + torch.manual_seed(42) + pts_list = (torch.rand(num_seeks) * duration).tolist() + + def _seek_decode(): + decoder = VideoDecoder( + video_path, + device=device, + seek_mode="exact", + num_ffmpeg_threads=1 if device != "cpu" else 0, + ) + result = decoder.get_frames_played_at(pts_list) + _gpu_sync() + return result + + t = benchmark.Timer( + stmt="_seek_decode()", + globals={"_seek_decode": _seek_decode}, + label="random_seek_decode", + sub_label=f"device={device}", + description=f"{num_seeks} random seeks", + ) + return t.blocked_autorange(min_run_time=min_run_time) + + +def bench_decode_and_resize( + video_path: str, + device: str, + num_frames: int = 100, + resize_h: int = 256, + resize_w: int = 256, + min_run_time: float = 10.0, +) -> benchmark.Measurement: + """Benchmark decode + resize pipeline (common in training).""" + import torchvision.transforms.v2.functional as F + + def _decode_resize(): + decoder = VideoDecoder( + video_path, + device=device, + seek_mode="approximate", + num_ffmpeg_threads=1 if device != "cpu" else 0, + ) + count = 0 + for frame in decoder: + resized = F.resize(frame.data, (resize_h, resize_w)) + count += 1 + if count >= num_frames: + break + return count + + t = benchmark.Timer( + stmt="_decode_resize()", + globals={"_decode_resize": _decode_resize}, + label="decode_and_resize", + sub_label=f"device={device}", + description=f"{num_frames} frames -> {resize_h}x{resize_w}", + ) + return t.blocked_autorange(min_run_time=min_run_time) + + +def bench_multithreaded_decode( + video_path: str, + device: str, + num_videos: int = 10, + num_threads: int = 4, + num_gpus: int = 1, + min_run_time: float = 10.0, +) -> tuple[benchmark.Measurement, int]: + """Benchmark concurrent video decoding across threads. + + When num_gpus > 1, distributes work across multiple GPUs to test + multi-GPU decode throughput. + """ + frame_count_holder = [0] + + def _decode_one(dev): + decoder = VideoDecoder( + video_path, + device=dev, + seek_mode="approximate", + num_ffmpeg_threads=1 if dev != "cpu" else 0, + ) + count = 0 + for frame in decoder: + count += 1 + return count + + def _decode_all(): + total = 0 + with ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [] + for i in range(num_videos): + dev = f"cuda:{i % num_gpus}" if device != "cpu" else "cpu" + futures.append(executor.submit(_decode_one, dev)) + for f in futures: + total += f.result() + frame_count_holder[0] = total + _gpu_sync() + return total + + # Warm up + _decode_all() + total_frames = frame_count_holder[0] + + gpu_label = f"{num_gpus} GPU{'s' if num_gpus > 1 else ''}" + t = benchmark.Timer( + stmt="_decode_all()", + globals={"_decode_all": _decode_all}, + label="multithreaded_decode", + sub_label=f"device={device}", + description=f"{num_videos} videos x {num_threads} threads x {gpu_label}", + ) + return t.blocked_autorange(min_run_time=min_run_time), total_frames + + +def measurement_to_result( + measurement: benchmark.Measurement, + test_name: str, + codec: str, + video_path: str, + video_info: dict, + num_frames: int, + device: str, + gpu_info: GPUInfo, + **extra_fields, +) -> BenchmarkResult: + """Convert a torch.benchmark.Measurement into our result format.""" + fps_median = num_frames / measurement.median + fps_p25 = num_frames / measurement._p75 # inverted: slower time = lower fps + fps_p75 = num_frames / measurement._p25 + + return BenchmarkResult( + test_name=test_name, + codec=codec, + video_file=Path(video_path).name, + video_resolution=f"{video_info['width']}x{video_info['height']}", + video_duration_s=video_info["duration_s"], + video_fps=video_info["fps"], + num_frames_decoded=num_frames, + device="gpu" if device != "cpu" else "cpu", + gpu_vendor=gpu_info.vendor, + gpu_name=gpu_info.name, + decode_fps_median=round(fps_median, 2), + decode_fps_p25=round(fps_p25, 2), + decode_fps_p75=round(fps_p75, 2), + time_median_s=round(measurement.median, 4), + time_iqr_s=round(measurement.iqr, 4), + **extra_fields, + ) + + +# --------------------------------------------------------------------------- +# Main benchmark runner +# --------------------------------------------------------------------------- +def run_all_benchmarks(args) -> dict: + gpu_info = detect_gpu() + print(f"\n{'='*70}") + print(f"TorchCodec GPU Benchmark") + print(f"{'='*70}") + print(f"GPU Vendor: {gpu_info.vendor}") + print(f"GPU Name: {gpu_info.name}") + print(f"Architecture: {gpu_info.architecture}") + print(f"VRAM: {gpu_info.vram_mb} MB") + print(f"Driver: {gpu_info.driver_version}") + print(f"Runtime: {gpu_info.runtime_version}") + print(f"Device Count: {gpu_info.device_count}") + print(f"PyTorch: {torch.__version__}") + print(f"TorchCodec: {torchcodec.__version__}") + print(f"Python: {platform.python_version()}") + print(f"{'='*70}\n") + + min_run_time = 5.0 if args.quick else 15.0 + codecs_to_test = args.codecs.split(",") if args.codecs else ["h264"] + + # Collect videos to benchmark + videos = {} + for codec in codecs_to_test: + if codec in BUNDLED_VIDEOS and BUNDLED_VIDEOS[codec].exists(): + videos[codec] = str(BUNDLED_VIDEOS[codec]) + else: + print(f"Warning: no test video for codec '{codec}', skipping") + + # Generate synthetic videos if requested + if args.synthetic: + synth_dir = Path(args.synthetic_dir) + synth_dir.mkdir(parents=True, exist_ok=True) + for spec in SYNTHETIC_SPECS: + path = generate_synthetic_video(spec, synth_dir) + if path: + videos[f"synthetic_{spec['name']}"] = str(path) + + if not videos: + print("ERROR: No test videos found. Ensure test/resources/ contains videos.") + sys.exit(1) + + devices = ["cuda:0", "cpu"] + all_results = [] + torch_results = [] # for Compare table + + # Smoke test: verify GPU decode works before running full benchmarks + if torch.cuda.is_available(): + first_video = next(iter(videos.values())) + print("Smoke test: decoding 1 frame on GPU...", end=" ", flush=True) + try: + test_decoder = VideoDecoder( + first_video, + device="cuda:0", + num_ffmpeg_threads=1, + ) + test_frame = next(iter(test_decoder)) + torch.cuda.synchronize() + del test_decoder + print(f"OK (shape={test_frame.data.shape})") + except Exception as e: + print(f"FAILED: {e}") + print("GPU decode not working, will only run CPU benchmarks.") + devices = ["cpu"] + + for codec_label, video_path in videos.items(): + video_info = get_video_info(video_path) + print(f"\n--- Video: {Path(video_path).name} ---") + print(f" Codec: {video_info['codec']} Resolution: {video_info['width']}x{video_info['height']}") + print(f" Duration: {video_info['duration_s']:.1f}s FPS: {video_info['fps']:.1f} Frames: {video_info['num_frames']}") + + for device in devices: + if device != "cpu" and not torch.cuda.is_available(): + continue + + device_label = f"{gpu_info.vendor} GPU" if device != "cpu" else "CPU" + print(f"\n [{device_label}] ({device})") + + # 1. Sequential decode (first N frames) + num_seq = min(200, video_info["num_frames"]) + print(f" Sequential decode ({num_seq} frames)...", end=" ", flush=True) + try: + m = bench_sequential_decode(video_path, device, num_seq, min_run_time=min_run_time) + torch_results.append(m) + r = measurement_to_result( + m, "sequential_decode", codec_label, video_path, + video_info, num_seq, device, gpu_info, + seek_mode="approximate", + ) + all_results.append(r) + print(f"{r.decode_fps_median:.1f} FPS (median)") + except Exception as e: + print(f"FAILED ({e})") + + # 2. Full video decode + print(f" Full video decode...", end=" ", flush=True) + try: + m, frame_count = bench_full_video_decode( + video_path, device, min_run_time=min_run_time + ) + torch_results.append(m) + r = measurement_to_result( + m, "full_video_decode", codec_label, video_path, + video_info, frame_count, device, gpu_info, + seek_mode="approximate", + ) + all_results.append(r) + print(f"{r.decode_fps_median:.1f} FPS (median), {frame_count} frames") + except Exception as e: + print(f"FAILED ({e})") + + # 3. Random seek + decode + num_seeks = 20 if args.quick else 50 + print(f" Random seek+decode ({num_seeks} seeks)...", end=" ", flush=True) + try: + m = bench_random_seek_decode(video_path, device, num_seeks, min_run_time=min_run_time) + torch_results.append(m) + r = measurement_to_result( + m, "random_seek_decode", codec_label, video_path, + video_info, num_seeks, device, gpu_info, + seek_mode="exact", + ) + all_results.append(r) + print(f"{r.decode_fps_median:.1f} FPS (median)") + except Exception as e: + print(f"FAILED ({e})") + + # 4. Decode + resize (training-like pipeline) + num_resize = min(100, video_info["num_frames"]) + print(f" Decode+resize ({num_resize} frames -> 256x256)...", end=" ", flush=True) + try: + m = bench_decode_and_resize( + video_path, device, num_resize, min_run_time=min_run_time + ) + torch_results.append(m) + r = measurement_to_result( + m, "decode_and_resize", codec_label, video_path, + video_info, num_resize, device, gpu_info, + resize="256x256", + ) + all_results.append(r) + print(f"{r.decode_fps_median:.1f} FPS (median)") + except Exception as e: + print(f"skipped ({e})") + + # 5. Multi-threaded decode (GPU only) + if device != "cpu" and args.threads > 1: + device_count = torch.cuda.device_count() + num_vids = args.threads * 2 + + # Single-GPU multi-threaded + print(f" Multi-threaded decode ({num_vids} videos, {args.threads} threads, 1 GPU)...", end=" ", flush=True) + try: + m, total_frames = bench_multithreaded_decode( + video_path, device, num_vids, args.threads, + num_gpus=1, + min_run_time=min_run_time, + ) + torch_results.append(m) + r = measurement_to_result( + m, "multithreaded_1gpu", codec_label, video_path, + video_info, total_frames, device, gpu_info, + num_threads=args.threads, + ) + all_results.append(r) + print(f"{r.decode_fps_median:.1f} FPS (median)") + except Exception as e: + print(f"FAILED ({e})") + + # Multi-GPU multi-threaded (if more than 1 GPU available) + if device_count >= 2: + num_gpus = min(device_count, args.threads) + print(f" Multi-threaded decode ({num_vids} videos, {args.threads} threads, {num_gpus} GPUs)...", end=" ", flush=True) + try: + m, total_frames = bench_multithreaded_decode( + video_path, device, num_vids, args.threads, + num_gpus=num_gpus, + min_run_time=min_run_time, + ) + torch_results.append(m) + r = measurement_to_result( + m, f"multithreaded_{num_gpus}gpu", codec_label, video_path, + video_info, total_frames, device, gpu_info, + num_threads=args.threads, + ) + all_results.append(r) + print(f"{r.decode_fps_median:.1f} FPS (median)") + except Exception as e: + print(f"FAILED ({e})") + + # Print comparison table + if torch_results: + print(f"\n{'='*70}") + print("PyTorch Benchmark Comparison Table") + print(f"{'='*70}") + compare = benchmark.Compare(torch_results) + compare.print() + + # Build output + output = { + "gpu_info": asdict(gpu_info), + "system_info": { + "cpu_count": os.cpu_count(), + "platform": platform.system(), + "machine": platform.machine(), + "python_version": platform.python_version(), + "torch_version": torch.__version__, + "torchcodec_version": torchcodec.__version__, + }, + "benchmark_config": { + "min_run_time_s": min_run_time, + "quick_mode": args.quick, + "codecs_tested": codecs_to_test, + }, + "results": [asdict(r) for r in all_results], + } + + return output + + +def print_summary(output: dict): + """Print a concise summary comparing GPU vs CPU performance.""" + results = output["results"] + gpu_info = output["gpu_info"] + + print(f"\n{'='*70}") + print(f"SUMMARY: {gpu_info['vendor']} {gpu_info['name']}") + print(f"{'='*70}") + print(f"{'Test':<30} {'GPU FPS':>10} {'CPU FPS':>10} {'Speedup':>10}") + print(f"{'-'*60}") + + # Group by (test_name, codec, video_file) + from collections import defaultdict + groups = defaultdict(dict) + for r in results: + key = (r["test_name"], r["codec"], r["video_file"]) + groups[key][r["device"]] = r["decode_fps_median"] + + for (test_name, codec, video_file), devs in sorted(groups.items()): + gpu_fps = devs.get("gpu", 0) + cpu_fps = devs.get("cpu", 0) + speedup = gpu_fps / cpu_fps if cpu_fps > 0 else float("inf") + label = f"{test_name} ({codec})" + print(f"{label:<30} {gpu_fps:>10.1f} {cpu_fps:>10.1f} {speedup:>9.1f}x") + + print() + + +def main(): + parser = argparse.ArgumentParser( + description="TorchCodec GPU Benchmark: ROCm vs CUDA comparison" + ) + parser.add_argument( + "--output", "-o", + type=str, + default=None, + help="Output JSON file path. Default: benchmark_results_{vendor}.json", + ) + parser.add_argument( + "--codecs", + type=str, + default="h264", + help="Comma-separated codecs to test: h264,h265,av1 (default: h264)", + ) + parser.add_argument( + "--quick", + action="store_true", + help="Quick mode: fewer iterations, faster turnaround", + ) + parser.add_argument( + "--threads", + type=int, + default=1, + help="Number of threads for multi-threaded decode benchmark (default: 1, disabled)", + ) + parser.add_argument( + "--synthetic", + action="store_true", + help="Also generate and benchmark synthetic videos (1080p, 4K)", + ) + parser.add_argument( + "--synthetic-dir", + type=str, + default="/tmp/torchcodec_benchmark_videos", + help="Directory for synthetic test videos", + ) + + args = parser.parse_args() + output = run_all_benchmarks(args) + + # Print summary + print_summary(output) + + # Write JSON + vendor = output["gpu_info"]["vendor"].lower() + gpu_name = output["gpu_info"]["name"].replace(" ", "_") + output_path = args.output or f"benchmark_results_{vendor}_{gpu_name}.json" + with open(output_path, "w") as f: + json.dump(output, f, indent=2) + print(f"Results written to: {output_path}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/run_gpu_benchmark.sh b/benchmarks/run_gpu_benchmark.sh new file mode 100755 index 000000000..1bda2d954 --- /dev/null +++ b/benchmarks/run_gpu_benchmark.sh @@ -0,0 +1,223 @@ +#!/bin/bash +# ============================================================================= +# TorchCodec GPU Benchmark Runner +# ============================================================================= +# +# Builds a Docker image (ROCm or CUDA) and runs the GPU benchmark inside it. +# Produces JSON results for cross-platform comparison. +# +# Usage: +# # Auto-detect GPU vendor and run +# bash benchmarks/run_gpu_benchmark.sh +# +# # Force ROCm build for MI300X +# bash benchmarks/run_gpu_benchmark.sh --rocm --arch gfx942 +# +# # Force CUDA build +# bash benchmarks/run_gpu_benchmark.sh --cuda +# +# # Quick benchmark (fewer iterations) +# bash benchmarks/run_gpu_benchmark.sh --quick +# +# # Test multiple codecs +# bash benchmarks/run_gpu_benchmark.sh --codecs h264,h265,av1 +# +# # Skip Docker build (run directly if torchcodec is already installed) +# bash benchmarks/run_gpu_benchmark.sh --no-docker +# +# ============================================================================= + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)" + +# Defaults +GPU_VENDOR="" +HIP_ARCH="" +PYTORCH_ROCM_IMAGE="rocm/pytorch:rocm7.2_ubuntu24.04_py3.12_pytorch_release_2.8.0" +CUDA_VERSION="12.6.3" +CODECS="h264" +QUICK="" +THREADS=1 +SYNTHETIC="" +NO_DOCKER=false +OUTPUT_DIR="${REPO_ROOT}/benchmark_results" + +# Parse arguments +while [[ $# -gt 0 ]]; do + case $1 in + --rocm) GPU_VENDOR="AMD"; shift ;; + --cuda) GPU_VENDOR="NVIDIA"; shift ;; + --arch) HIP_ARCH="$2"; shift 2 ;; + --rocm-image) PYTORCH_ROCM_IMAGE="$2"; shift 2 ;; + --cuda-version) CUDA_VERSION="$2"; shift 2 ;; + --codecs) CODECS="$2"; shift 2 ;; + --quick) QUICK="--quick"; shift ;; + --threads) THREADS="$2"; shift 2 ;; + --synthetic) SYNTHETIC="--synthetic"; shift ;; + --no-docker) NO_DOCKER=true; shift ;; + --output-dir) OUTPUT_DIR="$2"; shift 2 ;; + -h|--help) + echo "Usage: $0 [OPTIONS]" + echo "" + echo "Options:" + echo " --rocm Force ROCm (AMD) build" + echo " --cuda Force CUDA (NVIDIA) build" + echo " --arch ARCH HIP architecture (e.g., gfx942 for MI300X)" + echo " --rocm-image IMAGE rocm/pytorch base image tag" + echo " (default: rocm/pytorch:rocm7.2_ubuntu24.04_py3.12_pytorch_release_2.8.0)" + echo " --cuda-version VER CUDA version (default: 12.6.3)" + echo " --codecs LIST Comma-separated codecs (default: h264)" + echo " --quick Quick mode (fewer iterations)" + echo " --threads N Multi-threaded decode threads (default: 1)" + echo " --synthetic Generate and test synthetic videos" + echo " --no-docker Run directly without Docker" + echo " --output-dir DIR Output directory (default: benchmark_results/)" + echo " -h, --help Show this help" + exit 0 + ;; + *) echo "Unknown option: $1"; exit 1 ;; + esac +done + +# Auto-detect GPU vendor if not specified +if [[ -z "$GPU_VENDOR" ]]; then + if command -v rocminfo &>/dev/null && rocminfo 2>/dev/null | grep -q "gfx"; then + GPU_VENDOR="AMD" + echo "[*] Auto-detected AMD GPU (ROCm)" + elif command -v nvidia-smi &>/dev/null && nvidia-smi &>/dev/null; then + GPU_VENDOR="NVIDIA" + echo "[*] Auto-detected NVIDIA GPU (CUDA)" + else + echo "ERROR: Could not detect GPU. Use --rocm or --cuda to specify." + exit 1 + fi +fi + +mkdir -p "$OUTPUT_DIR" + +# --------------------------------------------------------------------------- +# Docker mode +# --------------------------------------------------------------------------- +if [[ "$NO_DOCKER" == false ]]; then + cd "$REPO_ROOT" + + # Dockerfiles clone from git — no build context needed. + # Pipe an empty context to avoid sending the whole repo. + if [[ "$GPU_VENDOR" == "AMD" ]]; then + IMAGE_TAG="torchcodec-bench:$(echo "$PYTORCH_ROCM_IMAGE" | sed 's|rocm/pytorch:||; s|_|-|g')" + echo "[*] Building ROCm image: $IMAGE_TAG" + echo "[*] Base image: $PYTORCH_ROCM_IMAGE" + + BUILD_ARGS=( + --build-arg "PYTORCH_ROCM_IMAGE=${PYTORCH_ROCM_IMAGE}" + --build-arg "CACHEBUST=$(date +%s)" + ) + if [[ -n "$HIP_ARCH" ]]; then + BUILD_ARGS+=(--build-arg "HIP_ARCHITECTURES=${HIP_ARCH}") + fi + + docker build \ + "${BUILD_ARGS[@]}" \ + -t "$IMAGE_TAG" \ + - < "$REPO_ROOT/docker/Dockerfile.rocm" + + echo "[*] Running benchmark in ROCm container..." + docker run --rm \ + --device /dev/kfd --device /dev/dri \ + --group-add video \ + -v "$OUTPUT_DIR:/results" \ + "$IMAGE_TAG" \ + python3 /workspace/benchmarks/decoders/rocm_vs_cuda_benchmark.py \ + --codecs "$CODECS" \ + --threads "$THREADS" \ + --output "/results/benchmark_rocm.json" \ + $QUICK $SYNTHETIC + + elif [[ "$GPU_VENDOR" == "NVIDIA" ]]; then + IMAGE_TAG="torchcodec-bench:cuda${CUDA_VERSION}" + echo "[*] Building CUDA image: $IMAGE_TAG" + + docker build \ + --build-arg "CUDA_VERSION=${CUDA_VERSION}" \ + --build-arg "CACHEBUST=$(date +%s)" \ + -t "$IMAGE_TAG" \ + - < "$REPO_ROOT/docker/Dockerfile.cuda" + + echo "[*] Running benchmark in CUDA container..." + docker run --rm \ + --gpus all \ + -e NVIDIA_DRIVER_CAPABILITIES=video,compute,utility \ + -v "$OUTPUT_DIR:/results" \ + "$IMAGE_TAG" \ + python3 /workspace/benchmarks/decoders/rocm_vs_cuda_benchmark.py \ + --codecs "$CODECS" \ + --threads "$THREADS" \ + --output "/results/benchmark_cuda.json" \ + $QUICK $SYNTHETIC + fi + +# --------------------------------------------------------------------------- +# Non-Docker mode (run directly) +# --------------------------------------------------------------------------- +else + echo "[*] Running benchmark directly (no Docker)..." + cd "$REPO_ROOT" + + TIMESTAMP=$(date +%Y%m%d_%H%M%S) + VENDOR_LOWER=$(echo "$GPU_VENDOR" | tr '[:upper:]' '[:lower:]') + OUTPUT_FILE="${OUTPUT_DIR}/benchmark_${VENDOR_LOWER}_${TIMESTAMP}.json" + + python3 benchmarks/decoders/rocm_vs_cuda_benchmark.py \ + --codecs "$CODECS" \ + --threads "$THREADS" \ + --output "$OUTPUT_FILE" \ + $QUICK $SYNTHETIC +fi + +echo "" +echo "[*] Benchmark complete. Results in: $OUTPUT_DIR/" +echo "" + +# If both ROCm and CUDA results exist, print comparison +ROCM_FILE="$OUTPUT_DIR/benchmark_rocm.json" +CUDA_FILE="$OUTPUT_DIR/benchmark_cuda.json" + +if [[ -f "$ROCM_FILE" && -f "$CUDA_FILE" ]]; then + echo "============================================================" + echo " Both ROCm and CUDA results found! Quick comparison:" + echo "============================================================" + python3 -c " +import json, sys + +with open('$ROCM_FILE') as f: + rocm = json.load(f) +with open('$CUDA_FILE') as f: + cuda = json.load(f) + +print(f\" ROCm GPU: {rocm['gpu_info']['name']}\") +print(f\" CUDA GPU: {cuda['gpu_info']['name']}\") +print() +print(f\"{'Test':<30} {'ROCm FPS':>10} {'CUDA FPS':>10} {'Ratio':>8}\") +print('-' * 60) + +rocm_by_key = {} +for r in rocm['results']: + if r['device'] == 'gpu': + rocm_by_key[(r['test_name'], r['codec'])] = r['decode_fps_median'] + +cuda_by_key = {} +for r in cuda['results']: + if r['device'] == 'gpu': + cuda_by_key[(r['test_name'], r['codec'])] = r['decode_fps_median'] + +for key in sorted(set(rocm_by_key) | set(cuda_by_key)): + r_fps = rocm_by_key.get(key, 0) + c_fps = cuda_by_key.get(key, 0) + ratio = r_fps / c_fps if c_fps > 0 else float('inf') + label = f\"{key[0]} ({key[1]})\" + print(f\"{label:<30} {r_fps:>10.1f} {c_fps:>10.1f} {ratio:>7.2f}x\") +print() +" +fi diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda new file mode 100644 index 000000000..318ebfc1f --- /dev/null +++ b/docker/Dockerfile.cuda @@ -0,0 +1,117 @@ +# ============================================================================= +# TorchCodec CUDA Dockerfile +# ============================================================================= +# +# Builds torchcodec with hardware-accelerated video decoding via NVDEC +# on NVIDIA GPUs. Companion to Dockerfile.rocm for cross-platform +# GPU benchmarking. +# +# Build args: +# CUDA_VERSION - CUDA version for base image (default: 12.6.3) +# UBUNTU_VERSION - Ubuntu version (default: 22.04) +# TORCHCODEC_REPO - Git repo URL to clone (default: Cemberk fork) +# TORCHCODEC_BRANCH - Git branch to build (default: main) +# +# Examples: +# # CUDA 12.6 +# docker build -f docker/Dockerfile.cuda \ +# -t torchcodec:cuda12.6 . +# +# # CUDA 13.0 +# docker build -f docker/Dockerfile.cuda \ +# --build-arg CUDA_VERSION=13.0.0 \ +# -t torchcodec:cuda13.0 . +# +# Run (requires NVIDIA GPU + nvidia-container-toolkit): +# docker run --rm -it --gpus all \ +# -e NVIDIA_DRIVER_CAPABILITIES=video,compute,utility \ +# torchcodec:cuda12.6 python3 -c "import torchcodec; print('OK')" +# +# ============================================================================= + +ARG CUDA_VERSION=12.6.3 +ARG UBUNTU_VERSION=22.04 +FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION} + +ARG TORCHCODEC_REPO=https://github.com/Cemberk/torchcodec.git +ARG TORCHCODEC_BRANCH=main + +ENV DEBIAN_FRONTEND=noninteractive + +# ============================================================================= +# 1. System packages +# ============================================================================= +RUN apt-get update && apt-get install -y --no-install-recommends \ + libavcodec-dev \ + libavformat-dev \ + libavutil-dev \ + libavdevice-dev \ + libavfilter-dev \ + libswscale-dev \ + libswresample-dev \ + pkg-config \ + cmake \ + ninja-build \ + g++ \ + git \ + python3-dev \ + python3-pip \ + python3-venv \ + && rm -rf /var/lib/apt/lists/* + +# ============================================================================= +# 2. Python virtual environment +# ============================================================================= +ENV VIRTUAL_ENV=/opt/venv +RUN python3 -m venv $VIRTUAL_ENV +ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +# ============================================================================= +# 3. PyTorch + ecosystem (CUDA wheels) +# ============================================================================= +RUN pip install --no-cache-dir --upgrade pip setuptools wheel && \ + pip install --no-cache-dir \ + torch torchvision torchaudio \ + --index-url https://download.pytorch.org/whl/cu126 + +# ============================================================================= +# 4. TorchCodec build-time Python dependencies +# ============================================================================= +RUN pip install --no-cache-dir pybind11 numpy + +# ============================================================================= +# 5. Clone and build TorchCodec +# ============================================================================= +# ARG before RUN busts Docker cache when CACHEBUST value changes. +ARG CACHEBUST=1 +RUN git clone --depth 1 --branch ${TORCHCODEC_BRANCH} ${TORCHCODEC_REPO} /build/torchcodec +WORKDIR /build/torchcodec + +ENV ENABLE_CUDA=1 +ENV I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1 +RUN export pybind11_DIR="$(python3 -c 'import pybind11; print(pybind11.get_cmake_dir())')" && \ + pip install --no-cache-dir . --no-build-isolation + +# ============================================================================= +# 6. Stage benchmarks and test resources, then clean up +# ============================================================================= +RUN mkdir -p /workspace && \ + cp -r /build/torchcodec/benchmarks /workspace/benchmarks && \ + cp -r /build/torchcodec/test/resources /workspace/test_resources && \ + rm -rf /build + +# ============================================================================= +# 7. Smoke test +# ============================================================================= +RUN python3 -c "\ +import torchcodec; \ +print(f'torchcodec {torchcodec.__version__} installed successfully'); \ +import torch; \ +print(f'torch {torch.__version__} (CUDA: {torch.version.cuda})'); \ +" + +# ============================================================================= +# 8. Default entrypoint +# ============================================================================= +WORKDIR /workspace +CMD ["python3"] diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm new file mode 100644 index 000000000..b15eacf42 --- /dev/null +++ b/docker/Dockerfile.rocm @@ -0,0 +1,127 @@ +# ============================================================================= +# TorchCodec ROCm Dockerfile +# ============================================================================= +# +# Builds torchcodec with hardware-accelerated video decoding via rocDecode +# on AMD GPUs. Uses the official rocm/pytorch images which ship with +# PyTorch, torchvision, and torchaudio pre-installed and tested against +# the matching ROCm stack. +# +# Build args: +# PYTORCH_ROCM_IMAGE - Full rocm/pytorch image tag (default below) +# HIP_ARCHITECTURES - Semicolon-separated GPU targets (default: broad list) +# TORCHCODEC_REPO - Git repo URL to clone (default: Cemberk fork) +# TORCHCODEC_BRANCH - Git branch to build (default: main) +# +# Examples: +# # ROCm 7.2 / PyTorch 2.8 on MI300X +# docker build -f docker/Dockerfile.rocm \ +# --build-arg HIP_ARCHITECTURES="gfx942" \ +# -t torchcodec:rocm7.2 . +# +# # ROCm 6.3 / PyTorch 2.6 on MI250 +# docker build -f docker/Dockerfile.rocm \ +# --build-arg PYTORCH_ROCM_IMAGE=rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.6.0 \ +# --build-arg HIP_ARCHITECTURES="gfx90a" \ +# -t torchcodec:rocm6.3 . +# +# Run (requires GPU device access): +# docker run --rm -it \ +# --device /dev/kfd --device /dev/dri \ +# --group-add video \ +# torchcodec:rocm7.2 python3 -c "import torchcodec; print('OK')" +# +# Available base images: https://hub.docker.com/r/rocm/pytorch/tags +# ============================================================================= + +ARG PYTORCH_ROCM_IMAGE=rocm/pytorch:rocm7.2_ubuntu24.04_py3.12_pytorch_release_2.8.0 +FROM ${PYTORCH_ROCM_IMAGE} + +ARG HIP_ARCHITECTURES="" +ARG TORCHCODEC_REPO=https://github.com/Cemberk/torchcodec.git +ARG TORCHCODEC_BRANCH=main + +ENV DEBIAN_FRONTEND=noninteractive + +# ============================================================================= +# 1. System packages (FFmpeg dev libs + build tools) +# ============================================================================= +RUN apt-get update && apt-get install -y --no-install-recommends \ + libavcodec-dev \ + libavformat-dev \ + libavutil-dev \ + libavdevice-dev \ + libavfilter-dev \ + libswscale-dev \ + libswresample-dev \ + pkg-config \ + cmake \ + ninja-build \ + g++ \ + git \ + && rm -rf /var/lib/apt/lists/* + +# ============================================================================= +# 2. rocDecode (AMD hardware video decoder) +# ============================================================================= +# rocDecode provides VCN-based hardware decoding on AMD datacenter GPUs. +# The package may or may not be available in the base image's ROCm repo. +# If it's not available, torchcodec still compiles (headers are vendored) +# and will fall back to CPU decoding at runtime. +RUN apt-get update && \ + (apt-get install -y --no-install-recommends rocdecode-dev 2>/dev/null || \ + apt-get install -y --no-install-recommends rocdecode 2>/dev/null || \ + echo "rocDecode package not found - will use vendored headers and dlopen at runtime") && \ + rm -rf /var/lib/apt/lists/* + +# ============================================================================= +# 3. TorchCodec build-time Python dependencies +# ============================================================================= +# PyTorch, torchvision, torchaudio are already in the base image. +RUN pip install --no-cache-dir pybind11 numpy + +# ============================================================================= +# 4. Clone and build TorchCodec +# ============================================================================= +# ARG before RUN busts Docker cache when CACHEBUST value changes. +ARG CACHEBUST=1 +RUN git clone --depth 1 --branch ${TORCHCODEC_BRANCH} ${TORCHCODEC_REPO} /build/torchcodec +WORKDIR /build/torchcodec + +# ENABLE_ROCM: auto-detected from torch.version.hip, but set explicitly +# for clarity in build logs. +# TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR: PyTorch's ROCm cmake config +# injects Clang-only flags (e.g. -Wno-duplicate-decl-specifier) +# that GCC rejects under -Werror. +ENV ENABLE_ROCM=ON +ENV I_CONFIRM_THIS_IS_NOT_A_LICENSE_VIOLATION=1 +ENV TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR=ON +RUN if [ -n "${HIP_ARCHITECTURES}" ]; then \ + export HIP_ARCHITECTURES="${HIP_ARCHITECTURES}"; \ + fi && \ + export pybind11_DIR="$(python3 -c 'import pybind11; print(pybind11.get_cmake_dir())')" && \ + pip install --no-cache-dir . --no-build-isolation + +# ============================================================================= +# 5. Stage benchmarks and test resources, then clean up +# ============================================================================= +RUN mkdir -p /workspace && \ + cp -r /build/torchcodec/benchmarks /workspace/benchmarks && \ + cp -r /build/torchcodec/test/resources /workspace/test_resources && \ + rm -rf /build + +# ============================================================================= +# 6. Smoke test +# ============================================================================= +RUN python3 -c "\ +import torchcodec; \ +print(f'torchcodec {torchcodec.__version__} installed successfully'); \ +import torch; \ +print(f'torch {torch.__version__} (HIP: {torch.version.hip})'); \ +" + +# ============================================================================= +# 7. Default entrypoint +# ============================================================================= +WORKDIR /workspace +CMD ["python3"] diff --git a/setup.py b/setup.py index 8e335b31d..19b492efe 100644 --- a/setup.py +++ b/setup.py @@ -112,6 +112,13 @@ def _build_all_extensions_with_cmake(self): torch_dir = Path(torch.utils.cmake_prefix_path) / "Torch" cmake_build_type = os.environ.get("CMAKE_BUILD_TYPE", "Release") enable_cuda = os.environ.get("ENABLE_CUDA", "") + enable_rocm = os.environ.get("ENABLE_ROCM", "") + + # Auto-detect ROCm if torch was built with HIP and neither flag is set + if not enable_cuda and not enable_rocm: + if hasattr(torch.version, "hip") and torch.version.hip is not None: + enable_rocm = "ON" + torchcodec_disable_compile_warning_as_error = os.environ.get( "TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR", "OFF" ) @@ -126,10 +133,35 @@ def _build_all_extensions_with_cmake(self): f"-DCMAKE_BUILD_TYPE={cmake_build_type}", f"-DPYTHON_VERSION={python_version.major}.{python_version.minor}", f"-DENABLE_CUDA={enable_cuda}", + f"-DENABLE_ROCM={enable_rocm}", f"-DTORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR={torchcodec_disable_compile_warning_as_error}", f"-DTORCHCODEC_DISABLE_HOMEBREW_RPATH={torchcodec_disable_homebrew_rpath}", ] + if enable_rocm: + rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm") + rocm_clang = os.path.join(rocm_path, "lib", "llvm", "bin", "clang++") + if os.path.exists(rocm_clang): + cmake_args.append(f"-DCMAKE_HIP_COMPILER={rocm_clang}") + cmake_args.append(f"-DCMAKE_PREFIX_PATH={rocm_path}") + + # Resolve GPU architectures: HIP_ARCHITECTURES > PYTORCH_ROCM_ARCH > torch.cuda.get_arch_list() + hip_archs = os.environ.get("HIP_ARCHITECTURES", "") + if not hip_archs: + pytorch_rocm_arch = os.environ.get("PYTORCH_ROCM_ARCH", "") + if pytorch_rocm_arch: + hip_archs = pytorch_rocm_arch.replace(" ", ";") + if not hip_archs: + try: + arch_list = torch.cuda.get_arch_list() if hasattr(torch.cuda, "get_arch_list") else [] + hip_archs = ";".join(a for a in arch_list if a.startswith("gfx")) + except Exception: + pass + if hip_archs: + cmake_args.append(f"-DHIP_ARCHITECTURES={hip_archs}") + if not os.environ.get("PYTORCH_ROCM_ARCH"): + os.environ["PYTORCH_ROCM_ARCH"] = hip_archs.replace(";", " ") + self.build_temp = os.getenv("TORCHCODEC_CMAKE_BUILD_DIR", self.build_temp) print(f"Using {self.build_temp = }", flush=True) Path(self.build_temp).mkdir(parents=True, exist_ok=True) diff --git a/src/torchcodec/_core/CMakeLists.txt b/src/torchcodec/_core/CMakeLists.txt index 67d5bb5e2..1810a4fff 100644 --- a/src/torchcodec/_core/CMakeLists.txt +++ b/src/torchcodec/_core/CMakeLists.txt @@ -138,10 +138,43 @@ function(make_torchcodec_libraries Metadata.cpp ) + if(ENABLE_CUDA AND ENABLE_ROCM) + message(FATAL_ERROR "Cannot enable both CUDA and ROCM simultaneously. Set only one of ENABLE_CUDA or ENABLE_ROCM.") + endif() + if(ENABLE_CUDA) list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp NVCUVIDRuntimeLoader.cpp) endif() + if(ENABLE_ROCM) + # Find HIP package - required for ROCm builds + # On ROCm PyTorch, HIP is available through the torch installation + if(NOT DEFINED ROCM_PATH) + if(DEFINED ENV{ROCM_PATH}) + set(ROCM_PATH "$ENV{ROCM_PATH}") + else() + set(ROCM_PATH "/opt/rocm") + endif() + endif() + list(APPEND CMAKE_PREFIX_PATH "${ROCM_PATH}") + find_package(hip QUIET) + + # C++ source files that don't contain HIP kernels + set(ROCM_CPP_SOURCES + RocDecodeDeviceInterface.cpp + RocDecodeCache.cpp + HIPCommon.cpp + RocDecodeRuntimeLoader.cpp + ) + list(APPEND core_sources ${ROCM_CPP_SOURCES}) + + # HIPColorspaceKernels.cpp contains __global__ HIP kernels and must + # be compiled with the ROCm Clang compiler with -x hip flag. + # We set ROCM_HIP_KERNEL_SOURCE so it can be linked in later. + set(ROCM_HIP_KERNEL_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/HIPColorspaceKernels.cpp") + endif() + set(core_library_dependencies ${ffmpeg_target} ${TORCH_LIBRARIES} @@ -154,6 +187,12 @@ function(make_torchcodec_libraries ) endif() + if(ENABLE_ROCM) + if(hip_FOUND) + list(APPEND core_library_dependencies hip::host) + endif() + endif() + make_torchcodec_sublibrary( "${core_library_name}" SHARED @@ -161,6 +200,66 @@ function(make_torchcodec_libraries "${core_library_dependencies}" ) + if(ENABLE_ROCM) + # Add HIP include dirs to core library for all ROCm .cpp files + target_include_directories(${core_library_name} + PRIVATE + "${ROCM_PATH}/include" + ) + + # Compile HIP kernel file with ROCm Clang and link it into the core library. + # We use add_custom_command because the kernel file needs -x hip and + # specific GPU architecture flags that are incompatible with the system + # C++ compiler (gcc). + set(HIP_KERNEL_OBJ "${CMAKE_CURRENT_BINARY_DIR}/HIPColorspaceKernels.o") + set(HIP_CLANG "${ROCM_PATH}/lib/llvm/bin/clang++") + + # GPU architectures for HIP kernel compilation. + # Set via -DHIP_ARCHITECTURES or PYTORCH_ROCM_ARCH env var. + if(NOT DEFINED HIP_ARCHITECTURES) + if(DEFINED ENV{PYTORCH_ROCM_ARCH}) + string(REPLACE " " ";" HIP_ARCHITECTURES "$ENV{PYTORCH_ROCM_ARCH}") + else() + message(FATAL_ERROR + "No GPU architectures specified for ROCm build. " + "Set PYTORCH_ROCM_ARCH env var or pass -DHIP_ARCHITECTURES.") + endif() + endif() + + # Build the --offload-arch flags + set(HIP_ARCH_FLAGS "") + foreach(arch ${HIP_ARCHITECTURES}) + list(APPEND HIP_ARCH_FLAGS "--offload-arch=${arch}") + endforeach() + + add_custom_command( + OUTPUT ${HIP_KERNEL_OBJ} + COMMAND ${HIP_CLANG} + -x hip + -c ${ROCM_HIP_KERNEL_SOURCE} + -o ${HIP_KERNEL_OBJ} + ${HIP_ARCH_FLAGS} + -fPIC + -std=c++17 + -O2 + -I${ROCM_PATH}/include + -I${CMAKE_CURRENT_SOURCE_DIR} + -I${CMAKE_CURRENT_SOURCE_DIR}/../../../ + -I${TORCH_INSTALL_PREFIX}/include + -I${TORCH_INSTALL_PREFIX}/include/torch/csrc/api/include + -D__HIP_PLATFORM_AMD__ + DEPENDS ${ROCM_HIP_KERNEL_SOURCE} + COMMENT "Compiling HIP kernels: HIPColorspaceKernels.cpp" + ) + add_custom_target(hip_kernels_${ffmpeg_major_version} DEPENDS ${HIP_KERNEL_OBJ}) + add_dependencies(${core_library_name} hip_kernels_${ffmpeg_major_version}) + target_link_libraries(${core_library_name} PUBLIC ${HIP_KERNEL_OBJ}) + # Link against amdhip64 for HIP runtime - use PRIVATE to avoid + # forcing transitive load of the HIP runtime which needs /dev/kfd. + # The HIP runtime is loaded lazily when actual GPU operations occur. + target_link_libraries(${core_library_name} PRIVATE "${ROCM_PATH}/lib/libamdhip64.so") + endif() + # 2. Create libtorchcodec_custom_opsN.{ext}. set(custom_ops_library_name "libtorchcodec_custom_ops${ffmpeg_major_version}") set(custom_ops_sources diff --git a/src/torchcodec/_core/HIPColorspaceKernels.cpp b/src/torchcodec/_core/HIPColorspaceKernels.cpp new file mode 100644 index 000000000..f46a0acde --- /dev/null +++ b/src/torchcodec/_core/HIPColorspaceKernels.cpp @@ -0,0 +1,283 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. +// +// HIP kernels for YUV<->RGB color conversion on AMD GPUs. +// These replace the NPP library functions used in the CUDA path. +// The color conversion math matches the CPU path (FFmpeg) as closely +// as possible for consistency. +// +// Adapted from rocDecode's colorspace_kernels.cpp with modifications +// for TorchCodec's output format requirements (HWC uint8 RGB24). + +// This file contains HIP GPU kernels. It is compiled as C++ but uses +// HIP intrinsics (__global__, threadIdx, etc.) which are available +// when compiling with the ROCm Clang compiler and -x hip flag. +// See CMakeLists.txt for the compile flags setup. +// +// NOTE: This file must NOT include HIPCommon.h or any ATen/torch headers +// because those pull in which is not available when compiling +// with the standalone ROCm Clang. Only is needed. +#include +#include + +namespace facebook::torchcodec { + +// ---- Constant memory for color conversion matrices ---- +__constant__ float d_yuv2rgb_mat[3][3]; +__constant__ float d_rgb2yuv_mat[3][3]; + +// ---- Helper: compute color matrix coefficients ---- +// Color space standard values match AVColorSpace enum values used in FFmpeg +enum HIPColorStandard { + HIP_CS_BT709 = 1, + HIP_CS_FCC = 4, + HIP_CS_BT470 = 5, + HIP_CS_BT601 = 6, + HIP_CS_SMPTE240M = 7, + HIP_CS_BT2020 = 9, +}; + +static void getColorMatCoefficients( + int colStandard, + float& wr, + float& wb, + int& black, + int& white, + int& maxVal) { + black = 16; + white = 235; + maxVal = 255; + + switch (colStandard) { + case HIP_CS_BT709: + default: + wr = 0.2126f; + wb = 0.0722f; + break; + case HIP_CS_FCC: + wr = 0.30f; + wb = 0.11f; + break; + case HIP_CS_BT470: + case HIP_CS_BT601: + wr = 0.2990f; + wb = 0.1140f; + break; + case HIP_CS_SMPTE240M: + wr = 0.212f; + wb = 0.087f; + break; + case HIP_CS_BT2020: + wr = 0.2627f; + wb = 0.0593f; + break; + } +} + +static void setMatYuv2Rgb(int colStandard, hipStream_t stream) { + float wr, wb; + int black, white, maxVal; + getColorMatCoefficients(colStandard, wr, wb, black, white, maxVal); + float mat[3][3] = { + {1.0f, 0.0f, (1.0f - wr) / 0.5f}, + {1.0f, + -wb * (1.0f - wb) / 0.5f / (1.0f - wb - wr), + -wr * (1.0f - wr) / 0.5f / (1.0f - wb - wr)}, + {1.0f, (1.0f - wb) / 0.5f, 0.0f}, + }; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + mat[i][j] = static_cast( + 1.0 * maxVal / (white - black) * mat[i][j]); + } + } + hipMemcpyToSymbolAsync( + d_yuv2rgb_mat, mat, sizeof(mat), 0, hipMemcpyHostToDevice, stream); +} + +static void setMatRgb2Yuv(int colStandard, hipStream_t stream) { + float wr, wb; + int black, white, maxVal; + getColorMatCoefficients(colStandard, wr, wb, black, white, maxVal); + float mat[3][3] = { + {wr, 1.0f - wb - wr, wb}, + {-0.5f * wr / (1.0f - wb), + -0.5f * (1.0f - wb - wr) / (1.0f - wb), + 0.5f}, + {0.5f, + -0.5f * (1.0f - wb - wr) / (1.0f - wr), + -0.5f * wb / (1.0f - wr)}, + }; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + mat[i][j] = static_cast( + 1.0 * (white - black) / maxVal * mat[i][j]); + } + } + hipMemcpyToSymbolAsync( + d_rgb2yuv_mat, mat, sizeof(mat), 0, hipMemcpyHostToDevice, stream); +} + +// ---- Device helpers ---- +template +__device__ static T clampVal(T x, T lower, T upper) { + return x < lower ? lower : (x > upper ? upper : x); +} + +// ---- NV12 -> RGB24 kernel ---- +// NV12 layout: Y plane (height rows of width bytes at pitch stride), +// UV plane (height/2 rows of width bytes, interleaved U,V) +// Output: HWC uint8 RGB (3 bytes per pixel, contiguous or strided) +__global__ void nv12ToRgb24Kernel( + const uint8_t* __restrict__ nv12, + int nv12Pitch, + uint8_t* __restrict__ rgb, + int rgbPitch, + int width, + int height, + int vPitch) { + // Each thread processes one pixel + int x = threadIdx.x + blockIdx.x * blockDim.x; + int y = threadIdx.y + blockIdx.y * blockDim.y; + if (x >= width || y >= height) { + return; + } + + // Y value + float fy = static_cast(nv12[y * nv12Pitch + x]) - 16.0f; + // UV values (subsampled 2x2) + int uvOffset = vPitch * nv12Pitch + (y / 2) * nv12Pitch + (x & ~1); + float fu = static_cast(nv12[uvOffset]) - 128.0f; + float fv = static_cast(nv12[uvOffset + 1]) - 128.0f; + + float r = d_yuv2rgb_mat[0][0] * fy + d_yuv2rgb_mat[0][1] * fu + + d_yuv2rgb_mat[0][2] * fv; + float g = d_yuv2rgb_mat[1][0] * fy + d_yuv2rgb_mat[1][1] * fu + + d_yuv2rgb_mat[1][2] * fv; + float b = d_yuv2rgb_mat[2][0] * fy + d_yuv2rgb_mat[2][1] * fu + + d_yuv2rgb_mat[2][2] * fv; + + int outIdx = y * rgbPitch + x * 3; + rgb[outIdx + 0] = static_cast(clampVal(r, 0.0f, 255.0f)); + rgb[outIdx + 1] = static_cast(clampVal(g, 0.0f, 255.0f)); + rgb[outIdx + 2] = static_cast(clampVal(b, 0.0f, 255.0f)); +} + +// ---- RGB24 -> NV12 kernel (for encoding support) ---- +// Input: HWC uint8 RGB (3 bytes per pixel) +// Output: NV12 (Y plane + interleaved UV plane) +__global__ void rgb24ToNv12Kernel( + const uint8_t* __restrict__ rgb, + int rgbPitch, + uint8_t* __restrict__ nv12Y, + int yPitch, + uint8_t* __restrict__ nv12UV, + int uvPitch, + int width, + int height) { + int x = threadIdx.x + blockIdx.x * blockDim.x; + int y = threadIdx.y + blockIdx.y * blockDim.y; + if (x >= width || y >= height) { + return; + } + + int rgbIdx = y * rgbPitch + x * 3; + float r = static_cast(rgb[rgbIdx + 0]); + float g = static_cast(rgb[rgbIdx + 1]); + float b = static_cast(rgb[rgbIdx + 2]); + + // Y + float yVal = d_rgb2yuv_mat[0][0] * r + d_rgb2yuv_mat[0][1] * g + + d_rgb2yuv_mat[0][2] * b + 16.0f; + nv12Y[y * yPitch + x] = static_cast(clampVal(yVal, 0.0f, 255.0f)); + + // UV (only for even coordinates - 2x2 subsampling) + if ((x & 1) == 0 && (y & 1) == 0) { + // Average 2x2 block for chroma + float rAvg = r, gAvg = g, bAvg = b; + int count = 1; + if (x + 1 < width) { + int idx2 = y * rgbPitch + (x + 1) * 3; + rAvg += static_cast(rgb[idx2 + 0]); + gAvg += static_cast(rgb[idx2 + 1]); + bAvg += static_cast(rgb[idx2 + 2]); + count++; + } + if (y + 1 < height) { + int idx3 = (y + 1) * rgbPitch + x * 3; + rAvg += static_cast(rgb[idx3 + 0]); + gAvg += static_cast(rgb[idx3 + 1]); + bAvg += static_cast(rgb[idx3 + 2]); + count++; + } + if (x + 1 < width && y + 1 < height) { + int idx4 = (y + 1) * rgbPitch + (x + 1) * 3; + rAvg += static_cast(rgb[idx4 + 0]); + gAvg += static_cast(rgb[idx4 + 1]); + bAvg += static_cast(rgb[idx4 + 2]); + count++; + } + rAvg /= count; + gAvg /= count; + bAvg /= count; + + float uVal = d_rgb2yuv_mat[1][0] * rAvg + d_rgb2yuv_mat[1][1] * gAvg + + d_rgb2yuv_mat[1][2] * bAvg + 128.0f; + float vVal = d_rgb2yuv_mat[2][0] * rAvg + d_rgb2yuv_mat[2][1] * gAvg + + d_rgb2yuv_mat[2][2] * bAvg + 128.0f; + + int uvIdx = (y / 2) * uvPitch + x; + nv12UV[uvIdx] = static_cast(clampVal(uVal, 0.0f, 255.0f)); + nv12UV[uvIdx + 1] = static_cast(clampVal(vVal, 0.0f, 255.0f)); + } +} + +// ---- Host launch functions ---- + +void launchNv12ToRgb24Kernel( + const uint8_t* nv12, + int nv12Pitch, + uint8_t* rgb, + int rgbPitch, + int width, + int height, + int colorStandard, + hipStream_t stream) { + setMatYuv2Rgb(colorStandard, stream); + + dim3 block(16, 16); + dim3 grid( + (width + block.x - 1) / block.x, + (height + block.y - 1) / block.y); + + // vPitch = height (the UV plane starts after height rows of Y data) + nv12ToRgb24Kernel<<>>( + nv12, nv12Pitch, rgb, rgbPitch, width, height, height); +} + +void launchRgb24ToNv12Kernel( + const uint8_t* rgb, + int rgbPitch, + uint8_t* nv12Y, + int yPitch, + uint8_t* nv12UV, + int uvPitch, + int width, + int height, + int colorStandard, + hipStream_t stream) { + setMatRgb2Yuv(colorStandard, stream); + + dim3 block(16, 16); + dim3 grid( + (width + block.x - 1) / block.x, + (height + block.y - 1) / block.y); + + rgb24ToNv12Kernel<<>>( + rgb, rgbPitch, nv12Y, yPitch, nv12UV, uvPitch, width, height); +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/HIPCommon.cpp b/src/torchcodec/_core/HIPCommon.cpp new file mode 100644 index 000000000..878d3404c --- /dev/null +++ b/src/torchcodec/_core/HIPCommon.cpp @@ -0,0 +1,174 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "HIPCommon.h" + +extern "C" { +#include +} + +namespace facebook::torchcodec { + +void initializeHIPContextWithPytorch(const torch::Device& device) { + // It is important for PyTorch itself to create the HIP/ROCm context. + // If some other library creates the context it may not be compatible + // with PyTorch. This is a dummy tensor to initialize the context. + // On ROCm PyTorch, torch::kCUDA maps to HIP devices. + torch::Tensor dummyTensorForHIPInitialization = torch::zeros( + {1}, torch::TensorOptions().dtype(torch::kUInt8).device(device)); +} + +// Map AVColorSpace to the color standard enum used by our HIP kernels. +// These values match the matrix_coefficients field in video signal description. +static int avColorSpaceToHIPStandard(int colorspace) { + switch (colorspace) { + case AVCOL_SPC_BT709: + return 1; // HIP_CS_BT709 + case AVCOL_SPC_BT470BG: + case AVCOL_SPC_SMPTE170M: + return 6; // HIP_CS_BT601 + case AVCOL_SPC_BT2020_NCL: + case AVCOL_SPC_BT2020_CL: + return 9; // HIP_CS_BT2020 + case AVCOL_SPC_SMPTE240M: + return 7; // HIP_CS_SMPTE240M + case AVCOL_SPC_FCC: + return 4; // HIP_CS_FCC + default: + return 6; // Default to BT.601 + } +} + +torch::Tensor convertNV12FrameToRGB_HIP( + uint8_t* nv12Data, + int nv12Pitch, + int width, + int height, + int colorspace, + int colorRange, + const torch::Device& device, + hipStream_t stream, + std::optional preAllocatedOutputTensor) { + auto frameDims = FrameDims(height, width); + torch::Tensor dst; + if (preAllocatedOutputTensor.has_value()) { + dst = preAllocatedOutputTensor.value(); + } else { + dst = allocateEmptyHWCTensor(frameDims, device); + } + + int colStandard = avColorSpaceToHIPStandard(colorspace); + + // The RGB output pitch is the tensor's row stride in bytes + // dst shape is [H, W, 3], stride(0) gives the row stride + int rgbPitch = static_cast(dst.stride(0) * dst.element_size()); + + launchNv12ToRgb24Kernel( + nv12Data, + nv12Pitch, + static_cast(dst.data_ptr()), + rgbPitch, + width, + height, + colStandard, + stream); + + return dst; +} + +NV12Frame convertRGBToNV12_HIP( + const torch::Tensor& rgbTensor, + int colorspace, + [[maybe_unused]] int colorRange, + hipStream_t stream) { + TORCH_CHECK( + rgbTensor.dim() == 3 && rgbTensor.size(2) == 3, + "Expected HWC RGB tensor with 3 channels, got shape: ", + rgbTensor.sizes()); + TORCH_CHECK( + rgbTensor.device().type() == torch::kCUDA, + "Expected tensor on CUDA/HIP device, got: ", + rgbTensor.device().str()); + + int height = static_cast(rgbTensor.size(0)); + int width = static_cast(rgbTensor.size(1)); + int rgbPitch = static_cast(rgbTensor.stride(0) * rgbTensor.element_size()); + + // Allocate NV12 buffer: Y plane (width * height) + UV plane (width * height/2) + int ySize = width * height; + int uvSize = width * (height / 2); + size_t totalSize = static_cast(ySize + uvSize); + + uint8_t* nv12Buffer = nullptr; + hipError_t err = hipMalloc( + reinterpret_cast(&nv12Buffer), totalSize); + TORCH_CHECK( + err == hipSuccess, + "Failed to allocate HIP memory for NV12: ", + hipGetErrorString(err)); + + uint8_t* yPlane = nv12Buffer; + uint8_t* uvPlane = nv12Buffer + ySize; + + int colStandard = avColorSpaceToHIPStandard(colorspace); + + launchRgb24ToNv12Kernel( + static_cast(rgbTensor.data_ptr()), + rgbPitch, + yPlane, + width, // Y pitch = width (tightly packed) + uvPlane, + width, // UV pitch = width + width, + height, + colStandard, + stream); + + return NV12Frame{yPlane, uvPlane, width, totalSize}; +} + +void validatePreAllocatedTensorShape_HIP( + const std::optional& preAllocatedOutputTensor, + int width, + int height) { + if (preAllocatedOutputTensor.has_value()) { + auto shape = preAllocatedOutputTensor.value().sizes(); + TORCH_CHECK( + (shape.size() == 3) && (shape[0] == height) && + (shape[1] == width) && (shape[2] == 3), + "Expected tensor of shape ", + height, + "x", + width, + "x3, got ", + shape); + } +} + +int getDeviceIndex_HIP(const torch::Device& device) { + int deviceIndex = static_cast(device.index()); + TORCH_CHECK( + deviceIndex >= -1 && deviceIndex < MAX_ROCM_GPUS, + "Invalid device index = ", + deviceIndex); + + if (deviceIndex == -1) { + // On ROCm PyTorch, hipGetDevice works through the CUDA->HIP mapping + TORCH_CHECK( + hipGetDevice(&deviceIndex) == hipSuccess, + "Failed to get current HIP device."); + } + return deviceIndex; +} + +// Provide the getDeviceIndex() definition required by Cache.h's +// PerGpuCache template. On CUDA builds this is defined in CUDACommon.cpp; +// on ROCm builds we provide the equivalent here. +int getDeviceIndex(const torch::Device& device) { + return getDeviceIndex_HIP(device); +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/HIPCommon.h b/src/torchcodec/_core/HIPCommon.h new file mode 100644 index 000000000..dd92b40fb --- /dev/null +++ b/src/torchcodec/_core/HIPCommon.h @@ -0,0 +1,95 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// HIP equivalents of the utilities in CUDACommon.h. +// On ROCm, PyTorch reuses the torch::kCUDA device type for HIP devices, +// and cudaStream_t maps to hipStream_t. This file provides the ROCm-specific +// implementations for device initialization, color conversion, and validation. + +#pragma once + +// On ROCm, we avoid including ATen/cuda headers directly with GCC because +// they reference which is only available when compiling with hipcc. +// Instead, we use HIP types directly and convert at the boundary. +#include + +#include "FFMPEGCommon.h" +#include "Frame.h" + +// Use the HIP-compatible types from the ROCm installation. +// hipStream_t and other HIP types are defined here. +#include + +namespace facebook::torchcodec { + +// Maximum number of GPUs supported (same as PyTorch's limit) +constexpr int MAX_ROCM_GPUS = 128; + +// Initialize the HIP/ROCm context through PyTorch to ensure compatibility. +// This must be called before any ROCm operations. +void initializeHIPContextWithPytorch(const torch::Device& device); + +// Convert an NV12 frame (in GPU memory) to an RGB HWC tensor on the same GPU. +// This uses custom HIP kernels instead of NPP. +torch::Tensor convertNV12FrameToRGB_HIP( + uint8_t* nv12Data, + int nv12Pitch, + int width, + int height, + int colorspace, // AVColorSpace enum value + int colorRange, // AVColorRange enum value + const torch::Device& device, + hipStream_t stream, + std::optional preAllocatedOutputTensor = std::nullopt); + +// Convert an RGB HWC tensor to NV12 format in GPU memory (for encoding). +// Returns a pair of device pointers: {Y plane, UV plane} and the pitch. +struct NV12Frame { + uint8_t* yPlane; + uint8_t* uvPlane; + int pitch; + size_t totalSize; +}; + +NV12Frame convertRGBToNV12_HIP( + const torch::Tensor& rgbTensor, + int colorspace, + int colorRange, + hipStream_t stream); + +// Validate pre-allocated tensor shape matches expected frame dimensions +void validatePreAllocatedTensorShape_HIP( + const std::optional& preAllocatedOutputTensor, + int width, + int height); + +// Get device index, handling the -1 (current device) case +int getDeviceIndex_HIP(const torch::Device& device); + +// HIP kernel launch functions (implemented in HIPColorspaceKernels.hip) +void launchNv12ToRgb24Kernel( + const uint8_t* nv12, + int nv12Pitch, + uint8_t* rgb, + int rgbPitch, + int width, + int height, + int colorStandard, + hipStream_t stream); + +void launchRgb24ToNv12Kernel( + const uint8_t* rgb, + int rgbPitch, + uint8_t* nv12Y, + int yPitch, + uint8_t* nv12UV, + int uvPitch, + int width, + int height, + int colorStandard, + hipStream_t stream); + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/RocDecodeCache.cpp b/src/torchcodec/_core/RocDecodeCache.cpp new file mode 100644 index 000000000..a185c23e4 --- /dev/null +++ b/src/torchcodec/_core/RocDecodeCache.cpp @@ -0,0 +1,49 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "RocDecodeCache.h" + +namespace facebook::torchcodec { + +RocDecodeCache& RocDecodeCache::getCache(const torch::Device& device) { + static RocDecodeCache cacheInstances[MAX_ROCM_GPUS]; + return cacheInstances[getDeviceIndex_HIP(device)]; +} + +UniqueRocDecoder RocDecodeCache::getDecoder( + const RocdecVideoFormat* videoFormat) { + CacheKey key(videoFormat); + std::lock_guard lock(cacheLock_); + + auto it = cache_.find(key); + if (it != cache_.end()) { + auto decoder = std::move(it->second); + cache_.erase(it); + return decoder; + } + + return nullptr; +} + +bool RocDecodeCache::returnDecoder( + const RocdecVideoFormat* videoFormat, + UniqueRocDecoder decoder) { + if (!decoder) { + return false; + } + + CacheKey key(videoFormat); + std::lock_guard lock(cacheLock_); + + if (cache_.size() >= MAX_CACHE_SIZE) { + return false; + } + + cache_[key] = std::move(decoder); + return true; +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/RocDecodeCache.h b/src/torchcodec/_core/RocDecodeCache.h new file mode 100644 index 000000000..28b41fa2b --- /dev/null +++ b/src/torchcodec/_core/RocDecodeCache.h @@ -0,0 +1,104 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Cache for rocDecode decoder instances, analogous to NVDECCache.h. +// Decoder creation is expensive, so we cache and reuse decoders with +// matching parameters across VideoDecoder instances. + +#pragma once + +#include +#include +#include + +#include + +#include "HIPCommon.h" +#include "RocDecodeRuntimeLoader.h" +#include "rocdecode_include/rocdecode.h" +#include "rocdecode_include/rocparser.h" + +namespace facebook::torchcodec { + +// Custom deleter for rocDecode decoder handles +struct RocDecoderDeleter { + void operator()(rocDecDecoderHandle* handlePtr) const { + if (handlePtr && *handlePtr) { + rocDecDestroyDecoder(*handlePtr); + delete handlePtr; + } + } +}; + +using UniqueRocDecoder = + std::unique_ptr; + +// A per-device cache for rocDecode decoders. There is one instance of this +// class per GPU device, and it is accessed through the static getCache() method. +class RocDecodeCache { + public: + static RocDecodeCache& getCache(const torch::Device& device); + + // Get decoder from cache - returns nullptr if none available + UniqueRocDecoder getDecoder(const RocdecVideoFormat* videoFormat); + + // Return decoder to cache - returns true if added to cache + bool returnDecoder( + const RocdecVideoFormat* videoFormat, + UniqueRocDecoder decoder); + + private: + // Cache key struct: a decoder can be reused only if all these parameters match + struct CacheKey { + rocDecVideoCodec codecType; + uint32_t width; + uint32_t height; + rocDecVideoChromaFormat chromaFormat; + uint32_t bitDepthLumaMinus8; + uint8_t numDecodeSurfaces; + + CacheKey() = delete; + + explicit CacheKey(const RocdecVideoFormat* videoFormat) + : codecType(videoFormat->codec), + width(videoFormat->coded_width), + height(videoFormat->coded_height), + chromaFormat(videoFormat->chroma_format), + bitDepthLumaMinus8(videoFormat->bit_depth_luma_minus8), + numDecodeSurfaces(videoFormat->min_num_decode_surfaces) {} + + CacheKey(const CacheKey&) = default; + CacheKey& operator=(const CacheKey&) = default; + + bool operator<(const CacheKey& other) const { + return std::tie( + codecType, + width, + height, + chromaFormat, + bitDepthLumaMinus8, + numDecodeSurfaces) < + std::tie( + other.codecType, + other.width, + other.height, + other.chromaFormat, + other.bitDepthLumaMinus8, + other.numDecodeSurfaces); + } + }; + + RocDecodeCache() = default; + ~RocDecodeCache() = default; + + std::map cache_; + std::mutex cacheLock_; + + // Max number of cached decoders per device + static constexpr int MAX_CACHE_SIZE = 20; +}; + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/RocDecodeDeviceInterface.cpp b/src/torchcodec/_core/RocDecodeDeviceInterface.cpp new file mode 100644 index 000000000..90d7be2c5 --- /dev/null +++ b/src/torchcodec/_core/RocDecodeDeviceInterface.cpp @@ -0,0 +1,841 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// Use c10/hip headers instead of c10/cuda to avoid dependency +// when compiling with GCC on ROCm. On ROCm PyTorch, the HIP stream +// API is the native interface. +#include +#include +#include +#include +#include + +#include "RocDecodeDeviceInterface.h" + +#include "DeviceInterface.h" +#include "FFMPEGCommon.h" +#include "RocDecodeCache.h" +#include "RocDecodeRuntimeLoader.h" + +#include "rocdecode_include/rocdecode.h" +#include "rocdecode_include/rocparser.h" + +extern "C" { +#include +} + +namespace facebook::torchcodec { + +namespace { + +// Register as the "beta" variant for CUDA device type. +// On ROCm PyTorch, torch::kCUDA maps to HIP devices. The build system +// ensures only one of NVDEC or rocDecode is compiled in, so there's +// no registration conflict. +static bool g_rocm_beta = registerDeviceInterface( + DeviceInterfaceKey(torch::kCUDA, /*variant=*/"beta"), + [](const torch::Device& device) { + return new RocDecodeDeviceInterface(device); + }); + +// Also register as the default "ffmpeg" CUDA variant so that +// device="cuda" without explicit variant selection works on ROCm. +static bool g_rocm_default = registerDeviceInterface( + DeviceInterfaceKey(torch::kCUDA), + [](const torch::Device& device) { + return new RocDecodeDeviceInterface(device); + }); + +// ---- Parser callbacks (C-style, forwarded to instance methods) ---- +static int ROCDECAPI +pfnSequenceCallback(void* pUserData, RocdecVideoFormat* videoFormat) { + auto decoder = static_cast(pUserData); + return decoder->streamPropertyChange(videoFormat); +} + +static int ROCDECAPI +pfnDecodePictureCallback(void* pUserData, RocdecPicParams* picParams) { + auto decoder = static_cast(pUserData); + return decoder->frameReadyForDecoding(picParams); +} + +static int ROCDECAPI +pfnDisplayPictureCallback(void* pUserData, RocdecParserDispInfo* dispInfo) { + auto decoder = static_cast(pUserData); + return decoder->frameReadyInDisplayOrder(dispInfo); +} + +// ---- Codec validation ---- +std::optional validateCodecSupport(AVCodecID codecId) { + switch (codecId) { + case AV_CODEC_ID_H264: + return rocDecVideoCodec_AVC; + case AV_CODEC_ID_HEVC: + return rocDecVideoCodec_HEVC; + case AV_CODEC_ID_AV1: + return rocDecVideoCodec_AV1; + case AV_CODEC_ID_VP9: + return rocDecVideoCodec_VP9; + case AV_CODEC_ID_VP8: + return rocDecVideoCodec_VP8; + case AV_CODEC_ID_MPEG4: + return rocDecVideoCodec_MPEG4; + case AV_CODEC_ID_MPEG1VIDEO: + return rocDecVideoCodec_MPEG1; + case AV_CODEC_ID_MPEG2VIDEO: + return rocDecVideoCodec_MPEG2; + default: + return std::nullopt; + } +} + +std::optional validateChromaSupport( + const AVPixFmtDescriptor* desc) { + TORCH_CHECK(desc != nullptr, "desc can't be null"); + + if (desc->nb_components == 1) { + return rocDecVideoChromaFormat_Monochrome; + } else if ( + desc->nb_components >= 3 && !(desc->flags & AV_PIX_FMT_FLAG_RGB)) { + if (desc->log2_chroma_w == 0 && desc->log2_chroma_h == 0) { + return rocDecVideoChromaFormat_444; + } else if (desc->log2_chroma_w == 1 && desc->log2_chroma_h == 1) { + return rocDecVideoChromaFormat_420; + } else if (desc->log2_chroma_w == 1 && desc->log2_chroma_h == 0) { + return rocDecVideoChromaFormat_422; + } + } + + return std::nullopt; +} + +bool nativeRocDecSupport( + const torch::Device& device, + const SharedAVCodecContext& codecContext) { + auto codecType = validateCodecSupport(codecContext->codec_id); + if (!codecType.has_value()) { + return false; + } + + const AVPixFmtDescriptor* desc = av_pix_fmt_desc_get(codecContext->pix_fmt); + if (!desc) { + return false; + } + + auto chromaFormat = validateChromaSupport(desc); + if (!chromaFormat.has_value()) { + return false; + } + + // Ensure rocDecode queries the correct GPU + int deviceIndex = getDeviceIndex_HIP(device); + hipSetDevice(deviceIndex); + + // Query decoder capabilities + RocdecDecodeCaps caps = {}; + caps.device_id = static_cast(deviceIndex); + caps.codec_type = codecType.value(); + caps.chroma_format = chromaFormat.value(); + caps.bit_depth_minus_8 = static_cast(desc->comp[0].depth - 8); + + rocDecStatus result = rocDecGetDecoderCaps(&caps); + if (result != ROCDEC_SUCCESS) { + return false; + } + + if (!caps.is_supported) { + return false; + } + + auto coded_width = static_cast(codecContext->coded_width); + auto coded_height = static_cast(codecContext->coded_height); + if (coded_width < caps.min_width || coded_height < caps.min_height || + coded_width > caps.max_width || coded_height > caps.max_height) { + return false; + } + + // Check NV12 output format support + bool supportsNV12 = + (caps.output_format_mask >> rocDecVideoSurfaceFormat_NV12) & 1; + if (!supportsNV12) { + return false; + } + + return true; +} + +// Callback for freeing HIP memory associated with AVFrame +void hipBufferFreeCallback(void* opaque, [[maybe_unused]] uint8_t* data) { + hipFree(opaque); +} + +} // namespace + +// ---- Constructor / Destructor ---- + +RocDecodeDeviceInterface::RocDecodeDeviceInterface(const torch::Device& device) + : DeviceInterface(device) { + TORCH_CHECK( + g_rocm_beta || g_rocm_default, + "RocDecodeDeviceInterface was not registered!"); + TORCH_CHECK( + device_.type() == torch::kCUDA, "Unsupported device: ", device_.str()); + + initializeHIPContextWithPytorch(device_); + + // Explicitly set the HIP device so the rocDecode library (loaded via + // dlopen) uses the same GPU as PyTorch. Without this, rocDecode may + // open a DRM fd to a different GPU, causing memory access faults in + // multi-GPU environments. + int deviceIndex = getDeviceIndex_HIP(device_); + hipError_t hipErr = hipSetDevice(deviceIndex); + TORCH_CHECK( + hipErr == hipSuccess, + "Failed to set HIP device ", + deviceIndex, + ": ", + hipGetErrorString(hipErr)); + + rocDecodeAvailable_ = loadRocDecodeLibrary(); +} + +RocDecodeDeviceInterface::~RocDecodeDeviceInterface() { + if (decoder_) { + flush(); + RocDecodeCache::getCache(device_).returnDecoder( + &videoFormat_, std::move(decoder_)); + } + + if (videoParser_) { + rocDecDestroyVideoParser(videoParser_); + videoParser_ = nullptr; + } +} + +// ---- Initialization ---- + +void RocDecodeDeviceInterface::initialize( + const AVStream* avStream, + const UniqueDecodingAVFormatContext& avFormatCtx, + [[maybe_unused]] const SharedAVCodecContext& codecContext) { + if (!rocDecodeAvailable_ || !nativeRocDecSupport(device_, codecContext)) { + cpuFallback_ = createDeviceInterface(torch::kCPU); + TORCH_CHECK( + cpuFallback_ != nullptr, "Failed to create CPU device interface"); + cpuFallback_->initialize(avStream, avFormatCtx, codecContext); + cpuFallback_->initializeVideo( + VideoStreamOptions(), + {}, + /*resizedOutputDims=*/std::nullopt); + return; + } + + TORCH_CHECK(avStream != nullptr, "AVStream cannot be null"); + timeBase_ = avStream->time_base; + frameRateAvgFromFFmpeg_ = avStream->r_frame_rate; + + const AVCodecParameters* codecPar = avStream->codecpar; + TORCH_CHECK(codecPar != nullptr, "CodecParameters cannot be null"); + + initializeBSF(codecPar, avFormatCtx); + + // Create parser + RocdecParserParams parserParams = {}; + auto codecType = validateCodecSupport(codecPar->codec_id); + TORCH_CHECK( + codecType.has_value(), + "This should never happen, we should be using the CPU fallback. " + "Please report a bug."); + parserParams.codec_type = codecType.value(); + parserParams.max_num_decode_surfaces = 8; + parserParams.max_display_delay = 0; + parserParams.clock_rate = 0; // default 10MHz + parserParams.user_data = this; + parserParams.pfn_sequence_callback = pfnSequenceCallback; + parserParams.pfn_decode_picture = pfnDecodePictureCallback; + parserParams.pfn_display_picture = pfnDisplayPictureCallback; + + rocDecStatus result = + rocDecCreateVideoParser(&videoParser_, &parserParams); + TORCH_CHECK( + result == ROCDEC_SUCCESS, + "Failed to create rocDecode video parser: ", + rocDecGetErrorName(result)); +} + +void RocDecodeDeviceInterface::initializeBSF( + const AVCodecParameters* codecPar, + const UniqueDecodingAVFormatContext& avFormatCtx) { + // Setup bitstream filters - identical logic to BetaCudaDeviceInterface + TORCH_CHECK(codecPar != nullptr, "codecPar cannot be null"); + TORCH_CHECK(avFormatCtx != nullptr, "AVFormatContext cannot be null"); + TORCH_CHECK( + avFormatCtx->iformat != nullptr, + "AVFormatContext->iformat cannot be null"); + std::string filterName; + + switch (codecPar->codec_id) { + case AV_CODEC_ID_H264: { + const std::string formatName = avFormatCtx->iformat->long_name + ? avFormatCtx->iformat->long_name + : ""; + if (formatName == "QuickTime / MOV" || + formatName == "FLV (Flash Video)" || + formatName == "Matroska / WebM" || formatName == "raw H.264 video") { + filterName = "h264_mp4toannexb"; + } + break; + } + case AV_CODEC_ID_HEVC: { + const std::string formatName = avFormatCtx->iformat->long_name + ? avFormatCtx->iformat->long_name + : ""; + if (formatName == "QuickTime / MOV" || + formatName == "FLV (Flash Video)" || + formatName == "Matroska / WebM" || formatName == "raw HEVC video") { + filterName = "hevc_mp4toannexb"; + } + break; + } + case AV_CODEC_ID_MPEG4: { + const std::string formatName = + avFormatCtx->iformat->name ? avFormatCtx->iformat->name : ""; + if (formatName == "avi") { + filterName = "mpeg4_unpack_bframes"; + } + break; + } + default: + break; + } + + if (filterName.empty()) { + return; + } + + const AVBitStreamFilter* avBSF = av_bsf_get_by_name(filterName.c_str()); + TORCH_CHECK( + avBSF != nullptr, "Failed to find bitstream filter: ", filterName); + + AVBSFContext* avBSFContext = nullptr; + int retVal = av_bsf_alloc(avBSF, &avBSFContext); + TORCH_CHECK( + retVal >= AVSUCCESS, + "Failed to allocate bitstream filter: ", + getFFMPEGErrorStringFromErrorCode(retVal)); + + bitstreamFilter_.reset(avBSFContext); + + retVal = avcodec_parameters_copy(bitstreamFilter_->par_in, codecPar); + TORCH_CHECK( + retVal >= AVSUCCESS, + "Failed to copy codec parameters: ", + getFFMPEGErrorStringFromErrorCode(retVal)); + + retVal = av_bsf_init(bitstreamFilter_.get()); + TORCH_CHECK( + retVal == AVSUCCESS, + "Failed to initialize bitstream filter: ", + getFFMPEGErrorStringFromErrorCode(retVal)); +} + +// ---- Decoder creation ---- + +UniqueRocDecoder RocDecodeDeviceInterface::createDecoder( + RocdecVideoFormat* videoFormat, + int deviceId) { + // Set HIP device before creating the decoder to ensure rocDecode + // uses the same GPU context as PyTorch/HIP. + hipSetDevice(deviceId); + + RocDecoderCreateInfo decoderParams = {}; + // device_id is the FIRST field in the real v1.5.0 struct (uint8_t) + decoderParams.device_id = static_cast(deviceId); + decoderParams.width = videoFormat->coded_width; + decoderParams.height = videoFormat->coded_height; + decoderParams.num_decode_surfaces = videoFormat->min_num_decode_surfaces; + decoderParams.codec_type = videoFormat->codec; + decoderParams.chroma_format = videoFormat->chroma_format; + decoderParams.bit_depth_minus_8 = videoFormat->bit_depth_luma_minus8; + decoderParams.intra_decode_only = 0; + decoderParams.max_width = videoFormat->coded_width; + decoderParams.max_height = videoFormat->coded_height; + decoderParams.display_rect.left = + static_cast(videoFormat->display_area.left); + decoderParams.display_rect.top = + static_cast(videoFormat->display_area.top); + decoderParams.display_rect.right = + static_cast(videoFormat->display_area.right); + decoderParams.display_rect.bottom = + static_cast(videoFormat->display_area.bottom); + // Request NV12 output - same as NVDEC path. 10bit videos will be + // automatically converted to 8bit by the VCN hardware. + decoderParams.output_format = rocDecVideoSurfaceFormat_NV12; + decoderParams.target_width = static_cast( + videoFormat->display_area.right - videoFormat->display_area.left); + decoderParams.target_height = static_cast( + videoFormat->display_area.bottom - videoFormat->display_area.top); + decoderParams.num_output_surfaces = 1; + + rocDecDecoderHandle* decoderHandle = new rocDecDecoderHandle(); + rocDecStatus result = rocDecCreateDecoder(decoderHandle, &decoderParams); + TORCH_CHECK( + result == ROCDEC_SUCCESS, + "Failed to create rocDecode decoder: ", + rocDecGetErrorName(result)); + + return UniqueRocDecoder(decoderHandle, RocDecoderDeleter{}); +} + +// ---- Parser callbacks ---- + +int RocDecodeDeviceInterface::streamPropertyChange( + RocdecVideoFormat* videoFormat) { + TORCH_CHECK(videoFormat != nullptr, "Invalid video format"); + + videoFormat_ = *videoFormat; + + if (videoFormat_.min_num_decode_surfaces == 0) { + videoFormat_.min_num_decode_surfaces = 20; + } + + if (!decoder_) { + decoder_ = RocDecodeCache::getCache(device_).getDecoder(videoFormat); + + if (!decoder_) { + decoder_ = createDecoder( + videoFormat, getDeviceIndex_HIP(device_)); + } + + TORCH_CHECK(decoder_, "Failed to get or create rocDecode decoder"); + } + + return static_cast(videoFormat_.min_num_decode_surfaces); +} + +int RocDecodeDeviceInterface::frameReadyForDecoding( + RocdecPicParams* picParams) { + TORCH_CHECK(picParams != nullptr, "Invalid picture parameters"); + TORCH_CHECK(decoder_, "Decoder not initialized before picture decode"); + + rocDecStatus result = rocDecDecodeFrame(*decoder_.get(), picParams); + + // 0 means error, 1 means success (same convention as NVCUVID) + return (result == ROCDEC_SUCCESS) ? 1 : 0; +} + +int RocDecodeDeviceInterface::frameReadyInDisplayOrder( + RocdecParserDispInfo* dispInfo) { + readyFrames_.push(*dispInfo); + return 1; // success +} + +// ---- Send/Receive pattern ---- + +int RocDecodeDeviceInterface::sendPacket(ReferenceAVPacket& packet) { + if (cpuFallback_) { + return cpuFallback_->sendPacket(packet); + } + + TORCH_CHECK( + packet.get() && packet->data && packet->size > 0, + "sendPacket received an empty packet, this is unexpected."); + + // Apply BSF if needed + AutoAVPacket filteredAutoPacket; + ReferenceAVPacket filteredPacket(filteredAutoPacket); + ReferenceAVPacket& packetToSend = applyBSF(packet, filteredPacket); + + RocdecSourceDataPacket rocPacket = {}; + rocPacket.payload = packetToSend->data; + rocPacket.payload_size = static_cast(packetToSend->size); + rocPacket.flags = ROCDEC_PKT_TIMESTAMP; + rocPacket.pts = static_cast(packetToSend->pts); + + return sendRocDecPacket(rocPacket); +} + +int RocDecodeDeviceInterface::sendEOFPacket() { + if (cpuFallback_) { + return cpuFallback_->sendEOFPacket(); + } + + RocdecSourceDataPacket rocPacket = {}; + rocPacket.flags = ROCDEC_PKT_ENDOFSTREAM; + eofSent_ = true; + + return sendRocDecPacket(rocPacket); +} + +int RocDecodeDeviceInterface::sendRocDecPacket( + RocdecSourceDataPacket& rocPacket) { + rocDecStatus result = rocDecParseVideoData(videoParser_, &rocPacket); + return result == ROCDEC_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL; +} + +ReferenceAVPacket& RocDecodeDeviceInterface::applyBSF( + ReferenceAVPacket& packet, + ReferenceAVPacket& filteredPacket) { + if (!bitstreamFilter_) { + return packet; + } + + int retVal = av_bsf_send_packet(bitstreamFilter_.get(), packet.get()); + TORCH_CHECK( + retVal >= AVSUCCESS, + "Failed to send packet to bitstream filter: ", + getFFMPEGErrorStringFromErrorCode(retVal)); + + retVal = av_bsf_receive_packet(bitstreamFilter_.get(), filteredPacket.get()); + TORCH_CHECK( + retVal >= AVSUCCESS, + "Failed to receive packet from bitstream filter: ", + getFFMPEGErrorStringFromErrorCode(retVal)); + + return filteredPacket; +} + +// ---- Frame retrieval ---- + +int RocDecodeDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) { + if (cpuFallback_) { + return cpuFallback_->receiveFrame(avFrame); + } + + if (readyFrames_.empty()) { + return eofSent_ ? AVERROR_EOF : AVERROR(EAGAIN); + } + + RocdecParserDispInfo dispInfo = readyFrames_.front(); + readyFrames_.pop(); + + // Ensure the correct HIP device is active before calling rocDecode APIs + hipSetDevice(getDeviceIndex_HIP(device_)); + + // rocDecGetVideoFrame is a blocking call that returns per-plane device + // pointers. Unlike NVDEC's map/unmap pattern, rocDecode gives us direct + // pointers to the decoded frame data. + RocdecProcParams procParams = {}; + procParams.progressive_frame = dispInfo.progressive_frame; + procParams.top_field_first = dispInfo.top_field_first; + + void* devMemPtr[3] = {nullptr, nullptr, nullptr}; + uint32_t pitch = 0; + + rocDecStatus result = rocDecGetVideoFrame( + *decoder_.get(), dispInfo.picture_index, devMemPtr, &pitch, &procParams); + if (result != ROCDEC_SUCCESS) { + return AVERROR_EXTERNAL; + } + + avFrame = convertRocDecFrameToAVFrame(devMemPtr, pitch, dispInfo); + + return AVSUCCESS; +} + +UniqueAVFrame RocDecodeDeviceInterface::convertRocDecFrameToAVFrame( + void* devMemPtr[3], + uint32_t pitch, + const RocdecParserDispInfo& dispInfo) { + TORCH_CHECK(devMemPtr[0] != nullptr, "Invalid decoded frame pointer"); + + int width = videoFormat_.display_area.right - videoFormat_.display_area.left; + int height = + videoFormat_.display_area.bottom - videoFormat_.display_area.top; + + TORCH_CHECK(width > 0 && height > 0, "Invalid frame dimensions"); + TORCH_CHECK( + pitch >= static_cast(width), "Pitch must be >= width"); + + // Unlike NVDEC which has an explicit map/unmap pattern to lock decode + // surfaces, rocDecGetVideoFrame returns raw pointers into the decoder's + // internal surface pool. These surfaces can be recycled as soon as the + // next frame is decoded, so we must copy the NV12 data to owned memory + // before returning. + int yHeight = height; + int uvHeight = height / 2; // NV12: UV plane is half height + size_t ySize = static_cast(pitch) * yHeight; + size_t uvSize = static_cast(pitch) * uvHeight; + size_t totalSize = ySize + uvSize; + + uint8_t* ownedBuffer = nullptr; + hipError_t err = + hipMalloc(reinterpret_cast(&ownedBuffer), totalSize); + TORCH_CHECK( + err == hipSuccess, + "Failed to allocate HIP memory for frame copy: ", + hipGetErrorString(err)); + + hipStream_t currentStream = + c10::hip::getCurrentHIPStream(device_.index()).stream(); + + // Copy Y plane from decoder surface to owned buffer + err = hipMemcpy2DAsync( + ownedBuffer, + pitch, + devMemPtr[0], + pitch, + width, + yHeight, + hipMemcpyDeviceToDevice, + currentStream); + TORCH_CHECK( + err == hipSuccess, + "Failed to copy Y plane: ", + hipGetErrorString(err)); + + // Copy UV plane from decoder surface to owned buffer + err = hipMemcpy2DAsync( + ownedBuffer + ySize, + pitch, + devMemPtr[1], + pitch, + width, + uvHeight, + hipMemcpyDeviceToDevice, + currentStream); + TORCH_CHECK( + err == hipSuccess, + "Failed to copy UV plane: ", + hipGetErrorString(err)); + + UniqueAVFrame avFrame(av_frame_alloc()); + TORCH_CHECK(avFrame.get() != nullptr, "Failed to allocate AVFrame"); + + avFrame->width = width; + avFrame->height = height; + avFrame->format = AV_PIX_FMT_CUDA; // CUDA format on ROCm = HIP device mem + avFrame->pts = static_cast(dispInfo.pts); + + setDuration(avFrame, computeSafeDuration(frameRateAvgFromFFmpeg_, timeBase_)); + + // Map matrix_coefficients to AVColorSpace + switch (videoFormat_.video_signal_description.matrix_coefficients) { + case 1: + avFrame->colorspace = AVCOL_SPC_BT709; + break; + case 6: + avFrame->colorspace = AVCOL_SPC_SMPTE170M; // BT.601 + break; + default: + avFrame->colorspace = AVCOL_SPC_SMPTE170M; + break; + } + + avFrame->color_range = + videoFormat_.video_signal_description.video_full_range_flag + ? AVCOL_RANGE_JPEG + : AVCOL_RANGE_MPEG; + + // Point AVFrame at our owned copy of the NV12 data + avFrame->data[0] = ownedBuffer; + avFrame->data[1] = ownedBuffer + ySize; + avFrame->data[2] = nullptr; + avFrame->data[3] = nullptr; + avFrame->linesize[0] = static_cast(pitch); + avFrame->linesize[1] = static_cast(pitch); + avFrame->linesize[2] = 0; + avFrame->linesize[3] = 0; + + // Register cleanup callback so HIP memory is freed when AVFrame is released + avFrame->opaque_ref = av_buffer_create( + nullptr, + 0, + hipBufferFreeCallback, + ownedBuffer, + 0); + TORCH_CHECK( + avFrame->opaque_ref != nullptr, + "Failed to create GPU memory cleanup reference"); + + return avFrame; +} + +// ---- Flush ---- + +void RocDecodeDeviceInterface::flush() { + if (cpuFallback_) { + cpuFallback_->flush(); + return; + } + + // Send EOF to flush remaining frames from parser + sendEOFPacket(); + eofSent_ = false; + + // Clear the ready frames queue + std::queue emptyQueue; + std::swap(readyFrames_, emptyQueue); +} + +// ---- CPU fallback: transfer CPU frame to GPU as NV12 ---- + +UniqueAVFrame RocDecodeDeviceInterface::transferCpuFrameToGpuNV12( + UniqueAVFrame& cpuFrame) { + TORCH_CHECK(cpuFrame != nullptr, "CPU frame cannot be null"); + + int width = cpuFrame->width; + int height = cpuFrame->height; + + // Intermediate NV12 CPU frame + UniqueAVFrame nv12CpuFrame(av_frame_alloc()); + TORCH_CHECK(nv12CpuFrame != nullptr, "Failed to allocate NV12 CPU frame"); + + nv12CpuFrame->format = AV_PIX_FMT_NV12; + nv12CpuFrame->width = width; + nv12CpuFrame->height = height; + + int ret = av_frame_get_buffer(nv12CpuFrame.get(), 0); + TORCH_CHECK( + ret >= 0, + "Failed to allocate NV12 CPU frame buffer: ", + getFFMPEGErrorStringFromErrorCode(ret)); + + SwsFrameContext swsFrameContext( + width, + height, + static_cast(cpuFrame->format), + width, + height); + + if (!swsContext_ || prevSwsFrameContext_ != swsFrameContext) { + swsContext_ = createSwsContext( + swsFrameContext, cpuFrame->colorspace, AV_PIX_FMT_NV12, SWS_BILINEAR); + prevSwsFrameContext_ = swsFrameContext; + } + + int convertedHeight = sws_scale( + swsContext_.get(), + cpuFrame->data, + cpuFrame->linesize, + 0, + height, + nv12CpuFrame->data, + nv12CpuFrame->linesize); + TORCH_CHECK( + convertedHeight == height, "sws_scale failed for CPU->NV12 conversion"); + + int ySize = width * height; + TORCH_CHECK(ySize % 2 == 0, "Y plane size must be even."); + int uvSize = ySize / 2; + size_t totalSize = static_cast(ySize + uvSize); + + uint8_t* hipBuffer = nullptr; + hipError_t err = + hipMalloc(reinterpret_cast(&hipBuffer), totalSize); + TORCH_CHECK( + err == hipSuccess, + "Failed to allocate HIP memory: ", + hipGetErrorString(err)); + + UniqueAVFrame gpuFrame(av_frame_alloc()); + TORCH_CHECK(gpuFrame != nullptr, "Failed to allocate GPU AVFrame"); + + gpuFrame->format = AV_PIX_FMT_CUDA; + gpuFrame->width = width; + gpuFrame->height = height; + gpuFrame->data[0] = hipBuffer; + gpuFrame->data[1] = hipBuffer + ySize; + gpuFrame->linesize[0] = width; + gpuFrame->linesize[1] = width; + + // Copy Y plane + err = hipMemcpy2D( + gpuFrame->data[0], + gpuFrame->linesize[0], + nv12CpuFrame->data[0], + nv12CpuFrame->linesize[0], + width, + height, + hipMemcpyHostToDevice); + TORCH_CHECK( + err == hipSuccess, + "Failed to copy Y plane to GPU: ", + hipGetErrorString(err)); + + // Copy UV plane + TORCH_CHECK(height % 2 == 0, "Height must be even."); + err = hipMemcpy2D( + gpuFrame->data[1], + gpuFrame->linesize[1], + nv12CpuFrame->data[1], + nv12CpuFrame->linesize[1], + width, + height / 2, + hipMemcpyHostToDevice); + TORCH_CHECK( + err == hipSuccess, + "Failed to copy UV plane to GPU: ", + hipGetErrorString(err)); + + ret = av_frame_copy_props(gpuFrame.get(), cpuFrame.get()); + TORCH_CHECK( + ret >= 0, + "Failed to copy frame properties: ", + getFFMPEGErrorStringFromErrorCode(ret)); + + // Associate cleanup callback to free HIP memory when AVFrame is freed + gpuFrame->opaque_ref = av_buffer_create( + nullptr, + 0, + hipBufferFreeCallback, + hipBuffer, + 0); + TORCH_CHECK( + gpuFrame->opaque_ref != nullptr, + "Failed to create GPU memory cleanup reference"); + + return gpuFrame; +} + +// ---- Color conversion: NV12 -> RGB tensor ---- + +void RocDecodeDeviceInterface::convertAVFrameToFrameOutput( + UniqueAVFrame& avFrame, + FrameOutput& frameOutput, + std::optional preAllocatedOutputTensor) { + UniqueAVFrame gpuFrame = + cpuFallback_ ? transferCpuFrameToGpuNV12(avFrame) : std::move(avFrame); + + TORCH_CHECK( + gpuFrame->format == AV_PIX_FMT_CUDA, + "Expected CUDA/HIP format frame from rocDecode interface"); + + validatePreAllocatedTensorShape_HIP( + preAllocatedOutputTensor, gpuFrame->width, gpuFrame->height); + + // Get current HIP stream for synchronization + hipStream_t currentStream = + c10::hip::getCurrentHIPStream(device_.index()).stream(); + + frameOutput.data = convertNV12FrameToRGB_HIP( + gpuFrame->data[0], + gpuFrame->linesize[0], + gpuFrame->width, + gpuFrame->height, + gpuFrame->colorspace, + gpuFrame->color_range, + device_, + currentStream, + preAllocatedOutputTensor); +} + +// ---- Details ---- + +std::string RocDecodeDeviceInterface::getDetails() { + std::string details = "rocDecode Device Interface."; + if (cpuFallback_) { + details += " Using CPU fallback."; + if (!rocDecodeAvailable_) { + details += " rocDecode not available!"; + } + } else { + details += " Using AMD VCN hardware decoder."; + } + return details; +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/RocDecodeDeviceInterface.h b/src/torchcodec/_core/RocDecodeDeviceInterface.h new file mode 100644 index 000000000..9f41cf869 --- /dev/null +++ b/src/torchcodec/_core/RocDecodeDeviceInterface.h @@ -0,0 +1,108 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// ROCm device interface for TorchCodec using AMD's rocDecode library. +// This is the AMD equivalent of BetaCudaDeviceInterface - it uses +// rocDecode's hardware video decoder (VCN) instead of NVIDIA's NVDEC. +// +// The design closely mirrors BetaCudaDeviceInterface to minimize divergence: +// - Same parser callback pattern (sequence/decode/display) +// - Same send/receive packet architecture +// - Same CPU fallback mechanism +// - Same frame ordering via display callback queue +// +// Key differences from the NVIDIA path: +// - Uses rocDecode API instead of NVCUVID +// - Uses HIP kernels for color conversion instead of NPP +// - rocDecGetVideoFrame returns per-plane device pointers (not CUdeviceptr) +// - rocDecGetVideoFrame is a blocking call (vs cuvidMapVideoFrame) +// - No map/unmap pattern - frames are copied or used directly + +#pragma once + +#include "DeviceInterface.h" +#include "FFMPEGCommon.h" +#include "HIPCommon.h" +#include "RocDecodeCache.h" + +#include "rocdecode_include/rocdecode.h" +#include "rocdecode_include/rocparser.h" + +#include +#include +#include + +namespace facebook::torchcodec { + +class RocDecodeDeviceInterface : public DeviceInterface { + public: + explicit RocDecodeDeviceInterface(const torch::Device& device); + virtual ~RocDecodeDeviceInterface(); + + void initialize( + const AVStream* avStream, + const UniqueDecodingAVFormatContext& avFormatCtx, + const SharedAVCodecContext& codecContext) override; + + void convertAVFrameToFrameOutput( + UniqueAVFrame& avFrame, + FrameOutput& frameOutput, + std::optional preAllocatedOutputTensor) override; + + int sendPacket(ReferenceAVPacket& packet) override; + int sendEOFPacket() override; + int receiveFrame(UniqueAVFrame& avFrame) override; + void flush() override; + + // rocDecode callback functions (must be public for C callbacks) + int streamPropertyChange(RocdecVideoFormat* videoFormat); + int frameReadyForDecoding(RocdecPicParams* picParams); + int frameReadyInDisplayOrder(RocdecParserDispInfo* dispInfo); + + std::string getDetails() override; + + private: + int sendRocDecPacket(RocdecSourceDataPacket& packet); + + void initializeBSF( + const AVCodecParameters* codecPar, + const UniqueDecodingAVFormatContext& avFormatCtx); + + ReferenceAVPacket& applyBSF( + ReferenceAVPacket& packet, + ReferenceAVPacket& filteredPacket); + + UniqueAVFrame convertRocDecFrameToAVFrame( + void* devMemPtr[3], + uint32_t pitch, + const RocdecParserDispInfo& dispInfo); + + UniqueAVFrame transferCpuFrameToGpuNV12(UniqueAVFrame& cpuFrame); + + static UniqueRocDecoder createDecoder( + RocdecVideoFormat* videoFormat, + int deviceId); + + RocdecVideoParser videoParser_ = nullptr; + UniqueRocDecoder decoder_; + RocdecVideoFormat videoFormat_ = {}; + + std::queue readyFrames_; + + bool eofSent_ = false; + + AVRational timeBase_ = {0, 1}; + AVRational frameRateAvgFromFFmpeg_ = {0, 1}; + + UniqueAVBSFContext bitstreamFilter_; + + std::unique_ptr cpuFallback_; + bool rocDecodeAvailable_ = false; + UniqueSwsContext swsContext_; + SwsFrameContext prevSwsFrameContext_; +}; + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/RocDecodeRuntimeLoader.cpp b/src/torchcodec/_core/RocDecodeRuntimeLoader.cpp new file mode 100644 index 000000000..daaf26e1e --- /dev/null +++ b/src/torchcodec/_core/RocDecodeRuntimeLoader.cpp @@ -0,0 +1,261 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "RocDecodeRuntimeLoader.h" + +#include "rocdecode_include/rocdecode.h" +#include "rocdecode_include/rocparser.h" + +#include +#include +#include + +#include +typedef void* tHandle; + +namespace facebook::torchcodec { + +/* clang-format off */ +// This file defines the logic to load the rocDecode library **at runtime**, +// along with the corresponding rocDecode functions that we'll need. +// +// We do this because we *do not want* to link (statically or dynamically) +// against librocdecode.so: it is not always available on the users machine! +// If we were to link against librocdecode.so, that would mean that our +// libtorchcodec_coreN.so would try to look for it when loaded at import time. +// And if it's not on the users machine, that causes `import torchcodec` to +// fail. +// +// This mirrors exactly the pattern used in NVCUVIDRuntimeLoader.cpp for +// libnvcuvid.so. See that file for a detailed explanation of the technique. +// +// At runtime, when the RocDecode device interface is first created, we call +// loadRocDecodeLibrary() which dlopen()s librocdecode.so and binds all +// function pointers via dlsym(). If the library is not found, we fall back +// to CPU decoding. + +// ---- Function pointer types for rocDecode API ---- +typedef rocDecStatus ROCDECAPI trocDecCreateDecoder(rocDecDecoderHandle*, RocDecoderCreateInfo*); +typedef rocDecStatus ROCDECAPI trocDecDestroyDecoder(rocDecDecoderHandle); +typedef rocDecStatus ROCDECAPI trocDecGetDecoderCaps(RocdecDecodeCaps*); +typedef rocDecStatus ROCDECAPI trocDecDecodeFrame(rocDecDecoderHandle, RocdecPicParams*); +typedef rocDecStatus ROCDECAPI trocDecGetDecodeStatus(rocDecDecoderHandle, int, RocdecDecodeStatus*); +typedef rocDecStatus ROCDECAPI trocDecReconfigureDecoder(rocDecDecoderHandle, RocdecReconfigureDecoderInfo*); +typedef rocDecStatus ROCDECAPI trocDecGetVideoFrame(rocDecDecoderHandle, int, void*[3], uint32_t*, RocdecProcParams*); +typedef const char* ROCDECAPI trocDecGetErrorName(rocDecStatus); + +// ---- Function pointer types for rocParser API ---- +typedef rocDecStatus ROCDECAPI trocDecCreateVideoParser(RocdecVideoParser*, RocdecParserParams*); +typedef rocDecStatus ROCDECAPI trocDecParseVideoData(RocdecVideoParser, RocdecSourceDataPacket*); +typedef rocDecStatus ROCDECAPI trocDecDestroyVideoParser(RocdecVideoParser); +/* clang-format on */ + +// Global function pointers - will be dynamically loaded +static trocDecCreateDecoder* dl_rocDecCreateDecoder = nullptr; +static trocDecDestroyDecoder* dl_rocDecDestroyDecoder = nullptr; +static trocDecGetDecoderCaps* dl_rocDecGetDecoderCaps = nullptr; +static trocDecDecodeFrame* dl_rocDecDecodeFrame = nullptr; +static trocDecGetDecodeStatus* dl_rocDecGetDecodeStatus = nullptr; +static trocDecReconfigureDecoder* dl_rocDecReconfigureDecoder = nullptr; +static trocDecGetVideoFrame* dl_rocDecGetVideoFrame = nullptr; +static trocDecGetErrorName* dl_rocDecGetErrorName = nullptr; + +static trocDecCreateVideoParser* dl_rocDecCreateVideoParser = nullptr; +static trocDecParseVideoData* dl_rocDecParseVideoData = nullptr; +static trocDecDestroyVideoParser* dl_rocDecDestroyVideoParser = nullptr; + +static tHandle g_rocdecode_handle = nullptr; +static tHandle g_rocparser_handle = nullptr; +static std::mutex g_rocdecode_mutex; + +bool isLoaded() { + return ( + g_rocdecode_handle && dl_rocDecCreateDecoder && dl_rocDecDestroyDecoder && + dl_rocDecGetDecoderCaps && dl_rocDecDecodeFrame && + dl_rocDecGetVideoFrame && dl_rocDecGetErrorName && + dl_rocDecCreateVideoParser && dl_rocDecParseVideoData && + dl_rocDecDestroyVideoParser); +} + +template +T* bindFunction(tHandle handle, const char* functionName) { + return reinterpret_cast(dlsym(handle, functionName)); +} + +bool _loadLibraries() { + // Try versioned names first (future-proof for soname bumps), + // then unversioned. This ensures we pick up whatever rocDecode + // is installed regardless of the ROCm version. + const char* decoder_lib_names[] = { + "librocdecode.so", // unversioned (typical for dev installs) + "librocdecode.so.0", // soversion 0 + "librocdecode.so.1", // future soversion bump + nullptr}; + + for (const char** name = decoder_lib_names; *name != nullptr; ++name) { + g_rocdecode_handle = dlopen(*name, RTLD_NOW); + if (g_rocdecode_handle != nullptr) { + break; + } + } + if (g_rocdecode_handle == nullptr) { + return false; + } + + // The parser functions may be in the same library or a separate one. + // Try the main library first, then try a separate parser library. + g_rocparser_handle = g_rocdecode_handle; + + return true; +} + +bool loadRocDecodeLibrary() { + std::lock_guard lock(g_rocdecode_mutex); + + if (isLoaded()) { + return true; + } + + if (!_loadLibraries()) { + return false; + } + + // Load decoder function pointers + dl_rocDecCreateDecoder = + bindFunction(g_rocdecode_handle, "rocDecCreateDecoder"); + dl_rocDecDestroyDecoder = + bindFunction(g_rocdecode_handle, "rocDecDestroyDecoder"); + dl_rocDecGetDecoderCaps = + bindFunction(g_rocdecode_handle, "rocDecGetDecoderCaps"); + dl_rocDecDecodeFrame = + bindFunction(g_rocdecode_handle, "rocDecDecodeFrame"); + dl_rocDecGetDecodeStatus = + bindFunction(g_rocdecode_handle, "rocDecGetDecodeStatus"); + dl_rocDecReconfigureDecoder = + bindFunction(g_rocdecode_handle, "rocDecReconfigureDecoder"); + dl_rocDecGetVideoFrame = + bindFunction(g_rocdecode_handle, "rocDecGetVideoFrame"); + dl_rocDecGetErrorName = + bindFunction(g_rocdecode_handle, "rocDecGetErrorName"); + + // Load parser function pointers + dl_rocDecCreateVideoParser = + bindFunction(g_rocparser_handle, "rocDecCreateVideoParser"); + dl_rocDecParseVideoData = + bindFunction(g_rocparser_handle, "rocDecParseVideoData"); + dl_rocDecDestroyVideoParser = + bindFunction(g_rocparser_handle, "rocDecDestroyVideoParser"); + + return isLoaded(); +} + +} // namespace facebook::torchcodec + +// Actual function definitions that forward to the dynamically loaded pointers. +// These are compiled against and called by the RocDecode device interface code. +extern "C" { + +rocDecStatus ROCDECAPI +rocDecCreateDecoder(rocDecDecoderHandle* decoder_handle, + RocDecoderCreateInfo* decoder_create_info) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecCreateDecoder, + "rocDecCreateDecoder called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecCreateDecoder( + decoder_handle, decoder_create_info); +} + +rocDecStatus ROCDECAPI +rocDecDestroyDecoder(rocDecDecoderHandle decoder_handle) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecDestroyDecoder, + "rocDecDestroyDecoder called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecDestroyDecoder(decoder_handle); +} + +rocDecStatus ROCDECAPI rocDecGetDecoderCaps(RocdecDecodeCaps* pdc) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecGetDecoderCaps, + "rocDecGetDecoderCaps called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecGetDecoderCaps(pdc); +} + +rocDecStatus ROCDECAPI +rocDecDecodeFrame(rocDecDecoderHandle decoder_handle, + RocdecPicParams* pic_params) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecDecodeFrame, + "rocDecDecodeFrame called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecDecodeFrame(decoder_handle, pic_params); +} + +rocDecStatus ROCDECAPI +rocDecGetDecodeStatus(rocDecDecoderHandle decoder_handle, int pic_idx, + RocdecDecodeStatus* decode_status) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecGetDecodeStatus, + "rocDecGetDecodeStatus called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecGetDecodeStatus( + decoder_handle, pic_idx, decode_status); +} + +rocDecStatus ROCDECAPI +rocDecReconfigureDecoder(rocDecDecoderHandle decoder_handle, + RocdecReconfigureDecoderInfo* reconfig_params) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecReconfigureDecoder, + "rocDecReconfigureDecoder called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecReconfigureDecoder( + decoder_handle, reconfig_params); +} + +rocDecStatus ROCDECAPI +rocDecGetVideoFrame(rocDecDecoderHandle decoder_handle, int pic_idx, + void* dev_mem_ptr[3], uint32_t* horizontal_pitch, + RocdecProcParams* vid_postproc_params) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecGetVideoFrame, + "rocDecGetVideoFrame called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecGetVideoFrame( + decoder_handle, pic_idx, dev_mem_ptr, horizontal_pitch, + vid_postproc_params); +} + +const char* ROCDECAPI rocDecGetErrorName(rocDecStatus rocdec_status) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecGetErrorName, + "rocDecGetErrorName called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecGetErrorName(rocdec_status); +} + +rocDecStatus ROCDECAPI +rocDecCreateVideoParser(RocdecVideoParser* parser_handle, + RocdecParserParams* params) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecCreateVideoParser, + "rocDecCreateVideoParser called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecCreateVideoParser( + parser_handle, params); +} + +rocDecStatus ROCDECAPI +rocDecParseVideoData(RocdecVideoParser parser_handle, + RocdecSourceDataPacket* packet) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecParseVideoData, + "rocDecParseVideoData called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecParseVideoData(parser_handle, packet); +} + +rocDecStatus ROCDECAPI +rocDecDestroyVideoParser(RocdecVideoParser parser_handle) { + TORCH_CHECK( + facebook::torchcodec::dl_rocDecDestroyVideoParser, + "rocDecDestroyVideoParser called but rocDecode not loaded!"); + return facebook::torchcodec::dl_rocDecDestroyVideoParser(parser_handle); +} + +} // extern "C" diff --git a/src/torchcodec/_core/RocDecodeRuntimeLoader.h b/src/torchcodec/_core/RocDecodeRuntimeLoader.h new file mode 100644 index 000000000..2edfbb5af --- /dev/null +++ b/src/torchcodec/_core/RocDecodeRuntimeLoader.h @@ -0,0 +1,17 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +namespace facebook::torchcodec { + +// Dynamically loads librocdecode.so and librocdec_parser.so at runtime. +// Returns true if all required functions were successfully loaded. +// See the corresponding .cpp for the full design rationale, which mirrors +// the NVCUVID runtime loader pattern. +bool loadRocDecodeLibrary(); + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/rocdecode_include/rocdecode.h b/src/torchcodec/_core/rocdecode_include/rocdecode.h new file mode 100644 index 000000000..050435546 --- /dev/null +++ b/src/torchcodec/_core/rocdecode_include/rocdecode.h @@ -0,0 +1,957 @@ +// Vendored rocDecode API header for TorchCodec ROCm backend. +// +// This is a copy of the rocDecode v1.5.0 public API header that allows +// compilation without requiring rocDecode headers to be installed at build +// time; the actual library is loaded at runtime via dlopen(). +// +// IMPORTANT: The struct layouts in this file MUST exactly match the installed +// rocDecode library (ABI compatibility). If rocDecode is updated, this file +// must be updated to match. Struct layout mismatches cause GPU memory access +// faults and crashes. +// +// Source: rocDecode v1.5.0 (ROCm 7.2) +// https://github.com/ROCm/rocDecode +// +// Copyright (c) 2023 - 2025 Advanced Micro Devices, Inc. All rights reserved. +// MIT License + +#pragma once + +#include + +#ifndef ROCDECAPI +#if defined(_WIN32) +#define ROCDECAPI __stdcall +#else +#define ROCDECAPI +#endif +#endif + +#if defined(__cplusplus) +extern "C" { +#endif + +/*********************************************************************************/ +// HANDLE of rocDecDecoder +// Used in subsequent API calls after rocDecCreateDecoder +/*********************************************************************************/ +typedef void* rocDecDecoderHandle; + +/*********************************************************************************/ +// rocDecoder return status enums +/*********************************************************************************/ +typedef enum rocDecStatus_enum { + ROCDEC_DEVICE_INVALID = -1, + ROCDEC_CONTEXT_INVALID = -2, + ROCDEC_RUNTIME_ERROR = -3, + ROCDEC_OUTOF_MEMORY = -4, + ROCDEC_INVALID_PARAMETER = -5, + ROCDEC_NOT_IMPLEMENTED = -6, + ROCDEC_NOT_INITIALIZED = -7, + ROCDEC_NOT_SUPPORTED = -8, + ROCDEC_SUCCESS = 0, +} rocDecStatus; + +/*********************************************************************************/ +// Video codec enums +/*********************************************************************************/ +typedef enum rocDecVideoCodec_enum { + rocDecVideoCodec_MPEG1 = 0, + rocDecVideoCodec_MPEG2, + rocDecVideoCodec_MPEG4, + rocDecVideoCodec_AVC, + rocDecVideoCodec_HEVC, + rocDecVideoCodec_AV1, + rocDecVideoCodec_VP8, + rocDecVideoCodec_VP9, + rocDecVideoCodec_JPEG, + rocDecVideoCodec_NumCodecs, + rocDecVideoCodec_YUV420 = + (('I' << 24) | ('Y' << 16) | ('U' << 8) | ('V')), + rocDecVideoCodec_YV12 = + (('Y' << 24) | ('V' << 16) | ('1' << 8) | ('2')), + rocDecVideoCodec_NV12 = + (('N' << 24) | ('V' << 16) | ('1' << 8) | ('2')), + rocDecVideoCodec_YUYV = + (('Y' << 24) | ('U' << 16) | ('Y' << 8) | ('V')), + rocDecVideoCodec_UYVY = + (('U' << 24) | ('Y' << 16) | ('V' << 8) | ('Y')) +} rocDecVideoCodec; + +/*********************************************************************************/ +// Video surface format enums used for output format of decoded output +/*********************************************************************************/ +typedef enum rocDecVideoSurfaceFormat_enum { + rocDecVideoSurfaceFormat_NV12 = 0, + rocDecVideoSurfaceFormat_P016 = 1, + rocDecVideoSurfaceFormat_YUV444 = 2, + rocDecVideoSurfaceFormat_YUV444_16Bit = 3, + rocDecVideoSurfaceFormat_YUV420 = 4, + rocDecVideoSurfaceFormat_YUV420_16Bit = 5, + rocDecVideoSurfaceFormat_YUV422 = 6, + rocDecVideoSurfaceFormat_YUV422_16Bit = 7, +} rocDecVideoSurfaceFormat; + +/*********************************************************************************/ +// Chroma format enums +/*********************************************************************************/ +typedef enum rocDecVideoChromaFormat_enum { + rocDecVideoChromaFormat_Monochrome = 0, + rocDecVideoChromaFormat_420, + rocDecVideoChromaFormat_422, + rocDecVideoChromaFormat_444 +} rocDecVideoChromaFormat; + +/*********************************************************************************/ +// Decode status enums +/*********************************************************************************/ +typedef enum rocDecodeStatus_enum { + rocDecodeStatus_Invalid = 0, + rocDecodeStatus_InProgress = 1, + rocDecodeStatus_Success = 2, + // 3 to 7 reserved for future use + rocDecodeStatus_Error = 8, + rocDecodeStatus_Error_Concealed = 9, + rocDecodeStatus_Displaying = 10, +} rocDecDecodeStatus; + +/*********************************************************************************/ +// RocdecDecodeCaps - used in rocDecGetDecoderCaps API +/*********************************************************************************/ +typedef struct _RocdecDecodeCaps { + uint8_t device_id; /**< IN: device id (0 for first, 1 for second, etc.) */ + rocDecVideoCodec codec_type; /**< IN: rocDecVideoCodec_XXX */ + rocDecVideoChromaFormat chroma_format; /**< IN: rocDecVideoChromaFormat_XXX */ + uint32_t bit_depth_minus_8; /**< IN: The Value "BitDepth minus 8" */ + uint32_t reserved_1[3]; /**< Reserved for future use - set to zero */ + uint8_t is_supported; /**< OUT: 1 if codec supported, 0 if not */ + uint8_t num_decoders; /**< OUT: Number of decoders supporting IN params */ + uint16_t output_format_mask; /**< OUT: each bit represents rocDecVideoSurfaceFormat enum */ + uint32_t max_width; /**< OUT: Max supported coded width in pixels */ + uint32_t max_height; /**< OUT: Max supported coded height in pixels */ + uint16_t min_width; /**< OUT: Min supported coded width in pixels */ + uint16_t min_height; /**< OUT: Min supported coded height in pixels */ + uint32_t reserved_2[6]; /**< Reserved for future use - set to zero */ +} RocdecDecodeCaps; + +/*********************************************************************************/ +// RocDecoderCreateInfo - used in rocDecCreateDecoder API +/*********************************************************************************/ +typedef struct _RocDecoderCreateInfo { + uint8_t device_id; /**< IN: device id (0 for first, 1 for second, etc.) */ + uint32_t width; /**< IN: Coded sequence width in pixels */ + uint32_t height; /**< IN: Coded sequence height in pixels */ + uint32_t num_decode_surfaces; /**< IN: Maximum number of internal decode surfaces */ + rocDecVideoCodec codec_type; /**< IN: rocDecVideoCodec_XXX */ + rocDecVideoChromaFormat chroma_format; /**< IN: rocDecVideoChromaFormat_XXX */ + uint32_t bit_depth_minus_8; /**< IN: The value "BitDepth minus 8" */ + uint32_t intra_decode_only; /**< IN: Set 1 only if video has all intra frames */ + uint32_t max_width; /**< IN: Coded sequence max width for reconfigure */ + uint32_t max_height; /**< IN: Coded sequence max height for reconfigure */ + struct { + int16_t left; + int16_t top; + int16_t right; + int16_t bottom; + } display_rect; /**< IN: area of the frame to display */ + rocDecVideoSurfaceFormat output_format; /**< IN: rocDecVideoSurfaceFormat_XXX */ + uint32_t target_width; /**< IN: Post-processed output width (aligned to 2) */ + uint32_t target_height; /**< IN: Post-processed output height (aligned to 2) */ + uint32_t num_output_surfaces; /**< IN: Max number of output surfaces mapped */ + struct { + int16_t left; + int16_t top; + int16_t right; + int16_t bottom; + } target_rect; /**< IN: target rectangle in output frame */ + uint32_t reserved_2[4]; /**< Reserved for future use - set to zero */ +} RocDecoderCreateInfo; + +/*********************************************************************************/ +// RocdecDecodeStatus - used in rocDecGetDecodeStatus API +/*********************************************************************************/ +typedef struct _RocdecDecodeStatus { + rocDecDecodeStatus decode_status; + uint32_t reserved[31]; + void* p_reserved[8]; +} RocdecDecodeStatus; + +/*********************************************************************************/ +// RocdecReconfigureDecoderInfo - used in rocDecReconfigureDecoder API +/*********************************************************************************/ +typedef struct _RocdecReconfigureDecoderInfo { + uint32_t width; /**< IN: Coded sequence width, MUST be <= max_width */ + uint32_t height; /**< IN: Coded sequence height, MUST be <= max_height */ + uint32_t target_width; /**< IN: Post processed output width */ + uint32_t target_height; /**< IN: Post processed output height */ + uint32_t num_decode_surfaces; /**< IN: Maximum number of internal decode surfaces */ + uint32_t bit_depth_minus_8; /**< IN: The Value "BitDepth minus 8" */ + uint32_t reserved_1[11]; /**< Reserved for future use. Set to zero */ + struct { + int16_t left; + int16_t top; + int16_t right; + int16_t bottom; + } display_rect; /**< IN: area of the frame to display */ + struct { + int16_t left; + int16_t top; + int16_t right; + int16_t bottom; + } target_rect; /**< IN: target rectangle in output frame */ + uint32_t reserved_2[11]; /**< Reserved for future use. Set to zero */ +} RocdecReconfigureDecoderInfo; + +/*********************************************************************************/ +// AVC/H.264 Picture Entry +/*********************************************************************************/ +typedef struct _RocdecAvcPicture { + int pic_idx; + uint32_t frame_idx; + uint32_t flags; + int32_t top_field_order_cnt; + int32_t bottom_field_order_cnt; + uint32_t reserved[4]; +} RocdecAvcPicture; + +#define RocdecAvcPicture_FLAGS_INVALID 0x00000001 +#define RocdecAvcPicture_FLAGS_TOP_FIELD 0x00000002 +#define RocdecAvcPicture_FLAGS_BOTTOM_FIELD 0x00000004 +#define RocdecAvcPicture_FLAGS_SHORT_TERM_REFERENCE 0x00000008 +#define RocdecAvcPicture_FLAGS_LONG_TERM_REFERENCE 0x00000010 +#define RocdecAvcPicture_FLAGS_NON_EXISTING 0x00000020 + +/*********************************************************************************/ +// HEVC Picture Entry +/*********************************************************************************/ +typedef struct _RocdecHevcPicture { + int pic_idx; + int poc; + uint32_t flags; + uint32_t reserved[4]; +} RocdecHevcPicture; + +#define RocdecHevcPicture_INVALID 0x00000001 +#define RocdecHevcPicture_FIELD_PIC 0x00000002 +#define RocdecHevcPicture_BOTTOM_FIELD 0x00000004 +#define RocdecHevcPicture_LONG_TERM_REFERENCE 0x00000008 +#define RocdecHevcPicture_RPS_ST_CURR_BEFORE 0x00000010 +#define RocdecHevcPicture_RPS_ST_CURR_AFTER 0x00000020 +#define RocdecHevcPicture_RPS_LT_CURR 0x00000040 + +/*********************************************************************************/ +// JPEG picture parameters (placeholder) +/*********************************************************************************/ +typedef struct _RocdecJPEGPicParams { + int reserved; +} RocdecJPEGPicParams; + +/*********************************************************************************/ +// MPEG2 QMatrix +/*********************************************************************************/ +typedef struct _RocdecMpeg2QMatrix { + int32_t load_intra_quantiser_matrix; + int32_t load_non_intra_quantiser_matrix; + int32_t load_chroma_intra_quantiser_matrix; + int32_t load_chroma_non_intra_quantiser_matrix; + uint8_t intra_quantiser_matrix[64]; + uint8_t non_intra_quantiser_matrix[64]; + uint8_t chroma_intra_quantiser_matrix[64]; + uint8_t chroma_non_intra_quantiser_matrix[64]; +} RocdecMpeg2QMatrix; + +/*********************************************************************************/ +// MPEG2 picture parameters +/*********************************************************************************/ +typedef struct _RocdecMpeg2PicParams { + uint16_t horizontal_size; + uint16_t vertical_size; + uint32_t forward_reference_pic; + uint32_t backward_reference_picture; + int32_t picture_coding_type; + int32_t f_code; + union { + struct { + uint32_t intra_dc_precision : 2; + uint32_t picture_structure : 2; + uint32_t top_field_first : 1; + uint32_t frame_pred_frame_dct : 1; + uint32_t concealment_motion_vectors : 1; + uint32_t q_scale_type : 1; + uint32_t intra_vlc_format : 1; + uint32_t alternate_scan : 1; + uint32_t repeat_first_field : 1; + uint32_t progressive_frame : 1; + uint32_t is_first_field : 1; + } bits; + uint32_t value; + } picture_coding_extension; + RocdecMpeg2QMatrix q_matrix; + uint32_t reserved[4]; +} RocdecMpeg2PicParams; + +/*********************************************************************************/ +// VC1 picture parameters (placeholder) +/*********************************************************************************/ +typedef struct _RocdecVc1PicParams { + int reserved; +} RocdecVc1PicParams; + +/*********************************************************************************/ +// AVC picture parameters (VA-API compatible) +/*********************************************************************************/ +typedef struct _RocdecAvcPicParams { + RocdecAvcPicture curr_pic; + RocdecAvcPicture ref_frames[16]; + uint16_t picture_width_in_mbs_minus1; + uint16_t picture_height_in_mbs_minus1; + uint8_t bit_depth_luma_minus8; + uint8_t bit_depth_chroma_minus8; + uint8_t num_ref_frames; + union { + struct { + uint32_t chroma_format_idc : 2; + uint32_t residual_colour_transform_flag : 1; + uint32_t gaps_in_frame_num_value_allowed_flag : 1; + uint32_t frame_mbs_only_flag : 1; + uint32_t mb_adaptive_frame_field_flag : 1; + uint32_t direct_8x8_inference_flag : 1; + uint32_t MinLumaBiPredSize8x8 : 1; + uint32_t log2_max_frame_num_minus4 : 4; + uint32_t pic_order_cnt_type : 2; + uint32_t log2_max_pic_order_cnt_lsb_minus4 : 4; + uint32_t delta_pic_order_always_zero_flag : 1; + } bits; + uint32_t value; + } seq_fields; + uint8_t num_slice_groups_minus1; + uint8_t slice_group_map_type; + uint16_t slice_group_change_rate_minus1; + int8_t pic_init_qp_minus26; + int8_t pic_init_qs_minus26; + int8_t chroma_qp_index_offset; + int8_t second_chroma_qp_index_offset; + union { + struct { + uint32_t entropy_coding_mode_flag : 1; + uint32_t weighted_pred_flag : 1; + uint32_t weighted_bipred_idc : 2; + uint32_t transform_8x8_mode_flag : 1; + uint32_t field_pic_flag : 1; + uint32_t constrained_intra_pred_flag : 1; + uint32_t pic_order_present_flag : 1; + uint32_t deblocking_filter_control_present_flag : 1; + uint32_t redundant_pic_cnt_present_flag : 1; + uint32_t reference_pic_flag : 1; + } bits; + uint32_t value; + } pic_fields; + uint16_t frame_num; + uint32_t reserved[8]; +} RocdecAvcPicParams; + +/*********************************************************************************/ +// AVC slice parameters (VA-API compatible) +/*********************************************************************************/ +typedef struct _RocdecAvcSliceParams { + uint32_t slice_data_size; + uint32_t slice_data_offset; + uint32_t slice_data_flag; + uint16_t slice_data_bit_offset; + uint16_t first_mb_in_slice; + uint8_t slice_type; + uint8_t direct_spatial_mv_pred_flag; + uint8_t num_ref_idx_l0_active_minus1; + uint8_t num_ref_idx_l1_active_minus1; + uint8_t cabac_init_idc; + int8_t slice_qp_delta; + uint8_t disable_deblocking_filter_idc; + int8_t slice_alpha_c0_offset_div2; + int8_t slice_beta_offset_div2; + RocdecAvcPicture ref_pic_list_0[32]; + RocdecAvcPicture ref_pic_list_1[32]; + uint8_t luma_log2_weight_denom; + uint8_t chroma_log2_weight_denom; + uint8_t luma_weight_l0_flag; + int16_t luma_weight_l0[32]; + int16_t luma_offset_l0[32]; + uint8_t chroma_weight_l0_flag; + int16_t chroma_weight_l0[32][2]; + int16_t chroma_offset_l0[32][2]; + uint8_t luma_weight_l1_flag; + int16_t luma_weight_l1[32]; + int16_t luma_offset_l1[32]; + uint8_t chroma_weight_l1_flag; + int16_t chroma_weight_l1[32][2]; + int16_t chroma_offset_l1[32][2]; + uint32_t reserved[4]; +} RocdecAvcSliceParams; + +/*********************************************************************************/ +// AVC Inverse Quantization Matrix (VA-API compatible) +/*********************************************************************************/ +typedef struct _RocdecAvcIQMatrix { + uint8_t scaling_list_4x4[6][16]; + uint8_t scaling_list_8x8[2][64]; + uint32_t reserved[4]; +} RocdecAvcIQMatrix; + +/*********************************************************************************/ +// HEVC picture parameters +/*********************************************************************************/ +typedef struct _RocdecHevcPicParams { + RocdecHevcPicture curr_pic; + RocdecHevcPicture ref_frames[15]; + uint16_t picture_width_in_luma_samples; + uint16_t picture_height_in_luma_samples; + union { + struct { + uint32_t chroma_format_idc : 2; + uint32_t separate_colour_plane_flag : 1; + uint32_t pcm_enabled_flag : 1; + uint32_t scaling_list_enabled_flag : 1; + uint32_t transform_skip_enabled_flag : 1; + uint32_t amp_enabled_flag : 1; + uint32_t strong_intra_smoothing_enabled_flag : 1; + uint32_t sign_data_hiding_enabled_flag : 1; + uint32_t constrained_intra_pred_flag : 1; + uint32_t cu_qp_delta_enabled_flag : 1; + uint32_t weighted_pred_flag : 1; + uint32_t weighted_bipred_flag : 1; + uint32_t transquant_bypass_enabled_flag : 1; + uint32_t tiles_enabled_flag : 1; + uint32_t entropy_coding_sync_enabled_flag : 1; + uint32_t pps_loop_filter_across_slices_enabled_flag : 1; + uint32_t loop_filter_across_tiles_enabled_flag : 1; + uint32_t pcm_loop_filter_disabled_flag : 1; + uint32_t no_pic_reordering_flag : 1; + uint32_t no_bi_pred_flag : 1; + uint32_t reserved_bits : 11; + } bits; + uint32_t value; + } pic_fields; + uint8_t sps_max_dec_pic_buffering_minus1; + uint8_t bit_depth_luma_minus8; + uint8_t bit_depth_chroma_minus8; + uint8_t pcm_sample_bit_depth_luma_minus1; + uint8_t pcm_sample_bit_depth_chroma_minus1; + uint8_t log2_min_luma_coding_block_size_minus3; + uint8_t log2_diff_max_min_luma_coding_block_size; + uint8_t log2_min_luma_transform_block_size_minus2; + uint8_t log2_diff_max_min_luma_transform_block_size; + uint8_t log2_min_pcm_luma_coding_block_size_minus3; + uint8_t log2_diff_max_min_pcm_luma_coding_block_size; + uint8_t max_transform_hierarchy_depth_intra; + uint8_t max_transform_hierarchy_depth_inter; + int8_t init_qp_minus26; + uint8_t diff_cu_qp_delta_depth; + int8_t pps_cb_qp_offset; + int8_t pps_cr_qp_offset; + uint8_t log2_parallel_merge_level_minus2; + uint8_t num_tile_columns_minus1; + uint8_t num_tile_rows_minus1; + uint16_t column_width_minus1[19]; + uint16_t row_height_minus1[21]; + union { + struct { + uint32_t lists_modification_present_flag : 1; + uint32_t long_term_ref_pics_present_flag : 1; + uint32_t sps_temporal_mvp_enabled_flag : 1; + uint32_t cabac_init_present_flag : 1; + uint32_t output_flag_present_flag : 1; + uint32_t dependent_slice_segments_enabled_flag : 1; + uint32_t pps_slice_chroma_qp_offsets_present_flag : 1; + uint32_t sample_adaptive_offset_enabled_flag : 1; + uint32_t deblocking_filter_override_enabled_flag : 1; + uint32_t pps_disable_deblocking_filter_flag : 1; + uint32_t slice_segment_header_extension_present_flag : 1; + uint32_t rap_pic_flag : 1; + uint32_t idr_pic_flag : 1; + uint32_t intra_pic_flag : 1; + uint32_t reserved_bits : 18; + } bits; + uint32_t value; + } slice_parsing_fields; + uint8_t log2_max_pic_order_cnt_lsb_minus4; + uint8_t num_short_term_ref_pic_sets; + uint8_t num_long_term_ref_pic_sps; + uint8_t num_ref_idx_l0_default_active_minus1; + uint8_t num_ref_idx_l1_default_active_minus1; + int8_t pps_beta_offset_div2; + int8_t pps_tc_offset_div2; + uint8_t num_extra_slice_header_bits; + uint32_t st_rps_bits; + uint32_t reserved[8]; +} RocdecHevcPicParams; + +/*********************************************************************************/ +// HEVC slice parameters +/*********************************************************************************/ +typedef struct _RocdecHevcSliceParams { + uint32_t slice_data_size; + uint32_t slice_data_offset; + uint32_t slice_data_flag; + uint32_t slice_data_byte_offset; + uint32_t slice_segment_address; + uint8_t ref_pic_list[2][15]; + union { + uint32_t value; + struct { + uint32_t last_slice_of_pic : 1; + uint32_t dependent_slice_segment_flag : 1; + uint32_t slice_type : 2; + uint32_t color_plane_id : 2; + uint32_t slice_sao_luma_flag : 1; + uint32_t slice_sao_chroma_flag : 1; + uint32_t mvd_l1_zero_flag : 1; + uint32_t cabac_init_flag : 1; + uint32_t slice_temporal_mvp_enabled_flag : 1; + uint32_t slice_deblocking_filter_disabled_flag : 1; + uint32_t collocated_from_l0_flag : 1; + uint32_t slice_loop_filter_across_slices_enabled_flag : 1; + uint32_t reserved : 18; + } fields; + } long_slice_flags; + uint8_t collocated_ref_idx; + uint8_t num_ref_idx_l0_active_minus1; + uint8_t num_ref_idx_l1_active_minus1; + int8_t slice_qp_delta; + int8_t slice_cb_qp_offset; + int8_t slice_cr_qp_offset; + int8_t slice_beta_offset_div2; + int8_t slice_tc_offset_div2; + uint8_t luma_log2_weight_denom; + int8_t delta_chroma_log2_weight_denom; + int8_t delta_luma_weight_l0[15]; + int8_t luma_offset_l0[15]; + int8_t delta_chroma_weight_l0[15][2]; + int8_t chroma_offset_l0[15][2]; + int8_t delta_luma_weight_l1[15]; + int8_t luma_offset_l1[15]; + int8_t delta_chroma_weight_l1[15][2]; + int8_t chroma_offset_l1[15][2]; + uint8_t five_minus_max_num_merge_cand; + uint16_t num_entry_point_offsets; + uint16_t entry_offset_to_subset_array; + uint16_t slice_data_num_emu_prevn_bytes; + uint32_t reserved[2]; +} RocdecHevcSliceParams; + +/*********************************************************************************/ +// HEVC IQ Matrix +/*********************************************************************************/ +typedef struct _RocdecHevcIQMatrix { + uint8_t scaling_list_4x4[6][16]; + uint8_t scaling_list_8x8[6][64]; + uint8_t scaling_list_16x16[6][64]; + uint8_t scaling_list_32x32[2][64]; + uint8_t scaling_list_dc_16x16[6]; + uint8_t scaling_list_dc_32x32[2]; + uint32_t reserved[4]; +} RocdecHevcIQMatrix; + +/*********************************************************************************/ +// VP9 picture parameters (VA-API compatible) +/*********************************************************************************/ +typedef struct _RocdecVp9PicParams { + uint16_t frame_width; + uint16_t frame_height; + uint32_t reference_frames[8]; + union { + struct { + uint32_t subsampling_x : 1; + uint32_t subsampling_y : 1; + uint32_t frame_type : 1; + uint32_t show_frame : 1; + uint32_t error_resilient_mode : 1; + uint32_t intra_only : 1; + uint32_t allow_high_precision_mv : 1; + uint32_t mcomp_filter_type : 3; + uint32_t frame_parallel_decoding_mode : 1; + uint32_t reset_frame_context : 2; + uint32_t refresh_frame_context : 1; + uint32_t frame_context_idx : 2; + uint32_t segmentation_enabled : 1; + uint32_t segmentation_temporal_update : 1; + uint32_t segmentation_update_map : 1; + uint32_t last_ref_frame : 3; + uint32_t last_ref_frame_sign_bias : 1; + uint32_t golden_ref_frame : 3; + uint32_t golden_ref_frame_sign_bias : 1; + uint32_t alt_ref_frame : 3; + uint32_t alt_ref_frame_sign_bias : 1; + uint32_t lossless_flag : 1; + } bits; + uint32_t value; + } pic_fields; + uint8_t filter_level; + uint8_t sharpness_level; + uint8_t log2_tile_rows; + uint8_t log2_tile_columns; + uint8_t frame_header_length_in_bytes; + uint16_t first_partition_size; + uint8_t mb_segment_tree_probs[7]; + uint8_t segment_pred_probs[3]; + uint8_t profile; + uint8_t bit_depth; + uint32_t va_reserved[8]; +} RocdecVp9PicParams; + +/*********************************************************************************/ +// VP9 Segmentation Parameter +/*********************************************************************************/ +typedef struct _RocdecVp9SegmentParameter { + union { + struct { + uint16_t segment_reference_enabled : 1; + uint16_t segment_reference : 2; + uint16_t segment_reference_skipped : 1; + } fields; + uint16_t value; + } segment_flags; + uint8_t filter_level[4][2]; + int16_t luma_ac_quant_scale; + int16_t luma_dc_quant_scale; + int16_t chroma_ac_quant_scale; + int16_t chroma_dc_quant_scale; + uint32_t va_reserved[4]; +} RocdecVp9SegmentParameter; + +/*********************************************************************************/ +// VP9 slice parameters (VA-API compatible) +/*********************************************************************************/ +typedef struct _RocdecVp9SliceParams { + uint32_t slice_data_size; + uint32_t slice_data_offset; + uint32_t slice_data_flag; + RocdecVp9SegmentParameter seg_param[8]; + uint32_t va_reserved[4]; +} RocdecVp9SliceParams; + +/*********************************************************************************/ +// AV1 Segmentation Information +/*********************************************************************************/ +typedef struct _RocdecAv1SegmentationStruct { + union { + struct { + uint32_t enabled : 1; + uint32_t update_map : 1; + uint32_t temporal_update : 1; + uint32_t update_data : 1; + uint32_t reserved : 28; + } bits; + uint32_t value; + } segment_info_fields; + int16_t feature_data[8][8]; + uint8_t feature_mask[8]; + uint32_t reserved[4]; +} RocdecAv1SegmentationStruct; + +/*********************************************************************************/ +// AV1 Film Grain Information +/*********************************************************************************/ +typedef struct _RocdecAv1FilmGrainStruct { + union { + struct { + uint32_t apply_grain : 1; + uint32_t chroma_scaling_from_luma : 1; + uint32_t grain_scaling_minus_8 : 2; + uint32_t ar_coeff_lag : 2; + uint32_t ar_coeff_shift_minus_6 : 2; + uint32_t grain_scale_shift : 2; + uint32_t overlap_flag : 1; + uint32_t clip_to_restricted_range : 1; + uint32_t reserved : 20; + } bits; + uint32_t value; + } film_grain_info_fields; + uint16_t grain_seed; + uint8_t num_y_points; + uint8_t point_y_value[14]; + uint8_t point_y_scaling[14]; + uint8_t num_cb_points; + uint8_t point_cb_value[10]; + uint8_t point_cb_scaling[10]; + uint8_t num_cr_points; + uint8_t point_cr_value[10]; + uint8_t point_cr_scaling[10]; + int8_t ar_coeffs_y[24]; + int8_t ar_coeffs_cb[25]; + int8_t ar_coeffs_cr[25]; + uint8_t cb_mult; + uint8_t cb_luma_mult; + uint16_t cb_offset; + uint8_t cr_mult; + uint8_t cr_luma_mult; + uint16_t cr_offset; + uint32_t reserved[4]; +} RocdecAv1FilmGrainStruct; + +typedef enum { + RocdecAv1TransformationIdentity = 0, + RocdecAv1TransformationTranslation = 1, + RocdecAv1TransformationRotzoom = 2, + RocdecAv1TransformationAffine = 3, + RocdecAv1TransformationCount +} RocdecAv1TransformationType; + +typedef struct _RocdecAv1WarpedMotionParams { + RocdecAv1TransformationType wmtype; + int32_t wmmat[8]; + uint8_t invalid; + uint32_t reserved[4]; +} RocdecAv1WarpedMotionParams; + +/*********************************************************************************/ +// AV1 picture parameters +/*********************************************************************************/ +typedef struct _RocdecAV1PicParams { + uint8_t profile; + uint8_t order_hint_bits_minus_1; + uint8_t bit_depth_idx; + uint8_t matrix_coefficients; + union { + struct { + uint32_t still_picture : 1; + uint32_t use_128x128_superblock : 1; + uint32_t enable_filter_intra : 1; + uint32_t enable_intra_edge_filter : 1; + uint32_t enable_interintra_compound : 1; + uint32_t enable_masked_compound : 1; + uint32_t enable_dual_filter : 1; + uint32_t enable_order_hint : 1; + uint32_t enable_jnt_comp : 1; + uint32_t enable_cdef : 1; + uint32_t mono_chrome : 1; + uint32_t color_range : 1; + uint32_t subsampling_x : 1; + uint32_t subsampling_y : 1; + uint32_t chroma_sample_position : 1; + uint32_t film_grain_params_present : 1; + uint32_t reserved : 16; + } fields; + uint32_t value; + } seq_info_fields; + int current_frame; + int current_display_picture; + uint8_t anchor_frames_num; + int* anchor_frames_list; + uint16_t frame_width_minus1; + uint16_t frame_height_minus1; + uint16_t output_frame_width_in_tiles_minus_1; + uint16_t output_frame_height_in_tiles_minus_1; + int ref_frame_map[8]; + uint8_t ref_frame_idx[7]; + uint8_t primary_ref_frame; + uint8_t order_hint; + RocdecAv1SegmentationStruct seg_info; + RocdecAv1FilmGrainStruct film_grain_info; + uint8_t tile_cols; + uint8_t tile_rows; + uint16_t width_in_sbs_minus_1[63]; + uint16_t height_in_sbs_minus_1[63]; + uint16_t tile_count_minus_1; + uint16_t context_update_tile_id; + union { + struct { + uint32_t frame_type : 2; + uint32_t show_frame : 1; + uint32_t showable_frame : 1; + uint32_t error_resilient_mode : 1; + uint32_t disable_cdf_update : 1; + uint32_t allow_screen_content_tools : 1; + uint32_t force_integer_mv : 1; + uint32_t allow_intrabc : 1; + uint32_t use_superres : 1; + uint32_t allow_high_precision_mv : 1; + uint32_t is_motion_mode_switchable : 1; + uint32_t use_ref_frame_mvs : 1; + uint32_t disable_frame_end_update_cdf : 1; + uint32_t uniform_tile_spacing_flag : 1; + uint32_t allow_warped_motion : 1; + uint32_t large_scale_tile : 1; + uint32_t reserved : 15; + } bits; + uint32_t value; + } pic_info_fields; + uint8_t superres_scale_denominator; + uint8_t interp_filter; + uint8_t filter_level[2]; + uint8_t filter_level_u; + uint8_t filter_level_v; + union { + struct { + uint8_t sharpness_level : 3; + uint8_t mode_ref_delta_enabled : 1; + uint8_t mode_ref_delta_update : 1; + uint8_t reserved : 3; + } bits; + uint8_t value; + } loop_filter_info_fields; + int8_t ref_deltas[8]; + int8_t mode_deltas[2]; + uint8_t base_qindex; + int8_t y_dc_delta_q; + int8_t u_dc_delta_q; + int8_t u_ac_delta_q; + int8_t v_dc_delta_q; + int8_t v_ac_delta_q; + union { + struct { + uint16_t using_qmatrix : 1; + uint16_t qm_y : 4; + uint16_t qm_u : 4; + uint16_t qm_v : 4; + uint16_t reserved : 3; + } bits; + uint16_t value; + } qmatrix_fields; + union { + struct { + uint32_t delta_q_present_flag : 1; + uint32_t log2_delta_q_res : 2; + uint32_t delta_lf_present_flag : 1; + uint32_t log2_delta_lf_res : 2; + uint32_t delta_lf_multi : 1; + uint32_t tx_mode : 2; + uint32_t reference_select : 1; + uint32_t reduced_tx_set_used : 1; + uint32_t skip_mode_present : 1; + uint32_t reserved : 20; + } bits; + uint32_t value; + } mode_control_fields; + uint8_t cdef_damping_minus_3; + uint8_t cdef_bits; + uint8_t cdef_y_strengths[8]; + uint8_t cdef_uv_strengths[8]; + union { + struct { + uint16_t yframe_restoration_type : 2; + uint16_t cbframe_restoration_type : 2; + uint16_t crframe_restoration_type : 2; + uint16_t lr_unit_shift : 2; + uint16_t lr_uv_shift : 1; + uint16_t reserved : 7; + } bits; + uint16_t value; + } loop_restoration_fields; + RocdecAv1WarpedMotionParams wm[7]; + uint32_t reserved[8]; +} RocdecAv1PicParams; + +/*********************************************************************************/ +// AV1 slice/tile parameters (VA-API compatible) +/*********************************************************************************/ +typedef struct _RocdecAv1SliceParams { + uint32_t slice_data_size; + uint32_t slice_data_offset; + uint32_t slice_data_flag; + uint16_t tile_row; + uint16_t tile_column; + uint16_t tg_start; + uint16_t tg_end; + uint8_t anchor_frame_idx; + uint16_t tile_idx_in_tile_list; + uint32_t reserved[4]; +} RocdecAv1SliceParams; + +/*********************************************************************************/ +// RocdecPicParams - Picture parameters for decoding +// Used in rocDecDecodeFrame API +/*********************************************************************************/ +typedef struct _RocdecPicParams { + int pic_width; /**< IN: Coded frame width */ + int pic_height; /**< IN: Coded frame height */ + int curr_pic_idx; /**< IN: Output index of the current picture */ + int field_pic_flag; /**< IN: 0=frame picture, 1=field picture */ + int bottom_field_flag; /**< IN: 0=top field, 1=bottom field */ + int second_field; /**< IN: Second field of a complementary field pair */ + // Bitstream data + uint32_t bitstream_data_len; /**< IN: Number of bytes in bitstream data buffer */ + const uint8_t* bitstream_data; /**< IN: Ptr to bitstream data for this picture */ + uint32_t num_slices; /**< IN: Number of slices in this picture */ + + int ref_pic_flag; /**< IN: This picture is a reference picture */ + int intra_pic_flag; /**< IN: This picture is entirely intra coded */ + uint32_t reserved[30]; /**< Reserved for future use */ + + // Codec-specific data + union { + RocdecMpeg2PicParams mpeg2; /**< Also used for MPEG-1 */ + RocdecAvcPicParams avc; + RocdecHevcPicParams hevc; + RocdecVc1PicParams vc1; + RocdecJPEGPicParams jpeg; + RocdecVp9PicParams vp9; + RocdecAv1PicParams av1; + uint32_t codec_reserved[256]; + } pic_params; + + // Variable size array - one slice param struct per slice + union { + RocdecAvcSliceParams* avc; + RocdecHevcSliceParams* hevc; + RocdecVp9SliceParams* vp9; + RocdecAv1SliceParams* av1; + } slice_params; + + union { + RocdecAvcIQMatrix avc; + RocdecHevcIQMatrix hevc; + } iq_matrix; +} RocdecPicParams; + +/*********************************************************************************/ +// RocdecProcParams - Picture parameters for postprocessing +// Used in rocDecGetVideoFrame API +/*********************************************************************************/ +typedef struct _RocdecProcParams { + int progressive_frame; /**< IN: Input is progressive */ + int top_field_first; /**< IN: Input frame is top field first */ + uint32_t reserved_flags[2]; /**< Reserved for future use (set to zero) */ + + // Fields below are used for raw YUV input + uint64_t raw_input_dptr; /**< IN: Input HIP device ptr for raw YUV */ + uint32_t raw_input_pitch; /**< IN: pitch in bytes of raw YUV input */ + uint32_t raw_input_format; /**< IN: Input YUV format (rocDecVideoCodec_enum) */ + uint64_t raw_output_dptr; /**< IN: Output HIP device mem ptr for raw YUV */ + uint32_t raw_output_pitch; /**< IN: pitch in bytes of raw YUV output */ + uint32_t raw_output_format; /**< IN: Output YUV format (rocDecVideoCodec_enum) */ + uint32_t reserved[16]; /**< Reserved for future use (set to zero) */ +} RocdecProcParams; + +// ---- API functions ---- + +extern rocDecStatus ROCDECAPI +rocDecCreateDecoder(rocDecDecoderHandle* decoder_handle, + RocDecoderCreateInfo* decoder_create_info); + +extern rocDecStatus ROCDECAPI +rocDecDestroyDecoder(rocDecDecoderHandle decoder_handle); + +extern rocDecStatus ROCDECAPI +rocDecGetDecoderCaps(RocdecDecodeCaps* decode_caps); + +extern rocDecStatus ROCDECAPI +rocDecDecodeFrame(rocDecDecoderHandle decoder_handle, + RocdecPicParams* pic_params); + +extern rocDecStatus ROCDECAPI +rocDecGetDecodeStatus(rocDecDecoderHandle decoder_handle, int pic_idx, + RocdecDecodeStatus* decode_status); + +extern rocDecStatus ROCDECAPI +rocDecReconfigureDecoder(rocDecDecoderHandle decoder_handle, + RocdecReconfigureDecoderInfo* reconfig_params); + +extern rocDecStatus ROCDECAPI +rocDecGetVideoFrame(rocDecDecoderHandle decoder_handle, int pic_idx, + void* dev_mem_ptr[3], uint32_t* horizontal_pitch, + RocdecProcParams* vid_postproc_params); + +extern const char* ROCDECAPI rocDecGetErrorName(rocDecStatus rocdec_status); + +#if defined(__cplusplus) +} +#endif diff --git a/src/torchcodec/_core/rocdecode_include/rocparser.h b/src/torchcodec/_core/rocdecode_include/rocparser.h new file mode 100644 index 000000000..495475b7f --- /dev/null +++ b/src/torchcodec/_core/rocdecode_include/rocparser.h @@ -0,0 +1,144 @@ +// Vendored rocParser API header for TorchCodec ROCm backend. +// Minimal subset of the rocDecode parser API. +// +// Original source: https://github.com/ROCm/rocm-systems/projects/rocdecode +// Copyright (c) 2023 - 2026 Advanced Micro Devices, Inc. All rights reserved. +// MIT License + +#pragma once + +#include "rocdecode.h" + +#if defined(__cplusplus) +extern "C" { +#endif + +// ---- Parser handle ---- +typedef void* RocdecVideoParser; +typedef uint64_t RocdecTimeStamp; + +// ---- Video format (from parser sequence callback) ---- +typedef struct { + rocDecVideoCodec codec; + struct { + uint32_t numerator; + uint32_t denominator; + } frame_rate; + uint8_t progressive_sequence; + uint8_t bit_depth_luma_minus8; + uint8_t bit_depth_chroma_minus8; + uint8_t min_num_decode_surfaces; + uint32_t coded_width; + uint32_t coded_height; + struct { + int left; + int top; + int right; + int bottom; + } display_area; + rocDecVideoChromaFormat chroma_format; + uint32_t bitrate; + struct { + int x; + int y; + } display_aspect_ratio; + struct { + uint8_t video_format : 3; + uint8_t video_full_range_flag : 1; + uint8_t reserved_zero_bits : 4; + uint8_t color_primaries; + uint8_t transfer_characteristics; + uint8_t matrix_coefficients; + } video_signal_description; + uint32_t seqhdr_data_length; + uint32_t reconfig_options; +} RocdecVideoFormat; + +// ---- Extended video format ---- +typedef struct { + RocdecVideoFormat format; + uint32_t max_width; + uint32_t max_height; + uint8_t raw_seqhdr_data[1024]; +} RocdecVideoFormatEx; + +// ---- Packet flags ---- +typedef enum { + ROCDEC_PKT_ENDOFSTREAM = 0x01, + ROCDEC_PKT_TIMESTAMP = 0x02, + ROCDEC_PKT_DISCONTINUITY = 0x04, + ROCDEC_PKT_ENDOFPICTURE = 0x08, + ROCDEC_PKT_NOTIFY_EOS = 0x10, +} RocdecVideoPacketFlags; + +// ---- Source data packet ---- +typedef struct _RocdecSourceDataPacket { + uint32_t flags; + uint32_t payload_size; + const uint8_t* payload; + RocdecTimeStamp pts; +} RocdecSourceDataPacket; + +// ---- Display info (from parser display callback) ---- +typedef struct _RocdecParserDispInfo { + int picture_index; + int progressive_frame; + int top_field_first; + int repeat_first_field; + RocdecTimeStamp pts; +} RocdecParserDispInfo; + +// ---- SEI message structures ---- +typedef struct _RocdecSeiMessage { + uint8_t sei_message_type; + uint8_t reserved[3]; + uint32_t sei_message_size; +} RocdecSeiMessage; + +typedef struct _RocdecSeiMessageInfo { + void* sei_data; + RocdecSeiMessage* sei_message; + uint32_t sei_message_count; + uint32_t picIdx; +} RocdecSeiMessageInfo; + +// ---- Parser callback function pointer types ---- +typedef int(ROCDECAPI* PFNVIDSEQUENCECALLBACK)(void*, RocdecVideoFormat*); +typedef int(ROCDECAPI* PFNVIDDECODECALLBACK)(void*, RocdecPicParams*); +typedef int(ROCDECAPI* PFNVIDDISPLAYCALLBACK)(void*, RocdecParserDispInfo*); +typedef int(ROCDECAPI* PFNVIDSEIMSGCALLBACK)(void*, RocdecSeiMessageInfo*); + +// ---- Parser params ---- +typedef struct _RocdecParserParams { + rocDecVideoCodec codec_type; + uint32_t max_num_decode_surfaces; + uint32_t clock_rate; + uint32_t error_threshold; + uint32_t max_display_delay; + uint32_t annex_b : 1; + uint32_t reserved : 31; + uint32_t reserved_1[4]; + void* user_data; + PFNVIDSEQUENCECALLBACK pfn_sequence_callback; + PFNVIDDECODECALLBACK pfn_decode_picture; + PFNVIDDISPLAYCALLBACK pfn_display_picture; + PFNVIDSEIMSGCALLBACK pfn_get_sei_msg; + void* reserved_2[5]; + RocdecVideoFormatEx* ext_video_info; +} RocdecParserParams; + +// ---- Parser API functions ---- +extern rocDecStatus ROCDECAPI +rocDecCreateVideoParser(RocdecVideoParser* parser_handle, + RocdecParserParams* params); + +extern rocDecStatus ROCDECAPI +rocDecParseVideoData(RocdecVideoParser parser_handle, + RocdecSourceDataPacket* packet); + +extern rocDecStatus ROCDECAPI +rocDecDestroyVideoParser(RocdecVideoParser parser_handle); + +#if defined(__cplusplus) +} +#endif