tests: Expand matmul tests

Author: jpsamaroo

test/array/linalg/matmul.jl

@@ -1,42 +1,41 @@
-@testset "Matmul" begin
+@testset "With/Without Transpose" begin
     X = rand(40, 40)
     tol = 1e-12
-
     X1 = distribute(X, Blocks(10, 20))
     X2 = X1'*X1
     X3 = X1*X1'
     X4 = X1*X1
-
     @test norm(collect(X2) - (X' * X)) < tol
     @test norm(collect(X3) - (X * X')) < tol
     @test norm(collect(X4) - (X * X)) < tol
-    @test chunks(X2) |> size == (2, 2)
-    @test chunks(X3) |> size == (4, 4)
-    @test chunks(X4) |> size == (4, 2)
-    @test map(x->size(x) == (20, 20), domainchunks(X2)) |> all
-    @test map(x->size(x) == (10, 10), domainchunks(X3)) |> all
-    @test map(x->size(x) == (10, 20), domainchunks(X4)) |> all
-
-    @testset "Powers" begin
-        x = rand(Blocks(4,4), 16, 16)
-        @test collect(x^1) == collect(x)
-        @test collect(x^2) == collect(x*x)
-        @test collect(x^3) == collect(x*x*x)
-    end
+end
 
-    @testset "GEMM: $T" for T in (Float32, Float64, ComplexF32, ComplexF64)
-        A = rand(T, 128, 128)
-        B = rand(T, 128, 128)
+@testset "Powers" begin
+    x = rand(Blocks(4,4), 16, 16)
+    @test collect(x^1) == collect(x)
+    @test collect(x^2) == collect(x*x)
+    @test collect(x^3) == collect(x*x*x)
+end
 
-        DA = view(A, Blocks(32, 32))
-        DB = view(B, Blocks(32, 32))
+function test_gemm!(T, szA, szB, partA, partB)
+    @assert szA[1] == szB[2]
+    szC = (szA[1], szA[1])
+    @assert partA.blocksize[1] == partB.blocksize[2]
+    partC = Blocks(partA.blocksize[1], partB.blocksize[2])
 
-        ## Out-of-place gemm
-        # No transA, No transB
-        DC = DA * DB
-        C = A * B
-        @test collect(DC) ≈ C
+    A = rand(T, szA...)
+    B = rand(T, szB...)
+
+    DA = distribute(A, partA)
+    DB = distribute(B, partB)
 
+    ## Out-of-place gemm
+    # No transA, No transB
+    DC = DA * DB
+    C = A * B
+    @test collect(DC) ≈ C
+
+    if szA == szB
         # No transA, transB
         DC = DA * DB'
         C = A * B'
@@ -46,41 +45,45 @@
         DC = DA' * DB
         C = A' * B
         @test collect(DC) ≈ C
+    end
 
-        # transA, transB
-        DC = DA' * DB'
-        C = A' * B'
-        @test collect(DC) ≈ C
+    # transA, transB
+    DC = DA' * DB'
+    C = A' * B'
+    @test collect(DC) ≈ C
 
-        ## In-place gemm
-        # No transA, No transB
-        C = zeros(T, 128, 128)
-        DC = view(C, Blocks(32, 32))
-        mul!(C, A, B)
-        mul!(DC, DA, DB)
-        @test collect(DC) ≈ C
+    ## In-place gemm
+    # No transA, No transB
+    C = zeros(T, szC...)
+    DC = distribute(C, partC)
+    mul!(C, A, B)
+    mul!(DC, DA, DB)
+    @test collect(DC) ≈ C
 
+    if szA == szB
         # No transA, transB
-        C = zeros(T, 128, 128)
-        DC = view(C, Blocks(32, 32))
+        C = zeros(T, szC...)
+        DC = distribute(C, partC)
         mul!(C, A, B')
         mul!(DC, DA, DB')
         @test collect(DC) ≈ C
 
         # transA, No transB
-        C = zeros(T, 128, 128)
-        DC = view(C, Blocks(32, 32))
+        C = zeros(T, szC...)
+        DC = distribute(C, partC)
         mul!(C, A', B)
         mul!(DC, DA', DB)
         @test collect(DC) ≈ C
+    end
 
-        # transA, transB
-        C = zeros(T, 128, 128)
-        DC = view(C, Blocks(32, 32))
-        mul!(C, A', B')
-        mul!(DC, DA', DB')
-        collect(DC) ≈ C
+    # transA, transB
+    C = zeros(T, szA[2], szA[2])
+    DC = distribute(C, partC)
+    mul!(C, A', B')
+    mul!(DC, DA', DB')
+    collect(DC) ≈ C
 
+    if szA == szB
         ## Out-of-place syrk
         # No trans, trans
         DC = DA * DA'
@@ -94,17 +97,49 @@
 
         ## In-place syrk
         # No trans, trans
-        C = zeros(T, 128, 128)
-        DC = distribute(C, Blocks(32, 32))
+        C = zeros(T, szC...)
+        DC = distribute(C, partC)
         mul!(C, A, A')
         mul!(DC, DA, DA')
         @test collect(DC) ≈ C
 
         # trans, No trans
-        C = zeros(T, 128, 128)
-        DC = distribute(C, Blocks(32, 32))
+        C = zeros(T, szC...)
+        DC = distribute(C, partC)
         mul!(C, A', A)
         mul!(DC, DA', DA)
         @test collect(DC) ≈ C
     end
 end
+
+_sizes_to_test = [
+    (4, 4),
+    (7, 7),
+    (12, 12),
+    (16, 16),
+]
+size_sets_to_test = map(_sizes_to_test) do sz
+    rows, cols = sz
+    return [
+        (rows, cols) => (cols, rows),
+        (rows ÷ 2, cols) => (cols, rows ÷ 2),
+        (rows, cols ÷ 2) => (cols ÷ 2, rows),
+    ]
+end
+sizes_to_test = vcat(size_sets_to_test...)
+part_sets_to_test = map(_sizes_to_test) do sz
+    rows, cols = sz
+    return [
+        Blocks(rows, cols) => Blocks(cols, rows),
+        Blocks(rows ÷ 2, cols) => Blocks(cols, rows ÷ 2),
+        Blocks(rows, cols ÷ 2) => Blocks(cols ÷ 2, rows),
+    ]
+end
+parts_to_test = vcat(part_sets_to_test...)
+@testset "Size=$szA*$szB" for (szA, szB) in sizes_to_test
+    @testset "Partitioning=$partA*$partB" for (partA,partB) in parts_to_test
+        @testset "T=$T" for T in (Float32, Float64, ComplexF32, ComplexF64)
+            test_gemm!(T, szA, szB, partA, partB)
+        end
+    end
+end