adt.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_ADT_H_
#define TVM_RELAY_ADT_H_

#include <tvm/ir/adt.h>
#include <tvm/ir/attrs.h>
#include <tvm/relay/base.h>
#include <tvm/relay/expr.h>
#include <tvm/relay/type.h>

#include <functional>
#include <string>
#include <utility>

namespace tvm {
namespace relay {

using Constructor = tvm::Constructor;
using ConstructorNode = tvm::ConstructorNode;

using TypeData = tvm::TypeData;
using TypeDataNode = tvm::TypeDataNode;

class PatternNode : public RelayNode {
 public:
  static constexpr const char* _type_key = "relay.Pattern";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_BASE_OBJECT_INFO(PatternNode, Object);
};

class Pattern : public ObjectRef {
 public:
  Pattern() {}
  explicit Pattern(ObjectPtr<tvm::Object> p) : ObjectRef(p) {}

  using ContainerType = PatternNode;
};

class PatternWildcard;
class PatternWildcardNode : public PatternNode {
 public:
  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("span", &span); }

  bool SEqualReduce(const PatternNode* other, SEqualReducer equal) const { return true; }

  void SHashReduce(SHashReducer hash_reduce) const {}

  static constexpr const char* _type_key = "relay.PatternWildcard";
  TVM_DECLARE_FINAL_OBJECT_INFO(PatternWildcardNode, PatternNode);
};

class PatternWildcard : public Pattern {
 public:
  /* \brief Overload the default constructors. */
  TVM_DLL PatternWildcard();
  explicit PatternWildcard(ObjectPtr<Object> n) : Pattern(n) {}
  /* \brief Copy constructor. */
  PatternWildcard(const PatternWildcard& pat) : PatternWildcard(pat.data_) {}
  /* \brief Move constructor. */
  PatternWildcard(PatternWildcard&& pat) : PatternWildcard(std::move(pat.data_)) {}
  /* \brief Copy assignment. */
  PatternWildcard& operator=(const PatternWildcard& other) {
    (*this).data_ = other.data_;
    return *this;
  }
  /* \brief Move assignment. */
  PatternWildcard& operator=(PatternWildcard&& other) {
    (*this).data_ = std::move(other.data_);
    return *this;
  }

  const PatternWildcardNode* operator->() const {
    return static_cast<const PatternWildcardNode*>(get());
  }

  using ContainerType = PatternWildcardNode;
};

class PatternVar;
class PatternVarNode : public PatternNode {
 public:
  tvm::relay::Var var;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("var", &var);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const PatternVarNode* other, SEqualReducer equal) const {
    return equal.DefEqual(var, other->var);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce.DefHash(var); }

  static constexpr const char* _type_key = "relay.PatternVar";
  TVM_DECLARE_FINAL_OBJECT_INFO(PatternVarNode, PatternNode);
};

class PatternVar : public Pattern {
 public:
  TVM_DLL explicit PatternVar(tvm::relay::Var var);

  TVM_DEFINE_OBJECT_REF_METHODS(PatternVar, Pattern, PatternVarNode);
};

class PatternConstructor;
class PatternConstructorNode : public PatternNode {
 public:
  Constructor constructor;
  tvm::Array<Pattern> patterns;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("constructor", &constructor);
    v->Visit("patterns", &patterns);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const PatternConstructorNode* other, SEqualReducer equal) const {
    return equal(constructor, other->constructor) && equal(patterns, other->patterns);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(constructor);
    hash_reduce(patterns);
  }

  static constexpr const char* _type_key = "relay.PatternConstructor";
  TVM_DECLARE_FINAL_OBJECT_INFO(PatternConstructorNode, PatternNode);
};

class PatternConstructor : public Pattern {
 public:
  TVM_DLL PatternConstructor(Constructor constructor, tvm::Array<Pattern> patterns);

  TVM_DEFINE_OBJECT_REF_METHODS(PatternConstructor, Pattern, PatternConstructorNode);
};

class PatternTuple;
class PatternTupleNode : public PatternNode {
 public:
  /* TODO(@jroesch): rename to field_pats */
  tvm::Array<Pattern> patterns;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("patterns", &patterns);
    v->Visit("span", &span);
  }

  bool SEqualReduce(const PatternTupleNode* other, SEqualReducer equal) const {
    return equal(patterns, other->patterns);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(patterns); }

  static constexpr const char* _type_key = "relay.PatternTuple";
  TVM_DECLARE_FINAL_OBJECT_INFO(PatternTupleNode, PatternNode);
};

class PatternTuple : public Pattern {
 public:
  TVM_DLL explicit PatternTuple(tvm::Array<Pattern> patterns);

  TVM_DEFINE_OBJECT_REF_METHODS(PatternTuple, Pattern, PatternTupleNode);
};

class Clause;
class ClauseNode : public Object {
 public:
  Pattern lhs;
  Expr rhs;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("lhs", &lhs);
    v->Visit("rhs", &rhs);
  }

  bool SEqualReduce(const ClauseNode* other, SEqualReducer equal) const {
    return equal(lhs, other->lhs) && equal(rhs, other->rhs);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(lhs);
    hash_reduce(rhs);
  }

  static constexpr const char* _type_key = "relay.Clause";
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(ClauseNode, Object);
};

class Clause : public ObjectRef {
 public:
  TVM_DLL explicit Clause(Pattern lhs, Expr rhs);

  TVM_DEFINE_OBJECT_REF_METHODS(Clause, ObjectRef, ClauseNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ClauseNode);
};

Clause WithFields(Clause clause, Optional<Pattern> opt_lhs = Optional<Pattern>(),
                  Optional<Expr> opt_rhs = Optional<Expr>());

class Match;
class MatchNode : public ExprNode {
 public:
  Expr data;

  tvm::Array<Clause> clauses;

  bool complete;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("data", &data);
    v->Visit("clauses", &clauses);
    v->Visit("complete", &complete);
    v->Visit("virtual_device_", &virtual_device_);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const MatchNode* other, SEqualReducer equal) const {
    equal->MarkGraphNode();
    return equal(data, other->data) && equal(clauses, other->clauses) &&
           equal(complete, other->complete);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(data);
    hash_reduce(clauses);
    hash_reduce(complete);
  }

  static constexpr const char* _type_key = "relay.Match";
  TVM_DECLARE_FINAL_OBJECT_INFO(MatchNode, ExprNode);
};

class Match : public Expr {
 public:
  TVM_DLL Match(Expr data, tvm::Array<Clause> clauses, bool complete = true, Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(Match, RelayExpr, MatchNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(MatchNode);
};

Match WithFields(Match match, Optional<Expr> opt_data = Optional<Expr>(),
                 Optional<Array<Clause>> opt_clauses = Optional<Array<Clause>>(),
                 Optional<Bool> opt_complete = Optional<Bool>(),
                 Optional<Span> opt_span = Optional<Span>());

}  // namespace relay
}  // namespace tvm

#endif  // TVM_RELAY_ADT_H_