Add sigmoid op (#469)

tenstorrent · Aug 23, 2024 · 855d49e · 855d49e
1 parent cfff32f
commit 855d49e
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 4 deletions.
diff --git a/docs/src/adding-an-op.md b/docs/src/adding-an-op.md
@@ -147,6 +147,7 @@ Invoked as part of the rewrite set:
 MatmulOpConversionPattern
 ```
 
+### Note:
 We also need to add this op to the C++ emitter,
 `lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp` see
 `populateTTNNToEmitCPatterns(...)`.

diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIROps.td b/include/ttmlir/Dialect/TTIR/IR/TTIROps.td
@@ -252,6 +252,13 @@ def TTIR_ReluOp : TTIR_ElementwiseUnaryOp<"relu"> {
     }];
 }
 
+def TTIR_SigmoidOp: TTIR_ElementwiseUnaryOp<"sigmoid"> {
+    let summary = "Eltwise sigmoid.";
+    let description = [{
+      Eltwise sigmoid operation.
+    }];
+}
+
 class TTIR_ReductionOp<string mnemonic, list<Trait> traits = []> : TTIR_DPSOp<mnemonic, traits> {
     let summary = "Reduction op.";
     let description = [{

diff --git a/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td b/include/ttmlir/Dialect/TTNN/IR/TTNNOps.td
@@ -104,6 +104,13 @@ def TTNN_ReluOp : TTNN_ElementwiseUnaryOp<"relu"> {
     }];
 }
 
+def TTNN_SigmoidOp : TTNN_ElementwiseUnaryOp<"sigmoid"> {
+    let summary = "Eltwise sigmoid.";
+    let description = [{
+      Eltwise sigmoid operation.
+    }];
+}
+
 def TTNN_AddOp : TTNN_ElementwiseBinaryOp<"add"> {
     let summary = "Eltwise add.";
     let description = [{

diff --git a/include/ttmlir/Target/TTNN/program.fbs b/include/ttmlir/Target/TTNN/program.fbs
@@ -36,7 +36,8 @@ enum EltwiseOpType: uint32 {
   Relu = 3,
   GreaterEqual = 4,
   Sqrt = 5,
-  Div = 6
+  Div = 6,
+  Sigmoid = 7,
 }
 
 table EltwiseOp {

diff --git a/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp b/lib/Conversion/TTIRToTTNN/TTIRToTTNN.cpp
@@ -174,6 +174,7 @@ void populateTTIRToTTNNPatterns(MLIRContext *ctx, RewritePatternSet &patterns,
            ElementwiseBinaryOpConversionPattern<ttir::GreaterEqualOp, ttnn::GreaterEqualOp>,
            ElementwiseBinaryOpConversionPattern<ttir::ReluOp, ttnn::ReluOp>,
            ElementwiseBinaryOpConversionPattern<ttir::SqrtOp, ttnn::SqrtOp>,
+           ElementwiseBinaryOpConversionPattern<ttir::SigmoidOp, ttnn::SigmoidOp>,
            ElementwiseBinaryOpConversionPattern<ttir::DivOp, ttnn::DivOp>,
            ReductionOpConversionPattern<ttir::SumOp, ttnn::SumOp>,
            ReductionOpConversionPattern<ttir::MeanOp, ttnn::MeanOp>,

diff --git a/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp b/lib/Conversion/TTNNToEmitC/TTNNToEmitC.cpp
@@ -171,7 +171,7 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   //
   patterns.add<DefaultOpConversionPattern<ttnn::ReluOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::SqrtOp>>(typeConverter, ctx);
-  patterns.add<DefaultOpConversionPattern<ttnn::SoftmaxOp>>(typeConverter, ctx);
+  patterns.add<DefaultOpConversionPattern<ttnn::SigmoidOp>>(typeConverter, ctx);
 
   // Eltwise binary ops
   //
@@ -197,6 +197,9 @@ void populateTTNNToEmitCPatterns(mlir::MLIRContext *ctx,
   //
   patterns.add<DefaultOpConversionPattern<ttnn::SumOp>>(typeConverter, ctx);
   patterns.add<DefaultOpConversionPattern<ttnn::MeanOp>>(typeConverter, ctx);
+
+  // Other ops
+  patterns.add<DefaultOpConversionPattern<ttnn::SoftmaxOp>>(typeConverter, ctx);
 }
 
 } // namespace mlir::tt
diff --git a/lib/Target/TTNN/TTNNToFlatbuffer.cpp b/lib/Target/TTNN/TTNNToFlatbuffer.cpp
@@ -139,6 +139,8 @@ createEltwiseOp(FlatbufferObjectCache &cache, EltwiseOp op) {
     type = ::tt::target::ttnn::EltwiseOpType::Sqrt;
   } else if constexpr (std::is_same_v<EltwiseOp, DivOp>) {
     type = ::tt::target::ttnn::EltwiseOpType::Div;
+  } else if constexpr (std::is_same_v<EltwiseOp, SigmoidOp>) {
+    type = ::tt::target::ttnn::EltwiseOpType::Sigmoid;
   } else {
     llvm_unreachable("unhandled EltwiseOp");
   }
@@ -242,6 +244,10 @@ emitTTNNOperation(FlatbufferObjectCache &cache, Operation *op,
   if (auto sqrtOp = dyn_cast<SqrtOp>(op); sqrtOp) {
     return createOperation(cache, createEltwiseOp(cache, sqrtOp), debugString);
   }
+  if (auto sigmoidOp = dyn_cast<SigmoidOp>(op); sigmoidOp) {
+    return createOperation(cache, createEltwiseOp(cache, sigmoidOp),
+                           debugString);
+  }
   if (auto divOp = dyn_cast<DivOp>(op); divOp) {
     return createOperation(cache, createEltwiseOp(cache, divOp), debugString);
   }

diff --git a/runtime/lib/ttnn/program.cpp b/runtime/lib/ttnn/program.cpp
@@ -256,6 +256,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::Sigmoid: {
+    assert(op->ins()->size() == 1 && "Unsupported number of inputs");
+    ::ttnn::Tensor &in = *liveTensors.at(op->ins()->Get(0)->global_id());
+    tensorPool.push_back(::ttnn::sigmoid(in));
+    liveTensors.insert_or_assign(op->out()->global_id(), &tensorPool.back());
+    break;
+  }
   }
 }
 
@@ -296,7 +303,7 @@ run(::tt::target::ttnn::ReductionOp const *op, ::ttnn::Device &device,
 }
 
 static void
-run(::tt::target::ttnn::SoftmaxOp const *op, ::ttnn::device::Device &device,
+run(::tt::target::ttnn::SoftmaxOp const *op, ::ttnn::Device &device,
     std::unordered_map<std::uint32_t, ::ttnn::Tensor *> &liveTensors,
     std::list<::ttnn::Tensor> &tensorPool) {
   ::ttnn::Tensor &in = *liveTensors.at(op->in()->global_id());
@@ -307,7 +314,7 @@ run(::tt::target::ttnn::SoftmaxOp const *op, ::ttnn::device::Device &device,
 }
 
 static void
-run(::tt::target::ttnn::TransposeOp const *op, ::ttnn::device::Device &device,
+run(::tt::target::ttnn::TransposeOp const *op, ::ttnn::Device &device,
     std::unordered_map<std::uint32_t, ::ttnn::Tensor *> &liveTensors,
     std::list<::ttnn::Tensor> &tensorPool) {
   ::ttnn::Tensor &in = *liveTensors.at(op->in()->global_id());

diff --git a/test/ttmlir/Dialect/TTNN/eltwise/unary/sigmoid/simple_sigmoid.mlir b/test/ttmlir/Dialect/TTNN/eltwise/unary/sigmoid/simple_sigmoid.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.sigmoid"[[C:.*]]
+    %1 = "ttir.sigmoid"(%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>
+  }
+}