added support for SkSL unpremul function

Change-Id: I970f1ad0dd0859448c874498fe02342f8abc3aa3
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/242897
Reviewed-by: Brian Salomon <bsalomon@google.com>
Commit-Queue: Ethan Nicholas <ethannicholas@google.com>

src/effects/SkHighContrastFilter.cpp

@@ -245,11 +245,7 @@ void GLHighContrastFilterEffect::emitCode(EmitArgs& args) {
 
     GrGLSLFPFragmentBuilder* fragBuilder = args.fFragBuilder;
 
-    fragBuilder->codeAppendf("half4 color = %s;", args.fInputColor);
+    fragBuilder->codeAppendf("half4 color = unpremul(%s);", args.fInputColor);
-
-    // Unpremultiply. The max() is to guard against 0 / 0.
-    fragBuilder->codeAppendf("half nonZeroAlpha = max(color.a, 0.0001);");
-    fragBuilder->codeAppendf("color = half4(color.rgb / nonZeroAlpha, nonZeroAlpha);");
 
     if (hcfe.linearize()) {
         fragBuilder->codeAppend("color.rgb = color.rgb * color.rgb;");

src/effects/SkTableColorFilter.cpp

@@ -301,8 +301,7 @@ void GLColorTableEffect::emitCode(EmitArgs& args) {
                                  kMaxValue, kMaxValue, kMaxValue, kMaxValue);
 
     } else {
-        fragBuilder->codeAppendf("\t\thalf nonZeroAlpha = max(%s.a, .0001);\n", args.fInputColor);
+        fragBuilder->codeAppendf("\t\thalf4 coord = unpremul(%s);\n",
-        fragBuilder->codeAppendf("\t\thalf4 coord = half4(%s.rgb / nonZeroAlpha, nonZeroAlpha);\n",
                                  args.fInputColor);
         fragBuilder->codeAppendf("\t\tcoord = coord * %f + half4(%f, %f, %f, %f);\n",
                                  kColorScaleFactor,

src/gpu/effects/GrColorMatrixFragmentProcessor.fp

@@ -18,10 +18,7 @@ layout(key) in bool premulOutput;
 void main() {
     half4 inputColor = sk_InColor;
     @if (unpremulInput) {
-        // The max() is to guard against 0 / 0 during unpremul when the incoming color is
+        inputColor = unpremul(inputColor);
-        // transparent black.
-        half nonZeroAlpha = max(inputColor.a, 0.0001);
-        inputColor = half4(inputColor.rgb / nonZeroAlpha, nonZeroAlpha);
     }
     sk_OutColor = m * inputColor + v;
     @if (clampRGBOutput) {

src/gpu/effects/GrSRGBEffect.cpp

@@ -47,8 +47,7 @@ public:
         // calculations to be performed with sufficient precision.
         fragBuilder->codeAppendf("float4 color = %s;", args.fInputColor);
         if (srgbe.alpha() == GrSRGBEffect::Alpha::kPremul) {
-            fragBuilder->codeAppendf("float nonZeroAlpha = max(color.a, 0.0001);");
+            fragBuilder->codeAppendf("color = unpremul(color);");
-            fragBuilder->codeAppendf("color = float4(color.rgb / nonZeroAlpha, color.a);");
         }
         fragBuilder->codeAppendf("color = float4(%s(half(color.r)), %s(half(color.g)), "
                                  "%s(half(color.b)), color.a);",

src/gpu/effects/generated/GrColorMatrixFragmentProcessor.cpp

@@ -37,11 +37,12 @@ public:
         mVar = args.fUniformHandler->addUniform(kFragment_GrShaderFlag, kHalf4x4_GrSLType, "m");
         vVar = args.fUniformHandler->addUniform(kFragment_GrShaderFlag, kHalf4_GrSLType, "v");
         fragBuilder->codeAppendf(
-                "half4 inputColor = %s;\n@if (%s) {\n    half nonZeroAlpha = max(inputColor.w, "
+                "    half4 unpremul0;     half unpremulNonZeroAlpha1;half4 inputColor = %s;\n@if "
-                "9.9999997473787516e-05);\n    inputColor = half4(inputColor.xyz / nonZeroAlpha, "
+                "(%s) {\n    inputColor = (unpremul0 = inputColor, unpremulNonZeroAlpha1 = "
-                "nonZeroAlpha);\n}\n%s = %s * inputColor + %s;\n@if (%s) {\n    %s = clamp(%s, "
+                "max(unpremul0.a, 0.0001), half4(unpremul0.rgb / unpremulNonZeroAlpha1, "
-                "0.0, 1.0);\n} else {\n    %s.w = clamp(%s.w, 0.0, 1.0);\n}\n@if (%s) {\n    "
+                "unpremulNonZeroAlpha1));\n}\n%s = %s * inputColor + %s;\n@if (%s) {\n    %s = "
-                "%s.xyz *= %s.w;\n}\n",
+                "clamp(%s, 0.0, 1.0);\n} else {\n    %s.w = clamp(%s.w, 0.0, 1.0);\n}\n@if (%s) "
+                "{\n    %s.xyz *= %s.w;\n}\n",
                 args.fInputColor, (_outer.unpremulInput ? "true" : "false"), args.fOutputColor,
                 args.fUniformHandler->getUniformCStr(mVar),
                 args.fUniformHandler->getUniformCStr(vVar),

src/gpu/glsl/GrGLSLShaderBuilder.cpp

@@ -159,8 +159,7 @@ void GrGLSLShaderBuilder::appendColorGamutXform(SkString* out,
         const GrShaderVar gColorXformArgs[] = { GrShaderVar("color", kHalf4_GrSLType) };
         SkString body;
         if (colorXformHelper->applyUnpremul()) {
-            body.append("half nonZeroAlpha = max(color.a, 0.0001);");
+            body.append("color = unpremul(color);");
-            body.append("color = half4(color.rgb / nonZeroAlpha, nonZeroAlpha);");
         }
         if (colorXformHelper->applySrcTF()) {
             body.appendf("color.r = %s(color.r);", srcTFFuncName.c_str());

src/sksl/SkSLGLSLCodeGenerator.cpp

@@ -472,6 +472,7 @@ void GLSLCodeGenerator::writeFunctionCall(const FunctionCall& c) {
         (*fFunctionClasses)["saturate"]    = FunctionClass::kSaturate;
         (*fFunctionClasses)["sample"]      = FunctionClass::kTexture;
         (*fFunctionClasses)["transpose"]   = FunctionClass::kTranspose;
+        (*fFunctionClasses)["unpremul"]    = FunctionClass::kUnpremul;
     }
 #ifndef SKSL_STANDALONE
     );
@@ -691,6 +692,24 @@ void GLSLCodeGenerator::writeFunctionCall(const FunctionCall& c) {
                     return;
                 }
                 break;
+            case FunctionClass::kUnpremul:
+                String tmpVar1 = "unpremul" + to_string(fVarCount++);
+                this->fFunctionHeader += String("    ") +
+                                         this->getTypePrecision(c.fArguments[0]->fType) +
+                                         this->getTypeName(c.fArguments[0]->fType) + " " + tmpVar1 +
+                                         ";";
+                String tmpVar2 = "unpremulNonZeroAlpha" + to_string(fVarCount++);
+                this->fFunctionHeader += String("    ") +
+                                         this->getTypePrecision(c.fArguments[0]->fType) + " " +
+                                         this->getTypeName(c.fArguments[0]->fType.componentType()) +
+                                         " " + tmpVar2 + ";";
+                this->write("(" + tmpVar1 + " = ");
+                this->writeExpression(*c.fArguments[0], kSequence_Precedence);
+                this->write(", " + tmpVar2 + " = max(" + tmpVar1 + ".a, " +
+                            to_string(SKSL_UNPREMUL_MIN) + "), " +
+                            this->getTypeName(*fContext.fHalf4_Type) + "(" + tmpVar1 + ".rgb / " +
+                            tmpVar2 + ", " + tmpVar2 + "))");
+                return;
         }
     }
     if (!nameWritten) {

src/sksl/SkSLGLSLCodeGenerator.h

@@ -254,7 +254,8 @@ protected:
         kPow,
         kSaturate,
         kTexture,
-        kTranspose
+        kTranspose,
+        kUnpremul
     };
     static std::unordered_map<StringFragment, FunctionClass>* fFunctionClasses;
 

src/sksl/SkSLMetalCodeGenerator.cpp

@@ -32,6 +32,7 @@ void MetalCodeGenerator::setupIntrinsics() {
     fIntrinsicMap[String("lessThanEqual")]      = METAL(LessThanEqual);
     fIntrinsicMap[String("greaterThan")]        = METAL(GreaterThan);
     fIntrinsicMap[String("greaterThanEqual")]   = METAL(GreaterThanEqual);
+    fIntrinsicMap[String("unpremul")]           = SPECIAL(Unpremul);
 }
 
 void MetalCodeGenerator::write(const char* s) {
@@ -69,50 +70,51 @@ void MetalCodeGenerator::writeExtension(const Extension& ext) {
     this->writeLine("#extension " + ext.fName + " : enable");
 }
 
-void MetalCodeGenerator::writeType(const Type& type) {
+String MetalCodeGenerator::getTypeName(const Type& type) {
     switch (type.kind()) {
         case Type::kStruct_Kind:
-            for (const Type* search : fWrittenStructs) {
+            return type.name();
-                if (*search == type) {
-                    // already written
-                    this->write(type.name());
-                    return;
-                }
-            }
-            fWrittenStructs.push_back(&type);
-            this->writeLine("struct " + type.name() + " {");
-            fIndentation++;
-            this->writeFields(type.fields(), type.fOffset);
-            fIndentation--;
-            this->write("}");
-            break;
         case Type::kVector_Kind:
-            this->writeType(type.componentType());
+            return this->getTypeName(type.componentType()) + to_string(type.columns());
-            this->write(to_string(type.columns()));
-            break;
         case Type::kMatrix_Kind:
-            this->writeType(type.componentType());
+            return this->getTypeName(type.componentType()) + to_string(type.columns()) + "x" +
-            this->write(to_string(type.columns()));
+                   to_string(type.rows());
-            this->write("x");
-            this->write(to_string(type.rows()));
-            break;
         case Type::kSampler_Kind:
-            this->write("texture2d<float> "); // FIXME - support other texture types;
+            return "texture2d<float>"; // FIXME - support other texture types;
-            break;
         default:
             if (type == *fContext.fHalf_Type) {
                 // FIXME - Currently only supporting floats in MSL to avoid type coercion issues.
-                this->write(fContext.fFloat_Type->name());
+                return fContext.fFloat_Type->name();
             } else if (type == *fContext.fByte_Type) {
-                this->write("char");
+                return "char";
             } else if (type == *fContext.fUByte_Type) {
-                this->write("uchar");
+                return "uchar";
             } else {
-                this->write(type.name());
+                return type.name();
             }
     }
 }
 
+void MetalCodeGenerator::writeType(const Type& type) {
+    if (type.kind() == Type::kStruct_Kind) {
+        for (const Type* search : fWrittenStructs) {
+            if (*search == type) {
+                // already written
+                this->write(this->getTypeName(type));
+                return;
+            }
+        }
+        fWrittenStructs.push_back(&type);
+        this->writeLine("struct " + type.name() + " {");
+        fIndentation++;
+        this->writeFields(type.fields(), type.fOffset);
+        fIndentation--;
+        this->write("}");
+    } else {
+        this->write(this->getTypeName(type));
+    }
+}
+
 void MetalCodeGenerator::writeExpression(const Expression& expr, Precedence parentPrecedence) {
     switch (expr.fKind) {
         case Expression::kBinary_Kind:
@@ -369,6 +371,23 @@ void MetalCodeGenerator::writeSpecialIntrinsic(const FunctionCall & c, SpecialIn
             this->writeExpression(*c.fArguments[1], kSequence_Precedence);
             this->write(")))");
             break;
+        case kUnpremul_SpecialIntrinsic: {
+            String tmpVar1 = "unpremul" + to_string(fVarCount++);
+            this->fFunctionHeader += String("    ") +
+                                     this->getTypeName(c.fArguments[0]->fType) + " " + tmpVar1 +
+                                     ";";
+            String tmpVar2 = "unpremulNonZeroAlpha" + to_string(fVarCount++);
+            this->fFunctionHeader += String("    ") +
+                                     this->getTypeName(c.fArguments[0]->fType.componentType()) +
+                                     " " + tmpVar2 + ";";
+            this->write("(" + tmpVar1 + " = ");
+            this->writeExpression(*c.fArguments[0], kSequence_Precedence);
+            this->write(", " + tmpVar2 + " = max(" + tmpVar1 + ".a, " +
+                        to_string(SKSL_UNPREMUL_MIN) + "), " +
+                        this->getTypeName(*fContext.fHalf4_Type) + "(" + tmpVar1 +
+                        ".rgb / " + tmpVar2 + ", " + tmpVar2 + "))");
+            return;
+        }
         default:
             ABORT("unsupported special intrinsic kind");
     }

src/sksl/SkSLMetalCodeGenerator.h

@@ -78,7 +78,7 @@ public:
     };
 
     MetalCodeGenerator(const Context* context, const Program* program, ErrorReporter* errors,
-                      OutputStream* out)
+                       OutputStream* out)
     : INHERITED(program, errors, out)
     , fReservedWords({"atan2", "rsqrt", "dfdx", "dfdy", "vertex", "fragment"})
     , fLineEnding("\n")
@@ -105,6 +105,7 @@ protected:
     enum SpecialIntrinsic {
         kTexture_SpecialIntrinsic,
         kMod_SpecialIntrinsic,
+        kUnpremul_SpecialIntrinsic,
     };
 
     enum MetalIntrinsic {
@@ -149,6 +150,8 @@ protected:
 
     void writePrecisionModifier();
 
+    String getTypeName(const Type& type);
+
     void writeType(const Type& type);
 
     void writeExtension(const Extension& ext);

src/sksl/SkSLSPIRVCodeGenerator.cpp

@@ -141,6 +141,7 @@ void SPIRVCodeGenerator::setupIntrinsics() {
                                                                 SpvOpUndef);
     fIntrinsicMap[String("EmitVertex")]       = ALL_SPIRV(EmitVertex);
     fIntrinsicMap[String("EndPrimitive")]     = ALL_SPIRV(EndPrimitive);
+    fIntrinsicMap[String("unpremul")]         = SPECIAL(Unpremul);
 // interpolateAt* not yet supported...
 }
 
@@ -1002,6 +1003,26 @@ SpvId SPIRVCodeGenerator::writeSpecialIntrinsic(const FunctionCall& c, SpecialIn
                                                GLSLstd450UClamp, spvArgs, out);
             break;
         }
+        case kUnpremul_SpecialIntrinsic: {
+            SpvId color = this->writeExpression(*c.fArguments[0], out);
+            SpvId a = this->writeSwizzle(*fContext.fHalf_Type, c.fArguments[0]->fType, color, { 3 },
+                                         out);
+            FloatLiteral min(fContext, -1, SKSL_UNPREMUL_MIN);
+            SpvId minId = this->writeFloatLiteral(min);
+            SpvId nonZeroAlpha = this->nextId();
+            this->writeGLSLExtendedInstruction(*fContext.fHalf_Type, nonZeroAlpha, GLSLstd450FMax,
+                                               SpvOpUndef, SpvOpUndef, { a, minId }, out);
+            SpvId rgb = this->writeSwizzle(*fContext.fHalf3_Type, *fContext.fHalf4_Type, color,
+                                           { 0, 1, 2 }, out);
+            SpvId scaled = this->writeBinaryExpression(*fContext.fHalf3_Type, rgb, Token::SLASH,
+                                                       *fContext.fFloat_Type, nonZeroAlpha,
+                                                       *fContext.fHalf3_Type, out);
+            this->writeOpCode(SpvOpCompositeConstruct, 5, out);
+            this->writeWord(this->getType(c.fType), out);
+            this->writeWord(result, out);
+            this->writeWord(scaled, out);
+            this->writeWord(nonZeroAlpha, out);
+        }
     }
     return result;
 }
@@ -1912,19 +1933,24 @@ SpvId SPIRVCodeGenerator::writeFieldAccess(const FieldAccess& f, OutputStream& o
 }
 
 SpvId SPIRVCodeGenerator::writeSwizzle(const Swizzle& swizzle, OutputStream& out) {
-    SpvId base = this->writeExpression(*swizzle.fBase, out);
+    return this->writeSwizzle(swizzle.fType, swizzle.fBase->fType,
+                              this->writeExpression(*swizzle.fBase, out), swizzle.fComponents, out);
+}
+
+SpvId SPIRVCodeGenerator::writeSwizzle(const Type& type, const Type& baseType, SpvId base,
+                                       const std::vector<int> components, OutputStream& out) {
     SpvId result = this->nextId();
-    size_t count = swizzle.fComponents.size();
+    size_t count = components.size();
     if (count == 1) {
-        this->writeInstruction(SpvOpCompositeExtract, this->getType(swizzle.fType), result, base,
+        this->writeInstruction(SpvOpCompositeExtract, this->getType(type), result, base,
-                               swizzle.fComponents[0], out);
+                               components[0], out);
     } else {
         this->writeOpCode(SpvOpVectorShuffle, 5 + (int32_t) count, out);
-        this->writeWord(this->getType(swizzle.fType), out);
+        this->writeWord(this->getType(type), out);
         this->writeWord(result, out);
         this->writeWord(base, out);
         SpvId other = base;
-        for (int c : swizzle.fComponents) {
+        for (int c : components) {
             if (c < 0) {
                 if (!fConstantZeroOneVector) {
                     FloatLiteral zero(fContext, -1, 0);
@@ -1944,11 +1970,11 @@ SpvId SPIRVCodeGenerator::writeSwizzle(const Swizzle& swizzle, OutputStream& out
             }
         }
         this->writeWord(other, out);
-        for (int component : swizzle.fComponents) {
+        for (int component : components) {
             if (component == SKSL_SWIZZLE_0) {
-                this->writeWord(swizzle.fBase->fType.columns(), out);
+                this->writeWord(baseType.columns(), out);
             } else if (component == SKSL_SWIZZLE_1) {
-                this->writeWord(swizzle.fBase->fType.columns() + 1, out);
+                this->writeWord(baseType.columns() + 1, out);
             } else {
                 this->writeWord(component, out);
             }
@@ -2068,6 +2094,11 @@ std::unique_ptr<Expression> create_literal_1(const Context& context, const Type&
 SpvId SPIRVCodeGenerator::writeBinaryExpression(const Type& leftType, SpvId lhs, Token::Kind op,
                                                 const Type& rightType, SpvId rhs,
                                                 const Type& resultType, OutputStream& out) {
+    // it's important to handle comma early, so we don't end up vectorizing the operands
+    if (op == Token::COMMA) {
+        return rhs;
+    }
+
     Type tmp("<invalid>");
     // overall type we are operating on: float2, int, uint4...
     const Type* operandType;
@@ -2260,8 +2291,6 @@ SpvId SPIRVCodeGenerator::writeBinaryExpression(const Type& leftType, SpvId lhs,
         case Token::BITWISEXOR:
             return this->writeBinaryOperation(resultType, *operandType, lhs, rhs, SpvOpUndef,
                                               SpvOpBitwiseXor, SpvOpBitwiseXor, SpvOpUndef, out);
-        case Token::COMMA:
-            return rhs;
         default:
             SkASSERT(false);
             return -1;

src/sksl/SkSLSPIRVCodeGenerator.h

@@ -137,6 +137,7 @@ private:
         kSampledImage_SpecialIntrinsic,
         kSubpassLoad_SpecialIntrinsic,
         kTexture_SpecialIntrinsic,
+        kUnpremul_SpecialIntrinsic,
     };
 
     enum class Precision {
@@ -253,6 +254,9 @@ private:
 
     SpvId writeSwizzle(const Swizzle& swizzle, OutputStream& out);
 
+    SpvId writeSwizzle(const Type& type, const Type& basetype, SpvId base,
+                       const std::vector<int> components, OutputStream& out);
+
     /**
      * Folds the potentially-vector result of a logical operation down to a single bool. If
      * operandType is a vector type, assumes that the intermediate result in id is a bvec of the

src/sksl/SkSLUtil.h

@@ -27,6 +27,8 @@ class GrShaderCaps;
 
 namespace SkSL {
 
+static constexpr float SKSL_UNPREMUL_MIN = 0.0001f;
+
 class OutputStream;
 class StringStream;
 

src/sksl/sksl_gpu.inc

@@ -102,6 +102,8 @@ $genType smoothstep($genType edge0, $genType edge1, $genType x);
 $genType smoothstep(float edge0, float edge1, $genType x);
 $genHType smoothstep($genHType edge0, $genHType edge1, $genHType x);
 $genHType smoothstep(half edge0, half edge1, $genHType x);
+half4 unpremul(half4 c);
+float4 unpremul(float4 c);
 $genBType isnan($genType x);
 $genBType isnan($genDType x);
 $genBType isinf($genType x);