@neighbors fails for 3D cube

[shy86]

When extending the official 2D @Stencil + @neighbors example to 3D (i.e. a cubic DArray), the code errors.

using Dagger
import Dagger: @stencil, Wrap

A = ones(Blocks(1, 1, 1), Int, 3, 3, 3)
A[:] = 1:length(A)
B = zeros(Blocks(1, 1, 1), Int, 3, 3, 3)

Dagger.spawn_datadeps() do
    @stencil begin
        B[idx] = sum(@neighbors(A[idx], 1, Wrap())) / 27
    end
end

ERROR: DTaskFailedException:
  Root Exception Type: AssertionError
  Root Exception:
AssertionError: Halo mismatch: edges=6 corners=6

[jpsamaroo]

Thanks for the report! That particular error was due to an incorrect indexing calculation (which thankfully this assertion caught), but there are other errors. Will update here once I have a solution.

[jpsamaroo]

It appears that the structure of the Dagger.HaloArray itself is incorrect for 3D problems (it fails to model non-edge, non-corner accesses, which we start to encounter only on >= 3 dimension accesses); I'm implementing a fix now that should scale to arbitrary dimensions.

[shy86]

@neighbors is a very cool macro. Compared with Dask’s map_overlap, I find its flexibility somewhat limited, for example in allowing specify different shared depths per dimension (although this can be handled by two-step indexing manually) and in boundary handling (e.g., no reflect type support, and NaN as a constant seems to fail). Improving these aspects would also be great.

[jpsamaroo]

allowing specify different shared depths per dimension

As in, having different neighborhood size per dimension? That should be easy enough, internally we already represent each neighborhood size dimension as its own value, we just don't have syntax to @neighbors for it. I can add that once we get this issue fixed 😄

no reflect type support

We should be able to support this with the boundary_transition and load_boundary_* functions that we can overload. It's probably very easy to implement, I'll add it to the Todo list!

NaN as a constant seems to fail

Interesting, I hadn't considered that as a desirable value, nor did I expect that it would fail. It's probably easy to support. Do you have a repro with error+stacktrace of it failing?

Also, for any of these, feel free to investigate and submit PRs for them! They just each need one test in test/array/stencil.jl.

[jpsamaroo]

NaN as a constant seems to fail

I tested this on #674 and it seems to work fine, as far as I can tell:

using Dagger

A = rand(Blocks(5, 5), 10, 10)
B = zeros(Blocks(5, 5), 10, 10)
Dagger.spawn_datadeps() do
    Dagger.@stencil begin
        B[idx] = sum(@neighbors(A[idx], (1, 2), Dagger.Pad(NaN)))
    end
end
display(B)

gives:

10x10 DMatrix{Float64} with 2x2 partitions of size 5x5:
 NaN  NaN  NaN        NaN        NaN        NaN        NaN        NaN        NaN  NaN
 NaN  NaN    7.60735    6.95167    7.09881    7.3372     6.91208    7.1308   NaN  NaN
 NaN  NaN    7.48192    5.97147    7.04651    6.49275    7.23797    6.62607  NaN  NaN
 NaN  NaN    8.66997    7.94676    8.37253    8.43407    8.54045    7.76147  NaN  NaN
 NaN  NaN    9.65397    8.69925    9.18463    9.22319    8.80354    8.21733  NaN  NaN
 NaN  NaN    9.64152    9.5245     9.14565    9.35475    7.82104    7.62884  NaN  NaN
 NaN  NaN    8.87506    9.7561     9.41126    9.52394    8.29834    7.83582  NaN  NaN
 NaN  NaN    7.95466    9.78878    9.12836    9.16269    9.53042    9.33113  NaN  NaN
 NaN  NaN    7.19826    8.9009     8.73386    9.09649    9.95296   10.3441   NaN  NaN
 NaN  NaN  NaN        NaN        NaN        NaN        NaN        NaN        NaN  NaN

[shy86]

Sorry for not replying promptly. In my use case, it may be a bug in NaNStatistics rather than Dagger.

using Dagger
using NaNStatistics

A = ones(Blocks(1, 1), Int, 3, 3)
B = zeros(Blocks(1, 1), Float64, 3, 3)

Dagger.spawn_datadeps() do
    Dagger.@stencil begin
        # Calculate the nan average of the 3x3 neighborhood (including the center)
        B[idx] = nansum(@neighbors(A[idx], 1, Dagger.Pad(NaN))) / 9.0
    end
end

B

3x3 DMatrix{Float64} with 3x3 partitions of size 1x1:
 NaN  NaN    NaN
 NaN    1.0  NaN
 NaN  NaN    NaN

The expected B should be

0.44 0.44 0.44
0.44 1.00 0.44
0.44 0.44 0.44

[jpsamaroo]

I can reproduce this issue locally. I'll look more into it when I get spare time, but if you want to look into it further, I suggest checking if it also fails with OffsetArrays, which behave similarly to our HaloArray.

[jpsamaroo]

Also I'm working on Reflect support and other boundary conditions, I should have something within a week.