From a47dcc7e76fbf0e4e76cb59f9303acfa51a544eb Mon Sep 17 00:00:00 2001 From: Sasa Vuckovic <134393361+svuckovicTT@users.noreply.github.com> Date: Thu, 21 Nov 2024 15:50:03 +0100 Subject: [PATCH] Fix unregistered dialects error message (ttmlir-translate) (#1341) --- docs/src/ttmlir-translate.md | 6 ++-- lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp | 33 +++++++++++++++++++ .../TTNNToEmitC/TTNNToEmitCPass.cpp | 22 ++++++++----- test/ttmlir/Silicon/TTNN/emitc/two_fns.mlir | 16 +++++++++ tools/ttnn-standalone/README.md | 2 +- tools/ttnn-standalone/ttnn-standalone.cpp | 6 ++-- 6 files changed, 69 insertions(+), 16 deletions(-) create mode 100644 test/ttmlir/Silicon/TTNN/emitc/two_fns.mlir diff --git a/docs/src/ttmlir-translate.md b/docs/src/ttmlir-translate.md index c82f7ee8f..ba9c69b3c 100644 --- a/docs/src/ttmlir-translate.md +++ b/docs/src/ttmlir-translate.md @@ -5,15 +5,15 @@ The `ttmlir-translate` translation utility. Unlike `ttmlir-opt` tool which is us ```bash # First, let's run `ttmlir-opt` to convert to proper dialect -./build/bin/ttmlir-opt --ttir-load-system-desc --ttir-layout --convert-ttir-to-ttnn --convert-ttnn-to-emitc test/ttmlir/Dialect/TTNN/simple_multiply.mlir -o c.mlir +./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/Dialect/TTNN/simple_multiply.mlir -o c.mlir # Now run `ttmlir-translate` to produce C++ code -./build/bin/ttmlir-translate -mlir-to-cpp c.mlir -allow-unregistered-dialect +./build/bin/ttmlir-translate --mlir-to-cpp c.mlir ``` Bonus: These two commands can be piped, to avoid writing a `mlir` file to disk, like so: ```bash -./build/bin/ttmlir-opt --ttir-load-system-desc --ttir-layout --convert-ttir-to-ttnn --convert-ttnn-to-emitc test/ttmlir/Dialect/TTNN/simple_multiply.mlir | ./build/bin/ttmlir-translate -mlir-to-cpp -allow-unregistered-dialect +./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/Dialect/TTNN/simple_multiply.mlir | ./build/bin/ttmlir-translate -mlir-to-cpp ``` ## Generate flatbuffer file from MLIR diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp index 9b7cf7fe8..92862cd9d 100644 --- a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp +++ b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp @@ -618,6 +618,35 @@ class DeallocateOpConversionPattern } }; +// Module Op conversion pattern +// +// This conversion pattern removes attributes from the ModuleOp. Previously, +// ttmlir-translate would complain when translating to C++ if there were any +// attributes from "unregistered" dialects. +// +class ModuleOpConversionPattern + : public TTNNToEmitCBaseOpConversionPattern { + +public: + ModuleOpConversionPattern(const TypeConverter &typeConverter, + MLIRContext *context, PatternBenefit benefit = 1) + : TTNNToEmitCBaseOpConversionPattern(typeConverter, + context, benefit) {} + + LogicalResult + matchAndRewrite(mlir::ModuleOp srcOp, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + rewriter.modifyOpInPlace(srcOp, [&]() { + for (const NamedAttribute &attr : srcOp->getAttrs()) { + srcOp->removeAttr(attr.getName()); + } + }); + + return success(); + } +}; + } // namespace namespace mlir::tt { @@ -720,6 +749,10 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx, // patterns.add>(typeConverter, ctx); + + // Module op + // + patterns.add(typeConverter, ctx); } } // namespace mlir::tt diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitCPass.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitCPass.cpp index 71a7c52b6..bd0c9044f 100644 --- a/lib/Conversion/TTNNToEmitC/TTNNToEmitCPass.cpp +++ b/lib/Conversion/TTNNToEmitC/TTNNToEmitCPass.cpp @@ -4,6 +4,11 @@ #include "ttmlir/Conversion/TTNNToEmitC/TTNNToEmitC.h" +#include "ttmlir/Dialect/TTNN/IR/TTNN.h" +#include "ttmlir/Dialect/TTNN/IR/TTNNOps.h" +#include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h" +#include "ttmlir/Dialect/TTNN/IR/TTNNOpsTypes.h" + #include "mlir/Dialect/EmitC/IR/EmitC.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/Transforms/FuncConversions.h" @@ -12,11 +17,6 @@ #include "mlir/Support/LogicalResult.h" #include "mlir/Transforms/DialectConversion.h" -#include "ttmlir/Dialect/TTNN/IR/TTNN.h" -#include "ttmlir/Dialect/TTNN/IR/TTNNOps.h" -#include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h" -#include "ttmlir/Dialect/TTNN/IR/TTNNOpsTypes.h" - using namespace mlir; using namespace mlir::tt; @@ -48,14 +48,20 @@ struct ConvertTTNNToEmitCPass void runOnOperation() override { mlir::ConversionTarget target(getContext()); + // EmitC is legal, TTNN is illegal + // target.addLegalDialect(); target.addIllegalDialect(); - target.addLegalOp(); + + // mlir::ModuleOp is legal only if no attributes are present on it + // + target.addDynamicallyLegalOp( + [&](mlir::ModuleOp op) { return op->getAttrs().empty(); }); // Add header imports to front of module // { - auto module = getOperation(); + mlir::ModuleOp module = getOperation(); OpBuilder builder(module); if (module.getBodyRegion().empty()) { @@ -107,7 +113,7 @@ struct ConvertTTNNToEmitCPass return; } } - }; + } }; } // namespace diff --git a/test/ttmlir/Silicon/TTNN/emitc/two_fns.mlir b/test/ttmlir/Silicon/TTNN/emitc/two_fns.mlir new file mode 100644 index 000000000..3f304969c --- /dev/null +++ b/test/ttmlir/Silicon/TTNN/emitc/two_fns.mlir @@ -0,0 +1,16 @@ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn + +#any_device = #tt.operand_constraint + +func.func @add(%arg0: tensor<32x32xbf16>, %arg1: tensor<32x32xbf16>) -> tensor<32x32xbf16> { + %0 = tensor.empty() : tensor<32x32xbf16> + %1 = "ttir.add"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<32x32xbf16>, tensor<32x32xbf16>, tensor<32x32xbf16>) -> tensor<32x32xbf16> + return %1 : tensor<32x32xbf16> +} + +func.func @subtract(%arg0: tensor<32x32xbf16>, %arg1: tensor<32x32xbf16>) -> tensor<32x32xbf16> { + %0 = tensor.empty() : tensor<32x32xbf16> + %1 = "ttir.subtract"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<32x32xbf16>, tensor<32x32xbf16>, tensor<32x32xbf16>) -> tensor<32x32xbf16> + return %1 : tensor<32x32xbf16> +} diff --git a/tools/ttnn-standalone/README.md b/tools/ttnn-standalone/README.md index 816cfe1cf..619e52d1c 100644 --- a/tools/ttnn-standalone/README.md +++ b/tools/ttnn-standalone/README.md @@ -14,7 +14,7 @@ Third party ML models (PyTorch, Jax, ONNX, ...) can be compiled to a set of TTNN ```bash # Compile a model to C++ code -./build/bin/ttmlir-opt --ttir-load-system-desc --ttir-implicit-device --ttir-layout --convert-ttir-to-ttnn --ttnn-decompose-layouts --ttnn-deallocate --convert-ttnn-to-emitc test/ttmlir/Silicon/TTNN/emitc/simple_add.mlir | ./build/bin/ttmlir-translate --mlir-to-cpp -allow-unregistered-dialect +./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/Silicon/TTNN/emitc/simple_add.mlir | ./build/bin/ttmlir-translate --mlir-to-cpp # Copy paste the generated function into `ttnn-standalone.cpp`. diff --git a/tools/ttnn-standalone/ttnn-standalone.cpp b/tools/ttnn-standalone/ttnn-standalone.cpp index dff9afff4..0dee60f13 100644 --- a/tools/ttnn-standalone/ttnn-standalone.cpp +++ b/tools/ttnn-standalone/ttnn-standalone.cpp @@ -5,11 +5,9 @@ #include "ttnn-precompiled.hpp" // To generate forward function, run: -// ./build/bin/ttmlir-opt --ttir-load-system-desc --ttir-implicit-device -// --ttir-layout --convert-ttir-to-ttnn --ttnn-decompose-layouts -// --ttnn-deallocate --convert-ttnn-to-emitc +// ./build/bin/ttmlir-opt --ttir-to-emitc-pipeline // test/ttmlir/Silicon/TTNN/emitc/simple_add.mlir | ./build/bin/ttmlir-translate -// --mlir-to-cpp -allow-unregistered-dialect +// --mlir-to-cpp ttnn::Tensor forward(ttnn::Tensor v1, ttnn::Tensor v2) { ttnn::Device *v3 = ttnn::DeviceGetter::getInstance();