diff --git a/include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h b/include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h index 7e5829873..48c723e1c 100644 --- a/include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h +++ b/include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h @@ -5,10 +5,11 @@ #ifndef TTMLIR_DIALECT_TTNN_PIPELINES_TTNNPIPELINES_H #define TTMLIR_DIALECT_TTNN_PIPELINES_TTNNPIPELINES_H -#include "mlir/Pass/PassOptions.h" #include "ttmlir/Dialect/TT/Utils/MemoryLayoutAnalysisParams.h" #include "ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h" +#include "mlir/Pass/PassOptions.h" + namespace mlir::tt::ttnn { // Options for the TTIR to TTNN backend pipeline. @@ -112,6 +113,11 @@ struct TTIRToTTNNBackendPipelineOptions *this, "mesh-shape", llvm::cl::desc("Set the multi-device mesh shape.")}; }; +// TTIR to EmitC pipeline options. +// Inherit from TTIRToTTNNBackendPipelineOptions to reuse the options. +// +struct TTIRToEmitCPipelineOptions : public TTIRToTTNNBackendPipelineOptions {}; + void createTTNNPipelineTTIRPasses( OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options); @@ -145,6 +151,9 @@ void createTTNNPipelineDeallocPassFromString(OpPassManager &pm, void createTTIRToTTNNBackendPipeline( OpPassManager &pm, const TTIRToTTNNBackendPipelineOptions &options); +void createTTIRToEmitCPipeline(OpPassManager &pm, + const TTIRToEmitCPipelineOptions &options); + /// Registers all pipelines for the `bufferization` dialect. Currently, /// this includes only the "ttir-to-ttnn-backend-pipeline". void registerTTNNPipelines(); diff --git a/lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp b/lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp index 7b5fcecf5..24980fb7c 100644 --- a/lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp +++ b/lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp @@ -4,13 +4,14 @@ #include "ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h" -#include "mlir/Pass/PassManager.h" - -#include "mlir/Transforms/Passes.h" #include "ttmlir/Conversion/Passes.h" +#include "ttmlir/Conversion/TTNNToEmitC/TTNNToEmitC.h" #include "ttmlir/Dialect/TTIR/Transforms/Passes.h" #include "ttmlir/Dialect/TTNN/Transforms/Passes.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/Passes.h" + namespace mlir::tt::ttnn { //===----------------------------------------------------------------------===// // Pipeline implementation. @@ -115,15 +116,32 @@ void createTTIRToTTNNBackendPipeline( createTTNNPipelineDeallocPass(pm, options); } +void createTTIRToEmitCPipeline(OpPassManager &pm, + const TTIRToEmitCPipelineOptions &options) { + createTTIRToTTNNBackendPipeline(pm, options); + pm.addPass(createConvertTTNNToEmitCPass()); +} + //===----------------------------------------------------------------------===// // Pipeline registration. //===----------------------------------------------------------------------===// void registerTTNNPipelines() { + // TTIR to TTNN backend pipeline. + // mlir::PassPipelineRegistration< mlir::tt::ttnn::TTIRToTTNNBackendPipelineOptions>( "ttir-to-ttnn-backend-pipeline", - "Pipeline lowering ttir to ttnn backend.", + "Pipeline lowering TTIR to TTNN backend.", mlir::tt::ttnn::createTTIRToTTNNBackendPipeline); + + // TTIR to EmitC pipeline. + // + mlir::PassPipelineRegistration( + "ttir-to-emitc-pipeline", + "Pipeline lowering TTIR to EmitC. Under the hood, it runs " + "--ttir-to-ttnn-backend-pipeline and then converts the resulting TTNN " + "dialect to EmitC.", + mlir::tt::ttnn::createTTIRToEmitCPipeline); } } // namespace mlir::tt::ttnn diff --git a/test/ttmlir/Dialect/TTNN/pipelines/ttir_to_emitc_add.mlir b/test/ttmlir/Dialect/TTNN/pipelines/ttir_to_emitc_add.mlir new file mode 100644 index 000000000..d59a59ea6 --- /dev/null +++ b/test/ttmlir/Dialect/TTNN/pipelines/ttir_to_emitc_add.mlir @@ -0,0 +1,14 @@ +// RUN: ttmlir-opt --ttir-to-emitc-pipeline="system-desc-path=%system_desc_path%" %s > %direct.mlir +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" --convert-ttnn-to-emitc %s > %indirect.mlir +// RUN: diff %direct.mlir %indirect.mlir +// +// This test checks that the (TTIR to EmitC pipeline) is equivalent to (TTIR to TTNN pipeline + dialect conversion from TTNN to EmitC). +// The `diff` command will return 0 if files are identical, otherwise it will return the diff, which will make `llvm-lit` treat the test as failed. + +#any_device_tile = #tt.operand_constraint + +func.func @add(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + %0 = tensor.empty() : tensor<64x128xf32> + %1 = "ttir.add"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device_tile, #any_device_tile, #any_device_tile]}> : (tensor<64x128xf32>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + return %1 : tensor<64x128xf32> +}