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

Signed-off-by: Dmitry Rogozhkin <[email protected]>

Author: dvrogozh

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", "")
         python_version = sys.version_info
         cmake_args = [
             f"-DCMAKE_INSTALL_PREFIX={self._install_prefix}",
@@ -120,6 +121,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}",
         ]
 
         Path(self.build_temp).mkdir(parents=True, exist_ok=True)

src/torchcodec/decoders/_core/CMakeLists.txt

@@ -9,6 +9,15 @@ find_package(Torch REQUIRED)
 find_package(Python3 ${PYTHON_VERSION} EXACT COMPONENTS Development)
 
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -pedantic -Werror ${TORCH_CXX_FLAGS}")
+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
@@ -61,12 +70,15 @@ function(make_torchcodec_libraries
         AVIOContextHolder.cpp
         FFMPEGCommon.cpp
         VideoDecoder.cpp
+	CPUOnlyDevice.cpp
     )
 
     if(ENABLE_CUDA)
         list(APPEND decoder_sources CudaDevice.cpp)
-    else()
-        list(APPEND decoder_sources CPUOnlyDevice.cpp)
+    endif()
+
+    if(ENABLE_XPU)
+        list(APPEND decoder_sources XpuDevice.cpp)
     endif()
 
     set(decoder_library_dependencies
@@ -81,6 +93,11 @@ function(make_torchcodec_libraries
         )
     endif()
 
+    if(ENABLE_XPU)
+        list(APPEND decoder_library_dependencies
+            PkgConfig::L0 PkgConfig::LIBVA)
+    endif()
+
     make_torchcodec_sublibrary(
         "${decoder_library_name}"
         SHARED

src/torchcodec/decoders/_core/CPUOnlyDevice.cpp

@@ -14,6 +14,7 @@ namespace facebook::torchcodec {
   TORCH_CHECK(false, "Unsupported device: " + device.str());
 }
 
+#ifndef ENABLE_CUDA
 void convertAVFrameToFrameOutputOnCuda(
     const torch::Device& device,
     [[maybe_unused]] const VideoDecoder::VideoStreamOptions& videoStreamOptions,
@@ -40,5 +41,35 @@ std::optional<const AVCodec*> findCudaCodec(
     [[maybe_unused]] const AVCodecID& codecId) {
   throwUnsupportedDeviceError(device);
 }
+#endif // ENABLE_CUDA
+
+#ifndef ENABLE_XPU
+void convertAVFrameToFrameOutputOnXpu(
+    const torch::Device& device,
+    [[maybe_unused]] const VideoDecoder::VideoStreamOptions& videoStreamOptions,
+    [[maybe_unused]] UniqueAVFrame& avFrame,
+    [[maybe_unused]] VideoDecoder::FrameOutput& frameOutput,
+    [[maybe_unused]] std::optional<torch::Tensor> preAllocatedOutputTensor) {
+  throwUnsupportedDeviceError(device);
+}
+
+void initializeContextOnXpu(
+    const torch::Device& device,
+    [[maybe_unused]] AVCodecContext* codecContext) {
+  throwUnsupportedDeviceError(device);
+}
+
+void releaseContextOnXpu(
+    const torch::Device& device,
+    [[maybe_unused]] AVCodecContext* codecContext) {
+  throwUnsupportedDeviceError(device);
+}
+
+std::optional<const AVCodec*> findXpuCodec(
+    const torch::Device& device,
+    [[maybe_unused]] const AVCodecID& codecId) {
+  throwUnsupportedDeviceError(device);
+}
+#endif // ENABLE_XPU
 
 } // namespace facebook::torchcodec

src/torchcodec/decoders/_core/DeviceInterface.h

@@ -29,19 +29,38 @@ void initializeContextOnCuda(
     const torch::Device& device,
     AVCodecContext* codecContext);
 
+void initializeContextOnXpu(
+    const torch::Device& device,
+    AVCodecContext* codecContext);
+
 void convertAVFrameToFrameOutputOnCuda(
     const torch::Device& device,
     const VideoDecoder::VideoStreamOptions& videoStreamOptions,
     UniqueAVFrame& avFrame,
     VideoDecoder::FrameOutput& frameOutput,
     std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
 
+void convertAVFrameToFrameOutputOnXpu(
+    const torch::Device& device,
+    const VideoDecoder::VideoStreamOptions& videoStreamOptions,
+    UniqueAVFrame& avFrame,
+    VideoDecoder::FrameOutput& frameOutput,
+    std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt);
+
 void releaseContextOnCuda(
     const torch::Device& device,
     AVCodecContext* codecContext);
 
+void releaseContextOnXpu(
+    const torch::Device& device,
+    AVCodecContext* codecContext);
+
 std::optional<const AVCodec*> findCudaCodec(
     const torch::Device& device,
     const AVCodecID& codecId);
 
+std::optional<const AVCodec*> findXpuCodec(
+    const torch::Device& device,
+    const AVCodecID& codecId);
+
 } // namespace facebook::torchcodec

src/torchcodec/decoders/_core/VideoDecoder.cpp

@@ -101,6 +101,8 @@ VideoDecoder::~VideoDecoder() {
     if (device.type() == torch::kCPU) {
     } else if (device.type() == torch::kCUDA) {
       releaseContextOnCuda(device, streamInfo.codecContext.get());
+    } else if (device.type() == torch::kXPU) {
+      releaseContextOnXpu(device, streamInfo.codecContext.get());
     } else {
       TORCH_CHECK(false, "Invalid device type: " + device.str());
     }
@@ -429,10 +431,16 @@ void VideoDecoder::addStream(
 
   // TODO_CODE_QUALITY it's pretty meh to have a video-specific logic within
   // addStream() which is supposed to be generic
-  if (mediaType == AVMEDIA_TYPE_VIDEO && device.type() == torch::kCUDA) {
-    avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
-        findCudaCodec(device, streamInfo.stream->codecpar->codec_id)
-            .value_or(avCodec));
+  if (mediaType == AVMEDIA_TYPE_VIDEO) {
+    if (device.type() == torch::kCUDA) {
+      avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
+          findCudaCodec(device, streamInfo.stream->codecpar->codec_id)
+              .value_or(avCodec));
+    } else if (device.type() == torch::kXPU) {
+      avCodec = makeAVCodecOnlyUseForCallingAVFindBestStream(
+          findXpuCodec(device, streamInfo.stream->codecpar->codec_id)
+              .value_or(avCodec));
+    }
   }
 
   AVCodecContext* codecContext = avcodec_alloc_context3(avCodec);
@@ -447,8 +455,12 @@ void VideoDecoder::addStream(
   streamInfo.codecContext->pkt_timebase = streamInfo.stream->time_base;
 
   // TODO_CODE_QUALITY same as above.
-  if (mediaType == AVMEDIA_TYPE_VIDEO && device.type() == torch::kCUDA) {
-    initializeContextOnCuda(device, codecContext);
+  if (mediaType == AVMEDIA_TYPE_VIDEO) {
+    if (device.type() == torch::kCUDA) {
+      initializeContextOnCuda(device, codecContext);
+    } else if (device.type() == torch::kXPU) {
+      initializeContextOnXpu(device, codecContext);
+    }
   }
 
   retVal = avcodec_open2(streamInfo.codecContext.get(), avCodec, nullptr);
@@ -476,7 +488,8 @@ void VideoDecoder::addVideoStream(
     const VideoStreamOptions& videoStreamOptions) {
   TORCH_CHECK(
       videoStreamOptions.device.type() == torch::kCPU ||
-          videoStreamOptions.device.type() == torch::kCUDA,
+          videoStreamOptions.device.type() == torch::kCUDA ||
+          videoStreamOptions.device.type() == torch::kXPU,
       "Invalid device type: " + videoStreamOptions.device.str());
 
   addStream(
@@ -1226,6 +1239,13 @@ VideoDecoder::FrameOutput VideoDecoder::convertAVFrameToFrameOutput(
         avFrame,
         frameOutput,
         preAllocatedOutputTensor);
+  } else if (streamInfo.videoStreamOptions.device.type() == torch::kXPU) {
+    convertAVFrameToFrameOutputOnXpu(
+        streamInfo.videoStreamOptions.device,
+        streamInfo.videoStreamOptions,
+        avFrame,
+        frameOutput,
+        preAllocatedOutputTensor);
   } else {
     TORCH_CHECK(
         false,

src/torchcodec/decoders/_core/VideoDecoderOps.cpp

@@ -210,10 +210,13 @@ void _add_video_stream(
     } else if (device.value().rfind("cuda", 0) == 0) { // starts with "cuda"
       std::string deviceStr(device.value());
       videoStreamOptions.device = torch::Device(deviceStr);
+    } else if (device.value().rfind("xpu", 0) == 0) { // starts with "xpu"
+      std::string deviceStr(device.value());
+      videoStreamOptions.device = torch::Device(deviceStr);
     } else {
       throw std::runtime_error(
           "Invalid device=" + std::string(device.value()) +
-          ". device must be either cpu or cuda.");
+          ". device must be either cpu, cuda or xpu.");
     }
   }