[Issue]: Model compilation time and size scales with dynamic dimensions' maximum size

[mferencevic]

Problem description

When MIGraphX compiles a model with dynamic dimensions, the compilation time heavily depends on the maximum size of the specified dynamic dimensions.
For the simple reproduction script below, it seems to scale linearly, but we've observed it to scale quadratically with the maximum size of the specified dynamic dimensions for larger models.

Also, the first run of the reproduction script below will produce compilation times (and sizes) as follows:

Compiling model with max size 8 ...
Compilation time: 4.881 s
Model size: 0.119063 MiB

Compiling model with max size 16 ...
Compilation time: 3.184 s
Model size: 0.236793 MiB

Compiling model with max size 32 ...
Compilation time: 7.989 s
Model size: 0.472382 MiB

Compiling model with max size 64 ...
Compilation time: 15.368 s
Model size: 0.943545 MiB

Compiling model with max size 128 ...
Compilation time: 30.448 s
Model size: 1.886622 MiB

Compiling model with max size 256 ...
Compilation time: 47.748 s
Model size: 3.777006 MiB

Compiling model with max size 512 ...
Compilation time: 96.688 s
Model size: 7.562881 MiB

While the second run will produce the following compilation times:

Compiling model with max size 8 ...
Compilation time: 0.396 s
Model size: 0.119063 MiB

Compiling model with max size 16 ...
Compilation time: 0.692 s
Model size: 0.236793 MiB

Compiling model with max size 32 ...
Compilation time: 1.334 s
Model size: 0.472382 MiB

Compiling model with max size 64 ...
Compilation time: 2.672 s
Model size: 0.943545 MiB

Compiling model with max size 128 ...
Compilation time: 5.298 s
Model size: 1.886622 MiB

Compiling model with max size 256 ...
Compilation time: 10.788 s
Model size: 3.777006 MiB

Compiling model with max size 512 ...
Compilation time: 22.186 s
Model size: 7.562881 MiB

We've traced the difference in compilation times to the cache stored in ~/.cache/comgr.
The cache doesn't help much in our specific scenario, because we compile all of our models exactly once.

Regardless, both series of compilation times are very bad and basically make the dynamic dimension support unusable for realistic computer-vision models.

As a side note, the model seems to be internally copied max_seq_len-times, which is very odd.
This is visible from the last print(migraphx_model) statement in the reproduction script.

Steps to reproduce

import math
import migraphx
import os
import time
import torch

DEVICE = "cuda:0"
EMBEDDING_COUNT = 32
EMBEDDING_DIM = 16
BATCH_SIZE = 4

torch.inference_mode(True)
torch.cuda.set_device(DEVICE)

model = torch.nn.Embedding(EMBEDDING_COUNT, EMBEDDING_DIM)
model.eval()
input_batch = torch.arange(math.ceil(EMBEDDING_COUNT / 2)).repeat(BATCH_SIZE, 1).contiguous()

torch.onnx.export(
    model,
    (input_batch,),
    "model.onnx",
    external_data=False,
    dynamo=True,
    dynamic_shapes=[
        {0: torch.export.Dim.DYNAMIC, 1: torch.export.Dim.DYNAMIC},
    ],
)

for max_seq_len in [8, 16, 32, 64, 128, 256, 512]:
    print("Compiling model with max size", max_seq_len, "...")
    migraphx_model = migraphx.parse_onnx("model.onnx", map_dyn_input_dims={
        "input": [
            migraphx.shape.dynamic_dimension(BATCH_SIZE, BATCH_SIZE, {BATCH_SIZE}),
            migraphx.shape.dynamic_dimension(1, max_seq_len, {1}),
        ],
    })
    compilation_time = -time.perf_counter()
    migraphx_model.compile(migraphx.get_target("gpu"), offload_copy=False)
    compilation_time += time.perf_counter()

    migraphx.save(migraphx_model, "model.mxr")
    mxr_size = os.path.getsize("model.mxr") / 1024 / 1024

    print("Compilation time:", f"{compilation_time:.03f}", "s")
    print("Model size:", f"{mxr_size:03f}", "MiB")
    print()

print(migraphx_model)

Environment

OS: Debian GNU/Linux 12 (bookworm)
CPU: AMD Ryzen 9 9950X
GPU: AMD Radeon AI PRO R9700
ROCm version: 7.2.1
MIGraphX version: 2.16.0.dev+20250912-17-406-gb91f1c0c0

[shivadbhavsar]

