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_RELAY_DATAFLOW_PATTERN_H_
#define TVM_RELAY_DATAFLOW_PATTERN_H_
 
#include <tvm/relay/expr.h>
#include <tvm/relay/type.h>
 
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
namespace tvm {
namespace relay {
 
class DFPatternNode : public Object {
 public:
  static constexpr const char* _type_key = "DFPatternNode";
  TVM_DECLARE_BASE_OBJECT_INFO(DFPatternNode, Object);
};
 
class DFPattern : public ObjectRef {
 public:
  DFPattern operator()(const std::vector<DFPattern>& args) const;
  DFPattern operator+(const DFPattern& other) const;
  DFPattern operator-(const DFPattern& other) const;
  DFPattern operator*(const DFPattern& other) const;
  DFPattern operator/(const DFPattern& other) const;
  DFPattern operator||(const DFPattern& other) const;
  DFPattern Optional(const std::function<DFPattern(const DFPattern&)>& func) const;
  DFPattern HasAttr(const Map<String, ObjectRef>& attrs) const;
  DFPattern HasType(const Type& type) const;
  DFPattern HasDtype(const DataType& dtype) const;
  DFPattern HasDtype(const std::string& dtype) const;
  DFPattern HasShape(const Array<PrimExpr> shape) const;
 
  TVM_DEFINE_OBJECT_REF_METHODS(DFPattern, ObjectRef, DFPatternNode);
};
 
class ExprPatternNode : public DFPatternNode {
 public:
  Expr expr;
 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("expr", &expr); }
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.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 VarPattern;
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 = "relay.dataflow_pattern.VarPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(VarPatternNode, DFPatternNode);
};
 
class VarPattern : public DFPattern {
 public:
  TVM_DLL VarPattern(String name_hint);
  TVM_DEFINE_OBJECT_REF_METHODS(VarPattern, DFPattern, VarPatternNode);
};
 
class ConstantPattern;
class ConstantPatternNode : public DFPatternNode {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) {}
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.ConstantPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantPatternNode, DFPatternNode);
};
 
class ConstantPattern : public DFPattern {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(ConstantPattern, DFPattern, ConstantPatternNode);
};
 
class CallPattern;
class CallPatternNode : public DFPatternNode {
 public:
  DFPattern op;
 
  tvm::Array<relay::DFPattern> args;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("op", &op);
    v->Visit("args", &args);
  }
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.CallPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(CallPatternNode, DFPatternNode);
};
 
class CallPattern : public DFPattern {
 public:
  TVM_DLL CallPattern(DFPattern op, Array<DFPattern> args);
  TVM_DEFINE_OBJECT_REF_METHODS(CallPattern, DFPattern, CallPatternNode);
};
 
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 = "relay.dataflow_pattern.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);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(FunctionPatternNode);
};
 
class LetPatternNode : public DFPatternNode {
 public:
  DFPattern var;
  DFPattern value;
  DFPattern body;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("var", &var);
    v->Visit("value", &value);
    v->Visit("body", &body);
  }
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.LetPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(LetPatternNode, DFPatternNode);
};
 
class LetPattern : public DFPattern {
 public:
  TVM_DLL LetPattern(DFPattern var, DFPattern value, DFPattern body);
 
  TVM_DEFINE_OBJECT_REF_METHODS(LetPattern, DFPattern, LetPatternNode);
};
 
class TuplePattern;
class TuplePatternNode : public DFPatternNode {
 public:
  tvm::Array<DFPattern> fields;
 
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.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 TupleGetItemPattern;
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 = "relay.dataflow_pattern.TupleGetItemPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleGetItemPatternNode, DFPatternNode);
};
 
class IfPatternNode : public DFPatternNode {
 public:
  DFPattern cond, true_branch, false_branch;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("cond", &cond);
    v->Visit("true_branch", &true_branch);
    v->Visit("false_branch", &false_branch);
  }
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.IfPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(IfPatternNode, DFPatternNode);
};
 
