Repository: invenia/NamedDims.jl (julia)

Issue:
Type instability when broadcasting
I was trying to chase down allocations that appeared in my code when converting from `Array` to `NamedDimsArray`. I think I have found a case where we may be failing to get type-stable code. Here is what I think is a minimal example:
```
julia> using NamedDims, BenchmarkTools

julia> function example(a,b,c,d)
           @. a -= 0.5 * d * b * c
           return a
       end
example (generic function with 1 method)

julia> a = NamedDimsArray{(:z,)}(rand(5));

julia> b = NamedDimsArray{(:z,)}(rand(5));

julia> c = rand();

julia> d = rand();

julia> @btime example($a,$b,$c,$d)
  733.862 ns (10 allocations: 352 bytes)
5-element NamedDimsArray{(:z,), Float64, 1, Vector{Float64}}:
 -1085.0390848895406
 -3944.198097644197
 -4564.398164918389
 -3641.011715928582
 -5293.772958207846
```
Using `Cthulhu.@descend` suggests that the problem is return type not being inferred for `broadcasted_names(bs...)` from inside `copyto!()` here
https://github.com/invenia/NamedDims.jl/blob/2fa07ee7d94f34767db849942828f97ba0232b28/src/broadcasting.jl#L68-L80
even though when I then step into `broadcasted_names()`, the return type is `Core.Const((:z,))`!

I think the example is minimal, because if I remove even the `0.5 *` from `example()`, `@btime` reports zero allocations.

I was using julia-1.6.4 and NamedDims-0.2.41. Any suggestions for fixes or further debugging very welcome!

Pull request:
Apparently less recursion helps the compiler to const-propagate.

Fixes #187.

Patch:
diff --git a/Project.toml b/Project.toml
index f82dce3..5542dae 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NamedDims"
 uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.2.41"
+version = "0.2.42"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
diff --git a/src/broadcasting.jl b/src/broadcasting.jl
index bfa98bb..a6b4d98 100644
--- a/src/broadcasting.jl
+++ b/src/broadcasting.jl
@@ -73,10 +73,18 @@ function Base.copyto!(dest::AbstractArray, bc::Broadcasted{NamedDimsStyle{S}}) w
 end
 
 broadcasted_names(bc::Broadcasted) = broadcasted_names(bc.args...)
-function broadcasted_names(a, bs...)
+function broadcasted_names(a, b)
     a_name = broadcasted_names(a)
-    b_name = broadcasted_names(bs...)
+    b_name = broadcasted_names(b)
     return unify_names_longest(a_name, b_name)
 end
+# Including two explicit arguments like this before starting recursion helps
+# const-propagation
+function broadcasted_names(a, b, cs...)
+    a_name = broadcasted_names(a)
+    b_name = broadcasted_names(b)
+    c_name = broadcasted_names(cs...)
+    return unify_names_longest(a_name, b_name, c_name)
+end
 broadcasted_names(a::AbstractArray) = dimnames(a)
 broadcasted_names(a) = tuple()


Test patch:
diff --git a/test/broadcasting.jl b/test/broadcasting.jl
index 5b2d135..9bfd663 100644
--- a/test/broadcasting.jl
+++ b/test/broadcasting.jl
@@ -144,3 +144,16 @@ end
     @test ((1,2,3,4) .+ nda) isa NamedDimsArray{(:foo,)}
     @test ((1,2,3,4) .+ nda) == [1, 2, 3, 4]
 end
+
+function foo(a, b, c, d)
+    @. a -= 0.5 * d * b * c
+    return a
+end
+@testset "Regression test against allocations in broadcasting #187" begin
+    # https://github.com/invenia/NamedDims.jl/issues/187
+    a = NamedDimsArray{(:z,)}(rand(5))
+    b = NamedDimsArray{(:z,)}(rand(5))
+    c = rand()
+    d = rand()
+    @test_modern 0 == @ballocated foo($a, $b, $c, $d)
+end