Adding silu to ops.cpp #2717

Author: tyoc213

test/cpp/test_aten_xla_tensor.cpp

@@ -3067,6 +3067,16 @@ TEST_F(AtenXlaTensorTest, TestLogsumexp) {
   }
 }
 
+TEST_F(AtenXlaTensorTest, TestSiLU) {
+  torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));
+  torch::Tensor b = torch::silu(a);
+  ForEachDevice([&](const torch::Device& device) {
+    torch::Tensor xla_a = CopyToDevice(a, device);
+    torch::Tensor xla_b = torch::silu(xla_a);
+    AllClose(b, xla_b, /*rtol=*/1e-3, /*atol=*/1e-5);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestSigmoid) {
   torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));
   torch::Tensor b = torch::sigmoid(a);

torch_xla/csrc/elementwise.cpp

@@ -180,6 +180,14 @@ xla::XlaOp BuildSigmoid(xla::XlaOp input) {
   return half + half * xla::Tanh(half * input);
 }
 
+
+xla::XlaOp BuildSiLU(xla::XlaOp input){
+  const xla::Shape& shape = XlaHelpers::ShapeOfXlaOp(input);
+  xla::XlaOp half = XlaHelpers::ScalarValue<float>(0.5, shape.element_type(),
+                                                   input.builder());
+  return input * (half + half * xla::Tanh(half * input));
+}
+
 xla::XlaOp BuildReciprocal(xla::XlaOp input) {
   const xla::Shape& shape = XlaHelpers::ShapeOfXlaOp(input);
   xla::XlaOp one = xla::One(input.builder(), shape.element_type());

torch_xla/csrc/elementwise.h

@@ -47,6 +47,10 @@ xla::XlaOp BuildLeakyReluBackward(xla::XlaOp grad_output, xla::XlaOp input,
 // Sigmoid(x) = (tanh(x ∗ 0.5) + 1) ∗ 0.5
 xla::XlaOp BuildSigmoid(xla::XlaOp input);
 
+// Computes silu function using Tanh
+// SiLU(x) = x * (tanh(x ∗ 0.5) + 1) ∗ 0.5
+xla::XlaOp BuildSiLU(xla::XlaOp input);
+
 // Computes the reciprocal function.
 // Reciprocal(x) = 1 / x
 xla::XlaOp BuildReciprocal(xla::XlaOp input);

torch_xla/csrc/ops/ops.cpp

@@ -206,6 +206,15 @@ NodePtr LogSigmoidBackward(const Value& grad_output, const Value& input,
   return grad_output * (Neg(max_deriv) - sign * (buffer - one) / buffer);
 }
 
+NodePtr SiLU(const Value& input) {
+  auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
+    xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
+    return node.ReturnOp(BuildSiLU(xla_input), loctx);
+  };
+  return GenericOp(OpKind(at::aten::silu), {input}, input.shape(),
+                   std::move(lower_fn));;
+}
+
 NodePtr Sigmoid(const Value& input) {
   auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
     xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));

torch_xla/csrc/ops/ops.h

@@ -126,6 +126,8 @@ NodePtr LogSigmoidBackward(const Value& grad_output, const Value& input,
 
 NodePtr Sigmoid(const Value& input);
 
+NodePtr SiLU(const Value& input);
+
 NodePtr SigmoidBackward(const Value& grad_output, const Value& output);
 
 NodePtr LogSoftmaxBackwardOp(const Value& grad_output, const Value& output,

torch_xla/csrc/tensor_methods.cpp

@@ -2337,9 +2337,8 @@ XLATensor XLATensor::select(const XLATensor& input, xla::int64 dim,
   return tensor_ops::Select(input, dim, index);
 }
 
-void XLATensor::silu_out(XLATensor& input, XLATensor& out){
-  auto value = input.GetIrValue() * ir::ops::Sigmoid(input.GetIrValue());
-  out.SetInPlaceIrValue(value);
+void XLATensor::silu_out(XLATensor& input, XLATensor& out) {
+  out.SetInPlaceIrValue(ir::ops::SiLU(input.GetIrValue()));
 }
 
 XLATensor XLATensor::sigmoid(const XLATensor& input) {