Python-generated schedules (just mirroring pipeline bundles) (#997)

Invoke either as `build/bin/tpp-sched` or `python -m mlir.tpp.sched`
(which requires `PYTHONPATH` to point to `build/python_packages` - the
`tools/tpp-sched/tpp-sched.py` script (installed as `tpp-sched`) takes
care of that).

When `tpp-sched` is invoked it just dumps the default pipeline schedule
to stdout. When a path is provided as a positional argument, i.e.
`tpp-sched.py PAYLOAD_FILE`, the schedule gets generated and immediately
applied to the payload and only the transformed payload gets dumped to
stdout. Use `--bundles=BUNDLE1,...,BUNDLEN` to control which bundles
make up the schedule. All the bundles' arguments available on `tpp-opt`
are available on `tpp-sched`.

When invoked as `tpp-sched --payload PAYLOAD_FILE`, it will just print
the payload and schedule together as one IR dump. This is meant to
facilitate working directly with the IR before passing it on to, e.g.,
`tpp-opt` (to apply the schedule to the payload use `build/bin/tpp-sched
--payload PAYLOAD_FILE | build/bin/tpp-opt --load-tpp-dialects
--transform-interpreter` -- as to why we need `--load-tpp-dialects`, see
[here](https://discourse.llvm.org/t/dynamic-pass-scheduling-and-dialect-loading/85166/2)).

Introduces Python-bindings for our own dialects, for now mostly to be
able to register them from Python before running passes that dependent
on them.

Requires MLIR to be build with `-DMLIR_ENABLE_BINDINGS_PYTHON=1`.

Adds RUN lines which use `tpp-sched` to each test which used `tpp-opt`
to apply a pipeline bundle.

For now we only have partial support for the GPU pipeline bundles.

Author: rolfmorel

CMakeLists.txt

@@ -43,6 +43,28 @@ endif()
 set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR})
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake/modules")
 
+# Nanobind for Python bindings
+if (CMAKE_VERSION VERSION_LESS 3.18)
+  set(PYTHON_DEV_MODULE Development)
+else()
+  set(PYTHON_DEV_MODULE Development.Module)
+endif()
+
+find_package(Python 3.8 COMPONENTS Interpreter ${PYTHON_DEV_MODULE} REQUIRED)
+
+# Nanobind docs say "Without this addition, binding code may run slowly and produce large binaries.":
+if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+  set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
+  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
+endif()
+
+# Detect the pip/conda-installed nanobind package and import it into CMake
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT)
+find_package(nanobind CONFIG REQUIRED)
+message (STATUS "Nanobind found")
+
 set(TPP_MAIN_INCLUDE_DIR "${PROJECT_SOURCE_DIR}/include")
 set(TPP_GEN_INCLUDE_DIR "${PROJECT_BINARY_DIR}/include")
 
@@ -71,6 +93,7 @@ include(sanitizers)
 # Sub directories
 add_subdirectory(include)
 add_subdirectory(lib)
+add_subdirectory(python)
 add_subdirectory(runtime)
 add_subdirectory(tools)
 add_subdirectory(test)

README.md

@@ -57,10 +57,11 @@ cmake -G Ninja ../llvm \
    -DLLVM_ENABLE_ASSERTIONS=ON \
    -DCMAKE_C_COMPILER=clang \
    -DCMAKE_CXX_COMPILER=clang++ \
+   -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
    -DLLVM_USE_LINKER=lld
 
 # Build
