Implement ImmutableArray

This rebases #31630 with several fixed and modifications.
After #31630, we had originally decided to hold off on said
PR in favor of implementing either more efficient layouts for
tuples or some sort of variable-sized struct type. However, in
the two years since, neither of those have happened (I had a go
at improving tuples and made some progress, but there is much
still to be done there). In the meantime, all across the package
ecosystem, we've seen an increasing creep of pre-allocation and
mutating operations, primarily caused by our lack of sufficiently
powerful immutable array abstractions and array optimizations.

This works fine for the individual packages in question, but it
causes a fair bit of trouble when trying to compose these packages
with transformation passes such as AD or domain specific optimizations,
since many of those passes do not play well with mutation. More
generally, we would like to avoid people needing to pierce
abstractions for performance reasons.

Given these developments, I think it's getting quite important
that we start to seriously look at arrays and try to provide
performant and well-optimized arrays in the language. More
importantly, I think this is somewhat independent from the
actual implementation details. To be sure, it would be nice
to move more of the array implementation into Julia by making
use of one of the abovementioned langugage features, but that
is a bit of an orthogonal concern and not absolutely required.

This PR provides an `ImmutableArray` type that is identical
in functionality and implementation to `Array`, except that
it is immutable. Two new intrinsics `Core.arrayfreeze` and
`Core.arraythaw` are provided which are semantically copies
and turn a mutable array into an immutable array and vice
versa.

In the original PR, I additionally provided generic functions
`freeze` and `thaw` that would simply forward to these
intrinsics. However, said generic functions have been omitted
from this PR in favor of simply using constructors to go
between mutable and immutable arrays at the high level.
Generic `freeze`/`thaw` functions can always be added later,
once we have a more complete picture of how these functions
would work on non-Array datatypes.

Some basic compiler support is provided to elide these copies
when the compiler can prove that the original object is
dead after the copy. For instance, in the following example:
```
function simple()
    a = Vector{Float64}(undef, 5)
    for i = 1:5
        a[i] = i
    end
    ImmutableArray(a)
end
```

the compiler will recognize that the array `a` is dead after
its use in `ImmutableArray` and the optimized implementation
will simply rewrite the type tag in the originally allocated
array to now mark it as immutable. It should be pointed out
however, that *semantically* there is still no mutation of the
original array, this is simply an optimization.

At the moment this compiler transform is rather limited, since
the analysis requires escape information in order to compute
whether or not the copy may be elided. However, more complete
escape analysis is being worked on at the moment, so hopefully
this analysis should become more powerful in the very near future.

I would like to get this cleaned up and merged resonably quickly,
and then crowdsource some improvements to the Array APIs more
generally. There are still a number of APIs that are quite bound
to the notion of mutable `Array`s. StaticArrays and other packages
have been inventing conventions for how to generalize those, but
we should form a view in Base what those APIs should look like and
harmonize them. Having the `ImmutableArray` in Base should help
with that.

Author: Keno

base/array.jl

@@ -147,12 +147,20 @@ function vect(X...)
     return copyto!(Vector{T}(undef, length(X)), X)
 end
 
-size(a::Array, d::Integer) = arraysize(a, convert(Int, d))
-size(a::Vector) = (arraysize(a,1),)
-size(a::Matrix) = (arraysize(a,1), arraysize(a,2))
-size(a::Array{<:Any,N}) where {N} = (@_inline_meta; ntuple(M -> size(a, M), Val(N))::Dims)
+const ImmutableArray = Core.ImmutableArray
+const IMArray{T,N} = Union{Array{T, N}, ImmutableArray{T,N}}
+const IMVector{T} = IMArray{T, 1}
+const IMMatrix{T} = IMArray{T, 2}
 
