type_relation.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_IR_TYPE_RELATION_H_
#define TVM_IR_TYPE_RELATION_H_
 
#include <tvm/ir/attrs.h>
#include <tvm/ir/diagnostic.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/module.h>
#include <tvm/ir/type.h>
#include <tvm/runtime/logging.h>
 
namespace tvm {
 
class TypeCallNode : public TypeNode {
 public:
  Type func;
  Array<Type> args;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("func", &func);
    v->Visit("args", &args);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const TypeCallNode* other, SEqualReducer equal) const {
    return equal(func, other->func) && equal(args, other->args);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(func);
    hash_reduce(args);
  }
 
  static constexpr const char* _type_key = "TypeCall";
  TVM_DECLARE_FINAL_OBJECT_INFO(TypeCallNode, TypeNode);
};
 
class TypeCall : public Type {
 public:
  TVM_DLL TypeCall(Type func, Array<Type> args);
 
  TVM_DEFINE_OBJECT_REF_METHODS(TypeCall, Type, TypeCallNode);
};
 
class TypeReporterNode : public Object {
 public:
  virtual ~TypeReporterNode() {}
  TVM_DLL virtual void Assign(const Type& dst, const Type& src) = 0;
 
  TVM_DLL virtual bool Assert(const PrimExpr& cond) = 0;
  TVM_DLL virtual bool AssertEQ(const PrimExpr& lhs, const PrimExpr& rhs) = 0;
 
  TVM_DLL virtual void SetSpan(const Span& span) = 0;
 
  TVM_DLL virtual Span GetSpan() = 0;
 
  TVM_DLL virtual DiagnosticContext GetDiagCtx() = 0;
 
  TVM_DLL virtual IRModule GetModule() = 0;
 
  // solver is not serializable.
  void VisitAttrs(AttrVisitor* v) {}
 
  static constexpr const char* _type_key = "TypeReporter";
  TVM_DECLARE_FINAL_OBJECT_INFO(TypeReporterNode, Object);
};
 
class TypeReporter : public ObjectRef {
 public:
  TypeReporter() {}
  explicit TypeReporter(ObjectPtr<Object> n) : ObjectRef(n) {}
  TypeReporterNode* operator->() const {
    return const_cast<TypeReporterNode*>(static_cast<const TypeReporterNode*>(get()));
  }
  using ContainerType = TypeReporterNode;
};
 
using TypeRelationFn = TypedEnvFunc<bool(const Array<Type>& args, int num_inputs,
                                         const Attrs& attrs, const TypeReporter& reporter)>;
 
class TypeRelationNode : public TypeConstraintNode {
 public:
  TypeRelationFn func;
  Array<Type> args;
  int num_inputs;
  Attrs attrs;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("func", &func);
    v->Visit("args", &args);
    v->Visit("num_inputs", &num_inputs);
    v->Visit("attrs", &attrs);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const TypeRelationNode* other, SEqualReducer equal) const {
    return equal(func, other->func) && equal(args, other->args) &&
           equal(num_inputs, other->num_inputs) && equal(attrs, other->attrs);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(func);
    hash_reduce(args);
    hash_reduce(num_inputs);
    hash_reduce(attrs);
  }
 
  static constexpr const char* _type_key = "TypeRelation";
  TVM_DECLARE_FINAL_OBJECT_INFO(TypeRelationNode, TypeConstraintNode);
};
 
class TypeRelation : public TypeConstraint {
 public:
  TVM_DLL TypeRelation(TypeRelationFn func, Array<Type> args, int num_inputs, Attrs attrs);
 
  TVM_DEFINE_OBJECT_REF_METHODS(TypeRelation, TypeConstraint, TypeRelationNode);
};
}  // namespace tvm
#endif  // TVM_IR_TYPE_RELATION_H_