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/ir/expr.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
#include <tvm/runtime/container/array.h>
#include <tvm/runtime/container/optional.h>
#include <tvm/support/with.h>
 
#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
 
namespace tvm {
 
namespace arith {
class Analyzer;
}
 
namespace relax {
 
class PatternSeq;
class CallPattern;
class OrPattern;
class AndPattern;
class NotPattern;
class ShapePattern;
class StructInfoPattern;
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 Object {
 public:
  static constexpr const char* _type_key = "DFPatternNode";
  TVM_DECLARE_BASE_OBJECT_INFO(DFPatternNode, Object);
};
 
class DFPattern : public 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 Map<String, ObjectRef>& attrs) const;
  TVM_DLL StructInfoPattern HasStructInfo(const StructInfo& struct_info) const;
  TVM_DLL TypePattern HasType(const Type& type) const;
  TVM_DLL DataTypePattern HasDtype(const DataType& dtype) const;
  TVM_DLL DataTypePattern HasDtype(const std::string& dtype) const;
  TVM_DLL ShapePattern HasShape(const Array<PrimExpr>& shape) const;
  TVM_DLL SameShapeConstraint HasSameShapeAs(const DFPattern& other) const;
  TVM_DLL DFPattern dup() const;
 
  TVM_DLL operator PatternSeq() const;
 
  TVM_DEFINE_OBJECT_REF_METHODS(DFPattern, 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 Object {
 public:
  virtual Array<DFPattern> GetDependentPatterns() const = 0;
 
  virtual std::tuple<PrimExpr, bool> AsPrimExpr(
      std::function<Optional<Var>(const DFPatternNode*)> match_state) const = 0;
 
  static constexpr const char* _type_key = "DFConstraintNode";
  static constexpr const uint32_t _type_child_slots = 1;
  TVM_DECLARE_BASE_OBJECT_INFO(DFConstraintNode, Object);
};
 
class DFConstraint : public ObjectRef {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(DFConstraint, ObjectRef, DFConstraintNode);
};
 
class PatternSeqNode final : public Object {
 public:
  tvm::Array<DFPattern> patterns;         
  std::vector<PairCons> pair_constraints; 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("patterns", &patterns); }
  static constexpr const char* _type_key = "relax.dpl.PatternSeq";
  TVM_DECLARE_BASE_OBJECT_INFO(PatternSeqNode, Object);
};
 
class PatternSeq final : public ObjectRef {
 public:
  TVM_DLL explicit PatternSeq(DFPattern init_pattern);
  TVM_DLL explicit PatternSeq(tvm::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_DEFINE_OBJECT_REF_METHODS(PatternSeq, ObjectRef, PatternSeqNode);
};
 
class PatternContextNode : public 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;
 
  static constexpr const char* _type_key = "relax.dpl.PatternContext";
  TVM_DECLARE_FINAL_OBJECT_INFO(PatternContextNode, Object);
};
 
class PatternContext : public ObjectRef {
 public:
  TVM_DLL explicit PatternContext(ObjectPtr<Object> n) : ObjectRef(n) {}
  TVM_DLL explicit PatternContext(bool incremental = false);
 
  const PatternContextNode* operator->() const {
    ICHECK(get() != nullptr);
    return static_cast<const PatternContextNode*>(get());
  }
 
  PatternContextNode* operator->() {
    ICHECK(get() != nullptr);
    return static_cast<PatternContextNode*>(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;
      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 Optional<PatternContext> Current();
 
  TVM_DLL void EnterWithScope() const;
  TVM_DLL void ExitWithScope() const;
 
 private:
  friend class With<PatternContext>;
};
 
class ExprPatternNode : public DFPatternNode {
 public:
  Expr expr; 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("expr", &expr); }
 
  static constexpr const char* _type_key = "relax.dpl.ExprPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(ExprPatternNode, DFPatternNode);
};
 
class ExprPattern : public DFPattern {
 public:
  TVM_DLL explicit ExprPattern(Expr expr);
  TVM_DEFINE_OBJECT_REF_METHODS(ExprPattern, DFPattern, ExprPatternNode);
};
 
class VarPatternNode : public DFPatternNode {
 public:
  String name;
  const String& name_hint() const { return name; }
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("name", &name); }
 
