V210Compute.hpp

#pragma once
#include <Gfx/Graph/Utils.hpp>
#include <Gfx/Graph/encoders/ColorSpaceOut.hpp>
#include <Gfx/Graph/encoders/ComputeEncoder.hpp>
#include <Gfx/Graph/encoders/GPUVideoEncoder.hpp>
 
#include <private/qrhi_p.h>
 
namespace score::gfx
{
 
struct V210ComputeEncoder final : ComputeEncoder
{
  // %1 = colorMatrixOut(...) shader defining convert_from_rgb(vec3)
  static constexpr const char* compute_shader = R"_(#version 450
    layout(local_size_x = 16, local_size_y = 1, local_size_z = 1) in;
 
    layout(binding = 0) uniform sampler2D src_tex;
    layout(std430, binding = 1) writeonly buffer V210Buf {
      uint v210[];
    };
    layout(std140, binding = 2) uniform Params {
      ivec2 src_size;
      ivec2 _pad0;
      uint  line_stride_words; // bytes-per-row / 4
      uint  groups_per_row;    // == src_width / 6
      uint  _pad1[2];
    };
 
    vec2 flip_y(vec2 tc) {
    // Only OpenGL. The rest of the engine puts its geometry through
    // renderer.clipSpaceCorrMatrix, which negates Y on Vulkan; this pass draws
    // a hardcoded triangle in raw NDC and does not, so the correction it needs
    // is not the same one. Flipping on Vulkan as well handed libav, GStreamer,
    // NDI and every other consumer an upside-down picture.
    #if defined(QSHADER_SPIRV) || defined(QSHADER_MSL) || defined(QSHADER_HLSL)
      return tc;
    #else
      return vec2(tc.x, 1.0 - tc.y);
    #endif
    }
    )_" "%1" R"_(
 
    uvec3 to_yuv10(vec3 rgb) {
      vec3 yuv = clamp(convert_from_rgb(rgb), 0.0, 1.0);
      return uvec3(yuv * 1023.0 + 0.5);
    }
 
    void main() {
      uint group_x = gl_GlobalInvocationID.x;
      uint y       = gl_GlobalInvocationID.y;
      if (group_x >= groups_per_row || int(y) >= src_size.y)
        return;
 
      // Y-flip on backends that need it (matches the fragment encoders):
      // only OpenGL. The rest of the engine negates Y through
      // renderer.clipSpaceCorrMatrix on Vulkan; this pass indexes texels
      // directly and does not, so it must not flip there.
      ivec2 srcSize = src_size;
    #if defined(QSHADER_SPIRV) || defined(QSHADER_MSL) || defined(QSHADER_HLSL)
      int src_y = int(y);
    #else
      int src_y = srcSize.y - 1 - int(y);
    #endif
 
      // Edge-clamp: the tail group at widths not divisible by 6 has fewer
      // than 6 source pixels (texelFetch does NOT clamp on its own).
      int xmax = src_size.x - 1;
      int x0i = int(group_x * 6u);
      uvec3 a = to_yuv10(texelFetch(src_tex, ivec2(min(x0i    , xmax), src_y), 0).rgb);
      uvec3 b = to_yuv10(texelFetch(src_tex, ivec2(min(x0i + 1, xmax), src_y), 0).rgb);
      uvec3 c = to_yuv10(texelFetch(src_tex, ivec2(min(x0i + 2, xmax), src_y), 0).rgb);
      uvec3 d = to_yuv10(texelFetch(src_tex, ivec2(min(x0i + 3, xmax), src_y), 0).rgb);
      uvec3 e = to_yuv10(texelFetch(src_tex, ivec2(min(x0i + 4, xmax), src_y), 0).rgb);
      uvec3 f = to_yuv10(texelFetch(src_tex, ivec2(min(x0i + 5, xmax), src_y), 0).rgb);
 
      uint cb01 = (a.y + b.y) >> 1;
      uint cr01 = (a.z + b.z) >> 1;
      uint cb23 = (c.y + d.y) >> 1;
      uint cr23 = (c.z + d.z) >> 1;
      uint cb45 = (e.y + f.y) >> 1;
      uint cr45 = (e.z + f.z) >> 1;
 
      // v210 packing: 6 pixels = 4 little-endian 32-bit words.
      uint w0 = cb01 | (a.x << 10) | (cr01 << 20);
      uint w1 = b.x  | (cb23 << 10) | (c.x << 20);
      uint w2 = cr23 | (d.x << 10) | (cb45 << 20);
      uint w3 = e.x  | (cr45 << 10) | (f.x << 20);
 
      uint base = y * line_stride_words + group_x * 4u;
      v210[base + 0u] = w0;
      v210[base + 1u] = w1;
      v210[base + 2u] = w2;
      v210[base + 3u] = w3;
    }
  )_";
 
