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

namespace tvm{
namespace te{
int TslPlaceholderOpNode::num_outputs() const { return 1; }
DataType TslPlaceholderOpNode::output_dtype(size_t i) const {
  CHECK_EQ(i, 0U);
  return dtype;
}
Array<PrimExpr> TslPlaceholderOpNode::output_unionshape(size_t i) const {
  CHECK_EQ(i, 0U);
  return union_shape;
}
Array<PrimExpr> TslPlaceholderOpNode::output_elemshape(size_t i) const {
  CHECK_EQ(i, 0U);
  return elem_shape;
}
TVM_REGISTER_NODE_TYPE(TslPlaceholderOpNode);

TslPlaceholderOp::TslPlaceholderOp(std::string name, Array<PrimExpr> union_shape,Array<PrimExpr> elem_shape, DataType dtype){
  auto n = make_object<TslPlaceholderOpNode>();
  n->name = name;
  n->union_shape=union_shape;
  n->elem_shape=elem_shape;
  n->dtype = dtype;
  data_ = std::move(n);
}
TensorUnion tsl_placeholder(Array<PrimExpr> union_shape,Array<PrimExpr> elem_shape, DataType dtype, std::string name) {
  return TslPlaceholderOp(name, union_shape,elem_shape, dtype).output(0);
}

TVM_REGISTER_GLOBAL("te.TslPlaceholder")
.set_body_typed([](Array<PrimExpr> union_shape,Array<PrimExpr> elem_shape, DataType dtype, std::string name) {
  return tsl_placeholder(union_shape,elem_shape, dtype, name);
});

}
}