  static constexpr const char* _type_key = "relax.dpl.VarPattern";
  TVM_DECLARE_BASE_OBJECT_INFO(VarPatternNode, DFPatternNode);
};
 
class VarPattern : public DFPattern {
 public:
  TVM_DLL VarPattern(String name_hint);
  TVM_DEFINE_OBJECT_REF_METHODS(VarPattern, DFPattern, VarPatternNode);
};
 
class DataflowVarPatternNode : public VarPatternNode {
 public:
  static constexpr const char* _type_key = "relax.dpl.DataflowVarPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(DataflowVarPatternNode, DFPatternNode);
};
 
class DataflowVarPattern : public DFPattern {
 public:
  TVM_DLL DataflowVarPattern(String name_hint);
  TVM_DEFINE_OBJECT_REF_METHODS(DataflowVarPattern, DFPattern, DataflowVarPatternNode);
};
 
class GlobalVarPatternNode : public VarPatternNode {
 public:
  static constexpr const char* _type_key = "relax.dpl.GlobalVarPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(GlobalVarPatternNode, DFPatternNode);
};
 
class GlobalVarPattern : public DFPattern {
 public:
  TVM_DLL GlobalVarPattern(String name_hint);
  TVM_DEFINE_OBJECT_REF_METHODS(GlobalVarPattern, DFPattern, GlobalVarPatternNode);
};
 
class ConstantPatternNode : public DFPatternNode {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) {}
 
  static constexpr const char* _type_key = "relax.dpl.ConstantPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantPatternNode, DFPatternNode);
};
 
class ConstantPattern : public DFPattern {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(ConstantPattern, DFPattern, ConstantPatternNode);
};
 
class CallPatternNode : public DFPatternNode {
 public:
  DFPattern op;               
  tvm::Array<DFPattern> args; 
  bool varg_default_wildcard; 
  // Todo(relax-team): Dataflow pattern for StructInfo, and match sinfo_args
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("op", &op);
    v->Visit("args", &args);
  }
 
  static constexpr const char* _type_key = "relax.dpl.CallPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(CallPatternNode, DFPatternNode);
};
 
class CallPattern : public DFPattern {
 public:
  TVM_DLL CallPattern(DFPattern op, Array<DFPattern> args, bool varg_default_wildcard = false);
  TVM_DEFINE_OBJECT_REF_METHODS(CallPattern, DFPattern, CallPatternNode);
};
 
class PrimArrPatternNode : public DFPatternNode {
 public:
  Array<PrimExpr> fields; 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
  static constexpr const char* _type_key = "relax.dpl.PrimArrPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(PrimArrPatternNode, DFPatternNode);
};
 
class PrimArrPattern : public DFPattern {
 public:
  TVM_DLL PrimArrPattern(Array<PrimExpr> arr);
  TVM_DEFINE_OBJECT_REF_METHODS(PrimArrPattern, DFPattern, PrimArrPatternNode);
};
 
class FunctionPatternNode : public DFPatternNode {
 public:
  tvm::Array<DFPattern> params; 
  DFPattern body; 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("params", &params);
    v->Visit("body", &body);
  }
 
  static constexpr const char* _type_key = "relax.dpl.FunctionPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(FunctionPatternNode, DFPatternNode);
};
 
class FunctionPattern : public DFPattern {
 public:
  TVM_DLL FunctionPattern(tvm::Array<DFPattern> params, DFPattern body);
 
  TVM_DEFINE_OBJECT_REF_METHODS(FunctionPattern, DFPattern, FunctionPatternNode);
};
 
class TuplePatternNode : public DFPatternNode {
 public:
  tvm::Array<DFPattern> fields; 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
 
  static constexpr const char* _type_key = "relax.dpl.TuplePattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(TuplePatternNode, DFPatternNode);
};
 
class TuplePattern : public DFPattern {
 public:
  TVM_DLL explicit TuplePattern(tvm::Array<DFPattern> fields);
  TVM_DEFINE_OBJECT_REF_METHODS(TuplePattern, DFPattern, TuplePatternNode);
};
 
class UnorderedTuplePatternNode : public DFPatternNode {
 public:
  tvm::Array<DFPattern> fields; 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
 
  static constexpr const char* _type_key = "relax.dpl.UnorderedTuplePattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(UnorderedTuplePatternNode, DFPatternNode);
};
 
