Add mutexes to protect preemtive start/exit of compile thread

results/parse_op_by_op_results.py

@@ -181,7 +181,6 @@ def process_json_files():
                     if result.returncode != 0:
                         error = result.stderr.split("\n")[0]
                         trace_dump = result.stderr
-                        print(error)
                 row_data = [
                     name,
                     input_shapes,

tests/models/codegen/test_codegen.py

@@ -31,10 +31,10 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(
+    reason="Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
+)
 def test_codegen(record_property, mode):
-    pytest.xfail(
-        "Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
-    )
     model_name = "codegen"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/flan_t5/test_flan_t5.py

@@ -31,10 +31,10 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(
+    reason="Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
+)
 def test_flan_t5(record_property, mode):
-    pytest.xfail(
-        "Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
-    )
     model_name = "FLAN-T5"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/gpt_neo/test_gpt_neo.py

@@ -39,10 +39,10 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(
+    reason="Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
+)
 def test_gpt_neo(record_property, mode):
-    pytest.xfail(
-        "Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
-    )
     model_name = "GPTNeo"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/hand_landmark/test_hand_landmark.py

@@ -41,8 +41,8 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(reason="Need to debud")
 def test_hand_landmark(record_property, mode):
-    pytest.xfail("Need to debug")
     model_name = "Hand Landmark"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/mobilenet_ssd/test_mobilenet_ssd.py

@@ -37,8 +37,8 @@ def _load_inputs(self):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(reason="Need to debug")
 def test_mobilenet_ssd(record_property, mode):
-    pytest.xfail("Need to debug")
     model_name = "MobileNetSSD"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/opt/test_opt.py

@@ -35,8 +35,8 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(reason="Need to debug")
 def test_opt(record_property, mode):
-    pytest.xfail("Need to debug")
     model_name = "OPT"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/t5/test_t5.py

@@ -28,11 +28,11 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(
+    reason="Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
+)
 @pytest.mark.parametrize("model_name", ["t5-small", "t5-base", "t5-large"])
 def test_t5(record_property, model_name, mode):
-    pytest.xfail(
-        "Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
-    )
     record_property("model_name", model_name)
     record_property("mode", mode)
 

tests/models/whisper/test_whisper.py

@@ -48,10 +48,10 @@ def set_model_eval(self, model):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(
+    reason="Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
+)
 def test_whisper(record_property, mode):
-    pytest.xfail(
-        "Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
-    )
     model_name = "Whisper"
     record_property("model_name", model_name)
     record_property("mode", mode)

tests/models/yolov5/test_yolov5.py

@@ -106,10 +106,8 @@ def _load_inputs(self):
     "mode",
     ["eval"],
 )
+@pytest.mark.xfail(reason="Fails due to pt2 compile issue")
 def test_yolov5(record_property, mode):
-    pytest.xfail(
-        "Fails due to pt2 compile issue when finishing generation, but we can still generate a graph"
-    )
     model_name = "YOLOv5"
     record_property("model_name", model_name)
     record_property("mode", mode)

tt_torch/dynamo/backend.py

@@ -101,27 +101,27 @@ def lower_to_stable_hlo(module, op=None):
         op.compilation_status = OpCompilationStatus.CONVERTED_TO_STABLE_HLO
 
 
-def compile_process(receiver, sender):
+def compile_process(receiver, sender, ttir_envent, ttnn_event):
     obj = receiver.get()
     faulthandler.disable()
     asm = obj["asm"]
     ttir = tt_mlir.compile_stable_hlo_to_ttir(asm)
     sender.put({"ttir": ttir})
-    time.sleep(0.1)
+    ttir_envent.wait()
     binary, ttnn = tt_mlir.compile_ttir_to_bytestream(ttir)
     sender.put({"binary": binary, "ttnn": ttnn})
-    time.sleep(0.1)
+    ttnn_event.wait()
     sys.exit(0)
 
 
-def execute_process(receiver, sender):
+def execute_process(receiver, sender, exec_event):
     obj = receiver.get()
     faulthandler.disable()
     binary = obj["binary"]
     inputs = obj["inputs"]
     outputs = tt_mlir.run(inputs, binary)
     sender.put({"outputs": outputs})
-    time.sleep(0.1)
+    exec_event.wait()
     sys.exit(0)
 
 
@@ -242,8 +242,12 @@ def compile_op(self, node, *inputs, **kwargs):
 
         sender = mp.Queue()
         receiver = mp.Queue()
+        ttir_event = mp.Event()
+        ttnn_event = mp.Event()
         obj = {"asm": module.operation.get_asm()}
-        process = mp.Process(target=compile_process, args=(sender, receiver))
+        process = mp.Process(
+            target=compile_process, args=(sender, receiver, ttir_event, ttnn_event)
+        )
         process.start()
         sender.put(obj)
         start = time.time()
@@ -254,10 +258,12 @@ def compile_op(self, node, *inputs, **kwargs):
                 if "ttir" in result:
                     op.compilation_status = OpCompilationStatus.CONVERTED_TO_TTIR
                     op.add_ttir_graph(result["ttir"])
+                    ttir_event.set()
                 if "binary" in result:
                     binary = result["binary"]
                     op.binary = binary
                     op.add_ttnn_graph(result["ttnn"])
+                    ttnn_event.set()
                     op.compilation_status = OpCompilationStatus.CONVERTED_TO_TTNN
                     break
             except mp.queues.Empty:
@@ -276,7 +282,10 @@ def run_op(self, binary, *inputs):
         receiver = mp.Queue()
         obj = {"binary": binary, "inputs": inputs}
 
-        process = mp.Process(target=execute_process, args=(sender, receiver))
+        exec_event = mp.Event()
+        process = mp.Process(
+            target=execute_process, args=(sender, receiver, exec_event)
+        )
         process.start()
         sender.put(obj)
         result = {}
@@ -286,6 +295,7 @@ def run_op(self, binary, *inputs):
                 break
             try:
                 result = receiver.get_nowait()
+                exec_event.set()
                 break
             except mp.queues.Empty:
                 pass