Reland "Runtime effects: Detect passthrough sample calls automatically"

As explained, this is *very* conservative. It only works when the child
is sampled from within main, and using a direct reference to the coords
parameter. If that parameter is ever modified (even after being used),
the optimization doesn't happen. For most cases, this is fine.

Reland changes the logic in GrSkSLFP slightly, to avoid turning all
samples into pass-through when a child is sampled with both pass-through
and explicit coordinates.

Bug: skia:11869
Change-Id: Iec18f059b4e78df0d2f53449aa0c2945c58a58f7
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/401677
Commit-Queue: Brian Osman <brianosman@google.com>
Reviewed-by: John Stiles <johnstiles@google.com>
Reviewed-by: Michael Ludwig <michaelludwig@google.com>
diff --git a/src/core/SkRuntimeEffect.cpp b/src/core/SkRuntimeEffect.cpp
index 34ff7eb..32d0411 100644
--- a/src/core/SkRuntimeEffect.cpp
+++ b/src/core/SkRuntimeEffect.cpp
@@ -159,7 +159,19 @@
     settings.fForceNoInline = options.forceNoInline;
     settings.fAllowNarrowingConversions = true;
 
-    const SkSL::FunctionDefinition* main = nullptr;
+    // Find 'main', then locate the sample coords parameter. (It might not be present.)
+    const SkSL::FunctionDefinition* main = SkSL::Program_GetFunction(*program, "main");
+    if (!main) {
+        RETURN_FAILURE("missing 'main' function");
+    }
+    const auto& mainParams = main->declaration().parameters();
+    auto iter = std::find_if(mainParams.begin(), mainParams.end(), [](const SkSL::Variable* p) {
+        return p->modifiers().fLayout.fBuiltin == SK_MAIN_COORDS_BUILTIN;
+    });
+    const SkSL::ProgramUsage::VariableCounts sampleCoordsUsage =
+            iter != mainParams.end() ? program->usage()->get(**iter)
+                                     : SkSL::ProgramUsage::VariableCounts{};
+
     uint32_t flags = 0;
     switch (kind) {
         case SkSL::ProgramKind::kRuntimeColorFilter: flags |= kAllowColorFilter_Flag; break;
@@ -168,7 +180,9 @@
                                                                kAllowShader_Flag);    break;
         default: SkUNREACHABLE;
     }
-    if (SkSL::Analysis::ReferencesSampleCoords(*program)) {
+
+
+    if (sampleCoordsUsage.fRead || sampleCoordsUsage.fWrite) {
         flags |= kUsesSampleCoords_Flag;
     }
 
@@ -180,7 +194,7 @@
     // this can be simpler. There is no way for color filters to refer to sk_FragCoord or sample
     // coords in that mode.
     if ((flags & kAllowColorFilter_Flag) &&
-        (SkSL::Analysis::ReferencesFragCoords(*program) || (flags & kUsesSampleCoords_Flag))) {
+        ((flags & kUsesSampleCoords_Flag) || SkSL::Analysis::ReferencesFragCoords(*program))) {
         flags &= ~kAllowColorFilter_Flag;
     }
 
@@ -203,7 +217,8 @@
             // Child effects that can be sampled ('shader' or 'colorFilter')
             if (varType.isEffectChild()) {
                 children.push_back(var.name());
-                sampleUsages.push_back(SkSL::Analysis::GetSampleUsage(*program, var));
+                sampleUsages.push_back(SkSL::Analysis::GetSampleUsage(
+                        *program, var, sampleCoordsUsage.fWrite != 0));
             }
             // 'uniform' variables
             else if (var.modifiers().fFlags & SkSL::Modifiers::kUniform_Flag) {
@@ -234,18 +249,6 @@
                 uniforms.push_back(uni);
             }
         }
-        // Functions
-        else if (elem->is<SkSL::FunctionDefinition>()) {
-            const auto& func = elem->as<SkSL::FunctionDefinition>();
-            const SkSL::FunctionDeclaration& decl = func.declaration();
-            if (decl.isMain()) {
-                main = &func;
-            }
-        }
-    }
-
-    if (!main) {
-        RETURN_FAILURE("missing 'main' function");
     }
 
 #undef RETURN_FAILURE
