Decompose aten.silu op

Decomposition of aten.silu.op is added as silu(x) = x * sigmoid(x).

Author: Prashant Kumar

e2e_testing/torchscript/basic.py

@@ -1345,3 +1345,21 @@ def forward(self, x):
 def HardswishRandomModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(128, 128, low=-10, high=10))
 
+# ==============================================================================
+
+class SiluModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([-1, -1], torch.float32, True),
+    ])
+    def forward(self, x):
+        return torch.ops.aten.silu(x)
+
+
+@register_test_case(module_factory=lambda: SiluModule())
+def SiluModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(128, 128, low=-10, high=10))

e2e_testing/torchscript/xfail_sets.py

@@ -129,4 +129,5 @@
     "NewOnesModuleFloat2D_basic",
     "NewOnesModuleFloat3D_basic",
     "NewOnesModuleFalsePinMemory_basic",
+    "SiluModule_basic",
 }

include/torch-mlir/Dialect/Torch/IR/GeneratedAtenOps.td

@@ -214,6 +214,34 @@ def Torch_AtenHardswish_Op : Torch_Op<"aten.hardswish_", [
   let assemblyFormat = "$self attr-dict `:` qualified(type($self)) `->` qualified(type($result))";
 }
 
+def Torch_AtenSiluOp : Torch_Op<"aten.silu", [
+    AllowsTypeRefinement,
+    HasValueSemantics
+  ]> {
+  let summary = "Generated op for `aten::silu : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` qualified(type($self)) `->` qualified(type($result))";
+}
+
+def Torch_AtenSilu_Op : Torch_Op<"aten.silu_", [
+    IsTrailingUnderscoreInplaceVariant,
+    AllowsTypeRefinement
+  ]> {
+  let summary = "Generated op for `aten::silu_ : (Tensor) -> (Tensor)`";
+  let arguments = (ins
+    AnyTorchTensorType:$self
+  );
+  let results = (outs
+    AnyTorchTensorType:$result
+  );
+  let assemblyFormat = "$self attr-dict `:` qualified(type($self)) `->` qualified(type($result))";
+}
+
 def Torch_AtenSinOp : Torch_Op<"aten.sin", [
     AllowsTypeRefinement,
     HasValueSemantics

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -740,6 +740,23 @@ class DecomposeAtenSquareOp : public OpRewritePattern<AtenSquareOp> {
 };
 } // namespace
 
+// Silu(x) = sigmoid(x) * x
+namespace {
+class DecomposeAtenSiluOp : public OpRewritePattern<AtenSiluOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(AtenSiluOp op,
+                                PatternRewriter &rewriter) const override {
+    Value self = op.self();
+    Value sigmoid =
+        rewriter.create<AtenSigmoidOp>(op.getLoc(), op.getType(), self);
+    rewriter.replaceOpWithNewOp<AtenMulTensorOp>(op, op.getType(), sigmoid,
+                                                 self);
+    return success();
+  }
+};
+} // namespace
+
 // Decompose aten.var into: sum(square(x - mean))/(numTensorElements-1)
 // for unbiased and mean(square(x - mean)) for biased case.
 namespace {
@@ -1318,6 +1335,8 @@ class DecomposeComplexOpsPass
     target.addIllegalOp<AtenHardsigmoidOp>();
     patterns.add<DecomposeAtenHardswishOp>(context);
     target.addIllegalOp<AtenHardswishOp>();
+    patterns.add<DecomposeAtenSiluOp>(context);
+    target.addIllegalOp<AtenSiluOp>();
     patterns.add<DecomposeConstantTensorNewLikeOp<AtenNewZerosOp, AtenZerosOp>>(
         context);
     target.addIllegalOp<AtenNewZerosOp>();

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -230,7 +230,7 @@ class TypeAnalyzer : public ForwardDataFlowAnalysis<ValueKnowledge> {
             AtenAbsOp, AtenThresholdOp, AtenSquareOp, PseudoAtenUniformOp,
             AtenCloneOp, AtenBernoulliOp, AtenBernoulli_FloatOp,
             PseudoAtenBernoulliFloatOp, PseudoAtenFillScalarOp,
-            AtenHardsigmoidOp, AtenHardswishOp>(op)) {
+            AtenHardsigmoidOp, AtenHardswishOp, AtenSiluOp>(op)) {
       return getLatticeElement(op->getResult(0)).join(*operands[0]);
     }
 

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -452,6 +452,7 @@ def emit_with_mutating_variants(key, **kwargs):
                 "aten::sigmoid : (Tensor) -> (Tensor)",
                 "aten::hardsigmoid : (Tensor) -> (Tensor)",
                 "aten::hardswish : (Tensor) -> (Tensor)",
+                "aten::silu : (Tensor) -> (Tensor)",
                 "aten::sin : (Tensor) -> (Tensor)",
                 "aten::exp : (Tensor) -> (Tensor)",
                 "aten::cos : (Tensor) -> (Tensor)",

test/Dialect/Torch/decompose-complex-ops.mlir

@@ -509,3 +509,14 @@ func @torch.aten.new_ones(%arg0: !torch.vtensor<[?,?],si64>) -> !torch.vtensor<[
   %1 = torch.aten.new_ones %arg0, %0, %none, %none, %none, %none : !torch.vtensor<[?,?],si64>, !torch.list<!torch.int>, !torch.none, !torch.none, !torch.none, !torch.none -> !torch.vtensor<[3,4],si64>
   return %1 : !torch.vtensor<[3,4],si64>
 }
+
+// -----
+// CHECK-LABEL:   func @torch.aten.silu(
+// CHECK-SAME:                  %[[INP:.*]]: !torch.vtensor<[?,?],f32>) -> !torch.vtensor {
+// CHECK:           %[[SIGMOID:.*]] = torch.aten.sigmoid %[[INP]] : !torch.vtensor<[?,?],f32> -> !torch.vtensor
+// CHECK:           %[[MUL:.*]] = torch.aten.mul.Tensor %[[SIGMOID]], %[[INP]] : !torch.vtensor, !torch.vtensor<[?,?],f32> -> !torch.vtensor
+// CHECK:           return %[[MUL]] : !torch.vtensor
+func @torch.aten.silu(%arg0: !torch.vtensor<[?,?],f32> loc(unknown)) -> !torch.vtensor {
+    %0 = torch.aten.silu %arg0 : !torch.vtensor<[?,?],f32> -> !torch.vtensor
+    return %0 : !torch.vtensor
+}