From b75d44d98fdba061b1c1da396cde9a42c07911a9 Mon Sep 17 00:00:00 2001 From: Jovan Serbedzija Date: Fri, 1 Nov 2024 12:17:49 +0100 Subject: [PATCH] Add log operation (#1123) --- include/ttmlir/Dialect/TTIR/IR/TTIROps.td | 7 +++++++ include/ttmlir/Dialect/TTNN/IR/TTNNOps.td | 7 +++++++ include/ttmlir/Target/TTNN/program.fbs | 3 ++- lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp | 1 + lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp | 3 ++- lib/Target/TTNN/TTNNToFlatbuffer.cpp | 5 +++++ runtime/lib/ttnn/operations/eltwise/unary/unary.cpp | 4 ++++ .../Silicon/TTNN/perf_unit/test_perf_log.mlir | 13 +++++++++++++ test/ttmlir/Silicon/TTNN/simple_eltwise.mlir | 8 ++++++++ 9 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 test/ttmlir/Silicon/TTNN/perf_unit/test_perf_log.mlir diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td index 674042348..f74341f4e 100644 --- a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td +++ b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td @@ -274,6 +274,13 @@ def TTIR_TypecastOp: TTIR_ElementwiseUnaryOp<"typecast"> { }]; } +def TTIR_LogOp: TTIR_ElementwiseUnaryOp<"log"> { + let summary = "Eltwise logarithm op."; + let description = [{ + Eltwise logarithm operation. Calculates log(x) for all elements x in input tensor. + }]; +} + class TTIR_ElementwiseBinaryOp traits = []> : TTIR_ElementwiseOp { let summary = "Eltwise binary op."; diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td index f878e7f21..e663dd018 100644 --- a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td +++ b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td @@ -237,6 +237,13 @@ def TTNN_SigmoidOp : TTNN_ElementwiseUnaryOp<"sigmoid"> { }]; } +def TTNN_LogOp : TTNN_ElementwiseUnaryOp<"log"> { + let summary = "Eltwise logarithm."; + let description = [{ + Eltwise logarithm operation. + }]; +} + def TTNN_AddOp : TTNN_ElementwiseBinaryOp<"add"> { let summary = "Eltwise add."; let description = [{ diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs index 27da085a8..8e3facea2 100644 --- a/include/ttmlir/Target/TTNN/program.fbs +++ b/include/ttmlir/Target/TTNN/program.fbs @@ -85,7 +85,8 @@ enum EltwiseOpType: uint32 { Minimum = 24, Ceil = 25, Sin = 26, - Cos = 27 + Cos = 27, + Log = 28 } union EltwiseOpParams { diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp index a873e0338..7cefdbb79 100644 --- a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp +++ b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp @@ -879,6 +879,7 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns, ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, + ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp index 4655bc0c3..70f6dd4d5 100644 --- a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp +++ b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp @@ -626,7 +626,8 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx, DefaultOpConversionPattern, DefaultOpConversionPattern, DefaultOpConversionPattern, - DefaultOpConversionPattern>(typeConverter, ctx); + DefaultOpConversionPattern, + DefaultOpConversionPattern>(typeConverter, ctx); // Eltwise binary ops // diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp index d731f4b5c..1cd080cc2 100644 --- a/lib/Target/TTNN/TTNNToFlatbuffer.cpp +++ b/lib/Target/TTNN/TTNNToFlatbuffer.cpp @@ -351,6 +351,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) { type = ::tt::target::ttnn::EltwiseOpType::Cos; } else if constexpr (std::is_same_v) { type = ::tt::target::ttnn::EltwiseOpType::Sin; + } else if constexpr (std::is_same_v) { + type = ::tt::target::ttnn::EltwiseOpType::Log; } else { llvm_unreachable("unhandled EltwiseOp"); } @@ -597,6 +599,9 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op, if (auto expOp = dyn_cast(op); expOp) { return createOperation(cache, createEltwiseOp(cache, expOp), debugString); } + if (auto logOp = dyn_cast(op); logOp) { + return createOperation(cache, createEltwiseOp(cache, logOp), debugString); + } if (auto sigmoidOp = dyn_cast(op); sigmoidOp) { return createOperation(cache, createEltwiseOp(cache, sigmoidOp), debugString); diff --git a/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp b/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp index 4b261d165..c5b5ff6e4 100644 --- a/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp +++ b/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp @@ -98,6 +98,10 @@ void run(const ::tt::target::ttnn::EltwiseOp *op, ProgramContext &context) { runEltwiseUnaryWithFastAndApproximateModeOP(op, tensorPool, ::ttnn::exp); break; } + case ::tt::target::ttnn::EltwiseOpType::Log: { + runEltwiseUnaryOP(op, tensorPool, ::ttnn::log); + break; + } default: throw std::invalid_argument("Unsupported unary operation"); } diff --git a/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_log.mlir b/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_log.mlir new file mode 100644 index 000000000..b1ca157c6 --- /dev/null +++ b/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_log.mlir @@ -0,0 +1,13 @@ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir +// RUN: FileCheck %s --input-file=%t.mlir +// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn +#any_device = #tt.operand_constraint +#any_device_tile = #tt.operand_constraint + +func.func @sqrt(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]] + %0 = tensor.empty() : tensor<64x128xf32> + // CHECK: %[[C:.*]] = "ttnn.log"[[C:.*]] + %1 = "ttir.log"(%arg0, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + return %1 : tensor<64x128xf32> +} diff --git a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir index b8f387d27..5b683e276 100644 --- a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir +++ b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir @@ -183,3 +183,11 @@ func.func @typecast(%arg0: tensor<64x128xf32>) -> tensor<64x128xbf16> { %1 = "ttir.typecast"(%arg0, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xbf16>) -> tensor<64x128xbf16> return %1 : tensor<64x128xbf16> } + +func.func @log(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> { + // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]] + %0 = tensor.empty() : tensor<64x128xf32> + // CHECK: %[[C:.*]] = "ttnn.log"[[C:.*]] + %1 = "ttir.log"(%arg0, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + return %1 : tensor<64x128xf32> +}