Codegen amin amax

[JackCaoG]

fix

example of codegen ops that takes bool and vector<int64>

currently C++ test failed with

354] MLIR V1 optimization pass is not enabled
2022-07-26 05:12:14.133348: I 1334741 tensorflow/compiler/jit/xla_device.cc:429] XLA_GPU and XLA_CPU devices are deprecated and will be removed in subsequent releases. Instead, use either @tf.function(jit_compile=True) for must-compile semantics, or run with TF_XLA_FLAGS=--tf_xla_auto_jit=2 for auto-clustering best-effort compilation.
unknown file: Failure
C++ exception with description "`InlinedVector::at(size_type) const` failed bounds check" thrown in the test body.

XLANativeFunction.cpp

at::Tensor XLANativeFunctions::amax(const at::Tensor& self, at::IntArrayRef dim,
                                    bool keepdim) {
  XLA_FN_COUNTER("xla::");
  auto common_device = torch_xla::bridge::GetXlaDevice(self);
  TORCH_INTERNAL_ASSERT(common_device);

  torch_xla::XLATensorPtr lazy_self =
      torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(self,
                                                              *common_device);
  torch::lazy::NodePtr node = torch::lazy::ReuseNode<Amax>(
      lazy_self->GetIrValue(), std::vector<int64_t>(dim.begin(), dim.end()),
      keepdim);
  if (!node) {
    auto self_meta = to_meta(self);
    auto out_meta = at::meta::amax(self_meta, dim, keepdim);

    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 = {self, dim, keepdim};
      const char* schema_str =
          "aten::amax(Tensor self, int[1] dim=[], bool keepdim=False) -> "
          "Tensor";
      applySymbolicShapesOnLT(schema_str, inputs, shapes);
    }

    node = torch::lazy::MakeNode<Amax>(
        lazy_self->GetIrValue(), std::vector<int64_t>(dim.begin(), dim.end()),
        keepdim, 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::amin(const at::Tensor& self, at::IntArrayRef dim,
                                    bool keepdim) {
  XLA_FN_COUNTER("xla::");
  auto common_device = torch_xla::bridge::GetXlaDevice(self);
  TORCH_INTERNAL_ASSERT(common_device);

  torch_xla::XLATensorPtr lazy_self =
      torch_xla::bridge::GetXlaTensorOrCreateForWrappedNumber(self,
                                                              *common_device);
  torch::lazy::NodePtr node = torch::lazy::ReuseNode<Amin>(
      lazy_self->GetIrValue(), std::vector<int64_t>(dim.begin(), dim.end()),
      keepdim);
  if (!node) {
    auto self_meta = to_meta(self);
    auto out_meta = at::meta::amin(self_meta, dim, keepdim);

    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 = {self, dim, keepdim};
      const char* schema_str =
          "aten::amin(Tensor self, int[1] dim=[], bool keepdim=False) -> "
          "Tensor";
      applySymbolicShapesOnLT(schema_str, inputs, shapes);
    }

    node = torch::lazy::MakeNode<Amin>(
        lazy_self->GetIrValue(), std::vector<int64_t>(dim.begin(), dim.end()),
        keepdim, std::move(shapes));
    CacheNode(node);
  }

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

LazyIr

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

  Amax(const torch::lazy::Value& self, const ::std::vector<int64_t>& dim,
       const bool& keepdim, std::vector<torch::lazy::Shape>&& shapes)
      : XlaNode(torch::lazy::OpKind(at::aten::amax), {self}, std::move(shapes),
                [&]() { return AmaxOutputShape(self, dim, keepdim); },
                /* num_outputs */ 1, torch::lazy::MHash(dim, keepdim)),
        dim(dim),
        keepdim(keepdim) {}

  std::string ToString() const override {
    std::stringstream ss;
    ss << XlaNode::ToString();
    ss << ", dim=" << dim;
    ss << ", keepdim=" << keepdim;
    return ss.str();
  }

  bool CanBeReused(const torch::lazy::Value& self,
                   const ::std::vector<int64_t>& dim,
                   const bool& keepdim) const {
    return false;
  }

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

  ::std::vector<int64_t> dim;
  bool keepdim;
};

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

  Amin(const torch::lazy::Value& self, const ::std::vector<int64_t>& dim,
       const bool& keepdim, std::vector<torch::lazy::Shape>&& shapes)
      : XlaNode(torch::lazy::OpKind(at::aten::amin), {self}, std::move(shapes),
                [&]() { return AminOutputShape(self, dim, keepdim); },
                /* num_outputs */ 1, torch::lazy::MHash(dim, keepdim)),
        dim(dim),
        keepdim(keepdim) {}

  std::string ToString() const override {
    std::stringstream ss;
    ss << XlaNode::ToString();
    ss << ", dim=" << dim;
    ss << ", keepdim=" << keepdim;
    return ss.str();
  }

  bool CanBeReused(const torch::lazy::Value& self,
                   const ::std::vector<int64_t>& dim,
                   const bool& keepdim) const {
    return false;
  }

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

  ::std::vector<int64_t> dim;
  bool keepdim;
};

[JackCaoG]

Oh I figured out what happened, it seems like codegen ops can't handle negative input, for example

t1 =torch.randn(4,3,4, device='xla:0')
t2 = torch.amax(t1, -2, True)

failed, but index -2 should be equivalent to index 1 here.

[JackCaoG]

This is blocked by pytorch/pytorch#82286

[JackCaoG]

@wonjoolee95 This one if ready for reviewing

[wonjoolee95]

Nice! 🚀