allow early returns

For posterity, here's my initial, wrong thinking:

    If we squint at "return foo" and read it as "result = foo; goto
    end_of_program", and then remember we can always skip forward jumps (and
    where's further forward than end of program?), early returns turn out to
    work just like a store.

The reason this is wrong is that by the time we reach a final return,
the entire mask stack has been popped back down to its original default
ffffffff (active) state.  But that return shouldn't override any prior
returns.  So that scheme isn't quite right.

Instead we accumulate the result by disabling updates to lanes that have
already returned.  By the time we're done, all lanes should have hit
_some_ active return, now asserted.

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

Author: Mike Klein

src/core/SkRuntimeEffect.cpp

@@ -361,6 +361,14 @@ static skvm::Color program_fn(skvm::Builder* p,
     std::vector<skvm::I32> cond_stack = { p->splat(0xffff'ffff) };
     std::vector<skvm::I32> mask_stack = cond_stack;
 
+    skvm::Color result = {
+        p->splat(0.0f),
+        p->splat(0.0f),
+        p->splat(0.0f),
+        p->splat(0.0f),
+    };
+    skvm::I32 result_locked_in = p->splat(0);
+
     for (const uint8_t *ip = fn.code(), *end = ip + fn.size(); ip != end; ) {
         using Inst = SkSL::ByteCodeInstruction;
 
@@ -637,21 +645,30 @@ static skvm::Color program_fn(skvm::Builder* p,
             } break;
 
             case Inst::kReturn: {
-                SkAssertResult(u8() == 4);
-                // We'd like to assert that (ip == end) -> there is only one return, but ByteCode
-                // always includes a kReturn/0 at the end of each function, as a precaution.
-                SkASSERT(stack.size() >= 4);
-                skvm::F32 a = pop(),
-                          b = pop(),
-                          g = pop(),
-                          r = pop();
-                return { r, g, b, a };
+                int count = u8();
+                SkAssertResult(count == 4 || count == 0);
+
+                if (count == 4) {
+                    SkASSERT(stack.size() >= 4);
+
+                    // Lane-by-lane, if we've already returned a value, that result is locked in;
+                    // later return instructions don't happen for that lane.
+                    skvm::I32 returns_here = bit_clear(mask_stack.back(),
+                                                       result_locked_in);
+
+                    result.a = select(returns_here, pop(), result.a);
+                    result.b = select(returns_here, pop(), result.b);
+                    result.g = select(returns_here, pop(), result.g);
+                    result.r = select(returns_here, pop(), result.r);
+
+                    result_locked_in |= returns_here;
+                }
             } break;
         }
     }
 
-    SkUNREACHABLE;
-    return {};
+    assert_true(result_locked_in);
+    return result;
 }
 
 static sk_sp<SkData> get_xformed_uniforms(const SkRuntimeEffect* effect,

src/core/SkVM.h

@@ -1219,6 +1219,9 @@ namespace skvm {
     static inline Q14& operator|=(Q14& x, Q14a y) { return (x = x | y); }
     static inline Q14& operator^=(Q14& x, Q14a y) { return (x = x ^ y); }
 
+    static inline I32 bit_clear(I32 x, I32a y) { return x->bit_clear(x,y); }
+    static inline I32 bit_clear(int x, I32  y) { return y->bit_clear(x,y); }
+
     static inline I32 select(I32 cond, I32a t, I32a f) { return cond->select(cond,t,f); }
     static inline F32 select(I32 cond, F32a t, F32a f) { return cond->select(cond,t,f); }
     static inline Q14 select(Q14 cond, Q14a t, Q14a f) { return cond->select(cond,t,f); }

src/sksl/SkSLByteCodeGenerator.cpp

@@ -1740,8 +1740,8 @@ void ByteCodeGenerator::writeIfStatement(const IfStatement& i) {
 }
 
 void ByteCodeGenerator::writeReturnStatement(const ReturnStatement& r) {
-    if (fLoopCount || fConditionCount) {
-        fErrors.error(r.fOffset, "return not allowed inside conditional or loop");
+    if (fLoopCount) {
+        fErrors.error(r.fOffset, "return not allowed inside loop");
         return;
     }
     int count = SlotCount(r.expression()->type());

tests/SkSLInterpreterTest.cpp

@@ -730,11 +730,45 @@ DEF_TEST(SkSLInterpreterRestrictFunctionCalls, r) {
     // Ensure that calls to undefined functions are not allowed (to prevent mutual recursion)
     expect_failure(r, "float foo(); float bar() { return foo(); } float foo() { return bar(); }");
 
-    // returns are not allowed inside conditionals (or loops, which are effectively the same thing)
-    expect_failure(r, "float main(float x, float y) { if (x < y) { return x; } return y; }");
+    // returns are not allowed inside loops
     expect_failure(r, "float main(float x) { while (x > 1) { return x; } return 0; }");
 }
 
+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; }";
+
+    SkSL::Compiler compiler;
+    std::unique_ptr<SkSL::Program> program = compiler.convertProgram(SkSL::Program::kGeneric_Kind,
+                                                                     SkSL::String(src),
+                                                                     SkSL::Program::Settings{});
+    REPORTER_ASSERT(r, program);
+
+    std::unique_ptr<SkSL::ByteCode> byteCode = compiler.toByteCode(*program);
+    REPORTER_ASSERT(r, byteCode);
+
+    // Our function should work fine via run().
+    const SkSL::ByteCodeFunction* fun = byteCode->getFunction("main");
+    float in[] = { 1.0f, 2.0f };
+    float ret;
+    REPORTER_ASSERT(r, byteCode->run(fun, in,2, &ret,1, nullptr,0));
+    REPORTER_ASSERT(r, ret == 1.0f);
+
+    in[0] = 3.0f;
+    REPORTER_ASSERT(r, byteCode->run(fun, in,2, &ret,1, nullptr,0));
+    REPORTER_ASSERT(r, ret == 2.0f);
+
+    // Now same again via runStriped(), won't quite work yet.
+    float xs[] = { 1.0f, 3.0f },
+          ys[] = { 2.0f, 2.0f };
+    float rets[2];
+    float* ins[] = { xs, ys };
+    float* outs[] = { rets+0, rets+1 } ;
+    REPORTER_ASSERT(r, byteCode->runStriped(fun,2, ins,2, outs,1, nullptr,0));
+    REPORTER_ASSERT(r, rets[0] == 1.0f);
+  //REPORTER_ASSERT(r, rets[1] == 2.0f);  // TODO: skia:10852, make striped early returns work.
+}
+
 DEF_TEST(SkSLInterpreterArrayBounds, r) {
     // Out of bounds array access at compile time
     expect_failure(r, "float main(float x[4]) { return x[-1]; }");