Julia - Linear combination of row-wise outer products

Here's another verison

# Parameters
n, m, p = 100, 10, 3

# Matrices & Vectors
A, B, x = randn(n, m), randn(n, p), randn(n)

function old_way(A, B, x)
    # Slow method
    result = zeros(m, p)
    for i in 1:n
        result += x[i] * (A[i, :] * B[i, :]')
    end
end

function another_way(A, B, x)
    sum(xi * (Arow * Brow') for (xi, Arow, Brow) in zip(x, eachrow(A), eachrow(B)))
end

And benchmarking:

julia> using BenchmarkTools

julia> @benchmark old_way(A, B, x)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):   83.495 μs …   2.653 ms  ┊ GC (min … max): 0.00% … 95.79%
 Time  (median):      87.500 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   101.496 μs ± 115.196 μs  ┊ GC (mean ± σ):  6.46% ±  5.53%

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

 Memory estimate: 153.48 KiB, allocs estimate: 1802.


julia> @benchmark another_way(A, B, x)
BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  25.850 μs … 923.032 μs  ┊ GC (min … max): 0.00% … 95.94%
 Time  (median):     27.477 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   31.851 μs ±  35.440 μs  ┊ GC (mean ± σ):  5.03% ±  4.49%

  ▇█▇▅▃▁           ▁▁▁                                         ▁
  ███████▇▇▆▇████████████▇█▇▇▇▆▇▆▅▆▆▆▅▆▆▅▆▆▇▆▅▅▅▆▅▆▅▅▅▄▃▅▅▅▄▄▄ █
  25.8 μs       Histogram: log(frequency) by time      77.4 μs <

 Memory estimate: 98.31 KiB, allocs estimate: 304.

So it's a little faster, and uses a little less memory.

Does this answer your question?

General tips to save time and memory:

1. Put code in a method instead of the global scope, and make sure every variable in that function comes from the arguments, not global variables. That way, Julia's compiler can infer the types of variables and optimize.

2. Reduce allocations where possible, and you have many opportunities. The changes here distinguish the old_way and new_way methods, and it causes a 5-6x speedup and reduction to 1 allocation.

   • When slicing an array, use @view to avoid default behavior of allocating a copy.
   • You can change result in-place with .+=. += allocates a new array and reassigns the variable result to it.
   • For elementwise operations like x[i] * ..., chaining dotted operators fuses the underlying elementwise loops and reduces allocations of intermediate arrays.
   • A matrix multiplication of a column (Mx1) vector and a row (1xN) vector can be simplified to elementwise multiplication.

n, m, p = 100, 10, 3
A, B, x = randn(n, m), randn(n, p), randn(n)

# Methods below do not use the above global variables

function old_way(A, B, x, n, m, p)
  result = zeros(m, p)
  for i in 1:n
    result += x[i] * (A[i, :] * B[i, :]')
  end
  result
end

function new_way(A, B, x, n, m, p)
  result = zeros(m, p)
  for i in 1:n
    result .+= x[i] .* ( @view(A[i, :]) .* @view(B[i, :])' )
  end
  result
end

using BenchmarkTools

@btime old_way(A, B, x, n, m, p);
# 36.753 μs (501 allocations: 125.33 KiB)

@btime new_way(A, B, x, n, m, p);
# 6.542 μs (1 allocation: 336 bytes)

old_way(A, B, x, n, m, p) == new_way(A, B, x, n, m, p)
# true

The example above avoided global variables so far, and the example below will show why. Even if you put your code in a method but still use global variables, not only is the performance just generally worse, trying to reduce allocations backfires:

# Methods below use n, m, p as global inputs

function old_oops(A, B, x)
  # same code as old_way(A, B, x, n, m, p)
end

function new_oops(A, B, x)
  # same code as new_way(A, B, x, n, m, p)
end

@btime old_oops(A, B, x);
# 95.317 μs (1802 allocations: 153.48 KiB)

@btime new_oops(A, B, x);
# 235.191 μs (1302 allocations: 81.61 KiB)

If your setup has the same structure as your MWE, using LinearAlgebra:

faster(A,B,x) = (diagm(x)*A)'*B

runs 4x faster:

using LinearAlgebra, BenchmarkTools

# Parameters
n, m, p = 100, 10, 3
# Matrices & Vectors
A, B, x = randn(n, m), randn(n, p), randn(n)
# Slow method

function slow(A,B,x)
    result = zeros(m, p)
    for i in 1:n
        result += x[i] * (A[i, :] * B[i, :]')
    end
    result
end

faster(A,B,x) = (diagm(x)*A)'*B

@assert(slow(A,B,x) ≈ faster(A,B,x))
@btime slow($A,$B,$x)     # 113.400 μs (1802 allocations: 139.39 KiB)
@btime faster($A,$B,$x)   # 28.100 μs (4 allocations: 86.41 KiB)