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/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/dtype.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/node_functor.h>
#include <tvm/runtime/base.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/var.h>
 
#include <algorithm>
#include <iostream>
#include <limits>
#include <string>
#include <unordered_map>
#include <utility>
 
namespace tvm {
namespace tirx {
 
using IntImmNode = tvm::IntImmNode;
using FloatImmNode = tvm::FloatImmNode;
 
class StringImmNode : public ExprNode {
 public:
  ffi::String value;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<StringImmNode>().def_ro("value", &StringImmNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.StringImm", StringImmNode, ExprNode);
};
 
class StringImm : public PrimExpr {
 public:
  TVM_DLL StringImm(ffi::String value, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StringImm, PrimExpr, StringImmNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(StringImmNode);
};
 
class CastNode : public ExprNode {
 public:
  PrimExpr value;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<CastNode>().def_ro("value", &CastNode::value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Cast", CastNode, ExprNode);
};
 
class Cast : public PrimExpr {
 public:
  TVM_DLL Cast(PrimType value_ty, PrimExpr value, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Cast, PrimExpr, CastNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(CastNode);
};
 
template <typename T>
class BinaryOpNode : public ExprNode {
 public:
  PrimExpr a;
  PrimExpr b;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<T>().def_ro("a", &T::a).def_ro("b", &T::b);
  }
  static const constexpr int _type_child_slots [[maybe_unused]] = 0;
  static const constexpr bool _type_final [[maybe_unused]] = true;
  TVM_FFI_DECLARE_OBJECT_INFO_PREDEFINED_TYPE_KEY(T, ExprNode);
};
 
class AddNode : public BinaryOpNode<AddNode> {
 public:
  static constexpr const char* _type_key = "tirx.Add";
};
 
class Add : public PrimExpr {
 public:
  TVM_DLL Add(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Add, PrimExpr, AddNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AddNode);
};
 
class SubNode : public BinaryOpNode<SubNode> {
 public:
  static constexpr const char* _type_key = "tirx.Sub";
};
 
class Sub : public PrimExpr {
 public:
  TVM_DLL Sub(PrimExpr a, PrimExpr b, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Sub, PrimExpr, SubNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SubNode);
};
 
class MulNode : public BinaryOpNode<MulNode> {
 public:
  static constexpr const char* _type_key = "tirx.Mul";
};
 
class Mul : public PrimExpr {
 public:
  TVM_DLL Mul(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Mul, PrimExpr, MulNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MulNode);
};
 
class DivNode : public BinaryOpNode<DivNode> {
 public:
  static constexpr const char* _type_key = "tirx.Div";
};
 
class Div : public PrimExpr {
 public:
  TVM_DLL Div(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Div, PrimExpr, DivNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(DivNode);
};
 
class ModNode : public BinaryOpNode<ModNode> {
 public:
  static constexpr const char* _type_key = "tirx.Mod";
};
 
class Mod : public PrimExpr {
 public:
  TVM_DLL Mod(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Mod, PrimExpr, ModNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ModNode);
};
 
class FloorDivNode : public BinaryOpNode<FloorDivNode> {
 public:
  static constexpr const char* _type_key = "tirx.FloorDiv";
};
 
class FloorDiv : public PrimExpr {
 public:
  TVM_DLL FloorDiv(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FloorDiv, PrimExpr, FloorDivNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FloorDivNode);
};
 
class FloorModNode : public BinaryOpNode<FloorModNode> {
 public:
  static constexpr const char* _type_key = "tirx.FloorMod";
};
 
class FloorMod : public PrimExpr {
 public:
  TVM_DLL FloorMod(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FloorMod, PrimExpr, FloorModNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FloorModNode);
};
 
class MinNode : public BinaryOpNode<MinNode> {
 public:
  static constexpr const char* _type_key = "tirx.Min";
};
 
class Min : public PrimExpr {
 public:
  TVM_DLL Min(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Min, PrimExpr, MinNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MinNode);
};
 
class MaxNode : public BinaryOpNode<MaxNode> {
 public:
  static constexpr const char* _type_key = "tirx.Max";
};
 
class Max : public PrimExpr {
 public:
  TVM_DLL Max(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Max, PrimExpr, MaxNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MaxNode);
};
 
template <typename T>
class CmpOpNode : public ExprNode {
 public:
  PrimExpr a;
  PrimExpr b;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<T>().def_ro("a", &T::a).def_ro("b", &T::b);
  }
  static const constexpr int _type_child_slots [[maybe_unused]] = 0;
  static const constexpr bool _type_final [[maybe_unused]] = true;
  TVM_FFI_DECLARE_OBJECT_INFO_PREDEFINED_TYPE_KEY(T, ExprNode);
};
 
class EQNode : public CmpOpNode<EQNode> {
 public:
  static constexpr const char* _type_key = "tirx.EQ";
};
 
class EQ : public PrimExpr {
 public:
  TVM_DLL EQ(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(EQ, PrimExpr, EQNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(EQNode);
};
 
class NENode : public CmpOpNode<NENode> {
 public:
  static constexpr const char* _type_key = "tirx.NE";
};
 
class NE : public PrimExpr {
 public:
  TVM_DLL NE(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(NE, PrimExpr, NENode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(NENode);
};
 
class LTNode : public CmpOpNode<LTNode> {
 public:
  static constexpr const char* _type_key = "tirx.LT";
};
 
class LT : public PrimExpr {
 public:
  TVM_DLL LT(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LT, PrimExpr, LTNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LTNode);
};
 
struct LENode : public CmpOpNode<LENode> {
 public:
  static constexpr const char* _type_key = "tirx.LE";
};
 
class LE : public PrimExpr {
 public:
  TVM_DLL LE(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LE, PrimExpr, LENode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LENode);
};
 
class GTNode : public CmpOpNode<GTNode> {
 public:
  static constexpr const char* _type_key = "tirx.GT";
};
 
class GT : public PrimExpr {
 public:
  TVM_DLL GT(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GT, PrimExpr, GTNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(GTNode);
};
 
class GENode : public CmpOpNode<GENode> {
 public:
  static constexpr const char* _type_key = "tirx.GE";
};
 
class GE : public PrimExpr {
 public:
  TVM_DLL GE(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(GE, PrimExpr, GENode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(GENode);
};
 
class AndNode : public ExprNode {
 public:
  PrimExpr a;
  PrimExpr b;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AndNode>().def_ro("a", &AndNode::a).def_ro("b", &AndNode::b);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.And", AndNode, ExprNode);
};
 
class And : public PrimExpr {
 public:
  TVM_DLL And(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(And, PrimExpr, AndNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(AndNode);
};
 
class OrNode : public ExprNode {
 public:
  PrimExpr a;
  PrimExpr b;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OrNode>().def_ro("a", &OrNode::a).def_ro("b", &OrNode::b);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Or", OrNode, ExprNode);
};
 
class Or : public PrimExpr {
 public:
  TVM_DLL Or(PrimExpr a, PrimExpr b, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Or, PrimExpr, OrNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(OrNode);
};
 
class NotNode : public ExprNode {
 public:
  PrimExpr a;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<NotNode>().def_ro("a", &NotNode::a);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Not", NotNode, ExprNode);
};
 
class Not : public PrimExpr {
 public:
  TVM_DLL Not(PrimExpr a, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Not, PrimExpr, NotNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(NotNode);
};
 
class SelectNode : public ExprNode {
 public:
  PrimExpr condition;
  PrimExpr true_value;
  PrimExpr false_value;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SelectNode>()
        .def_ro("condition", &SelectNode::condition)
        .def_ro("true_value", &SelectNode::true_value)
        .def_ro("false_value", &SelectNode::false_value);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Select", SelectNode, ExprNode);
};
 
class Select : public PrimExpr {
 public:
  TVM_DLL Select(PrimExpr condition, PrimExpr true_value, PrimExpr false_value, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Select, PrimExpr, SelectNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(SelectNode);
};
 
class BufferLoadNode : public ExprNode {
 public:
  Buffer buffer;
  ffi::Array<PrimExpr> indices;
  ffi::Optional<PrimExpr> predicate;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BufferLoadNode>()
        .def_ro("buffer", &BufferLoadNode::buffer)
        .def_ro("indices", &BufferLoadNode::indices)
        .def_ro("predicate", &BufferLoadNode::predicate);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferLoad", BufferLoadNode, ExprNode);
 
 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, ffi::Array<PrimExpr> indices,
                              ffi::Optional<PrimExpr> predicate = std::nullopt, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferLoad, PrimExpr, BufferLoadNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferLoadNode);
};
 
class ProducerLoadNode : public ExprNode {
 public:
  DataProducer producer;
  ffi::Array<PrimExpr> indices;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ProducerLoadNode>()
        .def_ro("producer", &ProducerLoadNode::producer)
        .def_ro("indices", &ProducerLoadNode::indices);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ProducerLoad", ProducerLoadNode, ExprNode);
};
 
class ProducerLoad : public PrimExpr {
 public:
  TVM_DLL explicit ProducerLoad(DataProducer producer, ffi::Array<PrimExpr> indices,
                                Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ProducerLoad, PrimExpr, ProducerLoadNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ProducerLoadNode);
};
 
class RampNode : public ExprNode {
 public:
  PrimExpr base;
  PrimExpr stride;
  PrimExpr lanes;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<RampNode>()
        .def_ro("base", &RampNode::base)
        .def_ro("stride", &RampNode::stride)
        .def_ro("lanes", &RampNode::lanes);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Ramp", RampNode, ExprNode);
};
 
class Ramp : public PrimExpr {
 public:
  TVM_DLL Ramp(PrimExpr base, PrimExpr stride, PrimExpr lanes, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Ramp, PrimExpr, RampNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(RampNode);
};
 
class BroadcastNode : public ExprNode {
 public:
  PrimExpr value;
  PrimExpr lanes;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BroadcastNode>()
        .def_ro("value", &BroadcastNode::value)
        .def_ro("lanes", &BroadcastNode::lanes);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Broadcast", BroadcastNode, ExprNode);
};
 
class Broadcast : public PrimExpr {
 public:
  TVM_DLL Broadcast(PrimExpr value, PrimExpr lanes, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Broadcast, PrimExpr, BroadcastNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BroadcastNode);
};
 
class LetNode : public ExprNode {
 public:
  Var var;
  PrimExpr value;
  PrimExpr body;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LetNode>()
        // TODO(tqchen): use SEqHashDefNonRecursive after the next pypi tvm-ffi release
        .def_ro("var", &LetNode::var, refl::AttachFieldFlag::SEqHashDefRecursive())
        .def_ro("value", &LetNode::value)
        .def_ro("body", &LetNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Let", LetNode, ExprNode);
};
 
class Let : public PrimExpr {
 public:
  TVM_DLL Let(Var var, PrimExpr value, PrimExpr body, Span span = Span());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Let, PrimExpr, LetNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(LetNode);
};
 
class ShuffleNode : public ExprNode {
 public:
  ffi::Array<PrimExpr> vectors;
  ffi::Array<PrimExpr> indices;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ShuffleNode>()
        .def_ro("vectors", &ShuffleNode::vectors)
        .def_ro("indices", &ShuffleNode::indices);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Shuffle", ShuffleNode, ExprNode);
};
 
class Shuffle : public PrimExpr {
 public:
  TVM_DLL Shuffle(ffi::Array<PrimExpr> vectors, ffi::Array<PrimExpr> indices, Span span = Span());
  TVM_DLL static PrimExpr Concat(ffi::Array<PrimExpr> vectors, Span span = Span());
  TVM_DLL static PrimExpr ExtractElement(PrimExpr vector, int index, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Shuffle, PrimExpr, ShuffleNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ShuffleNode);
};
 
// Reduce operator
class CommReducerNode : public ffi::Object {
 public:
  ffi::Array<PrimVar> lhs;
  ffi::Array<PrimVar> rhs;
  ffi::Array<PrimExpr> result;
  ffi::Array<PrimExpr> identity_element;
  ffi::Array<PrimExpr> operator()(ffi::Array<PrimExpr> a, ffi::Array<PrimExpr> b) const;
  mutable Span span;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<CommReducerNode>()
        .def_ro("lhs", &CommReducerNode::lhs, refl::AttachFieldFlag::SEqHashDefRecursive())
        .def_ro("rhs", &CommReducerNode::rhs, refl::AttachFieldFlag::SEqHashDefRecursive())
        .def_ro("result", &CommReducerNode::result)
        .def_ro("identity_element", &CommReducerNode::identity_element)
        .def_ro("span", &CommReducerNode::span, refl::AttachFieldFlag::SEqHashIgnore());
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.CommReducer", CommReducerNode, ffi::Object);
};
 
class CommReducer : public ffi::ObjectRef {
 public:
  TVM_DLL CommReducer(ffi::Array<PrimVar> lhs, ffi::Array<PrimVar> rhs, ffi::Array<PrimExpr> result,
                      ffi::Array<PrimExpr> identity_element, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CommReducer, ffi::ObjectRef, CommReducerNode);
};
 
class ReduceNode : public ExprNode {
 public:
  CommReducer combiner;
  ffi::Array<PrimExpr> source;
  ffi::Array<PrimExpr> init;
  ffi::Array<IterVar> axis;
  PrimExpr condition;
  int value_index;
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ReduceNode>()
        .def_ro("combiner", &ReduceNode::combiner)
        .def_ro("source", &ReduceNode::source)
        .def_ro("init", &ReduceNode::init)
        .def_ro("axis", &ReduceNode::axis)
        .def_ro("condition", &ReduceNode::condition)
        .def_ro("value_index", &ReduceNode::value_index);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Reduce", ReduceNode, ExprNode);
};
 
class Reduce : public PrimExpr {
 public:
  TVM_DLL Reduce(CommReducer combiner, ffi::Array<PrimExpr> src, ffi::Array<IterVar> rdom,
                 PrimExpr condition, int value_index, ffi::Array<PrimExpr> init,
                 Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Reduce, PrimExpr, ReduceNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ReduceNode);
};
 
/*
 * \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 ffi::Map<K, V>& dmap) {
  std::unordered_map<K, V> ret;
  for (auto kv : dmap) {
    ret[kv.first] = kv.second;
  }
  return ret;
}
}  // namespace tirx
 
namespace ffi {
 
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::StringImmNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::CastNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::AddNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::SubNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::MulNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::DivNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::ModNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::FloorDivNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::FloorModNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::MinNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::MaxNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::EQNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::NENode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::LTNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::LENode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::GTNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::GENode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::AndNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::OrNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::NotNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::SelectNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::BufferLoadNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::ProducerLoadNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::RampNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::BroadcastNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::LetNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::ShuffleNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, tirx::ReduceNode> = true;
 
template <>
inline constexpr bool use_default_type_traits_v<tvm::tirx::StringImm> = false;
 
template <>
struct TypeTraits<tvm::tirx::StringImm>
    : public ObjectRefWithFallbackTraitsBase<tvm::tirx::StringImm, ffi::String> {
  TVM_FFI_INLINE static tvm::tirx::StringImm ConvertFallbackValue(ffi::String value) {
    return tvm::tirx::StringImm(value);
  }
};
}  // namespace ffi
}  // namespace tvm
 
namespace std {
template <>
struct hash<::tvm::tirx::IterVar> : public ::tvm::ffi::ObjectPtrHash {};
}  // namespace std
#endif  // TVM_TIR_EXPR_H_