Yes this is a known problem with existing dynamic shapes, you identified the root cause correctly in that it will recompile the model for every single possible dim value in your dynamic dimension range. There is active work to allow a use case such as this to be compiled once and with roughly similar compile time as for a static input.

There is some preliminary work done for this workflow and can be enabled via MIGRAPHX_ENABLE_FULL_DYNAMIC. However, it is currently not very performant at runtime (and you might encounter some bugs at compile time as well since it is still work in progress and hidden behind an environment variable).

[mferencevic]

Thanks for the information.

We've tried setting MIGRAPHX_ENABLE_FULL_DYNAMIC=1, but it produced the error same as in #4844 even for this single dynamic dimension.

RuntimeError: src/shape.cpp:623: lens: SHAPE: lens() called on a dynamic shape

[shivadbhavsar]

Try with this PR, it should compile: #4853

[mferencevic]

I've tried #4853 and it does produce much better output.

Compiling model with max size 8 ...
Compilation time: 0.027 s
Model size: 0.006422 MiB

Compiling model with max size 16 ...
Compilation time: 0.005 s
Model size: 0.006422 MiB

Compiling model with max size 32 ...
Compilation time: 0.004 s
Model size: 0.006422 MiB

Compiling model with max size 64 ...
Compilation time: 0.004 s
Model size: 0.006422 MiB

Compiling model with max size 128 ...
Compilation time: 0.004 s
Model size: 0.006432 MiB

Compiling model with max size 256 ...
Compilation time: 0.004 s
Model size: 0.006447 MiB

Compiling model with max size 512 ...
Compilation time: 0.004 s
Model size: 0.006447 MiB

module: "main"
@0 = check_context::migraphx::gpu::context -> float_type, {}, {}
@1 = hip::hip_allocate_memory[shape=int8_type, {139264}, {1},id=main:scratch] -> int8_type, {139264}, {1}
@2 = load[offset=131072,end=139264](@1) -> dynamic, int32_type, {4, [1..512]}
input = @param:input -> dynamic, int64_type, {4, [1..512]}
@4 = gpu::dynamic_code_object_op[pre_op=gpu::precompile_op[op=pointwise,additional_args=1,ignore_modules=0,output_shape=nullopt]](input,@2), [main:pointwise0] -> dynamic, int32_type, {4, [1..512]}
@5 = hip::hip_copy_literal[id=main:@literal:0] -> float_type, {32, 16}, {16, 1}
@6 = load[offset=0,end=131072](@1) -> dynamic, float_type, {4, [1..512], 16}
@7 = gpu::dynamic_code_object_op[pre_op=gpu::precompile_op[op=gather[axis=0],additional_args=1,ignore_modules=0,output_shape=nullopt]](@5,@4,@6) -> dynamic, float_type, {4, [1..512], 16}
main:#output_0 = @param:main:#output_0 -> dynamic, float_type, {4, [1..512], 16}
@9 = gpu::dynamic_code_object_op[pre_op=gpu::precompile_op[op=pointwise,additional_args=1,ignore_modules=0,output_shape=nullopt]](@7,main:#output_0), [main:pointwise1] -> dynamic, float_type, {4, [1..512], 16}
@10 = @return(@9)

module: "main:pointwise0"
main:pointwise0:x0 = @param:x0 -> int64_type, {1}, {0}
main:pointwise0:@1 = convert[target_type=8](main:pointwise0:x0) -> int32_type, {1}, {0}
main:pointwise0:@2 = @return(main:pointwise0:@1)

module: "main:pointwise1"
main:pointwise1:x0 = @param:x0 -> float_type, {1}, {0}
main:pointwise1:@1 = convert[target_type=2](main:pointwise1:x0) -> float_type, {1}, {0}
main:pointwise1:@2 = @return(main:pointwise1:@1)

Although, I haven't executed the model to validate whether or not it produces correct results.

[mferencevic]

I've executed the model and the outputs are correct. That's a great initial step, although, for our actual models, this change makes no difference because we still can't compile any of them. This is just an initial list of errors that get triggered:

RuntimeError: src/shape.cpp:368: elements: SHAPE: elements() called on dynamic shape

RuntimeError: src/include/migraphx/op/concat.hpp:124: normalize_compute_shape: CONCAT: all input dimensions should match in axis 0

RuntimeError: src/onnx/checks.cpp:35: check_arg_empty: PARSE_TILE: dynamic shape is not supported

RuntimeError: src/shape.cpp:631: lens: SHAPE: lens() called on a dynamic shape