Potential for race conditions with shifted indices about tullio.jl HOT 6 CLOSED

Thanks for doing the test. Yes, I do think this can give incorrect gradients, sadly. Somehow I didn't think of this case before.

The obvious fix would be that, when i appears with two different shifts, it should be marked "unsafe" for the gradient calculation. (There are already some situations where the macro notices this.) However, in this example, marking both i and j unsafe would mean the entire gradient calculation has to run serially, which will be awful on the GPU.

A possible work-around would be to pretend that the different terms are separate arrays. Then it will accumulate the gradients into 𝛥x, 𝛥x2, 𝛥x3 and Zygote will combine them afterwards:


reg2(x) = begin
  x2 = x
  x3 = x
  sum(@tullio res[i, j] := abs2(x[i, j] - x2[i, j+1]) +  abs2(x[i, j] - x3[i+1, j]))
end


Perhaps it could solve for the index needed for the gradient, replacing

𝛥x[i + 1, j] = 𝛥x[i + 1, j] + 𝛥ℛ[1] * conj(-(((x[i, j] - x[i + 1, j]) + …


with

𝛥x[iplus1, j] = 𝛥x[iplus1, j] + 𝛥ℛ[1] * conj(-(((x[iplus1 - 1, j] - x[iplus1, j]) + …

But this would be pretty complicated, and the range of values allowed for iplus1 and i differ, so it wouldn’t be simple to fill 𝛥x in one safe loop like this.

from tullio.jl.

Yes. But I think it's much harder to get a wrong answer on the CPU, as there are fewer threads, which is probably why I didn't notice. For res above, if this testimage is divided 4 ways, and they each start at the top left of their blocks, they are very unlikely to fight each other at the overlapping edges.

There might be other functions which behave much worse, maybe something like x[end-i] instead of x[i+1]? But this doesn't seem to:

julia> f2(x) = sum(@tullio out[i] := x[i] + sqrt(x[20-i]) threads=2);  # 2 is the threshold to divide among threads

julia> Zygote.gradient(f2, 1:20)[1] ≈ ForwardDiff.gradient(f2, 1:20)   # all 4 threads work on these 20 elements
true

from tullio.jl.

So you mean the gradients could be wrong?

Because I'm observing certain artifacts in my deconvolution using the same regularizer on GPU vs on CPU. So something with the GPU might be wrong.

Here a MWE:

using Tullio, CUDA, KernelAbstractions, Zygote, TestImages

reg(x) = sum(@tullio res[i, j] := abs2(x[i, j] - x[i, j+1]) +  abs2(x[i, j] - x[i+1, j]))


x = Float32.(testimage("resolution_test_512"))
x_c = CuArray(x);

y1 = Zygote.gradient(reg, x)[1]
y2 = Zygote.gradient(reg, x_c)[1]
.≈(y1, Array(y2), rtol=0.01)[100:110, 100:110]

Returns

11×11 BitMatrix:
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1
 1  0  1  1  1  0  0  1  1  0  1

So the gradient might be broken, showing some column patterns.

from tullio.jl.

Ok, the workaround seems to work. I'll probably stick with it for now!
Thanks, your support is exceptional great! :)

from tullio.jl.

So this affects both CPU and GPU, correct?

from tullio.jl.

Yeah, I agree. Could be wrong but unlikely to happen and if, only affecting a minor amount of pixels.

from tullio.jl.