Skip to content

Commit 7f5c195

Browse files
NicolasHugNicolasHug
committed
Pass color-conversion matrix as kernel parameter
1 parent 04f1300 commit 7f5c195

3 files changed

Lines changed: 79 additions & 102 deletions

File tree

‎src/torchcodec/_core/color_conversion.cpp‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ void computeColorConversionMatrix(
118118
}
119119
}
120120

121-
bool maybeUpdateColorMatrix(
121+
void maybeUpdateColorMatrix(
122122
CachedColorMatrix& cachedColorMatrix,
123123
AVColorSpace colorspace,
124124
AVColorRange colorRange,
@@ -128,7 +128,7 @@ bool maybeUpdateColorMatrix(
128128
cachedColorMatrix.colorRange == colorRange &&
129129
cachedColorMatrix.bitDepth == bitDepth &&
130130
cachedColorMatrix.outScale == outScale) {
131-
return false;
131+
return;
132132
}
133133

134134
computeColorConversionMatrix(
@@ -138,7 +138,6 @@ bool maybeUpdateColorMatrix(
138138
cachedColorMatrix.bitDepth = bitDepth;
139139
cachedColorMatrix.outScale = outScale;
140140
cachedColorMatrix.valid = true;
141-
return true;
142141
}
143142

144143
void computeRGBToYUVMatrix(
@@ -221,7 +220,7 @@ torch::stable::Tensor convertYUVFrameToRGB(
221220
syncStreams(
222221
/*runningStream=*/nvdecStream, /*waitingStream=*/stream);
223222

224-
bool colorMatrixChanged = maybeUpdateColorMatrix(
223+
maybeUpdateColorMatrix(
225224
cachedColorMatrix,
226225
avFrame->colorspace,
227226
avFrame->color_range,
@@ -240,7 +239,6 @@ torch::stable::Tensor convertYUVFrameToRGB(
240239
validateInt64ToInt(dst.stride(0) * 2, "dst.stride(0)*2"),
241240
bitDepth,
242241
cachedColorMatrix.matrix,
243-
colorMatrixChanged,
244242
stream);
245243
} else {
246244
launchNV12ToRGBKernel(
@@ -253,7 +251,6 @@ torch::stable::Tensor convertYUVFrameToRGB(
253251
avFrame->linesize[1],
254252
validateInt64ToInt(dst.stride(0), "dst.stride(0)"),
255253
cachedColorMatrix.matrix,
256-
colorMatrixChanged,
257254
stream);
258255
}
259256

‎src/torchcodec/_core/color_conversion.cu‎

Lines changed: 75 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@ namespace facebook::torchcodec {
1717
// each. P016 uses 16-bit samples (uint16_t) with actual data in the most
1818
// significant bits (right-shifted by 16 - bitDepth).
1919

20-
// Color conversion matrix stored in constant memory for fast access.
21-
// It is available to every single thread.
22-
__constant__ float d_colorMatrix[3][4];
20+
// Wraps the 3x4 color-conversion matrix so it can be passed as a kernel param,
21+
// to ensure thread-safety (previous implem had it __global__, which wasn't
22+
// safe)
23+
struct ColorMatrix {
24+
float m[3][4];
25+
};
2326

2427
// Takes a pair of consecutive Y values, a pair of UV values, and writes the
2528
// corresponding two RGB values. Instead of writing each ra ga ba and rb gb bb
@@ -34,33 +37,34 @@ __device__ void writePairOfRGBPixels(
3437
float u,
3538
float v,
3639
T* rgbPlaneToWrite,
37-
int bitShift) {
40+
int bitShift,
41+
const ColorMatrix& cm) {
3842
constexpr float clampMax = sizeof(T) == 1 ? 255.0f : 65535.0f;
3943

4044
float ya = static_cast<float>(yayb.x >> bitShift);
4145
float yb = static_cast<float>(yayb.y >> bitShift);
4246

43-
float ra = d_colorMatrix[0][0] * ya + d_colorMatrix[0][1] * u +
44-
d_colorMatrix[0][2] * v + d_colorMatrix[0][3];
45-
float ga = d_colorMatrix[1][0] * ya + d_colorMatrix[1][1] * u +
46-
d_colorMatrix[1][2] * v + d_colorMatrix[1][3];
47+
float ra = cm.m[0][0] * ya + cm.m[0][1] * u +
48+
cm.m[0][2] * v + cm.m[0][3];
49+
float ga = cm.m[1][0] * ya + cm.m[1][1] * u +
50+
cm.m[1][2] * v + cm.m[1][3];
4751

4852
Vec2T raga = {static_cast<T>(fminf(fmaxf(ra, 0.0f), clampMax)),
4953
static_cast<T>(fminf(fmaxf(ga, 0.0f), clampMax))};
5054
*(reinterpret_cast<Vec2T*>(&rgbPlaneToWrite[0])) = raga;
5155

52-
float ba = d_colorMatrix[2][0] * ya + d_colorMatrix[2][1] * u +
53-
d_colorMatrix[2][2] * v + d_colorMatrix[2][3];
54-
float rb = d_colorMatrix[0][0] * yb + d_colorMatrix[0][1] * u +
55-
d_colorMatrix[0][2] * v + d_colorMatrix[0][3];
56+
float ba = cm.m[2][0] * ya + cm.m[2][1] * u +
57+
cm.m[2][2] * v + cm.m[2][3];
58+
float rb = cm.m[0][0] * yb + cm.m[0][1] * u +
59+
cm.m[0][2] * v + cm.m[0][3];
5660
Vec2T barb = {static_cast<T>(fminf(fmaxf(ba, 0.0f), clampMax)),
5761
static_cast<T>(fminf(fmaxf(rb, 0.0f), clampMax))};
5862
*(reinterpret_cast<Vec2T*>(&rgbPlaneToWrite[2])) = barb;
5963

60-
float gb = d_colorMatrix[1][0] * yb + d_colorMatrix[1][1] * u +
61-
d_colorMatrix[1][2] * v + d_colorMatrix[1][3];
62-
float bb = d_colorMatrix[2][0] * yb + d_colorMatrix[2][1] * u +
63-
d_colorMatrix[2][2] * v + d_colorMatrix[2][3];
64+
float gb = cm.m[1][0] * yb + cm.m[1][1] * u +
65+
cm.m[1][2] * v + cm.m[1][3];
66+
float bb = cm.m[2][0] * yb + cm.m[2][1] * u +
67+
cm.m[2][2] * v + cm.m[2][3];
6468
Vec2T gbbb = {static_cast<T>(fminf(fmaxf(gb, 0.0f), clampMax)),
6569
static_cast<T>(fminf(fmaxf(bb, 0.0f), clampMax))};
6670
*(reinterpret_cast<Vec2T*>(&rgbPlaneToWrite[4])) = gbbb;
@@ -95,7 +99,8 @@ __global__ void yuvToRgbKernel(
9599
int yPitchElements,
96100
int uvPitchElements,
97101
int rgbPitchElements,
98-
int bitShift) {
102+
int bitShift,
103+
const ColorMatrix cm) {
99104
// The kernel operates on 2x2 blocks, so it's called H / 2 * W / 2 times.
100105
// We have to multiply back by 2 to retrieve the output pixel coordinates x
101106
// and y.
@@ -120,33 +125,12 @@ __global__ void yuvToRgbKernel(
120125

121126
T* rgbPlaneToWrite = rgbOutput + y * rgbPitchElements + x * 3;
122127
writePairOfRGBPixels<T, Vec2T>(
123-
y1y2, u, v, rgbPlaneToWrite, bitShift);
128+
y1y2, u, v, rgbPlaneToWrite, bitShift, cm);
124129
rgbPlaneToWrite += rgbPitchElements; // go to next line
125130
writePairOfRGBPixels<T, Vec2T>(
126-
y3y4, u, v, rgbPlaneToWrite, bitShift);
131+
y3y4, u, v, rgbPlaneToWrite, bitShift, cm);
127132
}
128133

129-
namespace {
130-
131-
void maybeUploadColorMatrix(
132-
const float colorMatrix[3][4],
133-
bool colorMatrixChanged) {
134-
// We only send the color-matrix from CPU to GPU if it changed.
135-
// In practice it probably doesn't impact perf that much since decoding is
136-
// bottlenecked by NVDEC's frame mapping, not color-conversion. But it's a
137-
// simple optimization, so we do it anyway.
138-
if (colorMatrixChanged) {
139-
cudaMemcpyToSymbol(
140-
d_colorMatrix,
141-
colorMatrix,
142-
sizeof(float) * 12,
143-
0,
144-
cudaMemcpyHostToDevice);
145-
}
146-
}
147-
148-
} // namespace
149-
150134
void launchNV12ToRGBKernel(
151135
const uint8_t* yPlane,
152136
const uint8_t* uvPlane,
@@ -157,9 +141,9 @@ void launchNV12ToRGBKernel(
157141
int uvPitch,
158142
int rgbPitch,
159143
const float colorMatrix[3][4],
160-
bool colorMatrixChanged,
161144
cudaStream_t stream) {
162-
maybeUploadColorMatrix(colorMatrix, colorMatrixChanged);
145+
const auto& cm =
146+
*reinterpret_cast<const ColorMatrix*>(colorMatrix);
163147

164148
dim3 block(32, 2);
165149
dim3 grid(
@@ -175,7 +159,8 @@ void launchNV12ToRGBKernel(
175159
yPitch,
176160
uvPitch,
177161
rgbPitch,
178-
0); // bitShift = 0 for NV12
162+
0, // bitShift = 0 for NV12
163+
cm);
179164
}
180165

181166
void launchP016ToRGB16Kernel(
@@ -189,9 +174,9 @@ void launchP016ToRGB16Kernel(
189174
int rgbPitch,
190175
int bitDepth,
191176
const float colorMatrix[3][4],
192-
bool colorMatrixChanged,
193177
cudaStream_t stream) {
194-
maybeUploadColorMatrix(colorMatrix, colorMatrixChanged);
178+
const auto& cm =
179+
*reinterpret_cast<const ColorMatrix*>(colorMatrix);
195180

196181
int yPitchElements =
197182
yPitch / static_cast<int>(sizeof(uint16_t));
@@ -215,7 +200,8 @@ void launchP016ToRGB16Kernel(
215200
yPitchElements,
216201
uvPitchElements,
217202
rgbPitchElements,
218-
bitShift);
203+
bitShift,
204+
cm);
219205
}
220206

221207
// RGB -> NV12 kernel for encoding.
@@ -229,7 +215,8 @@ __global__ void rgbToNV12Kernel(
229215
int height,
230216
int rgbPitchElements,
231217
int yPitchElements,
232-
int uvPitchElements) {
218+
int uvPitchElements,
219+
const ColorMatrix cm) {
233220
int x = (blockIdx.x * blockDim.x + threadIdx.x) * 2;
234221
int y = (blockIdx.y * blockDim.y + threadIdx.y) * 2;
235222
if (x + 1 >= width || y + 1 >= height) {
@@ -247,18 +234,18 @@ __global__ void rgbToNV12Kernel(
247234
float r11 = row1[3], g11 = row1[4], b11 = row1[5];
248235

249236
// Compute Y for all 4 pixels
250-
float y00 = d_colorMatrix[0][0] * r00 +
251-
d_colorMatrix[0][1] * g00 + d_colorMatrix[0][2] * b00 +
252-
d_colorMatrix[0][3];
253-
float y01 = d_colorMatrix[0][0] * r01 +
254-
d_colorMatrix[0][1] * g01 + d_colorMatrix[0][2] * b01 +
255-
d_colorMatrix[0][3];
256-
float y10 = d_colorMatrix[0][0] * r10 +
257-
d_colorMatrix[0][1] * g10 + d_colorMatrix[0][2] * b10 +
258-
d_colorMatrix[0][3];
259-
float y11 = d_colorMatrix[0][0] * r11 +
260-
d_colorMatrix[0][1] * g11 + d_colorMatrix[0][2] * b11 +
261-
d_colorMatrix[0][3];
237+
float y00 = cm.m[0][0] * r00 +
238+
cm.m[0][1] * g00 + cm.m[0][2] * b00 +
239+
cm.m[0][3];
240+
float y01 = cm.m[0][0] * r01 +
241+
cm.m[0][1] * g01 + cm.m[0][2] * b01 +
242+
cm.m[0][3];
243+
float y10 = cm.m[0][0] * r10 +
244+
cm.m[0][1] * g10 + cm.m[0][2] * b10 +
245+
cm.m[0][3];
246+
float y11 = cm.m[0][0] * r11 +
247+
cm.m[0][1] * g11 + cm.m[0][2] * b11 +
248+
cm.m[0][3];
262249

263250
// Write Y plane
264251
uchar2 yRow0 = {
@@ -278,31 +265,31 @@ __global__ void rgbToNV12Kernel(
278265
&yPlane[(y + 1) * yPitchElements + x])) = yRow1;
279266

280267
// Compute U,V for all 4 pixels and average for 4:2:0 subsampling
281-
float u00 = d_colorMatrix[1][0] * r00 +
282-
d_colorMatrix[1][1] * g00 + d_colorMatrix[1][2] * b00 +
283-
d_colorMatrix[1][3];
284-
float u01 = d_colorMatrix[1][0] * r01 +
285-
d_colorMatrix[1][1] * g01 + d_colorMatrix[1][2] * b01 +
286-
d_colorMatrix[1][3];
287-
float u10 = d_colorMatrix[1][0] * r10 +
288-
d_colorMatrix[1][1] * g10 + d_colorMatrix[1][2] * b10 +
289-
d_colorMatrix[1][3];
290-
float u11 = d_colorMatrix[1][0] * r11 +
291-
d_colorMatrix[1][1] * g11 + d_colorMatrix[1][2] * b11 +
292-
d_colorMatrix[1][3];
293-
294-
float v00 = d_colorMatrix[2][0] * r00 +
295-
d_colorMatrix[2][1] * g00 + d_colorMatrix[2][2] * b00 +
296-
d_colorMatrix[2][3];
297-
float v01 = d_colorMatrix[2][0] * r01 +
298-
d_colorMatrix[2][1] * g01 + d_colorMatrix[2][2] * b01 +
299-
d_colorMatrix[2][3];
300-
float v10 = d_colorMatrix[2][0] * r10 +
301-
d_colorMatrix[2][1] * g10 + d_colorMatrix[2][2] * b10 +
302-
d_colorMatrix[2][3];
303-
float v11 = d_colorMatrix[2][0] * r11 +
304-
d_colorMatrix[2][1] * g11 + d_colorMatrix[2][2] * b11 +
305-
d_colorMatrix[2][3];
268+
float u00 = cm.m[1][0] * r00 +
269+
cm.m[1][1] * g00 + cm.m[1][2] * b00 +
270+
cm.m[1][3];
271+
float u01 = cm.m[1][0] * r01 +
272+
cm.m[1][1] * g01 + cm.m[1][2] * b01 +
273+
cm.m[1][3];
274+
float u10 = cm.m[1][0] * r10 +
275+
cm.m[1][1] * g10 + cm.m[1][2] * b10 +
276+
cm.m[1][3];
277+
float u11 = cm.m[1][0] * r11 +
278+
cm.m[1][1] * g11 + cm.m[1][2] * b11 +
279+
cm.m[1][3];
280+
281+
float v00 = cm.m[2][0] * r00 +
282+
cm.m[2][1] * g00 + cm.m[2][2] * b00 +
283+
cm.m[2][3];
284+
float v01 = cm.m[2][0] * r01 +
285+
cm.m[2][1] * g01 + cm.m[2][2] * b01 +
286+
cm.m[2][3];
287+
float v10 = cm.m[2][0] * r10 +
288+
cm.m[2][1] * g10 + cm.m[2][2] * b10 +
289+
cm.m[2][3];
290+
float v11 = cm.m[2][0] * r11 +
291+
cm.m[2][1] * g11 + cm.m[2][2] * b11 +
292+
cm.m[2][3];
306293

307294
float uAvg = (u00 + u01 + u10 + u11) * 0.25f;
308295
float vAvg = (v00 + v01 + v10 + v11) * 0.25f;
@@ -327,12 +314,8 @@ void launchRGBToNV12Kernel(
327314
int uvPitch,
328315
const float colorMatrix[3][4],
329316
cudaStream_t stream) {
330-
cudaMemcpyToSymbol(
331-
d_colorMatrix,
332-
colorMatrix,
333-
sizeof(float) * 12,
334-
0,
335-
cudaMemcpyHostToDevice);
317+
const auto& cm =
318+
*reinterpret_cast<const ColorMatrix*>(colorMatrix);
336319

337320
dim3 block(32, 2);
338321
dim3 grid(
@@ -347,7 +330,8 @@ void launchRGBToNV12Kernel(
347330
height,
348331
rgbPitch,
349332
yPitch,
350-
uvPitch);
333+
uvPitch,
334+
cm);
351335
}
352336

353337
} // namespace facebook::torchcodec

‎src/torchcodec/_core/color_conversion.h‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,9 +45,7 @@ void computeColorConversionMatrix(
4545
float outScale,
4646
float outMatrix[3][4]);
4747

48-
// Update the cached color matrix if any parameter changed.
49-
// Returns true if the matrix was recomputed (and needs re-uploading to GPU).
50-
bool maybeUpdateColorMatrix(
48+
void maybeUpdateColorMatrix(
5149
CachedColorMatrix& cachedColorMatrix,
5250
AVColorSpace colorspace,
5351
AVColorRange colorRange,
@@ -64,7 +62,6 @@ void launchNV12ToRGBKernel(
6462
int uvPitch,
6563
int rgbPitch,
6664
const float colorMatrix[3][4],
67-
bool colorMatrixChanged,
6865
cudaStream_t stream);
6966

7067
void launchP016ToRGB16Kernel(
@@ -78,7 +75,6 @@ void launchP016ToRGB16Kernel(
7875
int rgbPitch,
7976
int bitDepth,
8077
const float colorMatrix[3][4],
81-
bool colorMatrixChanged,
8278
cudaStream_t stream);
8379

8480
// Convert a YUV frame (NV12 or P016) on GPU to an interleaved RGB tensor.

0 commit comments

Comments
 (0)