[BACKEND] Add support for out of tree TTIR/TTGIR passes (#8401)

Follow on to: triton-lang/triton#8137. This PR
coupled with the Triton pass pipeline override adds the ability to call
out of tree TTIR and TTGIR MLIR passes. Integration has been done with
both the Python and triton-opt interfaces.

Co-authored w/ @plotfi

---------

Co-authored-by: Puyan Lotfi <[email protected]>

Author: CRobeck

.gitignore

@@ -43,6 +43,9 @@ pytest.ini
 # Instrumentation
 python/triton/instrumentation
 
+# MLIR Plugin
+python/triton/plugins
+
 # Python caches
 __pycache__/
 *.py[cod]

CMakeLists.txt

@@ -20,6 +20,8 @@ option(TRITON_BUILD_PYTHON_MODULE "Build Python Triton bindings" OFF)
 option(TRITON_BUILD_PROTON "Build the Triton Proton profiler" ON)
 option(TRITON_BUILD_UT "Build C++ Triton Unit Tests" ON)
 option(TRITON_BUILD_WITH_CCACHE "Build with ccache (if available)" ON)
+option(LLVM_BUILD_SHARED_LIBS
+  "Build all libraries as shared libraries instead of static" OFF)
 set(TRITON_CODEGEN_BACKENDS "" CACHE STRING "Enable different codegen backends")
 
 if(TRITON_BUILD_WITH_CCACHE)
@@ -60,6 +62,7 @@ else()
   set(CMAKE_EXE_LINKER_FLAGS_TRITONRELBUILDWITHASSERTS "/debug:fastlink /INCREMENTAL")
   set(CMAKE_MODULE_LINKER_FLAGS_TRITONRELBUILDWITHASSERTS "/debug:fastlink /INCREMENTAL")
   set(CMAKE_SHARED_LINKER_FLAGS_TRITONRELBUILDWITHASSERTS "/debug:fastlink /INCREMENTAL")
+  set(LLVM_BUILD_SHARED_LIBS "0")
 endif()
 
 # Default build type

Makefile

@@ -43,6 +43,8 @@ test-unit: all
 	$(PYTEST) --tb=short -vs python/examples/gluon/01-attention-forward.py
 	TRITON_ALWAYS_COMPILE=1 TRITON_DISABLE_LINE_INFO=0 LLVM_PASS_PLUGIN_PATH=python/triton/instrumentation/libGPUInstrumentationTestLib.so \
 		$(PYTEST) --capture=tee-sys -rfs -vvv python/test/unit/instrumentation/test_gpuhello.py
+	TRITON_PASS_PLUGIN_PATH=python/triton/plugins/libTritonPluginsTestLib.so \
+		$(PYTEST) -vvv python/test/unit/plugins/test_plugin.py
 	$(PYTEST) --tb=short -s -n $(NUM_PROCS) python/test/gluon
 
 .PHONY: test-distributed

README.md

@@ -272,7 +272,7 @@ def inspect_stages(_self, stages, options, language, capability):
     # inspect or modify add_stages here
 triton.knobs.runtime.add_stages_inspection_hook = inspect_stages
 ```
-
+Examples of how to use this for out of tree plugin passes is [here](lib/Plugins/README.md)
 # Changelog
 
 Version 2.0 is out! New features include:

bin/RegisterTritonDialects.h

@@ -45,6 +45,9 @@
 #include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
 #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
 
+#include "triton/Tools/PluginUtils.h"
+#include "triton/Tools/Sys/GetEnv.hpp"
+
 namespace mlir {
 namespace test {
 void registerTestAliasPass();
@@ -136,6 +139,21 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::triton::proton::gpu::registerScheduleBufferStorePass();
   mlir::triton::proton::gpu::registerAddSchedBarriersPass();
 
+  // Plugin passes
+  if (std::string filename =
+          mlir::triton::tools::getStrEnv("TRITON_PASS_PLUGIN_PATH");
+      !filename.empty()) {
+
+    TritonPlugin TP(filename);
+    std::vector<const char *> passNames;
+    if (auto result = TP.getPassHandles(passNames); !result)
+      llvm::report_fatal_error(result.takeError());
+
+    for (const char *passName : passNames)
+      if (auto result = TP.registerPass(passName); !result)
+        llvm::report_fatal_error(result.takeError());
+  }
+
   registry.insert<
       mlir::triton::TritonDialect, mlir::cf::ControlFlowDialect,
       mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect,

include/triton/Tools/PluginUtils.h

@@ -0,0 +1,84 @@
+#ifndef TRITON_PLUGIN_UTILS_H
+#define TRITON_PLUGIN_UTILS_H
+
+#include "mlir/Pass/PassManager.h"
+#include "llvm/Support/DynamicLibrary.h"
+#include "llvm/Support/Error.h"
+#include <cstdint>
+
+extern "C" {
+enum TritonPluginResult {
+  TP_SUCCESS = 0,
+  TP_GENERIC_FAILURE = 1,
+};
+};
+#define TRITON_PLUGIN_API                                                      \
+  extern "C" __attribute__((visibility("default"))) TritonPluginResult
+
+struct TritonPlugin {
+  TritonPlugin() = delete;
+  TritonPlugin(std::string filename) : filename(filename) {}
+
+private:
+  using enumeratePyBindHandlesType =
+      std::function<TritonPluginResult(uint32_t *, const char **)>;
+  using enumeratePyBindHandlesCType = TritonPluginResult (*)(uint32_t *,
+                                                             const char **);
+
+  // Put enumerate API names here, these can be involved with
+  // enumeratePyBindHandles
+  const std::string ENUMERATE_PASSES = "tritonEnumeratePluginPasses";
+
+  const std::string ADD_PASS = "tritonAddPluginPass";
+  using addPassType =
+      std::function<TritonPluginResult(mlir::PassManager *, const char *)>;
+  using addPassCType = TritonPluginResult (*)(mlir::PassManager *,
+                                              const char *);
+
+  const std::string REGISTER_PASS = "tritonRegisterPluginPass";
+  using registerPassType = std::function<TritonPluginResult(const char *)>;
+  using registerPassCType = TritonPluginResult (*)(const char *);
+
+  llvm::Error checkLibraryValid(const std::string &error) const;
+
+  llvm::Expected<intptr_t> getAddressOfSymbol(const std::string &symbol) const;
+
+  template <typename T, typename U>
+  llvm::Expected<T> getAPI(const std::string &symbol) const {
+    llvm::Expected<intptr_t> getDetailsFn = getAddressOfSymbol(symbol);
+    if (auto Err = getDetailsFn.takeError()) {
+      return Err;
+    }
+    auto func = reinterpret_cast<U>(*getDetailsFn);
+    return func;
+  }
+
+  llvm::Expected<TritonPluginResult> checkAPIResult(TritonPluginResult result,
+                                                    const char *handle) const;
+  llvm::Expected<TritonPluginResult>
+  enumeratePyBindHandles(enumeratePyBindHandlesType &enumeratePyBindHandles,
+                         std::vector<const char *> &passNames);
+
+public:
+  std::runtime_error err2exp(llvm::Error Err);
+
+  llvm::Error loadPlugin();
+
+  llvm::Expected<TritonPluginResult>
+  getPassHandles(std::vector<const char *> &handles);
+
+  llvm::Expected<TritonPluginResult> addPass(mlir::PassManager *pm,
+                                             const char *passHandle);
+
+  llvm::Expected<TritonPluginResult> registerPass(const char *passHandle);
+
+private:
+  std::string filename = "";
+  mutable llvm::sys::DynamicLibrary library;
+  enumeratePyBindHandlesType enumeratePassesAPI;
+  addPassType addPassAPI;
+  registerPassType registerPassAPI;
+  bool isLoaded = false;
+};
+
+#endif // TRITON_PLUGIN_UTILS_H

include/triton/Tools/Sys/GetEnv.hpp

@@ -45,6 +45,7 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_F32_DEFAULT",
     "TRITON_PREFER_TMEM_16x256_LAYOUT",
     "TRITON_ENABLE_EXPERIMENTAL_CONSAN",
+    "TRITON_PASS_PLUGIN_PATH"
     // clang-format on
 };
 

lib/CMakeLists.txt

@@ -4,3 +4,4 @@ add_subdirectory(Dialect)
 add_subdirectory(Target)
 add_subdirectory(Tools)
 add_subdirectory(Instrumentation)
+add_subdirectory(Plugins)

lib/Plugins/CMakeLists.txt

@@ -0,0 +1,43 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name Plugins)
+add_public_tablegen_target(TritonPluginsIncGen)
+
+llvm_canonicalize_cmake_booleans(
+  MLIR_ENABLE_BINDINGS_PYTHON
+)
+
+set(TRITON_PLUGIN_PASSES
+    TritonPluginsTestLib
+    )
+
+set(TritonPluginsTestLib_SOURCES
+    TritonPlugin.cpp
+    )
+
+
+foreach( plugin ${TRITON_PLUGIN_PASSES} )
+    add_mlir_library(${plugin}
+        ${${plugin}_SOURCES}
+        SHARED
+
+        ADDITIONAL_HEADER_DIRS
+        ${PROJECT_BINARY_DIR}/lib
+
+        DEPENDS
+        TritonTableGen
+        TritonCanonicalizeIncGen
+        TritonPluginsIncGen
+
+        LINK_LIBS PUBLIC
+        MLIRPass
+        LLVMSupport
+        MLIRSupport
+        "$<$<PLATFORM_ID:Darwin>:-undefined dynamic_lookup>"
+        )
+
+    set_target_properties(${plugin} PROPERTIES
+          LIBRARY_OUTPUT_DIRECTORY
+      "${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/../plugins")
+
+    target_compile_options(${plugin} PRIVATE -fvisibility=hidden)
+endforeach()

lib/Plugins/Passes.td

@@ -0,0 +1,9 @@
+#ifndef TRITONGPU_PLUGIN_PASSES
+#define TRITONGPU_PLUGIN_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def TritonGPUMLIRPlugin : Pass<"tritongpu-plugin", "mlir::ModuleOp"> {
+  let summary = "Triton MLIR Plugin Pass";
+}
+#endif