diff --git a/src/gallium/drivers/d3d12/d3d12_video_enc_hevc.cpp b/src/gallium/drivers/d3d12/d3d12_video_enc_hevc.cpp index 5837d463f03..48c19f196d4 100644 --- a/src/gallium/drivers/d3d12/d3d12_video_enc_hevc.cpp +++ b/src/gallium/drivers/d3d12/d3d12_video_enc_hevc.cpp @@ -1071,6 +1071,17 @@ d3d12_video_encoder_build_codec_headers_hevc(struct d3d12_video_encoder *pD3D12E static_cast(pD3D12Enc->m_upBitstreamBuilder.get()); assert(pHEVCBitstreamBuilder); + size_t writtenAUDBytesCount = 0; + bool forceWriteAUD = (pD3D12Enc->m_currentEncodeConfig.m_ConfigDirtyFlags & d3d12_video_encoder_config_dirty_flag_aud_header); + if (forceWriteAUD) + { + pHEVCBitstreamBuilder->write_aud(pD3D12Enc->m_BitstreamHeadersBuffer, + pD3D12Enc->m_BitstreamHeadersBuffer.begin(), + currentPicParams.pHEVCPicData->FrameType, + writtenAUDBytesCount); + pWrittenCodecUnitsSizes.push_back(writtenAUDBytesCount); + } + uint32_t active_seq_parameter_set_id = pHEVCBitstreamBuilder->get_active_sps().sps_seq_parameter_set_id; uint32_t active_video_parameter_set_id = pHEVCBitstreamBuilder->get_active_vps().vps_video_parameter_set_id; @@ -1086,7 +1097,7 @@ d3d12_video_encoder_build_codec_headers_hevc(struct d3d12_video_encoder *pD3D12E gopHasBFrames, active_video_parameter_set_id, pD3D12Enc->m_BitstreamHeadersBuffer, - pD3D12Enc->m_BitstreamHeadersBuffer.begin(), + pD3D12Enc->m_BitstreamHeadersBuffer.begin() + writtenAUDBytesCount, writtenVPSBytesCount); pHEVCBitstreamBuilder->set_active_vps(vps); pWrittenCodecUnitsSizes.push_back(writtenVPSBytesCount); @@ -1110,7 +1121,7 @@ d3d12_video_encoder_build_codec_headers_hevc(struct d3d12_video_encoder *pD3D12E *codecConfigDesc.pHEVCConfig, pD3D12Enc->m_currentEncodeConfig.m_encoderGOPConfigDesc.m_HEVCGroupOfPictures, pD3D12Enc->m_BitstreamHeadersBuffer, - pD3D12Enc->m_BitstreamHeadersBuffer.begin() + writtenVPSBytesCount, + pD3D12Enc->m_BitstreamHeadersBuffer.begin() + writtenAUDBytesCount + writtenVPSBytesCount, writtenSPSBytesCount); pHEVCBitstreamBuilder->set_active_sps(sps); pWrittenCodecUnitsSizes.push_back(writtenSPSBytesCount); @@ -1129,8 +1140,8 @@ d3d12_video_encoder_build_codec_headers_hevc(struct d3d12_video_encoder *pD3D12E bool forceWritePPS = (pD3D12Enc->m_currentEncodeConfig.m_ConfigDirtyFlags & d3d12_video_encoder_config_dirty_flag_picture_header); if (forceWritePPS || d3d12_video_encoder_needs_new_pps_hevc(pD3D12Enc, writeNewSPS, tentative_pps, active_pps)) { pHEVCBitstreamBuilder->set_active_pps(tentative_pps); - pD3D12Enc->m_BitstreamHeadersBuffer.resize(writtenSPSBytesCount + writtenVPSBytesCount + writtenPPSBytesCount); - memcpy(&pD3D12Enc->m_BitstreamHeadersBuffer.data()[(writtenSPSBytesCount + writtenVPSBytesCount)], pD3D12Enc->m_StagingHeadersBuffer.data(), writtenPPSBytesCount); + pD3D12Enc->m_BitstreamHeadersBuffer.resize(writtenAUDBytesCount + writtenVPSBytesCount + writtenSPSBytesCount + writtenPPSBytesCount); + memcpy(&pD3D12Enc->m_BitstreamHeadersBuffer.data()[(writtenAUDBytesCount + writtenVPSBytesCount + writtenSPSBytesCount)], pD3D12Enc->m_StagingHeadersBuffer.data(), writtenPPSBytesCount); pWrittenCodecUnitsSizes.push_back(writtenPPSBytesCount); } else { writtenPPSBytesCount = 0; @@ -1138,8 +1149,8 @@ d3d12_video_encoder_build_codec_headers_hevc(struct d3d12_video_encoder *pD3D12E } // Shrink buffer to fit the headers - if (pD3D12Enc->m_BitstreamHeadersBuffer.size() > (writtenPPSBytesCount + writtenSPSBytesCount + writtenVPSBytesCount)) { - pD3D12Enc->m_BitstreamHeadersBuffer.resize(writtenPPSBytesCount + writtenSPSBytesCount + writtenVPSBytesCount); + if (pD3D12Enc->m_BitstreamHeadersBuffer.size() > (writtenAUDBytesCount + writtenVPSBytesCount + writtenSPSBytesCount + writtenPPSBytesCount)) { + pD3D12Enc->m_BitstreamHeadersBuffer.resize(writtenAUDBytesCount + writtenVPSBytesCount + writtenSPSBytesCount + writtenPPSBytesCount); } assert(std::accumulate(pWrittenCodecUnitsSizes.begin(), pWrittenCodecUnitsSizes.end(), 0u) == diff --git a/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.cpp b/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.cpp index fdaa6ecb641..8713a11dd7c 100644 --- a/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.cpp +++ b/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.cpp @@ -552,6 +552,15 @@ d3d12_video_bitstream_builder_hevc::write_end_of_sequence_nalu(std::vector & headerBitstream, + std::vector::iterator placingPositionStart, + D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC frameType, + size_t & writtenBytes) +{ + m_hevcEncoder.write_aud(headerBitstream, placingPositionStart, frameType, writtenBytes); +} + void d3d12_video_bitstream_builder_hevc::print_vps(const HevcVideoParameterSet& VPS) { diff --git a/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.h b/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.h index 601dca31b08..31b1118db43 100644 --- a/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.h +++ b/src/gallium/drivers/d3d12/d3d12_video_encoder_bitstream_builder_hevc.h @@ -72,6 +72,11 @@ class d3d12_video_bitstream_builder_hevc : public d3d12_video_bitstream_builder_ std::vector::iterator placingPositionStart, size_t & writtenBytes); + void write_aud(std::vector & headerBitstream, + std::vector::iterator placingPositionStart, + D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC frameType, + size_t & writtenBytes); + void print_vps(const HevcVideoParameterSet& vps); void print_sps(const HevcSeqParameterSet& sps); void print_pps(const HevcPicParameterSet& pps); diff --git a/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.cpp b/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.cpp index a14c71b1565..3046b3ad55e 100644 --- a/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.cpp +++ b/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.cpp @@ -90,6 +90,51 @@ d3d12_video_nalu_writer_hevc::write_end_of_sequence_nalu(std::vector &head generic_write_bytes(headerBitstream, placingPositionStart, writtenBytes, &endOfSeqNALU); } +void +d3d12_video_nalu_writer_hevc::write_aud(std::vector &headerBitstream, + std::vector::iterator placingPositionStart, + D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC frameType, + size_t &writtenBytes) { + HevcAccessUnitDelimiter AUD = {}; + + AUD.nalu = + { + // uint8_t forbidden_zero_bit; + static_cast(0u), + // uint8_t nal_unit_type + static_cast(HEVC_NALU_AUD_NUT), + // uint8_t nuh_layer_id + static_cast(0u), + // uint8_t nuh_temporal_id_plus1 + static_cast(1u) + }; + + // pic_type slice_type values that may be present in the coded picture + switch (frameType) + { + case D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC_IDR_FRAME: + case D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC_I_FRAME: + { + AUD.pic_type = 0u; // 0 I + } break; + case D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC_P_FRAME: + { + AUD.pic_type = 1u; // 1 P, I + } break; + case D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC_B_FRAME: + { + AUD.pic_type = 2u; // 2 B, P, I + } break; + default: + { + debug_printf("d3d12_video_nalu_writer_hevc::write_aud failed: Invalid D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC frameType argument \n"); + assert(false); + } break; + } + + generic_write_bytes(headerBitstream, placingPositionStart, writtenBytes, &AUD); +} + void d3d12_video_nalu_writer_hevc::generic_write_bytes( std::vector &headerBitstream, std::vector::iterator placingPositionStart, @@ -161,6 +206,10 @@ d3d12_video_nalu_writer_hevc::write_bytes_from_struct(d3d12_video_encoder_bitstr // Do nothing for these two, just the header suffices return 1; } break; + case HEVC_NALU_AUD_NUT: + { + return write_aud_bytes(pBitstream, (HevcAccessUnitDelimiter *) pData); + } break; default: { unreachable("Unsupported NALU value"); @@ -168,6 +217,20 @@ d3d12_video_nalu_writer_hevc::write_bytes_from_struct(d3d12_video_encoder_bitstr } } +uint32_t +d3d12_video_nalu_writer_hevc::write_aud_bytes(d3d12_video_encoder_bitstream *pBitstream, HevcAccessUnitDelimiter *pAUD) +{ + int32_t iBytesWritten = pBitstream->get_byte_count(); + + pBitstream->put_bits(3, pAUD->pic_type); + + rbsp_trailing(pBitstream); + pBitstream->flush(); + + iBytesWritten = pBitstream->get_byte_count() - iBytesWritten; + return (uint32_t) iBytesWritten; +} + uint32_t d3d12_video_nalu_writer_hevc::write_vps_bytes(d3d12_video_encoder_bitstream *pBitstream, HevcVideoParameterSet *vps) { diff --git a/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.h b/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.h index 6359143cc0c..03b13e65eb7 100644 --- a/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.h +++ b/src/gallium/drivers/d3d12/d3d12_video_encoder_nalu_writer_hevc.h @@ -318,6 +318,11 @@ struct HevcVideoParameterSet { uint8_t vps_extension_data_flag; }; +struct HevcAccessUnitDelimiter { + HEVCNaluHeader nalu; + uint8_t pic_type; +}; + class d3d12_video_nalu_writer_hevc { public: @@ -362,12 +367,17 @@ public: placingPositionStart, size_t &writtenBytes); + void write_aud(std::vector & headerBitstream, + std::vector::iterator placingPositionStart, + D3D12_VIDEO_ENCODER_FRAME_TYPE_HEVC frameType, + size_t & writtenBytes); private: // Writes from structure into bitstream with RBSP trailing but WITHOUT NAL unit wrap (eg. nal_idc_type, etc) uint32_t write_vps_bytes(d3d12_video_encoder_bitstream *pBitstream, HevcVideoParameterSet *pSPS); uint32_t write_sps_bytes(d3d12_video_encoder_bitstream *pBitstream, HevcSeqParameterSet *pSPS); uint32_t write_pps_bytes(d3d12_video_encoder_bitstream *pBitstream, HevcPicParameterSet *pPPS); + uint32_t write_aud_bytes(d3d12_video_encoder_bitstream *pBitstream, HevcAccessUnitDelimiter *pAUD); // Adds NALU wrapping into structures and ending NALU control bits uint32_t wrap_rbsp_into_nalu(d3d12_video_encoder_bitstream *pNALU, d3d12_video_encoder_bitstream *pRBSP, HEVCNaluHeader *pHeader);