Bug fix multiple runtimes (#330)

- For some reason the cmake variables for toggling different runtimes
  were not working for me.  Only the TTNN runtime was able to be
  enabled.
- We also need to set the runtime before we open the device, otherwise
  the device and the rest of the runtime will mismatch.
- For now we also just hardcode ttrt to be single device

CMakeLists.txt

@@ -6,6 +6,7 @@ if (NOT DEFINED ENV{TTMLIR_ENV_ACTIVATED})
 endif()
 
 option(TT_RUNTIME_ENABLE_PERF_TRACE "Enable performance mode" OFF)
+option(TTMLIR_ENABLE_RUNTIME "Enable runtime" OFF)
 
 set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
 

lib/Dialect/TTMetal/Transforms/SerializeToBinary.cpp

@@ -119,7 +119,7 @@ class TTMetalSerializeToBinary
       cqBuilder.name = entry.getSymName().data();
 
       auto argumentAllocations = mlir::cast<ArrayAttr>(
-          entry->getDiscardableAttr("argument_allocations"));
+          entry->getDiscardableAttr(ArgumentAllocationAttr::name));
       assert(argumentAllocations && "expected argument_allocations attribute");
       for (auto &input : entry.getBody().getArguments()) {
         auto argAlloc = mlir::cast<tt::ArgumentAllocationAttr>(

runtime/CMakeLists.txt

@@ -1,12 +1,7 @@
 # Options
 option(TTMLIR_ENABLE_RUNTIME_TESTS "Enable runtime tests" OFF)
-option(TT_RUNTIME_ENABLE_TTNN "Enable TTNN Runtime" OFF)
-option(TT_RUNTIME_ENABLE_TTMETAL "Enable TTMetal Runtime" OFF)
-
-if (NOT TT_RUNTIME_ENABLE_TTNN AND NOT TT_RUNTIME_ENABLE_TTMETAL)
-  # Default to TTNN
-  set(TT_RUNTIME_ENABLE_TTNN ON)
-endif()
+option(TT_RUNTIME_ENABLE_TTNN "Enable TTNN Runtime" ON)
+option(TT_RUNTIME_ENABLE_TTMETAL "Enable TTMetal Runtime" ON)
 
 add_subdirectory(lib)
 add_subdirectory(tools)

runtime/lib/CMakeLists.txt

@@ -1,7 +1,3 @@
-if (TT_RUNTIME_ENABLE_TTMETAL AND TTMLIR_ENABLE_RUNTIME_TESTS)
-  message(FATAL_ERROR "TTMETAL runtime does not support runtime tests yet")
-endif()
-
 if (TTMLIR_ENABLE_RUNTIME)
   if (TT_RUNTIME_ENABLE_TTNN)
     add_subdirectory(ttnn)
@@ -18,7 +14,7 @@ else()
   add_library(TTRuntimeTTMetal INTERFACE)
 endif()
 
-message(STATUS "TT_RUNTIME_ENABLE_TTNN: ${TT_RUNTIME_ENABLE_TTNN} ${TT_RUNTIME_ENABLE_TTMETAL}")
+message(STATUS "Runtimes Enabled: TTNN[${TT_RUNTIME_ENABLE_TTNN}] TTMETAL[${TT_RUNTIME_ENABLE_TTMETAL}]")
 add_library(TTRuntime STATIC runtime.cpp)
 set_property(TARGET TTRuntime PROPERTY CXX_STANDARD 20)
 if (TTMLIR_ENABLE_RUNTIME AND TT_RUNTIME_ENABLE_TTNN)

runtime/tools/python/ttrt/common/api.py

@@ -178,13 +178,14 @@ def run(args):
 
     # execution
     print("executing action for all provided flatbuffers")
-    device = ttrt.runtime.open_device(device_ids)
+    # TODO: sort the flatbuffers by runtime type, for now just take the first one
+    ttrt.runtime.set_compatible_runtime(fbb_list[0][1])
+    device = ttrt.runtime.open_device([device_ids[0]])
     atexit.register(lambda: ttrt.runtime.close_device(device))
 
     torch.manual_seed(args.seed)
 
     for (binary_name, fbb, fbb_dict, program_indices) in fbb_list:
-        ttrt.runtime.set_compatible_runtime(fbb)
         torch_inputs[binary_name] = {}
         torch_outputs[binary_name] = {}
         for program_index in program_indices: