Support native acceleration with Vulkan for h264 videos

Author: khamilowicz

examples/vp8_to_h264.exs

@@ -1,23 +1,40 @@
-Mix.install([
-  :membrane_file_plugin,
-  :membrane_ivf_plugin,
-  {:membrane_transcoder_plugin, path: __DIR__ |> Path.join("..") |> Path.expand()}
-])
+use_native = System.argv() |> Enum.member?("--native")
+
+vk_dep =
+  if use_native &&
+       match?({_, 0}, System.cmd("pkg-config", ["--exists", "vulkan"], stderr_to_stdout: true)) do
+    [{:membrane_vk_video_plugin, "~> 0.2.0"}]
+  else
+    []
+  end
+
+Mix.install(
+  vk_dep ++
+    [
+      :membrane_file_plugin,
+      :membrane_ivf_plugin,
+      {:membrane_transcoder_plugin, path: __DIR__ |> Path.join("..") |> Path.expand()}
+    ]
+)
 
 defmodule Example do
-  alias Membrane.RCPipeline
-  require Membrane.RCPipeline, as: RCPipeline
+  alias Membrane.{H264, RCPipeline}
+  require RCPipeline
+
   import Membrane.ChildrenSpec
 
-  def convert(input_file, output_file) do
+  def convert(input_file, output_file, native_acceleration) do
     pipeline = RCPipeline.start_link!()
 
     spec =
       child(%Membrane.File.Source{
         location: input_file
       })
       |> child(:deserializer, Membrane.IVF.Deserializer)
-      |> child(:transcoder, %Membrane.Transcoder{output_stream_format: Membrane.H264})
+      |> child(:transcoder, %Membrane.Transcoder{
+        output_stream_format: H264,
+        native_acceleration: native_acceleration
+      })
       |> child(:sink, %Membrane.File.Sink{location: output_file})
 
     RCPipeline.subscribe(pipeline, _any)
@@ -27,5 +44,25 @@ defmodule Example do
   end
 end
 
-File.mkdir("tmp")
-Example.convert(Path.join("./test/fixtures", "video_vp8.ivf"), Path.join("./tmp", "video.h264"))
+File.mkdir_p!("tmp")
+
+if use_native && !Membrane.Transcoder.vulkan_available?() do
+  raise "Vulkan is not available. Cannot run with --native flag."
+end
+
+acceleration = if use_native, do: :if_available, else: :never
+
+output_file = Path.join(__DIR__, "tmp/video.h264")
+
+if use_native do
+  IO.puts("Vulkan available: true")
+  IO.puts("Using native_acceleration: :if_available")
+end
+
+Example.convert(
+  Path.join(__DIR__, "../test/fixtures/video_vp8.ivf"),
+  output_file,
+  acceleration
+)
+
+IO.puts("Done. Output written to #{output_file}")

lib/transcoder.ex

@@ -26,6 +26,9 @@ defmodule Membrane.Transcoder do
   Now, the only supported stream parameters are:
   * `:pixel_format` in `Membrane.RawVideo`
   * `:alignment` and `:stream_structure` in `Membrane.H264` and `Membrane.H265`
+
+  When the `membrane_vk_video_plugin` dependency is present and Vulkan hardware is available,
+  H.264 encode/decode can be offloaded to the GPU by setting `native_acceleration: :if_available`.
   """
   use Membrane.Bin
 
@@ -139,6 +142,17 @@ defmodule Membrane.Transcoder do
 
                 If nil or not set, the input stream format won't be overriden.
                 """
+              ],
+              native_acceleration: [
+                spec: :never | :if_available,
+                default: :never,
+                description: """
+                Specifies whether to use Vulkan hardware acceleration for video transcoding.
+
+                Can be:
+                * `:never` - Always use software-based transcoding (default)
+                * `:if_available` - Use Vulkan acceleration when available on the system
+                """
               ]
 
   @impl true
@@ -155,12 +169,29 @@ defmodule Membrane.Transcoder do
       opts
       |> Map.from_struct()
       |> Map.merge(%{
-        input_stream_format: nil
+        input_stream_format: nil,
+        use_vulkan?: should_use_vulkan?(opts.native_acceleration)
       })
 
     {[spec: spec], state}
   end
 
+  defp should_use_vulkan?(:if_available), do: vulkan_available?()
+  defp should_use_vulkan?(_native_acceleration), do: false
+
+  @doc """
+  Returns `true` if Vulkan hardware acceleration is available on this machine.
+
+  Use this to guard code that sets `native_acceleration: :if_available`, for example
+  to print a diagnostic or pick a different pipeline branch.
+  """
+  @spec vulkan_available?() :: boolean()
+  def vulkan_available?() do
+    Code.ensure_loaded?(Membrane.VKVideo.Decoder) and
+      Code.ensure_loaded?(Membrane.VKVideo.Encoder) and
+      Code.ensure_loaded?(Membrane.VKVideo.Native)
+  end
+
   defp maybe_plug_stream_format_changer(builder, nil), do: builder
 
   defp maybe_plug_stream_format_changer(builder, enforced_stream_format) do
