//
// Created by bb on 2020/12/25.
//
#include <tvm/runtime/registry.h>
#include <tvm/te/operation.h>
#include <tvm/te/tensor.h>

#include <tvm/tsl/te/operation.h>
#include <tvm/tsl/te/tensor.h>

#include <tvm/tsl/tir/expr.h>

namespace tvm {
namespace te {

TensorUnion::TensorUnion(Array<PrimExpr> union_shape, Array<PrimExpr> elem_shape, DataType dtype,
                         TslOperation op, int value_index) {
  auto n = make_object<TensorUnionNode>();
  n->union_shape = std::move(union_shape);
  n->elem_shape = std::move(elem_shape);
  n->dtype = dtype;
  n->op = op;
  n->value_index = value_index;
  data_ = std::move(n);
}

String TensorUnionNode::GetNameHint() const {
  return op->num_outputs() == 1 ? op->name : (op->name + ".v" + std::to_string(value_index));
}

PrimExpr TensorUnion::operator()(Array<Var> union_indices) const {
  Array<PrimExpr> arr(union_indices.begin(), union_indices.end());
  return operator()(arr);
}

TensorUnion TslOperation::output(size_t i) const {
  auto node = make_object<TensorUnionNode>();
  node->op = *this;
  node->value_index = i;
  node->dtype = (*this)->output_dtype(i);
  node->union_shape = (*this)->output_unionshape(i);
  node->elem_shape = (*this)->output_elemshape(i);
  return TensorUnion(node);
}

PrimExpr TensorUnion::operator()(Array<PrimExpr> union_indices) const {
  if (uniondim() != 0) {
    CHECK_EQ(uniondim(), union_indices.size())
        << "Tensor dimension mismatch in read"
        << "ndim = " << uniondim() << ", indices.size=" << union_indices.size();
  }
  return TULoad((*this), union_indices);
}
}
}