Add support for isfinite and floor op. (#956)

* Add end-to-end implementation of the ops
* Add stablehlo to ttir conversion

include/ttmlir/Dialect/TTIR/IR/TTIROps.td

@@ -211,15 +211,17 @@ def TTIR_CosOp: TTIR_ElementwiseUnaryOp<"cos"> {
     }];
 }
 
-def TTIR_SignOp: TTIR_ElementwiseUnaryOp<"sign"> {
-    let summary = "Eltwise sign operation.";
+def TTIR_FloorOp: TTIR_ElementwiseUnaryOp<"floor"> {
+    let summary = "Eltwise floor op.";
     let description = [{
-      Returns the sign of the `operand` element-wise and produces a `result`
-      tensor.
+      Eltwise floor operation.
+    }];
+}
 
-      Example:
-        %a: [[3, -2, 0], [1, -4, 4]]
-        "ttir.sign"(%a, %out) -> %out: [[1, -1, 0], [1, -1, 1]]
+def TTIR_IsFiniteOp: TTIR_ElementwiseUnaryOp<"isfinite"> {
+    let summary = "Eltwise isfinite op.";
+    let description = [{
+      Eltwise isfinite operation.
     }];
 }
 
@@ -265,6 +267,18 @@ def TTIR_SigmoidOp: TTIR_ElementwiseUnaryOp<"sigmoid"> {
     }];
 }
 
+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_SinOp: TTIR_ElementwiseUnaryOp<"sin"> {
     let summary = "Eltwise sine.";
     let description = [{

include/ttmlir/Dialect/TTNN/IR/TTNNOps.td

@@ -193,6 +193,20 @@ def TTNN_ExpOp : TTNN_ElementwiseUnaryOp<"exp"> {
     }];
 }
 
+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 = [{

include/ttmlir/Target/TTNN/program.fbs

@@ -90,7 +90,9 @@ enum EltwiseOpType: uint32 {
   Log1p = 29,
   Expm1 = 30,
   Sign = 31,
-  Remainder = 32
+  Remainder = 32,
+  IsFinite = 33,
+  Floor = 34,
 }
 
 union EltwiseOpParams {

lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp

@@ -911,6 +911,11 @@ void addElementwiseUnaryOpsConversionPatterns(MLIRContext *ctx,
       mlir::stablehlo::CosineOp, mlir::tt::ttir::CosOp>>(typeConverter, 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<

lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp

@@ -860,6 +860,8 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
            ElementwiseOpConversionPattern<ttir::AbsOp, ttnn::AbsOp>,
            ElementwiseOpConversionPattern<ttir::AddOp, ttnn::AddOp>,
            ElementwiseOpConversionPattern<ttir::CbrtOp, ttnn::CbrtOp>,
+           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>,

lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp

@@ -616,6 +616,8 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   //
   patterns.add<DefaultOpConversionPattern<ttnn::AbsOp>,
                DefaultOpConversionPattern<ttnn::CbrtOp>,
+               DefaultOpConversionPattern<ttnn::FloorOp>,
+               DefaultOpConversionPattern<ttnn::IsFiniteOp>,
                DefaultOpConversionPattern<ttnn::LogicalNotOp>,
                DefaultOpConversionPattern<ttnn::NegOp>,
                DefaultOpConversionPattern<ttnn::ReluOp>,

lib/Target/TTNN/TTNNToFlatbuffer.cpp

@@ -303,6 +303,10 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
     type = ::tt::target::ttnn::EltwiseOpType::Add;
   } else if constexpr (std::is_same_v<EltwiseOp, CbrtOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Cbrt;
+  } 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>) {
@@ -546,6 +550,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);
   }

runtime/lib/ttnn/operations/eltwise/unary/unary.cpp

@@ -62,6 +62,14 @@ void run(const ::tt::target::ttnn::EltwiseOp *op, ProgramContext &context) {
     runEltwiseUnaryOP(op, tensorPool, ::ttnn::cos);
     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;

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<32x32x3xf32>) -> tensor<32x32x3xf32> {
+    %0 = stablehlo.floor %arg0 : tensor<32x32x3xf32>
+    // CHECK: %[[C:.*]] = tensor.empty[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttir.floor"[[C:.*]]
+    return %0 : tensor<32x32x3xf32>
+  }
+}

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<32x32x3xf32>) -> tensor<32x32x3xi1> {
+    // CHECK: %[[E:.*]] = tensor.empty() : tensor<32x32x3xbf16>
+    // CHECK: %[[C:.*]] = "ttir.isfinite"(%arg0, %[[E]])
+    // CHECK-SAME: (tensor<32x32x3xf32>, tensor<32x32x3xbf16>) -> tensor<32x32x3xbf16>
+    %0 = stablehlo.is_finite %arg0 : (tensor<32x32x3xf32>) -> tensor<32x32x3xi1>
+    // CHECK: return %[[C]] : tensor<32x32x3xbf16>
+    return %0 : tensor<32x32x3xi1>
+  }
+}

