Add op lib selection in test device

Author: JewelRoam

graph_net/torch/test_reference_device.py

@@ -27,6 +27,15 @@ def get_reference_output_path(reference_dir, model_path):
     return os.path.join(reference_dir, f"{model_name}.pth")
 
 
+def register_op_lib(op_lib):
+    if op_lib == "flaggems":
+        import flag_gems
+
+        flag_gems.enable()
+    else:
+        pass
+
+
 def test_single_model(args):
     ref_log = get_reference_log_path(args.reference_dir, args.model_path)
     ref_dump = get_reference_output_path(args.reference_dir, args.model_path)
@@ -51,6 +60,10 @@ def test_single_model(args):
                 test_compiler.get_compile_framework_version(args),
             )
 
+            test_compiler_util.print_with_log_prompt(
+                "[Config] op_lib:", args.op_lib, args.log_prompt
+            )
+
             success = False
             time_stats = {}
             try:
@@ -99,6 +112,7 @@ def test_multi_models(args):
                     f"--model-path {model_path}",
                     f"--compiler {args.compiler}",
                     f"--device {args.device}",
+                    f"--op-lib {args.op_lib}",
                     f"--warmup {args.warmup}",
                     f"--trials {args.trials}",
                     f"--log-prompt {args.log_prompt}",
@@ -136,6 +150,7 @@ def main(args):
     ref_dump_dir.mkdir(parents=True, exist_ok=True)
 
     if path_utils.is_single_model_dir(args.model_path):
+        register_op_lib(args.op_lib)
         test_single_model(args)
     else:
         test_multi_models(args)
@@ -163,6 +178,13 @@ def main(args):
         default="cuda",
         help="Device for testing the compiler (e.g., 'cpu' or 'cuda')",
     )
+    parser.add_argument(
+        "--op-lib",
+        type=str,
+        required=False,
+        default="default",
+        help="Customized operator library (eg. default, flaggems)",
+    )
     parser.add_argument(
         "--warmup", type=int, required=False, default=5, help="Number of warmup steps"
     )

graph_net/torch/test_target_device.py

@@ -143,6 +143,7 @@ def test_multi_models(args):
                     f"-m graph_net.torch.{module_name}",
                     f"--model-path {model_path}",
                     f"--device {args.device}",
+                    f"--op-lib {args.op_lib}",
                     f"--log-prompt {args.log_prompt}",
                     f"--reference-dir {args.reference_dir}",
                 ]
@@ -167,6 +168,18 @@ def main(args):
     assert args.device in ["cuda", "dcu", "xpu", "cpu"]
 
     if path_utils.is_single_model_dir(args.model_path):
+        if args.op_lib == "origin":
+            ref_log = test_reference_device.get_reference_log_path(
+                args.reference_dir, args.model_path
+            )
+            config = parse_config_from_reference_log(ref_log)
+            vars(args)["op_lib"] = config.get("op_lib")
+            test_compiler_util.print_with_log_prompt(
+                "[Config] op_lib:", args.op_lib, args.log_prompt
+            )
+        else:
+            test_reference_device.register_op_lib(args.op_lib)
+
         args = update_args_and_set_seed(args, args.model_path)
         test_single_model(args)
     else:
@@ -194,6 +207,13 @@ def main(args):
         default="cuda",
         help="Device for testing the compiler (e.g., 'cpu' or 'cuda')",
     )
+    parser.add_argument(
+        "--op-lib",
+        type=str,
+        required=False,
+        default="default",
+        help="Customized operator library (eg. default, flaggems or origin)",
+    )
     parser.add_argument(
         "--log-prompt",
         type=str,