  QRhiBuffer* m_paramsUBO{};
  QRhiSampler* m_sampler{};
  QRhiShaderResourceBindings* m_srb{};
  QRhiComputePipeline* m_pipeline{};
  int m_width{};
  int m_height{};
  int m_groupsPerRow{};
  uint32_t m_lineStrideBytes{};
 
  // Output buffer is owned by the caller. Pass it in init().
  bool init(
      QRhi& rhi, const RenderState& state, QRhiTexture* inputRGBA, int width,
      int height, QRhiBuffer* outputBuffer,
      const QString& colorConversion) override
  {
    if(!outputBuffer || width % 2 != 0)
      return false;
    if(!rhi.isFeatureSupported(QRhi::Compute))
      return false;
 
    m_width = width;
    m_height = height;
    // Ceil: the tail group of a width not divisible by 6 (1280, 2048-DCI)
    // still occupies a full 4-word group on the padded wire row; the shader
    // edge-clamps its source fetches.
    m_groupsPerRow = (width + 5) / 6;
    m_lineStrideBytes = ((width + 47) / 48) * 128;
 
    // Params UBO (std140, 32 bytes).
    m_paramsUBO = rhi.newBuffer(QRhiBuffer::Dynamic, QRhiBuffer::UniformBuffer, 32);
    m_paramsUBO->setName("V210ComputeEncoder::params");
    if(!m_paramsUBO->create())
      return false;
 
    m_sampler = rhi.newSampler(
        QRhiSampler::Nearest, QRhiSampler::Nearest, QRhiSampler::None,
        QRhiSampler::ClampToEdge, QRhiSampler::ClampToEdge);
    if(!m_sampler->create())
      return false;
 
    m_srb = rhi.newShaderResourceBindings();
    m_srb->setBindings({
        QRhiShaderResourceBinding::sampledTexture(
            0, QRhiShaderResourceBinding::ComputeStage, inputRGBA, m_sampler),
        QRhiShaderResourceBinding::bufferStore(
            1, QRhiShaderResourceBinding::ComputeStage, outputBuffer),
        QRhiShaderResourceBinding::uniformBuffer(
            2, QRhiShaderResourceBinding::ComputeStage, m_paramsUBO),
    });
    if(!m_srb->create())
      return false;
 
    QShader cs = makeCompute(
        state, QString::fromLatin1(compute_shader).arg(colorConversion));
    m_pipeline = rhi.newComputePipeline();
    m_pipeline->setShaderStage({QRhiShaderStage::Compute, cs});
    m_pipeline->setShaderResourceBindings(m_srb);
    if(!m_pipeline->create())
      return false;
 
    return true;
  }
 
  // Update params UBO + dispatch the compute pass. Caller has wrapped this
  // inside an offscreen frame; the resource update batch must already be
  // active (typical pattern: pass the batch in, encoder writes the UBO).
  void exec(
      QRhi& rhi, QRhiCommandBuffer& cb, QRhiResourceUpdateBatch* res) override
  {
    struct alignas(16) ParamsData
    {
      int32_t srcW, srcH;
      int32_t pad0[2];
      uint32_t lineStrideWords;
      uint32_t groupsPerRow;
      uint32_t pad1[2];
    } p{
        m_width, m_height,
        {0, 0},
        m_lineStrideBytes / 4,
        static_cast<uint32_t>(m_groupsPerRow),
        {0, 0}};
    res->updateDynamicBuffer(m_paramsUBO, 0, sizeof(p), &p);
 
    cb.beginComputePass(res);
    cb.setComputePipeline(m_pipeline);
    cb.setShaderResources(m_srb);
    cb.dispatch(
        (m_groupsPerRow + 15) / 16, // local_size_x = 16
        m_height,
        1);
    cb.endComputePass();
  }
 
  void release() override
  {
    delete m_pipeline;   m_pipeline = nullptr;
    delete m_srb;        m_srb = nullptr;
    delete m_sampler;    m_sampler = nullptr;
    delete m_paramsUBO;  m_paramsUBO = nullptr;
  }
};
 
} // namespace score::gfx