Sync FE and BE apis

python/TTNNModule.cpp

@@ -164,6 +164,10 @@ void populateTTNNModule(py::module &m) {
                                }
                                return static_cast<uint32_t>(
                                    self.getMemLayout().getValue());
-                             });
+                             })
+      .def_property_readonly("is_tiled", &tt::ttnn::TTNNLayoutAttr::isTiled)
+      // TODO fix
+      .def_property_readonly("data_type",
+                             &tt::ttnn::TTNNLayoutAttr::getDataType);
 }
 } // namespace mlir::ttmlir::python

python/Util.cpp

@@ -3,6 +3,8 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include "ttmlir/Bindings/Python/TTMLIRModule.h"
+#include <pybind11/pytypes.h>
+#include <variant>
 
 namespace mlir::ttmlir::python {
 
@@ -17,25 +19,26 @@ void populateUtilModule(py::module &m) {
     return source;
   });
 
-  m.def("get_loc_name", [](MlirLocation _loc) -> std::string {
-    mlir::Location loc = unwrap(_loc);
-    if (mlir::isa<mlir::NameLoc>(loc)) {
-      mlir::NameLoc nameLoc = mlir::cast<mlir::NameLoc>(loc);
-      return nameLoc.getName().str();
-    }
-    return "-";
-  });
+  m.def("get_loc_name",
+        [](MlirLocation _loc) -> std::variant<std::string, py::object> {
+          mlir::Location loc = unwrap(_loc);
+          if (mlir::isa<mlir::NameLoc>(loc)) {
+            mlir::NameLoc nameLoc = mlir::cast<mlir::NameLoc>(loc);
+            return nameLoc.getName().str();
+          }
+          return py::none();
+        });
 
-  m.def("get_loc_full", [](MlirLocation _loc) -> std::string {
-    mlir::Location loc = unwrap(_loc);
-    if (mlir::isa<mlir::FileLineColLoc>(loc)) {
-      mlir::FileLineColLoc fileLoc = mlir::cast<mlir::FileLineColLoc>(loc);
-      return fileLoc.getFilename().str() + ":" +
-             std::to_string(fileLoc.getLine()) + ":" +
-             std::to_string(fileLoc.getColumn());
-    }
-    return "-";
-  });
+  m.def("get_loc_full",
+        [](MlirLocation _loc) -> std::variant<std::string, py::object> {
+          mlir::Location loc = unwrap(_loc);
+
+          std::string locationStr;
+          llvm::raw_string_ostream output(locationStr);
+          loc.print(output);
+
+          return locationStr;
+        });
 }
 
 } // namespace mlir::ttmlir::python

tools/explorer/CMakeLists.txt

@@ -3,7 +3,7 @@ include(ExternalProject)
 set(TT_EXPLORER_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/run.py)
 set(TTMLIR_BUILD_BIN_DIR ${TTMLIR_BINARY_DIR}/bin)
 
-set(MODEL_EXPLORER_VERSION "ca884d5eb3291507e7f4e76776957e231b2d9b6d")
+set(MODEL_EXPLORER_VERSION "8ec112eaee8006301039ee34c55d3751a1b82c14")
 ExternalProject_Add(
   model-explorer
   PREFIX ${CMAKE_CURRENT_SOURCE_DIR}/model-explorer

tools/explorer/test/run_tests.py

@@ -16,8 +16,9 @@
     "test/ttmlir/Dialect/TTNN/optimizer/mnist_sharding.mlir",
     "tools/explorer/test/models/*.mlir",
 ]
+MNIST_SHARDING_PATH = "test/ttmlir/Silicon/TTNN/optimizer/mnist_sharding.mlir"
 TEST_EXECUTE_MODEL_PATHS = [
-    "test/ttmlir/Silicon/TTNN/optimizer/mnist_sharding.mlir",
+    MNIST_SHARDING_PATH,
 ]
 
 
@@ -28,7 +29,7 @@ def get_test_files(paths):
     return files
 
 
