lower sigmoid to mlir, and adding test with varying shapes of inputs (#…

…236)

forge/csrc/passes/lower_to_mlir.cpp

@@ -539,6 +539,7 @@ class MLIRGenerator
             lowering_handler_map["greater_equal"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::GreaterEqualOp>;
             lowering_handler_map["unsqueeze"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::UnsqueezeOp>;
             lowering_handler_map["conv2d"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::Conv2dOp>;
+            lowering_handler_map["sigmoid"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SigmoidOp>;
         }
 };
 }

forge/test/mlir/test_ops.py

@@ -410,3 +410,33 @@ def forward(self, x):
 
     co_out = [co.to("cpu") for co in co_out]
     assert compare_with_golden_pcc(golden=fw_out, calculated=co_out[0], pcc=0.99)
+
+
+@pytest.mark.parametrize("shape", [
+    (7,),           # 1D tensor
+    (32,),          # 1D tensor
+    (7, 32),        # 2D tensor
+    (32, 41),       # 2D tensor
+    (1, 7, 32),     # 3D tensor
+    (1, 32, 41),    # 3D tensor
+    (1, 7, 32, 41), # 4D tensor
+    (2, 7, 32, 41)  # 4D tensor
+])
+def test_sigmoid(shape):
+    class Sigmoid(nn.Module):
+        def __init__(self):
+            super().__init__()
+        def forward(self, x):
+            return torch.sigmoid(x)
+
+    inputs = [
+        torch.rand(*shape),
+    ]
+    framework_model = Sigmoid()
+    fw_out = framework_model(*inputs)
+
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs)
+    co_out = compiled_model(*inputs)
+
+    co_out = [co.to("cpu") for co in co_out]
+    assert compare_with_golden_pcc(golden=fw_out, calculated=co_out[0], pcc=0.99)