From 6531f76c30b4c39ffcfc5f0f9f9f1487a3a21686 Mon Sep 17 00:00:00 2001
From: Muhammad Asif Manzoor <mmanzoor@tenstorrent.com>
Date: Mon, 21 Oct 2024 13:05:27 +0000
Subject: [PATCH] Add support for isfinite and floor op.

* Add end-to-end implementation of the ops
* Add stablehlo to ttir conversion
---
 include/ttmlir/Dialect/TTIR/IR/TTIROps.td        | 14 ++++++++++++++
 include/ttmlir/Dialect/TTNN/IR/TTNNOps.td        | 14 ++++++++++++++
 include/ttmlir/Target/TTNN/program.fbs           |  2 ++
 .../StableHLOToTTIR/StableHLOToTTIRPatterns.cpp  |  5 +++++
 lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp         |  8 +++++---
 lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp       | 11 +++++++----
 lib/Target/TTNN/TTNNToFlatbuffer.cpp             | 11 +++++++++++
 runtime/lib/ttnn/operations/eltwise/unary.cpp    |  8 ++++++++
 .../Conversion/StableHLOToTTIR/floor_op.mlir     | 11 +++++++++++
 .../Conversion/StableHLOToTTIR/isfinite_op.mlir  | 13 +++++++++++++
 .../TTNN/eltwise/unary/floor/simple_floor.mlir   | 11 +++++++++++
 .../eltwise/unary/isfinite/simple_isfinite.mlir  | 11 +++++++++++
 test/ttmlir/Silicon/TTNN/simple_eltwise.mlir     | 16 ++++++++++++++++
 13 files changed, 128 insertions(+), 7 deletions(-)
 create mode 100644 test/ttmlir/Conversion/StableHLOToTTIR/floor_op.mlir
 create mode 100644 test/ttmlir/Conversion/StableHLOToTTIR/isfinite_op.mlir
 create mode 100644 test/ttmlir/Dialect/TTNN/eltwise/unary/floor/simple_floor.mlir
 create mode 100644 test/ttmlir/Dialect/TTNN/eltwise/unary/isfinite/simple_isfinite.mlir

diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td
index aee213b5af..5eb2d76bf5 100644
--- a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td
+++ b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td
@@ -217,6 +217,20 @@ def TTIR_NegOp: TTIR_ElementwiseUnaryOp<"neg"> {
     }];
 }
 
+def TTIR_FloorOp: TTIR_ElementwiseUnaryOp<"floor"> {
+    let summary = "Eltwise floor op.";
+    let description = [{
+      Eltwise floor operation.
+    }];
+}
+
+def TTIR_IsFiniteOp: TTIR_ElementwiseUnaryOp<"isfinite"> {
+    let summary = "Eltwise isfinite op.";
+    let description = [{
+      Eltwise isfinite operation.
+    }];
+}
+
 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 46d0fcf52a..c84182ba19 100644
--- a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td
+++ b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td
@@ -146,6 +146,20 @@ def TTNN_NegOp : TTNN_ElementwiseUnaryOp<"neg"> {
     }];
 }
 