-def send_command(command, model_path, settings):
+def send_command(command, model_path, settings={}):
     cmd = {
         "extensionId": "tt_adapter",
         "cmdId": command,
@@ -51,7 +52,7 @@ def execute_command(model_path, settings):
 def wait_for_execution_to_finish(timeout):
     for _ in range(timeout):
         try:
-            response = send_command("status_check", "", {})
+            response = send_command("status_check", "")
             if response.status_code == 200 and response.json().get("graphs")[0].get(
                 "isDone"
             ):
@@ -60,9 +61,7 @@ def wait_for_execution_to_finish(timeout):
             print(f"Request failed: {e}")
             raise Exception("Status check request failed")
         time.sleep(1)
-    raise RuntimeError(
-        f"Execution did not finish within {MODEL_EXECUTION_TIMEOUT} seconds"
-    )
+    raise RuntimeError(f"Execution did not finish within {timeout} seconds")
 
 
 def execute_command_and_wait(model_path, settings, timeout):
@@ -107,7 +106,7 @@ def server_shutdown():
 
 @pytest.mark.parametrize("model_path", get_test_files(TEST_LOAD_MODEL_PATHS))
 def test_load_model(model_path):
-    result = send_command("convert", model_path, {})
+    result = send_command("convert", model_path)
     assert result.ok
     if "error" in result.json():
         print(result.json())
@@ -119,22 +118,25 @@ def test_execute_model(model_path):
     execute_command_and_wait(
         model_path, {"optimizationPolicy": "DF Sharding"}, timeout=60
     )
+    send_command("convert", model_path)
 
 
 def test_execute_mnist_l1_interleaved():
     execute_command_and_wait(
-        "test/ttmlir/Silicon/TTNN/optimizer/mnist_sharding.mlir",
+        MNIST_SHARDING_PATH,
         {"optimizationPolicy": "Greedy L1 Interleaved"},
         timeout=60,
     )
+    send_command("convert", MNIST_SHARDING_PATH)
 
 
 def test_execute_mnist_optimizer_disabled():
     execute_command_and_wait(
-        "test/ttmlir/Silicon/TTNN/optimizer/mnist_sharding.mlir",
+        MNIST_SHARDING_PATH,
         {"optimizationPolicy": "Optimizer Disabled"},
         timeout=60,
     )
+    send_command("convert", MNIST_SHARDING_PATH)
 
 
 def test_execute_model_invalid_policy():

tools/explorer/tt_adapter/src/tt_adapter/main.py

@@ -44,25 +44,39 @@ def __init__(self):
     def convert(
         self, model_path: str, settings: Dict
     ) -> model_explorer.ModelExplorerGraphs:
-        perf_trace = None
-        if optimized_model_path := self.model_runner.get_optimized_model_path():
+        if optimized_model_path := self.model_runner.get_optimized_model_path(
+            model_path
+        ):
             print(f"Using optimized model: {optimized_model_path}")
-            model_path = optimized_model_path
-
             # Get performance results.
-            perf_trace = self.model_runner.get_perf_trace()
+            perf_trace = self.model_runner.get_perf_trace(model_path)
+
+            module = utils.parse_mlir_file(optimized_model_path)
+
+            # Convert TTIR to Model Explorer Graphs and Display/Return
+            graph, perf_data = mlir.build_graph(module, perf_trace)
+            if perf_data:
+                # TODO(odjuricic) We can probably edit the actual graph response or create our own instead of just adding to dataclass.s
+                graph = utils.add_to_dataclass(
+                    graph, "overlays", {"Performance Trace": perf_data.graphsData}
+                )
+
+            if overrides := self.model_runner.get_overrides(model_path):
+                graph = utils.add_to_dataclass(graph, "overrides", overrides)
+        else:
+            module = utils.parse_mlir_file(model_path)
 
-        module = utils.parse_mlir_file(model_path)
+            # Convert TTIR to Model Explorer Graphs and Display/Return
+            graph, _ = mlir.build_graph(module)
 
-        # Convert TTIR to Model Explorer Graphs and Display/Return
-        graph, perf_data = mlir.build_graph(module, perf_trace)
-        if perf_data:
-            graph = utils.add_to_dataclass(graph, "perf_data", perf_data.graphsData)
         return {"graphs": [graph]}
 
     def execute(
         self, model_path: str, settings: Dict
     ) -> model_explorer.ModelExplorerGraphs:
+
+        print("SETTINGS: ", settings)
+
         override_handler = optimizer_overrides.OptimizerOverridesHandler()
         override_handler.set_system_desc_path(
             f"{self.model_runner.get_artifacts_dir()}/system_desc.ttsys"
@@ -84,7 +98,9 @@ def execute(
                 OPTIMIZATION_POLICIES[optimization_policy]
             )
 
-        self.model_runner.run(model_path, override_handler.to_string())
+        self.model_runner.run(
+            model_path, override_handler.to_string(), settings.get("overrides", None)
+        )
 
         return {"graphs": []}