Merge pull request #35 from alexmul1114/Faster-MTTKRP

Faster MTTKRP

benchmark/benchmarks.jl

@@ -1,7 +1,11 @@
 using BenchmarkTools
 
 # Benchmark suite modules
-const SUITE_MODULES = Dict("gcp" => :BenchmarkGCP, "mttkrp" => :BenchmarkMTTKRP)
+const SUITE_MODULES = Dict(
+    "gcp" => :BenchmarkGCP,
+    "mttkrp" => :BenchmarkMTTKRP,
+    "mttkrp-large" => :BenchmarkMTTKRPLarge,
+)
 
 # Create top-level suite including only sub-suites
 # specified by ENV variable "GCP_BENCHMARK_SUITES"

benchmark/suites/mttkrp-large.jl

@@ -0,0 +1,41 @@
+module BenchmarkMTTKRPLarge
+
+using BenchmarkTools, GCPDecompositions
+using Random
+
+const SUITE = BenchmarkGroup()
+
+# Collect setups
+const SETUPS = []
+
+## Balanced order-4 tensors
+append!(
+    SETUPS,
+    [
+        (; size = sz, rank = r, mode = n) for sz in [ntuple(n -> In, 4) for In in 20:20:80],
+        r in 20:20:120, n in 1:4
+    ],
+)
+
+## Imbalanced order-4 tensors
+append!(
+    SETUPS,
+    [
+        (; size = sz, rank = r, mode = n) for sz in [(20, 40, 80, 500), (500, 80, 40, 20)],
+        r in 100:100:300, n in 1:4
+    ],
+)
+
+# Generate random benchmarks
+for SETUP in SETUPS
+    Random.seed!(0)
+    X = randn(SETUP.size)
+    U = [randn(In, SETUP.rank) for In in SETUP.size]
+    SUITE["size=$(SETUP.size), rank=$(SETUP.rank), mode=$(SETUP.mode)"] = @benchmarkable(
+        GCPDecompositions.mttkrp($X, $U, $(SETUP.mode)),
+        seconds = 2,
+        samples = 5,
+    )
+end
+
+end

src/gcp-opt.jl

