

/*!
 * \file tvm/tsl/tir/expr.h
 * \brief TIR expressions.
 */
// Acknowledgement: Many low-level IR nodes originate from Halide.
#ifndef TVM_TSL_TIR_EXPR_H_
#define TVM_TSL_TIR_EXPR_H_

#include <tvm/ir/expr.h>
#include <tvm/node/container.h>
#include <tvm/node/functor.h>
#include <tvm/node/node.h>
#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/data_type.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/var.h>
#include <tvm/tsl/tir/buffer.h>

#include <algorithm>
#include <iostream>
#include <limits>
#include <string>
#include <unordered_map>
#include <utility>

namespace tvm {
namespace tir {

class TULoadNode : public PrimExprNode {
 public:
  TslDataProducer producer;
  Array<PrimExpr> union_indices;
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &(this->dtype));
    v->Visit("producer", &producer);
    v->Visit("union_indices", &union_indices);
  }
  // TODO:investigate SEQUAL/SHASH
  bool SEqualReduce(const TULoadNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(producer, other->producer) &&
           equal(union_indices, other->union_indices);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(producer);
    hash_reduce(union_indices);
  }

  static constexpr const char* _type_key = "tir.TULoad";
  TVM_DECLARE_FINAL_OBJECT_INFO(TULoadNode, PrimExprNode);
};

class TULoad : public PrimExpr {
 public:
  TVM_DLL explicit TULoad(TslDataProducer producer, Array<PrimExpr> union_indices);

  TVM_DEFINE_OBJECT_REF_METHODS(TULoad, PrimExpr, TULoadNode);
};

/* OpNode start (yuxguo)
 * PrimExprNode -> (TslUnaryOpNode, TslBinaryOpNode) -> (TslTGemmOpNode, TslTAddOpNode, ...)
 */

template <typename T>
class TslBinaryOpNode : public PrimExprNode {
 public:
  /*! \brief The left operand. */
  PrimExpr a;
  /*! \brief The right operand. */
  PrimExpr b;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &(this->dtype));
    v->Visit("a", &a);
    v->Visit("b", &b);
  }

  bool SEqualReduce(const T* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(a, other->a) && equal(b, other->b);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(a);
    hash_reduce(b);
  }
  TVM_DECLARE_FINAL_OBJECT_INFO(T, PrimExprNode);
};

template <typename T>
class TslUnaryOpNode : public PrimExprNode {
 public:
  /*! \brief The operand. */
  PrimExpr a;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &(this->dtype));
    v->Visit("a", &a);
  }

  bool SEqualReduce(const T* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(a, other->a);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(a);
  }
  TVM_DECLARE_FINAL_OBJECT_INFO(T, PrimExprNode);
};

class TslTGemmNode : public TslBinaryOpNode<TslTGemmNode> {
 public:
  static constexpr const char* _type_key = "tir.TslTGemm";
};

class TslTAddNode : public TslBinaryOpNode<TslTAddNode> {
 public:
  static constexpr const char* _type_key = "tir.TslTAdd";
};

class TslTWriteNode : public TslBinaryOpNode<TslTWriteNode> {
 public:
  static constexpr const char* _type_key = "tir.TslTWrite";
};

class TslTStoreNode : public TslBinaryOpNode<TslTStoreNode> {
 public:
  static constexpr const char* _type_key = "tir.TslTStore";
};



/* Op start (yuxguo) manage opnode
 * TslOperation (TslTensorGemmOp, TslTensorAddOp, ...)
 * usage: using TslComputeOp to specify compute type and expanded in scheduleOps pass
 */

class TslTGemm : public PrimExpr {
 public:
  TVM_DLL TslTGemm(PrimExpr a, PrimExpr b);
  TVM_DEFINE_OBJECT_REF_METHODS(TslTGemm, PrimExpr, TslTGemmNode);
};


class TslTAdd : public PrimExpr {
 public:
  TVM_DLL TslTAdd(PrimExpr a, PrimExpr b);
  TVM_DEFINE_OBJECT_REF_METHODS(TslTAdd, PrimExpr, TslTAddNode);
};

class TslTWrite : public PrimExpr {
 public:
  TVM_DLL TslTWrite(PrimExpr a, PrimExpr b);
  TVM_DEFINE_OBJECT_REF_METHODS(TslTWrite, PrimExpr, TslTWriteNode);
};

class TslTStore : public PrimExpr {
 public:
  TVM_DLL TslTStore(PrimExpr a, PrimExpr b);
  TVM_DEFINE_OBJECT_REF_METHODS(TslTStore, PrimExpr, TslTStoreNode);
};

}  // namespace tir
}  // namespace tvm

#endif  // TVM_TSL_TIR_EXPR_H_
