Skip to content

Full codegen for sgn, sign #3577

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
May 17, 2022
Merged

Full codegen for sgn, sign #3577

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
77ec0fe
Full codegen for sgn and sign
miladm May 17, 2022
9339748
cleanup
miladm May 17, 2022
6894ff2
linter
miladm May 17, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 0 additions & 10 deletions torch_xla/csrc/aten_xla_type.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2949,16 +2949,6 @@ at::Tensor XLANativeFunctions::sigmoid_backward(const at::Tensor& grad_output,
bridge::GetXlaTensor(grad_output), bridge::GetXlaTensor(output)));
}

at::Tensor XLANativeFunctions::sgn(const at::Tensor& self) {
XLA_FN_COUNTER("xla::");
return bridge::AtenFromXlaTensor(XLATensor::sgn(bridge::GetXlaTensor(self)));
}

at::Tensor XLANativeFunctions::sign(const at::Tensor& self) {
XLA_FN_COUNTER("xla::");
return bridge::AtenFromXlaTensor(XLATensor::sign(bridge::GetXlaTensor(self)));
}

at::Tensor XLANativeFunctions::sin(const at::Tensor& self) {
XLA_FN_COUNTER("xla::");
return bridge::AtenFromXlaTensor(XLATensor::sin(bridge::GetXlaTensor(self)));
Expand Down
31 changes: 8 additions & 23 deletions torch_xla/csrc/ops/ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ PTXLA_BINARY_OP(Atan2, at::aten::atan2, xla::Atan2);
torch::lazy::NodePtr Trunc(const XlaValue& input) {
std::vector<torch::lazy::Shape> shapes;
return Floor(torch::lazy::MakeNode<Abs>(input, std::move(shapes))) *
SignOp(input);
torch::lazy::MakeNode<Sign>(input, std::vector<torch::lazy::Shape>());
}

torch::lazy::NodePtr FracOp(const XlaValue& input) {
Expand Down Expand Up @@ -129,26 +129,6 @@ torch::lazy::NodePtr ReciprocalOp(const XlaValue& input) {
input.xla_shape(), std::move(lower_fn));
}

torch::lazy::NodePtr SgnOp(const XlaValue& input) {
auto lower_fn = [](const XlaNode& node,
LoweringContext* loctx) -> XlaOpVector {
xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
return node.ReturnOp(BuildSgn(xla_input), loctx);
};
return GenericOp(torch::lazy::OpKind(at::aten::sgn), {input},
input.xla_shape(), std::move(lower_fn));
}

torch::lazy::NodePtr SignOp(const XlaValue& input) {
auto lower_fn = [](const XlaNode& node,
LoweringContext* loctx) -> XlaOpVector {
xla::XlaOp xla_input = loctx->GetOutputOp(node.operand(0));
return node.ReturnOp(BuildSign(xla_input), loctx);
};
return GenericOp(torch::lazy::OpKind(at::aten::sign), {input},
input.xla_shape(), std::move(lower_fn));
}

torch::lazy::NodePtr ReluOp(const XlaValue& input) {
auto lower_fn = [](const XlaNode& node,
LoweringContext* loctx) -> XlaOpVector {
Expand Down Expand Up @@ -783,8 +763,13 @@ torch::lazy::NodePtr Remainder(const XlaValue& input, const XlaValue& divisor) {
torch::lazy::NodePtr f = Fmod(
input,
torch::lazy::MakeNode<Abs>(divisor, std::vector<torch::lazy::Shape>()));
return f + divisor * ComparisonOp(at::aten::lt, SignOp(f) * SignOp(divisor),
ScalarOp(0, input.xla_shape()));
return f + divisor * ComparisonOp(
at::aten::lt,
torch::lazy::MakeNode<Sign>(
f, std::vector<torch::lazy::Shape>()) *
torch::lazy::MakeNode<Sign>(
divisor, std::vector<torch::lazy::Shape>()),
Copy link
Collaborator Author

@miladm miladm May 17, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This kind of call pattern starts to get very messy. I suggest we think about a cleaner solution. Thinking...

Copy link
Collaborator

@JackCaoG JackCaoG May 17, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we need to clean up all of the Node Level lowering. meaning a Node should not use another Node to do the lowering directly. The solution is to wrote a lower function for all these ops and in this case call BuildSign.

ScalarOp(0, input.xla_shape()));
}

torch::lazy::NodePtr MaxUnary(const XlaValue& input) {
Expand Down
4 changes: 0 additions & 4 deletions torch_xla/csrc/ops/ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,10 +71,6 @@ torch::lazy::NodePtr Tanh(const XlaValue& input);

torch::lazy::NodePtr Neg(const XlaValue& input);

torch::lazy::NodePtr SgnOp(const XlaValue& input);

torch::lazy::NodePtr SignOp(const XlaValue& input);

torch::lazy::NodePtr ReluOp(const XlaValue& input);

torch::lazy::NodePtr Min(const XlaValue& input, const XlaValue& other);
Expand Down
10 changes: 10 additions & 0 deletions torch_xla/csrc/ops/ops_lower_fn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,4 +48,14 @@ torch_xla::XlaOpVector Maximum::Lower(LoweringContext* loctx) const {
return ReturnOp(xla::Max(promoted.first, promoted.second), loctx);
}

torch_xla::XlaOpVector Sgn::Lower(LoweringContext* loctx) const {
xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
return ReturnOp(BuildSgn(xla_input), loctx);
}

torch_xla::XlaOpVector Sign::Lower(LoweringContext* loctx) const {
xla::XlaOp xla_input = loctx->GetOutputOp(operand(0));
return ReturnOp(BuildSign(xla_input), loctx);
}

} // namespace torch_xla
4 changes: 4 additions & 0 deletions torch_xla/csrc/ops/ops_xla_shape_fn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,4 +28,8 @@ xla::Shape MaximumOutputShape(const XlaValue& input, const XlaValue& other) {
lower_for_shape_fn);
}

xla::Shape SgnOutputShape(const XlaValue& input) { return input.xla_shape(); }

xla::Shape SignOutputShape(const XlaValue& input) { return input.xla_shape(); }

} // namespace torch_xla
4 changes: 4 additions & 0 deletions torch_xla/csrc/ops/ops_xla_shape_fn.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,4 +19,8 @@ xla::Shape AtanhOutputShape(const XlaValue& input);

xla::Shape MaximumOutputShape(const XlaValue& input, const XlaValue& other);

xla::Shape SgnOutputShape(const XlaValue& input);

xla::Shape SignOutputShape(const XlaValue& input);

} // namespace torch_xla
4 changes: 0 additions & 4 deletions torch_xla/csrc/tensor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1033,10 +1033,6 @@ class XLATensor : public c10::intrusive_ptr_target {
static XLATensor sigmoid_backward(const XLATensor& grad_output,
const XLATensor& output);

static XLATensor sgn(const XLATensor& input);

static XLATensor sign(const XLATensor& input);

static XLATensor sin(const XLATensor& input);

static XLATensor sinh(const XLATensor& input);
Expand Down
8 changes: 0 additions & 8 deletions torch_xla/csrc/tensor_methods.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2511,14 +2511,6 @@ XLATensor XLATensor::sigmoid_backward(const XLATensor& grad_output,
SigmoidBackward(grad_output.GetIrValue(), output.GetIrValue()));
}

XLATensor XLATensor::sgn(const XLATensor& input) {
return input.CreateFrom(SgnOp(input.GetIrValue()));
}

XLATensor XLATensor::sign(const XLATensor& input) {
return input.CreateFrom(SignOp(input.GetIrValue()));
}

XLATensor XLATensor::sin(const XLATensor& input) {
return input.CreateFrom(Sin(input.GetIrValue()));
}
Expand Down
4 changes: 2 additions & 2 deletions xla_native_functions.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ full_codegen:
- atan
- atanh
- maximum
- sgn
- sign
supported:
- __ilshift__.Scalar
- __ilshift__.Tensor
Expand Down Expand Up @@ -264,10 +266,8 @@ supported:
- select.int
- selu
- selu_
- sgn
- sigmoid
- sigmoid_backward
- sign
- silu.out
- silu_backward
- sin
Expand Down