Initial implementation for simple scheduler (#589)

* Initial implementation for simple scheduler

* Scheduler assumes that input is `FuncOp` which contains single
  `FuncOp`
* Scheduler constructs dependency graph using TTIR ops.
* This excludes ops like `ttnn.open_device`, `ttnn.close_device`, we
  need to find some way how to handle them i.e either scheduler puts
  them at the start/end of schedule or we handle them in some different way.
* Snapshotting creates deep copy of scheduler members
* Tested by creating pass which printed dependency graph of some more
  complex graph that was handcrafted (didn't include pass nor test code
  in PR but I can do if needed).

* Trying to fix ci

* Debugging Ci

* Second attempt to fix ci

* Redeclare arg

---------

Co-authored-by: Milan Topalovic <mtopalovic994@gmail.com>

.github/Dockerfile.ci

@@ -11,6 +11,9 @@ ENV TTMLIR_TOOLCHAIN_DIR=/opt/ttmlir-toolchain
 RUN mkdir -p $BUILD_DIR && \
     mkdir -p $TTMLIR_TOOLCHAIN_DIR
 
+
+ARG GIT_SHA
+
 # Clone the project and update submodules
 RUN git clone https://github.com/tenstorrent/$PROJECT_NAME.git $BUILD_DIR/$PROJECT_NAME && \
     cd $BUILD_DIR/$PROJECT_NAME && \

env/CMakeLists.txt

@@ -50,6 +50,7 @@ ExternalProject_Add(
     # ======================
     -DLLVM_ENABLE_PIC=ON
     -DLLVM_BUILD_LLVM_DYLIB=ON
+    -DLLVM_INSTALL_GTEST=ON
     -DLLVM_LINK_LLVM_DYLIB=ON
     -DMLIR_BUILD_MLIR_C_DYLIB=ON
     -DMLIR_LINK_MLIR_DYLIB=ON

include/ttmlir/Scheduler/Scheduler.h

@@ -0,0 +1,55 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_SCHEDULER_SCHEDULER_H
+#define TTMLIR_SCHEDULER_SCHEDULER_H
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Operation.h"
+
+namespace mlir::tt::scheduler {
+
+class Scheduler {
+public:
+  // Constructor taking an MLIR Operation (or a module)
+  Scheduler(func::FuncOp *root);
+
+  // Copy constructor
+  Scheduler(const Scheduler &scheduler);
+
+  // Method to get the next set of schedulable operations
+  llvm::SmallVector<mlir::Operation *> getScheduleableOps();
+
+  // Method to check if an operation can be scheduled
+  bool canSchedule(mlir::Operation *op);
+
+  // Method to schedule an operation
+  void scheduleOp(mlir::Operation *op);
+
+  // Method to take a snapshot of the scheduler
+  std::unique_ptr<Scheduler> snapshot();
+
+  // Method to get the scheduled operations
+  llvm::SmallVector<mlir::Operation *> getSchedule() const;
+
+  // Method to check if there are unscheduled operations
+  bool hasUnscheduledOps() const;
+
+private:
+  // Map of scheduled operations
+  llvm::DenseSet<mlir::Operation *> scheduledOpsMap;
+  // Operation schedule in order of execution
+  llvm::SmallVector<mlir::Operation *> schedule;
+  // Set of unscheduled operations
+  llvm::DenseSet<mlir::Operation *> unscheduledOps;
+  // Map of dependencies
+  llvm::DenseMap<mlir::Operation *, llvm::SmallVector<mlir::Operation *>>
+      dependencies;
+};
+
+} // namespace mlir::tt::scheduler
+
+#endif

lib/CMakeLists.txt

@@ -5,6 +5,7 @@ add_subdirectory(CAPI)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
 add_subdirectory(Target)
+add_subdirectory(Scheduler)
 
 # Shared library will include runtime code
 # so we only build it if runtime is enabled

lib/Scheduler/CMakeLists.txt

@@ -0,0 +1,9 @@
+add_mlir_library(MLIRScheduler
+    Scheduler.cpp
+
+    ADDITIONAL_HEADER_DIRS
+    ${PROJECT_SOURCE_DIR}/include/ttmlir/Scheduler
+
+    LINK_LIBS PUBLIC
+    MLIR
+)

lib/Scheduler/Scheduler.cpp

@@ -0,0 +1,86 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "ttmlir/Scheduler/Scheduler.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "ttmlir/Dialect/TTIR/IR/TTIROpsDialect.h.inc"
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/IR/Operation.h>
+
+namespace mlir::tt::scheduler {
+
+bool isTTIROp(mlir::Operation *op) {
+  return isa<ttir::TTIRDialect>(op->getDialect());
+}
+
+// Init the dependencies map of all ops which are TTIR ops
+Scheduler::Scheduler(func::FuncOp *func) {
+  for (auto &op : func->getOps()) {
+    if (isTTIROp(&op)) {
+      dependencies[&op] = {};
+      unscheduledOps.insert(&op);
+    }
+  }
+
+  for (auto &op : func->getOps()) {
+    // Skip non TTIR operations
+    // Skip operations which do not implement DestinationStyleOpInterface
+    if (!isTTIROp(&op)) {
+      continue;
+    }
+
+    OpResult result = op.getResult(0);
+
+    for (mlir::Operation *use : result.getUsers()) {
+      // Skip non TTIR operations
+      // Skip operations which set the result
+      if (isTTIROp(use) && use->getResult(0) != result) {
+        dependencies[use].push_back(&op);
+      }
+    }
+  }
+}
+
+Scheduler::Scheduler(const Scheduler &scheduler)
+    : scheduledOpsMap(scheduler.scheduledOpsMap), schedule(scheduler.schedule),
+      unscheduledOps(scheduler.unscheduledOps),
+      dependencies(scheduler.dependencies) {}
+
+llvm::SmallVector<mlir::Operation *> Scheduler::getScheduleableOps() {
+  llvm::SmallVector<mlir::Operation *> scheduleableOps;
+  for (auto &op : unscheduledOps) {
+    if (canSchedule(op)) {
+      scheduleableOps.push_back(op);
+    }
+  }
+
+  return scheduleableOps;
+}
+
+bool Scheduler::canSchedule(mlir::Operation *op) {
+  for (mlir::Operation *dep : dependencies[op]) {
+    if (!scheduledOpsMap.count(dep)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+void Scheduler::scheduleOp(mlir::Operation *op) {
+  scheduledOpsMap.insert(op);
+  unscheduledOps.erase(op);
+  schedule.push_back(op);
+}
+
+std::unique_ptr<Scheduler> Scheduler::snapshot() {
+  return std::make_unique<Scheduler>(*this);
+}
+
+llvm::SmallVector<mlir::Operation *> Scheduler::getSchedule() const {
+  return schedule;
+}
+
+bool Scheduler::hasUnscheduledOps() const { return !unscheduledOps.empty(); }
+} // namespace mlir::tt::scheduler

test/CMakeLists.txt

@@ -1,28 +1,38 @@
+add_subdirectory(unittests)
+
 llvm_canonicalize_cmake_booleans(
-        MLIR_ENABLE_BINDINGS_PYTHON
+    MLIR_ENABLE_BINDINGS_PYTHON
+)
+
+configure_lit_site_cfg(
+    ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
+    ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
+    MAIN_CONFIG
+    ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
 )
 
 configure_lit_site_cfg(
-        ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in
-        ${CMAKE_CURRENT_BINARY_DIR}/lit.site.cfg.py
-        MAIN_CONFIG
-        ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
+    ${CMAKE_CURRENT_SOURCE_DIR}/unittests/lit.site.cfg.py.in
+    ${CMAKE_CURRENT_BINARY_DIR}/unittests/lit.site.cfg.py
+    MAIN_CONFIG
+    ${CMAKE_CURRENT_SOURCE_DIR}/unittests/lit.cfg.py
 )
 
 set(TTMLIR_TEST_DEPENDS
-        FileCheck count not
-        ttmlir-opt
-        ttmlir-translate
-        )
+    FileCheck count not
+    ttmlir-opt
+    MLIRUnitTests
+    ttmlir-translate
+)
 if(MLIR_ENABLE_BINDINGS_PYTHON AND TTMLIR_ENABLE_BINDINGS_PYTHON)
-  list(APPEND TTMLIR_TEST_DEPENDS TTMLIRPythonModules)
+    list(APPEND TTMLIR_TEST_DEPENDS TTMLIRPythonModules)
 endif()
 
 add_lit_testsuite(check-ttmlir "Running the ttmlir regression tests"
-        ${CMAKE_CURRENT_BINARY_DIR}
-        --xunit-xml-output report.xml
-        DEPENDS ${TTMLIR_TEST_DEPENDS}
-        )
+    ${CMAKE_CURRENT_BINARY_DIR}
+    --xunit-xml-output report.xml
+    DEPENDS ${TTMLIR_TEST_DEPENDS}
+)
 set_target_properties(check-ttmlir PROPERTIES FOLDER "Tests")
 
 add_lit_testsuites(TTMLIRStatic ${CMAKE_CURRENT_SOURCE_DIR} DEPENDS ${TTMLIR_TEST_DEPENDS})

test/unittests/CMakeLists.txt

@@ -0,0 +1,8 @@
+add_custom_target(MLIRUnitTests)
+set_target_properties(MLIRUnitTests PROPERTIES FOLDER "MLIR Tests")
+
+function(add_mlir_unittest test_dirname)
+  add_unittest(MLIRUnitTests ${test_dirname} ${ARGN})
+endfunction()
+
+add_subdirectory(TestScheduler)

test/unittests/TestScheduler/CMakeLists.txt

@@ -0,0 +1,11 @@
+add_mlir_unittest(SchedulerTests
+    TestScheduler.cpp
+)
+
+target_link_libraries(SchedulerTests
+    PRIVATE
+    MLIR
+    MLIRTTDialect
+    MLIRTTIRDialect
+    MLIRScheduler
+)