class UnorderedTuplePattern : public DFPattern {
 public:
  TVM_DLL explicit UnorderedTuplePattern(tvm::Array<DFPattern> fields);
  TVM_DEFINE_OBJECT_REF_METHODS(UnorderedTuplePattern, DFPattern, UnorderedTuplePatternNode);
};
 
class TupleGetItemPatternNode : public DFPatternNode {
 public:
  DFPattern tuple; 
  int index;       
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("tuple", &tuple);
    v->Visit("index", &index);
  }
 
  static constexpr const char* _type_key = "relax.dpl.TupleGetItemPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleGetItemPatternNode, DFPatternNode);
};
 
class TupleGetItemPattern : public DFPattern {
 public:
  TVM_DLL TupleGetItemPattern(DFPattern tuple, int index);
  TVM_DEFINE_OBJECT_REF_METHODS(TupleGetItemPattern, DFPattern, TupleGetItemPatternNode);
};
 
class AndPatternNode : public DFPatternNode {
 public:
  DFPattern left;  
  DFPattern right; 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("left", &left);
    v->Visit("right", &right);
  }
 
  static constexpr const char* _type_key = "relax.dpl.AndPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(AndPatternNode, DFPatternNode);
};
 
class AndPattern : public DFPattern {
 public:
  TVM_DLL AndPattern(DFPattern lhs, DFPattern rhs);
  TVM_DEFINE_OBJECT_REF_METHODS(AndPattern, DFPattern, AndPatternNode);
};
 
class OrPatternNode : public DFPatternNode {
 public:
  DFPattern left;  
  DFPattern right; 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("left", &left);
    v->Visit("right", &right);
  }
 
  static constexpr const char* _type_key = "relax.dpl.OrPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(OrPatternNode, DFPatternNode);
};
 
class OrPattern : public DFPattern {
 public:
  TVM_DLL OrPattern(DFPattern left, DFPattern right);
  TVM_DEFINE_OBJECT_REF_METHODS(OrPattern, DFPattern, OrPatternNode);
};
 
class NotPatternNode : public DFPatternNode {
 public:
  DFPattern reject; 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("reject", &reject); }
 
  static constexpr const char* _type_key = "relax.dpl.NotPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(NotPatternNode, DFPatternNode);
};
 
class NotPattern : public DFPattern {
 public:
  TVM_DLL NotPattern(DFPattern reject);
  TVM_DEFINE_OBJECT_REF_METHODS(NotPattern, DFPattern, NotPatternNode);
};
 
class WildcardPatternNode : public DFPatternNode {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) {}
 
  static constexpr const char* _type_key = "relax.dpl.WildcardPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(WildcardPatternNode, DFPatternNode);
};
 
class WildcardPattern : public DFPattern {
 public:
  WildcardPattern();
 
  // Declaring WildcardPattern declared as non-nullable avoids the
  // default zero-parameter constructor for ObjectRef with `data_ =
  // nullptr`.  This allows a zero-parameter constructor to be
  // declared here, to create a valid wildcard instance.
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(WildcardPattern, DFPattern, WildcardPatternNode);
};
 
class TypePatternNode : public DFPatternNode {
 public:
  DFPattern pattern; 
  Type type;         
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pattern", &pattern);
    v->Visit("type", &type);
  }
 
  static constexpr const char* _type_key = "relax.dpl.TypePattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(TypePatternNode, DFPatternNode);
};
 
class TypePattern : public DFPattern {
 public:
  TVM_DLL TypePattern(DFPattern pattern, Type type);
  TVM_DEFINE_OBJECT_REF_METHODS(TypePattern, DFPattern, TypePatternNode);
};
 
class StructInfoPatternNode : public DFPatternNode {
 public:
  DFPattern pattern;      
  StructInfo struct_info; 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pattern", &pattern);
    v->Visit("struct_info", &struct_info);
  }
 
  static constexpr const char* _type_key = "relax.dpl.StructInfoPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(StructInfoPatternNode, DFPatternNode);
};
 
class StructInfoPattern : public DFPattern {
 public:
  TVM_DLL StructInfoPattern(DFPattern pattern, StructInfo struct_info);
  TVM_DEFINE_OBJECT_REF_METHODS(StructInfoPattern, DFPattern, StructInfoPatternNode);
};
 
