[BUG]: Slow Max Graphs Compilation

[taalhaataahir0102]

Bug description

MAX graphs takes along time to compile, around 300-500 ms even for smaller graphs as well. The compilation becomes a bit faster on next consecutive code execution due to caching but still the time compilation time for graphs remains above 250 ms. Note: enable_model_ir_cache is already true in modular.cfg file.

Steps to reproduce

Code example:

fn main() raises:
    var session = engine.InferenceSession()
    start = now()
    var graph1 = Graph(in_types=List[Type](TensorType(DType.float32, "a","m")))
    var transposed = ops.transpose(graph1[0],-1,-2)
    graph1.output(transposed)
    graph1.verify()
    var transpose = session.load(graph1)
    end1 = now()
    print("total compile time for graph1: ",(end1 - start)/1000000000)

    var graph2 = Graph(in_types=List[Type](TensorType(DType.float32, "m", "n"), TensorType(DType.float32, "n","x")))
    var out3 = graph2[0] @ graph2[1]
    graph2.output(out3)
    graph2.verify()
    var multiplication = session.load(graph2)
    end2 = now()
    print("total compile time for graph2: ",(end2 - end1)/1000000000)

The output is following:

total compile time for graph1:  0.31278938899999997
total compile time for graph2:  0.451578655

System information

OS: Ubuntu 22.04.3 LTS
Mojo version: mojo 24.5.0 (e8aacb95)
max 24.5.0 (e8aacb95)
Modular version 24.5.0

[ehsanmok]

For the record, here is the discussion in our community discord.