Set the bitstream buffer size to match the input image size.

For higher image resolutions and bitrates, the fixed 2MB buffer size is inadequate. To handle the worst-case scenario, ensure the bitstream size is the same as the input image size. Signed-off-by: Raju Konda <kraju@nvidia.com>
nvpro-samples · Mar 11, 2025 · 84b2703 · 84b2703
1 parent b5d0d26
commit 84b2703
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 9 deletions.
diff --git a/vk_video_encoder/libs/VkVideoEncoder/VkVideoEncoder.cpp b/vk_video_encoder/libs/VkVideoEncoder/VkVideoEncoder.cpp
@@ -539,11 +539,20 @@ VkResult VkVideoEncoder::AssembleBitstreamData(VkSharedBaseObj<VkVideoEncodeFram
     VkDeviceSize maxSize;
     uint8_t* data = encodeFrameInfo->outputBitstreamBuffer->GetDataPtr(0, maxSize);
 
-    size_t vcl = fwrite(data + encodeResult.bitstreamStartOffset, 1, encodeResult.bitstreamSize,
-                        m_encoderConfig->outputFileHandler.GetFileHandle());
+    size_t totalBytesWritten = 0;
+    while (totalBytesWritten < encodeResult.bitstreamSize) { // handle partial writes
+        size_t remainingBytes = encodeResult.bitstreamSize - totalBytesWritten;
+        size_t bytesWritten = fwrite(data + encodeResult.bitstreamStartOffset + totalBytesWritten, 1, remainingBytes,
+                                    m_encoderConfig->outputFileHandler.GetFileHandle());
+        if (bytesWritten == 0) {
+            std::cerr << "Error writing VCL data" << std::endl;
+            return VK_ERROR_OUT_OF_HOST_MEMORY;
+        }
+        totalBytesWritten += bytesWritten;
+    }
 
     if (m_encoderConfig->verboseFrameStruct) {
-        std::cout << "       == Output VCL data " << (vcl ? "SUCCESS" : "FAIL") << " with size: " << encodeResult.bitstreamSize
+        std::cout << "       == Output VCL data " << ((totalBytesWritten == encodeResult.bitstreamSize) ? "SUCCESS" : "FAIL") << " with size: " << encodeResult.bitstreamSize
                   << " and offset: " << encodeResult.bitstreamStartOffset
                   << ", Input Order: " << encodeFrameInfo->gopPosition.inputOrder
                   << ", Encode  Order: " << encodeFrameInfo->gopPosition.encodeOrder << std::endl << std::flush;
@@ -687,6 +696,7 @@ VkResult VkVideoEncoder::InitEncoder(VkSharedBaseObj<EncoderConfig>& encoderConf
     }
 
     m_maxCodedExtent = { encoderConfig->encodeMaxWidth, encoderConfig->encodeMaxHeight }; // max coded size
+    m_streamBufferSize = std::max(m_minStreamBufferSize, (size_t)encoderConfig->input.fullImageSize); // use worst case size
 
     const uint32_t maxActiveReferencePicturesCount = encoderConfig->videoCapabilities.maxActiveReferencePictures;
     const uint32_t maxDpbPicturesCount = std::min<uint32_t>(m_maxDpbPicturesCount, encoderConfig->videoCapabilities.maxDpbSlots);

diff --git a/vk_video_encoder/libs/VkVideoEncoder/VkVideoEncoderAV1.cpp b/vk_video_encoder/libs/VkVideoEncoder/VkVideoEncoderAV1.cpp
@@ -955,12 +955,29 @@ VkResult VkVideoEncoderAV1::AssembleBitstreamData(VkSharedBaseObj<VkVideoEncodeF
         }
 
         for (const auto& curIndex : m_batchFramesIndxSetToAssemble) {
-            if (frameIdx == curIndex) {
-                fwrite(data + encodeResult.bitstreamStartOffset, 1, encodeResult.bitstreamSize,
-                                   m_encoderConfig->outputFileHandler.GetFileHandle());
-            } else {
-                fwrite(m_bitstream[curIndex].data(), 1, m_bitstream[curIndex].size(),
-                                   m_encoderConfig->outputFileHandler.GetFileHandle());
+            const uint8_t* writeData = (frameIdx == curIndex) ? (data + encodeResult.bitstreamStartOffset) : m_bitstream[curIndex].data();
+            const size_t bytesToWrite = (frameIdx == curIndex) ? encodeResult.bitstreamSize : m_bitstream[curIndex].size();
+
+            // Write data in chunks to handle partial writes
+            size_t totalBytesWritten = 0;
+            while (totalBytesWritten < bytesToWrite) {
+                const size_t remainingBytes = bytesToWrite - totalBytesWritten;
+                const size_t bytesWritten = fwrite(writeData + totalBytesWritten, 1, 
+                                                 remainingBytes,
+                                                 m_encoderConfig->outputFileHandler.GetFileHandle());
+
+                if (bytesWritten == 0) {
+                    std::cerr << "Failed to write bitstream data" << std::endl;
+                    return VK_ERROR_OUT_OF_HOST_MEMORY;
+                }
+
+                totalBytesWritten += bytesWritten;
+            }
+
+            // Verify complete write
+            if (totalBytesWritten != bytesToWrite) {
+                std::cerr << "Warning: Incomplete write - expected " << bytesToWrite << " bytes but wrote " << totalBytesWritten << " bytes\n";
+                return VK_ERROR_OUT_OF_HOST_MEMORY;
             }
         }
         // reset the batch frames to assemble counter