Aggressive constprop in sparse * dense

[jishnub]

Since tA is usually known from the types, this helps eliminate branches. Strangely, Cthulhu suggests that LinearAlgebra.wrap call within spdensemul! is still inferred as a Union, but this doesn't matter much as there's a function barrier.

[codecov]

Codecov Report

Merging #460 (634ff3f) into main (bb86364) will increase coverage by 0.06%.
Report is 2 commits behind head on main.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #460      +/-   ##
==========================================
+ Coverage   85.43%   85.50%   +0.06%     
==========================================
  Files          13       13              
  Lines        8733     8773      +40     
==========================================
+ Hits         7461     7501      +40     
  Misses       1272     1272              

Files	Coverage Δ
src/linalg.jl	91.55% <100.00%> (ø)
src/sparsematrix.jl	95.75% <100.00%> (+0.06%)	⬆️
src/sparsevector.jl	95.64% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[dkarrasch]

Could you please share some quick piece of code about how you check for elision of branches. I think I tried that with _generic_matmatmul! in LinearAlgebra (the one big BLAS method), but wasn't sure how to see that exactly.

[jishnub]

I was checking with Cthulhu. On main, using

julia> S = sprand(5, 5, 0.2);

julia> @descend_code_warntype S * ones(size(S,2))

and descending into the spdensemul! call, I see

spdensemul!(C, tA, tB, A, B, _add) @ SparseArrays ~/Dropbox/JuliaPackages/SparseArrays.jl/src/linalg.jl:38
38 function spdensemul!(C::Vector{Float64}, tA::Char, tB::Char, A::SparseMatrixCSC{Float64, Int64}, B::Vector{Float64}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})::Vector{Float64}
39     if tA == 'N'
40         _spmatmul!(C::Vector{Float64}, A::SparseMatrixCSC{Float64, Int64}, LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)::Union{LinearAlgebra.Adjoint{Float64, Vector{Float64}}, LinearAlgebra.Transpose{Float64, Vector{Float64}}, Vector{Float64}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha::Bool, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta::Bool)
41     elseif tA == 'T'
42         _At_or_Ac_mul_B!(transpose, C::Vector{Float64}, A::SparseMatrixCSC{Float64, Int64}, LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)::Union{LinearAlgebra.Adjoint{Float64, Vector{Float64}}, LinearAlgebra.Transpose{Float64, Vector{Float64}}, Vector{Float64}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha::Bool, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta::Bool)
43     elseif tA == 'C'
44         _At_or_Ac_mul_B!(adjoint, C::Vector{Float64}, A::SparseMatrixCSC{Float64, Int64}, LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)::Union{LinearAlgebra.Adjoint{Float64, Vector{Float64}}, LinearAlgebra.Transpose{Float64, Vector{Float64}}, Vector{Float64}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha::Bool, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta::Bool)
45     elseif (tA::Char in ('S', 's', 'H', 'h'))::Bool && tB == 'N'
46         rangefun::Union{typeof(SparseArrays.nzrangelo), typeof(SparseArrays.nzrangeup)} = isuppercase(tA::Char)::Bool ? nzrangeup : nzrangelo
47         diagop::Union{typeof(identity), typeof(real)} = (tA::Char in ('S', 's'))::Bool ? identity : real
48         odiagop::Union{typeof(adjoint), typeof(transpose)} = (tA::Char in ('S', 's'))::Bool ? transpose : adjoint
49         T::Core.Const(Float64) = eltype(C::Vector{Float64})::Type{Float64}
50         _mul!(rangefun::Union{typeof(SparseArrays.nzrangelo), typeof(SparseArrays.nzrangeup)}, diagop::Union{typeof(identity), typeof(real)}, odiagop::Union{typeof(adjoint), typeof(transpose)}, C::Vector{Float64}, A::SparseMatrixCSC{Float64, Int64}, B::Vector{Float64}, T::Type{Float64}(_add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha::Bool)::Float64, T::Type{Float64}(_add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta::Bool)::Float64)
51     else
52         LinearAlgebra._generic_matmatmul!(C::Vector{Float64}, 'N', 'N', LinearAlgebra.wrap(A::SparseMatrixCSC{Float64, Int64}, tA::Char)::Union{LinearAlgebra.Adjoint{Float64, SparseMatrixCSC{Float64, Int64}}, LinearAlgebra.Hermitian{Float64, SparseMatrixCSC{Float64, Int64}}, LinearAlgebra.Symmetric{Float64, SparseMatrixCSC{Float64, Int64}}, LinearAlgebra.Transpose{Float64, SparseMatrixCSC{Float64, Int64}}, SparseMatrixCSC{Float64, Int64}}, LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)::Union{LinearAlgebra.Adjoint{Float64, Vector{Float64}}, LinearAlgebra.Transpose{Float64, Vector{Float64}}, Vector{Float64}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
53     end
54     return C::Vector{Float64}
55 end
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • %5 = ==(::Char,::Char)::Bool
   LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)
   _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha
   _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta
   %11 = call → _spmatmul!(…,…,…,…,…)::Any
   %13 = ==(::Char,::Char)::Bool
   LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)
   _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha
   _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta
