Rework cat implementation + add zero! (#18)

This PR enables a proper implementation/specialization point for implementing concatenations.
Simultaneously, we define an overloadable `zero!` function.

Author: lkdvos

.github/workflows/IntegrationTestRequest.yml

@@ -0,0 +1,14 @@
+name: "Integration Test Request"
+
+on:
+  issue_comment:
+    types: [created]
+
+jobs:
+  integrationrequest:
+    if: |
+      github.event.issue.pull_request &&
+      contains(fromJSON('["OWNER", "COLLABORATOR", "MEMBER"]'), github.event.comment.author_association)
+    uses: ITensor/ITensorActions/.github/workflows/IntegrationTestRequest.yml@main
+    with:
+      localregistry: https://github.com/ITensor/ITensorRegistry.git

Project.toml

@@ -1,11 +1,12 @@
 name = "DerivableInterfaces"
 uuid = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.3.17"
+version = "0.4.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ExproniconLite = "55351af7-c7e9-48d6-89ff-24e801d99491"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
@@ -15,6 +16,7 @@ TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 [compat]
 Adapt = "4.1.1"
 ArrayLayouts = "1.11.0"
+Compat = "3.47,4.10"
 ExproniconLite = "0.10.13"
 LinearAlgebra = "1.10"
 MLStyle = "0.4.17"

docs/Project.toml

@@ -4,6 +4,6 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 
 [compat]
-DerivableInterfaces = "0.3"
+DerivableInterfaces = "0.4"
 Documenter = "1"
 Literate = "2"

docs/make.jl

@@ -16,7 +16,7 @@ makedocs(;
     edit_link="main",
     assets=String[],
   ),
-  pages=["Home" => "index.md"],
+  pages=["Home" => "index.md", "Reference" => "reference.md"],
 )
 
 deploydocs(;

docs/src/reference.md

@@ -0,0 +1,5 @@
+# Reference
+
+```@autodocs
+Modules = [DerivableInterfaces, DerivableInterfaces.Concatenate]
+```

src/DerivableInterfaces.jl

@@ -1,11 +1,16 @@
 module DerivableInterfaces
 
+export Concatenate
+
 include("interface_function.jl")
 include("abstractinterface.jl")
 include("derive_macro.jl")
 include("interface_macro.jl")
 include("wrappedarrays.jl")
+
+include("zero.jl")
 include("abstractarrayinterface.jl")
+include("concatenate.jl")
 include("defaultarrayinterface.jl")
 include("traits.jl")
 

src/abstractarrayinterface.jl

@@ -153,21 +153,12 @@ end
   @interface interface map!(Returns(value), a, a)
 end
 
-using ArrayLayouts: zero!
+# TODO: should this be recursive? `map!(zero!, A, A)` might also work?
+@interface ::AbstractArrayInterface DerivableInterfaces.zero!(A::AbstractArray) =
+  fill!(A, zero(eltype(A)))
 
-# `zero!` isn't defined in `Base`, but it is defined in `ArrayLayouts`
-# and is useful for sparse array logic, since it can be used to empty
-# the sparse array storage.
-# We use a single function definition to minimize method ambiguities.
-@interface interface::AbstractArrayInterface function ArrayLayouts.zero!(a::AbstractArray)
-  # More generally, the first codepath could be taking if `zero(eltype(a))`
-  # is defined and the elements are immutable.
-  f = eltype(a) <: Number ? Returns(zero(eltype(a))) : zero!
-  return @interface interface map!(f, a, a)
-end
-
-# Specialized version of `Base.zero` written in terms of `ArrayLayouts.zero!`.
-# This is friendlier for sparse arrays since `ArrayLayouts.zero!` makes it easier
+# Specialized version of `Base.zero` written in terms of `zero!`.
+# This is friendlier for sparse arrays since `zero!` makes it easier
 # to handle the logic of dropping all elements of the sparse array when possible.
 # We use a single function definition to minimize method ambiguities.
 @interface interface::AbstractArrayInterface function Base.zero(a::AbstractArray)
@@ -250,102 +241,3 @@ end
 ## @interface ::AbstractMatrixInterface function Base.*(a1, a2)
 ##   return ArrayLayouts.mul(a1, a2)
 ## end
-
-# Concatenation
-
-axis_cat(a1::AbstractUnitRange, a2::AbstractUnitRange) = Base.OneTo(length(a1) + length(a2))
-function axis_cat(
-  a1::AbstractUnitRange, a2::AbstractUnitRange, a_rest::AbstractUnitRange...
-)
-  return axis_cat(axis_cat(a1, a2), a_rest...)
-end
-
-unval(x) = x
-unval(::Val{x}) where {x} = x
-
-function cat_axes(as::AbstractArray...; dims)
-  return ntuple(length(first(axes.(as)))) do dim
-    return if dim in unval(dims)
-      axis_cat(map(axes -> axes[dim], axes.(as))...)
-    else
-      axes(first(as))[dim]
-    end
-  end
-end
-
-function cat! end
-
-# Represents concatenating `args` over `dims`.
-struct Cat{Args<:Tuple{Vararg{AbstractArray}},dims}
-  args::Args
-end
-to_cat_dims(dim::Integer) = Int(dim)
-to_cat_dims(dim::Int) = (dim,)
-to_cat_dims(dims::Val) = to_cat_dims(unval(dims))
-to_cat_dims(dims::Tuple) = dims
-Cat(args::AbstractArray...; dims) = Cat{typeof(args),to_cat_dims(dims)}(args)
-cat_dims(::Cat{<:Any,dims}) where {dims} = dims
-
-function Base.axes(a::Cat)
-  return cat_axes(a.args...; dims=cat_dims(a))
-end
-Base.eltype(a::Cat) = promote_type(eltype.(a.args)...)
-function Base.similar(a::Cat)
-  ax = axes(a)
-  elt = eltype(a)
-  # TODO: This drops GPU information, maybe use MemoryLayout?
-  return similar(arraytype(interface(a.args...), elt), ax)
-end
-
-# https://github.com/JuliaLang/julia/blob/v1.11.1/base/abstractarray.jl#L1748-L1857
-# https://docs.julialang.org/en/v1/base/arrays/#Concatenation-and-permutation
-# This is very similar to the `Base.cat` implementation but handles zero values better.
-function cat_offset!(
-  a_dest::AbstractArray, offsets, a1::AbstractArray, a_rest::AbstractArray...; dims
-)
-  inds = ntuple(ndims(a_dest)) do dim
-    dim in unval(dims) ? offsets[dim] .+ axes(a1, dim) : axes(a_dest, dim)
-  end
-  a_dest[inds...] = a1
-  new_offsets = ntuple(ndims(a_dest)) do dim
-    dim in unval(dims) ? offsets[dim] + size(a1, dim) : offsets[dim]
-  end
-  cat_offset!(a_dest, new_offsets, a_rest...; dims)
-  return a_dest
-end
-function cat_offset!(a_dest::AbstractArray, offsets; dims)
-  return a_dest
-end
-
-@interface ::AbstractArrayInterface function cat!(
-  a_dest::AbstractArray, as::AbstractArray...; dims
-)
-  offsets = ntuple(zero, ndims(a_dest))
-  # TODO: Fill `a_dest` with zeros if needed using `zero!`.
-  cat_offset!(a_dest, offsets, as...; dims)
-  return a_dest
-end
-
-function cat_along(dims, as::AbstractArray...)
-  return @interface interface(as...) cat_along(dims, as...)
-end
-
-@interface interface::AbstractArrayInterface function cat_along(dims, as::AbstractArray...)
-  a_dest = similar(Cat(as...; dims))
-  @interface interface cat!(a_dest, as...; dims)
-  return a_dest
-end
-
-@interface interface::AbstractArrayInterface function Base.cat(as::AbstractArray...; dims)
-  return @interface interface cat_along(dims, as...)
-end
-
-# TODO: Use `@derive` instead:
-# ```julia
-# @derive (T=AbstractArray,) begin
-#   cat!(a_dest::AbstractArray, as::T...; dims)
-# end
-# ```
-function cat!(a_dest::AbstractArray, as::AbstractArray...; dims)
-  return @interface interface(as...) cat!(a_dest, as...; dims)
-end

src/concatenate.jl

@@ -0,0 +1,140 @@
+"""
+    module Concatenate
+
+Alternative implementation for `Base.cat` through [`cat(!)`](@ref cat).
+
+This is mostly a copy of the Base implementation, with the main difference being
+that the destination is chosen based on all inputs instead of just the first.
+
+Additionally, we have an intermediate representation in terms of a Concatenated object,
+reminiscent of how Broadcast works.
+
+The various entry points for specializing behavior are:
+
+* Destination selection can be achieved through
+
+    Base.similar(concat::Concatenated{Interface}, ::Type{T}, axes) where {Interface}
+
+* Custom implementations:
+
+    Base.copy(concat::Concatenated{Interface}) # custom implementation of cat
+    Base.copyto!(dest, concat::Concatenated{Interface}) # custom implementation of cat! based on interface
+    Base.copyto!(dest, concat::Concatenated{Nothing}) # custom implementation of cat! based on typeof(dest)
+"""
+module Concatenate
+
+using Compat: @compat
+export concatenate
+@compat public Concatenated, cat, cat!, concatenated
+
+using Base: promote_eltypeof
+using ..DerivableInterfaces:
+  DerivableInterfaces, AbstractInterface, interface, zero!, arraytype
+
+"""
+    Concatenated{Interface,Dims,Args<:Tuple}
+
+Lazy representation of the concatenation of various `Args` along `Dims`, in order to provide
+hooks to customize the implementation.
+"""
+struct Concatenated{Interface,Dims,Args<:Tuple}
+  interface::Interface
+  dims::Val{Dims}
+  args::Args
+
+  function Concatenated(
+    interface::Union{Nothing,AbstractInterface}, dims::Val{Dims}, args::Tuple
+  ) where {Dims}
+    return new{typeof(interface),Dims,typeof(args)}(interface, dims, args)
+  end
+  function Concatenated(dims, args::Tuple)
+    return Concatenated(interface(args...), dims, args)
+  end
+  function Concatenated{Interface}(dims, args) where {Interface}
+    return Concatenated(Interface(), dims, args)
+  end
+  function Concatenated{Interface,Dims}(args) where {Interface,Dims}
+    return new{Interface,Dims,typeof(args)}(Interface(), Val(Dims), args)
+  end
+end
+
+dims(::Concatenated{A,D}) where {A,D} = D
+DerivableInterfaces.interface(concat::Concatenated) = concat.interface
+
+concatenated(dims, args...) = concatenated(Val(dims), args...)
+concatenated(dims::Val, args...) = Concatenated(dims, args)
+
+function Base.convert(
+  ::Type{Concatenated{NewInterface}}, concat::Concatenated{<:Any,Dims,Args}
+) where {NewInterface,Dims,Args}
+  return Concatenated{NewInterface}(
+    concat.dims, concat.args
+  )::Concatenated{NewInterface,Dims,Args}
+end
+
+# allocating the destination container
+# ------------------------------------
+Base.similar(concat::Concatenated) = similar(concat, eltype(concat))
+Base.similar(concat::Concatenated, ::Type{T}) where {T} = similar(concat, T, axes(concat))
+function Base.similar(concat::Concatenated, ::Type{T}, ax) where {T}
+  return similar(arraytype(interface(concat), T), ax)
+end
+
+Base.eltype(concat::Concatenated) = promote_eltypeof(concat.args...)
+
+# For now, simply couple back to base implementation
+function Base.axes(concat::Concatenated)
+  catdims = Base.dims2cat(dims(concat))
+  return Base.cat_size_shape(catdims, concat.args...)
+end
+
+# Main logic
+# ----------
+"""
+    concatenate(dims, args...)
+
+Concatenate the supplied `args` along dimensions `dims`.
+
+See also [`cat`] and [`cat!`](@ref).
+"""
+concatenate(dims, args...) = Base.materialize(concatenated(dims, args...))
+
+"""
+    Concatenate.cat(args...; dims)
+
+Concatenate the supplied `args` along dimensions `dims`.
+
+See also [`concatenate`] and [`cat!`](@ref).
+"""
+cat(args...; dims) = concatenate(dims, args...)
+Base.materialize(concat::Concatenated) = copy(concat)
+
+"""
+    Concatenate.cat!(dest, args...; dims)
+
+Concatenate the supplied `args` along dimensions `dims`, placing the result into `dest`.
+"""
+function cat!(dest, args...; dims)
+  Base.materialize!(dest, concatenated(dims, args...))
+  return dest
+end
+Base.materialize!(dest, concat::Concatenated) = copyto!(dest, concat)
+
+Base.copy(concat::Concatenated) = copyto!(similar(concat), concat)
+
+# default falls back to replacing interface with Nothing
+# this permits specializing on typeof(dest) without ambiguities
+# Note: this needs to be defined for AbstractArray specifically to avoid ambiguities with Base.
+@inline Base.copyto!(dest::AbstractArray, concat::Concatenated) =
+  copyto!(dest, convert(Concatenated{Nothing}, concat))
+
+# couple back to Base implementation if no specialization exists:
+# https://github.com/JuliaLang/julia/blob/29da86bb983066dd076439c2c7bc5e28dbd611bb/base/abstractarray.jl#L1852
+function Base.copyto!(dest::AbstractArray, concat::Concatenated{Nothing})
+  catdims = Base.dims2cat(dims(concat))
+  shape = Base.cat_size_shape(catdims, concat.args...)
+  count(!iszero, catdims)::Int > 1 && zero!(dest)
+  return Base.__cat(dest, shape, catdims, concat.args...)
+end
+
+end

src/traits.jl

@@ -40,12 +40,11 @@ function derive(::Val{:AbstractArrayOps}, type)
     Base.iszero(::$type)
     Base.real(::$type)
     Base.fill!(::$type, ::Any)
-    ArrayLayouts.zero!(::$type)
+    DerivableInterfaces.zero!(::$type)
     Base.zero(::$type)
     Base.permutedims!(::Any, ::$type, ::Any)
     Broadcast.BroadcastStyle(::Type{<:$type})
     Base.copyto!(::$type, ::Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{0}})
-    Base.cat(::$type...; kwargs...)
     ArrayLayouts.MemoryLayout(::Type{<:$type})
     LinearAlgebra.mul!(::AbstractMatrix, ::$type, ::$type, ::Number, ::Number)
   end

src/zero.jl

@@ -0,0 +1,6 @@
+"""
+    zero!(x::AbstractArray)
+
+In-place version of `Base.zero`.
+"""
+function zero! end