Silu backward (#3195)

* SiLU backward lowering

* Torch pin

* Delete .torch_pin

Author: JackCaoG

test/cpp/test_aten_xla_tensor.cpp

@@ -3428,6 +3428,20 @@ TEST_F(AtenXlaTensorTest, TestSiLU) {
   ExpectCounterChanged("xla::silu_out", cpp_test::GetIgnoredCounters());
 }
 
+TEST_F(AtenXlaTensorTest, TestSiLUBackward) {
+  auto testfn = [&](const std::vector<torch::Tensor>& inputs) -> torch::Tensor {
+    return torch::silu(inputs[0]);
+  };
+  ForEachDevice([&](const torch::Device& device) {
+    TestBackward(
+        {torch::rand({2, 2},
+                     torch::TensorOptions(torch::kFloat).requires_grad(true))},
+        device, testfn);
+  });
+  ExpectCounterNotChanged("aten::.*", cpp_test::GetIgnoredCounters());
+  ExpectCounterChanged("xla::silu_backward", cpp_test::GetIgnoredCounters());
+}
+
 TEST_F(AtenXlaTensorTest, TestSigmoid) {
   torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat));
   torch::Tensor b = torch::sigmoid(a);

torch_xla/csrc/aten_xla_type.cpp

@@ -2894,6 +2894,15 @@ at::Tensor& XLANativeFunctions::silu_out(const at::Tensor& self,
   return out;
 }
 
+at::Tensor XLANativeFunctions::silu_backward(const at::Tensor& grad_output,
+                                             const at::Tensor& self) {
+  XLA_FN_COUNTER("xla::");
+  XLATensor grad_output_tensor = bridge::GetXlaTensor(grad_output);
+  XLATensor self_tensor = bridge::GetXlaTensor(self);
+  return bridge::AtenFromXlaTensor(
+      XLATensor::silu_backward(grad_output_tensor, self_tensor));
+}
+
 at::Tensor XLANativeFunctions::sigmoid(const at::Tensor& self) {
   XLA_FN_COUNTER("xla::");
   return bridge::AtenFromXlaTensor(

torch_xla/csrc/elementwise.cpp

@@ -182,6 +182,13 @@ xla::XlaOp BuildSigmoid(xla::XlaOp input) {
   return half + half * xla::Tanh(half * input);
 }
 
+xla::XlaOp BuildSiLUBackward(xla::XlaOp grad_output, xla::XlaOp input) {
+  const xla::Shape& shape = XlaHelpers::ShapeOfXlaOp(input);
+  xla::XlaOp one = xla::One(input.builder(), shape.element_type());
+  xla::XlaOp input_sigmoid = BuildSigmoid(input);
+  return grad_output * (input_sigmoid * (one + input * (one - input_sigmoid)));
+}
+
 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

@@ -48,6 +48,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 the backward of Silu
+// grad_output * (sigmoid(input) * (1 + input * (1 - sigmoid(input))))
+xla::XlaOp BuildSiLUBackward(xla::XlaOp grad_output, xla::XlaOp input);
+
 // Computes the reciprocal function.
 // Reciprocal(x) = 1 / x
 xla::XlaOp BuildReciprocal(xla::XlaOp input);

torch_xla/csrc/ops/ops.cpp

@@ -219,6 +219,25 @@ NodePtr SiLU(const Value& input) {
                    std::move(lower_fn));
 }
 
+NodePtr SiLUBackward(const Value& grad_output, const Value& input) {
+  auto lower_fn = [](const Node& node, LoweringContext* loctx) -> XlaOpVector {
+    xla::XlaOp xla_grad_output = loctx->GetOutputOp(node.operand(0));
+    xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(1));
+    return node.ReturnOp(BuildSiLUBackward(xla_grad_output, xla_input), loctx);
+  };
+  auto lower_for_shape_fn =
+      [](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
+    return BuildSiLUBackward(operands[0], operands[1]);
+  };
+  return GenericOp(OpKind(at::aten::silu_backward), {grad_output, input},
+                   [&]() {
+                     return InferOutputShape(
+                         {grad_output.shape(), input.shape()},
+                         lower_for_shape_fn);
+                   },
+                   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

@@ -131,6 +131,8 @@ NodePtr Sigmoid(const Value& input);
 
 NodePtr SiLU(const Value& input);
 
+NodePtr SiLUBackward(const Value& grad_output, 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.h

@@ -972,6 +972,7 @@ class XLATensor {
                           xla::int64_t index);
 
   static void silu_out(XLATensor& input, XLATensor& out);
+  static XLATensor silu_backward(XLATensor& grad_output, XLATensor& input);
   static XLATensor sigmoid(const XLATensor& input);
   static XLATensor sigmoid_backward(const XLATensor& grad_output,
                                     const XLATensor& output);

torch_xla/csrc/tensor_methods.cpp

@@ -2311,6 +2311,11 @@ void XLATensor::silu_out(XLATensor& input, XLATensor& out) {
   out.SetInPlaceIrValue(ir::ops::SiLU(input.GetIrValue()));
 }
 
+XLATensor XLATensor::silu_backward(XLATensor& grad_output, XLATensor& input) {
+  return input.CreateFrom(
+      ir::ops::SiLUBackward(grad_output.GetIrValue(), input.GetIrValue()));
+}
+
 XLATensor XLATensor::sigmoid(const XLATensor& input) {
   return input.CreateFrom(ir::ops::Sigmoid(input.GetIrValue()));
 }

xla_native_functions.yaml

@@ -261,6 +261,7 @@ supported:
   - sigmoid_backward
   - sign
   - silu.out
+  - silu_backward
   - sin
   - sinh
   - slice.Tensor