Enable Intel GPU support in torchcodec on Linux (xpu device)

This commit enables support for Intel GPUs in torchcodec. It adds:
* ffmpeg-vaapi for decoding
* VAAPI based color space conversion (decoding output to RGBA)
* RGBA surface import as torch tensor (on torch xpu device)
* RGBA to RGB24 tensor slicing

To build torchcodec with Intel GPU support:
* Install pytorch with XPU backend support. For example, with:
```
pip3 install torch --index-url https://download.pytorch.org/whl/xpu
```
* Install oneAPI development environment following
  https://github.com/pytorch/pytorch?tab=readme-ov-file#intel-gpu-support
* Build and install FFmpeg with `--enable-vaapi`
* Install torcheval (for tests): `pip3 install torcheval`
* Build torchcodec with: `ENABLE_XPU=1 python3 setup.py devel`

Notes:
* RGB24 is not supported color format on current Intel GPUs (as it
  is considered to be suboptimal due to odd alignments)
* Intel media and compute APIs can't seamlessly work with the
  memory from each other. For example, Intel computes's Unified
  Shared Memory pointers are not recognized by media APIs. Thus,
  lower level sharing via dma fds is needed. This alos makes this
  part of the solution OS dependent.
* Color space conversion algoriths might be quite different as it
  happens for Intel. This requires to check PSNR values instead of
  per-pixel atol/rtol differences.
* Installing oneAPI environment is neded due to
  pytorch/pytorch#149075

This commit was primary verfied on Intel Battlemage G21 (0xe20b) and
Intel Data Center GPU Flex (0x56c0).

Co-authored-by: Edgar Romo Montiel <[email protected]>
Signed-off-by: Edgar Romo Montiel <[email protected]>
Signed-off-by: Dmitry Rogozhkin <[email protected]>

Author: 2 people

setup.py

@@ -112,6 +112,7 @@ 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_xpu = os.environ.get("ENABLE_XPU", "")
         torchcodec_disable_compile_warning_as_error = os.environ.get(
             "TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR", "OFF"
         )
@@ -123,6 +124,7 @@ 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_XPU={enable_xpu}",
             f"-DTORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR={torchcodec_disable_compile_warning_as_error}",
         ]
 

src/torchcodec/_core/CMakeLists.txt

@@ -34,6 +34,15 @@ else()
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -pedantic ${TORCHCODEC_WERROR_OPTION} ${TORCH_CXX_FLAGS}")
 endif()
 
+if(ENABLE_CUDA)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_CUDA")
+endif()
+if(ENABLE_XPU)
+    find_package(PkgConfig REQUIRED)
+    pkg_check_modules(L0 REQUIRED IMPORTED_TARGET level-zero)
+    pkg_check_modules(LIBVA REQUIRED IMPORTED_TARGET libva)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DENABLE_XPU")
+endif()
 
 function(make_torchcodec_sublibrary
     library_name
@@ -109,7 +118,11 @@ function(make_torchcodec_libraries
     )
 
     if(ENABLE_CUDA)
-	    list(APPEND core_sources CudaDeviceInterface.cpp)
+        list(APPEND core_sources CudaDeviceInterface.cpp)
+    endif()
+
+    if(ENABLE_XPU)
+        list(APPEND core_sources XpuDeviceInterface.cpp)
     endif()
 
     set(core_library_dependencies
@@ -124,6 +137,11 @@ function(make_torchcodec_libraries
         )
     endif()
 
+    if(ENABLE_XPU)
+        list(APPEND core_library_dependencies
+            PkgConfig::L0 PkgConfig::LIBVA)
+    endif()
+
     make_torchcodec_sublibrary(
         "${core_library_name}"
         SHARED