Fix measure type error on AMD GPU by passing type on loc function, and also fix broadcasting

[b-fg]

When running on LUMI, I came across an error type on measure, where loc was being used without passing the type. I think in general we should always pass the type when calling loc so I have fixed this across the codebase.

Also, I came across a broadcasting problem on the AMD GPU in LUMI, where operation such as a[i,:] .= ... within a kernel are not allowed. To solve this I explicitly wrote the necessary loops in the measure function, and the only other place where broadcasting is used within a kernel is in

WaterLily.jl/src/Metrics.jl

Line 98 in d23ac35

@loop df[I,:] .= p[I]*nds(body,loc(0,I),t) over I ∈ inside(p)

So probably this would fail too. I will first check any performance regression for measure and if all is good we can rewrite the forces routines too.

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

Attention: Patch coverage is 55.55556% with 8 lines in your changes missing coverage. Please review.

❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Files	Coverage Δ
src/util.jl	69.66% <66.66%> (-16.36%)	⬇️
src/Body.jl	77.77% <72.72%> (-22.23%)	⬇️
src/Metrics.jl	0.00% <0.00%> (-82.61%)	⬇️

... and 7 files with indirect coverage changes

[b-fg]

All tests passing locally on GPU as well.

[b-fg]

These are the benchmarks for master v1.1 before new mutable Simulation struct (68b7c6b), master v1.2 (d23ac35) and this PR (852dfdd) (all with the remeasure=true, which needs to be fixed). Not much difference on the large cases
Benchmarks also relevant for #149

Benchmark environment: jelly sim_step! (max_steps=100)
▶ log2p = 5
┌─────────┬───────────┬────────┬───────────┬─────────────┬────────┬──────────┬──────────┐
│ Backend │ WaterLily │ Julia  │ Precision │ Allocations │ GC [%] │ Time [s] │ Speed-up │
├─────────┼───────────┼────────┼───────────┼─────────────┼────────┼──────────┼──────────┤
│     GPU │   68b7c6b │ 1.10.4 │   Float32 │     6490205 │   0.00 │     1.54 │     1.00 │
│     GPU │   852dfdd │ 1.10.4 │   Float32 │     6486503 │   0.00 │     1.57 │     0.98 │
│     GPU │   d23ac35 │ 1.10.4 │   Float32 │     6486686 │   0.00 │     1.54 │     1.00 │
└─────────┴───────────┴────────┴───────────┴─────────────┴────────┴──────────┴──────────┘
▶ log2p = 6
┌─────────┬───────────┬────────┬───────────┬─────────────┬────────┬──────────┬──────────┐
│ Backend │ WaterLily │ Julia  │ Precision │ Allocations │ GC [%] │ Time [s] │ Speed-up │
├─────────┼───────────┼────────┼───────────┼─────────────┼────────┼──────────┼──────────┤
│     GPU │   68b7c6b │ 1.10.4 │   Float32 │     7776425 │   0.00 │     4.62 │     1.00 │
│     GPU │   852dfdd │ 1.10.4 │   Float32 │     7774282 │   0.00 │     4.79 │     0.96 │
│     GPU │   d23ac35 │ 1.10.4 │   Float32 │     7768925 │   0.00 │     4.69 │     0.98 │
└─────────┴───────────┴────────┴───────────┴─────────────┴────────┴──────────┴──────────┘
▶ log2p = 7
┌─────────┬───────────┬────────┬───────────┬─────────────┬────────┬──────────┬──────────┐
│ Backend │ WaterLily │ Julia  │ Precision │ Allocations │ GC [%] │ Time [s] │ Speed-up │
├─────────┼───────────┼────────┼───────────┼─────────────┼────────┼──────────┼──────────┤
│     GPU │   68b7c6b │ 1.10.4 │   Float32 │     6941539 │   0.00 │    33.07 │     1.00 │
│     GPU │   852dfdd │ 1.10.4 │   Float32 │     6937895 │   0.00 │    33.29 │     0.99 │
│     GPU │   d23ac35 │ 1.10.4 │   Float32 │     6940665 │   0.00 │    33.54 │     0.99 │
└─────────┴───────────┴────────┴───────────┴─────────────┴────────┴──────────┴──────────┘

[weymouth]

Looks good. If this is passing the AutoDiff test then you must be doing it right!

[weymouth]

This is the same as fill, which must maintain the type. So do these actually need to be specified?

[b-fg]

I don't think they need but this is more correct imho.

[weymouth]

how can it be more correct if they compile to the same thing and waste characters? ;-)

[b-fg]

x)

[b-fg]

The ForwardDiff test is only passing for Array, which is actually hard-coded in tests. So this still does not solve the AD for GPU issue. Also the autodiff test is not using measure since it contains a NoBody body.

[weymouth]

I'll run a test on the spinning cylinder example and add a body into the test suite.

[weymouth]

Doesn't work. Failed on apply!, didn't even get to measure!...

[b-fg]

Huh, I will test on that then and try to reproduce it.

[weymouth]

Let me finish the test case first and push it up

[weymouth]

Seems to be working now. I must have had a bug in the bigger version of the code. Current version is a single spinning cylinder. Tested with finite differences (although the FD accuracy is terrible for this case, it converges to the AD result).

[b-fg]

OK great. Note that with this PR the sdf needs to be defined with the same precision as the Simulation, at least when using AMDGPU.jl kernels. I think the CUDA.jl kernels are a bit more forgiving.