Improve inlining, faster Lazy mul getindex (#268)

* Improve inlining, faster Lazy mul getindex

* Fix _permutedims

* fix instantiate for lazy arrays

* Update lazyconcat.jl

* Require Julia v1.9

.github/workflows/ci.yml

@@ -12,8 +12,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
-          - '1' # Latest stable release
+          - '1.9'
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "LazyArrays"
 uuid = "5078a376-72f3-5289-bfd5-ec5146d43c02"
-version = "1.5.2"
+version = "1.6"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -24,7 +24,7 @@ FillArrays = "1.0"
 MacroTools = "0.5"
 MatrixFactorizations = "1.0, 2.0"
 StaticArrays = "1.0"
-julia = "1.6"
+julia = "1.9"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"

src/lazyconcat.jl

@@ -14,16 +14,19 @@ Vcat(A...) = ApplyArray(vcat, A...)
 Vcat{T}(A...) where T = ApplyArray{T}(vcat, A...)
 Vcat() = Vcat{Any}()
 
-Vcat(A::AbstractVector...) = ApplyVector(vcat, A...)
-Vcat{T}(A::AbstractVector...) where T = ApplyVector{T}(vcat, A...)
-
-function instantiate(A::Applied{DefaultApplyStyle,typeof(vcat)})
-    isempty(A.args) && return A
-    m = size(A.args[1],2)
-    for k=2:length(A.args)
-        size(A.args[k],2) == m || throw(ArgumentError("number of columns of each array must match (got $(map(x->size(x,2), A)))"))
+@inline Vcat(A::AbstractVector...) = ApplyVector(vcat, A...)
+@inline Vcat(A::AbstractVector{T}...) where T = ApplyVector{T}(vcat, A...)
+@inline Vcat{T}(A::AbstractVector...) where T = ApplyVector{T}(vcat, A...)
+
+@inline function applied_instantiate(::typeof(vcat), args...)
+    iargs = map(instantiate, args)
+    if !isempty(iargs)
+        m = size(iargs[1],2)
+        for k=2:length(iargs)
+            size(iargs[k],2) == m || throw(ArgumentError("number of columns of each array must match (got $(map(x->size(x,2), args)))"))
+        end
     end
-    Applied{DefaultApplyStyle}(A.f,map(instantiate,A.args))
+    vcat, iargs
 end
 
 @inline applied_eltype(::typeof(vcat)) = Any

src/linalg/mul.jl

@@ -208,6 +208,14 @@ __mul_args_cols(jr, z, y...) =
     (__mul_args_cols(_mul_args_colsupport(z,jr), y...)..., jr)
 _mul_args_cols(jr, z, y...) = __mul_args_cols(_mul_args_colsupport(z,jr), y...)
 
+_transposefirst(a, b...) = (transpose(a), b...)
+
+function _mul_getindex(args::Tuple, k::Int, j::Int)
+    kjr = intersect.(_mul_args_rows(k, args...), _mul_args_cols(j, reverse(args)...))
+    any(isempty, kjr) && return zero(mapreduce(eltype, promote_type, args))
+    *(_transposefirst(map(getindex, args, (k, kjr...), (kjr..., j))...)...)
+end
+
 sublayout(::ApplyLayout{typeof(*)}, _...) = ApplyLayout{typeof(*)}()
 # matrix-indexing loses the multiplication structure as we don't support tensor multiplication
 sublayout(::ApplyLayout{typeof(*)}, ::Type{<:Tuple{AbstractMatrix}}) = UnknownLayout()
@@ -235,7 +243,7 @@ _vec_mul_arguments(args, (kr,jr)::Tuple{AbstractVector,Number}) =
 
 # this is a row-vector view
 _vec_mul_arguments(args, (kr,jr)::Tuple{Number,AbstractVector}) =
-    _vec_mul_arguments(reverse(map(permutedims, args)), (jr,kr))
+    _vec_mul_arguments(reverse(map(transpose, args)), (jr,kr))
 
 _mat_mul_arguments(V) = _mat_mul_arguments(arguments(parent(V)), parentindices(V))
 _vec_mul_arguments(V) = _vec_mul_arguments(arguments(parent(V)), parentindices(V))

test/lazymultests.jl

@@ -16,7 +16,7 @@ Base.convert(::Type{AbstractMatrix{T}}, A::MyMatrix) where T = MyMatrix(convert(
 Base.convert(::Type{MyMatrix{T}}, A::AbstractArray{T}) where T = MyMatrix{T}(A)
 Base.convert(::Type{MyMatrix{T}}, A::AbstractArray) where T = MyMatrix{T}(convert(AbstractArray{T}, A))
 Base.convert(::Type{MyMatrix}, A::AbstractArray{T}) where T = MyMatrix{T}(A)
-Base.getindex(A::MyMatrix, kj...) = A.A[kj...]
+Base.getindex(A::MyMatrix, k::Integer, j::Integer) = A.A[k, j]
 Base.getindex(A::MyMatrix, k::Integer, ::Colon) = A.A[k,:]
 Base.getindex(A::MyMatrix, ::Colon, j::Integer) = A.A[:,j]
 Base.getindex(A::MyMatrix, ::Colon, j::AbstractVector) = MyMatrix(A.A[:,j])