diff --git a/src/gpu/effects/GrSkSLFP.cpp b/src/gpu/effects/GrSkSLFP.cpp
index aba6a96..fcdbbc2 100644
--- a/src/gpu/effects/GrSkSLFP.cpp
+++ b/src/gpu/effects/GrSkSLFP.cpp
@@ -96,6 +96,22 @@
             }
 
             String sampleChild(int index, String coords) override {
+                // If the child was sampled using the coords passed to main (and they are never
+                // modified), then we will have marked the child as PassThrough. The code generator
+                // doesn't know that, and still supplies coords. Inside invokeChild, we assert that
+                // any coords passed for a PassThrough child match args.fSampleCoords exactly.
+                //
+                // Normally, this is valid. Here, we *copied* the sample coords to a local variable
+                // (so that they're mutable in the runtime effect SkSL). Thus, the coords string we
+                // get here is the name of the local copy, and fSampleCoords still points to the
+                // unmodified original (which might be a varying, for example).
+                // To prevent the assert, we pass the empty string in this case. Note that for
+                // children sampled like this, invokeChild doesn't even use the coords parameter,
+                // except for that assert.
+                const GrFragmentProcessor* child = fArgs.fFp.childProcessor(index);
+                if (child && !child->isSampledWithExplicitCoords()) {
+                    coords.clear();
+                }
                 return String(fSelf->invokeChild(index, fInputColor, fArgs, coords).c_str());
             }
 
