From 0604190a159f9ee7f02ec63801e15230a6b1332d Mon Sep 17 00:00:00 2001 From: Kristijan Mitrovic Date: Wed, 30 Oct 2024 16:49:26 +0000 Subject: [PATCH] Added support for stablehlo.remainder op. Added tests. Tested with ttrt. --- include/ttmlir/Dialect/TTIR/IR/TTIROps.td | 15 +++++++++++++++ include/ttmlir/Dialect/TTNN/IR/TTNNOps.td | 15 +++++++++++++++ include/ttmlir/Target/TTNN/program.fbs | 3 ++- .../StableHLOToTTIR/StableHLOToTTIRPatterns.cpp | 2 ++ lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp | 1 + lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp | 4 +++- lib/Target/TTNN/TTNNToFlatbuffer.cpp | 6 ++++++ .../eltwise/binary/binary_composite.cpp | 4 ++++ .../StableHLOToTTIR/binary/remainder_op.mlir | 12 ++++++++++++ .../binary/remainder/simple_remainder.mlir | 12 ++++++++++++ .../TTNN/perf_unit/test_perf_remainder.mlir | 14 ++++++++++++++ test/ttmlir/Silicon/TTNN/simple_eltwise.mlir | 9 +++++++++ 12 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 test/ttmlir/Conversion/StableHLOToTTIR/binary/remainder_op.mlir create mode 100644 test/ttmlir/Dialect/TTNN/eltwise/binary/remainder/simple_remainder.mlir create mode 100644 test/ttmlir/Silicon/TTNN/perf_unit/test_perf_remainder.mlir diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td index 92f491b03d..7934dd4e4e 100644 --- a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td +++ b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td @@ -408,6 +408,21 @@ def TTIR_SubtractOp : TTIR_ElementwiseBinaryOp<"subtract"> { }]; } +def TTIR_RemainderOp : TTIR_ElementwiseBinaryOp<"remainder"> { + let summary = "Eltwise remainder."; + let description = [{ + Performs element-wise remainder of dividend lhs and divisor rhs tensors and produces a + result tensor. + + Example: + + // %lhs: [17, -17, 17, -17] + // %rhs: [3, 3, -3, -3] + %result = "ttir.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64> + // %result: [2, -2, 2, -2] + }]; +} + class TTIR_ReductionOp traits = []> : TTIR_DPSOp { diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td index 4dcf05046a..74140d3e25 100644 --- a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td +++ b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td @@ -376,6 +376,21 @@ def TTNN_SubtractOp : TTNN_ElementwiseBinaryOp<"subtract"> { }]; } +def TTNN_RemainderOp : TTNN_ElementwiseBinaryOp<"remainder"> { + let summary = "Eltwise remainder."; + let description = [{ + Performs element-wise remainder of dividend lhs and divisor rhs tensors and produces a + result tensor. + + Example: + + // %lhs: [17, -17, 17, -17] + // %rhs: [3, 3, -3, -3] + %result = "ttnn.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64> + // %result: [2, -2, 2, -2] + }]; +} + class TTNN_ReductionOp traits = []> : TTNN_Op { let summary = "Reduction op."; let description = [{ diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs index 11e781c626..ca49e99f06 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, + Remainder = 31 } union EltwiseOpParams { diff --git a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp index 2986f486a5..2faf430119 100644 --- a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp +++ b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp @@ -882,6 +882,8 @@ void addElementwiseBinaryOpsConversionPatterns(MLIRContext *ctx, patterns.add>(typeConverter, ctx); + patterns.add>(typeConverter, ctx); } void addReduceOpsConversionPatterns(MLIRContext *ctx, diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp index b0cfc36349..2244e17051 100644 --- a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp +++ b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp @@ -886,6 +886,7 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns, ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, ElementwiseOpConversionPattern, + ElementwiseOpConversionPattern, ReductionOpConversionPattern, ReductionOpConversionPattern, ReductionOpConversionPattern, diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp index f78ed65e11..178176b8fa 100644 --- a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp +++ b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp @@ -646,7 +646,9 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx, DefaultOpConversionPattern, DefaultOpConversionPattern, DefaultOpConversionPattern, - DefaultOpConversionPattern>(typeConverter, ctx); + DefaultOpConversionPattern, + DefaultOpConversionPattern>(typeConverter, + ctx); // Tensor manipulation ops // diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp index cc41042cbf..5d4b392e07 100644 --- a/lib/Target/TTNN/TTNNToFlatbuffer.cpp +++ b/lib/Target/TTNN/TTNNToFlatbuffer.cpp @@ -357,6 +357,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) { type = ::tt::target::ttnn::EltwiseOpType::Log; } else if constexpr (std::is_same_v) { type = ::tt::target::ttnn::EltwiseOpType::Expm1; + } else if constexpr (std::is_same_v) { + type = ::tt::target::ttnn::EltwiseOpType::Remainder; } else { llvm_unreachable("unhandled EltwiseOp"); } @@ -623,6 +625,10 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op, if (auto divOp = dyn_cast(op); divOp) { return createOperation(cache, createEltwiseOp(cache, divOp), debugString); } + if (auto remainderOp = dyn_cast(op); remainderOp) { + return createOperation(cache, createEltwiseOp(cache, remainderOp), + debugString); + } if (auto matmulOp = dyn_cast(op); matmulOp) { return createOperation(cache, createOp(cache, matmulOp), debugString); } diff --git a/runtime/lib/ttnn/operations/eltwise/binary/binary_composite.cpp b/runtime/lib/ttnn/operations/eltwise/binary/binary_composite.cpp index e0dbddf8f2..09b154e6eb 100644 --- a/runtime/lib/ttnn/operations/eltwise/binary/binary_composite.cpp +++ b/runtime/lib/ttnn/operations/eltwise/binary/binary_composite.cpp @@ -38,6 +38,10 @@ void run(const ::tt::target::ttnn::EltwiseOp *op, ProgramContext &context) { runEltwiseBinaryCompositeOP(op, tensorPool, ::ttnn::minimum); break; } + case ::tt::target::ttnn::EltwiseOpType::Remainder: { + runEltwiseBinaryCompositeOP(op, tensorPool, ::ttnn::remainder); + break; + } default: throw std::invalid_argument( "Unsupported Eltwise Binary Composite operation"); diff --git a/test/ttmlir/Conversion/StableHLOToTTIR/binary/remainder_op.mlir b/test/ttmlir/Conversion/StableHLOToTTIR/binary/remainder_op.mlir new file mode 100644 index 0000000000..bbca3a3f99 --- /dev/null +++ b/test/ttmlir/Conversion/StableHLOToTTIR/binary/remainder_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_remainder attributes {} { + func.func public @test_remainder(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>) -> tensor<32x32xf32> { + %0 = stablehlo.remainder %arg0, %arg1 : tensor<32x32xf32> + // CHECK: %[[EMPTY:[0-9]+]] = tensor.empty() : tensor<32x32xf32> + // CHECK: %[[REM:[0-9]+]] = "ttir.remainder"(%arg0, %arg1, %[[EMPTY]]){{.*}} -> tensor<32x32xf32> + return %0 : tensor<32x32xf32> + // CHECK: return %[[REM]] : tensor<32x32xf32> + } +} diff --git a/test/ttmlir/Dialect/TTNN/eltwise/binary/remainder/simple_remainder.mlir b/test/ttmlir/Dialect/TTNN/eltwise/binary/remainder/simple_remainder.mlir new file mode 100644 index 0000000000..281dccfdd2 --- /dev/null +++ b/test/ttmlir/Dialect/TTNN/eltwise/binary/remainder/simple_remainder.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 @remainder(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>) -> tensor<32x32xf32> { + %0 = tensor.empty() : tensor<32x32xf32> + // CHECK: %[[EMPTY:.*]] = "ttnn.empty"{{.*}} -> tensor<32x32xf32, {{.*}} + %1 = "ttir.remainder"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<32x32xf32>, tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32> + // CHECK: %[[REM:[0-9]+]] = "ttnn.remainder"({{.*}}, {{.*}}, %[[EMPTY]]){{.*}} -> tensor<32x32xf32, {{.*}} + return %1 : tensor<32x32xf32> + // CHECK: return {{.*}} : tensor<32x32xf32, {{.*}} + } +} diff --git a/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_remainder.mlir b/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_remainder.mlir new file mode 100644 index 0000000000..68375a9e06 --- /dev/null +++ b/test/ttmlir/Silicon/TTNN/perf_unit/test_perf_remainder.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 @remainder(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>) -> tensor<32x32xf32> { + %0 = tensor.empty() : tensor<32x32xf32> + // CHECK: %[[EMPTY:.*]] = "ttnn.empty"{{.*}} -> tensor<32x32xf32, {{.*}} + %1 = "ttir.remainder"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<32x32xf32>, tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32> + // CHECK: %[[REM:[0-9]+]] = "ttnn.remainder"({{.*}}, {{.*}}, %[[EMPTY]]){{.*}} -> tensor<32x32xf32, {{.*}} + return %1 : tensor<32x32xf32> + // CHECK: return {{.*}} : tensor<32x32xf32, {{.*}} +} diff --git a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir index 816a3c5862..01559b7401 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 @remainder(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>) -> tensor<32x32xf32> { + %0 = tensor.empty() : tensor<32x32xf32> + // CHECK: %[[EMPTY:.*]] = "ttnn.empty"{{.*}} -> tensor<32x32xf32, {{.*}} + %1 = "ttir.remainder"(%arg0, %arg1, %0) <{operandSegmentSizes = array, operand_constraints = [#any_device, #any_device, #any_device]}> : (tensor<32x32xf32>, tensor<32x32xf32>, tensor<32x32xf32>) -> tensor<32x32xf32> + // CHECK: %[[REM:[0-9]+]] = "ttnn.remainder"({{.*}}, {{.*}}, %[[EMPTY]]){{.*}} -> tensor<32x32xf32, {{.*}} + return %1 : tensor<32x32xf32> + // CHECK: return {{.*}} : tensor<32x32xf32, {{.*}} +}