void ShaderInfo::generateFragmentSkSL()

in src/gpu/graphite/ShaderInfo.cpp [829:1228]


void ShaderInfo::generateFragmentSkSL(const Caps* caps,
                                      const ShaderCodeDictionary* dict,
                                      const RenderStep* step,
                                      UniquePaintParamsID paintID,
                                      bool useStorageBuffers,
                                      Swizzle writeSwizzle,
                                      skia_private::TArray<SamplerDesc>* outDescs) {
    PaintParamsKey key = dict->lookup(paintID);
    SkASSERT(key.isValid());  // invalid keys should have been caught by invalid paint ID earlier

    std::string label = key.toString(dict).c_str();

    // Two varyings are reserved for 1) the SSBO indices and 2) local coordinates.
    constexpr int kFixedVaryings = 2;
    const int availableVaryings = caps->maxVaryings() - kFixedVaryings - step->varyings().size();
    fRootNodes = key.getRootNodes(caps, dict, &fShaderNodeAlloc, availableVaryings);

    // TODO(b/366220690): aggregateSnippetData() goes away entirely once the VulkanGraphicsPipeline
    // is updated to use the extracted SamplerDescs directly.
    for (const ShaderNode* root : fRootNodes) {
        this->aggregateSnippetData(root);
    }

#if defined(SK_DEBUG)
    // Validate the root node structure of the key.
    SkASSERT(fRootNodes.size() == 2 || fRootNodes.size() == 3);
    // First node produces the source color (all snippets return a half4), so we just require that
    // its signature takes no extra args or just local coords.
    const ShaderSnippet* srcSnippet = dict->getEntry(fRootNodes[0]->codeSnippetId());
    // TODO(b/349997190): Once SkEmptyShader doesn't use the passthrough snippet, we can assert
    // that srcSnippet->needsPriorStageOutput() is false.
    SkASSERT(!srcSnippet->needsBlenderDstColor());
    // Second node is the final blender, so it must take both the src color and dst color, and not
    // any local coordinate.
    const ShaderSnippet* blendSnippet = dict->getEntry(fRootNodes[1]->codeSnippetId());
    SkASSERT(blendSnippet->needsPriorStageOutput() && blendSnippet->needsBlenderDstColor());
    SkASSERT(!blendSnippet->needsLocalCoords());

    const ShaderSnippet* clipSnippet =
            fRootNodes.size() > 2 ? dict->getEntry(fRootNodes[2]->codeSnippetId()) : nullptr;
    SkASSERT(!clipSnippet ||
             (!clipSnippet->needsPriorStageOutput() && !clipSnippet->needsBlenderDstColor()));
#endif

    // The RenderStep should be performing shading since otherwise there's no need to generate a
    // fragment shader program at all.
    SkASSERT(step->performsShading());

    // Check for unexpected corruption / illegal instructions occurring in the wild.
    SkASSERTF_RELEASE(fRootNodes.size() == 2 || fRootNodes.size() == 3,
                      "root node size = %zu, label = %s",
                      fRootNodes.size(),
                      label.c_str());

    // Extract the root nodes for clarity
    const ShaderNode* const srcColorRoot = fRootNodes[0];
    const ShaderNode* const finalBlendRoot = fRootNodes[1];
    const int32_t finalBlendRootSnippetId = finalBlendRoot->codeSnippetId();
    const ShaderNode* const clipRoot = fRootNodes.size() > 2 ? fRootNodes[2] : nullptr;

    // Determine the algorithm for final blending: direct HW blending, coverage-modified HW
    // blending (w/ or w/o dual-source blending) or via dst-read requirement.
    Coverage finalCoverage = step->coverage();
    if (finalCoverage == Coverage::kNone && SkToBool(clipRoot)) {
        finalCoverage = Coverage::kSingleChannel;
    }

    // Initialize the final blend mode to the final snippet's blend mode. It may be changed based
    // upon whether or not we can use hardware blending.
    std::optional<SkBlendMode> finalBlendMode;
    if (finalBlendRootSnippetId < kBuiltInCodeSnippetIDCount &&
        finalBlendRootSnippetId >= kFixedBlendIDOffset) {
        finalBlendMode = static_cast<SkBlendMode>(finalBlendRootSnippetId - kFixedBlendIDOffset);
    }
    const bool useHardwareBlending = CanUseHardwareBlending(caps, finalBlendMode, finalCoverage);
    if (useHardwareBlending) {
        // If we can use hardware blending, update the dstReadStrategy to be kNoneRequired to ensure
        // that ShaderInfo properly informs PipelineInfo of the pipeline's dst read requirement.
        fDstReadStrategy = DstReadStrategy::kNoneRequired;
    } else {
        // If we cannot use hardware blending, then we must perform a dst read within the shader.
        // Therefore we should assert that a valid strategy to do so was passed in. Later operations
        // also expect the blend mode to be kSrc, so update that here.
        SkASSERT(fDstReadStrategy != DstReadStrategy::kNoneRequired);
        finalBlendMode = SkBlendMode::kSrc;
    }

    const bool hasStepUniforms = step->numUniforms() > 0 && step->usesUniformsInFragmentSkSL();
    const bool useStepStorageBuffer = useStorageBuffers && hasStepUniforms;
    const bool useShadingStorageBuffer = useStorageBuffers && step->performsShading();

    auto allReqFlags = srcColorRoot->requiredFlags() | finalBlendRoot->requiredFlags();
    if (clipRoot) {
        allReqFlags |= clipRoot->requiredFlags();
    }
    const bool useGradientStorageBuffer = caps->gradientBufferSupport() &&
                                          (allReqFlags & SnippetRequirementFlags::kGradientBuffer);

    const bool useDstSampler = fDstReadStrategy == DstReadStrategy::kTextureCopy ||
                               fDstReadStrategy == DstReadStrategy::kTextureSample;

    const std::vector<LiftedExpression> liftedExpr = collect_lifted_expressions(fRootNodes);

    // The uniforms are mangled by having their index in 'fEntries' as a suffix (i.e., "_%d")
    const ResourceBindingRequirements& bindingReqs = caps->resourceBindingRequirements();

    std::string preamble;
    int numPaintUniforms = 0;
    int numUnliftedPaintUniforms = 0;
    bool wrotePaintColor = false;
    if (useShadingStorageBuffer) {
        preamble = emit_paint_params_storage_buffer(bindingReqs.fUniformsSetIdx,
                                                    bindingReqs.fPaintParamsBufferBinding,
                                                    fRootNodes,
                                                    &numPaintUniforms,
                                                    &numUnliftedPaintUniforms,
                                                    &wrotePaintColor);
        SkSL::String::appendf(&preamble, "uint %s;\n", this->shadingSsboIndex());
    } else {
        preamble = emit_paint_params_uniforms(bindingReqs.fUniformsSetIdx,
                                              bindingReqs.fPaintParamsBufferBinding,
                                              bindingReqs.fUniformBufferLayout,
                                              fRootNodes,
                                              &numPaintUniforms,
                                              &numUnliftedPaintUniforms,
                                              &wrotePaintColor);
    }
    fHasPaintUniforms = numPaintUniforms > 0;
    fHasLiftedPaintUniforms = numPaintUniforms - numUnliftedPaintUniforms > 0;
    const bool hasUnliftedPaintUniforms = numUnliftedPaintUniforms > 0;

    fHasSsboIndicesVarying = useShadingStorageBuffer &&
                             (hasUnliftedPaintUniforms || hasStepUniforms);
    const bool defineLocalCoordsVarying = this->needsLocalCoords();
    preamble += emit_varyings(step,
                              /*direction=*/"in",
                              liftedExpr,
                              fHasSsboIndicesVarying,
                              defineLocalCoordsVarying);

    preamble += emit_intrinsic_constants(bindingReqs);

    if (fDstReadStrategy == DstReadStrategy::kReadFromInput) {
        // If this shader reads the dst texture as an input attachment, assert that a valid set
        // index has been assigned within ResourceBindingRequirements.
        SkASSERT(bindingReqs.fInputAttachmentSetIdx != ResourceBindingRequirements::kUnassigned);
        // TODO: The following SkSL depends upon the fact that Vulkan is currently the only backend
        // that utilizes DstReadStrategy::kReadFromInput. Update accordingly if other backends add
        // support for this DstReadStrategy.
        SkSL::String::appendf(
                &preamble,
                "layout (vulkan, input_attachment_index=%d, set=%d, binding=%d) "
                "subpassInput DstTextureInput;\n",
                /*input attachment idx within set=*/0,
                /*input attachment set idx=*/bindingReqs.fInputAttachmentSetIdx,
                /*binding=*/0);
    }

    if (hasStepUniforms) {
        if (useStepStorageBuffer) {
            preamble += emit_render_step_storage_buffer(bindingReqs.fUniformsSetIdx,
                                                        bindingReqs.fRenderStepBufferBinding,
                                                        step->uniforms());
        } else {
            preamble += emit_render_step_uniforms(bindingReqs.fUniformsSetIdx,
                                                  bindingReqs.fRenderStepBufferBinding,
                                                  bindingReqs.fUniformBufferLayout,
                                                  step->uniforms());
        }
    }

    if (useGradientStorageBuffer) {
        SkSL::String::appendf(&preamble,
                              "layout (set=%d, binding=%d) readonly buffer FSGradientBuffer {\n"
                              "    float %s[];\n"
                              "};\n",
                              bindingReqs.fUniformsSetIdx,
                              bindingReqs.fGradientBufferBinding,
                              ShaderInfo::kGradientBufferName);
        fHasGradientBuffer = true;
    }

    {
        int binding = 0;
        preamble += emit_textures_and_samplers(bindingReqs, fRootNodes, &binding, outDescs);
        int paintTextureCount = binding;
        if (step->hasTextures()) {
            preamble += step->texturesAndSamplersSkSL(bindingReqs, &binding);
            if (outDescs) {
                // Determine how many render step samplers were used by comparing the binding value
                // against paintTextureCount, taking into account the binding requirements. We
                // assume and do not anticipate the render steps to use immutable samplers.
                int renderStepSamplerCount = bindingReqs.fSeparateTextureAndSamplerBinding
                                                     ? (binding - paintTextureCount) / 2
                                                     : binding - paintTextureCount;
                // Add default SamplerDescs for all the dynamic samplers used by the render step so
                // the size of outDescs will be equivalent to the total number of samplers.
                outDescs->push_back_n(renderStepSamplerCount);
            }
        }
        if (useDstSampler) {
            preamble += EmitSamplerLayout(bindingReqs, &binding);
            preamble += " sampler2D dstSampler;";
            // Add default SamplerDesc for the intrinsic dstSampler to stay consistent with
            // `fNumFragmentTexturesAndSamplers`.
            if (outDescs) {
                outDescs->push_back({});
            }
        }

        // Record how many textures and samplers are used.
        fNumFragmentTexturesAndSamplers = binding;
    }

    // Emit preamble declarations and helper functions required for snippets. In the default case
    // this adds functions that bind a node's specific mangled uniforms to the snippet's
    // implementation in the SkSL modules.
    emit_preambles(*this, fRootNodes, /*treeLabel=*/"", &preamble);

    std::string mainBody = "void main() {";

    if (fHasSsboIndicesVarying) {
        SkSL::String::appendf(&mainBody,
                              "%s = %s.y;\n",
                              this->shadingSsboIndex(),
                              RenderStep::ssboIndicesVarying());
    }

    if (step->emitsPrimitiveColor()) {
        mainBody += "half4 primitiveColor;";
        mainBody += step->fragmentColorSkSL();
    } else {
        SkASSERT(!(fRootNodes[0]->requiredFlags() & SnippetRequirementFlags::kPrimitiveColor));
    }

    // Using kDefaultArgs as the initial value means it will refer to undefined variables, but the
    // root nodes should--at most--be depending on the coordinate when "needsLocalCoords" is true.
    // If the PaintParamsKey violates that structure, this will produce SkSL compile errors.
    ShaderSnippet::Args args = ShaderSnippet::kDefaultArgs;
    args.fFragCoord = "localCoordsVar";  // the varying added in emit_varyings()
    // TODO(b/349997190): The paint root node should not depend on any prior stage's output, but
    // it can happen with how SkEmptyShader is currently mapped to `sk_passthrough`. In this case
    // it requires that prior stage color to be transparent black. When SkEmptyShader can instead
    // cause the draw to be skipped, this can go away.
    args.fPriorStageOutput = "half4(0)";

    // Calculate the src color and stash its output variable in `args`
    args.fPriorStageOutput = srcColorRoot->invokeAndAssign(*this, args, &mainBody);

    // If not using hardware blending, we perform a dst read in the shader and must add SkSL
    // accordingly.
    const bool readDstInShader = !useHardwareBlending;
    if (readDstInShader) {
        // Get the current dst color into a local variable, it may be used later on for coverage
        // blending as well as the final blend.
        mainBody += "half4 dstColor;";
        if (useDstSampler) {
            // dstReadBounds is in frag coords and already includes the replay translation. The
            // reciprocol of the dstCopy dimensions are in ZW.
            mainBody += "dstColor = sample(dstSampler,"
                                          "dstReadBounds.zw*(sk_FragCoord.xy - dstReadBounds.xy));";
        } else if (fDstReadStrategy == DstReadStrategy::kReadFromInput) {
            // The dst texture should have been written to with the appropriate write swizzle, so we
            // do not need to worry about the read swizzle when accessing that value for blending.
            mainBody += "// Read color from input attachment\n";
            mainBody += "dstColor = subpassLoad(DstTextureInput);\n";
        } else {
            SkASSERT(fDstReadStrategy == DstReadStrategy::kFramebufferFetch);
            mainBody += "dstColor = sk_LastFragColor;";
        }

        args.fBlenderDstColor = "dstColor";
        args.fPriorStageOutput = finalBlendRoot->invokeAndAssign(*this, args, &mainBody);
    }

    if (writeSwizzle != Swizzle::RGBA()) {
        SkSL::String::appendf(&mainBody, "%s = %s.%s;", args.fPriorStageOutput.c_str(),
                                                        args.fPriorStageOutput.c_str(),
                                                        writeSwizzle.asString().c_str());
    }

    if (finalCoverage == Coverage::kNone) {
        // Either direct HW blending or a dst-read w/o any extra coverage. In both cases we just
        // need to assign directly to sk_FragCoord and update the HW blend info to finalBlendMode.
        SkASSERT(finalBlendMode.has_value());

        fBlendInfo = gBlendTable[static_cast<int>(*finalBlendMode)];
        SkSL::String::appendf(&mainBody, "sk_FragColor = %s;", args.fPriorStageOutput.c_str());
    } else {
        // Accumulate the output coverage. This will either modify the src color and secondary
        // outputs for dual-source blending, or be combined directly with the in-shader blended
        // final color if a dst-readback was required.
        if (useStepStorageBuffer) {
            SkSL::String::appendf(&mainBody,
                                  "uint stepSsboIndex = %s.x;\n",
                                  RenderStep::ssboIndicesVarying());
            mainBody +=
                    emit_uniforms_from_storage_buffer("step", "stepSsboIndex", step->uniforms());
        }

        mainBody += "half4 outputCoverage = half4(1);";
        mainBody += step->fragmentCoverageSkSL();

        if (clipRoot) {
            // The clip block node is invoked with device coords, not local coords like the main
            // shading root node. However sk_FragCoord includes any replay translation and we
            // need to recover the original device coordinate.
            mainBody += "float2 devCoord = sk_FragCoord.xy - viewport.xy;";
            args.fFragCoord = "devCoord";
            std::string clipBlockOutput = clipRoot->invokeAndAssign(*this, args, &mainBody);
            SkSL::String::appendf(&mainBody, "outputCoverage *= %s.a;", clipBlockOutput.c_str());
        }

        const char* outColor = args.fPriorStageOutput.c_str();
        if (readDstInShader) {
            // If this draw uses a non-coherent dst read, we want to keep the existing dst color (or
            // whatever has been previously drawn) when there's no coverage. This helps for batching
            // text draws that need to read from a dst copy for blends. However, this only helps the
            // case where the outer bounding boxes of each letter overlap and not two actual parts
            // of the text.
            if (useDstSampler) {
                // We don't think any shaders actually output negative coverage, but just as a
                // safety check for floating point precision errors, we compare with <= here. We
                // just check the RGB values of the coverage, since the alpha may not have been set
                // when using LCD. If we are using single-channel coverage, alpha will be equal to
                // RGB anyway.
                mainBody +=
                    "if (all(lessThanEqual(outputCoverage.rgb, half3(0)))) {"
                        "discard;"
                    "}";
            }

            // Use kSrc HW BlendInfo and do the coverage blend with dst in the shader.
            SkASSERT(finalBlendMode.has_value() && finalBlendMode.value() == SkBlendMode::kSrc);
            fBlendInfo = gBlendTable[static_cast<int>(*finalBlendMode)];
            SkSL::String::appendf(
                    &mainBody,
                    "sk_FragColor = %s * outputCoverage + dstColor * (1.0 - outputCoverage);",
                    outColor);
            if (finalCoverage == Coverage::kLCD) {
                SkSL::String::appendf(
                        &mainBody,
                        "half3 lerpRGB = mix(dstColor.aaa, %s.aaa, outputCoverage.rgb);"
                        "sk_FragColor.a = max(max(lerpRGB.r, lerpRGB.g), lerpRGB.b);",
                        outColor);
            }
        } else {
            // Adjust the shader output(s) to incorporate the coverage so that HW blending produces
            // the correct output.
            if (finalBlendMode > SkBlendMode::kLastCoeffMode) {
                SkASSERT(finalCoverage == Coverage::kSingleChannel);
                fBlendInfo = gBlendTable[static_cast<int>(*finalBlendMode)];
                mainBody += emit_advanced_blend_color_output("sk_FragColor", outColor);
            } else {
                // Porter-Duff blend modes can utilize BlendFormula.
                // TODO: Determine whether draw is opaque and pass that to GetBlendFormula.
                BlendFormula coverageBlendFormula =
                        finalCoverage == Coverage::kLCD
                                ? skgpu::GetLCDBlendFormula(*finalBlendMode)
                                : skgpu::GetBlendFormula(
                                        /*isOpaque=*/false, /*hasCoverage=*/true, *finalBlendMode);
                fBlendInfo = {coverageBlendFormula.equation(),
                              coverageBlendFormula.srcCoeff(),
                              coverageBlendFormula.dstCoeff(),
                              SK_PMColor4fTRANSPARENT,
                              coverageBlendFormula.modifiesDst()};

                if (finalCoverage == Coverage::kLCD) {
                    mainBody += "outputCoverage.a = max(max(outputCoverage.r, "
                                                           "outputCoverage.g), "
                                                   "outputCoverage.b);";
                }

                mainBody += emit_color_output(coverageBlendFormula.primaryOutput(),
                                              "sk_FragColor",
                                              outColor);
                if (coverageBlendFormula.hasSecondaryOutput()) {
                    SkASSERT(caps->shaderCaps()->fDualSourceBlendingSupport);
                    mainBody += emit_color_output(coverageBlendFormula.secondaryOutput(),
                                                  "sk_SecondaryFragColor",
                                                  outColor);
                }
            }
        }
    }
    mainBody += "}\n";

    fFragmentSkSL = preamble + "\n" + mainBody;

    fFSLabel = writeSwizzle.asString().c_str();
    fFSLabel += " + ";
    fFSLabel = step->name();
    fFSLabel += " + ";
    fFSLabel += label;
    if (fDstReadStrategy != DstReadStrategy::kNoneRequired) {
        fFSLabel += " + Dst Read (";
        fFSLabel += dst_read_strategy_to_str(fDstReadStrategy);
        fFSLabel += ")";
    }
}