Support function calls in SkSL-to-SkVM

Bug: skia:10680
Change-Id: I8697bdc157d250f3c390c7f49074318aa8c7bdab
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/351918
Commit-Queue: Brian Osman <brianosman@google.com>
Reviewed-by: Mike Klein <mtklein@google.com>

src/core/SkRuntimeEffect.cpp

@@ -78,16 +78,9 @@ private:
 
             fCompiler = new SkSL::Compiler(fCaps.get());
 
-            // Using an inline threshold of zero would stop all inlining, and cause us to re-emit
+            // Using an inline threshold of zero stops all inlining, and causes us to re-emit SkSL
-            // SkSL that is nearly identical to what was ingested. That would be in the spirit of
+            // that is nearly identical to what was ingested.
-            // applying no workarounds, but causes problems (today). On the CPU backend, we only
+            fInlineThreshold = 0;
-            // compile the user SkSL once, then emit directly to skvm. The CPU backend doesn't
-            // support function calls, so some tests only work because of inlining. This needs to
-            // be addressed robustly - by adding function call support and/or forcing inlining,
-            // but for now, we use defaults that let the majority of our test cases work on all
-            // backends. (Note that there are other control flow constructs that don't work on the
-            // CPU backend, this is a special case of a more general problem.) skbug.com/10680
-            fInlineThreshold = SkSL::Program::Settings().fInlineThreshold;
         }
 
         SkSL::ShaderCapsPointer fCaps;

src/sksl/SkSLVMGenerator.cpp

@@ -74,6 +74,8 @@ struct Value {
     ValRef    operator[](int i)       { return fVals[i]; }
     skvm::Val operator[](int i) const { return fVals[i]; }
 
+    SkSpan<skvm::Val> asSpan() { return fVals; }
+
 private:
     SkSTArray<4, skvm::Val, true> fVals;
 };
@@ -209,7 +211,12 @@ private:
         return result;
     }
 
-    skvm::I32 mask() { return fMask; }
+    skvm::I32 mask() {
+        // As we encounter (possibly conditional) return statements, fReturned is updated to store
+        // the lanes that have already returned. For the remainder of the current function, those
+        // lanes should be disabled.
+        return fMask & ~currentFunction().fReturned;
+    }
 
     Value writeExpression(const Expression& expr);
     Value writeBinaryExpression(const BinaryExpression& b);
@@ -248,12 +255,18 @@ private:
     std::unordered_map<const Variable*, Slot> fVariableMap;
     std::vector<skvm::Val> fSlots;
 
+    // Conditional execution mask (changes are managed by AutoMask, and tied to control-flow scopes)
+    skvm::I32 fMask;
+
     //
     // State that's local to the generation of a single function:
     //
