diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td index 92f491b03..c87770bef 100644 --- a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td +++ b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td @@ -211,6 +211,18 @@ def TTIR_CosOp: TTIR_ElementwiseUnaryOp<"cos"> { }]; } +def TTIR_SignOp: TTIR_ElementwiseUnaryOp<"sign"> { + let summary = "Eltwise sign operation."; + let description = [{ + Returns the sign of the `operand` element-wise and produces a `result` + tensor. + + Example: + %a: [[3, -2, 0], [1, -4, 4]] + "ttir.sign"(%a, %out) -> %out: [[1, -1, 0], [1, -1, 1]] + }]; +} + def TTIR_LogicalNotOp: TTIR_ElementwiseUnaryOp<"logical_not"> { let summary = "Eltwise logical not op."; let description = [{ diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td index 4dcf05046..01ebae803 100644 --- a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td +++ b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td @@ -167,6 +167,18 @@ def TTNN_CeilOp : TTNN_ElementwiseUnaryOp<"ceil"> { }]; } +def TTNN_SignOp: TTNN_ElementwiseUnaryOp<"sign"> { + let summary = "Eltwise sign operation."; + let description = [{ + Returns the sign of the `operand` element-wise and produces a `result` + tensor. + + Example: + %a: [[3, -2, 0], [1, -4, 4]] + "ttnn.sign"(%a, %out) -> %out: [[1, -1, 0], [1, -1, 1]] + }]; +} + def TTNN_CosOp : TTNN_ElementwiseUnaryOp<"cos"> { let summary = "Eltwise cosine."; let description = [{ diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs index 11e781c62..cea35447c 100644 --- a/include/ttmlir/Target/TTNN/program.fbs +++ b/include/ttmlir/Target/TTNN/program.fbs @@ -88,7 +88,8 @@ enum EltwiseOpType: uint32 { Cos = 27, Log = 28, Log1p = 29, - Expm1 = 30 + Expm1 = 30, + Sign = 31 } union EltwiseOpParams { diff --git a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp index 2986f486a..15b1f086b 100644 --- a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp +++ b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp @@ -863,6 +863,8 @@ void addElementwiseUnaryOpsConversionPatterns(MLIRContext *ctx, mlir::stablehlo::Log1pOp, mlir::tt::ttir::Log1pOp>>(typeConverter, ctx); patterns.add>(typeConverter, ctx); + patterns.add>(typeConverter, ctx); } void addElementwiseBinaryOpsConversionPatterns(MLIRContext *ctx, diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp index b0cfc3634..42d834634 100644 --- a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp +++ b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp @@ -876,6 +876,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 f78ed65e1..0582dce37 100644 --- a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp +++ b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp @@ -621,6 +621,7 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx, DefaultOpConversionPattern, DefaultOpConversionPattern, DefaultOpConversionPattern, + DefaultOpConversionPattern, DefaultOpConversionPattern, DefaultOpConversionPattern, DefaultOpConversionPattern, diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp index cc41042cb..71d793a00 100644 --- a/lib/Target/TTNN/TTNNToFlatbuffer.cpp +++ b/lib/Target/TTNN/TTNNToFlatbuffer.cpp @@ -337,6 +337,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) { type = ::tt::target::ttnn::EltwiseOpType::Sqrt; } else if constexpr (std::is_same_v) { type = ::tt::target::ttnn::EltwiseOpType::Rsqrt; + } else if constexpr (std::is_same_v) { + type = ::tt::target::ttnn::EltwiseOpType::Sign; } else if constexpr (std::is_same_v) { type = ::tt::target::ttnn::EltwiseOpType::Reciprocal; } else if constexpr (std::is_same_v) { @@ -600,6 +602,9 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op, if (auto rsqrtOp = dyn_cast(op); rsqrtOp) { return createOperation(cache, createEltwiseOp(cache, rsqrtOp), debugString); } + if (auto signOp = dyn_cast(op); signOp) { + return createOperation(cache, createEltwiseOp(cache, signOp), debugString); + } if (auto expOp = dyn_cast(op); expOp) { return createOperation(cache, createEltwiseOp(cache, expOp), debugString); } diff --git a/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp b/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp index 77a89cf07..c7dba6e95 100644 --- a/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp +++ b/runtime/lib/ttnn/operations/eltwise/unary/unary.cpp @@ -96,6 +96,10 @@ void run(const ::tt::target::ttnn::EltwiseOp *op, ProgramContext &context) { runEltwiseUnaryOP(op, tensorPool, ::ttnn::reciprocal); break; } + case ::tt::target::ttnn::EltwiseOpType::Sign: { + runEltwiseUnaryOP(op, tensorPool, ::ttnn::sign); + break; + } case ::tt::target::ttnn::EltwiseOpType::Exp: { runEltwiseUnaryWithFastAndApproximateModeOP(op, tensorPool, ::ttnn::exp); break; diff --git a/test/ttmlir/Conversion/StableHLOToTTIR/sign_op.mlir b/test/ttmlir/Conversion/StableHLOToTTIR/sign_op.mlir new file mode 100644 index 000000000..0bf4a1bca --- /dev/null +++ b/test/ttmlir/Conversion/StableHLOToTTIR/sign_op.mlir @@ -0,0 +1,12 @@ +// REQUIRES: stablehlo +// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | FileCheck %s +#any_device = #tt.operand_constraint +module @jit_eltwise_sign attributes {} { + func.func public @test_sign(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> { + %0 = stablehlo.sign %arg0 : tensor<13x21x3xf32> + // CHECK: [[VAL0:%[0-9]+]] = tensor.empty() : [[TENSOR_SIZE:tensor<[0-9]+x[0-9]+x[0-9]+xf[0-9]+>]] + // CHECK: [[VAL1:%[0-9]+]] = "ttir.sign"(%arg0, [[VAL0]]) <{operandSegmentSizes = array, operand_constraints = [#any_device_tile, #any_device_tile]}> : ([[TENSOR_SIZE]], [[TENSOR_SIZE]]) -> [[TENSOR_SIZE]] + return %0 : tensor<13x21x3xf32> + // CHECK: return [[VAL1]] : [[TENSOR_SIZE]] + } +} diff --git a/test/ttmlir/Dialect/TTNN/eltwise/unary/sign/simple_sign.mlir b/test/ttmlir/Dialect/TTNN/eltwise/unary/sign/simple_sign.mlir new file mode 100644 index 000000000..c82547bff --- /dev/null +++ b/test/ttmlir/Dialect/TTNN/eltwise/unary/sign/simple_sign.mlir @@ -0,0 +1,12 @@ +// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s +#any_device = #tt.operand_constraint +module attributes {} { + func.func @forward(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + %0 = tensor.empty() : tensor<64x128xf32> + // CHECK: [[VAL0:%[0-9]+]] = "ttnn.empty"(%{{[0-9]+}}) <{dtype = {{.*}}, layout = {{.*}}, memory_config = {{.*}}, <{{.*}}>>, shape = #ttnn.shape<[[TENSOR_SHAPE:[0-9]+x[0-9]+]]>}> + %1 = "ttir.sign"(%arg0, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + // CHECK: %{{[0-9]+}} = "ttnn.sign"(%{{[0-9]+}}, [[VAL0]]) <{operandSegmentSizes = array}> : (tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}>, tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}) -> tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}> + return %1 : tensor<64x128xf32> + // CHECK: return %{{[0-9]+}} : tensor<[[TENSOR_SHAPE]]xf32, {{.*}}> + } +} diff --git a/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_sign.mlir b/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_sign.mlir new file mode 100644 index 000000000..543a54d3e --- /dev/null +++ b/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_sign.mlir @@ -0,0 +1,14 @@ +// 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 @sign(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + %0 = tensor.empty() : tensor<64x128xf32> + // CHECK: [[VAL0:%[0-9]+]] = "ttnn.empty"(%{{[0-9]+}}) <{dtype = {{.*}}, layout = {{.*}}, memory_config = {{.*}}, <{{.*}}>>, shape = #ttnn.shape<[[TENSOR_SHAPE:[0-9]+x[0-9]+]]>}> + %1 = "ttir.sign"(%arg0, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + // CHECK: %{{[0-9]+}} = "ttnn.sign"(%{{[0-9]+}}, [[VAL0]]) <{operandSegmentSizes = array}> : (tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}>, tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}) -> tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}> + return %1 : tensor<64x128xf32> + // CHECK: return %{{[0-9]+}} : tensor<[[TENSOR_SHAPE]]xf32, {{.*}}> +} diff --git a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir index 816a3c586..4b3cd93bf 100644 --- a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir +++ b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir @@ -209,3 +209,12 @@ func.func @expm1(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { return %1 : tensor<64x128xf32> // CHECK: return %{{[0-9]+}} : tensor<[[TENSOR_SHAPE]]xf32, {{.*}}> } + +func.func @sign(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> { + %0 = tensor.empty() : tensor<64x128xf32> + // CHECK: [[VAL0:%[0-9]+]] = "ttnn.empty"(%{{[0-9]+}}) <{dtype = {{.*}}, layout = {{.*}}, memory_config = {{.*}}, <{{.*}}>>, shape = #ttnn.shape<[[TENSOR_SHAPE:[0-9]+x[0-9]+]]>}> + %1 = "ttir.sign"(%arg0, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32> + // CHECK: %{{[0-9]+}} = "ttnn.sign"(%{{[0-9]+}}, [[VAL0]]) <{operandSegmentSizes = array}> : (tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}>, tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}) -> tensor<[[TENSOR_SHAPE]]x{{.*}}, {{.*}}> + return %1 : tensor<64x128xf32> + // CHECK: return %{{[0-9]+}} : tensor<[[TENSOR_SHAPE]]xf32, {{.*}}> +}