-asize_from(a::Array, n) = n > ndims(a) ? () : (arraysize(a,n), asize_from(a, n+1)...)
+ImmutableArray(a::Array) = Core.arrayfreeze(a)
+Array(a::ImmutableArray) = Core.arraythaw(a)
+
+size(a::IMArray, d::Integer) = arraysize(a, convert(Int, d))
+size(a::IMVector) = (arraysize(a,1),)
+size(a::IMMatrix) = (arraysize(a,1), arraysize(a,2))
+size(a::IMArray{<:Any,N}) where {N} = (@_inline_meta; ntuple(M -> size(a, M), Val(N))::Dims)
+
+asize_from(a::IMArray, n) = n > ndims(a) ? () : (arraysize(a,n), asize_from(a, n+1)...)
 
 allocatedinline(T::Type) = (@_pure_meta; ccall(:jl_stored_inline, Cint, (Any,), T) != Cint(0))
 
@@ -223,6 +231,13 @@ function isassigned(a::Array, i::Int...)
     ccall(:jl_array_isassigned, Cint, (Any, UInt), a, ii) == 1
 end
 
+function isassigned(a::ImmutableArray, i::Int...)
+    @_inline_meta
+    ii = (_sub2ind(size(a), i...) % UInt) - 1
+    @boundscheck ii < length(a) % UInt || return false
+    ccall(:jl_array_isassigned, Cint, (Any, UInt), a, ii) == 1
+end
+
 ## copy ##
 
 """
@@ -895,6 +910,9 @@ function getindex end
 @eval getindex(A::Array, i1::Int) = arrayref($(Expr(:boundscheck)), A, i1)
 @eval getindex(A::Array, i1::Int, i2::Int, I::Int...) = (@_inline_meta; arrayref($(Expr(:boundscheck)), A, i1, i2, I...))
 
+@eval getindex(A::ImmutableArray, i1::Int) = arrayref($(Expr(:boundscheck)), A, i1)
+@eval getindex(A::ImmutableArray, i1::Int, i2::Int, I::Int...) = (@_inline_meta; arrayref($(Expr(:boundscheck)), A, i1, i2, I...))
+
 # Faster contiguous indexing using copyto! for UnitRange and Colon
 function getindex(A::Array, I::AbstractUnitRange{<:Integer})
     @_inline_meta

base/compiler/optimize.jl

@@ -307,7 +307,7 @@ function run_passes(ci::CodeInfo, sv::OptimizationState)
     ir = adce_pass!(ir)
     #@Base.show ("after_adce", ir)
     @timeit "type lift" ir = type_lift_pass!(ir)
-    @timeit "compact 3" ir = compact!(ir)
+    ir = memory_opt!(ir)
     #@Base.show ir
     if JLOptions().debug_level == 2
         @timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))

base/compiler/ssair/ir.jl

@@ -319,6 +319,13 @@ function setindex!(x::IRCode, @nospecialize(repl), s::SSAValue)
     return x
 end
 
+function ssadominates(ir::IRCode, domtree::DomTree, ssa1::Int, ssa2::Int)
+    bb1 = block_for_inst(ir.cfg, ssa1)
+    bb2 = block_for_inst(ir.cfg, ssa2)
+    bb1 == bb2 && return ssa1 < ssa2
+    return dominates(domtree, bb1, bb2)
+end
+
 # SSA values that need renaming
 struct OldSSAValue
     id::Int

base/compiler/ssair/passes.jl

@@ -1255,3 +1255,77 @@ function cfg_simplify!(ir::IRCode)
     compact.active_result_bb = length(bb_starts)
     return finish(compact)
 end
+
+function is_allocation(stmt)
+    isexpr(stmt, :foreigncall) || return false
+    s = stmt.args[1]
+    isa(s, QuoteNode) && (s = s.value)
+    return s === :jl_alloc_array_1d
+end
+
+function memory_opt!(ir::IRCode)
+    compact = IncrementalCompact(ir, false)
+    uses = IdDict{Int, Vector{Int}}()
+    relevant = IdSet{Int}()
+    revisit = Int[]
+    function mark_val(val)
+        isa(val, SSAValue) || return
+        val.id in relevant && pop!(relevant, val.id)
+    end
+    for ((_, idx), stmt) in compact
+        if isa(stmt, ReturnNode)
+            isdefined(stmt, :val) || continue
+            val = stmt.val
+            if isa(val, SSAValue) && val.id in relevant
+                (haskey(uses, val.id)) || (uses[val.id] = Int[])
+                push!(uses[val.id], idx)
+            end
+            continue
+        end
+        (isexpr(stmt, :call) || isexpr(stmt, :foreigncall)) || continue
+        if is_allocation(stmt)
+            push!(relevant, idx)
+            # TODO: Mark everything else here
+            continue
+        end
+        # TODO: Replace this by interprocedural escape analysis
+        if is_known_call(stmt, arrayset, compact)
+            # The value being set escapes, everything else doesn't
+            mark_val(stmt.args[4])
+            arr = stmt.args[3]
+            if isa(arr, SSAValue) && arr.id in relevant
+                (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+                push!(uses[arr.id], idx)
+            end
+        elseif is_known_call(stmt, Core.arrayfreeze, compact) && isa(stmt.args[2], SSAValue)
+            push!(revisit, idx)
+        else
+            # For now we assume everything escapes
+            # TODO: We could handle PhiNodes specially and improve this
+            for ur in userefs(stmt)
+                mark_val(ur[])
+            end
+        end
+    end
+    ir = finish(compact)
+    isempty(revisit) && return ir
+    domtree = construct_domtree(ir.cfg.blocks)
+    for idx in revisit
+        # Make sure that the value we reference didn't escape
+        id = ir.stmts[idx][:inst].args[2].id
+        (id in relevant) || continue
+
+        # We're ok to steal the memory if we don't dominate any uses
+        ok = true
+        for use in uses[id]
+            if ssadominates(ir, domtree, idx, use)
+                ok = false
+                break
+            end
+        end
+        ok || continue
+
+        ir.stmts[idx][:inst].args[1] = Core.mutating_arrayfreeze
+    end
+    return ir
+end

base/compiler/tfuncs.jl

@@ -1532,6 +1532,21 @@ function builtin_tfunction(interp::AbstractInterpreter, @nospecialize(f), argtyp
                            sv::Union{InferenceState,Nothing})
     if f === tuple
         return tuple_tfunc(argtypes)
+    elseif f === Core.arrayfreeze || f === Core.arraythaw
+        if length(argtypes) != 1
+            isva && return Any
+            return Bottom
+        end
+        a = widenconst(argtypes[1])
+        at = (f === Core.arrayfreeze ? Array : ImmutableArray)
+        rt = (f === Core.arrayfreeze ? ImmutableArray : Array)
+        if a <: at
+            unw = unwrap_unionall(a)
+            if isa(unw, DataType)
+                return rewrap_unionall(rt{unw.parameters[1], unw.parameters[2]}, a)
+            end
+        end
+        return rt
     end
     if isa(f, IntrinsicFunction)
         if is_pure_intrinsic_infer(f) && _all(@nospecialize(a) -> isa(a, Const), argtypes)

base/dict.jl

@@ -372,7 +372,7 @@ end
 function setindex!(h::Dict{K,V}, v0, key0) where V where K
     key = convert(K, key0)
     if !isequal(key, key0)
-        throw(ArgumentError("$(limitrepr(key0)) is not a valid key for type $K"))
+        throw(KeyTypeError(K, key0))
     end
     setindex!(h, v0, key)
 end

base/experimental.jl

@@ -11,6 +11,8 @@ module Experimental
 
 using Base: Threads, sync_varname
 using Base.Meta
+using Base: ImmutableArray
+
 
 """
     Const(A::Array)