@@ -174,21 +174,85 @@ function _gcp(
     return CPD(λ, Tuple(U))
 end
 
-# inefficient but simple
+"""
+    mttkrp(X, (U1, U2, ..., UN), n)
+
+Compute the Matricized Tensor Times Khatri-Rao Product (MTTKRP)
+of an N-way tensor X with the matrices U1, U2, ..., UN along mode n.
+
+Algorithm is based on Section III-B of the paper:
+> **Fast Alternating LS Algorithms for High Order CANDECOMP/PARAFAC Tensor Factorizations**.
+> Anh-Huy Phan, Petr Tichavský, Andrzej Cichocki.
+> *IEEE Transactions on Signal Processing*, 2013.
+> DOI: 10.1109/TSP.2013.2269903
+"""
 function mttkrp(X, U, n)
     # Dimensions
     N, I, r = length(U), Tuple(size.(U, 1)), (only∘unique)(size.(U, 2))
     (N == ndims(X) && I == size(X)) || throw(DimensionMismatch("`X` and `U` do not have matching dimensions"))
 
-    # Matricized tensor (in mode n)
-    Xn = reshape(permutedims(X, [n; setdiff(1:N, n)]), size(X, n), :)
+    # Allocate output array G
+    G = similar(U[n])
 
-    # Khatri-Rao product (in mode n)
-    Zn = similar(U[1], prod(I[setdiff(1:N, n)]), r)
-    for j in 1:r
-        Zn[:, j] = reduce(kron, [view(U[i], :, j) for i in reverse(setdiff(1:N, n))])
+    # Choose appropriate multiplication order:
+    # + n == 1: no splitting required
+    # + n == N: no splitting required
+    # + 1 < n < N: better to multiply "bigger" side out first
+    #   + prod(I[1:n]) > prod(I[n:N]): better to multiply left-to-right
+    #   + prod(I[1:n]) < prod(I[n:N]): better to multiply right-to-left
+    if n == 1
+        mul!(G, reshape(X, I[1], :), khatrirao(U[reverse(2:N)]...))
+    elseif n == N
+        mul!(G, transpose(reshape(X, :, I[N])), khatrirao(U[reverse(1:N-1)]...))
+    elseif prod(I[1:n]) > prod(I[n:N])
+        # Inner multiplication: left side
+        kr_left = khatrirao(U[reverse(1:n-1)]...)
+        L = reshape(transpose(reshape(X, :, prod(I[n:N]))) * kr_left, (I[n:N]..., r))
+
+        # Outer multiplication: right side
+        kr_right = khatrirao(U[reverse(n+1:N)]...)
+        for j in 1:r
+            mul!(
+                view(G, :, j),
+                reshape(selectdim(L, ndims(L), j), I[n], :),
+                view(kr_right, :, j),
+            )
+        end
+    else
+        # Inner multiplication: right side
+        kr_right = khatrirao(U[reverse(n+1:N)]...)
+        R = reshape(reshape(X, prod(I[1:n]), :) * kr_right, (I[1:n]..., r))
+
+        # Outer multiplication: left side
+        kr_left = khatrirao(U[reverse(1:n-1)]...)
+        for j in 1:r
+            mul!(
+                view(G, :, j),
+                transpose(reshape(selectdim(R, ndims(R), j), :, I[n])),
+                view(kr_left, :, j),
+            )
+        end
     end
+    return G
+end
 
-    # MTTKRP (in mode n)
-    return Xn * Zn
+"""
+    khatrirao(A1, A2, ...)
+
+Compute the Khatri-Rao product (i.e., the column-wise Kronecker product)
+of the matrices `A1`, `A2`, etc.
+"""
+function khatrirao(A::Vararg{T,N}) where {T<:AbstractMatrix,N}
+    # Special case: N = 1
+    if N == 1
+        return A[1]
+    end
+
+    # General case: N > 1
+    r = (only ∘ unique)(size.(A, 2))
+    K = similar(A[1], prod(size.(A, 1)), r)
+    for j in 1:r
+        K[:, j] = reduce(kron, [view(A[i], :, j) for i in 1:N])
+    end
+    return K
 end

test/items/gcp-opt.jl

@@ -35,7 +35,7 @@ end
 @testitem "LeastSquaresLoss" begin
     using Random
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         X = [M[I] for I in CartesianIndices(size(M))]
@@ -55,7 +55,7 @@ end
 @testitem "NonnegativeLeastSquaresLoss" begin
     using Random
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         X = [M[I] for I in CartesianIndices(size(M))]
@@ -73,7 +73,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(fill(10.0, r), rand.(sz, r))
         X = [rand(Poisson(M[I])) for I in CartesianIndices(size(M))]
@@ -100,7 +100,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), randn.(sz, r))
         X = [rand(Poisson(exp(M[I]))) for I in CartesianIndices(size(M))]
@@ -127,7 +127,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         k = 1.5
@@ -155,7 +155,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         X = [rand(Rayleigh(M[I]/(sqrt(pi/2)))) for I in CartesianIndices(size(M))]
@@ -182,7 +182,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         X = [rand(Bernoulli(M[I]/(M[I] + 1))) for I in CartesianIndices(size(M))]
@@ -209,7 +209,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         X = [rand(Bernoulli(exp(M[I])/(exp(M[I]) + 1))) for I in CartesianIndices(size(M))]
@@ -236,7 +236,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         num_failures = 5
@@ -265,7 +265,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+    @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         X = [M[I] for I in CartesianIndices(size(M))]
@@ -294,7 +294,7 @@ end
     using Random
     using Distributions
 
-    @testset "size(X)=$sz, rank(X)=$r, β" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2, β in [0, 0.5, 1]
+    @testset "size(X)=$sz, rank(X)=$r, β" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2, β in [0, 0.5, 1]
         Random.seed!(0)
         M = CPD(ones(r), rand.(sz, r))
         # May want to consider other distributions depending on value of β
@@ -342,7 +342,7 @@ end
     using Random, Distributions, IntervalSets
 
     @testset "Least Squares" begin
-        @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+        @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
             Random.seed!(0)
             M = CPD(ones(r), randn.(sz, r))
             X = [M[I] for I in CartesianIndices(size(M))]
@@ -366,7 +366,7 @@ end
     end
 
     @testset "Poisson" begin
-        @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (30, 40, 50)], r in 1:2
+        @testset "size(X)=$sz, rank(X)=$r" for sz in [(15, 20, 25), (50, 40, 30)], r in 1:2
             Random.seed!(0)
             M = CPD(fill(10.0, r), rand.(sz, r))
             X = [rand(Poisson(M[I])) for I in CartesianIndices(size(M))]