Codegen selu and silu

Author: JackCaoG

test/cpp/test_aten_xla_tensor.cpp

@@ -6386,8 +6386,8 @@ TEST_F(AtenXlaTensorTest, TestSeluInPlace) {
     AllClose(input, xla_input);
   });
 
+  // selu_ uses elu_ instead of selu
   ExpectCounterNotChanged("aten::.*", cpp_test::GetIgnoredCounters());
-  ExpectCounterChanged("xla::selu_", cpp_test::GetIgnoredCounters());
 }
 
 TEST_F(AtenXlaTensorTest, TestCelu) {

torch_xla/csrc/aten_xla_type.cpp

@@ -2785,32 +2785,6 @@ at::Tensor XLANativeFunctions::select(const at::Tensor& self, int64_t dim,
       XLATensor::select(bridge::GetXlaTensor(self), dim, index));
 }
 
-at::Tensor XLANativeFunctions::selu(const at::Tensor& self) {
-  XLA_FN_COUNTER("xla::");
-  return bridge::AtenFromXlaTensor(XLATensor::selu(bridge::GetXlaTensor(self)));
-}
-
-at::Tensor& XLANativeFunctions::selu_(at::Tensor& self) {
-  XLA_FN_COUNTER("xla::");
-  XLATensorPtr self_tensor = bridge::GetXlaTensor(self);
-  XLATensor::selu_(self_tensor);
-  return self;
-}
-
-at::Tensor XLANativeFunctions::silu(const at::Tensor& self) {
-  XLA_FN_COUNTER("xla::");
-  return bridge::AtenFromXlaTensor(XLATensor::silu(bridge::GetXlaTensor(self)));
-}
-
-at::Tensor XLANativeFunctions::silu_backward(const at::Tensor& grad_output,
-                                             const at::Tensor& self) {
-  XLA_FN_COUNTER("xla::");
-  XLATensorPtr grad_output_tensor = bridge::GetXlaTensor(grad_output);
-  XLATensorPtr 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/ops/ops_lower_fn.cpp

@@ -216,6 +216,11 @@ torch_xla::XlaOpVector Rsqrt::Lower(LoweringContext* loctx) const {
   return ReturnOp(xla::Rsqrt(xla_input), loctx);
 }
 
+torch_xla::XlaOpVector Selu::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  return ReturnOp(BuildSelu(xla_input), loctx);
+}
+
 torch_xla::XlaOpVector Sgn::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   return ReturnOp(BuildSgn(xla_input), loctx);
@@ -226,6 +231,17 @@ torch_xla::XlaOpVector Sign::Lower(LoweringContext* loctx) const {
   return ReturnOp(BuildSign(xla_input), loctx);
 }
 
+torch_xla::XlaOpVector Silu::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
+  return ReturnOp(xla_input * BuildSigmoid(xla_input), loctx);
+}
+
+torch_xla::XlaOpVector SiluBackward::Lower(LoweringContext* loctx) const {
+  xla::XlaOp xla_grad_output = loctx->GetOutputOp(operand(0));
+  xla::XlaOp xla_input = loctx->GetOutputOp(operand(1));
+  return ReturnOp(BuildSiLUBackward(xla_grad_output, xla_input), loctx);
+}
+
 torch_xla::XlaOpVector Sin::Lower(LoweringContext* loctx) const {
   xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
   return ReturnOp(xla::Sin(xla_input), loctx);

torch_xla/csrc/ops/ops_xla_shape_fn.cpp

@@ -2,6 +2,7 @@
 
 #include "tensorflow/compiler/xla/client/lib/logdet.h"
 #include "tensorflow/compiler/xla/shape_util.h"
+#include "torch_xla/csrc/elementwise.h"
 #include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/pooling.h"
 
@@ -208,6 +209,10 @@ xla::Shape RsqrtOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }
 
+xla::Shape SeluOutputShape(const torch::lazy::Value& input) {
+  return GetXlaShape(input);
+}
+
 xla::Shape SgnOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }
@@ -216,6 +221,20 @@ xla::Shape SignOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }
 
+xla::Shape SiluOutputShape(const torch::lazy::Value& input) {
+  return GetXlaShape(input);
+}
+
+xla::Shape SiluBackwardOutputShape(const torch::lazy::Value& grad_output,
+                                   const torch::lazy::Value& input) {
+  auto lower_for_shape_fn =
+      [](absl::Span<const xla::XlaOp> operands) -> xla::XlaOp {
+    return BuildSiLUBackward(operands[0], operands[1]);
+  };
+  return InferOutputShape({GetXlaShape(grad_output), GetXlaShape(input)},
+                          lower_for_shape_fn);
+}
+
 xla::Shape SinOutputShape(const torch::lazy::Value& input) {
   return GetXlaShape(input);
 }

torch_xla/csrc/ops/ops_xla_shape_fn.h

@@ -84,10 +84,17 @@ xla::Shape RoundOutputShape(const torch::lazy::Value& input);
 
 xla::Shape RsqrtOutputShape(const torch::lazy::Value& input);
 
+xla::Shape SeluOutputShape(const torch::lazy::Value& input);
+
 xla::Shape SgnOutputShape(const torch::lazy::Value& input);
 
 xla::Shape SignOutputShape(const torch::lazy::Value& input);
 
+xla::Shape SiluOutputShape(const torch::lazy::Value& input);
+
+xla::Shape SiluBackwardOutputShape(const torch::lazy::Value& grad_output,
+                                   const torch::lazy::Value& input);
+
 xla::Shape SinOutputShape(const torch::lazy::Value& input);
 
 xla::Shape SinhOutputShape(const torch::lazy::Value& input);

torch_xla/csrc/tensor_methods.cpp

@@ -2420,24 +2420,6 @@ XLATensorPtr XLATensor::select(const XLATensorPtr& input, int64_t dim,
   return tensor_ops::Select(input, dim, index);
 }
 
-XLATensorPtr XLATensor::selu(const XLATensorPtr& input) {
-  return input->CreateFrom(Selu(input->GetIrValue()));
-}
-
-void XLATensor::selu_(XLATensorPtr& input) {
-  input->SetInPlaceIrValue(Selu(input->GetIrValue()));
-}
-
-XLATensorPtr XLATensor::silu(const XLATensorPtr& input) {
-  return input->CreateFrom(SiLU(input->GetIrValue()));
-}
-
-XLATensorPtr XLATensor::silu_backward(XLATensorPtr& grad_output,
-                                      XLATensorPtr& input) {
-  return input->CreateFrom(
-      SiLUBackward(grad_output->GetIrValue(), input->GetIrValue()));
-}
-
 XLATensorPtr XLATensor::sigmoid(const XLATensorPtr& input) {
   return input->CreateFrom(Sigmoid(input->GetIrValue()));
 }

xla_native_functions.yaml

@@ -34,8 +34,11 @@ full_codegen:
   - reciprocal
   - round
   - rsqrt
+  - selu
   - sgn
   - sign
+  - silu
+  - silu_backward
   - sin
   - sinh
   - tan
@@ -271,12 +274,8 @@ supported:
   - scatter.value_reduce
   - scatter_add
   - select.int
-  - selu
-  - selu_
   - sigmoid
   - sigmoid_backward
-  - silu
-  - silu_backward
   - slice.Tensor
   - slogdet
   - smooth_l1_loss