[Issue]: Compilation of swin_transformer fails when MIGRAPHX_MLIR_USE_SPECIFIC_OPS=attention is set,

[menglcai]

Problem Description

Compilation of swin_transformer fails when MIGRAPHX_MLIR_USE_SPECIFIC_OPS=attention is set, but succeeds when the flag is not set.

[O] PYTHONPATH=/opt/rocm/lib:$PYTHONPATH python test.py
[X] PYTHONPATH=/opt/rocm/lib:$PYTHONPATH MIGRAPHX_MLIR_USE_SPECIFIC_OPS=attention python test.py

The error message is as follows: run_high_level_pipeline: Invalid MLIR created: Error: Cannot trace first gemm index for linalg.generic op

Tested with the following system configuration on MI250:

• docker image: rocm/pytorch:rocm7.2_ubuntu24.04_py3.12_pytorch_release_2.9.1
• torch_migraphx : master branch, ROCm/torch_migraphx@b94b985
• python version : 2.9.1+rocm7.2.0.lw.git7e1940d4
• [X] MIGraphX Version: 2.15.0.20250912-17-195-g1afd1b89c (apt install)
• [X] MIGraphX Version: 2.15.0.20250912-17-220-g07ef4ba4f (rbuild, rocm-7.2.3 branch)

Operating System

Ubuntu 24.04.3 LTS

CPU

AMD EPYC 7763 64-Core Processor

GPU

AMD Instinct MI250

Other

No response

ROCm Version

ROCm 6.0.0

Steps to Reproduce

The test code:

import argparse
import torch
import torchvision

def compile_mgx(model):
    import torch_migraphx
    return torch.compile(model, backend="migraphx", options={"deallocate": True, "fp16": True}, dynamic=False)

if __name__ == "__main__":
    weights = torchvision.models.Swin_T_Weights.DEFAULT
    model = torchvision.models.swin_t(weights=weights)
    model.to("cuda").eval()
    model = compile_mgx(model)

    # Inference
    with torch.no_grad():
        tensor = torch.rand(1, 3, 224, 224).to("cuda")  # fake input
        output = model(tensor)

    print("finish")

Run following command:

# [O] MIGraphX compile pass
PYTHONPATH=/opt/rocm/lib:$PYTHONPATH python test.py

# [X] MIGraphX compile fail
MIGRAPHX_MLIR_USE_SPECIFIC_OPS=attention  PYTHONPATH=/opt/rocm/lib:$PYTHONPATH python test.py

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response