src/Makefile

@@ -261,7 +261,7 @@ $(BUILDDIR)/interpreter.o $(BUILDDIR)/interpreter.dbg.obj: $(SRCDIR)/builtin_pro
 $(BUILDDIR)/jitlayers.o $(BUILDDIR)/jitlayers.dbg.obj: $(SRCDIR)/jitlayers.h $(SRCDIR)/codegen_shared.h
 $(BUILDDIR)/jltypes.o $(BUILDDIR)/jltypes.dbg.obj: $(SRCDIR)/builtin_proto.h
 $(build_shlibdir)/libllvmcalltest.$(SHLIB_EXT): $(SRCDIR)/codegen_shared.h $(BUILDDIR)/julia_version.h
-$(BUILDDIR)/llvm-alloc-opt.o $(BUILDDIR)/llvm-alloc-opt.dbg.obj: $(SRCDIR)/codegen_shared.h
+$(BUILDDIR)/llvm-alloc-opt.o $(BUILDDIR)/llvm-alloc-opt.dbg.obj: $(SRCDIR)/codegen_shared.h $(SRCDIR)/llvm-pass-helpers.h
 $(BUILDDIR)/llvm-final-gc-lowering.o $(BUILDDIR)/llvm-final-gc-lowering.dbg.obj: $(SRCDIR)/llvm-pass-helpers.h
 $(BUILDDIR)/llvm-gc-invariant-verifier.o $(BUILDDIR)/llvm-gc-invariant-verifier.dbg.obj: $(SRCDIR)/codegen_shared.h
 $(BUILDDIR)/llvm-late-gc-lowering.o $(BUILDDIR)/llvm-late-gc-lowering.dbg.obj: $(SRCDIR)/llvm-pass-helpers.h

