Optimize WAV decoding with buffered read

[Dan-Flores]

This PR enables the optimized file reading from this TODO:

TODO WavDecoder: Optimize decoding and converison by processing samples in
fixed size buffer. For now, we use a simpler implementation: create a
tensor from raw bytes and convert tensor to float32

The optimization: read samples and convert them through a fixed size buffer, rather than reading all values at once.

• In WavDecoder::getSamplesInRange, we create a buffer close to DEFAULT_CHUNK_BUFFER_SIZE bytes.
  • The actual buffer size is adjusted based on the numBytesPerSample, to ensure we do not read partial samples.
• In WavDecoder::getSamplesInRange, we process samples according to the buffer size and create the output samples tensor.
• WavDecoder::convertSamplesToFloat does the actual conversion to float32.

Additional changes:
readValue is split into safeReadValue and readValue. This is necessary for performance, see inline comments for more details.

Benchmarks:

Details

Benchmarking code

Audio Decoding Benchmark

File: test/resources/sine_stereo_s32_10min.wav
Warmup: 3, Trials: 10

AudioDecoder 0.249s
WavDecoder 0.148s
SoundFile 0.146s

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/meta-pytorch/torchcodec/1320

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 8480eec with merge base 7e8c399 ():

NEW FAILURE - The following job has failed:

• Lint / mypy (3.12) (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Dan-Flores]

Here we use readValue() without the additional TORCH_CHECKs in safeReadValue() for performance.

It is safe to read values without the bounds checks because the TORCH_CHECK above ensures that bufferData is large enough for totalSamples.

[mollyxu]

From my understanding, numSamples = dataSize / numBytesPerSample. Would we want to validate that dataSize in case the headerValue is some arbitrarily large number due to corruption?

[Dan-Flores]

For integer overflow, we do not need to validate before doing that division, because int32_t / uint16_t will always fit inside int64_t.

For general overallocation concerns, this is a bit trickier to validate. Since dataSize is int32_t, the maximum valid value it can have is substantial but not always incorrect. I'll add a clamp to ensure the tensor allocated is not larger than the actual file, to prevent an incorrect header from causing us to overallocate.

[Dan-Flores]

To address the other dimension of our allocated tensor, header_.numChannels: there is no equivalent ground truth as the filesystem's fileSize is to dataSize.

Instead, I've added a validation in validateHeader() for numBytesPerSample that the WAV specification already requires (labelled here as BytePerBloc):

numBytesPerSample == numChannels * bitsPerSample / 8

We can substitute the values used to calculate numSamples to see that the number of elements in the tensor is bounded by the file size:

samples tensor elements = numSamples * header_.numChannels
	= (std::min(header_.dataSize, fileSize_) / header_.numBytesPerSample)) * header_.numChannels
	= (fileSize_ / header_.numBytesPerSample) * header_.numChannels
	= (fileSize_ / (numChannels * bitsPerSample / 8)) * header_.numChannels
	= fileSize_ / (bitsPerSample / 8)

As an aside: If the header is corrupted such that numSamples is calculated incorrectly, we were already going to fail in WavDecoder::getSamplesInRange() when we call safeReadFile() to read data beyond what exists. Clamping dataSize to fileSize doesn't prevent that failure either, but it does prevent us from overallocating before failing.

[NicolasHug]

I think we should avoid naming this a "chunk".

A "chunk" in the context of wav files refer to the "format chunk" or "data chunk". This thing is just a temporary buffer that 'views' the data chunk, but it's not a chunk of its own (i.e. it's not a part of the wav specs).

[NicolasHug]

Can you show me where this is validated?

By 'above' I think you mean 'by the caller', but I don't see where. There should also be a comment about this expectation at the top of the function, similar to the one that exists for outputPtr.

[Dan-Flores]

The chain of intermediate variables passed to readValue() from getSamplesInRange() is complex, so I am less confident in this claim. Since we create the buffer and determine the number of samples we read, we should be safe to read from the buffer without checking the length each time, but this could probably use some refactoring to make this more clear.