TTKernel to emitc pipeline

tenstorrent · Jan 1, 2025 · cb20078 · cb20078
1 parent 90c30ba
commit cb20078
Show file tree

Hide file tree

Showing 8 changed files with 83 additions and 3 deletions.
diff --git a/include/ttmlir/Conversion/TTKernelToEmitC/TTKernelToEmitC.h b/include/ttmlir/Conversion/TTKernelToEmitC/TTKernelToEmitC.h
@@ -19,9 +19,9 @@ namespace mlir::tt {
 // Runs a conversion pass to EmitC dialect on a func op containing given
 // region's body. Also, it adds boilerplate code such as includes and namespace
 // declarations.
-LogicalResult convertTTKernelRegionToEmitC(
-    OpBuilder &builder, Region *region,
-    const ttkernel::KernelConfigInterface &kernelConfig);
+LogicalResult
+convertTTKernelRegionToEmitC(OpBuilder &builder, Region *region,
+                             const ttkernel::ThreadType &threadType);
 
 // Converts given region to EmitC dialect and translates it to C++ code.
 LogicalResult

diff --git a/include/ttmlir/Dialect/TTKernel/Pipelines/TTKernelPipelines.h b/include/ttmlir/Dialect/TTKernel/Pipelines/TTKernelPipelines.h
@@ -0,0 +1,34 @@
+// 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 {
+// Options for the TTKernel to EmitC pipeline.
+//
+// struct TTIRToTTMetalBackendPipelineOptions
+//     : public PassPipelineOptions<TTIRToTTMetalBackendPipelineOptions> {
+//   ListOption<int64_t> meshShape{
+//       *this, "mesh-shape", llvm::cl::desc("Set the multi-device mesh
+//       shape.")};
+
+//   // Option to provide a system descriptor flatbuffer file to compile
+//   // against.
+//   //
+//   Option<std::string> systemDescPath{
+//       *this, "system-desc-path",
+//       llvm::cl::desc(
+//           "Pass in a system descriptor flatbuffer to compile against."),
+//       llvm::cl::init("")};
+// };
+
+void createTTKernelToEmitCPipeline(OpPassManager &pm);
+
+void registerTTKernelPipelines();
+} // namespace mlir::tt::ttkernel
+
+#endif
diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
@@ -27,6 +27,7 @@ MLIRTTMetalDialect
 MLIRTTIRPipelines
 MLIRTTNNPipelines
 MLIRTTMetalPipelines
+MLIRTTKernelPipelines
 )
 
 if (TTMLIR_ENABLE_STABLEHLO)

diff --git a/lib/Dialect/TTKernel/CMakeLists.txt b/lib/Dialect/TTKernel/CMakeLists.txt
@@ -1 +1,2 @@
 add_subdirectory(IR)
+add_subdirectory(Pipelines)
diff --git a/lib/Dialect/TTKernel/Pipelines/CMakeLists.txt b/lib/Dialect/TTKernel/Pipelines/CMakeLists.txt
@@ -0,0 +1,11 @@
+add_mlir_dialect_library(MLIRTTKernelPipelines
+  TTKernelPipelines.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${PROJECT_SOURCE_DIR}/include/ttmlir
+
+  LINK_LIBS PUBLIC
+  MLIRTTIRDialect
+  MLIRTTKernelDialect
+  MLIRTransforms
+)
diff --git a/lib/Dialect/TTKernel/Pipelines/TTKernelPipelines.cpp b/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
diff --git a/lib/RegisterAll.cpp b/lib/RegisterAll.cpp
@@ -14,6 +14,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"
@@ -76,4 +77,5 @@ void mlir::tt::registerAllPasses() {
   mlir::tt::ttir::registerTTIRPipelines();
   mlir::tt::ttnn::registerTTNNPipelines();
   mlir::tt::ttmetal::registerTTMetalPipelines();
+  mlir::tt::ttkernel::registerTTKernelPipelines();
 }
diff --git a/lib/SharedLib/CMakeLists.txt b/lib/SharedLib/CMakeLists.txt
@@ -24,6 +24,7 @@ set(TTMLIR_LIBS
     MLIRTTNNAnalysis
     MLIRTTNNPipelines
     MLIRTTMetalPipelines
+    MLIRTTKernelPipelines
     TTMLIRTTNNToEmitC
 )
Original file line number	Diff line number	Diff line change
		@@ -1 +1,2 @@
		add_subdirectory(IR)
		add_subdirectory(Pipelines)