@@ -187,7 +218,8 @@ defmodule Membrane.Transcoder do
       |> plug_transcoding(
         format,
         state.output_stream_format,
-        state.transcoding_policy
+        state.transcoding_policy,
+        state.use_vulkan?
       )
       |> get_child(:output_funnel)
 
@@ -232,15 +264,15 @@ defmodule Membrane.Transcoder do
     end
   end
 
-  defp plug_transcoding(builder, input_format, output_format, transcoding_policy)
+  defp plug_transcoding(builder, input_format, output_format, transcoding_policy, _use_vulkan?)
        when Audio.is_audio_format(input_format) do
     builder
     |> Audio.plug_audio_transcoding(input_format, output_format, transcoding_policy)
   end
 
-  defp plug_transcoding(builder, input_format, output_format, transcoding_policy)
+  defp plug_transcoding(builder, input_format, output_format, transcoding_policy, use_vulkan?)
        when Video.is_video_format(input_format) do
     builder
-    |> Video.plug_video_transcoding(input_format, output_format, transcoding_policy)
+    |> Video.plug_video_transcoding(input_format, output_format, transcoding_policy, use_vulkan?)
   end
 end

lib/transcoder/video.ex

@@ -33,29 +33,50 @@ defmodule Membrane.Transcoder.Video do
           ChildrenSpec.builder(),
           video_stream_format(),
           video_stream_format(),
+          :always | :if_needed | :never,
           boolean()
         ) :: ChildrenSpec.builder()
-  def plug_video_transcoding(builder, input_format, output_format, transcoding_policy)
+  def plug_video_transcoding(
+        builder,
+        input_format,
+        output_format,
+        transcoding_policy,
+        use_vulkan?
+      )
       when is_video_format(input_format) and is_video_format(output_format) do
-    do_plug_video_transcoding(builder, input_format, output_format, transcoding_policy)
+    do_plug_video_transcoding(
+      builder,
+      input_format,
+      output_format,
+      transcoding_policy,
+      use_vulkan?
+    )
   end
 
   defp do_plug_video_transcoding(
          builder,
          %RemoteStream{content_format: h26x},
          output_format,
-         transcoding_policy
+         transcoding_policy,
+         use_vulkan?
        )
        when h26x in [H264, H265] do
     do_plug_video_transcoding(
       builder,
       struct!(h26x),
       output_format,
-      transcoding_policy
+      transcoding_policy,
+      use_vulkan?
     )
   end
 
-  defp do_plug_video_transcoding(builder, %H264{}, %H264{} = output_format, transcoding_policy)
+  defp do_plug_video_transcoding(
+         builder,
+         %H264{},
+         %H264{} = output_format,
+         transcoding_policy,
+         _use_vulkan?
+       )
        when transcoding_policy in [:if_needed, :never] do
     builder
     |> child(:h264_parser, %H264.Parser{
@@ -64,7 +85,13 @@ defmodule Membrane.Transcoder.Video do
     })
   end
 