test/ttmlir/Dialect/TTNN/eltwise/unary/floor/simple_floor.mlir

@@ -0,0 +1,15 @@
+// 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: %{{[0-9]+}} = "ttnn.empty"
+    // CHECK-SAME: [[TENSOR:tensor<64x128xf32,]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %{{[0-9]+}} = "ttnn.floor"
+    // CHECK-SAME: [[TENSOR]]
+    // CHECK-SAME: [[TENSOR]]
+    // CHECK-SAME: -> [[TENSOR]]
+    %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>
+  }
+}

test/ttmlir/Dialect/TTNN/eltwise/unary/isfinite/simple_isfinite.mlir

@@ -0,0 +1,15 @@
+// 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"
+    // CHECK-SAME: [[TENSOR:tensor<64x128xbf16,]]
+    %0 = tensor.empty() : tensor<64x128xbf16>
+    // CHECK: %[[C:.*]] = "ttnn.isfinite"
+    // CHECK-SAME: tensor<64x128xf32,
+    // CHECK-SAME: [[TENSOR]]
+    // CHECK-SAME: -> [[TENSOR]]
+    %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>
+  }
+}

test/ttmlir/Silicon/TTNN/perf_unit/test_perf_floor.mlir

@@ -0,0 +1,17 @@
+// 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<dram|l1|scalar|tile|any_device|any_device_tile>
+#any_device_tile = #tt.operand_constraint<dram|l1|tile|any_device_tile>
+
+func.func @floor(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+  // CHECK: %{{[0-9]+}} = "ttnn.empty"
+  // CHECK-SAME: [[TENSOR:tensor<64x128xf32,]]
+  %0 = tensor.empty() : tensor<64x128xf32>
+  // CHECK: %{{[0-9]+}} = "ttnn.floor"
+  // CHECK-SAME: [[TENSOR]]
+  // CHECK-SAME: [[TENSOR]]
+  // CHECK-SAME: -> [[TENSOR]]
+  %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>
+}

test/ttmlir/Silicon/TTNN/perf_unit/test_perf_isfinite.mlir

@@ -0,0 +1,17 @@
+// 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<dram|l1|scalar|tile|any_device|any_device_tile>
+#any_device_tile = #tt.operand_constraint<dram|l1|tile|any_device_tile>
+
+func.func @is_finite(%arg0: tensor<64x128xf32>) -> tensor<64x128xbf16> {
+  // CHECK: %[[C:.*]] = "ttnn.empty"
+  // CHECK-SAME: [[TENSOR:tensor<64x128xbf16,]]
+  %0 = tensor.empty() : tensor<64x128xbf16>
+  // CHECK: %[[C:.*]] = "ttnn.isfinite"
+  // CHECK-SAME: tensor<64x128xf32,
+  // CHECK-SAME: [[TENSOR]]
+  // CHECK-SAME: -> [[TENSOR]]
+  %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>
+}

test/ttmlir/Silicon/TTNN/simple_eltwise.mlir

@@ -42,6 +42,30 @@ func.func @div(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<6
   return %1 : tensor<64x128xf32>
 }
 
+func.func @floor(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+  // CHECK: %{{[0-9]+}} = "ttnn.empty"
+  // CHECK-SAME: [[TENSOR:tensor<64x128xf32,]]
+  %0 = tensor.empty() : tensor<64x128xf32>
+  // CHECK: %{{[0-9]+}} = "ttnn.floor"
+  // CHECK-SAME: [[TENSOR]]
+  // CHECK-SAME: [[TENSOR]]
+  // CHECK-SAME: -> [[TENSOR]]
+  %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<64x128xbf16> {
+  // CHECK: %[[C:.*]] = "ttnn.empty"
+  // CHECK-SAME: [[TENSOR:tensor<64x128xbf16,]]
+  %0 = tensor.empty() : tensor<64x128xbf16>
+  // CHECK: %[[C:.*]] = "ttnn.isfinite"
+  // CHECK-SAME: tensor<64x128xf32,
+  // CHECK-SAME: [[TENSOR]]
+  // CHECK-SAME: -> [[TENSOR]]
+  %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>
+}
+
 func.func @minimum(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> {
   // CHECK: %[[C:.*]] = "ttnn.empty"
   // CHECK-SAME: [[TENSOR:tensor<64x128xf32,]]