-    SkSpan<skvm::Val> fReturnValue;
+    struct Function {
-    skvm::I32 fMask;
+        const SkSpan<skvm::Val> fReturnValue;
-    skvm::I32 fReturned;
+        skvm::I32               fReturned;
+    };
+    std::vector<Function> fFunctionStack;
+    Function& currentFunction() { return fFunctionStack.back(); }
 
     class AutoMask {
     public:
@@ -362,8 +375,7 @@ SkVMGenerator::SkVMGenerator(const Program& program,
 
             { "sample", Intrinsic::kSample },
         } {
-    fMask     = fBuilder->splat(0xffff'ffff);
+    fMask = fBuilder->splat(0xffff'ffff);
-    fReturned = fBuilder->splat(0);
 
     // Now, add storage for each global variable (including uniforms) to fSlots, and entries in
     // fVariableMap to remember where every variable is stored.
@@ -420,11 +432,10 @@ SkVMGenerator::SkVMGenerator(const Program& program,
 void SkVMGenerator::writeFunction(const FunctionDefinition& function,
                                   SkSpan<skvm::Val> arguments,
                                   SkSpan<skvm::Val> outReturn) {
-
     const FunctionDeclaration& decl = function.declaration();
+    SkASSERT(slot_count(decl.returnType()) == outReturn.size());
 
-    fReturnValue = outReturn;
+    fFunctionStack.push_back({outReturn, /*returned=*/fBuilder->splat(0)});
-    SkASSERT(slot_count(decl.returnType()) == fReturnValue.size());
 
     // For all parameters, copy incoming argument IDs to our vector of (all) variable IDs
     size_t argIdx = 0;
@@ -455,6 +466,8 @@ void SkVMGenerator::writeFunction(const FunctionDefinition& function,
         argIdx += nslots;
     }
     SkASSERT(argIdx == arguments.size());
+
+    fFunctionStack.pop_back();
 }
 
 SkVMGenerator::Slot SkVMGenerator::getSlot(const Variable& v) {
@@ -1122,13 +1135,53 @@ Value SkVMGenerator::writeIntrinsicCall(const FunctionCall& c) {
 }
 
 Value SkVMGenerator::writeFunctionCall(const FunctionCall& f) {
-    // TODO: Support calling other functions (by recursively generating their programs, eg inlining)
+    if (f.function().isBuiltin() && !f.function().definition()) {
-    if (f.function().isBuiltin()) {
         return this->writeIntrinsicCall(f);
     }
 
-    SkDEBUGFAIL("Function calls not supported yet");
+    const FunctionDeclaration& decl = f.function();
-    return {};
+
+    // Evaluate all arguments, gather the results into a contiguous list of IDs
+    std::vector<skvm::Val> argVals;
+    for (const auto& arg : f.arguments()) {
+        Value v = this->writeExpression(*arg);
+        for (size_t i = 0; i < v.slots(); ++i) {
+            argVals.push_back(v[i]);
+        }
+    }
+
+    // Create storage for the return value
+    size_t nslots = slot_count(f.type());
+    Value result(nslots);
+    for (size_t i = 0; i < nslots; ++i) {
+        result[i] = fBuilder->splat(0.0f);
+    }
+
+    {
+        // This AutoMask merges currentFunction().fReturned into fMask. Lanes that conditionally
+        // returned in the current function would otherwise resume execution within the child.
+        AutoMask m(this, ~currentFunction().fReturned);
+        this->writeFunction(*f.function().definition(), argVals, result.asSpan());
+    }
+
+    // Propagate new values of any 'out' params back to the original arguments
+    const std::unique_ptr<Expression>* argIter = f.arguments().begin();
+    size_t valIdx = 0;
+    for (const Variable* p : decl.parameters()) {
+        size_t nslots = slot_count(p->type());
+        if (p->modifiers().fFlags & Modifiers::kOut_Flag) {
+            Value v(nslots);
+            for (size_t i = 0; i < nslots; ++i) {
+                v[i] = argVals[valIdx + i];
+            }
+            const std::unique_ptr<Expression>& arg = *argIter;
+            this->writeStore(*arg, v);
+        }
+        valIdx += nslots;
+        argIter++;
+    }
+
+    return result;
 }
 
 Value SkVMGenerator::writePrefixExpression(const PrefixExpression& p) {
@@ -1307,18 +1360,19 @@ void SkVMGenerator::writeIfStatement(const IfStatement& i) {
 }
 
 void SkVMGenerator::writeReturnStatement(const ReturnStatement& r) {
-    // TODO: Can we suppress other side effects for lanes that have returned? fMask needs to
+    skvm::I32 returnsHere = this->mask();
-    // fold in knowledge of conditional returns earlier in the function.
-    skvm::I32 returnsHere = bit_clear(this->mask(), fReturned);
 
-    // TODO: returns with no expression
+    if (r.expression()) {
-    Value val = this->writeExpression(*r.expression());
+        Value val = this->writeExpression(*r.expression());
 
-    for (size_t i = 0; i < val.slots(); ++i) {
+        int i = 0;
-        fReturnValue[i] = select(returnsHere, f32(val[i]), f32(fReturnValue[i])).id;
+        for (skvm::Val& slot : currentFunction().fReturnValue) {
+            slot = select(returnsHere, f32(val[i]), f32(slot)).id;
+            i++;
+        }
     }
 
-    fReturned |= returnsHere;
+    currentFunction().fReturned |= returnsHere;
 }
 
 void SkVMGenerator::writeVarDeclaration(const VarDeclaration& decl) {

tests/SkSLInterpreterTest.cpp

@@ -20,6 +20,10 @@ struct ProgramBuilder {
     ProgramBuilder(skiatest::Reporter* r, const char* src)
             : fCaps(GrContextOptions{}), fCompiler(&fCaps) {
         SkSL::Program::Settings settings;
+        // The SkSL inliner is well tested in other contexts. Here, we disable inlining entirely,
+        // to stress-test the VM generator's handling of function calls with varying signatures.
+        settings.fInlineThreshold = 0;
+
         fProgram = fCompiler.convertProgram(SkSL::Program::kGeneric_Kind, SkSL::String(src),
                                             settings);
         if (!fProgram) {
@@ -290,15 +294,12 @@ DEF_TEST(SkSLInterpreterAnd, r) {
             "color = half4(color.a); }", 1, 1, 0, 3, 1, 1, 0, 3);
     test(r, "void main(inout half4 color) { if (color.r > color.g && color.g > color.b) "
             "color = half4(color.a); }", 2, 1, 1, 3, 2, 1, 1, 3);
-    // TODO: SkVM function call support
     test(r, "int global; bool update() { global = 123; return true; }"
             "void main(inout half4 color) { global = 0;  if (color.r > color.g && update()) "
-            "color = half4(color.a); color.a = global; }", 2, 1, 1, 3, 3, 3, 3, 123,
+            "color = half4(color.a); color.a = global; }", 2, 1, 1, 3, 3, 3, 3, 123);
-            /*testWithSkVM=*/false);
     test(r, "int global; bool update() { global = 123; return true; }"
             "void main(inout half4 color) { global = 0;  if (color.r > color.g && update()) "
-            "color = half4(color.a); color.a = global; }", 1, 1, 1, 3, 1, 1, 1, 0,
+            "color = half4(color.a); color.a = global; }", 1, 1, 1, 3, 1, 1, 1, 0);
-            /*testWithSkVM=*/false);
 }
 
 DEF_TEST(SkSLInterpreterOr, r) {
@@ -308,15 +309,12 @@ DEF_TEST(SkSLInterpreterOr, r) {
             "color = half4(color.a); }", 1, 1, 0, 3, 3, 3, 3, 3);
     test(r, "void main(inout half4 color) { if (color.r > color.g || color.g > color.b) "
             "color = half4(color.a); }", 1, 1, 1, 3, 1, 1, 1, 3);
-    // TODO: SkVM function call support
     test(r, "int global; bool update() { global = 123; return true; }"
             "void main(inout half4 color) { global = 0;  if (color.r > color.g || update()) "
-            "color = half4(color.a); color.a = global; }", 1, 1, 1, 3, 3, 3, 3, 123,
+            "color = half4(color.a); color.a = global; }", 1, 1, 1, 3, 3, 3, 3, 123);
-            /*testWithSkVM=*/false);
     test(r, "int global; bool update() { global = 123; return true; }"
             "void main(inout half4 color) { global = 0;  if (color.r > color.g || update()) "
-            "color = half4(color.a); color.a = global; }", 2, 1, 1, 3, 3, 3, 3, 0,
+            "color = half4(color.a); color.a = global; }", 2, 1, 1, 3, 3, 3, 3, 0);
-            /*testWithSkVM=*/false);
 }
 
 DEF_TEST(SkSLInterpreterMatrix, r) {
@@ -738,6 +736,33 @@ DEF_TEST(SkSLInterpreterRestrictFunctionCalls, r) {
                       "{ for (int i = 0; i < 1; i++) { if (x > 2) { return x; } } return 0; }");
 }
 
+DEF_TEST(SkSLInterpreterReturnThenCall, r) {
+    // Test that early returns disable execution in subsequently called functions
+    const char* src = R"(
+        float y;
+        void inc () { ++y; }
+        void maybe_inc() { if (y < 0) return; inc(); }
+        void main(inout float x) { y = x; maybe_inc(); x = y; }
+    )";
+
+    ProgramBuilder program(r, src);
+    const SkSL::FunctionDefinition* main = SkSL::Program_GetFunction(*program, "main");
+    REPORTER_ASSERT(r, main);
+
+    skvm::Builder b;
+    SkSL::ProgramToSkVM(*program, *main, &b);
+    skvm::Program p = b.done();
+
+    float xs[] = { -2.0f, 0.0f, 3.0f, -1.0f };
+    const void* uniforms = nullptr;
+    p.eval(4, uniforms, xs);
+
+    REPORTER_ASSERT(r, xs[0] == -2.0f);
+    REPORTER_ASSERT(r, xs[1] ==  1.0f);
+    REPORTER_ASSERT(r, xs[2] ==  4.0f);
+    REPORTER_ASSERT(r, xs[3] == -1.0f);
+}
+
 DEF_TEST(SkSLInterpreterEarlyReturn, r) {
     // Unlike returns in loops, returns in conditionals should work.
     const char* src = "float main(float x, float y) { if (x < y) { return x; } return y; }";