Long-running shaders result in webgl context lost

TensorFlow.js version

0.10.0

Browser version

Chrome 64.0.3282.167 (Linux)

Describe the problem or feature request

Consider the following toy example of matrix exponentiation:

const size = 10000;
const iters = 10;

let tensor, tmp;
tensor = tf.zeros([size, size]);
for (let i=0; i<iters; i++) {
  tmp = tensor.matMul(tensor);
  tensor.dispose();
  tensor = tmp;
}
tensor.max().print();

Expected:

The program should always print 0, and neither the behavior nor memory usage should vary as iters increases, since it results in multiplying the exact same matrix, and the tensor is disposed after being used.

Actual:

With iters = 5, the program behaves as expected. With iters = 10 I have seen the following three behaviors:

• Sometimes it prints NaN, and Chrome reports that “WebGL hit a snag”
• Sometimes it errors with “Couldn’t compile vertex shader”
• Sometimes it errors with the long failed script compilation found here

Reproduction

Navigate to https://js.tensorflow.org/, open the console, and post the example script there. If it works as expected, try increasing the number of iterations and see if it breaks.

[dsmilkov]

Hey Dandelion, thanks for the detailed report! And hi again!

This is consistent with our observation of executing long-running shaders. A matmul of matrices of size 10Kx10K is ~2TFLOPS. The hypothesis is that Chrome is trying to protect itself (other tabs) by prematurely killing the WebGL context, if a single program takes more than x seconds, especially after several consecutive programs. This results in either a NaN or WebGL hit a Snag error.

The solution is non-trivial, but the idea is to use some divide and conquer where we internally divide our matmul job into several smaller matmul shader calls, each on the order GLOPS, instead of TFLOPS, which will keep Chrome and other browsers happy.

We won’t be prioritizing this yet, since it hasn’t come up in real use-cases (existing ML models), but happy to revisit this later, or take contributions!

[dsmilkov]

Thanks! We won’t be able to fix this for several reasons:

• The fix for this is non-trivial and likely impossible without breaking the API since we would need to “release the GPU and the main thread”, but we probably won’t be able to keep the entire matmul operation synchronous (tf.matmul().dataSync() will break).
• Huge matmuls like this do not show up in practical ML models, thus making this a low priority for the time being.