Adding reciprocal op in ttmlir e2e (#472)

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

@@ -245,6 +245,13 @@ def TTIR_SqrtOp : TTIR_ElementwiseUnaryOp<"sqrt"> {
     }];
 }
 
+def TTIR_ReciprocalOp : TTIR_ElementwiseUnaryOp<"reciprocal"> {
+    let summary = "Eltwise reciprocal.";
+    let description = [{
+      Eltwise reciprocal operation.
+    }];
+}
+
 def TTIR_ReluOp : TTIR_ElementwiseUnaryOp<"relu"> {
     let summary = "Eltwise ReLU.";
     let description = [{

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

@@ -97,6 +97,13 @@ def TTNN_SqrtOp : TTNN_ElementwiseUnaryOp<"sqrt"> {
     }];
 }
 
+def TTNN_ReciprocalOp : TTNN_ElementwiseUnaryOp<"reciprocal"> {
+    let summary = "Eltwise reciprocal.";
+    let description = [{
+      Eltwise reciprocal operation.
+    }];
+}
+
 def TTNN_ReluOp : TTNN_ElementwiseUnaryOp<"relu"> {
     let summary = "Eltwise ReLU.";
     let description = [{

include/ttmlir/Target/TTNN/program.fbs

@@ -38,6 +38,7 @@ enum EltwiseOpType: uint32 {
   Sqrt = 5,
   Div = 6,
   Sigmoid = 7,
+  Reciprocal = 8,
 }
 
 table EltwiseOp {

lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp

@@ -175,6 +175,7 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
            ElementwiseBinaryOpConversionPattern<ttir::ReluOp, ttnn::ReluOp>,
            ElementwiseBinaryOpConversionPattern<ttir::SqrtOp, ttnn::SqrtOp>,
            ElementwiseBinaryOpConversionPattern<ttir::SigmoidOp, ttnn::SigmoidOp>,
+           ElementwiseBinaryOpConversionPattern<ttir::ReciprocalOp, ttnn::ReciprocalOp>,
            ElementwiseBinaryOpConversionPattern<ttir::DivOp, ttnn::DivOp>,
            ReductionOpConversionPattern<ttir::SumOp, ttnn::SumOp>,
            ReductionOpConversionPattern<ttir::MeanOp, ttnn::MeanOp>,

lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp

@@ -172,6 +172,8 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   patterns.add<DefaultOpConversionPattern<ttnn::ReluOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SqrtOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SigmoidOp>>(typeConverter, ctx);
+  patterns.add<DefaultOpConversionPattern<ttnn::ReciprocalOp>>(typeConverter,
+                                                               ctx);
 
   // Eltwise binary ops
   //

lib/Target/TTNN/TTNNToFlatbuffer.cpp

@@ -137,6 +137,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
     type = ::tt::target::ttnn::EltwiseOpType::Relu;
   } else if constexpr (std::is_same_v<EltwiseOp, SqrtOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Sqrt;
+  } else if constexpr (std::is_same_v<EltwiseOp, ReciprocalOp>) {
+    type = ::tt::target::ttnn::EltwiseOpType::Reciprocal;
   } else if constexpr (std::is_same_v<EltwiseOp, DivOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Div;
   } else if constexpr (std::is_same_v<EltwiseOp, SigmoidOp>) {
@@ -248,6 +250,10 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
     return createOperation(cache, createEltwiseOp(cache, sigmoidOp),
                            debugString);
   }
+  if (auto reciprocalOp = dyn_cast<ReciprocalOp>(op); reciprocalOp) {
+    return createOperation(cache, createEltwiseOp(cache, reciprocalOp),
+                           debugString);
+  }
   if (auto divOp = dyn_cast<DivOp>(op); divOp) {
     return createOperation(cache, createEltwiseOp(cache, divOp), debugString);
   }

runtime/lib/ttnn/program.cpp

@@ -10,6 +10,7 @@
 
 #include "tt/runtime/detail/ttnn.h"
 #include "tt/runtime/runtime.h"
+#include "ttmlir/Target/TTNN/program_generated.h"
 #include "ttnn/tensor/types.hpp"
 #include "ttnn/types.hpp"
 #include "types_generated.h"
@@ -266,6 +267,13 @@ run(::tt::target::ttnn::EltwiseOp const *op, ::ttnn::Device &device,
     liveTensors.insert_or_assign(op->out()->global_id(), &tensorPool.back());
     break;
   }
+  case ::tt::target::ttnn::EltwiseOpType::Reciprocal: {
+    assert(op->ins()->size() == 1 && "Unsupported number of inputs");
+    ::ttnn::Tensor &in = *liveTensors.at(op->ins()->Get(0)->global_id());
+    tensorPool.push_back(::ttnn::reciprocal(in));
+    liveTensors.insert_or_assign(op->out()->global_id(), &tensorPool.back());
+    break;
+  }
   }
 }
 

test/ttmlir/Dialect/TTNN/eltwise/unary/reciprocal/simple_reciprocal.mlir

@@ -0,0 +1,13 @@
+// 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 @forward(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+    // CHECK: %[[C:.*]] = "ttnn.open_device"[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.reciprocal"[[C:.*]]
+    %1 = "ttir.reciprocal"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+    // CHECK: "ttnn.close_device"[[C:.*]]
+    return %1 : tensor<64x128xf32>
+  }
+}

test/ttmlir/Silicon/TTNN/eltwise/unary/simple_reciprocal.mlir

@@ -0,0 +1,15 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline --ttir-load-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>
+module attributes {} {
+  func.func @forward(%arg0: tensor<64x128xf32>) -> tensor<64x128xf32> {
+    // CHECK: %[[C:.*]] = "ttnn.open_device"[[C:.*]]
+    // CHECK: %[[C:.*]] = "ttnn.empty"[[C:.*]]
+    %0 = tensor.empty() : tensor<64x128xf32>
+    // CHECK: %[[C:.*]] = "ttnn.reciprocal"[[C:.*]]
+    %1 = "ttir.reciprocal"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>, operand_constraints = [#any_device, #any_device]}> : (tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
+    // CHECK: "ttnn.close_device"[[C:.*]]
+    return %1 : tensor<64x128xf32>
+  }
+}