Diagonal Matrix Mult is Slower than Dense

[andrew-saydjari]

In trying a MWE for my core functionality, I have noticed that a simple matrix vector multiply using a Diagonal matrix is more than 5x slower than the dense matrix-vector multiply. One can of course convert this to a broadcasted vector multiply, where by I can get Reactant to within 20% of my standard CPU performance, but it would be nice to keep the same syntax and have Reactant just handle Diagonal types more seemlessly. My motivation here is that my actual production code uses custom types that are low rank representations of the matrices with overloaded definitions of matrix multiplications, that often involve diagonal components.

using Reactant, Random, LinearAlgebra, PrettyChairmarks

function core_mul_dense(AinvVIinvtX,AinvVIinv,VAinvVIinvtX,Ainv,V,x,out)
    mul!(AinvVIinvtX,AinvVIinv',x)
    mul!(VAinvVIinvtX,V,AinvVIinvtX,-1,0)
    VAinvVIinvtX .+= x
    mul!(out,Ainv,VAinvVIinvtX)
    return
end

rng = Xoshiro(122)
n = 8000
ln = 50
x = randn(rng,n)
AinvVIinv = randn(rng,n,ln)
V = randn(rng,n,ln)
Ax = randn(rng,n)
Ainv = Diagonal(Ax)
AinvDense = Matrix(Ainv)
Ainv_vec = reshape(Ax,:,1)
AinvVIinvtX = zeros(ln)
VAinvVIinvtX = zeros(n)
out = zeros(n);

## Diagonal matrix-vector multiply
@bs core_mul_dense(AinvVIinvtX,AinvVIinv,VAinvVIinvtX,Ainv,V,x,out) seconds=3

# Chairmarks.Benchmark: 9585 samples with 1 evaluation.
#  Range (min … max):  234.913 μs …   4.769 ms  ┊ GC (min … max): 0.00% … 0.00%
#  Time  (median):     288.765 μs               ┊ GC (median):    0.00%
#  Time  (mean ± σ):   298.449 μs ± 124.596 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

## Fully dense matrix-vector multiply
@bs core_mul_dense(AinvVIinvtX,AinvVIinv,VAinvVIinvtX,AinvDense,V,x,out) seconds=10

# Chairmarks.Benchmark: 202 samples with 1 evaluation.
#  Range (min … max):  44.806 ms … 56.304 ms  ┊ GC (min … max): 0.00% … 0.00%
#  Time  (median):     46.955 ms              ┊ GC (median):    0.00%
#  Time  (mean ± σ):   47.470 ms ±  1.750 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

Reactant.set_default_backend("cpu")

xR = Reactant.ConcreteRArray(x)
AinvVIinvR = Reactant.ConcreteRArray(AinvVIinv)
VR = Reactant.ConcreteRArray(V)
AinvDenseR = Reactant.ConcreteRArray(AinvDense)
AinvVIinvtXR = Reactant.ConcreteRArray(AinvVIinvtX)
VAinvVIinvtXR = Reactant.ConcreteRArray(VAinvVIinvtX)
outR = Reactant.ConcreteRArray(out);

f = @compile sync=true core_mul_dense(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvDenseR,VR,xR,outR)

## Reactant fully dense matrix-vector multiply
@bs f(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvDenseR,VR,xR,outR) seconds=10

# Chairmarks.Benchmark: 204 samples with 1 evaluation.
#  Range (min … max):  42.615 ms … 60.736 ms  ┊ GC (min … max): 0.00% … 0.00%
#  Time  (median):     44.233 ms              ┊ GC (median):    0.00%
#  Time  (mean ± σ):   46.662 ms ±  4.344 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

## Reactant diagonal matrix-vector multiply
AinvR = Reactant.to_rarray(Ainv);

f = @compile sync=true core_mul_dense(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvR,VR,xR,outR)

@bs f(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvR,VR,xR,outR) seconds=10

# Chairmarks.Benchmark: 37 samples with 1 evaluation.
#  Range (min … max):  239.017 ms … 394.568 ms  ┊ GC (min … max): 0.00% … 0.00%
#  Time  (median):     246.455 ms               ┊ GC (median):    0.00%
#  Time  (mean ± σ):   273.044 ms ±  49.257 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

[andrew-saydjari]

A nice summary.

LA Type	CPU	Reactant CPU	Reactant GPU
Dense	40.211 ms	44.233 ms	501.665 μs
Diagonal	273.977 μs	246.455 ms	859.558 μs
Vector	273.977 μs	326.525 μs	86.524 μs

[wsmoses]

ok fixed the accidental large constant here: #1487

pending optimizations that need implementation to help:

• Dot general of reshape Enzyme-JAX#1075
• Add (x * -1), y -> y - x Enzyme-JAX#1178
• Matrix multiply of diagonal (dot general of scatter of zeros) Enzyme-JAX#1179

[avik-pal]

After EnzymeAD/Enzyme-JAX#1194

julia> @bs f(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvDenseR,VR,xR,outR) seconds=10
       
Chairmarks.Benchmark: 1234 samples with 1 evaluation.
 Range (min … max):  5.006 ms … 12.469 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     7.599 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   7.681 ms ±  1.137 ms  ┊ GC (mean ± σ):  0.00% ± 0.00%

                ▂   ▃▂▃█▇▅▅▄▅▃▄▁▂▃▁                           
  ▂▂▃▃▄▅▅▄▅▅▄▄▅████████████████████▆█▅▅▇▅▆▄▄▄▃▃▂▄▄▃▁▃▃▃▃▂▃▂▃ ▄
  5.01 ms        Histogram: frequency by time        11.2 ms <

 Memory estimate: 752.0 bytes, allocs estimate: 27.


julia> @bs f(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvR,VR,xR,outR) seconds=10
       
Chairmarks.Benchmark: 39290 samples with 1 evaluation.
 Range (min … max):   85.719 μs … 163.830 ms  ┊ GC (min … max): 0.00% … 21.41%
 Time  (median):     192.969 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   217.589 μs ± 840.476 μs  ┊ GC (mean ± σ):  0.00% ±  0.11%

    ▁   ▆█▄▄▁                                                    
  ▁▄█▅▆██████▆▄▄▃▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  85.7 μs          Histogram: frequency by time          938 μs <

 Memory estimate: 752.0 bytes, allocs estimate: 27.

[andrew-saydjari]

Awesome! I am happy to pull the new version and benchmark, but I think maybe the most relevant comparison now that the huge slow down is fixed is Diagonal CPU versus Diagonal Reactant CPU, to understand if the slowdown I was seeing (see the Vector row above) persists, or is fixed.

[wsmoses]

also fyi the new features will be available once this JuliaPackaging/Yggdrasil#11739 finishes building [and is bumped here]

[wsmoses]

now with these enabled (specifically this pr: #1510), Reactant seems to win in all cases on my laptop (cpu only, macbook pro)

## Normal diag matrix-vector multiply
julia> @bs core_mul_dense(AinvVIinvtX,AinvVIinv,VAinvVIinvtX,Ainv,V,x,out) seconds=3
Chairmarks.Benchmark: 27602 samples with 1 evaluation.
 Range (min … max):  101.375 μs … 179.000 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     102.541 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   103.133 μs ±   1.899 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

     ▃▅▇██▇▆▄▂▁▁▁▁    ▁▂▃▄▄▃▂▂▂▁▁▁                              ▃
  ▂▂███████████████▇███████████████▇▇██▇█▇▇▇▇▇▇▇█▆▇▇▆▇▅▆▇▇▇█▆▆▅ █
  101 μs        Histogram: log(frequency) by time        111 μs <

 Memory estimate: 0.0 bytes, allocs estimate: 0.


## Normal dense matrix-vector multiply
julia> @bs core_mul_dense(AinvVIinvtX,AinvVIinv,VAinvVIinvtX,AinvDense,V,x,out) seconds=10
Chairmarks.Benchmark: 1924 samples with 1 evaluation.
 Range (min … max):  3.256 ms …   8.806 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     4.979 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   4.940 ms ± 601.351 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                          ▁▁▄▃▁▄▂▁▂      ▁▆▆▆█▇▂▁▄▂            
  ▁▁▂▂▂▂▂▃▃▄▃▂▂▃▃▂▄▄▄▄▆▆▆██████████▆▅▅▄▄███████████▅▄▂▂▂▂▁▁▁▁ ▄
  3.26 ms         Histogram: frequency by time        6.15 ms <


## Reactant dense matrix-vector multiply
julia> @bs f(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvDenseR,VR,xR,outR) seconds=10
Chairmarks.Benchmark: 3691 samples with 1 evaluation.
 Range (min … max):  2.193 ms …   4.603 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     2.539 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   2.601 ms ± 192.530 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                 █▆▂            ▄▄                             
  ▁▂▂▂▂▂▁▁▁▁▁▁▁▂▇████▅▄▂▂▂▂▂▂▂▂▆███▅▃▃▂▂▂▂▂▂▂▁▁▁▁▁▁▂▁▁▁▁▁▁▁▁▁ ▂
  2.19 ms         Histogram: frequency by time        3.17 ms <

 Memory estimate: 864.0 bytes, allocs estimate: 30.

## Reactant diagonal matrix-vector multiply
julia> @bs f(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvR,VR,xR,outR) seconds=10
Chairmarks.Benchmark: 104486 samples with 1 evaluation.
 Range (min … max):  78.458 μs … 477.708 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     88.875 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   89.392 μs ±   4.163 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                           ▃█▅▅▅▅▄                              
  ▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▂▂▄▄▄▃▂▃▄▇██████████▇▅▄▄▃▃▃▃▃▃▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁ ▃
  78.5 μs         Histogram: frequency by time         99.6 μs <

 Memory estimate: 864.0 bytes, allocs estimate: 30.


julia> @code_hlo core_mul_dense(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvR,VR,xR,outR)
module @reactant_core_mu... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<50xf64> {tf.aliasing_output = 0 : i32}, %arg1: tensor<50x8000xf64>, %arg2: tensor<8000xf64> {tf.aliasing_output = 1 : i32}, %arg3: tensor<8000xf64>, %arg4: tensor<50x8000xf64>, %arg5: tensor<8000xf64>, %arg6: tensor<8000xf64> {tf.aliasing_output = 2 : i32}) -> (tensor<50xf64>, tensor<8000xf64>, tensor<8000xf64>) {
    %0 = stablehlo.dot_general %arg1, %arg5, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<50x8000xf64>, tensor<8000xf64>) -> tensor<50xf64>
    %1 = stablehlo.dot_general %arg4, %0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<50x8000xf64>, tensor<50xf64>) -> tensor<8000xf64>
    %2 = stablehlo.subtract %arg5, %1 : tensor<8000xf64>
    %3 = stablehlo.multiply %arg3, %2 : tensor<8000xf64>
    return %0, %2, %3 : tensor<50xf64>, tensor<8000xf64>, tensor<8000xf64>
  }
}