diff --git a/src/sksl/SkSLAnalysis.cpp b/src/sksl/SkSLAnalysis.cpp
index 852476d..2a079f3 100644
--- a/src/sksl/SkSLAnalysis.cpp
+++ b/src/sksl/SkSLAnalysis.cpp
@@ -74,8 +74,8 @@
 // Visitor that determines the merged SampleUsage for a given child 'fp' in the program.
 class MergeSampleUsageVisitor : public ProgramVisitor {
 public:
-    MergeSampleUsageVisitor(const Context& context, const Variable& fp)
-            : fContext(context), fFP(fp) {}
+    MergeSampleUsageVisitor(const Context& context, const Variable& fp, bool writesToSampleCoords)
+            : fContext(context), fFP(fp), fWritesToSampleCoords(writesToSampleCoords) {}
 
     SampleUsage visit(const Program& program) {
         fUsage = SampleUsage(); // reset to none
@@ -86,6 +86,7 @@
 protected:
     const Context& fContext;
     const Variable& fFP;
+    const bool fWritesToSampleCoords;
     SampleUsage fUsage;
 
     bool visitExpression(const Expression& e) override {
@@ -97,7 +98,18 @@
                 if (fc.arguments().size() >= 2) {
                     const Expression* coords = fc.arguments()[1].get();
                     if (coords->type() == *fContext.fTypes.fFloat2) {
-                        fUsage.merge(SampleUsage::Explicit());
+                        // If the coords are a direct reference to the program's sample-coords,
+                        // and those coords are never modified, we can conservatively turn this
+                        // into PassThrough sampling. In all other cases, we consider it Explicit.
+                        if (!fWritesToSampleCoords && coords->is<VariableReference>() &&
+                            coords->as<VariableReference>()
+                                            .variable()
+                                            ->modifiers()
+                                            .fLayout.fBuiltin == SK_MAIN_COORDS_BUILTIN) {
+                            fUsage.merge(SampleUsage::PassThrough());
+                        } else {
+                            fUsage.merge(SampleUsage::Explicit());
+                        }
                     } else if (coords->type() == *fContext.fTypes.fFloat3x3) {
                         // Determine the type of matrix for this call site
                         if (coords->isConstantOrUniform()) {
@@ -570,8 +582,10 @@
 ////////////////////////////////////////////////////////////////////////////////
 // Analysis
 
-SampleUsage Analysis::GetSampleUsage(const Program& program, const Variable& fp) {
-    MergeSampleUsageVisitor visitor(*program.fContext, fp);
+SampleUsage Analysis::GetSampleUsage(const Program& program,
+                                     const Variable& fp,
+                                     bool writesToSampleCoords) {
+    MergeSampleUsageVisitor visitor(*program.fContext, fp, writesToSampleCoords);
     return visitor.visit(program);
 }
 
diff --git a/src/sksl/SkSLAnalysis.h b/src/sksl/SkSLAnalysis.h
index 945c287..ee61671 100644
--- a/src/sksl/SkSLAnalysis.h
+++ b/src/sksl/SkSLAnalysis.h
@@ -33,7 +33,14 @@
  * Provides utilities for analyzing SkSL statically before it's composed into a full program.
  */
 struct Analysis {
-    static SampleUsage GetSampleUsage(const Program& program, const Variable& fp);
+    /**
+     * Determines how `program` samples `fp`. By default, assumes that the sample coords
+     * (SK_MAIN_COORDS_BUILTIN) might be modified, so `sample(fp, sampleCoords)` is treated as
+     * Explicit. If writesToSampleCoords is false, treats that as PassThrough, instead.
+     */
+    static SampleUsage GetSampleUsage(const Program& program,
+                                      const Variable& fp,
+                                      bool writesToSampleCoords = true);
 
     static bool ReferencesBuiltin(const Program& program, int builtin);
 
diff --git a/src/sksl/ir/SkSLProgram.h b/src/sksl/ir/SkSLProgram.h
index e44497e..152c6a6 100644
--- a/src/sksl/ir/SkSLProgram.h
+++ b/src/sksl/ir/SkSLProgram.h
@@ -199,6 +199,8 @@
         return result;
     }
 
+    const ProgramUsage* usage() const { return fUsage.get(); }
+
     std::unique_ptr<String> fSource;
     std::unique_ptr<ProgramConfig> fConfig;
     std::shared_ptr<Context> fContext;
diff --git a/tests/SkRuntimeEffectTest.cpp b/tests/SkRuntimeEffectTest.cpp
index 1055c04..edb4733 100644
--- a/tests/SkRuntimeEffectTest.cpp
+++ b/tests/SkRuntimeEffectTest.cpp
@@ -16,6 +16,7 @@
 #include "src/core/SkColorSpacePriv.h"
 #include "src/core/SkTLazy.h"
 #include "src/gpu/GrColor.h"
+#include "src/gpu/GrFragmentProcessor.h"
 #include "tests/Test.h"
 
 #include <algorithm>
@@ -525,3 +526,51 @@
         REPORTER_ASSERT(r, filter && !filter->isAlphaUnchanged());
     }
 }
+
+DEF_TEST(SkRuntimeShaderSampleUsage, r) {
+    auto test = [&](const char* src, bool expectExplicit) {
+        auto [effect, err] =
+                SkRuntimeEffect::MakeForShader(SkStringPrintf("uniform shader child; %s", src));
+        REPORTER_ASSERT(r, effect);
+
+        auto child = GrFragmentProcessor::MakeColor({ 1, 1, 1, 1 });
+        auto fp = effect->makeFP(nullptr, &child, 1);
+        REPORTER_ASSERT(r, fp);
+
+        REPORTER_ASSERT(r, fp->childProcessor(0)->isSampledWithExplicitCoords() == expectExplicit);
+    };
+
+    // This test verifies that we detect calls to sample where the coords are the same as those
+    // passed to main. In those cases, it's safe to turn the "explicit" sampling into "passthrough"
+    // sampling. This optimization is implemented very conservatively.
+
+    // Cases where our optimization is valid, and works:
+
+    // Direct use of passed-in coords
+    test("half4 main(float2 xy) { return sample(child, xy); }", false);
+    // Sample with passed-in coords, read (but don't write) sample coords elsewhere
+    test("half4 main(float2 xy) { return sample(child, xy) + sin(xy.x); }", false);
+
+    // Cases where our optimization is not valid, and does not happen:
+
+    // Sampling with values completely unrelated to passed-in coords
+    test("half4 main(float2 xy) { return sample(child, float2(0, 0)); }", true);
+    // Use of expression involving passed in coords
+    test("half4 main(float2 xy) { return sample(child, xy * 0.5); }", true);
+    // Use of coords after modification
+    test("half4 main(float2 xy) { xy *= 2; return sample(child, xy); }", true);
+    // Use of coords after modification via out-param call
+    test("void adjust(inout float2 xy) { xy *= 2; }"
+         "half4 main(float2 xy) { adjust(xy); return sample(child, xy); }", true);
+
+    // There should (must) not be any false-positive cases. There are false-negatives.
+    // In all of these cases, our optimization would be valid, but does not happen:
+
+    // Direct use of passed-in coords, modified after use
+    test("half4 main(float2 xy) { half4 c = sample(child, xy); xy *= 2; return c; }", true);
+    // Passed-in coords copied to a temp variable
+    test("half4 main(float2 xy) { float2 p = xy; return sample(child, p); }", true);
+    // Use of coords passed to helper function
+    test("half4 helper(float2 xy) { return sample(child, xy); }"
+         "half4 main(float2 xy) { return helper(xy); }", true);
+}