TTKernel to emitc pipeline

include/ttmlir/Dialect/TTKernel/Pipelines/TTKernelPipelines.h

@@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_DIALECT_TTKERNEL_PIPELINES_TTKERNELPIPELINES_H
+#define TTMLIR_DIALECT_TTKERNEL_PIPELINES_TTKERNELPIPELINES_H
+
+#include "mlir/Pass/PassOptions.h"
+
+namespace mlir::tt::ttkernel {
+
+void createTTKernelToEmitCPipeline(OpPassManager &pm);
+
+void registerTTKernelPipelines();
+} // namespace mlir::tt::ttkernel
+
+#endif

lib/CMakeLists.txt

@@ -27,6 +27,7 @@ MLIRTTMetalDialect
 MLIRTTIRPipelines
 MLIRTTNNPipelines
 MLIRTTMetalPipelines
+MLIRTTKernelPipelines
 )
 
 if (TTMLIR_ENABLE_STABLEHLO)

lib/Dialect/TTKernel/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(IR)
+add_subdirectory(Pipelines)

lib/Dialect/TTKernel/Pipelines/CMakeLists.txt

@@ -0,0 +1,9 @@
+add_mlir_dialect_library(MLIRTTKernelPipelines
+  TTKernelPipelines.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${PROJECT_SOURCE_DIR}/include/ttmlir
+
+  LINK_LIBS PUBLIC
+  MLIRTTKernelDialect
+)

lib/Dialect/TTKernel/Pipelines/TTKernelPipelines.cpp

@@ -0,0 +1,30 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "ttmlir/Dialect/TTKernel/Pipelines/TTKernelPipelines.h"
+
+#include "mlir/Pass/PassManager.h"
+
+#include "ttmlir/Conversion/Passes.h"
+#include "ttmlir/Dialect/TTIR/Transforms/Passes.h"
+
+namespace mlir::tt::ttkernel {
+//===----------------------------------------------------------------------===//
+// Pipeline implementation.
+//===----------------------------------------------------------------------===//
+
+void createTTKernelToEmitCPipeline(OpPassManager &pm) {
+  pm.addPass(createConvertTTKernelToEmitC());
+}
+
+//===----------------------------------------------------------------------===//
+// Pipeline registration.
+//===----------------------------------------------------------------------===//
+
+void registerTTKernelPipelines() {
+  mlir::PassPipelineRegistration<>(
+      "ttkernel-to-emitc-pipeline", "Pipeline lowering ttkernel to emitc.",
+      mlir::tt::ttkernel::createTTKernelToEmitCPipeline);
+}
+} // namespace mlir::tt::ttkernel

lib/RegisterAll.cpp

@@ -10,6 +10,7 @@
 #include "ttmlir/Dialect/TTIR/Pipelines/TTIRPipelines.h"
 #include "ttmlir/Dialect/TTIR/Transforms/Passes.h"
 #include "ttmlir/Dialect/TTKernel/IR/TTKernel.h"
+#include "ttmlir/Dialect/TTKernel/Pipelines/TTKernelPipelines.h"
 #include "ttmlir/Dialect/TTMetal/Pipelines/TTMetalPipelines.h"
 #include "ttmlir/Dialect/TTMetal/Transforms/Passes.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNN.h"
@@ -82,4 +83,5 @@ void mlir::tt::registerAllPasses() {
   mlir::tt::ttir::registerTTIRPipelines();
   mlir::tt::ttnn::registerTTNNPipelines();
   mlir::tt::ttmetal::registerTTMetalPipelines();
+  mlir::tt::ttkernel::registerTTKernelPipelines();
 }

lib/SharedLib/CMakeLists.txt

@@ -25,6 +25,7 @@ set(TTMLIR_LIBS
     MLIRTTNNAnalysis
     MLIRTTNNPipelines
     MLIRTTMetalPipelines
+    MLIRTTKernelPipelines
     TTMLIRTTNNToEmitC
 )
 

python/Passes.cpp

@@ -144,6 +144,29 @@ void populatePassesModule(py::module &m) {
         }
       },
       py::arg("module"), py::arg("options") = "");
+  m.def(
+      "ttkernel_to_emitc_backend_pipeline",
+      [](MlirModule module, std::string options = "") {
+        mlir::Operation *moduleOp = unwrap(mlirModuleGetOperation(module));
+        mlir::PassManager pm(moduleOp->getName());
+        mlir::DialectRegistry registry;
+        mlir::tt::registerAllDialects(registry);
+        mlir::tt::registerAllExtensions(registry);
+        mlir::MLIRContext *ctx = unwrap(mlirModuleGetContext(module));
+        ctx->appendDialectRegistry(registry);
+        const auto *pipeline =
+            mlir::PassPipelineInfo::lookup("ttkernel-to-emitc-pipeline");
+        mlir::function_ref<mlir::LogicalResult(const llvm::Twine &)>
+            err_handler =
+                [](const llvm::Twine &loc) { return mlir::failure(); };
+        if (mlir::failed(pipeline->addToPipeline(pm, options, err_handler))) {
+          throw std::runtime_error("Failed to add pipeline to pass manager");
+        }
+        if (mlir::failed(pm.run(moduleOp))) {
+          throw std::runtime_error("Failed to run pass manager");
+        }
+      },
+      py::arg("module"), py::arg("options") = "");
 
   py::class_<std::shared_ptr<void>>(m, "SharedVoidPtr")
       .def(py::init<>())

python/pykernel/pykernel_ast.py

@@ -7,6 +7,47 @@
 import ast
 import inspect
 import functools
+from ttmlir.passes import ttkernel_to_emitc_backend_pipeline
+
+
+def ttkernel_to_emitc(
+    module,
+    dump_to_file: bool = True,
+    output_file_name: str = "test.mlir",
+):
+    """
+    Converts TTKernel module `module` to EmitC module and optionally dumps to file.
+
+    Wrapper around `ttkernel_to_emitc_backend_pipeline` pybound pass.
+
+    Arguments
+    ---------
+    module: ???
+        TTKernel module to convert to EmitC module
+
+    dump_to_file: bool
+        Flag which indicates that generated EmitC module will be dumped to file.
+
+    output_file_name: str
+        Name of the output file.
+
+    Returns
+    -------
+    MLIR module containing MLIR op graph defined by `module` and instance of TTKernelBuilder.
+    """
+
+    # Now, pass it through the TTIR to TTMetal pipeline. Module gets
+    # modified in place.
+    ttkernel_to_emitc_backend_pipeline(module)
+
+    print("`ttkernel_to_emitc_backend_pipeline` passed successfully.")
+
+    # Optionally dump to file.
+    if dump_to_file:
+        with open(output_file_name, "w") as f:
+            f.write(str(module))
+
+    return module
 
 
 def get_supported_nodes():
@@ -352,7 +393,8 @@ def _wrapper(*args, **kwargs):
         # print(ast.dump(m, indent=4) + "\n")
         b.visit(m)
         print(b.module)
-
+        # ttkernel_to_emitc(b.module)
+        # print(b.module)
         # Check if generated IR is valid
         b.module.operation.verify()