julia> @code_hlo core_mul_dense(AinvVIinvtXR,AinvVIinvR,VAinvVIinvtXR,AinvDenseR,VR,xR,outR)
module @reactant_core_mu... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<50xf64> {tf.aliasing_output = 0 : i32}, %arg1: tensor<50x8000xf64>, %arg2: tensor<8000xf64> {tf.aliasing_output = 1 : i32}, %arg3: tensor<8000x8000xf64>, %arg4: tensor<50x8000xf64>, %arg5: tensor<8000xf64>, %arg6: tensor<8000xf64> {tf.aliasing_output = 2 : i32}) -> (tensor<50xf64>, tensor<8000xf64>, tensor<8000xf64>) {
    %0 = stablehlo.dot_general %arg1, %arg5, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<50x8000xf64>, tensor<8000xf64>) -> tensor<50xf64>
    %1 = stablehlo.dot_general %arg4, %0, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<50x8000xf64>, tensor<50xf64>) -> tensor<8000xf64>
    %2 = stablehlo.subtract %arg5, %1 : tensor<8000xf64>
    %3 = stablehlo.dot_general %arg3, %2, contracting_dims = [0] x [0], precision = [DEFAULT, DEFAULT] : (tensor<8000x8000xf64>, tensor<8000xf64>) -> tensor<8000xf64>
    return %0, %2, %3 : tensor<50xf64>, tensor<8000xf64>, tensor<8000xf64>
  }
}

[andrew-saydjari]

Awesome, I will give this a test drive with my complicated types, that have nested diagonals (as woodbury updates) to see if Reactant descends those recursive rules correctly. Thanks so much!