+def TTNN_FloorOp: TTNN_ElementwiseUnaryOp<"floor"> {
+    let summary = "Eltwise floor op.";
+    let description = [{
+      Eltwise floor operation.
+    }];
+}
+
+def TTNN_IsFiniteOp: TTNN_ElementwiseUnaryOp<"isfinite"> {
+    let summary = "Eltwise isfinite op.";
+    let description = [{
+      Eltwise isfinite operation.
+    }];
+}
+
 def TTNN_LogicalNotOp: TTNN_ElementwiseUnaryOp<"logical_not"> {
     let summary = "Eltwise logical not op.";
     let description = [{
diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs
index bc98664dbb..c3146867c5 100644
--- a/include/ttmlir/Target/TTNN/program.fbs
+++ b/include/ttmlir/Target/TTNN/program.fbs
@@ -73,6 +73,8 @@ enum EltwiseOpType: uint32 {
   LogicalAnd = 20,
   LogicalOr = 21,
   LogicalNot = 22,
+  IsFinite = 23,
+  Floor = 24,
 }
 
 table EltwiseOp {
diff --git a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp
index 4890e0ad20..6fa9b5f128 100644
--- a/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp
+++ b/lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp
@@ -863,6 +863,11 @@ void addElementwiseUnaryOpsConversionPatterns(MLIRContext *ctx,
                                                                ctx);
   patterns.add<StableHLOToTTIROpDefaultConversionPattern<
       mlir::stablehlo::ExpOp, mlir::tt::ttir::ExpOp>>(typeConverter, ctx);
+  patterns.add<StableHLOToTTIROpDefaultConversionPattern<
+      mlir::stablehlo::FloorOp, mlir::tt::ttir::FloorOp>>(typeConverter, ctx);
+  patterns.add<StableHLOToTTIROpDefaultConversionPattern<
+      mlir::stablehlo::IsFiniteOp, mlir::tt::ttir::IsFiniteOp>>(typeConverter,
+                                                                ctx);
   patterns.add<StableHLOToTTIROpDefaultConversionPattern<
       mlir::stablehlo::NegOp, mlir::tt::ttir::NegOp>>(typeConverter, ctx);
   patterns.add<StableHLOToTTIROpDefaultConversionPattern<
diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp
index c9052c8661..b2b68fc2d0 100644
--- a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp
+++ b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp
@@ -729,11 +729,15 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
            ToLayoutOpConversionPattern,
            ElementwiseOpConversionPattern<ttir::AbsOp, ttnn::AbsOp>,
            ElementwiseOpConversionPattern<ttir::AddOp, ttnn::AddOp>,
+           ElementwiseOpConversionPattern<ttir::DivOp, ttnn::DivOp>,
+           ElementwiseOpConversionPattern<ttir::EqualOp, ttnn::EqualOp>,
+           ElementwiseOpConversionPattern<ttir::ExpOp, ttnn::ExpOp>,
+           ElementwiseOpConversionPattern<ttir::FloorOp, ttnn::FloorOp>,
+           ElementwiseOpConversionPattern<ttir::IsFiniteOp, ttnn::IsFiniteOp>,
            ElementwiseOpConversionPattern<ttir::LogicalAndOp, ttnn::LogicalAndOp>,
            ElementwiseOpConversionPattern<ttir::LogicalOrOp, ttnn::LogicalOrOp>,
            ElementwiseOpConversionPattern<ttir::LogicalNotOp, ttnn::LogicalNotOp>,
            ElementwiseOpConversionPattern<ttir::MultiplyOp, ttnn::MultiplyOp>,
-           ElementwiseOpConversionPattern<ttir::EqualOp, ttnn::EqualOp>,
            ElementwiseOpConversionPattern<ttir::NotEqualOp, ttnn::NotEqualOp>,
            ElementwiseOpConversionPattern<ttir::GreaterEqualOp, ttnn::GreaterEqualOp>,
            ElementwiseOpConversionPattern<ttir::GreaterThanOp, ttnn::GreaterThanOp>,
@@ -747,8 +751,6 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
            ElementwiseOpConversionPattern<ttir::SigmoidOp, ttnn::SigmoidOp>,
            ElementwiseOpConversionPattern<ttir::TypecastOp, ttnn::TypecastOp>,
            ElementwiseOpConversionPattern<ttir::ReciprocalOp, ttnn::ReciprocalOp>,
-           ElementwiseOpConversionPattern<ttir::ExpOp, ttnn::ExpOp>,
-           ElementwiseOpConversionPattern<ttir::DivOp, ttnn::DivOp>,
            ReductionOpConversionPattern<ttir::SumOp, ttnn::SumOp>,
            ReductionOpConversionPattern<ttir::MeanOp, ttnn::MeanOp>,
            ReductionOpConversionPattern<ttir::MaxOp, ttnn::MaxOp>,
diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp
index 946862c233..9a989c7a0e 100644
--- a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp
+++ b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp
@@ -527,15 +527,18 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   // Eltwise unary ops
   //
   patterns.add<DefaultOpConversionPattern<ttnn::AbsOp>,
+               DefaultOpConversionPattern<ttnn::ExpOp>,
+               DefaultOpConversionPattern<ttnn::FloorOp>,
+               DefaultOpConversionPattern<ttnn::IsFiniteOp>,
                DefaultOpConversionPattern<ttnn::LogicalNotOp>,
                DefaultOpConversionPattern<ttnn::NegOp>,
+               DefaultOpConversionPattern<ttnn::ReciprocalOp>,
                DefaultOpConversionPattern<ttnn::ReluOp>,
-               DefaultOpConversionPattern<ttnn::SqrtOp>,
                DefaultOpConversionPattern<ttnn::RsqrtOp>,
                DefaultOpConversionPattern<ttnn::SigmoidOp>,
-               DefaultOpConversionPattern<ttnn::TypecastOp>,
-               DefaultOpConversionPattern<ttnn::ReciprocalOp>,
-               DefaultOpConversionPattern<ttnn::ExpOp>>(typeConverter, ctx);
+               DefaultOpConversionPattern<ttnn::SqrtOp>,
+               DefaultOpConversionPattern<ttnn::TypecastOp>>(typeConverter,
+                                                             ctx);
 
   // Eltwise binary ops
   //
diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp
index 40ed701fd8..1b84c553b9 100644
--- a/lib/Target/TTNN/TTNNToFlatbuffer.cpp
+++ b/lib/Target/TTNN/TTNNToFlatbuffer.cpp
@@ -258,6 +258,10 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
     type = ::tt::target::ttnn::EltwiseOpType::Abs;
   } else if constexpr (std::is_same_v<EltwiseOp, AddOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Add;
+  } else if constexpr (std::is_same_v<EltwiseOp, FloorOp>) {
+    type = ::tt::target::ttnn::EltwiseOpType::Floor;
+  } else if constexpr (std::is_same_v<EltwiseOp, IsFiniteOp>) {
+    type = ::tt::target::ttnn::EltwiseOpType::IsFinite;
   } else if constexpr (std::is_same_v<EltwiseOp, LogicalAndOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::LogicalAnd;
   } else if constexpr (std::is_same_v<EltwiseOp, LogicalNotOp>) {
@@ -481,6 +485,13 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
   if (auto addOp = dyn_cast<AddOp>(op); addOp) {
     return createOperation(cache, createEltwiseOp(cache, addOp), debugString);
   }
+  if (auto floorOp = dyn_cast<FloorOp>(op); floorOp) {
+    return createOperation(cache, createEltwiseOp(cache, floorOp), debugString);
+  }
+  if (auto isFiniteOp = dyn_cast<IsFiniteOp>(op); isFiniteOp) {
+    return createOperation(cache, createEltwiseOp(cache, isFiniteOp),
+                           debugString);
+  }
   if (auto andOp = dyn_cast<LogicalAndOp>(op); andOp) {
     return createOperation(cache, createEltwiseOp(cache, andOp), debugString);
   }
diff --git a/runtime/lib/ttnn/operations/eltwise/unary.cpp b/runtime/lib/ttnn/operations/eltwise/unary.cpp
index 0566732744..7775b1129a 100644
--- a/runtime/lib/ttnn/operations/eltwise/unary.cpp
+++ b/runtime/lib/ttnn/operations/eltwise/unary.cpp
@@ -76,6 +76,14 @@ void run(const ::tt::target::ttnn::EltwiseOp *op, ProgramContext &context) {
     runEltwiseUnaryOP(op, tensorPool, ::ttnn::abs);
     break;
   }
+  case ::tt::target::ttnn::EltwiseOpType::Floor: {
+    runEltwiseUnaryOP(op, tensorPool, ::ttnn::floor);
+    break;
+  }
+  case ::tt::target::ttnn::EltwiseOpType::IsFinite: {
+    runEltwiseUnaryOP(op, tensorPool, ::ttnn::isfinite);
+    break;
+  }
   case ::tt::target::ttnn::EltwiseOpType::LogicalNot: {
     runEltwiseUnaryOP(op, tensorPool, ::ttnn::logical_not);
     break;
diff --git a/test/ttmlir/Conversion/StableHLOToTTIR/floor_op.mlir b/test/ttmlir/Conversion/StableHLOToTTIR/floor_op.mlir
new file mode 100644
index 0000000000..c67473ead1
--- /dev/null
+++ b/test/ttmlir/Conversion/StableHLOToTTIR/floor_op.mlir
@@ -0,0 +1,11 @@
+// REQUIRES: stablehlo
+// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | FileCheck %s
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module @jit_eltwise_floor attributes {} {
+  func.func public @test_floor(%arg0: tensor<13x21x3xf32>) -> tensor<13x21x3xf32> {
+    %0 = stablehlo.floor %arg0 : tensor<13x21x3xf32>
+    // CHECK: %[[C:.*]] = tensor.empty[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttir.floor"[[C:.*]]
+    return %0 : tensor<13x21x3xf32>
+  }
+}
diff --git a/test/ttmlir/Conversion/StableHLOToTTIR/isfinite_op.mlir b/test/ttmlir/Conversion/StableHLOToTTIR/isfinite_op.mlir
new file mode 100644
index 0000000000..b87dae6231
--- /dev/null
+++ b/test/ttmlir/Conversion/StableHLOToTTIR/isfinite_op.mlir
@@ -0,0 +1,13 @@
+// REQUIRES: stablehlo
+// RUN: ttmlir-opt --stablehlo-to-ttir-pipeline %s | FileCheck %s
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module @jit_eltwise_isfinite attributes {} {
+  func.func public @test_isfinite(%arg0: tensor<13x31x3xf32>) -> tensor<13x31x3xi1> {
+    // CHECK: %[[E:.*]] = tensor.empty() : tensor<13x31x3xbf16>
+    // CHECK: %[[C:.*]] = "ttir.isfinite"(%arg0, %[[E]])
+    // CHECK-SAME: (tensor<13x31x3xf32>, tensor<13x31x3xbf16>) -> tensor<13x31x3xbf16>
+    %0 = stablehlo.is_finite %arg0 : (tensor<13x31x3xf32>) -> tensor<13x31x3xi1>
+    // CHECK: return %[[C]] : tensor<13x31x3xbf16>
+    return %0 : tensor<13x31x3xi1>
+  }
+}
diff --git a/test/ttmlir/Dialect/TTNN/eltwise/unary/floor/simple_floor.mlir b/test/ttmlir/Dialect/TTNN/eltwise/unary/floor/simple_floor.mlir
new file mode 100644
index 0000000000..4bcf02eadc
--- /dev/null
+++ b/test/ttmlir/Dialect/TTNN/eltwise/unary/floor/simple_floor.mlir
@@ -0,0 +1,11 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module attributes {} {
+  func.func @floor(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.floor"[[C:.*]]
+    %1 = "ttir.floor"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+    return %1 : tensor<64x128xf32>
+  }
+}
diff --git a/test/ttmlir/Dialect/TTNN/eltwise/unary/isfinite/simple_isfinite.mlir b/test/ttmlir/Dialect/TTNN/eltwise/unary/isfinite/simple_isfinite.mlir
new file mode 100644
index 0000000000..26be8fbece
--- /dev/null
+++ b/test/ttmlir/Dialect/TTNN/eltwise/unary/isfinite/simple_isfinite.mlir
@@ -0,0 +1,11 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s
+#any_device = #tt.operand_constraint<dram|l1|scalar|tile|any_device|any_device_tile>
+module attributes {} {
+  func.func @is_finite(%arg0: tensor<64x128xf32>) -> tensor<64x128xbf16> {
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xbf16>
+    // CHECK: %[[C:.*]] = "ttnn.isfinite"[[C:.*]]
+    %1 = "ttir.isfinite"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xbf16>) -> tensor<64x128xbf16>
+    return %1 : tensor<64x128xbf16>
+  }
+}
diff --git a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir
index 3d5cbba155..8e30d75c42 100644
--- a/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir
+++ b/test/ttmlir/Silicon/TTNN/simple_eltwise.mlir
@@ -20,6 +20,22 @@ func.func @div(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<6
   return %1 : tensor<64x128xf32>
 }
 
+func.func @floor(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+  // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+  %0 = tensor.empty() : tensor<64x128xf32>
+  // CHECK: %[[C:.*]] = "ttnn.floor"[[C:.*]]
+  %1 = "ttir.floor"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+  return %1 : tensor<64x128xf32>
+}
+
+func.func @is_finite(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+  // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+  %0 = tensor.empty() : tensor<64x128xf32>
+  // CHECK: %[[C:.*]] = "ttnn.isfinite"[[C:.*]]
+  %1 = "ttir.isfinite"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+  return %1 : tensor<64x128xf32>
+}
+
 func.func @multiply(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> {
   // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
   %0 = tensor.empty() : tensor<64x128xf32>