src/builtin_proto.h

@@ -51,6 +51,9 @@ DECLARE_BUILTIN(typeassert);
 DECLARE_BUILTIN(_typebody);
 DECLARE_BUILTIN(typeof);
 DECLARE_BUILTIN(_typevar);
+DECLARE_BUILTIN(arrayfreeze);
+DECLARE_BUILTIN(arraythaw);
+DECLARE_BUILTIN(mutating_arrayfreeze);
 
 JL_CALLABLE(jl_f_invoke_kwsorter);
 JL_CALLABLE(jl_f__structtype);

src/builtins.c

@@ -1330,7 +1330,9 @@ JL_CALLABLE(jl_f__typevar)
 JL_CALLABLE(jl_f_arraysize)
 {
     JL_NARGS(arraysize, 2, 2);
-    JL_TYPECHK(arraysize, array, args[0]);
+    if (!jl_is_arrayish(args[0])) {
+        jl_type_error("arraysize", (jl_value_t*)jl_array_type, args[0]);
+    }
     jl_array_t *a = (jl_array_t*)args[0];
     size_t nd = jl_array_ndims(a);
     JL_TYPECHK(arraysize, long, args[1]);
@@ -1369,7 +1371,9 @@ JL_CALLABLE(jl_f_arrayref)
 {
     JL_NARGSV(arrayref, 3);
     JL_TYPECHK(arrayref, bool, args[0]);
-    JL_TYPECHK(arrayref, array, args[1]);
+    if (!jl_is_arrayish(args[1])) {
+        jl_type_error("arrayref", (jl_value_t*)jl_array_type, args[1]);
+    }
     jl_array_t *a = (jl_array_t*)args[1];
     size_t i = array_nd_index(a, &args[2], nargs - 2, "arrayref");
     return jl_arrayref(a, i);
@@ -1645,6 +1649,54 @@ JL_CALLABLE(jl_f__equiv_typedef)
     return equiv_type(args[0], args[1]) ? jl_true : jl_false;
 }
 
+JL_CALLABLE(jl_f_arrayfreeze)
+{
+    JL_NARGSV(arrayfreeze, 1);
+    JL_TYPECHK(arrayfreeze, array, args[0]);
+    jl_array_t *a = (jl_array_t*)args[0];
+    jl_datatype_t *it = (jl_datatype_t *)jl_apply_type2((jl_value_t*)jl_immutable_array_type,
+        jl_tparam0(jl_typeof(a)), jl_tparam1(jl_typeof(a)));
+    JL_GC_PUSH1(&it);
+    // The idea is to elide this copy if the compiler or runtime can prove that
+    // doing so is safe to do.
+    jl_array_t *na = jl_array_copy(a);
+    jl_set_typeof(na, it);
+    JL_GC_POP();
+    return (jl_value_t*)na;
+}
+
+JL_CALLABLE(jl_f_mutating_arrayfreeze)
+{
+    // N.B.: These error checks pretend to be arrayfreeze since this is a drop
+    // in replacement and we don't want to change the visible error type in the
+    // optimizer
+    JL_NARGSV(arrayfreeze, 1);
+    JL_TYPECHK(arrayfreeze, array, args[0]);
+    jl_array_t *a = (jl_array_t*)args[0];
+    jl_datatype_t *it = (jl_datatype_t *)jl_apply_type2((jl_value_t*)jl_immutable_array_type,
+        jl_tparam0(jl_typeof(a)), jl_tparam1(jl_typeof(a)));
+    jl_set_typeof(a, it);
+    return (jl_value_t*)a;
+}
+
+JL_CALLABLE(jl_f_arraythaw)
+{
+    JL_NARGSV(arraythaw, 1);
+    if (((jl_datatype_t*)jl_typeof(args[0]))->name != jl_immutable_array_typename) {
+        jl_type_error("arraythaw", (jl_value_t*)jl_immutable_array_type, args[0]);
+    }
+    jl_array_t *a = (jl_array_t*)args[0];
+    jl_datatype_t *it = (jl_datatype_t *)jl_apply_type2((jl_value_t*)jl_array_type,
+        jl_tparam0(jl_typeof(a)), jl_tparam1(jl_typeof(a)));
+    JL_GC_PUSH1(&it);
+    // The idea is to elide this copy if the compiler or runtime can prove that
+    // doing so is safe to do.
+    jl_array_t *na = jl_array_copy(a);
+    jl_set_typeof(na, it);
+    JL_GC_POP();
+    return (jl_value_t*)na;
+}
+
 // IntrinsicFunctions ---------------------------------------------------------
 
 static void (*runtime_fp[num_intrinsics])(void);
@@ -1797,6 +1849,10 @@ void jl_init_primitives(void) JL_GC_DISABLED
     jl_builtin_arrayset = add_builtin_func("arrayset", jl_f_arrayset);
     jl_builtin_arraysize = add_builtin_func("arraysize", jl_f_arraysize);
 
+    jl_builtin_arrayfreeze = add_builtin_func("arrayfreeze", jl_f_arrayfreeze);
+    jl_builtin_mutating_arrayfreeze = add_builtin_func("mutating_arrayfreeze", jl_f_mutating_arrayfreeze);
+    jl_builtin_arraythaw = add_builtin_func("arraythaw", jl_f_arraythaw);
+
     // method table utils
     jl_builtin_applicable = add_builtin_func("applicable", jl_f_applicable);
     jl_builtin_invoke = add_builtin_func("invoke", jl_f_invoke);
@@ -1868,6 +1924,7 @@ void jl_init_primitives(void) JL_GC_DISABLED
     add_builtin("AbstractArray", (jl_value_t*)jl_abstractarray_type);
     add_builtin("DenseArray", (jl_value_t*)jl_densearray_type);
     add_builtin("Array", (jl_value_t*)jl_array_type);
+    add_builtin("ImmutableArray", (jl_value_t*)jl_immutable_array_type);
 
     add_builtin("Expr", (jl_value_t*)jl_expr_type);
     add_builtin("LineNumberNode", (jl_value_t*)jl_linenumbernode_type);