add atomic operators for globals, memory, and setonce (#52868)

This implements `AtomicMemory`, atomic operations for globals, and the
class of atomic function for setting a field or global or memory exactly
once (setting undef => value). It is quite similar to an
`@atomicreplace`, but where there was not a previous value. This is not
needed for pointers, since it is simply calling C_NULL => pointer in
that case, since there is no "undef" value that throws after load.

plenty of future work still, for a later PR:
- syntax lowering for `@atomic global x = y` and `@atomic closedover =
y`
   - adding `Core.AtomicBox` for closure capture
   - generic functions for `atomicsetindex` & friends
   - macro for `@atomic x = y` and `@atomic x[] = y`, etc
   - design for `@atomiconce f()`

base/Base.jl

@@ -91,6 +91,36 @@ function replaceproperty!(x, f::Symbol, expected, desired, success_order::Symbol
  val = desired isa ty ? desired : convert(ty, desired)
  return Core.replacefield!(x, f, expected, val, success_order, fail_order)
 end
+function setpropertyonce!(x, f::Symbol, desired, success_order::Symbol=:not_atomic, fail_order::Symbol=success_order)
+ @inline
+ ty = fieldtype(typeof(x), f)
+ val = desired isa ty ? desired : convert(ty, desired)
+ return Core.setfieldonce!(x, f, val, success_order, fail_order)
+end
+
+function swapproperty!(x::Module, f::Symbol, v, order::Symbol=:not_atomic)
+ @inline
+ ty = Core.get_binding_type(x, f)
+ val = v isa ty ? v : convert(ty, v)
+ return Core.swapglobal!(x, f, val, order)
+end
+function modifyproperty!(x::Module, f::Symbol, op, v, order::Symbol=:not_atomic)
+ @inline
+ return Core.modifyglobal!(x, f, op, v, order)
+end
+function replaceproperty!(x::Module, f::Symbol, expected, desired, success_order::Symbol=:not_atomic, fail_order::Symbol=success_order)
+ @inline
+ ty = Core.get_binding_type(x, f)
+ val = desired isa ty ? desired : convert(ty, desired)
+ return Core.replaceglobal!(x, f, expected, val, success_order, fail_order)
+end
+function setpropertyonce!(x::Module, f::Symbol, desired, success_order::Symbol=:not_atomic, fail_order::Symbol=success_order)
+ @inline
+ ty = Core.get_binding_type(x, f)
+ val = desired isa ty ? desired : convert(ty, desired)
+ return Core.setglobalonce!(x, f, val, success_order, fail_order)
+end
+
 
 convert(::Type{Any}, Core.@nospecialize x) = x
 convert(::Type{T}, x::T) where {T} = x

base/boot.jl

@@ -65,7 +65,7 @@
 #end
 
 # struct GenericMemoryRef{kind::Symbol, T, AS::AddrSpace}
-# mem::Memory{kind, T, AS}
+# mem::GenericMemory{kind, T, AS}
 # data::Ptr{Cvoid} # make this GenericPtr{addrspace, Cvoid}
 #end
 
@@ -191,8 +191,8 @@ export
  Tuple, Type, UnionAll, TypeVar, Union, Nothing, Cvoid,
  AbstractArray, DenseArray, NamedTuple, Pair,
  # special objects
- Function, Method, Array, Memory, MemoryRef, GenericMemory, GenericMemoryRef,
- Module, Symbol, Task, UndefInitializer, undef, WeakRef, VecElement,
+ Function, Method, Module, Symbol, Task, UndefInitializer, undef, WeakRef, VecElement,
+ Array, Memory, MemoryRef, AtomicMemory, AtomicMemoryRef, GenericMemory, GenericMemoryRef,
  # numeric types
  Number, Real, Integer, Bool, Ref, Ptr,
  AbstractFloat, Float16, Float32, Float64,
@@ -209,10 +209,10 @@ export
  # AST representation
  Expr, QuoteNode, LineNumberNode, GlobalRef,
  # object model functions
- fieldtype, getfield, setfield!, swapfield!, modifyfield!, replacefield!,
+ fieldtype, getfield, setfield!, swapfield!, modifyfield!, replacefield!, setfieldonce!,
  nfields, throw, tuple, ===, isdefined, eval,
  # access to globals
- getglobal, setglobal!,
+ getglobal, setglobal!, swapglobal!, modifyglobal!, replaceglobal!, setglobalonce!,
  # ifelse, sizeof # not exported, to avoid conflicting with Base
  # type reflection
  <:, typeof, isa, typeassert,
@@ -516,22 +516,24 @@ const undef = UndefInitializer()
 (self::Type{GenericMemory{kind,T,addrspace}})(::UndefInitializer, d::NTuple{1,Int}) where {T,kind,addrspace} = self(undef, getfield(d,1))
 # empty vector constructor
 (self::Type{GenericMemory{kind,T,addrspace}})() where {T,kind,addrspace} = self(undef, 0)
-# copy constructors
 
-const Memory{T} = GenericMemory{:not_atomic, T, CPU}
-const MemoryRef{T} = GenericMemoryRef{:not_atomic, T, CPU}
 GenericMemoryRef(mem::GenericMemory) = memoryref(mem)
 GenericMemoryRef(ref::GenericMemoryRef, i::Integer) = memoryref(ref, Int(i), @_boundscheck)
 GenericMemoryRef(mem::GenericMemory, i::Integer) = memoryref(memoryref(mem), Int(i), @_boundscheck)
-MemoryRef(mem::Memory) = memoryref(mem)
-MemoryRef(ref::MemoryRef, i::Integer) = memoryref(ref, Int(i), @_boundscheck)
-MemoryRef(mem::Memory, i::Integer) = memoryref(memoryref(mem), Int(i), @_boundscheck)
-MemoryRef{T}(mem::Memory{T}) where {T} = memoryref(mem)
-MemoryRef{T}(ref::MemoryRef{T}, i::Integer) where {T} = memoryref(ref, Int(i), @_boundscheck)
-MemoryRef{T}(mem::Memory{T}, i::Integer) where {T} = memoryref(memoryref(mem), Int(i), @_boundscheck)
+GenericMemoryRef{kind,<:Any,AS}(mem::GenericMemory{kind,<:Any,AS}) where {kind,AS} = memoryref(mem)
+GenericMemoryRef{kind,<:Any,AS}(ref::GenericMemoryRef{kind,<:Any,AS}, i::Integer) where {kind,AS} = memoryref(ref, Int(i), @_boundscheck)
+GenericMemoryRef{kind,<:Any,AS}(mem::GenericMemory{kind,<:Any,AS}, i::Integer) where {kind,AS} = memoryref(memoryref(mem), Int(i), @_boundscheck)
+GenericMemoryRef{kind,T,AS}(mem::GenericMemory{kind,T,AS}) where {kind,T,AS} = memoryref(mem)
+GenericMemoryRef{kind,T,AS}(ref::GenericMemoryRef{kind,T,AS}, i::Integer) where {kind,T,AS} = memoryref(ref, Int(i), @_boundscheck)
+GenericMemoryRef{kind,T,AS}(mem::GenericMemory{kind,T,AS}, i::Integer) where {kind,T,AS} = memoryref(memoryref(mem), Int(i), @_boundscheck)
+
+const Memory{T} = GenericMemory{:not_atomic, T, CPU}
+const MemoryRef{T} = GenericMemoryRef{:not_atomic, T, CPU}
+const AtomicMemory{T} = GenericMemory{:atomic, T, CPU}
+const AtomicMemoryRef{T} = GenericMemoryRef{:atomic, T, CPU}
 
 # construction helpers for Array
-new_as_memoryref(self::Type{GenericMemoryRef{isatomic,T,addrspace}}, m::Int) where {T,isatomic,addrspace} = memoryref(fieldtype(self, :mem)(undef, m))
+new_as_memoryref(self::Type{GenericMemoryRef{kind,T,addrspace}}, m::Int) where {T,kind,addrspace} = memoryref(fieldtype(self, :mem)(undef, m))
 
 # checked-multiply intrinsic function for dimensions
 _checked_mul_dims() = 1, false

base/compiler/abstractinterpretation.jl

@@ -1113,12 +1113,12 @@ function const_prop_function_heuristic(interp::AbstractInterpreter, @nospecializ
  if !still_nothrow || ismutabletype(arrty)
  return false
  end
- elseif ⊑(𝕃ᵢ, arrty, Array)
+ elseif ⊑(𝕃ᵢ, arrty, Array) || ⊑(𝕃ᵢ, arrty, GenericMemory)
  return false
  end
  elseif istopfunction(f, :iterate)
  itrty = argtypes[2]
- if ⊑(𝕃ᵢ, itrty, Array)
+ if ⊑(𝕃ᵢ, itrty, Array) || ⊑(𝕃ᵢ, itrty, GenericMemory)
  return false
  end
  end
@@ -1501,7 +1501,7 @@ function abstract_iteration(interp::AbstractInterpreter, @nospecialize(itft), @n
  # Return Bottom if this is not an iterator.
  # WARNING: Changes to the iteration protocol must be reflected here,
  # this is not just an optimization.
- # TODO: this doesn't realize that Array, SimpleVector, Tuple, and NamedTuple do not use the iterate protocol
+ # TODO: this doesn't realize that Array, GenericMemory, SimpleVector, Tuple, and NamedTuple do not use the iterate protocol
  stateordonet === Bottom && return AbstractIterationResult(Any[Bottom], AbstractIterationInfo(CallMeta[CallMeta(Bottom, Any, call.effects, info)], true))
  valtype = statetype = Bottom
  ret = Any[]
@@ -2076,8 +2076,8 @@ function abstract_call_known(interp::AbstractInterpreter, @nospecialize(f),
  return abstract_apply(interp, argtypes, si, sv, max_methods)
  elseif f === invoke
  return abstract_invoke(interp, arginfo, si, sv)
- elseif f === modifyfield!
- return abstract_modifyfield!(interp, argtypes, si, sv)
+ elseif f === modifyfield! || f === Core.modifyglobal! || f === Core.memoryrefmodify! || f === atomic_pointermodify
+ return abstract_modifyop!(interp, f, argtypes, si, sv)
  elseif f === Core.finalizer
  return abstract_finalizer(interp, argtypes, sv)
  elseif f === applicable

base/compiler/ssair/inlining.jl

@@ -1286,10 +1286,18 @@ function process_simple!(todo::Vector{Pair{Int,Any}}, ir::IRCode, idx::Int, stat
  end
  end
 
- if (sig.f !== Core.invoke && sig.f !== Core.finalizer && sig.f !== modifyfield!) &&
- is_builtin(optimizer_lattice(state.interp), sig)
- # No inlining for builtins (other invoke/apply/typeassert/finalizer)
- return nothing
+ if is_builtin(optimizer_lattice(state.interp), sig)
+ let f = sig.f
+ if (f !== Core.invoke &&
+ f !== Core.finalizer &&
+ f !== modifyfield! &&
+ f !== Core.modifyglobal! &&
+ f !== Core.memoryrefmodify! &&
+ f !== atomic_pointermodify)
+ # No inlining defined for most builtins (just invoke/apply/typeassert/finalizer), so attempt an early exit for them
+ return nothing
+ end
+ end
  end
 
  # Special case inliners for regular functions
@@ -1571,7 +1579,7 @@ function handle_opaque_closure_call!(todo::Vector{Pair{Int,Any}},
  return nothing
 end
 
-function handle_modifyfield!_call!(ir::IRCode, idx::Int, stmt::Expr, info::ModifyFieldInfo, state::InliningState)
+function handle_modifyop!_call!(ir::IRCode, idx::Int, stmt::Expr, info::ModifyOpInfo, state::InliningState)
  info = info.info
  info isa MethodResultPure && (info = info.info)
  info isa ConstCallInfo && (info = info.call)
@@ -1679,8 +1687,8 @@ function assemble_inline_todo!(ir::IRCode, state::InliningState)
  # handle special cased builtins
  if isa(info, OpaqueClosureCallInfo)
  handle_opaque_closure_call!(todo, ir, idx, stmt, info, flag, sig, state)
- elseif isa(info, ModifyFieldInfo)
- handle_modifyfield!_call!(ir, idx, stmt, info, state)
+ elseif isa(info, ModifyOpInfo)
+ handle_modifyop!_call!(ir, idx, stmt, info, state)
  elseif isa(info, InvokeCallInfo)
  handle_invoke_call!(todo, ir, idx, stmt, info, flag, sig, state)
  elseif isa(info, FinalizerInfo)

base/compiler/stmtinfo.jl

@@ -222,13 +222,18 @@ struct FinalizerInfo <: CallInfo
 end
 
 """
- info::ModifyFieldInfo <: CallInfo
+ info::ModifyOpInfo <: CallInfo
 
-Represents a resolved all of `modifyfield!(obj, name, op, x, [order])`.
-`info.info` wraps the call information of `op(getfield(obj, name), x)`.
+Represents a resolved call of one of:
+ - `modifyfield!(obj, name, op, x, [order])`
+ - `modifyglobal!(mod, var, op, x, order)`
+ - `memoryrefmodify!(memref, op, x, order, boundscheck)`
+ - `Intrinsics.atomic_pointermodify(ptr, op, x, order)`
+
+`info.info` wraps the call information of `op(getval(), x)`.
 """
-struct ModifyFieldInfo <: CallInfo
- info::CallInfo # the callinfo for the `op(getfield(obj, name), x)` call
+struct ModifyOpInfo <: CallInfo
+ info::CallInfo # the callinfo for the `op(getval(), x)` call
 end
 
 @specialize