Why is the parameterized version of this function slower than the vanilla one?

[RicRax]

Hi, I wrote some benchmarks to see how mojo performs in matmul, following as a guide this: https://docs.modular.com/mojo/notebooks/Matmul.
However, I noticed that my version with parameters is slower than the one without.
The vanilla version benches with 0.2 seconds, while the one with parameters in 1.6 seconds.
The code for the Matrix struct is the same as the one on the docs. This is the vanilla function:

fn matrix_getitem(self: object, i: object) raises -> object:
    return self.value[i]


fn matrix_setitem(self: object, i: object, value: object) raises -> object:
    self.value[i] = value
    return None


fn matrix_append(self: object, value: object) raises -> object:
    self.value.append(value)
    return None


fn matrix_init(rows: Int, cols: Int) raises -> object:
    var value = object([])
    return object(
        Attr("value", value), Attr("__getitem__", matrix_getitem), Attr("__setitem__", matrix_setitem),
        Attr("rows", rows), Attr("cols", cols), Attr("append", matrix_append),
    )

-----

from matrix import matrix_init
from matrix_types import Matrix
import benchmark

def kernel_3mm(ni, nj, nk, nl, nm, A, B, C, D, E, F, G):
    for i in range(ni):
        for j in range(nj):
            for k in range(nk):
                E[i][j] += A[i][k] * B[k][j]                


    for i in range(nj):
        for j in range(nl):
            for k in range(nm):
                F[i][j] += C[i][k] * D[k][j]

    for i in range(ni):
        for j in range(nl):
            for k in range(nj):
                G[i][j] += E[i][k] * F[k][j]




def benchmark_t3mm(ni: Int ,nj: Int, nk: Int,nl: Int, nm: Int):
    A = matrix_init(ni, nk)
    B = matrix_init(nk, nj)
    C = matrix_init(nj, nm)
    D = matrix_init(nm, nl)
    E = matrix_init(ni, nj)
    F = matrix_init(nj, nl)
    G = matrix_init(ni, nl)

    for i in range(ni):
        a_row = object([])
        for j in range(nk):
            a_row.append((i * j) / ni)
        A.append(a_row)
    
    for i in range(nk):
        b_row = object([])
        for j in range(nj):
            b_row.append((i * (j + 1)) / nj)
        B.append(b_row)

    for i in range(nj):
        c_row = object([])
        for j in range(nm):
            c_row.append((i * (j + 3)) / nl)
        C.append(c_row)

    for i in range(nm):
        d_row = object([])
        for j in range(nl):
            d_row.append((i * (j + 2)) / nk)
        D.append(d_row)

    for i in range(ni):
        e_row = object([])
        for j in range(nj):
            e_row.append(0.0)
        E.append(e_row)

    for i in range(nj):
        f_row = object([])
        for j in range(nl):
            f_row.append(0.0)
        F.append(f_row)

    for i in range(ni):
        g_row = object([])
        for j in range(nl):
            g_row.append(0.0)
        G.append(g_row)

    @parameter
    fn test_fn():
        try:
            _ = kernel_3mm(ni, nj, nk, nl, nm, A, B, C, D, E, F, G )
        except:
            print("error occurred")

    var secs = benchmark.run[test_fn](max_runtime_secs=0.5).mean()
    _ = (A, B, C, D, E, F, G)
    return secs

The parameterized version, again very similar to the docs one, is this:

from matrix_types import Matrix
import benchmark

fn kernel_3mm_types(ni: Int, nj: Int, nk: Int, nl: Int, nm: Int, A:Matrix, B:Matrix, C:Matrix, D:Matrix, E:Matrix, F:Matrix, G:Matrix):
    for i in range(ni):
        for j in range(nj):
            for k in range(nk):
                E[i,j] += A[i,k] * B[k,j]                

    for i in range(nj):
        for j in range(nl):
            for k in range(nm):
                F[i,j] += C[i,k] * D[k,j]

    for i in range(ni):
        for j in range(nl):
            for k in range(nj):
                G[i,j] += E[i,k] * F[k,j]


alias ni = 10
alias nj = 10
alias nk = 10
alias nl = 10
alias nm = 10

@always_inline
fn benchmark_t3mm_typed[
    func: fn (Int, Int, Int, Int, Int, Matrix, Matrix, Matrix, Matrix, Matrix, Matrix, Matrix) -> None]() -> object:
    var A = Matrix[ni,nk]()
    # make method in matrix to initialize values the way it is done in the no typed version
    var B = Matrix[nk,nj]()
    var C = Matrix[nj,nm]()
    var D = Matrix[nm,nl]()
    var E = Matrix[ni,nj]()
    var F = Matrix[nj,nl]()
    var G = Matrix[ni,nl]()

    for i in range(ni):
        for j in range(nk):
            A[i,j] = (i * j) / ni
            
    for i in range(nk):
        for j in range(nj):
            B[i,j] = (i * (j + 1)) / nj
            
    for i in range(nj):
        for j in range(nm):
            C[i,j] = (i * (j + 3)) / nl
            
    for i in range(nm):
        for j in range(nl):
            D[i,j] = (i * (j + 2)) / nk

    @always_inline
    @parameter
    fn test_fn():
        _ = func(ni, nj, nk, nl, nm, A, B, C, D, E, F, G)

    var secs = benchmark.run[test_fn](max_runtime_secs=0.5).mean()

    A.data.free()
    B.data.free()
    C.data.free()
    D.data.free()
    E.data.free()
    F.data.free()
    G.data.free()

    return secs

Anyone has some advice on why this is the case? Help would be much appreciated

[Moosems]

That’s very interesting. Would you try this once or twice more and see if this can be replicated? I noticed that in the vanilla test function you have a try except block, why is that no longer in the parameterized one?

[RicRax]

Hi, thank you for answering! I did not include a try except block because it was not there in the example, but then I modified the function to raise an exception and included it later, although it does not make any difference on the result. I have replicated the result many times on my machine running Linux, what other steps should I take?

[Moosems]

Im actually not part of the Mojo developer team but this certainly is interesting. Could you provide some details of your machine? Soft ping @jackos

[RicRax]

I ran this of Linux Mint 21.1, x86 arch, Ryzen 7 4000 series processor

[Moosems]

Interesting. I hope this can get figured out by the devs sometime soon :)

[hogepodge]

I tried to reproduce this, and was getting the opposite result (parameterized running faster than vanilla). I had to make some inference about the code you were using, and I was wondering if you could share a complete example as a single file that demonstrates the behavior, just to be sure that we're working from the same starting point. Thanks!

[RicRax]

example.zip
Hi, here as a single file I included the classes and the functions I am using. Thank you a lot for answering, I am new to Mojo so maybe I am making some naive mistake. Would you share the code you used to get that result too? Thank you!

[hogepodge]

To my eyes, it looks like you're not running the same functions. There's a lot more happening in the second function, including casting to object which is going to have a performance penalty.

[RicRax]

Hi, could you point out what more is happening in the second function? Where am I casting to object?

[RicRax]

Can you please share the code you inferred before? It would really help me out to understand how to do this