class ShapePatternNode : public DFPatternNode {
 public:
  DFPattern pattern;     
  Array<PrimExpr> shape; 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pattern", &pattern);
    v->Visit("shape", &shape);
  }
 
  static constexpr const char* _type_key = "relax.dpl.ShapePattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(ShapePatternNode, DFPatternNode);
};
 
class ShapePattern : public DFPattern {
 public:
  TVM_DLL ShapePattern(DFPattern pattern, Array<PrimExpr> type);
  TVM_DEFINE_OBJECT_REF_METHODS(ShapePattern, DFPattern, ShapePatternNode);
};
 
class SameShapeConstraintNode : public DFConstraintNode {
 public:
  Array<DFPattern> args; 
  Array<DFPattern> GetDependentPatterns() const override { return args; }
 
  std::tuple<PrimExpr, bool> AsPrimExpr(
      std::function<Optional<Var>(const DFPatternNode*)> match_state) const override;
 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("args", &args); }
 
  static constexpr const char* _type_key = "relax.dpl.SameShapeConstraint";
  TVM_DECLARE_FINAL_OBJECT_INFO(SameShapeConstraintNode, DFConstraintNode);
};
 
class SameShapeConstraint : public DFConstraint {
 public:
  TVM_DLL SameShapeConstraint(Array<DFPattern> args);
  TVM_DEFINE_OBJECT_REF_METHODS(SameShapeConstraint, DFConstraint, SameShapeConstraintNode);
};
 
class DataTypePatternNode : public DFPatternNode {
 public:
  DFPattern pattern; 
  DataType dtype;    
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pattern", &pattern);
    v->Visit("dtype", &dtype);
  }
 
  static constexpr const char* _type_key = "relax.dpl.DataTypePattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(DataTypePatternNode, DFPatternNode);
};
 
class DataTypePattern : public DFPattern {
 public:
  TVM_DLL DataTypePattern(DFPattern pattern, DataType dtype);
  TVM_DEFINE_OBJECT_REF_METHODS(DataTypePattern, DFPattern, DataTypePatternNode);
};
 
class AttrPatternNode : public DFPatternNode {
 public:
  DFPattern pattern; 
  DictAttrs attrs;   
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pattern", &pattern);
    v->Visit("attrs", &attrs);
  }
 
  static constexpr const char* _type_key = "relax.dpl.AttrPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(AttrPatternNode, DFPatternNode);
};
 
class AttrPattern : public DFPattern {
 public:
  TVM_DLL AttrPattern(DFPattern pattern, DictAttrs attrs);
  TVM_DEFINE_OBJECT_REF_METHODS(AttrPattern, DFPattern, AttrPatternNode);
};
 
class ExternFuncPatternNode : public DFPatternNode {
 public:
  String global_symbol_; 
  const String& global_symbol() const { return global_symbol_; }
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("global_symbol", &global_symbol_); }
 
  static constexpr const char* _type_key = "relax.dpl.ExternFuncPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(ExternFuncPatternNode, DFPatternNode);
};
 
class ExternFuncPattern : public DFPattern {
 public:
  TVM_DLL ExternFuncPattern(String global_symbol);
  TVM_DEFINE_OBJECT_REF_METHODS(ExternFuncPattern, DFPattern, ExternFuncPatternNode);
};
 
VarPattern IsVar(const String& name);
ConstantPattern IsConst();
WildcardPattern Wildcard();
ExprPattern IsExpr(const Expr& expr);
ExprPattern IsOp(const String& op_name);
// Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
CallPattern IsCallTIR(const String& name, Optional<TuplePattern> args = NullOpt);
CallPattern IsCallTIR(const String& name, TuplePattern var_args);
CallPattern IsCallDPSPacked(const String& name, Optional<TuplePattern> args = NullOpt);
CallPattern IsCallDPSPacked(const String& name, TuplePattern var_args);
DFPattern IsTuple(const 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(GetRef<DFPattern>(this->get()),
                     Array<DFPattern>({std::forward<Args>(args)...}));
}
 
}  // namespace relax
}  // namespace tvm
#endif  // TVM_RELAX_DATAFLOW_PATTERN_H_