

/*!
 * \file expr.cc
 */
#include <tvm/runtime/registry.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tsl/tir/expr.h>

#include <limits>
#include <memory>

// #include "../../support/str_escape.h"

namespace tvm {
namespace tir {

#define TVM_DEFINE_BINOP_CONSTRUCTOR(Name)                             \
  Name::Name(PrimExpr a, PrimExpr b) {                                 \
    using T = Name::ContainerType;                                     \
    ICHECK(a.defined()) << "ValueError: a is undefined\n";             \
    ICHECK(b.defined()) << "ValueError: b is undefined\n";             \
    ICHECK(a.dtype() == b.dtype()) << "TypeError: mismatched types\n"; \
    ObjectPtr<T> node = make_object<T>();                              \
    node->dtype = a.dtype();                                           \
    node->a = std::move(a);                                            \
    node->b = std::move(b);                                            \
    data_ = std::move(node);                                           \
  }

TULoad::TULoad(TslDataProducer producer, Array<PrimExpr> union_indices) {
  ObjectPtr<TULoadNode> node = make_object<TULoadNode>();
  node->dtype = producer->GetDataType();
  node->producer = std::move(producer);
  node->union_indices = std::move(union_indices);
  data_ = std::move(node);
}
TVM_REGISTER_GLOBAL("tir.TULoad")
    .set_body_typed([](DataProducer producer, Array<PrimExpr> union_indices) {
      return ProducerLoad(producer, union_indices);
    });
TVM_REGISTER_NODE_TYPE(TULoadNode);

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
    .set_dispatch<TULoadNode>([](const ObjectRef& node, ReprPrinter* p) {
      auto* op = static_cast<const TULoadNode*>(node.get());
      p->stream << op->producer->GetNameHint() << "[";
      for (size_t i = 0; i < op->union_indices.size(); ++i) {
        p->Print(op->union_indices[i]);
        if (i < op->union_indices.size() - 1) {
          p->stream << ", ";
        }
      }
      p->stream << "]";
    });

// TslTGemm
TVM_DEFINE_BINOP_CONSTRUCTOR(TslTGemm);

TVM_REGISTER_GLOBAL("tir.TslTGemm").set_body_typed([](PrimExpr a, PrimExpr b) {
  return TslTGemm(a, b);
});

TVM_REGISTER_NODE_TYPE(TslTGemmNode);

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
    .set_dispatch<TslTGemmNode>([](const ObjectRef& node, ReprPrinter* p) {
      auto* op = static_cast<const TslTGemmNode*>(node.get());
      p->stream << "TslTGemm(";
      p->Print(op->a);
      p->stream << ", ";
      p->Print(op->b);
      p->stream << ')';
    });

// TslTAdd
TVM_DEFINE_BINOP_CONSTRUCTOR(TslTAdd);

TVM_REGISTER_GLOBAL("tir.TslTAdd").set_body_typed([](PrimExpr a, PrimExpr b) {
  return TslTAdd(a, b);
});

TVM_REGISTER_NODE_TYPE(TslTAddNode);

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
    .set_dispatch<TslTAddNode>([](const ObjectRef& node, ReprPrinter* p) {
      auto* op = static_cast<const TslTAddNode*>(node.get());
      p->stream << "TslTAdd(";
      p->Print(op->a);
      p->stream << ", ";
      p->Print(op->b);
      p->stream << ')';
    });


// TslTWrite
TVM_DEFINE_BINOP_CONSTRUCTOR(TslTWrite);

TVM_REGISTER_GLOBAL("tir.TslTWrite").set_body_typed([](PrimExpr a, PrimExpr b) {
  return TslTWrite(a, b);
});

TVM_REGISTER_NODE_TYPE(TslTWriteNode);

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
    .set_dispatch<TslTWriteNode>([](const ObjectRef& node, ReprPrinter* p) {
      auto* op = static_cast<const TslTWriteNode*>(node.get());
      p->stream << "TslTWrite(";
      p->Print(op->a);
      p->stream << ", ";
      p->Print(op->b);
      p->stream << ')';
    });

// TslTStore
TVM_DEFINE_BINOP_CONSTRUCTOR(TslTStore);

TVM_REGISTER_GLOBAL("tir.TslTStore").set_body_typed([](PrimExpr a, PrimExpr b) {
  return TslTStore(a, b);
});

TVM_REGISTER_NODE_TYPE(TslTStoreNode);

TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
    .set_dispatch<TslTStoreNode>([](const ObjectRef& node, ReprPrinter* p) {
      auto* op = static_cast<const TslTStoreNode*>(node.get());
      p->stream << "TslTStore(";
      p->Print(op->a);
      p->stream << ", ";
      p->Print(op->b);
      p->stream << ')';
    });



}  // namespace tir
}  // namespace tvm
