Skip to content
Merged
Show file tree
Hide file tree
Changes from 32 commits
Commits
Show all changes
36 commits
Select commit Hold shift + click to select a range
78ab058
Let's just commit 3k loc in a single commit
NicolasHug Sep 25, 2025
b45decc
Fixes
NicolasHug Sep 26, 2025
316f218
Merge branch 'main' of github.com:pytorch/torchcodec into aeaenjfjanef
NicolasHug Sep 30, 2025
d0192ec
GetCache -> getCache
NicolasHug Sep 30, 2025
515deb5
Make UniqueCUvideodecoder a pointer on CUvideodecoder, not void
NicolasHug Sep 30, 2025
13fad10
Make device and device_variant have a default instead of being std::o…
NicolasHug Sep 30, 2025
eb8de72
Remove old registerDeviceInterface
NicolasHug Sep 30, 2025
4f7a4fb
Call std::memset
NicolasHug Sep 30, 2025
dcf3124
remove unnecessary cuda_runtime.h include, update cmake accordingly
NicolasHug Sep 30, 2025
0ad7370
abstract frameBuffer_ into a FrameBuffer class
NicolasHug Sep 30, 2025
aad142e
Cleanup BSF logic
NicolasHug Sep 30, 2025
2592888
Return int in callback instead of unsigned char
NicolasHug Sep 30, 2025
b5fe9bc
define width and height as unsigned int
NicolasHug Sep 30, 2025
5605c90
Rework frame ordering and pts matching
NicolasHug Oct 1, 2025
7494259
Merge branch 'main' of github.com:pytorch/torchcodec into aeaenjfjanef
NicolasHug Oct 1, 2025
560b376
Fix cuda context initialization
NicolasHug Oct 1, 2025
88196c5
Merge branch 'aeaenjfjanef' into nvdec-rework-frame-ordering
NicolasHug Oct 1, 2025
2a78b84
Renaming
NicolasHug Oct 1, 2025
5d194e5
Comment
NicolasHug Oct 1, 2025
d1e51b3
Merge branch 'main' of github.com:pytorch/torchcodec into aeaenjfjanef
NicolasHug Oct 2, 2025
f9c7297
Skip equality check on ffmepg 4
NicolasHug Oct 2, 2025
b7bbfb2
Merge branch 'aeaenjfjanef' into nvdec-rework-frame-ordering
NicolasHug Oct 2, 2025
390fd7c
Refac, simplify
NicolasHug Oct 2, 2025
f55dcc0
Update comment
NicolasHug Oct 2, 2025
7e4dd10
Define constant, add TODO for AVRational
NicolasHug Oct 2, 2025
f614846
Use uint32_t types
NicolasHug Oct 2, 2025
aa6e253
Create packet.reset() and add P0 TODO
NicolasHug Oct 2, 2025
186eaa4
Add TODO
NicolasHug Oct 2, 2025
1cb4890
Merge branch 'aeaenjfjanef' into nvdec-rework-frame-ordering
NicolasHug Oct 2, 2025
c5b32a4
Merge branch 'main' of github.com:pytorch/torchcodec into nvdec-rewor…
NicolasHug Oct 2, 2025
70873bf
lint
NicolasHug Oct 2, 2025
8e73bcf
Add TODOs and more explicit initialization
NicolasHug Oct 3, 2025
12c75e7
Add h265 support
NicolasHug Oct 2, 2025
9b63504
Merge branch 'main' of github.com:pytorch/torchcodec into nvdec-param…
NicolasHug Oct 3, 2025
718a3e3
Merge branch 'nvdec-params-and-todos' into nvdec-h265
NicolasHug Oct 3, 2025
c42388f
put initializeBSF below
NicolasHug Oct 3, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
220 changes: 115 additions & 105 deletions src/torchcodec/_core/BetaCudaDeviceInterface.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -41,10 +41,17 @@ pfnSequenceCallback(void* pUserData, CUVIDEOFORMAT* videoFormat) {
}

