Codegen HardshrinkBackward and SoftshrinkBackward

[vanbasten23]

LazyIr.h:

class HardshrinkBackward : public XlaNode {
 public:
  static torch::lazy::OpKind ClassOpKind() {
    return torch::lazy::OpKind(at::aten::hardshrink_backward);
  }

  HardshrinkBackward(const torch::lazy::Value& grad_out,
                     const torch::lazy::Value& self,
                     const torch::lazy::Value& lambd,
                     std::vector<torch::lazy::Shape>&& shapes)
      : XlaNode(torch::lazy::OpKind(at::aten::hardshrink_backward),
                {grad_out, self, lambd}, std::move(shapes),
                [&]() {
                  return HardshrinkBackwardOutputShape(grad_out, self, lambd);
                },
                /* num_outputs */ 1, torch::lazy::MHash()) {}

  std::string ToString() const override {
    std::stringstream ss;
    ss << XlaNode::ToString();

    return ss.str();
  }

  bool CanBeReused(const torch::lazy::Value& grad_out,
                   const torch::lazy::Value& self,
                   const torch::lazy::Value& lambd) const {
    return false;
  }

  torch_xla::XlaOpVector Lower(LoweringContext* loctx) const override;
};

class SoftshrinkBackward : public XlaNode {
 public:
  static torch::lazy::OpKind ClassOpKind() {
    return torch::lazy::OpKind(at::aten::softshrink_backward);
  }

  SoftshrinkBackward(const torch::lazy::Value& grad_output,
                     const torch::lazy::Value& self,
                     const torch::lazy::Value& lambd,
                     std::vector<torch::lazy::Shape>&& shapes)
      : XlaNode(torch::lazy::OpKind(at::aten::softshrink_backward),
                {grad_output, self, lambd}, std::move(shapes),
                [&]() {
                  return SoftshrinkBackwardOutputShape(grad_output, self,
                                                       lambd);
                },
                /* num_outputs */ 1, torch::lazy::MHash()) {}

  std::string ToString() const override {
    std::stringstream ss;
    ss << XlaNode::ToString();

    return ss.str();
  }

  bool CanBeReused(const torch::lazy::Value& grad_output,
                   const torch::lazy::Value& self,
                   const torch::lazy::Value& lambd) const {
    return false;
  }

  torch_xla::XlaOpVector Lower(LoweringContext* loctx) const override;
};

XlaNativeFunctions.cpp:

at::Tensor XLANativeFunctions::hardshrink_backward(const at::Tensor& grad_out,
                                                   const at::Tensor& self,
                                                   const at::Scalar& lambd) {
  XLA_FN_COUNTER("xla::");
  auto common_device = torch_xla::bridge::GetXlaDevice(grad_out, self);
  TORCH_INTERNAL_ASSERT(common_device);

  torch_xla::XLATensorPtr lazy_grad_out =
      torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(grad_out,
                                                              *common_device);
  torch_xla::XLATensorPtr lazy_self =
      torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(self,
                                                              *common_device);
  auto node_lambd =
      torch::lazy::LazyGraphExecutor::Get()->GetIrValueForScalarFromCodegen(
          lambd, *common_device);
  torch::lazy::NodePtr node = torch::lazy::ReuseNode<HardshrinkBackward>(
      lazy_grad_out->GetIrValue(), lazy_self->GetIrValue(), node_lambd);
  if (!node) {
    auto grad_out_meta = to_meta(grad_out);
    auto self_meta = to_meta(self);
    auto out_meta =
        at::meta::hardshrink_backward(grad_out_meta, self_meta, lambd);

    std::vector<torch::lazy::Shape> shapes{
        torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};
    TORCH_INTERNAL_ASSERT(shapes.size() == 1);
    if (torch::lazy::symbolicShapeEnabled()) {
      std::vector<torch::jit::IValue> inputs = {grad_out, self, lambd};
      const char* schema_str =
          "aten::hardshrink_backward(Tensor grad_out, Tensor self, Scalar "
          "lambd) -> Tensor";
      applySymbolicShapesOnLT(schema_str, inputs, shapes);
    }

    node = torch::lazy::MakeNode<HardshrinkBackward>(
        lazy_grad_out->GetIrValue(), lazy_self->GetIrValue(), node_lambd,
        std::move(shapes));
    CacheNode(node);
  }

  auto result = torch_xla::bridge::AtenFromXlaTensor(
      torch_xla::XLATensor::Create(std::move(node), *common_device));
  return result;
}

at::Tensor XLANativeFunctions::softshrink_backward(
    const at::Tensor& grad_output, const at::Tensor& self,
    const at::Scalar& lambd) {
  XLA_FN_COUNTER("xla::");
  auto common_device = torch_xla::bridge::GetXlaDevice(grad_output, self);
  TORCH_INTERNAL_ASSERT(common_device);

  torch_xla::XLATensorPtr lazy_grad_output =
      torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(grad_output,
                                                              *common_device);
  torch_xla::XLATensorPtr lazy_self =
      torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(self,
                                                              *common_device);
  auto node_lambd =
      torch::lazy::LazyGraphExecutor::Get()->GetIrValueForScalarFromCodegen(
          lambd, *common_device);
  torch::lazy::NodePtr node = torch::lazy::ReuseNode<SoftshrinkBackward>(
      lazy_grad_output->GetIrValue(), lazy_self->GetIrValue(), node_lambd);
  if (!node) {
    auto grad_output_meta = to_meta(grad_output);
    auto self_meta = to_meta(self);
    auto out_meta =
        at::meta::softshrink_backward(grad_output_meta, self_meta, lambd);

    std::vector<torch::lazy::Shape> shapes{
        torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};
    TORCH_INTERNAL_ASSERT(shapes.size() == 1);
    if (torch::lazy::symbolicShapeEnabled()) {
      std::vector<torch::jit::IValue> inputs = {grad_output, self, lambd};
      const char* schema_str =
          "aten::softshrink_backward(Tensor grad_output, Tensor self, Scalar "
          "lambd) -> Tensor";
      applySymbolicShapesOnLT(schema_str, inputs, shapes);
    }

    node = torch::lazy::MakeNode<SoftshrinkBackward>(
        lazy_grad_output->GetIrValue(), lazy_self->GetIrValue(), node_lambd,
        std::move(shapes));
    CacheNode(node);
  }

  auto result = torch_xla::bridge::AtenFromXlaTensor(
      torch_xla::XLATensor::Create(std::move(node), *common_device));
  return result;
}

[wonjoo-wj]

@vanbasten23, any reason why this is in draft?

[wonjoo-wj]

nit: we may want to do MaybeCast similar to 4ccfc24

[vanbasten23]

done

[vanbasten23]

@vanbasten23, any reason why this is in draft?

I realized I haven't done the MaybeCast thing. So I'll fix it and then mark this pr as ready for review. :p

[wonjoo-wj]

Thanks! 👍