Lower celu, celu_, selu, and selu_

[wonjoolee95]

Lower celu, celu_, selu, and selu_.

Trying to fix the newly introduced test failures due to #3539. This PR lowers celu, celu_, selu, and selu_ into separate lowerings.

Once this PR merges, I'll rebase the above PR and fix its tests.

[wonjoolee95]

This should be ready for a quick review, @JackCaoG. Thanks!

[JackCaoG]

mostly lgtm

[wonjoolee95]

Thanks for the comments, @JackCaoG! For some reason, if I remove celu_ and selu_, the codegen doesn't seem to take care of the in-place functions. The CI tests on the previous commit failed with celu and selu xla counters not changing when the in-place functions were called. Also I don't see the celu_ and selu_ being generated at RegisterXLA.cpp. I've added back the in-place operators for now.

[JackCaoG]

Can you check torch_xla/csrc/RegisterXLA.cpp after removing inplace?

For example for acos I see

at::Tensor wrapper__acos(const at::Tensor & self) {
    // No device check


  // DeviceGuard omitted
  return torch_xla::XLANativeFunctions::acos(self);
}

} // anonymous namespace
at::Tensor & wrapper_out_acos_out(const at::Tensor & self, at::Tensor & out) {
  auto wrapper_out_acos_out_tmp = wrapper__acos(self);
  at::_copy_from_and_resize(wrapper_out_acos_out_tmp, out);
  return out;
}
at::Tensor & wrapper__acos_(at::Tensor & self) {
  auto wrapper__acos__tmp = wrapper__acos(self);
  at::_copy_from(wrapper__acos__tmp, self);
  return self;
} 

this allow us to implement one lowering.

[wonjoolee95]

Interestingly I don't see the in-place version getting automatically generated. Here is how the RegisterXLA.cpp looks like for celu and selu:

namespace {

at::Tensor wrapper__selu(const at::Tensor & self) {
    // No device check


  // DeviceGuard omitted
  return torch_xla::XLANativeFunctions::selu(self);
}

} // anonymous namespace
namespace {

at::Tensor wrapper__celu(const at::Tensor & self, const at::Scalar & alpha) {
    // No device check


  // DeviceGuard omitted
  return torch_xla::XLANativeFunctions::celu(self, alpha);
}

} // anonymous namespace

However, I don't see one auto-generated for celu_ or selu_.

[JackCaoG]

@bdhirsh do we know what's the prerequisite for a inplace op to be generated?

[bdhirsh]

oh, it's because celu and selu don't have out= variants. The codegen needs out= variants in order to group the different variants properly so it can generate the inplace kernels.

Manually writing the two celu_() and selu_() inplace kernel lowerings for XLA seems like a good short term solution?

Caveat: The functionalization pass in core will still correctly functionalize celu_() and selu_(), so once XLA is finally using it then you won't need to worry about manually writing them and we can delete them. I know I keep making it sound like some beautiful end state in the distant future 😛

[JackCaoG]

oh, ok. Let's keep celu_ and selu_ for now.

[wonjoolee95]

Thanks for the clarification, @bdhirsh! Just did a hard reset to the commit before I removed the in-place ops.

[JackCaoG]

Thanks!