From 6171aed7dd7b2125fc0b72c56e3d3773a888a525 Mon Sep 17 00:00:00 2001
From: Nikola Vukobrat <124874832+nvukobratTT@users.noreply.github.com>
Date: Tue, 5 Nov 2024 17:27:48 +0100
Subject: [PATCH] Option to setup compiler configurations using pytest
 arguments (#601)

Fix #600
---
 forge/test/conftest.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/forge/test/conftest.py b/forge/test/conftest.py
index 722254c6..90431b71 100644
--- a/forge/test/conftest.py
+++ b/forge/test/conftest.py
@@ -29,6 +29,8 @@
 from forge.verify.config import TestKind
 from forge.torch_compile import reset_state
 
+from forge.config import _get_global_compiler_config
+
 collect_ignore = ["legacy_tests"]
 
 
@@ -98,6 +100,9 @@ def clear_forge():
 
 
 def pytest_addoption(parser):
+    parser.addoption(
+        "--generate-op-tests", action="store_true", default=False, help="Generate op tests for the given model"
+    )
     parser.addoption(
         "--silicon-only", action="store_true", default=False, help="run silicon tests only, skip golden/model"
     )
@@ -166,6 +171,25 @@ def runslow(request):
 """
 
 
+@pytest.fixture(autouse=True, scope="session")
+def initialize_global_compiler_configuration_based_on_pytest_args(pytestconfig):
+    """
+    Set the global compiler config options for the test session
+    which will generate op tests for the given model.
+    """
+    compiler_cfg = _get_global_compiler_config()
+
+    compiler_cfg.tvm_generate_op_tests = pytestconfig.getoption("--generate-op-tests")
+
+    if compiler_cfg.tvm_generate_op_tests:
+        # For running standalone tests, we need to retain the generated python files
+        # together with stored model parameters
+        compiler_cfg.retain_tvm_python_files = True
+
+        # Required to prevent early tensor deallocation
+        compiler_cfg.enable_op_level_comparision = True
+
+
 @pytest.hookimpl(tryfirst=True)
 def pytest_cmdline_preparse(config, args):