-ninja 
+ninja
 
 popd
 ```
@@ -97,7 +98,7 @@ popd
 
 To enable experimental GPU support see: [GPU/README.md](lib/TPP/GPU/README.md)
 
-In the example above, we are building both LLVM/MLIR and tpp-mlir in relese mode. You can easily change the build type by adopting the `-DCMAKE_BUILD_TYPE` option, e.g. `=DCMAKE_BUILD_TYPE=RelWithDebInfo`. 
+In the example above, we are building both LLVM/MLIR and tpp-mlir in relese mode. You can easily change the build type by adopting the `-DCMAKE_BUILD_TYPE` option, e.g. `=DCMAKE_BUILD_TYPE=RelWithDebInfo`.
 
 ### Conda Environment
 

include/TPP-CAPI/Dialects.h

@@ -0,0 +1,18 @@
+#ifndef TPP_CAPI_DIALECTS_H
+#define TPP_CAPI_DIALECTS_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Check, check);
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Xsmm, xsmm);
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Perf, perf);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // TPP_CAPI_DIALECTS_H

include/TPP/Passes.h

@@ -22,6 +22,10 @@ namespace arith {
 class ArithDialect;
 } // namespace arith
 
+namespace async {
+class AsyncDialect;
+} // namespace async
+
 namespace check {
 class CheckDialect;
 } // namespace check
@@ -56,6 +60,10 @@ namespace memref {
 class MemRefDialect;
 } // namespace memref
 
+namespace omp {
+class OpenMPDialect;
+} // namespace omp
+
 namespace perf {
 class PerfDialect;
 } // namespace perf

include/TPP/Passes.td

@@ -11,6 +11,23 @@
 
 include "mlir/Pass/PassBase.td"
 
+def LoadTppDialects : Pass<"load-tpp-dialects", "ModuleOp"> {
+  let summary = "Pre-load all TPP-specific dialects";
+  let description = [{
+    Pre-load dialects that -transform-interpreter would try to load at runtime.
+
+    The issue is that -transform-interpreter runs inside the multi-threaded
+    passmanager. Hence when the interpreter dynamically tries to load dependent
+    dialects this triggers an assert as loading during multi-threaded execution
+    could lead to concurrency issues.
+  }];
+  let dependentDialects = ["xsmm::XsmmDialect",
+                           "check::CheckDialect",
+                           "perf::PerfDialect",
+                           "omp::OpenMPDialect",
+                           "async::AsyncDialect"];
+}
+
 def ConvertLinalgToXsmm : Pass<"convert-linalg-to-xsmm", "func::FuncOp"> {
   let summary = "Convert linalg to xsmm";
   let description = [{

lib/CMakeLists.txt

@@ -1,3 +1,4 @@
 set(CMAKE_COMPILE_WARNING_AS_ERROR ON)
 
 add_subdirectory(TPP)
+add_subdirectory(TPP-CAPI)

lib/TPP-CAPI/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_mlir_public_c_api_library(TPPCAPI
+  Dialects.cpp
+  LINK_LIBS PUBLIC
+  TPPCheckDialect
+  TPPPerfDialect
+  TPPXsmmDialect
+  TPPTransforms
+  TPPPipeline
+  tpp_xsmm_runner_utils
+)

lib/TPP-CAPI/Dialects.cpp

@@ -0,0 +1,12 @@
+#include "mlir/CAPI/Registration.h"
+
+#include "TPP-CAPI/Dialects.h"
+#include "TPP/Dialect/Check/CheckDialect.h"
+#include "TPP/Dialect/Perf/PerfDialect.h"
+#include "TPP/Dialect/Xsmm/XsmmDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Check, check, mlir::check::CheckDialect)
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Xsmm, xsmm, mlir::xsmm::XsmmDialect)
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Perf, perf, mlir::perf::PerfDialect)

lib/TPP/CMakeLists.txt

@@ -12,6 +12,7 @@ get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 add_mlir_library(TPPPipeline
   DefaultPipeline.cpp
   DefaultTppPasses.cpp
+  LoadTppDialects.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/TPP

lib/TPP/LoadTppDialects.cpp

@@ -0,0 +1,38 @@
+//===- LoadTppDialects.cpp -----------------------------------------*- C++-*-===//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Pass is a no-op as it is only used for the side-effect of loading dialects.
+//
+//===----------------------------------------------------------------------===//
+#include "mlir/Pass/Pass.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "TPP/Dialect/Check/CheckDialect.h"
+#include "TPP/Dialect/Perf/PerfDialect.h"
+#include "TPP/Dialect/Xsmm/XsmmDialect.h"
+#include "mlir/Dialect/Async/IR/Async.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+
+
+namespace mlir {
+namespace tpp {
+#define GEN_PASS_DEF_LOADTPPDIALECTS
+#include "TPP/Passes.h.inc"
+} // namespace tpp
+} // namespace mlir
+
+using namespace mlir;
+using namespace std;
+
+namespace mlir {
+namespace tpp {
+struct LoadTppDialects
+    : public impl::LoadTppDialectsBase<LoadTppDialects> {
+  void runOnOperation() override {}
+};
+} // namespace tpp
+} // namespace mlir
+