static int CUDAAPI
pfnDecodePictureCallback(void* pUserData, CUVIDPICPARAMS* pPicParams) {
pfnDecodePictureCallback(void* pUserData, CUVIDPICPARAMS* picParams) {
BetaCudaDeviceInterface* decoder =
static_cast<BetaCudaDeviceInterface*>(pUserData);
return decoder->frameReadyForDecoding(pPicParams);
return decoder->frameReadyForDecoding(picParams);
}

static int CUDAAPI
pfnDisplayPictureCallback(void* pUserData, CUVIDPARSERDISPINFO* dispInfo) {
BetaCudaDeviceInterface* decoder =
static_cast<BetaCudaDeviceInterface*>(pUserData);
return decoder->frameReadyInDisplayOrder(dispInfo);
}

static UniqueCUvideodecoder createDecoder(CUVIDEOFORMAT* videoFormat) {
Expand Down Expand Up @@ -131,6 +138,24 @@ static UniqueCUvideodecoder createDecoder(CUVIDEOFORMAT* videoFormat) {
return UniqueCUvideodecoder(decoder, CUvideoDecoderDeleter{});
}

cudaVideoCodec validateCodecSupport(AVCodecID codecId) {
switch (codecId) {
case AV_CODEC_ID_H264:
return cudaVideoCodec_H264;
case AV_CODEC_ID_HEVC:
return cudaVideoCodec_HEVC;
// TODONVDEC P0: support more codecs
// case AV_CODEC_ID_AV1: return cudaVideoCodec_AV1;
// case AV_CODEC_ID_MPEG4: return cudaVideoCodec_MPEG4;
// case AV_CODEC_ID_VP8: return cudaVideoCodec_VP8;
// case AV_CODEC_ID_VP9: return cudaVideoCodec_VP9;
// case AV_CODEC_ID_MJPEG: return cudaVideoCodec_JPEG;
default: {
TORCH_CHECK(false, "Unsupported codec type: ", avcodec_get_name(codecId));
}
}
}

} // namespace

BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
Expand All @@ -142,7 +167,7 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)

BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
// TODONVDEC P0: we probably need to free the frames that have been decoded by
// NVDEC but not yet "mapped" - i.e. those that are still in frameBuffer_?
// NVDEC but not yet "mapped" - i.e. those that are still in readyFrames_?

if (decoder_) {
NVDECCache::getCache(device_.index())
Expand All @@ -156,29 +181,62 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
}
}

void BetaCudaDeviceInterface::initializeInterface(AVStream* avStream) {
torch::Tensor dummyTensorForCudaInitialization = torch::empty(
{1}, torch::TensorOptions().dtype(torch::kUInt8).device(device_));
void BetaCudaDeviceInterface::initializeBSF(
const AVCodecParameters* codecPar,
const UniqueDecodingAVFormatContext& avFormatCtx) {
// Setup bit stream filters (BSF):
// https://ffmpeg.org/doxygen/7.0/group__lavc__bsf.html
// This is only needed for some formats, like H264 or HEVC.

TORCH_CHECK(avStream != nullptr, "AVStream cannot be null");
timeBase_ = avStream->time_base;
TORCH_CHECK(codecPar != nullptr, "codecPar cannot be null");
TORCH_CHECK(avFormatCtx != nullptr, "AVFormatContext cannot be null");
TORCH_CHECK(
avFormatCtx->iformat != nullptr,
"AVFormatContext->iformat cannot be null");
std::string filterName;

// Matching logic is taken from DALI
switch (codecPar->codec_id) {
case AV_CODEC_ID_H264: {
const std::string formatName = avFormatCtx->iformat->long_name
? avFormatCtx->iformat->long_name
: "";

if (formatName == "QuickTime / MOV" ||
formatName == "FLV (Flash Video)" ||
formatName == "Matroska / WebM" || formatName == "raw H.264 video") {
filterName = "h264_mp4toannexb";
}
break;
}

const AVCodecParameters* codecpar = avStream->codecpar;
TORCH_CHECK(codecpar != nullptr, "CodecParameters cannot be null");
case AV_CODEC_ID_HEVC: {
const std::string formatName = avFormatCtx->iformat->long_name
? avFormatCtx->iformat->long_name
: "";

TORCH_CHECK(
// TODONVDEC P0 support more
avStream->codecpar->codec_id == AV_CODEC_ID_H264,
"Can only do H264 for now");
if (formatName == "QuickTime / MOV" ||
formatName == "FLV (Flash Video)" ||
formatName == "Matroska / WebM" || formatName == "raw HEVC video") {
filterName = "hevc_mp4toannexb";
}
break;
}

// Setup bit stream filters (BSF):
// https://ffmpeg.org/doxygen/7.0/group__lavc__bsf.html
// This is only needed for some formats, like H264 or HEVC. TODONVDEC P1: For
// now we apply BSF unconditionally, but it should be optional and dependent
// on codec and container.
const AVBitStreamFilter* avBSF = av_bsf_get_by_name("h264_mp4toannexb");
default:
// No bitstream filter needed for other codecs
// TODONVDEC P1 MPEG4 will need one!
break;
}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: I would prefer putting this switching logic into a function - I know we have a style difference there. :) My rationale is that I find it useful to think in terms of pure functions when I can, and this can definitely be a pure function, and then we can simply say in this scope:

std::string filterName = toFilterName(codecPar->codec_id);

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIUC we'd need toFilterName to return "" when no match exist, and it would be up to the caller to check against that?


if (filterName.empty()) {
// Only initialize BSF if we actually need one
return;
}

const AVBitStreamFilter* avBSF = av_bsf_get_by_name(filterName.c_str());
TORCH_CHECK(
avBSF != nullptr, "Failed to find h264_mp4toannexb bitstream filter");
avBSF != nullptr, "Failed to find bitstream filter: ", filterName);

