
/*!
 * \file tvm/tsl/te/operation.h
 * \brief Operation node can generate one or multiple Tensors
 */
#ifndef TVM_TSL_TE_OPERATION_H_
#define TVM_TSL_TE_OPERATION_H_

#include <tvm/arith/analyzer.h>
#include <tvm/te/operation.h>
#include <tvm/te/schedule.h>
#include <tvm/te/tensor.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
#include <tvm/tsl/te/tensor.h>

#include <string>
#include <unordered_map>
#include <vector>

namespace tvm {
/*! \brief Tensor expression language DSL. */
namespace te {

class TVM_DLL TslOperationNode : public Object {
 public:
  /*! \brief optional name of the operation */
  std::string name;
  /*! \brief optional tag of the operation */
  std::string tag;
  /*! \brief additional attributes of the operation*/
  Map<String, ObjectRef> attrs;
  // virtual destructor.
  virtual ~TslOperationNode() {}
  /*! \return number of outputs */
  virtual int num_outputs() const = 0;
  /*!
   * \return The list of iteration variable at root
   * \note root_iter_vars decides the shape of the outputs.
   */
  virtual Array<IterVar> root_iter_vars() const = 0;
  /*!
   * \brief Get data type. i-th output tensor.
   * \param i The output index.
   * \return type of i-th output.
   */
  virtual DataType output_dtype(size_t i) const = 0;

  virtual Array<PrimExpr> output_unionshape(size_t i) const = 0;
  virtual Array<PrimExpr> output_elemshape(size_t i) const = 0;
  /*!
   * \brief List all the input Tensors.
   * \return List of input tensors.
   */
  virtual Array<TensorUnion> InputTensorUnions() const = 0;

  static constexpr const char* _type_key = "TslOperation";

  TVM_DECLARE_BASE_OBJECT_INFO(TslOperationNode, Object);
};

class TslPlaceholderOpNode : public TslOperationNode {
 public:
  Array<PrimExpr> union_shape;
  Array<PrimExpr> elem_shape;
  DataType dtype;
  int num_outputs() const final;
  Array<IterVar> root_iter_vars() const final { return {}; }
  DataType output_dtype(size_t i) const final;
  Array<PrimExpr> output_unionshape(size_t i) const final;
  Array<PrimExpr> output_elemshape(size_t i) const final;
  Array<TensorUnion> InputTensorUnions() const final { return {}; }
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("name", &name);
    v->Visit("tag", &tag);
    v->Visit("attrs", &attrs);
    v->Visit("union_shape", &union_shape);
    v->Visit("elem_shape", &elem_shape);
    v->Visit("dtype", &dtype);
  }
  static constexpr const char* _type_key = "TslPlaceholderOp";
  TVM_DECLARE_FINAL_OBJECT_INFO(TslPlaceholderOpNode, TslOperationNode);
};

class TslPlaceholderOp : public TslOperation {
 public:
  TVM_DLL TslPlaceholderOp(std::string name, Array<PrimExpr> union_shape,
                           Array<PrimExpr> elem_shape, DataType dtype);

  TVM_DEFINE_OBJECT_REF_METHODS(TslPlaceholderOp, TslOperation, TslPlaceholderOpNode);
};

TVM_DLL TensorUnion tsl_placeholder(Array<PrimExpr> union_shape, Array<PrimExpr> elem_shape,
                                    DataType dtype, std::string name);

inline const TslOperationNode* TslOperation::operator->() const {
  return static_cast<const TslOperationNode*>(get());
}

class TVM_DLL TslBaseComputeOpNode : public TslOperationNode {
 public:
  Array<IterVar> union_axis;
  Array<IterVar> union_reduce_axis;
  Array<IterVar> root_iter_vars() const final;  // root union itervars
  Array<PrimExpr> output_unionshape(size_t idx) const final;
  Array<PrimExpr> output_elemshape(size_t idx) const final { return {}; }  // TODO:deal with this
  static constexpr const char* _type_key = "TslBaseComputeOp";
  TVM_DECLARE_BASE_OBJECT_INFO(TslBaseComputeOpNode, TslOperationNode);
};

class TVM_DLL TslComputeOpNode : public TslBaseComputeOpNode {
 public:
  Array<PrimExpr> body;
  TslComputeOpNode() {}
  int num_outputs() const final;
  DataType output_dtype(size_t i) const final;
  Array<TensorUnion> InputTensorUnions() const final;
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("name", &name);
    v->Visit("tag", &tag);
    v->Visit("attrs", &attrs);
    v->Visit("union_axis", &union_axis);
    v->Visit("union_reduce_axis", &union_reduce_axis);
    v->Visit("body", &body);
  }
  static constexpr const char* _type_key = "TslComputeOp";
  TVM_DECLARE_FINAL_OBJECT_INFO(TslComputeOpNode, TslBaseComputeOpNode);
};

class TslComputeOp : public TslOperation {
 public:
  TVM_DLL TslComputeOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
                       Array<IterVar> union_axis, Array<PrimExpr> body);
  TVM_DEFINE_OBJECT_REF_METHODS(TslComputeOp, TslOperation, ComputeOpNode);
};

TVM_DLL TensorUnion Tslcompute(Array<PrimExpr> union_shape, FCompute fcompute,
                               std::string name = "TU", std::string tag = "",
                               Map<String, ObjectRef> attrs = {});

}  // namespace te
}  // namespace tvm
#endif  // TVM_TSL_TE_OPERATION_H_
