Unify output tensor allocation across all callers and devices #321

NicolasHug · 2024-10-30T10:23:06Z

This PR unifies the output tensor frame allocation within a single util: allocateEmptyHWCTensorForStream.

All high-level decoding entry-points are now responsible for allocating that output tensor.

All lower-level decoding entry-points now accept and expect this pre-allocated tensor to be passed. These low-level entry-points do not allocate anything that gets returned.

TODO, as a follow-up: try to fold the preAllocatedOutputTensor parameter within the rawOutput struct.

NicolasHug · 2024-10-31T14:18:17Z

src/torchcodec/decoders/_core/VideoDecoder.cpp

-      ptsSeconds(torch::empty({numFrames}, {torch::kFloat64})),
-      durationSeconds(torch::empty({numFrames}, {torch::kFloat64})) {}
+torch::Tensor VideoDecoder::allocateEmptyHWCTensorForStream(
+    int streamIndex,


Please double check this:

To allocate a tensor, we need its height and width. To get those, we need a streamIndex[1]. To get the streamIndex, we need either:

it to be passed as a parameter

for the "NoDemux" functions, we need at least a RawOutput.

[1] well, not strictly true, there are some cases where it's not needed, but the whole point of this PR is to unify that logic.

In the most generic case this isn't true because a single stream can have frames with different heights and widths:

#312

But ignoring that fact, to get the frame dimensions before decoding you do need the stream_index

NicolasHug · 2024-10-31T14:19:36Z

src/torchcodec/decoders/_core/VideoDecoder.cpp

+  }
+}
+
+VideoDecoder::BatchDecodedOutput VideoDecoder::allocateBatchDecodedOutput(


The previous BatchDecodedOutput constructor wasn't a member of VideoDecoder, and so I got compilation errors because we're calling allocateEmptyHWCTensorForStream which itself is a member method (because it accesses attributes blahblahblah).

So I had to make the "constructor" a member itself.

I understand that but by making it a member we are letting it access a lot more information than it strictly needs. It only needs height, width, device, and not all the other stuff in VideoDecoder. What's worse is it can accidentally change the state of the VideoDecoder because it's not const

I think it may be better to pass in the height and width after looking it up in the caller.

A general principle I like to apply is that classes should know how to create correct objects, and ideally, constructed objects should be valid. That means classes should have constructors, and the construtor is responsible for making the object valid. So I think we should keep the BatchDecodedOutput explicit constructor. If VideoDecoder then needs to construct BatchDecodedOutput objects using some of its own internal state, then VideoDecoder should figure out what's relevant to BatchDecodedOutput and pass it in.

We may need a VideoDecoder::allocatedBatchDecodedOutput() method, but if we do, it should call BatchDecodedOutput's explicit constructor. What this could potentially look like:

VideoDecoder::BatchDecodedOutput::BatchDecodedOutput( int64_t numFrames, torch::Tensor frames) : frames(frames), ptsSeconds(torch::empty({numFrames}, {torch::kFloat64})), durationSeconds(torch::empty({numFrames}, {torch::kFloat64})) {} VideoDecoder::BatchDecodedOutput VideoDecoder::allocateBatchDecodedOutput( int streamIndex, int64_t numFrames) { return BatchDecodedOutput( numFrames, allocateEmptyHWCTensorForStream(streamIndex, numFrames) ); }

We can take advantage of tensors have reasonable move-like semantics.

ahmadsharif1

#312 is closely related to this PR, hence I am bringing it up. You can probably address that next.

For this PR, for non-batch functions we should allocate the tensor using the frame's own width and height. That way users can get all the frames in a video whose stream has varying dimensions by calling non-batch functions.

For batch functions we can use streamIndex and if the user wants to get frames that are of different sizes, we can either return an error (for now) or resize them later.

ahmadsharif1 · 2024-10-31T14:42:21Z

src/torchcodec/decoders/_core/VideoDecoder.cpp

-      ptsSeconds(torch::empty({numFrames}, {torch::kFloat64})),
-      durationSeconds(torch::empty({numFrames}, {torch::kFloat64})) {}
+torch::Tensor VideoDecoder::allocateEmptyHWCTensorForStream(
+    int streamIndex,


In the most generic case this isn't true because a single stream can have frames with different heights and widths:

#312

But ignoring that fact, to get the frame dimensions before decoding you do need the stream_index

ahmadsharif1 · 2024-10-31T14:48:07Z

src/torchcodec/decoders/_core/VideoDecoder.cpp

+  }
+}
+
+VideoDecoder::BatchDecodedOutput VideoDecoder::allocateBatchDecodedOutput(


I understand that but by making it a member we are letting it access a lot more information than it strictly needs. It only needs height, width, device, and not all the other stuff in VideoDecoder. What's worse is it can accidentally change the state of the VideoDecoder because it's not const

I think it may be better to pass in the height and width after looking it up in the caller.

ahmadsharif1 · 2024-10-31T14:54:04Z

src/torchcodec/decoders/_core/VideoDecoder.cpp

-  auto output = convertAVFrameToDecodedOutput(rawOutput);
-  output.frame = MaybePermuteHWC2CHW(output.streamIndex, output.frame);
+  auto streamIndex = rawOutput.streamIndex;
+  auto preAllocatedOutputTensor = allocateEmptyHWCTensorForStream(streamIndex);


In the light of this: #312, it would be more correct to use rawOutput.frame's dimensions, and not the stream's dimensions

ahmadsharif1 · 2024-10-31T14:55:13Z

src/torchcodec/decoders/_core/VideoDecoder.cpp

-  output.frame = MaybePermuteHWC2CHW(output.streamIndex, output.frame);
+  auto rawOutput = getNextRawDecodedOutputNoDemux();
+  auto streamIndex = rawOutput.streamIndex;
+  auto preAllocatedOutputTensor = allocateEmptyHWCTensorForStream(streamIndex);


ahmadsharif1 · 2024-10-31T15:00:09Z

src/torchcodec/decoders/_core/VideoDecoder.h

+  // --------------------------------------------------------------------------
+  // Tensor (frames) manipulation APIs
+  // --------------------------------------------------------------------------
+  torch::Tensor MaybePermuteHWC2CHW(int streamIndex, torch::Tensor& hwcTensor);


This doesn't need to be in the header at all.

This should just be a utility function in an anonymous namespace in the cpp file. You can pass it the height and width.

Same goes for function below.

NicolasHug · 2024-11-04T14:18:48Z

I've discussed this PR with both of you @ahmadsharif1 @scotts and the conclusion is that bubbling-up the tensor allocation might not be the best move at this time, because it would change the source of where we retrieve height and width from.

I'll follow-up soon with some alternative improvements. I want to at least document why the existing code is the way it is, and hopefully make the logic a bit clearer via dedicated utils.

NicolasHug added 10 commits October 29, 2024 13:10

Remove call from convertAVFrameToDecodedOutputOnCPU

7c37e9a

Nit

f273b52

Slighlty better

639952a

No more if

3ae34c7

Make getNextDecodedOutputNoDemux private

8019219

Fix CUDA tests?

dc40ad6

getNextFrame -> getNextFrameNoDemux

6bd363d

getNextDecodedOutputNoDemux -> getNextFrameOutputNoDemuxInternal

0b6590d

Merge branch 'main' of github.com:pytorch/torchcodec into bubble-up

1c01ec7

very WIP

9efa08f

facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 30, 2024

NicolasHug and others added 8 commits October 31, 2024 10:21

Merge branch 'main' of github.com:pytorch/torchcodec into bubble-up

3ab551e

getFrameAtIndexInternal now *must* take preAllocatedTensor

0568482

Also handle getFramePlayedAtTimestampNoDemux

649401e

Handle getNextFrameNoDemux

7cf3a3e

Unify allocation within allocateEmptyHWCTensorForStream

8c11e43

pre-alloc tensor now mandatory in convertAVFrameToDecodedOutputOnCPU

5c6e23c

Handle GPU?

fcbe4a3

Fix CUDA

7322f9e

NicolasHug changed the title ~~[WPI] Output tensor allocation must now happen in high-level entry-points~~ Output tensor allocation must now happen in high-level entry-points Oct 31, 2024

NicolasHug changed the title ~~Output tensor allocation must now happen in high-level entry-points~~ Unify output tensor allocation across all callers and devices Oct 31, 2024

NicolasHug marked this pull request as ready for review October 31, 2024 14:12

NicolasHug commented Oct 31, 2024

View reviewed changes

ahmadsharif1 reviewed Oct 31, 2024

View reviewed changes

NicolasHug closed this Nov 4, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unify output tensor allocation across all callers and devices #321

Unify output tensor allocation across all callers and devices #321

NicolasHug commented Oct 30, 2024 •

edited

Loading

NicolasHug Oct 31, 2024

ahmadsharif1 Oct 31, 2024

NicolasHug Oct 31, 2024

ahmadsharif1 Oct 31, 2024 •

edited

Loading

scotts Oct 31, 2024

ahmadsharif1 left a comment

ahmadsharif1 Oct 31, 2024

ahmadsharif1 Oct 31, 2024 •

edited

Loading

ahmadsharif1 Oct 31, 2024

ahmadsharif1 Oct 31, 2024

ahmadsharif1 Oct 31, 2024

NicolasHug commented Nov 4, 2024 •

edited

Loading

Unify output tensor allocation across all callers and devices #321

Unify output tensor allocation across all callers and devices #321

Conversation

NicolasHug commented Oct 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ahmadsharif1 Oct 31, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ahmadsharif1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ahmadsharif1 Oct 31, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

NicolasHug commented Nov 4, 2024 • edited Loading

NicolasHug commented Oct 30, 2024 •

edited

Loading

ahmadsharif1 Oct 31, 2024 •

edited

Loading

ahmadsharif1 Oct 31, 2024 •

edited

Loading

NicolasHug commented Nov 4, 2024 •

edited

Loading