函数向输入参数添加维度时的类型稳定性问题答案

【问题标题】：type-stability issue when a function adds a dimension to the input argument函数向输入参数添加维度时的类型稳定性问题
【发布时间】：2018-02-28 10:00:59
【问题描述】：

我有一个函数，它返回一个元素类型与输入数组相同但具有额外维度的数组。这是一个简单的例子：

function myfun(a::Array{T,N}) where {T,N}
   b = Array{T,N+1}(size(a)...,2)
   b[:] = 42
   return b
end

在 2x2 数组上调用此函数时，它返回一个 2x2x2 数组。

myfun(zeros(2,2))
2×2×2 Array{Float64,3}:
[:, :, 1] =
 42.0  42.0
 42.0  42.0

[:, :, 2] =
 42.0  42.0
 42.0  42.0

但是，此函数不是类型稳定的。根据@code_warntype，b 的类型为Any。

即使在b 上使用类型注释，结果也不是类型稳定的关于维数：

function myfun(a::Array{T,N}) where {T,N}
      b = Array{T,N+1}(size(a)...,2) :: Array{T,N+1}
      b[:] = T(42)
      return b
end

@code_warntype myfun(zeros(2,2)) 现在返回Array{Float64,_} where _ 类型为b。当输入参数为 2 维时，Julia 是否无法计算出维数为 3？

我正在使用 julia 0.6.2（在 linux 上）。

【问题讨论】：

标签： types julia generic-programming

【解决方案1】：

这是由于构造函数(Array{T,N+1}(size(a)...,2))是在运行时执行的，你可以在编译时使用@generated functions预计算N：

julia> @generated function myfun(a::Array{T,N}) where {T,N}
           NN = N+1
           quote 
               b = Array{$T,$NN}(size(a)...,2)
               b[:] = 42
               return b
           end
       end
myfun (generic function with 1 method)

julia> @code_warntype myfun(zeros(2,2))
Variables:
  #self# <optimized out>
  a::Array{Float64,2}
  b::Array{Float64,3}

Body:
  begin  # line 2:
      # meta: location REPL[1] # line 4:
      SSAValue(2) = (Base.arraysize)(a::Array{Float64,2}, 1)::Int64
      SSAValue(1) = (Base.arraysize)(a::Array{Float64,2}, 2)::Int64
      b::Array{Float64,3} = $(Expr(:foreigncall, :(:jl_alloc_array_3d), Array{Float64,3}, svec(Any, Int64, Int64, Int64), Array{Float64,3}, 0, SSAValue(2), 0, SSAValue(1), 0, :($(QuoteNode(2))), 0)) # line 5:
      $(Expr(:invoke, MethodInstance for fill!(::Array{Float64,3}, ::Int64), :(Base.fill!), :(b), 42))
      # meta: pop location
      return b::Array{Float64,3}
  end::Array{Float64,3}

julia> myfun(zeros(2,2))
2×2×2 Array{Float64,3}:
[:, :, 1] =
 42.0  42.0
 42.0  42.0

[:, :, 2] =
 42.0  42.0
 42.0  42.0

【讨论】：