expr.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

// Acknowledgement: Many low-level IR nodes originate from Halide.
#ifndef TVM_TIR_EXPR_H_
#define TVM_TIR_EXPR_H_

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

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

namespace tvm {
namespace tir {

using IntImmNode = tvm::IntImmNode;
using FloatImmNode = tvm::FloatImmNode;

class StringImmNode : public PrimExprNode {
 public:
  String value;

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

  bool SEqualReduce(const StringImmNode* other, SEqualReducer equal) const {
    return equal(value, other->value);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(value); }

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

class StringImm : public PrimExpr {
 public:
  TVM_DLL StringImm(String value, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(StringImm, PrimExpr, StringImmNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(StringImmNode);
};

class CastNode : public PrimExprNode {
 public:
  PrimExpr value;

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

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

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(value);
  }

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

class Cast : public PrimExpr {
 public:
  TVM_DLL Cast(DataType dtype, PrimExpr value, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Cast, PrimExpr, CastNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(CastNode);
};

template <typename T>
class BinaryOpNode : public PrimExprNode {
 public:
  PrimExpr a;
  PrimExpr b;

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

  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);
};

class AddNode : public BinaryOpNode<AddNode> {
 public:
  static constexpr const char* _type_key = "tir.Add";
};

class Add : public PrimExpr {
 public:
  TVM_DLL Add(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Add, PrimExpr, AddNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AddNode);
};

class SubNode : public BinaryOpNode<SubNode> {
 public:
  static constexpr const char* _type_key = "tir.Sub";
};

class Sub : public PrimExpr {
 public:
  TVM_DLL Sub(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Sub, PrimExpr, SubNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SubNode);
};

class MulNode : public BinaryOpNode<MulNode> {
 public:
  static constexpr const char* _type_key = "tir.Mul";
};

class Mul : public PrimExpr {
 public:
  TVM_DLL Mul(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Mul, PrimExpr, MulNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MulNode);
};

class DivNode : public BinaryOpNode<DivNode> {
 public:
  static constexpr const char* _type_key = "tir.Div";
};

class Div : public PrimExpr {
 public:
  TVM_DLL Div(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Div, PrimExpr, DivNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DivNode);
};

class ModNode : public BinaryOpNode<ModNode> {
 public:
  static constexpr const char* _type_key = "tir.Mod";
};

class Mod : public PrimExpr {
 public:
  TVM_DLL Mod(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Mod, PrimExpr, ModNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ModNode);
};

class FloorDivNode : public BinaryOpNode<FloorDivNode> {
 public:
  static constexpr const char* _type_key = "tir.FloorDiv";
};

class FloorDiv : public PrimExpr {
 public:
  TVM_DLL FloorDiv(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(FloorDiv, PrimExpr, FloorDivNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FloorDivNode);
};

class FloorModNode : public BinaryOpNode<FloorModNode> {
 public:
  static constexpr const char* _type_key = "tir.FloorMod";
};

class FloorMod : public PrimExpr {
 public:
  TVM_DLL FloorMod(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(FloorMod, PrimExpr, FloorModNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FloorModNode);
};

class MinNode : public BinaryOpNode<MinNode> {
 public:
  static constexpr const char* _type_key = "tir.Min";
};

class Min : public PrimExpr {
 public:
  TVM_DLL Min(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Min, PrimExpr, MinNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MinNode);
};

class MaxNode : public BinaryOpNode<MaxNode> {
 public:
  static constexpr const char* _type_key = "tir.Max";
};

class Max : public PrimExpr {
 public:
  TVM_DLL Max(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Max, PrimExpr, MaxNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MaxNode);
};

template <typename T>
class CmpOpNode : public PrimExprNode {
 public:
  PrimExpr a;
  PrimExpr b;

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

  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);
};

class EQNode : public CmpOpNode<EQNode> {
 public:
  static constexpr const char* _type_key = "tir.EQ";
};

class EQ : public PrimExpr {
 public:
  TVM_DLL EQ(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(EQ, PrimExpr, EQNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(EQNode);
};

class NENode : public CmpOpNode<NENode> {
 public:
  static constexpr const char* _type_key = "tir.NE";
};

class NE : public PrimExpr {
 public:
  TVM_DLL NE(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(NE, PrimExpr, NENode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(NENode);
};

class LTNode : public CmpOpNode<LTNode> {
 public:
  static constexpr const char* _type_key = "tir.LT";
};

class LT : public PrimExpr {
 public:
  TVM_DLL LT(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(LT, PrimExpr, LTNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LTNode);
};

struct LENode : public CmpOpNode<LENode> {
 public:
  static constexpr const char* _type_key = "tir.LE";
};

class LE : public PrimExpr {
 public:
  TVM_DLL LE(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(LE, PrimExpr, LENode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LENode);
};

class GTNode : public CmpOpNode<GTNode> {
 public:
  static constexpr const char* _type_key = "tir.GT";
};

class GT : public PrimExpr {
 public:
  TVM_DLL GT(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(GT, PrimExpr, GTNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(GTNode);
};

class GENode : public CmpOpNode<GENode> {
 public:
  static constexpr const char* _type_key = "tir.GE";
};

class GE : public PrimExpr {
 public:
  TVM_DLL GE(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(GE, PrimExpr, GENode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(GENode);
};

class AndNode : public PrimExprNode {
 public:
  PrimExpr a;
  PrimExpr b;

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

  bool SEqualReduce(const AndNode* 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);
  }

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

class And : public PrimExpr {
 public:
  TVM_DLL And(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(And, PrimExpr, AndNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AndNode);
};

class OrNode : public PrimExprNode {
 public:
  PrimExpr a;
  PrimExpr b;

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

  bool SEqualReduce(const OrNode* 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);
  }

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

class Or : public PrimExpr {
 public:
  TVM_DLL Or(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Or, PrimExpr, OrNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(OrNode);
};

class NotNode : public PrimExprNode {
 public:
  PrimExpr a;

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

  bool SEqualReduce(const NotNode* 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);
  }

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

class Not : public PrimExpr {
 public:
  TVM_DLL Not(PrimExpr a, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Not, PrimExpr, NotNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(NotNode);
};

class SelectNode : public PrimExprNode {
 public:
  PrimExpr condition;
  PrimExpr true_value;
  PrimExpr false_value;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("condition", &condition);
    v->Visit("true_value", &true_value);
    v->Visit("false_value", &false_value);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const SelectNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(condition, other->condition) &&
           equal(true_value, other->true_value) && equal(false_value, other->false_value);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(condition);
    hash_reduce(true_value);
    hash_reduce(false_value);
  }

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

class Select : public PrimExpr {
 public:
  TVM_DLL Select(PrimExpr condition, PrimExpr true_value, PrimExpr false_value, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Select, PrimExpr, SelectNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SelectNode);
};

class BufferLoadNode : public PrimExprNode {
 public:
  Buffer buffer;
  Array<PrimExpr> indices;

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

  bool SEqualReduce(const BufferLoadNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(buffer, other->buffer) &&
           equal(indices, other->indices);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(buffer);
    hash_reduce(indices);
  }

  static constexpr const char* _type_key = "tir.BufferLoad";
  TVM_DECLARE_FINAL_OBJECT_INFO(BufferLoadNode, PrimExprNode);

 private:
  void LegalizeDType();
  friend class BufferLoad;
  friend class CustomDatatypesLowerer;
  friend class VectorTypeRewriter;
  friend class Vectorizer;
};

class BufferLoad : public PrimExpr {
 public:
  TVM_DLL explicit BufferLoad(Buffer buffer, Array<PrimExpr> indices, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(BufferLoad, PrimExpr, BufferLoadNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferLoadNode);
};

class ProducerLoadNode : public PrimExprNode {
 public:
  DataProducer producer;
  Array<PrimExpr> indices;

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

  bool SEqualReduce(const ProducerLoadNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(producer, other->producer) &&
           equal(indices, other->indices);
  }

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

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

class ProducerLoad : public PrimExpr {
 public:
  TVM_DLL explicit ProducerLoad(DataProducer producer, Array<PrimExpr> indices, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(ProducerLoad, PrimExpr, ProducerLoadNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ProducerLoadNode);
};

class LoadNode : public PrimExprNode {
 public:
  Var buffer_var;
  PrimExpr index;
  PrimExpr predicate;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("buffer_var", &buffer_var);
    v->Visit("index", &index);
    v->Visit("predicate", &predicate);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const LoadNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(buffer_var, other->buffer_var) &&
           equal(index, other->index) && equal(predicate, other->predicate);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(buffer_var);
    hash_reduce(index);
    hash_reduce(predicate);
  }

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

class Load : public PrimExpr {
 public:
  TVM_DLL Load(DataType dtype, Var buffer_var, PrimExpr index, PrimExpr predicate,
               Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Load, PrimExpr, LoadNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LoadNode);
};

class RampNode : public PrimExprNode {
 public:
  PrimExpr base;
  PrimExpr stride;
  int lanes;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("base", &base);
    v->Visit("stride", &stride);
    v->Visit("lanes", &lanes);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const RampNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(base, other->base) && equal(stride, other->stride) &&
           equal(lanes, other->lanes);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(base);
    hash_reduce(stride);
    hash_reduce(lanes);
  }

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

class Ramp : public PrimExpr {
 public:
  TVM_DLL Ramp(PrimExpr base, PrimExpr stride, int lanes, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Ramp, PrimExpr, RampNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(RampNode);
};

class BroadcastNode : public PrimExprNode {
 public:
  PrimExpr value;
  int lanes;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("value", &value);
    v->Visit("lanes", &lanes);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const BroadcastNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(value, other->value) && equal(lanes, other->lanes);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(value);
    hash_reduce(lanes);
  }

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

class Broadcast : public PrimExpr {
 public:
  TVM_DLL Broadcast(PrimExpr value, int lanes, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Broadcast, PrimExpr, BroadcastNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BroadcastNode);
};

class LetNode : public PrimExprNode {
 public:
  Var var;
  PrimExpr value;
  PrimExpr body;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("var", &var);
    v->Visit("value", &value);
    v->Visit("body", &body);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const LetNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal.DefEqual(var, other->var) &&
           equal(value, other->value) && equal(body, other->body);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce.DefHash(var);
    hash_reduce(value);
    hash_reduce(body);
  }

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

class Let : public PrimExpr {
 public:
  TVM_DLL Let(Var var, PrimExpr value, PrimExpr body, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Let, PrimExpr, LetNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LetNode);
};

class CallNode : public PrimExprNode {
 public:
  RelayExpr op;

  Array<PrimExpr> args;
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("op", &op);
    v->Visit("args", &args);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const CallNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(op, other->op) && equal(args, other->args);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(op);
    hash_reduce(args);
  }

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

class Call : public PrimExpr {
 public:
  TVM_DLL Call(DataType dtype, RelayExpr op, Array<PrimExpr> args, Span span = Span());
  TVM_DEFINE_OBJECT_REF_METHODS(Call, PrimExpr, CallNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(CallNode);
};

class ShuffleNode : public PrimExprNode {
 public:
  Array<PrimExpr> vectors;
  Array<PrimExpr> indices;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("vectors", &vectors);
    v->Visit("indices", &indices);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const ShuffleNode* other, SEqualReducer equal) const {
    return equal(dtype, other->dtype) && equal(vectors, other->vectors) &&
           equal(indices, other->indices);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(vectors);
    hash_reduce(indices);
  }

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

class Shuffle : public PrimExpr {
 public:
  TVM_DLL Shuffle(Array<PrimExpr> vectors, Array<PrimExpr> indices, Span span = Span());
  TVM_DLL static PrimExpr Concat(Array<PrimExpr> vectors, Span span = Span());
  TVM_DLL static PrimExpr ExtractElement(PrimExpr vector, int index, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Shuffle, PrimExpr, ShuffleNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ShuffleNode);
};

// Reduce operator
class CommReducerNode : public Object {
 public:
  Array<Var> lhs;
  Array<Var> rhs;
  Array<PrimExpr> result;
  Array<PrimExpr> identity_element;
  Array<PrimExpr> operator()(Array<PrimExpr> a, Array<PrimExpr> b) const;
  mutable Span span;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("lhs", &lhs);
    v->Visit("rhs", &rhs);
    v->Visit("result", &result);
    v->Visit("identity_element", &identity_element);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const CommReducerNode* other, SEqualReducer equal) const {
    return equal.DefEqual(lhs, other->lhs) && equal.DefEqual(rhs, other->rhs) &&
           equal(result, other->result) && equal(identity_element, other->identity_element);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(lhs);
    hash_reduce.DefHash(rhs);
    hash_reduce(result);
    hash_reduce(identity_element);
  }

  static constexpr const char* _type_key = "tir.CommReducer";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(CommReducerNode, Object);
};

class CommReducer : public ObjectRef {
 public:
  TVM_DLL CommReducer(Array<Var> lhs, Array<Var> rhs, Array<PrimExpr> result,
                      Array<PrimExpr> identity_element, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(CommReducer, ObjectRef, CommReducerNode);
};

class ReduceNode : public PrimExprNode {
 public:
  CommReducer combiner;
  Array<PrimExpr> source;
  Array<PrimExpr> init;
  Array<IterVar> axis;
  PrimExpr condition;
  int value_index;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("combiner", &combiner);
    v->Visit("source", &source);
    v->Visit("init", &init);
    v->Visit("axis", &axis);
    v->Visit("condition", &condition);
    v->Visit("value_index", &value_index);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const ReduceNode* other, SEqualReducer equal) const {
    // check axis first so IterVars can define the necessary variables.
    return equal(dtype, other->dtype) && equal(axis, other->axis) &&
           equal(combiner, other->combiner) && equal(source, other->source) &&
           equal(init, other->init) && equal(condition, other->condition) &&
           equal(value_index, other->value_index);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(dtype);
    hash_reduce(axis);
    hash_reduce(combiner);
    hash_reduce(source);
    hash_reduce(init);
    hash_reduce(condition);
    hash_reduce(value_index);
  }

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

class Reduce : public PrimExpr {
 public:
  TVM_DLL Reduce(CommReducer combiner, Array<PrimExpr> src, Array<IterVar> rdom, PrimExpr condition,
                 int value_index, Array<PrimExpr> init, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Reduce, PrimExpr, ReduceNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ReduceNode);
};

class AnyNode : public PrimExprNode {
 public:
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("dtype", &dtype);
    v->Visit("span", &span);
  }

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

  void SHashReduce(SHashReducer hash_reduce) const {}

  Var ToVar() const { return Var("any_dim", DataType::Int(32)); }

  SizeVar ToSizeVar() const { return SizeVar("any_dim", DataType::Int(32)); }

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

class Any : public PrimExpr {
 public:
  TVM_DLL Any(Span span = Span());

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Any, PrimExpr, AnyNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AnyNode);
};

/*
 * \brief Template function to convert Map to unordered_map
 *  Sometimes useful for API gluing when internal uses unordered_map
 * \param dmap The container map
 * \return The corresponding unordered_map.
 * \tparam K the key of the Map.
 * \tparam V the value of the Map.
 */
template <typename K, typename V>
inline std::unordered_map<K, V> as_unordered_map(const Map<K, V>& dmap) {
  std::unordered_map<K, V> ret;
  for (auto kv : dmap) {
    ret[kv.first] = kv.second;
  }
  return ret;
}
}  // namespace tir
}  // namespace tvm

namespace std {
template <>
struct hash<::tvm::tir::IterVar> : public ::tvm::ObjectPtrHash {};
}  // namespace std
#endif  // TVM_TIR_EXPR_H_