

/*!
 * \file tvm/tsl/te/tensor.h
 * \brief Dataflow tensor object
 */
#ifndef TVM_TSL_TE_TENSOR_H_
#define TVM_TSL_TE_TENSOR_H_

#include <tvm/arith/bound.h>
#include <tvm/node/container.h>
#include <tvm/te/tensor.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
#include <tvm/tsl/tir/buffer.h>

#include <string>
#include <type_traits>
#include <utility>
#include <vector>

namespace tvm {
namespace te {

class TslOperationNode;
class TensorUnion;

class TslOperation : public ObjectRef {
 public:
  TslOperation() {}
  explicit TslOperation(ObjectPtr<Object> n) : ObjectRef(n) {}

  inline const TslOperationNode* operator->() const;
  TVM_DLL TensorUnion output(size_t i) const;

  using ContainerType = TslOperationNode;
};

class TensorUnionNode : public tvm::tir::TslDataProducerNode {
 public:
  Array<PrimExpr> union_shape;
  Array<PrimExpr> elem_shape;
  DataType dtype;
  TslOperation op;
  int value_index{0};
  TensorUnionNode() {}
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("union_shape", &union_shape);
    v->Visit("elem_shape", &elem_shape);
    v->Visit("dtype", &dtype);
    v->Visit("op", &op);
    v->Visit("value_index", &value_index);
  }
  Array<PrimExpr> GetUnionShape() const { return union_shape; }
  Array<PrimExpr> GetElemShape() const { return elem_shape; }
  DataType GetDataType() const final { return dtype; }
  TVM_DLL String GetNameHint() const final;
  static constexpr const char* _type_key = "TensorUnion";
  TVM_DECLARE_FINAL_OBJECT_INFO(TensorUnionNode, TslDataProducerNode);
};

class TensorUnion : public tvm::tir::TslDataProducer {
 public:
  TVM_DLL TensorUnion(Array<PrimExpr> union_shape, Array<PrimExpr> elem_shape, DataType dtype,
                      TslOperation op, int value_index);
  inline bool operator==(const TensorUnion& other) const;

  inline bool operator!=(const TensorUnion& other) const;

  inline size_t uniondim() const;
  inline size_t elemdim() const;
  template <typename... Args>
  inline PrimExpr operator()(Args&&... args) const {
    Array<PrimExpr> indices{std::forward<Args>(args)...};
    return operator()(indices);
  }
  /*!
   * \brief Take elements from the tensor
   * \param indices the indices.
   * \return the result expression representing tensor read.
   */
  TVM_DLL PrimExpr operator()(Array<PrimExpr> indices) const;
  /*!
   * \brief Take elements from the tensor
   * \param indices the indices.
   * \return the result expression representing tensor read.
   */
  TVM_DLL PrimExpr operator()(Array<Var> indices) const;
  TVM_DEFINE_OBJECT_REF_METHODS(TensorUnion, TslDataProducer, TensorUnionNode);
};

inline size_t TensorUnion::uniondim() const { return (*this)->union_shape.size(); }
inline size_t TensorUnion::elemdim() const { return (*this)->elem_shape.size(); }

inline bool TensorUnion::operator==(const TensorUnion& other) const {
  if (get() == other.get()) return true;
  if (get() == nullptr || other.get() == nullptr) return false;
  if ((*this)->op.defined() || other->op.defined()) {
    return (*this)->op == other->op && (*this)->value_index == other->value_index;
  } else {
    return false;
  }
}
inline bool TensorUnion::operator!=(const TensorUnion& other) const { return !(*this == other); }

}  // namespace te
}  // namespace tvm

namespace std {
template <>
struct hash<::tvm::te::TslOperation> : public ::tvm::ObjectPtrHash {};

template <>
struct hash<::tvm::te::TensorUnion> {
  std::size_t operator()(const ::tvm::te::TensorUnion& k) const {
    ::tvm::ObjectPtrHash hasher;
    if (k.defined() && k->op.defined()) {
      return hasher(k->op);
    } else {
      return hasher(k);
    }
  }
};

};  // namespace std

#endif  // TVM_TSL_TE_TENSOR_H_
