1/*
2 * Copyright 2018 Google Inc.
3 *
4 * Use of this source code is governed by a BSD-style license that can be
5 * found in the LICENSE file.
6 */
7
8#include "src/gpu/effects/GrYUVtoRGBEffect.h"
9
10#include "include/core/SkYUVAInfo.h"
11#include "src/core/SkYUVMath.h"
12#include "src/gpu/GrTexture.h"
13#include "src/gpu/GrYUVATextureProxies.h"
14#include "src/gpu/effects/GrMatrixEffect.h"
15#include "src/gpu/effects/GrTextureEffect.h"
16#include "src/gpu/glsl/GrGLSLFragmentShaderBuilder.h"
17#include "src/gpu/glsl/GrGLSLProgramBuilder.h"
18#include "src/sksl/SkSLUtil.h"
19
20static void border_colors(const GrYUVATextureProxies& yuvaProxies, float planeBorders[4][4]) {
21    float m[20];
22    SkColorMatrix_RGB2YUV(yuvaProxies.yuvaInfo().yuvColorSpace(), m);
23    for (int i = 0; i < SkYUVAInfo::kYUVAChannelCount; ++i) {
24        auto [plane, channel] = yuvaProxies.yuvaLocations()[i];
25        if (plane == -1) {
26            return;
27        }
28        auto c = static_cast<int>(channel);
29        planeBorders[plane][c] = m[i*5 + 4];
30    }
31}
32
33std::unique_ptr<GrFragmentProcessor> GrYUVtoRGBEffect::Make(const GrYUVATextureProxies& yuvaProxies,
34                                                            GrSamplerState samplerState,
35                                                            const GrCaps& caps,
36                                                            const SkMatrix& localMatrix,
37                                                            const SkRect* subset,
38                                                            const SkRect* domain) {
39    SkASSERT(!subset || SkRect::Make(yuvaProxies.yuvaInfo().dimensions()).contains(*subset));
40
41    int numPlanes = yuvaProxies.yuvaInfo().numPlanes();
42    if (!yuvaProxies.isValid()) {
43        return nullptr;
44    }
45
46    bool usesBorder = samplerState.wrapModeX() == GrSamplerState::WrapMode::kClampToBorder ||
47                      samplerState.wrapModeY() == GrSamplerState::WrapMode::kClampToBorder;
48    float planeBorders[4][4] = {};
49    if (usesBorder) {
50        border_colors(yuvaProxies, planeBorders);
51    }
52
53    bool snap[2] = {false, false};
54    std::unique_ptr<GrFragmentProcessor> planeFPs[SkYUVAInfo::kMaxPlanes];
55    for (int i = 0; i < numPlanes; ++i) {
56        bool useSubset = SkToBool(subset);
57        GrSurfaceProxyView view = yuvaProxies.makeView(i);
58        SkMatrix planeMatrix = yuvaProxies.yuvaInfo().originMatrix();
59        // The returned matrix is a view matrix but we need a local matrix.
60        SkAssertResult(planeMatrix.invert(&planeMatrix));
61        SkRect planeSubset;
62        SkRect planeDomain;
63        bool makeLinearWithSnap = false;
64        auto [ssx, ssy] = yuvaProxies.yuvaInfo().planeSubsamplingFactors(i);
65        SkASSERT(ssx > 0 && ssx <= 4);
66        SkASSERT(ssy > 0 && ssy <= 2);
67        float scaleX = 1.f;
68        float scaleY = 1.f;
69        if (ssx > 1 || ssy > 1) {
70            scaleX = 1.f/ssx;
71            scaleY = 1.f/ssy;
72            // We would want to add a translation to this matrix to handle other sitings.
73            SkASSERT(yuvaProxies.yuvaInfo().sitingX() == SkYUVAInfo::Siting::kCentered);
74            SkASSERT(yuvaProxies.yuvaInfo().sitingY() == SkYUVAInfo::Siting::kCentered);
75            planeMatrix.postConcat(SkMatrix::Scale(scaleX, scaleY));
76            if (subset) {
77                planeSubset = {subset->fLeft  *scaleX,
78                               subset->fTop   *scaleY,
79                               subset->fRight *scaleX,
80                               subset->fBottom*scaleY};
81            } else {
82                planeSubset = SkRect::Make(view.dimensions());
83            }
84            if (domain) {
85                planeDomain = {domain->fLeft  *scaleX,
86                               domain->fTop   *scaleY,
87                               domain->fRight *scaleX,
88                               domain->fBottom*scaleY};
89            }
90            // If the image is not a multiple of the subsampling then the subsampled plane needs to
91            // be tiled at less than its full width/height. This only matters when the mode is not
92            // clamp.
93            if (samplerState.wrapModeX() != GrSamplerState::WrapMode::kClamp) {
94                int dx = (ssx*view.width() - yuvaProxies.yuvaInfo().width());
95                float maxRight = view.width() - dx*scaleX;
96                if (planeSubset.fRight > maxRight) {
97                    planeSubset.fRight = maxRight;
98                    useSubset = true;
99                }
100            }
101            if (samplerState.wrapModeY() != GrSamplerState::WrapMode::kClamp) {
102                int dy = (ssy*view.height() - yuvaProxies.yuvaInfo().height());
103                float maxBottom = view.height() - dy*scaleY;
104                if (planeSubset.fBottom > maxBottom) {
105                    planeSubset.fBottom = maxBottom;
106                    useSubset = true;
107                }
108            }
109            // This promotion of nearest to linear filtering for UV planes exists to mimic
110            // libjpeg[-turbo]'s do_fancy_upsampling option. We will filter the subsampled plane,
111            // however we want to filter at a fixed point for each logical image pixel to simulate
112            // nearest neighbor.
113            if (samplerState.filter() == GrSamplerState::Filter::kNearest) {
114                bool snapX = (ssx != 1),
115                     snapY = (ssy != 1);
116                makeLinearWithSnap = snapX || snapY;
117                snap[0] |= snapX;
118                snap[1] |= snapY;
119                if (domain) {
120                    // The outer YUVToRGB effect will ensure sampling happens at pixel centers
121                    // within this plane.
122                    planeDomain = {std::floor(planeDomain.fLeft)   + 0.5f,
123                                   std::floor(planeDomain.fTop)    + 0.5f,
124                                   std::floor(planeDomain.fRight)  + 0.5f,
125                                   std::floor(planeDomain.fBottom) + 0.5f};
126                }
127            }
128        } else {
129            if (subset) {
130                planeSubset = *subset;
131            }
132            if (domain) {
133                planeDomain = *domain;
134            }
135        }
136        if (useSubset) {
137            if (makeLinearWithSnap) {
138                // The plane is subsampled and we have an overall subset on the image. We're
139                // emulating do_fancy_upsampling using linear filtering but snapping look ups to the
140                // y-plane pixel centers. Consider a logical image pixel at the edge of the subset.
141                // When computing the logical pixel color value we should use a 50/50 blend of two
142                // values from the subsampled plane. Depending on where the subset edge falls in
143                // actual subsampled plane, one of those values may come from outside the subset.
144                // Hence, we use this custom inset factory which applies the wrap mode to
145                // planeSubset but allows linear filtering to read pixels from the plane that are
146                // just outside planeSubset.
147                SkRect* domainRect = domain ? &planeDomain : nullptr;
148                planeFPs[i] = GrTextureEffect::MakeCustomLinearFilterInset(std::move(view),
149                                                                           kUnknown_SkAlphaType,
150                                                                           planeMatrix,
151                                                                           samplerState.wrapModeX(),
152                                                                           samplerState.wrapModeY(),
153                                                                           planeSubset,
154                                                                           domainRect,
155                                                                           {scaleX/2.f, scaleY/2.f},
156                                                                           caps,
157                                                                           planeBorders[i]);
158            } else if (domain) {
159                planeFPs[i] = GrTextureEffect::MakeSubset(std::move(view),
160                                                          kUnknown_SkAlphaType,
161                                                          planeMatrix,
162                                                          samplerState,
163                                                          planeSubset,
164                                                          planeDomain,
165                                                          caps,
166                                                          planeBorders[i]);
167            } else {
168                planeFPs[i] = GrTextureEffect::MakeSubset(std::move(view),
169                                                          kUnknown_SkAlphaType,
170                                                          planeMatrix,
171                                                          samplerState,
172                                                          planeSubset,
173                                                          caps,
174                                                          planeBorders[i]);
175            }
176        } else {
177            GrSamplerState planeSampler = samplerState;
178            if (makeLinearWithSnap) {
179                planeSampler.setFilterMode(GrSamplerState::Filter::kLinear);
180            }
181            planeFPs[i] = GrTextureEffect::Make(std::move(view),
182                                                kUnknown_SkAlphaType,
183                                                planeMatrix,
184                                                planeSampler,
185                                                caps,
186                                                planeBorders[i]);
187        }
188    }
189    std::unique_ptr<GrFragmentProcessor> fp(
190            new GrYUVtoRGBEffect(planeFPs,
191                                 numPlanes,
192                                 yuvaProxies.yuvaLocations(),
193                                 snap,
194                                 yuvaProxies.yuvaInfo().yuvColorSpace()));
195    return GrMatrixEffect::Make(localMatrix, std::move(fp));
196}
197
198static SkAlphaType alpha_type(const SkYUVAInfo::YUVALocations locations) {
199    return locations[SkYUVAInfo::YUVAChannels::kA].fPlane >= 0 ? kPremul_SkAlphaType
200                                                               : kOpaque_SkAlphaType;
201}
202
203GrYUVtoRGBEffect::GrYUVtoRGBEffect(std::unique_ptr<GrFragmentProcessor> planeFPs[4],
204                                   int numPlanes,
205                                   const SkYUVAInfo::YUVALocations& locations,
206                                   const bool snap[2],
207                                   SkYUVColorSpace yuvColorSpace)
208        : GrFragmentProcessor(kGrYUVtoRGBEffect_ClassID,
209                              ModulateForClampedSamplerOptFlags(alpha_type(locations)))
210        , fLocations(locations)
211        , fYUVColorSpace(yuvColorSpace) {
212    std::copy_n(snap, 2, fSnap);
213
214    if (fSnap[0] || fSnap[1]) {
215        // Need this so that we can access coords in SKSL to perform snapping.
216        this->setUsesSampleCoordsDirectly();
217        for (int i = 0; i < numPlanes; ++i) {
218            this->registerChild(std::move(planeFPs[i]), SkSL::SampleUsage::Explicit());
219        }
220    } else {
221        for (int i = 0; i < numPlanes; ++i) {
222            this->registerChild(std::move(planeFPs[i]));
223        }
224    }
225}
226
227#if GR_TEST_UTILS
228SkString GrYUVtoRGBEffect::onDumpInfo() const {
229    SkString str("(");
230    for (int i = 0; i < SkYUVAInfo::kYUVAChannelCount; ++i) {
231        str.appendf("Locations[%d]=%d %d, ",
232                    i, fLocations[i].fPlane, static_cast<int>(fLocations[i].fChannel));
233    }
234    str.appendf("YUVColorSpace=%d, snap=(%d, %d))",
235                static_cast<int>(fYUVColorSpace), fSnap[0], fSnap[1]);
236    return str;
237}
238#endif
239
240std::unique_ptr<GrFragmentProcessor::ProgramImpl> GrYUVtoRGBEffect::onMakeProgramImpl() const {
241    class Impl : public ProgramImpl {
242    public:
243        void emitCode(EmitArgs& args) override {
244            GrGLSLFPFragmentBuilder* fragBuilder = args.fFragBuilder;
245            const GrYUVtoRGBEffect& yuvEffect = args.fFp.cast<GrYUVtoRGBEffect>();
246
247            int numPlanes = yuvEffect.numChildProcessors();
248
249            const char* sampleCoords = "";
250            if (yuvEffect.fSnap[0] || yuvEffect.fSnap[1]) {
251                fragBuilder->codeAppendf("float2 snappedCoords = %s;", args.fSampleCoord);
252                if (yuvEffect.fSnap[0]) {
253                    fragBuilder->codeAppend("snappedCoords.x = floor(snappedCoords.x) + 0.5;");
254                }
255                if (yuvEffect.fSnap[1]) {
256                    fragBuilder->codeAppend("snappedCoords.y = floor(snappedCoords.y) + 0.5;");
257                }
258                sampleCoords = "snappedCoords";
259            }
260
261            fragBuilder->codeAppendf("half4 color;");
262            const bool hasAlpha = yuvEffect.fLocations[SkYUVAInfo::YUVAChannels::kA].fPlane >= 0;
263
264            for (int planeIdx = 0; planeIdx < numPlanes; ++planeIdx) {
265                std::string colorChannel;
266                std::string planeChannel;
267                for (int locIdx = 0; locIdx < (hasAlpha ? 4 : 3); ++locIdx) {
268                    auto [yuvPlane, yuvChannel] = yuvEffect.fLocations[locIdx];
269                    if (yuvPlane == planeIdx) {
270                        colorChannel.push_back("rgba"[locIdx]);
271                        planeChannel.push_back("rgba"[static_cast<int>(yuvChannel)]);
272                    }
273                }
274
275                SkASSERT(colorChannel.size() == planeChannel.size());
276                if (!colorChannel.empty()) {
277                    fragBuilder->codeAppendf(
278                            "color.%s = (%s).%s;",
279                            colorChannel.c_str(),
280                            this->invokeChild(planeIdx, args, sampleCoords).c_str(),
281                            planeChannel.c_str());
282                }
283            }
284
285            if (!hasAlpha) {
286                fragBuilder->codeAppendf("color.a = 1;");
287            }
288
289            if (kIdentity_SkYUVColorSpace != yuvEffect.fYUVColorSpace) {
290                fColorSpaceMatrixVar = args.fUniformHandler->addUniform(&yuvEffect,
291                        kFragment_GrShaderFlag, kHalf3x3_GrSLType, "colorSpaceMatrix");
292                fColorSpaceTranslateVar = args.fUniformHandler->addUniform(&yuvEffect,
293                        kFragment_GrShaderFlag, kHalf3_GrSLType, "colorSpaceTranslate");
294                fragBuilder->codeAppendf(
295                        "color.rgb = saturate(color.rgb * %s + %s);",
296                        args.fUniformHandler->getUniformCStr(fColorSpaceMatrixVar),
297                        args.fUniformHandler->getUniformCStr(fColorSpaceTranslateVar));
298            }
299            if (hasAlpha) {
300                // premultiply alpha
301                fragBuilder->codeAppendf("color.rgb *= color.a;");
302            }
303            fragBuilder->codeAppendf("return color;");
304        }
305
306    private:
307        void onSetData(const GrGLSLProgramDataManager& pdman,
308                       const GrFragmentProcessor& proc) override {
309            const GrYUVtoRGBEffect& yuvEffect = proc.cast<GrYUVtoRGBEffect>();
310
311            if (yuvEffect.fYUVColorSpace != kIdentity_SkYUVColorSpace) {
312                SkASSERT(fColorSpaceMatrixVar.isValid());
313                float yuvM[20];
314                SkColorMatrix_YUV2RGB(yuvEffect.fYUVColorSpace, yuvM);
315                // We drop the fourth column entirely since the transformation
316                // should not depend on alpha. The fifth column is sent as a separate
317                // vector. The fourth row is also dropped entirely because alpha should
318                // never be modified.
319                SkASSERT(yuvM[3] == 0 && yuvM[8] == 0 && yuvM[13] == 0 && yuvM[18] == 1);
320                SkASSERT(yuvM[15] == 0 && yuvM[16] == 0 && yuvM[17] == 0 && yuvM[19] == 0);
321                float mtx[9] = {
322                    yuvM[ 0], yuvM[ 1], yuvM[ 2],
323                    yuvM[ 5], yuvM[ 6], yuvM[ 7],
324                    yuvM[10], yuvM[11], yuvM[12],
325                };
326                float v[3] = {yuvM[4], yuvM[9], yuvM[14]};
327                pdman.setMatrix3f(fColorSpaceMatrixVar, mtx);
328                pdman.set3fv(fColorSpaceTranslateVar, 1, v);
329            }
330        }
331
332        UniformHandle fColorSpaceMatrixVar;
333        UniformHandle fColorSpaceTranslateVar;
334    };
335
336    return std::make_unique<Impl>();
337}
338
339SkString GrYUVtoRGBEffect::getShaderDfxInfo() const
340{
341    SkString format;
342    format.printf("ShaderDfx_GrYUVtoRGBEffect");
343    for (auto [plane, channel] : fLocations) {
344        format.appendf("_%d_%d", plane, channel);
345    }
346    format.appendf("_%d_%d_%d", fYUVColorSpace, fSnap[0], fSnap[1]);
347    return format;
348}
349
350void GrYUVtoRGBEffect::onAddToKey(const GrShaderCaps& caps, GrProcessorKeyBuilder* b) const {
351    uint32_t packed = 0;
352    int i = 0;
353    for (auto [plane, channel] : fLocations) {
354        if (plane < 0) {
355            continue;
356        }
357
358        uint8_t chann = static_cast<int>(channel);
359
360        SkASSERT(plane < 4 && chann < 4);
361
362        packed |= (plane | (chann << 2)) << (i++ * 4);
363    }
364    if (fYUVColorSpace == kIdentity_SkYUVColorSpace) {
365        packed |= 1 << 16;
366    }
367    if (fSnap[0]) {
368        packed |= 1 << 17;
369    }
370    if (fSnap[1]) {
371        packed |= 1 << 18;
372    }
373    b->add32(packed);
374}
375
376bool GrYUVtoRGBEffect::onIsEqual(const GrFragmentProcessor& other) const {
377    const GrYUVtoRGBEffect& that = other.cast<GrYUVtoRGBEffect>();
378
379    return fLocations == that.fLocations            &&
380           std::equal(fSnap, fSnap + 2, that.fSnap) &&
381           fYUVColorSpace == that.fYUVColorSpace;
382}
383
384GrYUVtoRGBEffect::GrYUVtoRGBEffect(const GrYUVtoRGBEffect& src)
385        : GrFragmentProcessor(src)
386        , fLocations((src.fLocations))
387        , fYUVColorSpace(src.fYUVColorSpace) {
388    std::copy_n(src.fSnap, 2, fSnap);
389}
390
391std::unique_ptr<GrFragmentProcessor> GrYUVtoRGBEffect::clone() const {
392    return std::unique_ptr<GrFragmentProcessor>(new GrYUVtoRGBEffect(*this));
393}
394