dataflow_pattern.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.
 */
 
#ifndef TVM_RELAX_DATAFLOW_PATTERN_H_
#define TVM_RELAX_DATAFLOW_PATTERN_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/with_context.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
 
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
 
namespace tvm {
 
namespace sym {
class AnalyzerObj;
class Analyzer;
}  // namespace sym
 
namespace relax {
 
class PatternSeq;
class CallPattern;
class OrPattern;
class AndPattern;
class NotPattern;
class ShapePattern;
class TypePattern;
class DataTypePattern;
class AttrPattern;
class SameShapeConstraint;
 
TVM_DLL PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index = -1);
TVM_DLL PatternSeq operator^(const PatternSeq& lhs, const PatternSeq& rhs);
 
TVM_DLL PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index = -1);
TVM_DLL PatternSeq operator>>(const PatternSeq& lhs, const PatternSeq& rhs);
 
class DFPatternNode : public ffi::Object {
 public:
  static constexpr const uint32_t _type_child_slots = 21;
  TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFPattern", DFPatternNode, ffi::Object);
};
 
class DFPattern : public ffi::ObjectRef {
 public:
  template <typename... Args>
  CallPattern operator()(Args&&... args) const;
  TVM_DLL CallPattern operator()(const std::vector<DFPattern>& args) const;
  TVM_DLL OrPattern operator|(const DFPattern& other) const;
  TVM_DLL AndPattern operator&(const DFPattern& other) const;
  TVM_DLL NotPattern operator~() const;
  TVM_DLL AttrPattern HasAttr(const ffi::Map<ffi::String, Any>& attrs) const;
  TVM_DLL TypePattern HasType(const Type& ty) const;
  TVM_DLL DataTypePattern HasDtype(DLDataType dtype) const;
  TVM_DLL DataTypePattern HasDtype(const std::string& dtype) const;
  TVM_DLL ShapePattern HasShape(const ffi::Array<PrimExpr>& shape) const;
  TVM_DLL SameShapeConstraint HasSameShapeAs(const DFPattern& other) const;
  TVM_DLL DFPattern dup() const;
 
  TVM_DLL operator PatternSeq() const;
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DFPattern, ffi::ObjectRef, DFPatternNode);
};
 
struct PairCons {
  enum Type {
    kUsedBy,     
    kOnlyUsedBy, 
  } type = kUsedBy;
  int index = -1; 
  TVM_DLL explicit PairCons(Type t, int index = -1) : type(t), index(index) {}
 
  bool operator==(const PairCons& other) const {
    return type == other.type && index == other.index;
  }
};
 
class DFConstraintNode : public ffi::Object {
 public:
  virtual ffi::Array<DFPattern> GetDependentPatterns() const = 0;
 
  virtual std::tuple<PrimExpr, bool> AsCondition(
      std::function<ffi::Optional<Var>(const DFPatternNode*)> match_state) const = 0;
 
  static constexpr const uint32_t _type_child_slots = 1;
  TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.DFConstraint", DFConstraintNode, ffi::Object);
};
 
class DFConstraint : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DFConstraint, ffi::ObjectRef, DFConstraintNode);
};
 
class PatternSeqNode final : public ffi::Object {
 public:
  tvm::ffi::Array<DFPattern> patterns;    
  std::vector<PairCons> pair_constraints; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PatternSeqNode>().def_ro("patterns", &PatternSeqNode::patterns);
  }
  TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.PatternSeq", PatternSeqNode, ffi::Object);
};
 
class PatternSeq final : public ffi::ObjectRef {
 public:
  TVM_DLL explicit PatternSeq(DFPattern init_pattern);
  TVM_DLL explicit PatternSeq(tvm::ffi::Array<DFPattern> patterns, bool only_used_by = false);
 
  PatternSeq UsedBy(PatternSeq other, int index = -1) const;
  PatternSeq OnlyUsedBy(PatternSeq other, int index = -1) const;
 
  PatternSeq dup() const;
 
  // friend functions
  friend PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index);
  friend PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PatternSeq, ffi::ObjectRef, PatternSeqNode);
};
 
class PatternContextObj : public ffi::Object {
 public:
  enum ExternUse {
    kMay,     
    kMustNot, 
  } allow_extern_use = kMay;
 