-  defp do_plug_video_transcoding(builder, %H265{}, %H265{} = output_format, transcoding_policy)
+  defp do_plug_video_transcoding(
+         builder,
+         %H265{},
+         %H265{} = output_format,
+         transcoding_policy,
+         _use_vulkan?
+       )
        when transcoding_policy in [:if_needed, :never] do
     builder
     |> child(:h265_parser, %H265.Parser{
@@ -77,7 +104,19 @@ defmodule Membrane.Transcoder.Video do
          builder,
          %RawVideo{} = input_format,
          %RawVideo{} = output_format,
-         _transcoding_policy
+         _transcoding_policy,
+         true
+       ) do
+    builder
+    |> maybe_plug_swscale_converter_vulkan(input_format, output_format)
+  end
+
+  defp do_plug_video_transcoding(
+         builder,
+         %RawVideo{} = input_format,
+         %RawVideo{} = output_format,
+         _transcoding_policy,
+         false
        ) do
     builder
     |> maybe_plug_swscale_converter(input_format, output_format)
@@ -87,7 +126,8 @@ defmodule Membrane.Transcoder.Video do
          builder,
          %format_module{},
          %format_module{},
-         transcoding_policy
+         transcoding_policy,
+         _use_vulkan?
        )
        when transcoding_policy in [:if_needed, :never] do
     Membrane.Logger.debug("""
@@ -97,20 +137,40 @@ defmodule Membrane.Transcoder.Video do
     builder
   end
 
-  defp do_plug_video_transcoding(_builder, input_format, output_format, :never) do
-    raise """
-    Cannot convert input format #{inspect(input_format)} to output format #{inspect(output_format)} \
-    with :transcoding_policy option set to :never.
-    """
+  defp do_plug_video_transcoding(_builder, input_format, output_format, :never, _use_vulkan?),
+    do:
+      raise("""
+      Cannot convert input format #{inspect(input_format)} to output format #{inspect(output_format)} \
+      with :transcoding_policy option set to :never.
+      """)
+
+  defp do_plug_video_transcoding(builder, input_format, output_format, _transcoding_policy, true) do
+    builder
+    |> maybe_plug_parser_and_decoder_vulkan(input_format)
+    |> maybe_plug_swscale_converter_vulkan(input_format, output_format)
+    |> maybe_plug_encoder_and_parser_vulkan(output_format)
   end
 
-  defp do_plug_video_transcoding(builder, input_format, output_format, _transcoding_policy) do
+  defp do_plug_video_transcoding(builder, input_format, output_format, _transcoding_policy, false) do
     builder
     |> maybe_plug_parser_and_decoder(input_format)
     |> maybe_plug_swscale_converter(input_format, output_format)
     |> maybe_plug_encoder_and_parser(output_format)
   end
 
+  # VK-specific decoder: child name :vk_h264_decoder distinguishes it from FFmpeg's :h264_decoder
+  defp maybe_plug_parser_and_decoder_vulkan(builder, %H264{}) do
+    builder
+    |> child(:h264_input_parser, %H264.Parser{
+      output_stream_structure: :annexb,
+      output_alignment: :au
+    })
+    |> child(:vk_h264_decoder, Membrane.VKVideo.Decoder)
+  end
+
+  defp maybe_plug_parser_and_decoder_vulkan(builder, format),
+    do: maybe_plug_parser_and_decoder(builder, format)
+
   defp maybe_plug_parser_and_decoder(builder, %H264{}) do
     builder
     |> child(:h264_input_parser, %H264.Parser{
@@ -143,10 +203,41 @@ defmodule Membrane.Transcoder.Video do
     builder |> child(:vp8_decoder, decoder_module)
   end
 
-  defp maybe_plug_parser_and_decoder(builder, %RawVideo{}) do
-    builder
+  defp maybe_plug_parser_and_decoder(builder, %RawVideo{}), do: builder
+
+  # VK decoder outputs NV12; these clauses handle the NV12 <-> other-format conversion
+  # when mixing VK decode/encode with different pixel formats.
+
+  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %RawVideo{pixel_format: nil}),
+    do: builder
+
+  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %RawVideo{pixel_format: :NV12}),
+    do: builder
+
+  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %RawVideo{} = output_format) do
+    builder |> child(:raw_video_converter, %SWScale.Converter{format: output_format.pixel_format})
+  end
+
+  defp maybe_plug_swscale_converter_vulkan(builder, %RawVideo{pixel_format: :NV12}, %H264{}),
+    do: builder
+
+  defp maybe_plug_swscale_converter_vulkan(builder, %RawVideo{}, %H264{}) do
+    builder |> child(:raw_video_converter, %SWScale.Converter{format: :NV12})
   end
 
+  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, %H264{}), do: builder
+
+  defp maybe_plug_swscale_converter_vulkan(builder, %H264{}, _output_format) do
+    builder |> child(:raw_video_converter, %SWScale.Converter{format: :I420})
+  end
+
+  defp maybe_plug_swscale_converter_vulkan(builder, _input_format, %H264{}) do
+    builder |> child(:raw_video_converter, %SWScale.Converter{format: :NV12})
+  end
+
+  defp maybe_plug_swscale_converter_vulkan(builder, input_format, output_format),
+    do: maybe_plug_swscale_converter(builder, input_format, output_format)
+
   defp maybe_plug_swscale_converter(builder, input_format, %RawVideo{} = output_format) do
     case input_format do
       _any when output_format.pixel_format == nil ->
@@ -175,10 +266,20 @@ defmodule Membrane.Transcoder.Video do
     end
   end
 
-  defp maybe_plug_swscale_converter(builder, _input_format, _output_format) do
+  defp maybe_plug_swscale_converter(builder, _input_format, _output_format), do: builder
+
+  defp maybe_plug_encoder_and_parser_vulkan(builder, %H264{} = h264) do
     builder
+    |> child(:vk_h264_encoder, Membrane.VKVideo.Encoder)
+    |> child(:h264_output_parser, %H264.Parser{
+      output_stream_structure: stream_structure_type(h264),
+      output_alignment: h264.alignment
+    })
   end
 
+  defp maybe_plug_encoder_and_parser_vulkan(builder, format),
+    do: maybe_plug_encoder_and_parser(builder, format)
+
   defp maybe_plug_encoder_and_parser(builder, %H264{} = h264) do
     builder
     |> child(:h264_encoder, %H264.FFmpeg.Encoder{preset: :ultrafast})
@@ -219,9 +320,7 @@ defmodule Membrane.Transcoder.Video do
     builder |> child(:vp8_encoder, %VP9.Encoder{g_threads: number_of_threads, cpu_used: 15})
   end
 
-  defp maybe_plug_encoder_and_parser(builder, %RawVideo{}) do
-    builder
-  end
+  defp maybe_plug_encoder_and_parser(builder, %RawVideo{}), do: builder
 
   defp stream_structure_type(%h26x{stream_structure: stream_structure})
        when h26x in [H264, H265] do