Inline the whole broadcast expression to avoid allocation (#30986)

* Inline the whole broadcast expression to avoid allocation

`map` has an inline limit of 16. To make sure that the whole broadcast
tree gets inlined properly, I added the `_inlined_map` function. I am not
sure if it is a good idea, but worth trying.

This PR solves the issue which I have mentioned in
2693778#issuecomment-461248258

```julia
julia> @allocated foo(tmp, uprev, k1, k2, k3, k4, k5, k6, k7, k8, k9, k10, k11, k12, k13,
                    k14, k15, k16, k17, k18, k19, k20, k21, k22, k23, k24, k25,
                    k26, k27, k28, k29, k30, k31, k32, k33, k34)
0
```

* Fix CI failure

* Stricter test (vectorization & no allocation for a 9-array bc)

* rm `_inlined_map`

Author: YingboMa

base/broadcast.jl

@@ -464,7 +464,8 @@ julia> Broadcast.combine_axes(1, 1, 1)
 ()
 ```
 """
-@inline combine_axes(A, B...) = broadcast_shape(axes(A), combine_axes(B...))
+@inline combine_axes(A, B, C...) = broadcast_shape(axes(A), combine_axes(B, C...))
+@inline combine_axes(A, B) = broadcast_shape(axes(A), axes(B))
 combine_axes(A) = axes(A)
 
 # shape (i.e., tuple-of-indices) inputs
@@ -502,7 +503,7 @@ function check_broadcast_shape(shp, Ashp::Tuple)
     _bcsm(shp[1], Ashp[1]) || throw(DimensionMismatch("array could not be broadcast to match destination"))
     check_broadcast_shape(tail(shp), tail(Ashp))
 end
-check_broadcast_axes(shp, A) = check_broadcast_shape(shp, axes(A))
+@inline check_broadcast_axes(shp, A) = check_broadcast_shape(shp, axes(A))
 # comparing many inputs
 @inline function check_broadcast_axes(shp, A, As...)
     check_broadcast_axes(shp, A)
@@ -911,7 +912,7 @@ _is_static_broadcast_28126(dest::AbstractArray, x::AbstractArray{<:Any,1}) = axe
 _is_static_broadcast_28126(dest::AbstractArray, x::AbstractArray) = axes(dest) == axes(x) # This can be better with other missing dimensions
 
 @inline _is_static_broadcast_28126_args(dest, args::Tuple) = _is_static_broadcast_28126(dest, args[1]) && _is_static_broadcast_28126_args(dest, tail(args))
-_is_static_broadcast_28126_args(dest, args::Tuple{Any}) = _is_static_broadcast_28126(dest, args[1])
+@inline _is_static_broadcast_28126_args(dest, args::Tuple{Any}) = _is_static_broadcast_28126(dest, args[1])
 _is_static_broadcast_28126_args(dest, args::Tuple{}) = true
 
 struct _NonExtruded28126{T}

test/broadcast.jl

@@ -791,14 +791,15 @@ let
 end
 
 @testset "large fusions vectorize and don't allocate (#28126)" begin
-    u, k1, k2, k3, k4, k5, k6, k7 = (ones(1000) for i in 1:8)
-    function goo(u, k1, k2, k3, k4, k5, k6, k7)
-        @. u = 0.1*(0.1*k1 + 0.2*k2 + 0.3*k3 + 0.4*k4 + 0.5*k5 + 0.6*k6 + 0.7*k7)
+    using InteractiveUtils: code_llvm
+    u, uprev, k1, k2, k3, k4, k5, k6, k7 = (ones(1000) for i in 1:9)
+    function goo(u, uprev, k1, k2, k3, k4, k5, k6, k7)
+        @. u = uprev + 0.1*(0.1*k1 + 0.2*k2 + 0.3*k3 + 0.4*k4 + 0.5*k5 + 0.6*k6 + 0.7*k7)
         nothing
     end
-    @allocated goo(u, k1, k2, k3, k4, k5, k6, k7)
-    @test @allocated(goo(u, k1, k2, k3, k4, k5, k6, k7)) == 0
-    @test occursin("vector.body", sprint(code_llvm, goo, NTuple{8, Vector{Float32}}))
+    @allocated goo(u, uprev, k1, k2, k3, k4, k5, k6, k7)
+    @test @allocated(goo(u, uprev, k1, k2, k3, k4, k5, k6, k7)) == 0
+    @test occursin("vector.body", sprint(code_llvm, goo, NTuple{9, Vector{Float32}}))
 end
 
 # Broadcasted iterable/indexable APIs