  // src node -> <dst node, constraint type> constraints.
  // Dst nodes are kept in a vector to keep them ordered.
  std::map<DFPattern, std::vector<std::pair<DFPattern, std::vector<PairCons>>>> edge_constraints;
 
  // Underlying DFPattern nodes which the edge constraints may reference
  // Kept as a separate vector of patterns to process constraints in a fixed order.
  std::vector<DFPattern> src_ordered;
 
  // Non-edge constraints
  std::vector<DFConstraint> validation_constraints;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.PatternContext", PatternContextObj, ffi::Object);
};
 
class PatternContext : public ffi::ObjectRef {
 public:
  explicit PatternContext(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
  static constexpr bool _type_is_nullable = false;
  TVM_DLL explicit PatternContext(bool incremental = false);
 
  const PatternContextObj* operator->() const {
    TVM_FFI_ICHECK(get() != nullptr);
    return static_cast<const PatternContextObj*>(get());
  }
 
  PatternContextObj* operator->() {
    TVM_FFI_ICHECK(get() != nullptr);
    return static_cast<PatternContextObj*>(get_mutable());
  }
 
  void add_constraint(DFPattern producer, DFPattern consumer, PairCons cons) {
    auto& pairs = (*this)->edge_constraints[producer];
    auto it = std::find_if(pairs.begin(), pairs.end(),
                           [consumer](auto p) { return p.first == consumer; });
    if (it == pairs.end()) {
      pairs.emplace_back(consumer, std::vector{cons});
    } else {
      auto& vec = it->second;
      TVM_FFI_ICHECK(std::find(vec.cbegin(), vec.cend(), cons) == vec.cend())
          << "Constraint already exists";
      vec.push_back(cons);
    }
 
    auto& patterns = (*this)->src_ordered;
    if (std::find(patterns.begin(), patterns.end(), producer) == patterns.end()) {
      patterns.push_back(producer);
    }
  }
 
  void add_constraint(DFConstraint constraint) {
    (*this)->validation_constraints.push_back(constraint);
  }
 
  TVM_DLL static ffi::Optional<PatternContext> Current();
 
  TVM_DLL void EnterWithScope() const;
  TVM_DLL void ExitWithScope() const;
 
 private:
  friend class With<PatternContext>;
};
 
class ExprPatternNode : public DFPatternNode {
 public:
  explicit ExprPatternNode(Expr expr) : expr(std::move(expr)) {}
  explicit ExprPatternNode(ffi::UnsafeInit) : expr(ffi::UnsafeInit{}) {}
 
  Expr expr; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExprPatternNode>().def_ro("expr", &ExprPatternNode::expr);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExprPattern", ExprPatternNode, DFPatternNode);
};
 
class ExprPattern : public DFPattern {
 public:
  TVM_DLL explicit ExprPattern(Expr expr);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExprPattern, DFPattern, ExprPatternNode);
};
 
class VarPatternNode : public DFPatternNode {
 public:
  ffi::String name;
  const ffi::String& name_hint() const { return name; }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<VarPatternNode>().def_ro("name", &VarPatternNode::name);
  }
 
  static constexpr const uint32_t _type_child_slots = 1;
  TVM_FFI_DECLARE_OBJECT_INFO("relax.dpl.VarPattern", VarPatternNode, DFPatternNode);
};
 
class VarPattern : public DFPattern {
 public:
  TVM_DLL VarPattern(ffi::String name_hint);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VarPattern, DFPattern, VarPatternNode);
};
 
class DataflowVarPatternNode : public VarPatternNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataflowVarPatternNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataflowVarPattern", DataflowVarPatternNode,
                                    VarPatternNode);
};
 
class DataflowVarPattern : public DFPattern {
 public:
  TVM_DLL DataflowVarPattern(ffi::String name_hint);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataflowVarPattern, DFPattern,
                                                DataflowVarPatternNode);
};
 
class GlobalVarPatternNode : public VarPatternNode {
 public:
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.GlobalVarPattern", GlobalVarPatternNode,
                                    DFPatternNode);
};
 
class GlobalVarPattern : public DFPattern {
 public:
  TVM_DLL GlobalVarPattern(ffi::String name_hint);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(GlobalVarPattern, DFPattern, GlobalVarPatternNode);
};
 
class ConstantPatternNode : public DFPatternNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ConstantPatternNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ConstantPattern", ConstantPatternNode,
                                    DFPatternNode);
};
 
