From 26a499323fc1ff0ef0c6fea3ed7b28aac8ce8d3b Mon Sep 17 00:00:00 2001
From: Ted Themistokleous
 <107195283+TedThemistokleous@users.noreply.github.com>
Date: Wed, 21 Aug 2024 10:22:42 -0400
Subject: [PATCH] [MIGraphX EP Support] Update migx scripts (#21806)

### Description
<!-- Describe your changes. -->
No code changes to the EP only changes to the scripts whihc invoke
MIGraphX EP

- One case be explicit to set MIGraphX EP when running gpt2 testing
- The other to ensure we turn off optimizations like tensorRT and allow
MIGraphX to handle graph optimizations


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
MIGraphX has moved away from using rocBLAS and without this, some cases
used in CI shall fail as optmizations will attempt to use rocBLAS
kernels instead of MIGraphx EP directly.
---
 onnxruntime/python/tools/transformers/benchmark.py       | 9 +++++++++
 .../tools/transformers/models/gpt2/convert_to_onnx.py    | 6 +++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/python/tools/transformers/benchmark.py b/onnxruntime/python/tools/transformers/benchmark.py
index 4800c48744236..dc0bb55212e28 100644
--- a/onnxruntime/python/tools/transformers/benchmark.py
+++ b/onnxruntime/python/tools/transformers/benchmark.py
@@ -135,6 +135,15 @@ def run_onnxruntime(
             )
             return results
 
+    if provider == "migraphx":
+        optimizer_info = OptimizerInfo.NOOPT
+        warm_up_repeat = 5
+        if "MIGraphXExecutionProvider" not in onnxruntime.get_available_providers():
+            logger.error(
+                "Please install onnxruntime-rocm package, and use a machine with GPU for testing gpu performance."
+            )
+            return results
+
     if optimizer_info == OptimizerInfo.NOOPT:
         logger.warning(
             f"OptimizerInfo is set to {optimizer_info}, graph optimizations specified in FusionOptions are not applied."
diff --git a/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py b/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py
index 0ab26308295a9..75887cc744081 100644
--- a/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py
+++ b/onnxruntime/python/tools/transformers/models/gpt2/convert_to_onnx.py
@@ -370,8 +370,12 @@ def main(argv=None, experiment_name: str = "", run_id: str = "0", csv_filename:
     logger.info(f"Output path: {output_path}")
     model_size_in_MB = int(get_onnx_model_size(output_path, args.use_external_data_format) / 1024 / 1024)  # noqa: N806
 
+    provider = args.provider
+    if args.provider == "migraphx":
+        provider = "MIGraphXExecutionProvider"
+
     session = create_onnxruntime_session(
-        output_path, args.use_gpu, args.provider, enable_all_optimization=True, verbose=args.verbose
+        output_path, args.use_gpu, provider, enable_all_optimization=True, verbose=args.verbose
     )
     if args.model_class == "GPT2LMHeadModel" and session is not None:
         parity_result = gpt2helper.test_parity(