optimizer: compute side-effect-freeness for array allocations (#43565)

This would be useful for Julia-level optimizations on arrays.
Initially I want to have this in order to add array primitives support
in EscapeAnalysis.jl, which should help us implement a variety of array
optimizations including dead array allocation elimination, copy-elision
from `Array` to `ImmutableArray` conversion (#42465), etc., but I found
this might be already useful for us since this enables some DCE in very
simple cases like:
```julia
julia> function simple!(x::T) where T
           d = IdDict{T,T}() # dead alloc
           # ... computations that don't use `d` at all
           return nothing
       end
simple! (generic function with 1 method)

julia> @code_typed simple!("foo")
CodeInfo(
1 ─     return Main.nothing
) => Nothing
```

This enhancement is super limited though, e.g. DCE won't happen when
array allocation involves other primitive operations like `arrayset`:
```julia
julia> code_typed() do
           a = Int[0,1,2]
           nothing
       end

1-element Vector{Any}:
 CodeInfo(
1 ─ %1 = $(Expr(:foreigncall, :(:jl_alloc_array_1d), Vector{Int64}, svec(Any, Int64), 0, :(:ccall), Vector{Int64}, 3, 3))::Vector{Int64}
│        Base.arrayset(false, %1, 0, 1)::Vector{Int64}
│        Base.arrayset(false, %1, 1, 2)::Vector{Int64}
│        Base.arrayset(false, %1, 2, 3)::Vector{Int64}
└──      return Main.nothing
) => Nothing
```

Further enhancement o optimize cases like above will be based on top of
incoming EA.jl (Julia-level escape analysis) or LLVM-level escape analysis.

Author: aviatesk

base/array.jl

@@ -193,17 +193,17 @@ end
 
 
 """
-    Base.bitsunionsize(U::Union)
+    Base.bitsunionsize(U::Union) -> Int
 
 For a `Union` of [`isbitstype`](@ref) types, return the size of the largest type; assumes `Base.isbitsunion(U) == true`.
 
 # Examples
 ```jldoctest
 julia> Base.bitsunionsize(Union{Float64, UInt8})
-0x0000000000000008
+8
 
 julia> Base.bitsunionsize(Union{Float64, UInt8, Int128})
-0x0000000000000010
+16
 ```
 """
 function bitsunionsize(u::Union)

base/boot.jl

@@ -444,7 +444,6 @@ unsafe_convert(::Type{T}, x::T) where {T} = x
 
 const NTuple{N,T} = Tuple{Vararg{T,N}}
 
-
 ## primitive Array constructors
 struct UndefInitializer end
 const undef = UndefInitializer()

base/compiler/optimize.jl

@@ -236,6 +236,8 @@ function stmt_effect_free(@nospecialize(stmt), @nospecialize(rt), src::Union{IRC
                 eT ⊑ fT || return false
             end
             return true
+        elseif head === :foreigncall
+            return foreigncall_effect_free(stmt, rt, src)
         elseif head === :new_opaque_closure
             length(args) < 5 && return false
             typ = argextype(args[1], src)
@@ -260,6 +262,76 @@ function stmt_effect_free(@nospecialize(stmt), @nospecialize(rt), src::Union{IRC
     return true
 end
 
+function foreigncall_effect_free(stmt::Expr, @nospecialize(rt), src::Union{IRCode,IncrementalCompact})
+    args = stmt.args
+    name = args[1]
+    isa(name, QuoteNode) && (name = name.value)
+    isa(name, Symbol) || return false
+    ndims = alloc_array_ndims(name)
+    if ndims !== nothing
+        if ndims == 0
+            return new_array_no_throw(args, src)
+        else
+            return alloc_array_no_throw(args, ndims, src)
+        end
+    end
+    return false
+end
+
+function alloc_array_ndims(name::Symbol)
+    if name === :jl_alloc_array_1d
+        return 1
+    elseif name === :jl_alloc_array_2d
+        return 2
+    elseif name === :jl_alloc_array_3d
+        return 3
+    elseif name === :jl_new_array
+        return 0
+    end
+    return nothing
+end
+
+function alloc_array_no_throw(args::Vector{Any}, ndims::Int, src::Union{IRCode,IncrementalCompact})
+    length(args) ≥ ndims+6 || return false
+    atype = widenconst(argextype(args[6], src))
+    isType(atype) || return false
+    atype = atype.parameters[1]
+    dims = Csize_t[]
+    for i in 1:ndims
+        dim = argextype(args[i+6], src)
+        isa(dim, Const) || return false
+        dimval = dim.val
+        isa(dimval, Int) || return false
+        push!(dims, reinterpret(Csize_t, dimval))
+    end
+    return _new_array_no_throw(atype, ndims, dims)
+end
+
+function new_array_no_throw(args::Vector{Any}, src::Union{IRCode,IncrementalCompact})
+    length(args) ≥ 7 || return false
+    atype = widenconst(argextype(args[6], src))
+    isType(atype) || return false
+    atype = atype.parameters[1]
+    dims = argextype(args[7], src)
+    isa(dims, Const) || return dims === Tuple{}
+    dimsval = dims.val
+    isa(dimsval, Tuple{Vararg{Int}}) || return false
+    ndims = nfields(dimsval)
+    isa(ndims, Int) || return false
+    dims = Csize_t[reinterpret(Csize_t, dimval) for dimval in dimsval]
+    return _new_array_no_throw(atype, ndims, dims)
+end
+
+function _new_array_no_throw(@nospecialize(atype), ndims::Int, dims::Vector{Csize_t})
+    isa(atype, DataType) || return false
+    eltype = atype.parameters[1]
+    iskindtype(typeof(eltype)) || return false
+    elsz = aligned_sizeof(eltype)
+    return ccall(:jl_array_validate_dims, Cint,
+        (Ptr{Csize_t}, Ptr{Csize_t}, UInt32, Ptr{Csize_t}, Csize_t),
+        #=nel=#RefValue{Csize_t}(), #=tot=#RefValue{Csize_t}(), ndims, dims, elsz) == 0
+end
+
 """
     argextype(x, src::Union{IRCode,IncrementalCompact}) -> t
     argextype(x, src::CodeInfo, sptypes::Vector{Any}) -> t

base/compiler/tfuncs.jl

@@ -341,7 +341,7 @@ function sizeof_nothrow(@nospecialize(x))
     exact || return false # Could always be the type Bottom at runtime, for example, which throws
     t === DataType && return true # DataType itself has a size
     if isa(x, Union)
-        isinline, sz, _ = uniontype_layout(x)
+        isinline = uniontype_layout(x)[1]
         return isinline # even any subset of this union would have a size
     end
     isa(x, DataType) || return false
@@ -381,7 +381,7 @@ function sizeof_tfunc(@nospecialize(x),)
         # Normalize the query to ask about that type.
         x = unwrap_unionall(t)
         if exact && isa(x, Union)
-            isinline, sz, _ = uniontype_layout(x)
+            isinline = uniontype_layout(x)[1]
             return isinline ? Const(Int(Core.sizeof(x))) : Bottom
         end
         isa(x, DataType) || return Int
@@ -1470,30 +1470,27 @@ function _opaque_closure_tfunc(@nospecialize(arg), @nospecialize(isva),
 end
 
 # whether getindex for the elements can potentially throw UndefRef
-function array_type_undefable(@nospecialize(a))
-    if isa(a, Union)
-        return array_type_undefable(a.a) || array_type_undefable(a.b)
-    elseif isa(a, UnionAll)
+function array_type_undefable(@nospecialize(arytype))
+    if isa(arytype, Union)
+        return array_type_undefable(arytype.a) || array_type_undefable(arytype.b)
+    elseif isa(arytype, UnionAll)
         return true
     else
-        etype = (a::DataType).parameters[1]
-        return !(etype isa Type && (isbitstype(etype) || isbitsunion(etype)))
+        elmtype = (arytype::DataType).parameters[1]
+        return !(elmtype isa Type && (isbitstype(elmtype) || isbitsunion(elmtype)))
     end
 end
 
-function array_builtin_common_nothrow(argtypes::Array{Any,1}, first_idx_idx::Int)
+function array_builtin_common_nothrow(argtypes::Vector{Any}, first_idx_idx::Int)
     length(argtypes) >= 4 || return false
-    atype = argtypes[2]
-    (argtypes[1] ⊑ Bool && atype ⊑ Array) || return false
-    for i = first_idx_idx:length(argtypes)
-        argtypes[i] ⊑ Int || return false
-    end
+    boundcheck = argtypes[1]
+    arytype = argtypes[2]
+    array_builtin_common_typecheck(boundcheck, arytype, argtypes, first_idx_idx) || return false
     # If we could potentially throw undef ref errors, bail out now.
-    atype = widenconst(atype)
-    array_type_undefable(atype) && return false
+    arytype = widenconst(arytype)
+    array_type_undefable(arytype) && return false
     # If we have @inbounds (first argument is false), we're allowed to assume
     # we don't throw bounds errors.
-    boundcheck = argtypes[1]
     if isa(boundcheck, Const)
         !(boundcheck.val::Bool) && return true
     end
@@ -1503,19 +1500,33 @@ function array_builtin_common_nothrow(argtypes::Array{Any,1}, first_idx_idx::Int
     return false
 end
 
+function array_builtin_common_typecheck(
+    @nospecialize(boundcheck), @nospecialize(arytype),
+    argtypes::Vector{Any}, first_idx_idx::Int)
+    (boundcheck ⊑ Bool && arytype ⊑ Array) || return false
+    for i = first_idx_idx:length(argtypes)
+        argtypes[i] ⊑ Int || return false
+    end
+    return true
+end
+
+function arrayset_typecheck(@nospecialize(arytype), @nospecialize(elmtype))
+    # Check that we can determine the element type
+    arytype = widenconst(arytype)
+    isa(arytype, DataType) || return false
+    elmtype_expected = arytype.parameters[1]
+    isa(elmtype_expected, Type) || return false
+    # Check that the element type is compatible with the element we're assigning
+    elmtype ⊑ elmtype_expected || return false
+    return true
+end
+
 # Query whether the given builtin is guaranteed not to throw given the argtypes
 function _builtin_nothrow(@nospecialize(f), argtypes::Array{Any,1}, @nospecialize(rt))
     if f === arrayset
         array_builtin_common_nothrow(argtypes, 4) || return true
         # Additionally check element type compatibility
-        a = widenconst(argtypes[2])
-        # Check that we can determine the element type
-        isa(a, DataType) || return false
-        ap1 = a.parameters[1]
-        isa(ap1, Type) || return false
-        # Check that the element type is compatible with the element we're assigning
-        argtypes[3] ⊑ ap1 || return false
-        return true
+        return arrayset_typecheck(argtypes[2], argtypes[3])
     elseif f === arrayref || f === const_arrayref
         return array_builtin_common_nothrow(argtypes, 3)
     elseif f === Core._expr

base/reflection.jl

@@ -351,22 +351,24 @@ function datatype_alignment(dt::DataType)
     return Int(alignment)
 end
 
-function uniontype_layout(T::Type)
+function uniontype_layout(@nospecialize T::Type)
     sz = RefValue{Csize_t}(0)
     algn = RefValue{Csize_t}(0)
     isinline = ccall(:jl_islayout_inline, Cint, (Any, Ptr{Csize_t}, Ptr{Csize_t}), T, sz, algn) != 0
-    (isinline, sz[], algn[])
+    (isinline, Int(sz[]), Int(algn[]))
 end
 
+LLT_ALIGN(x, sz) = (x + sz - 1) & -sz
+
 # amount of total space taken by T when stored in a container
-function aligned_sizeof(T::Type)
+function aligned_sizeof(@nospecialize T::Type)
     @_pure_meta
     if isbitsunion(T)
         _, sz, al = uniontype_layout(T)
-        return (sz + al - 1) & -al
+        return LLT_ALIGN(sz, al)
     elseif allocatedinline(T)
         al = datatype_alignment(T)
-        return (Core.sizeof(T) + al - 1) & -al
+        return LLT_ALIGN(Core.sizeof(T), al)
     else
         return Core.sizeof(Ptr{Cvoid})
     end

src/array.c

@@ -76,27 +76,40 @@ typedef uint64_t wideint_t;
 
 #define MAXINTVAL (((size_t)-1)>>1)
 
+JL_DLLEXPORT int jl_array_validate_dims(size_t *nel, size_t *tot, uint32_t ndims, size_t *dims, size_t elsz)
+{
+    size_t i;
+    size_t _nel = 1;
+    for(i=0; i < ndims; i++) {
+        size_t di = dims[i];
+        wideint_t prod = (wideint_t)_nel * (wideint_t)di;
+        if (prod >= (wideint_t) MAXINTVAL || di >= MAXINTVAL)
+            return 1;
+        _nel = prod;
+    }
+    wideint_t prod = (wideint_t)elsz * (wideint_t)_nel;
+    if (prod >= (wideint_t) MAXINTVAL)
+        return 2;
+    *nel = _nel;
+    *tot = (size_t)prod;
+    return 0;
+}
+
 static jl_array_t *_new_array_(jl_value_t *atype, uint32_t ndims, size_t *dims,
-                               int8_t isunboxed, int8_t hasptr, int8_t isunion, int8_t zeroinit, int elsz)
+                               int8_t isunboxed, int8_t hasptr, int8_t isunion, int8_t zeroinit, size_t elsz)
 {
     jl_task_t *ct = jl_current_task;
-    size_t i, tot, nel=1;
+    size_t i, tot, nel;
     void *data;
     jl_array_t *a;
-
-    for(i=0; i < ndims; i++) {
-        size_t di = dims[i];
-        wideint_t prod = (wideint_t)nel * (wideint_t)di;
-        if (prod > (wideint_t) MAXINTVAL || di > MAXINTVAL)
-            jl_exceptionf(jl_argumenterror_type, "invalid Array dimensions");
-        nel = prod;
-    }
+    assert(isunboxed || elsz == sizeof(void*));
     assert(atype == NULL || isunion == jl_is_uniontype(jl_tparam0(atype)));
+    int validated = jl_array_validate_dims(&nel, &tot, ndims, dims, elsz);
+    if (validated == 1)
+        jl_exceptionf(jl_argumenterror_type, "invalid Array dimensions");
+    else if (validated == 2)
+        jl_error("invalid Array size");
     if (isunboxed) {
-        wideint_t prod = (wideint_t)elsz * (wideint_t)nel;
-        if (prod > (wideint_t) MAXINTVAL)
-            jl_error("invalid Array size");
-        tot = prod;
         if (elsz == 1 && !isunion) {
             // extra byte for all julia allocated byte arrays
             tot++;
@@ -106,12 +119,6 @@ static jl_array_t *_new_array_(jl_value_t *atype, uint32_t ndims, size_t *dims,
             tot += nel;
         }
     }
-    else {
-        wideint_t prod = (wideint_t)sizeof(void*) * (wideint_t)nel;
-        if (prod > (wideint_t) MAXINTVAL)
-            jl_error("invalid Array size");
-        tot = prod;
-    }
 
     int ndimwords = jl_array_ndimwords(ndims);
     int tsz = sizeof(jl_array_t) + ndimwords*sizeof(size_t);
@@ -196,7 +203,7 @@ static inline jl_array_t *_new_array(jl_value_t *atype, uint32_t ndims, size_t *
 jl_array_t *jl_new_array_for_deserialization(jl_value_t *atype, uint32_t ndims, size_t *dims,
                                              int isunboxed, int hasptr, int isunion, int elsz)
 {
-    return _new_array_(atype, ndims, dims, isunboxed, hasptr, isunion, 0, elsz);
+    return _new_array_(atype, ndims, dims, isunboxed, hasptr, isunion, 0, (size_t)elsz);
 }
 
 #ifndef JL_NDEBUG

src/jl_exported_funcs.inc

@@ -45,6 +45,7 @@
     XX(jl_array_to_string) \
     XX(jl_array_typetagdata) \
     XX(jl_arrayunset) \
+    XX(jl_array_validate_dims) \
     XX(jl_atexit_hook) \
     XX(jl_atomic_bool_cmpswap_bits) \
     XX(jl_atomic_cmpswap_bits) \
@@ -564,4 +565,3 @@
     YY(LLVMExtraAddGCInvariantVerifierPass) \
     YY(LLVMExtraAddDemoteFloat16Pass) \
     YY(LLVMExtraAddCPUFeaturesPass) \
-

src/julia.h

@@ -1530,6 +1530,7 @@ JL_DLLEXPORT void jl_array_sizehint(jl_array_t *a, size_t sz);
 JL_DLLEXPORT void jl_array_ptr_1d_push(jl_array_t *a, jl_value_t *item);
 JL_DLLEXPORT void jl_array_ptr_1d_append(jl_array_t *a, jl_array_t *a2);
 JL_DLLEXPORT jl_value_t *jl_apply_array_type(jl_value_t *type, size_t dim);
+JL_DLLEXPORT int jl_array_validate_dims(size_t *nel, size_t *tot, uint32_t ndims, size_t *dims, size_t elsz);
 // property access
 JL_DLLEXPORT void *jl_array_ptr(jl_array_t *a);
 JL_DLLEXPORT void *jl_array_eltype(jl_value_t *a);

test/cmdlineargs.jl

@@ -344,11 +344,7 @@ let exename = `$(Base.julia_cmd()) --startup-file=no --color=no`
         rm(memfile)
         @test popfirst!(got) == "        0 g(x) = x + 123456"
         @test popfirst!(got) == "        - function f(x)"
-        if Sys.WORD_SIZE == 64
-            @test popfirst!(got) == "       48     []"
-        else
-            @test popfirst!(got) == "       32     []"
-        end
+        @test popfirst!(got) == "        -     []"
         if Sys.WORD_SIZE == 64
             # P64 pools with 64 bit tags
             @test popfirst!(got) == "       16     Base.invokelatest(g, 0)"

test/compiler/inline.jl

@@ -151,12 +151,13 @@ end
 end
 
 function fully_eliminated(f, args)
+    @nospecialize f args
     let code = code_typed(f, args)[1][1].code
         return length(code) == 1 && isa(code[1], ReturnNode)
     end
 end
-
 function fully_eliminated(f, args, retval)
+    @nospecialize f args
     let code = code_typed(f, args)[1][1].code
         return length(code) == 1 && isa(code[1], ReturnNode) && code[1].val == retval
     end