class ConstantPattern : public DFPattern {
 public:
  explicit ConstantPattern(ffi::UnsafeInit tag, ffi::ObjectPtr<ConstantPatternNode> n)
      : DFPattern(tag) {
    data_ = std::move(n);
  }
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ConstantPattern, DFPattern, ConstantPatternNode);
};
 
class CallPatternNode : public DFPatternNode {
 public:
  explicit CallPatternNode(DFPattern op) : op(std::move(op)) {}
  explicit CallPatternNode(ffi::UnsafeInit) : op(ffi::UnsafeInit{}) {}
 
  DFPattern op;                    
  tvm::ffi::Array<DFPattern> args; 
  bool varg_default_wildcard; 
  // Todo(relax-team): Dataflow pattern for Type, and match ty_args
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<CallPatternNode>()
        .def_ro("op", &CallPatternNode::op)
        .def_ro("args", &CallPatternNode::args);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.CallPattern", CallPatternNode, DFPatternNode);
};
 
class CallPattern : public DFPattern {
 public:
  TVM_DLL CallPattern(DFPattern op, ffi::Array<DFPattern> args, bool varg_default_wildcard = false);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CallPattern, DFPattern, CallPatternNode);
};
 
class PrimArrPatternNode : public DFPatternNode {
 public:
  ffi::Array<PrimExpr> fields; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PrimArrPatternNode>().def_ro("fields", &PrimArrPatternNode::fields);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.PrimArrPattern", PrimArrPatternNode, DFPatternNode);
};
 
class PrimArrPattern : public DFPattern {
 public:
  TVM_DLL PrimArrPattern(ffi::Array<PrimExpr> arr);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimArrPattern, DFPattern, PrimArrPatternNode);
};
 
class FunctionPatternNode : public DFPatternNode {
 public:
  explicit FunctionPatternNode(DFPattern body) : body(std::move(body)) {}
  explicit FunctionPatternNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {}
 
  tvm::ffi::Array<DFPattern> params; 
  DFPattern body; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FunctionPatternNode>()
        .def_ro("params", &FunctionPatternNode::params)
        .def_ro("body", &FunctionPatternNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.FunctionPattern", FunctionPatternNode,
                                    DFPatternNode);
};
 
class FunctionPattern : public DFPattern {
 public:
  TVM_DLL FunctionPattern(tvm::ffi::Array<DFPattern> params, DFPattern body);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionPattern, DFPattern, FunctionPatternNode);
};
 
class TuplePatternNode : public DFPatternNode {
 public:
  tvm::ffi::Array<DFPattern> fields; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TuplePatternNode>().def_ro("fields", &TuplePatternNode::fields);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TuplePattern", TuplePatternNode, DFPatternNode);
};
 
class TuplePattern : public DFPattern {
 public:
  TVM_DLL explicit TuplePattern(tvm::ffi::Array<DFPattern> fields);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TuplePattern, DFPattern, TuplePatternNode);
};
 
class UnorderedTuplePatternNode : public DFPatternNode {
 public:
  tvm::ffi::Array<DFPattern> fields; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<UnorderedTuplePatternNode>().def_ro("fields",
                                                        &UnorderedTuplePatternNode::fields);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.UnorderedTuplePattern", UnorderedTuplePatternNode,
                                    DFPatternNode);
};
 
class UnorderedTuplePattern : public DFPattern {
 public:
  TVM_DLL explicit UnorderedTuplePattern(tvm::ffi::Array<DFPattern> fields);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(UnorderedTuplePattern, DFPattern,
                                                UnorderedTuplePatternNode);
};
 
class TupleGetItemPatternNode : public DFPatternNode {
 public:
  explicit TupleGetItemPatternNode(DFPattern tuple) : tuple(std::move(tuple)) {}
  explicit TupleGetItemPatternNode(ffi::UnsafeInit) : tuple(ffi::UnsafeInit{}) {}
 
  DFPattern tuple; 
  int index;       
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TupleGetItemPatternNode>()
        .def_ro("tuple", &TupleGetItemPatternNode::tuple)
        .def_ro("index", &TupleGetItemPatternNode::index);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TupleGetItemPattern", TupleGetItemPatternNode,
                                    DFPatternNode);
};
 