class IfPattern : public DFPattern {
 public:
  TVM_DLL IfPattern(DFPattern cond, DFPattern then_clause, DFPattern else_clause);
  TVM_DEFINE_OBJECT_REF_METHODS(IfPattern, DFPattern, IfPatternNode);
};
 
class TupleGetItemPattern : public DFPattern {
 public:
  TVM_DLL TupleGetItemPattern(DFPattern tuple, int index);
  TVM_DEFINE_OBJECT_REF_METHODS(TupleGetItemPattern, DFPattern, TupleGetItemPatternNode);
};
 
class AltPattern;
class AltPatternNode : 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 = "relay.dataflow_pattern.AltPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(AltPatternNode, DFPatternNode);
};
 
class AltPattern : public DFPattern {
 public:
  TVM_DLL AltPattern(DFPattern left, DFPattern right);
  TVM_DEFINE_OBJECT_REF_METHODS(AltPattern, DFPattern, AltPatternNode);
};
 
class WildcardPatternNode : public DFPatternNode {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) {}
 
  Optional<DFPattern> pattern{nullptr};
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.WildcardPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(WildcardPatternNode, DFPatternNode);
};
 
class WildcardPattern : public DFPattern {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(WildcardPattern, DFPattern, WildcardPatternNode);
 
  void redirect_to(DFPattern pat) const;
};
 
class TypePattern;
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 = "relay.dataflow_pattern.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 ShapePattern;
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 = "relay.dataflow_pattern.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 DataTypePattern;
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 = "relay.dataflow_pattern.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 AttrPattern;
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 = "relay.dataflow_pattern.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 DominatorPattern;
class DominatorPatternNode : public DFPatternNode {
 public:
  DFPattern parent;
  DFPattern path;
  DFPattern child;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("parent", &parent);
    v->Visit("path", &path);
    v->Visit("child", &child);
  }
 
  static constexpr const char* _type_key = "relay.dataflow_pattern.DominatorPattern";
  TVM_DECLARE_FINAL_OBJECT_INFO(DominatorPatternNode, DFPatternNode);
};
 
class DominatorPattern : public DFPattern {
 public:
  TVM_DLL DominatorPattern(DFPattern parent, DFPattern path, DFPattern child);
  TVM_DEFINE_OBJECT_REF_METHODS(DominatorPattern, DFPattern, DominatorPatternNode);
};
 
DFPattern IsVar(const String& name);
DFPattern IsConstant();
DFPattern IsWildcard();
DFPattern IsExpr(const Expr& expr);
DFPattern IsOp(const String& op_name);
DFPattern IsTuple(const Array<DFPattern>& fields);
DFPattern IsTupleGetItem(const DFPattern tuple, int index = -1);
 
class DFPatternPrinter : public ReprPrinter {
 public:
  std::stringstream string_stream{};
 
  std::unordered_map<DFPattern, std::pair<size_t, std::string>, ObjectPtrHash, ObjectPtrEqual>
      memo_{};
  std::vector<DFPattern> auxiliary_patterns{};
 
  DFPatternPrinter(std::ostream& stream)  // NOLINT(*)
      : ReprPrinter(stream) {}
  TVM_DLL void Print(const ObjectRef& node);
  using FType = NodeFunctor<void(const ObjectRef&, DFPatternPrinter*)>;
  TVM_DLL static FType& vtable();
};
 
inline std::ostream& operator<<(std::ostream& os,
                                const DFPattern& n) {  // NOLINT(*)
  std::stringstream string_stream{}, tmp_stream{};
  DFPatternPrinter printer{tmp_stream};
  printer.Print(n);
  string_stream << "Main pattern:" << std::endl;
  string_stream << printer.string_stream.str();
  string_stream << std::endl;
  string_stream << "Auxiliary patterns:";
  for (const DFPattern& pat : printer.auxiliary_patterns) {
    string_stream << std::endl;
    string_stream << printer.memo_[pat].second;
  }
  os << string_stream.str();
  return os;
}
 
String PrettyPrint(const DFPattern& pattern);
 
}  // namespace relay
}  // namespace tvm
#endif  // TVM_RELAY_DATAFLOW_PATTERN_H_