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

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

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

}  // namespace relay
}  // namespace tvm
#endif  // TVM_RELAY_DATAFLOW_PATTERN_H_