AVBSFContext* avBSFContext = nullptr;
int retVal = av_bsf_alloc(avBSF, &avBSFContext);
Expand All @@ -189,7 +247,7 @@ void BetaCudaDeviceInterface::initializeInterface(AVStream* avStream) {

bitstreamFilter_.reset(avBSFContext);

retVal = avcodec_parameters_copy(bitstreamFilter_->par_in, codecpar);
retVal = avcodec_parameters_copy(bitstreamFilter_->par_in, codecPar);
TORCH_CHECK(
retVal >= AVSUCCESS,
"Failed to copy codec parameters: ",
Expand All @@ -200,18 +258,33 @@ void BetaCudaDeviceInterface::initializeInterface(AVStream* avStream) {
retVal == AVSUCCESS,
"Failed to initialize bitstream filter: ",
getFFMPEGErrorStringFromErrorCode(retVal));
}

void BetaCudaDeviceInterface::initializeInterface(
const AVStream* avStream,
const UniqueDecodingAVFormatContext& avFormatCtx) {
torch::Tensor dummyTensorForCudaInitialization = torch::empty(
{1}, torch::TensorOptions().dtype(torch::kUInt8).device(device_));

TORCH_CHECK(avStream != nullptr, "AVStream cannot be null");
timeBase_ = avStream->time_base;

const AVCodecParameters* codecPar = avStream->codecpar;
TORCH_CHECK(codecPar != nullptr, "CodecParameters cannot be null");

initializeBSF(codecPar, avFormatCtx);

// Create parser. Default values that aren't obvious are taken from DALI.
CUVIDPARSERPARAMS parserParams = {};
parserParams.CodecType = cudaVideoCodec_H264;
parserParams.CodecType = validateCodecSupport(codecPar->codec_id);
parserParams.ulMaxNumDecodeSurfaces = 8;
parserParams.ulMaxDisplayDelay = 0;
// Callback setup, all are triggered by the parser within a call
// to cuvidParseVideoData
parserParams.pUserData = this;
parserParams.pfnSequenceCallback = pfnSequenceCallback;
parserParams.pfnDecodePicture = pfnDecodePictureCallback;
parserParams.pfnDisplayPicture = nullptr;
parserParams.pfnDisplayPicture = pfnDisplayPictureCallback;

CUresult result = cuvidCreateVideoParser(&videoParser_, &parserParams);
TORCH_CHECK(
Expand Down Expand Up @@ -267,10 +340,6 @@ int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
cuvidPacket.flags = CUVID_PKT_TIMESTAMP;
cuvidPacket.timestamp = packet->pts;

// Like DALI: store packet PTS in queue to later assign to frames as they
// come out
packetsPtsQueue.push(packet->pts);

} else {
// End of stream packet
cuvidPacket.flags = CUVID_PKT_ENDOFSTREAM;
Expand Down Expand Up @@ -322,70 +391,38 @@ void BetaCudaDeviceInterface::applyBSF(ReferenceAVPacket& packet) {
// ready to be decoded, i.e. the parser received all the necessary packets for a
// given frame. It means we can send that frame to be decoded by the hardware
// NVDEC decoder by calling cuvidDecodePicture which is non-blocking.
int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* pPicParams) {
int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* picParams) {
if (isFlushing_) {
return 0;
}

TORCH_CHECK(pPicParams != nullptr, "Invalid picture parameters");
TORCH_CHECK(picParams != nullptr, "Invalid picture parameters");
TORCH_CHECK(decoder_, "Decoder not initialized before picture decode");

// Send frame to be decoded by NVDEC - non-blocking call.
CUresult result = cuvidDecodePicture(*decoder_.get(), pPicParams);
if (result != CUDA_SUCCESS) {
return 0; // Yes, you're reading that right, 0 mean error.
}
CUresult result = cuvidDecodePicture(*decoder_.get(), picParams);

// The frame was sent to be decoded on the NVDEC hardware. Now we store some
// relevant info into our frame buffer so that we can retrieve the decoded
// frame later when receiveFrame() is called.
// Importantly we need to 'guess' the PTS of that frame. The heuristic we use
// (like in DALI) is that the frames are ready to be decoded in the same order
// as the packets were sent to the parser. So we assign the PTS of the frame
// by popping the PTS of the oldest packet in our packetsPtsQueue (note:
// oldest doesn't necessarily mean lowest PTS!).
// Yes, you're reading that right, 0 means error, 1 means success
return (result == CUDA_SUCCESS);
}

TORCH_CHECK(
// TODONVDEC P0 the queue may be empty, handle that.
!packetsPtsQueue.empty(),
"PTS queue is empty when decoding a frame");
int64_t guessedPts = packetsPtsQueue.front();
packetsPtsQueue.pop();

// Field values taken from DALI
CUVIDPARSERDISPINFO dispInfo = {};
dispInfo.picture_index = pPicParams->CurrPicIdx;
dispInfo.progressive_frame = !pPicParams->field_pic_flag;
dispInfo.top_field_first = pPicParams->bottom_field_flag ^ 1;
dispInfo.repeat_first_field = 0;
dispInfo.timestamp = guessedPts;

FrameBuffer::Slot* slot = frameBuffer_.findEmptySlot();
slot->dispInfo = dispInfo;
slot->guessedPts = guessedPts;
slot->occupied = true;

return 1;
int BetaCudaDeviceInterface::frameReadyInDisplayOrder(
CUVIDPARSERDISPINFO* dispInfo) {
readyFrames_.push(*dispInfo);
return 1; // success
}

// Moral equivalent of avcodec_receive_frame(). Here, we look for a decoded
// frame with the exact desired PTS in our frame buffer. This logic is only
// valid in exact seek_mode, for now.
int BetaCudaDeviceInterface::receiveFrame(
UniqueAVFrame& avFrame,
int64_t desiredPts) {
FrameBuffer::Slot* slot = frameBuffer_.findFrameWithExactPts(desiredPts);
if (slot == nullptr) {
// Moral equivalent of avcodec_receive_frame().
int BetaCudaDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) {
if (readyFrames_.empty()) {
// No frame found, instruct caller to try again later after sending more
// packets.
return AVERROR(EAGAIN);
}

slot->occupied = false;
slot->guessedPts = -1;
CUVIDPARSERDISPINFO dispInfo = readyFrames_.front();
readyFrames_.pop();

CUVIDPROCPARAMS procParams = {};
CUVIDPARSERDISPINFO dispInfo = slot->dispInfo;
procParams.progressive_frame = dispInfo.progressive_frame;
procParams.top_field_first = dispInfo.top_field_first;
procParams.unpaired_field = dispInfo.repeat_first_field < 0;
Expand Down Expand Up @@ -445,7 +482,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame(
avFrame->width = width;
avFrame->height = height;
avFrame->format = AV_PIX_FMT_CUDA;
avFrame->pts = dispInfo.timestamp; // == guessedPts
avFrame->pts = dispInfo.timestamp;

// TODONVDEC P0: Zero division error!!!
// TODONVDEC P0: Move AVRational arithmetic to FFMPEGCommon, and put the
Expand Down Expand Up @@ -511,13 +548,8 @@ void BetaCudaDeviceInterface::flush() {

isFlushing_ = false;

for (auto& slot : frameBuffer_) {
slot.occupied = false;
slot.guessedPts = -1;
}

std::queue<int64_t> empty;
packetsPtsQueue.swap(empty);
std::queue<CUVIDPARSERDISPINFO> emptyQueue;
std::swap(readyFrames_, emptyQueue);

eofSent_ = false;
}
Expand Down Expand Up @@ -551,26 +583,4 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
preAllocatedOutputTensor);
}

BetaCudaDeviceInterface::FrameBuffer::Slot*
BetaCudaDeviceInterface::FrameBuffer::findEmptySlot() {
for (auto& slot : frameBuffer_) {
if (!slot.occupied) {
return &slot;
}
}
frameBuffer_.emplace_back();
return &frameBuffer_.back();
}

BetaCudaDeviceInterface::FrameBuffer::Slot*
BetaCudaDeviceInterface::FrameBuffer::findFrameWithExactPts(
int64_t desiredPts) {
for (auto& slot : frameBuffer_) {
if (slot.occupied && slot.guessedPts == desiredPts) {
return &slot;
}
}
return nullptr;
}

} // namespace facebook::torchcodec
Loading
Loading