class TupleGetItemPattern : public DFPattern {
 public:
  TVM_DLL TupleGetItemPattern(DFPattern tuple, int index);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleGetItemPattern, DFPattern,
                                                TupleGetItemPatternNode);
};
 
class AndPatternNode : public DFPatternNode {
 public:
  explicit AndPatternNode(DFPattern left, DFPattern right)
      : left(std::move(left)), right(std::move(right)) {}
  explicit AndPatternNode(ffi::UnsafeInit) : left(ffi::UnsafeInit{}), right(ffi::UnsafeInit{}) {}
 
  DFPattern left;  
  DFPattern right; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AndPatternNode>()
        .def_ro("left", &AndPatternNode::left)
        .def_ro("right", &AndPatternNode::right);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AndPattern", AndPatternNode, DFPatternNode);
};
 
class AndPattern : public DFPattern {
 public:
  TVM_DLL AndPattern(DFPattern lhs, DFPattern rhs);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AndPattern, DFPattern, AndPatternNode);
};
 
class OrPatternNode : public DFPatternNode {
 public:
  explicit OrPatternNode(DFPattern left, DFPattern right)
      : left(std::move(left)), right(std::move(right)) {}
  explicit OrPatternNode(ffi::UnsafeInit) : left(ffi::UnsafeInit{}), right(ffi::UnsafeInit{}) {}
 
  DFPattern left;  
  DFPattern right; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OrPatternNode>()
        .def_ro("left", &OrPatternNode::left)
        .def_ro("right", &OrPatternNode::right);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.OrPattern", OrPatternNode, DFPatternNode);
};
 
class OrPattern : public DFPattern {
 public:
  TVM_DLL OrPattern(DFPattern left, DFPattern right);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OrPattern, DFPattern, OrPatternNode);
};
 
class NotPatternNode : public DFPatternNode {
 public:
  explicit NotPatternNode(DFPattern reject) : reject(std::move(reject)) {}
  explicit NotPatternNode(ffi::UnsafeInit) : reject(ffi::UnsafeInit{}) {}
 
  DFPattern reject; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<NotPatternNode>().def_ro("reject", &NotPatternNode::reject);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.NotPattern", NotPatternNode, DFPatternNode);
};
 
class NotPattern : public DFPattern {
 public:
  TVM_DLL NotPattern(DFPattern reject);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(NotPattern, DFPattern, NotPatternNode);
};
 
class WildcardPatternNode : public DFPatternNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<WildcardPatternNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.WildcardPattern", WildcardPatternNode,
                                    DFPatternNode);
};
 
class WildcardPattern : public DFPattern {
 public:
  WildcardPattern();
  explicit WildcardPattern(ffi::ObjectPtr<WildcardPatternNode> data)
      : DFPattern(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
 
  // Declaring WildcardPattern declared as non-nullable avoids the
  // default zero-parameter constructor for ffi::ObjectRef with `data_ =
  // nullptr`.  This allows a zero-parameter constructor to be
  // declared here, to create a valid wildcard instance.
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WildcardPattern, DFPattern, WildcardPatternNode);
};
 
class TypePatternNode : public DFPatternNode {
 public:
  explicit TypePatternNode(DFPattern pattern) : pattern(std::move(pattern)) {}
  explicit TypePatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {}
 
  DFPattern pattern;         
  Type ty = Type::Missing(); 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TypePatternNode>()
        .def_ro("pattern", &TypePatternNode::pattern)
        .def_ro("ty", &TypePatternNode::ty);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.TypePattern", TypePatternNode, DFPatternNode);
};
 
class TypePattern : public DFPattern {
 public:
  TVM_DLL TypePattern(DFPattern pattern, Type ty);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TypePattern, DFPattern, TypePatternNode);
};
 
class ShapePatternNode : public DFPatternNode {
 public:
  explicit ShapePatternNode(DFPattern pattern) : pattern(std::move(pattern)) {}
  explicit ShapePatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {}
 
  DFPattern pattern;          
  ffi::Array<PrimExpr> shape; 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ShapePatternNode>()
        .def_ro("pattern", &ShapePatternNode::pattern)
        .def_ro("shape", &ShapePatternNode::shape);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ShapePattern", ShapePatternNode, DFPatternNode);
};
 
class ShapePattern : public DFPattern {
 public:
  TVM_DLL ShapePattern(DFPattern pattern, ffi::Array<PrimExpr> type);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapePattern, DFPattern, ShapePatternNode);
};
 
