[PROTON-DEV] Add Proton Lowering Pass (#5847)

Author: CRobeck

bin/RegisterTritonDialects.h

@@ -3,6 +3,7 @@
 #include "amd/include/TritonAMDGPUTransforms/Passes.h"
 #include "third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h"
 #include "third_party/proton/dialect/include/Dialect/Proton/IR/Dialect.h"
+#include "third_party/proton/dialect/include/TritonProtonToLLVM/Passes.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
@@ -68,6 +69,9 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::triton::registerTritonAMDGPUInsertInstructionSchedHints();
   mlir::triton::registerTritonAMDGPULowerInstructionSchedHints();
 
+  // Proton passes
+  mlir::triton::registerProtonLoweringPass();
+
   // TODO: register Triton & TritonGPU passes
   registry
       .insert<mlir::triton::TritonDialect, mlir::cf::ControlFlowDialect,

test/Proton/ops.mlir

@@ -1,7 +1,6 @@
-// RUN: triton-opt --split-input-file %s -cse -canonicalize | FileCheck %s
+// RUN: triton-opt --split-input-file %s -cse -canonicalize --proton-lowering-pass | FileCheck %s
 
 module {
-  // CHECK-LABEL: proton_record
   tt.func @proton_record() {
     // CHECK: proton.record() {isStart = true, regionId = 1 : i32}
     // CHECK-NEXT: proton.record() {isStart = false, regionId = 1 : i32}

third_party/proton/dialect/include/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(Dialect)
+add_subdirectory(TritonProtonToLLVM)

third_party/proton/dialect/include/TritonProtonToLLVM/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonProtonToLLVM)
+add_public_tablegen_target(TritonProtonConversionPassIncGen)

third_party/proton/dialect/include/TritonProtonToLLVM/Passes.h

@@ -0,0 +1,23 @@
+#ifndef TRITON_THIRD_PARTY_PROTON_INCLUDE_TRITONPROTONGPUTOLLVM_PASSES_H_
+#define TRITON_THIRD_PARTY_PROTON_INCLUDE_TRITONPROTONGPUTOLLVM_PASSES_H_
+
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+
+class ModuleOp;
+template <typename T> class OperationPass;
+
+} // namespace mlir
+
+namespace mlir::triton {
+std::unique_ptr<OperationPass<ModuleOp>> createProtonLoweringPass();
+
+#define GEN_PASS_REGISTRATION
+#include "../third_party/proton/dialect/include/TritonProtonToLLVM/Passes.h.inc"
+
+} // namespace mlir::triton
+
+#endif

third_party/proton/dialect/include/TritonProtonToLLVM/Passes.td

@@ -0,0 +1,10 @@
+#ifndef TRITONPROTON_CONVERSION_PASSES
+#define TRITONPROTON_CONVERSION_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def ProtonLoweringPass : Pass<"proton-lowering-pass", "mlir::ModuleOp"> {
+    let constructor = "mlir::triton::createProtonLoweringPass()";
+
+}
+#endif

third_party/proton/dialect/lib/TritonProtonToLLVM/CMakeLists.txt

@@ -1,5 +1,9 @@
 add_triton_library(TritonProtonToLLVM
     RecordOpToLLVM.cpp
+    ProtonLoweringPass.cpp
+
+    DEPENDS
+    TritonProtonConversionPassIncGen
 
     LINK_LIBS PUBLIC
     ProtonIR

third_party/proton/dialect/lib/TritonProtonToLLVM/ProtonLoweringPass.cpp

@@ -0,0 +1,49 @@
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "triton/Analysis/Allocation.h"
+#include "triton/Analysis/Utility.h"
+#include "triton/Conversion/TritonGPUToLLVM/Passes.h"
+#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+
+#include "third_party/proton/dialect/include/Dialect/Proton/IR/Dialect.h"
+
+using namespace mlir;
+using namespace mlir::triton;
+
+namespace mlir {
+namespace triton {
+#define GEN_PASS_DEF_PROTONLOWERINGPASS
+#include "../third_party/proton/dialect/include/TritonProtonToLLVM/Passes.h.inc"
+} // namespace triton
+} // namespace mlir
+
+namespace {
+struct ProtonLoweringPass
+    : public mlir::triton::impl::ProtonLoweringPassBase<ProtonLoweringPass> {
+  void runOnOperation() override {
+    ModuleOp mod = getOperation();
+    ModuleAllocation allocation(mod);
+
+    OpBuilder b(mod.getBodyRegion());
+    MLIRContext *context = &getContext();
+    auto loc = mod.getLoc();
+
+    /*Add Proton Op Lowerings Here*/
+  }
+};
+
+} // namespace
+
+namespace mlir {
+
+namespace triton {
+
+std::unique_ptr<OperationPass<ModuleOp>> createProtonLoweringPass() {
+  return std::make_unique<ProtonLoweringPass>();
+}
+
+} // namespace triton
+
+} // namespace mlir

third_party/proton/dialect/triton_proton.cc

@@ -5,10 +5,23 @@
 #include <pybind11/stl.h>
 #include <pybind11/stl_bind.h>
 
+#include "../third_party/proton/dialect/include/TritonProtonToLLVM/Passes.h"
+
 namespace py = pybind11;
 
+namespace {
+
+void init_triton_proton_passes_ttgpuir(py::module &&m) {
+  using namespace mlir::triton;
+  m.def("add_proton_lowering_pass",
+        [](mlir::PassManager &pm) { pm.addPass(createProtonLoweringPass()); });
+}
+} // namespace
+
 void init_triton_proton(py::module &&m) {
+  m.doc() = "Python bindings to the Proton backend";
   auto passes = m.def_submodule("passes");
+  init_triton_proton_passes_ttgpuir(passes.def_submodule("ttgpuir"));
 
   // load dialects
   m.def("load_dialects", [](mlir::MLIRContext &context) {