v  %20 = call → _At_or_Ac_mul_B!(…,…,…,…,…,…)::Any

whereas, on this PR, I obtain

spdensemul!(C, tA, tB, A, B, _add) @ SparseArrays ~/Dropbox/JuliaPackages/SparseArrays.jl/src/linalg.jl:38
┌ Warning: Some line information is missing, type-assignment may be incomplete
└ @ Cthulhu ~/.julia/packages/Cthulhu/9kV71/src/codeview.jl:144
38 function spdensemul!(C::Vector{Float64}, tA::Char, tB::Char, A::SparseMatrixCSC{Float64, Int64}, B::Vector{Float64}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})::Vector{Float64}
39     if tA == 'N'
40         _spmatmul!(C::Vector{Float64}, A::SparseMatrixCSC{Float64, Int64}, LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)::Union{LinearAlgebra.Adjoint{Float64, Vector{Float64}}, LinearAlgebra.Transpose{Float64, Vector{Float64}}, Vector{Float64}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha::Bool, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta::Bool)
41     elseif tA == 'T'
42         _At_or_Ac_mul_B!(transpose, C, A, LinearAlgebra.wrap(B, tB), _add.alpha, _add.beta)
43     elseif tA == 'C'
44         _At_or_Ac_mul_B!(adjoint, C, A, LinearAlgebra.wrap(B, tB), _add.alpha, _add.beta)
45     elseif tA in ('S', 's', 'H', 'h') && tB == 'N'
46         rangefun = isuppercase(tA) ? nzrangeup : nzrangelo
47         diagop = tA in ('S', 's') ? identity : real
48         odiagop = tA in ('S', 's') ? transpose : adjoint
49         T = eltype(C)
50         _mul!(rangefun, diagop, odiagop, C, A, B, T(_add.alpha), T(_add.beta))
51     else
52         LinearAlgebra._generic_matmatmul!(C, 'N', 'N', LinearAlgebra.wrap(A, tA), LinearAlgebra.wrap(B, tB), _add)
53     end
54     return C::Vector{Float64}
55 end
Select a call to descend into or ↩ to ascend. [q]uit. [b]ookmark.
Toggles: [w]arn, [h]ide type-stable statements, [t]ype annotations, [s]yntax highlight for Source/LLVM/Native, [j]ump to source always.
Show: [S]ource code, [A]ST, [T]yped code, [L]LVM IR, [N]ative code
Actions: [E]dit source code, [R]evise and redisplay
 • %6 = < concrete eval > ==(::Core.Const('N'),::Core.Const('N'))::Core.Const(true)
   LinearAlgebra.wrap(B::Vector{Float64}, tB::Char)
   _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.alpha
   _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool}.beta
   %12 = call → _spmatmul!(…,…,…,…,…)::Any
   ↩

This indicates that only one branch (the tA == 'N' one) is being inferred, as the equality check is evaluated at compile time.

[jishnub]

Should we go ahead with this?

[dkarrasch]

I'd say yes! Do you think there could be any benefit from adding those macros to the _generic_mat***mul! in LinearAlgebra? See JuliaLang/julia#51812 for some interesting discussion.

[dkarrasch]

I think we should even backport this to v1.10... Opinions? Objections?

[PallHaraldsson]

This was backported here to release-1.10, but not to JuliaLang? I'm just trying to understand, I think it's not in1.10-rc1, and I think at least this line needs changing:

https://github.com/JuliaLang/julia/blob/5aaa94854367ca875375e38ae14f369f124e7315/stdlib/SparseArrays.version#L2

[dkarrasch]

Yes, likely this needs to be bumped on the release-1.10 branch there.

[ViralBShah]

I think @KristofferC may not want this in 1.10.