class SameShapeConstraintNode : public DFConstraintNode {
 public:
  ffi::Array<DFPattern> args; 
  ffi::Array<DFPattern> GetDependentPatterns() const override { return args; }
 
  std::tuple<PrimExpr, bool> AsCondition(
      std::function<ffi::Optional<Var>(const DFPatternNode*)> match_state) const override;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SameShapeConstraintNode>().def_ro("args", &SameShapeConstraintNode::args);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.SameShapeConstraint", SameShapeConstraintNode,
                                    DFConstraintNode);
};
 
class SameShapeConstraint : public DFConstraint {
 public:
  TVM_DLL SameShapeConstraint(ffi::Array<DFPattern> args);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(SameShapeConstraint, DFConstraint,
                                                SameShapeConstraintNode);
};
 
class DataTypePatternNode : public DFPatternNode {
 public:
  explicit DataTypePatternNode(DFPattern pattern) : pattern(std::move(pattern)) {}
  explicit DataTypePatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {}
 
  DFPattern pattern; 
  DLDataType dtype;  
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataTypePatternNode>()
        .def_ro("pattern", &DataTypePatternNode::pattern)
        .def_ro("dtype", &DataTypePatternNode::dtype);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.DataTypePattern", DataTypePatternNode,
                                    DFPatternNode);
};
 
class DataTypePattern : public DFPattern {
 public:
  TVM_DLL DataTypePattern(DFPattern pattern, DLDataType dtype);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DataTypePattern, DFPattern, DataTypePatternNode);
};
 
class AttrPatternNode : public DFPatternNode {
 public:
  explicit AttrPatternNode(DFPattern pattern) : pattern(std::move(pattern)) {}
  explicit AttrPatternNode(ffi::UnsafeInit) : pattern(ffi::UnsafeInit{}) {}
 
  DFPattern pattern; 
  DictAttrs attrs;   
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AttrPatternNode>()
        .def_ro("pattern", &AttrPatternNode::pattern)
        .def_ro("attrs", &AttrPatternNode::attrs);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.AttrPattern", AttrPatternNode, DFPatternNode);
};
 
class AttrPattern : public DFPattern {
 public:
  TVM_DLL AttrPattern(DFPattern pattern, DictAttrs attrs);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrPattern, DFPattern, AttrPatternNode);
};
 
class ExternFuncPatternNode : public DFPatternNode {
 public:
  ffi::String global_symbol_; 
  const ffi::String& global_symbol() const { return global_symbol_; }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExternFuncPatternNode>().def_ro("global_symbol",
                                                    &ExternFuncPatternNode::global_symbol_);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.dpl.ExternFuncPattern", ExternFuncPatternNode,
                                    DFPatternNode);
};
 
class ExternFuncPattern : public DFPattern {
 public:
  TVM_DLL ExternFuncPattern(ffi::String global_symbol);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternFuncPattern, DFPattern,
                                                ExternFuncPatternNode);
};
 
VarPattern IsVar(const ffi::String& name);
ConstantPattern IsConst();
WildcardPattern Wildcard();
ExprPattern IsExpr(const Expr& expr);
ExprPattern IsOp(const ffi::String& op_name);
// Todo(relax-team): Dataflow pattern for Type, and match out_ty
CallPattern IsCallTIR(const ffi::String& name, ffi::Optional<TuplePattern> args = std::nullopt);
CallPattern IsCallTIR(const ffi::String& name, TuplePattern var_args);
CallPattern IsCallDPSPacked(const ffi::String& name,
                            ffi::Optional<TuplePattern> args = std::nullopt);
CallPattern IsCallDPSPacked(const ffi::String& name, TuplePattern var_args);
DFPattern IsTuple(const ffi::Array<DFPattern>& fields, bool unordered = false);
TupleGetItemPattern IsTupleGetItem(const DFPattern tuple, int index = -1);
 
template <typename... Args>
CallPattern DFPattern::operator()(Args&&... args) const {
  return CallPattern(ffi::GetRef<DFPattern>(this->get()),
                     ffi::Array<DFPattern>({std::forward<Args>(args)...}));
}
 
}  // namespace relax
}  // namespace tvm
#endif  // TVM_RELAX_DATAFLOW_PATTERN_H_