struct_info.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_STRUCT_INFO_H_
#define TVM_RELAX_STRUCT_INFO_H_
 
#include <tvm/ir/env_func.h>
#include <tvm/ir/source_map.h>
#include <tvm/node/node.h>
#include <tvm/relax/block_builder.h>
#include <tvm/relax/expr.h>
#include <tvm/relax/type.h>
 
namespace tvm {
namespace relax {
 
class ObjectStructInfoNode : public StructInfoNode {
 public:
  void VisitAttrs(AttrVisitor* v) { v->Visit("span", &span); }
 
  bool SEqualReduce(const ObjectStructInfoNode* other, SEqualReducer equal) const { return true; }
 
  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(0); }
 
  static constexpr const char* _type_key = "relax.ObjectStructInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(ObjectStructInfoNode, StructInfoNode);
};
 
class ObjectStructInfo : public StructInfo {
 public:
  TVM_DLL ObjectStructInfo(Span span = Span());
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ObjectStructInfo, StructInfo, ObjectStructInfoNode);
};
 
class PrimStructInfoNode : public StructInfoNode {
 public:
  Optional<PrimExpr> value;
 
  DataType dtype;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("value", &value);
    v->Visit("dtype", &dtype);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const PrimStructInfoNode* other, SEqualReducer equal) const {
    return equal(value, other->value) && equal(dtype, other->dtype);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(value);
    hash_reduce(dtype);
  }
 
  static constexpr const char* _type_key = "relax.PrimStructInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(PrimStructInfoNode, StructInfoNode);
};
 
class PrimStructInfo : public StructInfo {
 public:
  /* Construct a PrimStructInfo with a known dtype, but unknown value */
  TVM_DLL PrimStructInfo(DataType dtype, Span span = Span());
 
  /* Construct a PrimStructInfo with a known value */
  TVM_DLL PrimStructInfo(PrimExpr value, Span span = Span());
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(PrimStructInfo, StructInfo, PrimStructInfoNode);
};
 
class ShapeStructInfoNode : public StructInfoNode {
 public:
  Optional<Array<PrimExpr>> values;
  int ndim;
 
  bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("values", &values);
    v->Visit("ndim", &ndim);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const ShapeStructInfoNode* other, SEqualReducer equal) const {
    return equal(values, other->values) && equal(ndim, other->ndim);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(values);
    hash_reduce(ndim);
  }
 
  static constexpr const char* _type_key = "relax.ShapeStructInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(ShapeStructInfoNode, StructInfoNode);
};
 
class ShapeStructInfo : public StructInfo {
 public:
  TVM_DLL ShapeStructInfo(Array<PrimExpr> values, Span span = Span());
  TVM_DLL ShapeStructInfo(int ndim, Span span = Span());
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ShapeStructInfo, StructInfo, ShapeStructInfoNode);
};
 
class TensorStructInfoNode : public StructInfoNode {
 public:
  Optional<Expr> shape;
  Optional<VDevice> vdevice;
  DataType dtype;
  int ndim;
 
  bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
 
  bool IsUnknownDtype() const { return dtype.is_void(); }
 
  Optional<Array<PrimExpr>> GetShape() const {
    if (!shape.defined()) return {};
    ShapeStructInfo shape_sinfo = Downcast<ShapeStructInfo>(this->shape.value()->struct_info_);
    return shape_sinfo->values;
  }
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("shape", &shape);
    v->Visit("dtype", &dtype);
    v->Visit("vdevice", &vdevice);
    v->Visit("ndim", &ndim);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const TensorStructInfoNode* other, SEqualReducer equal) const {
    return equal(shape, other->shape) && equal(ndim, other->ndim) &&
           equal(vdevice, other->vdevice) && equal(dtype, other->dtype);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(shape);
    hash_reduce(dtype);
    hash_reduce(vdevice);
    hash_reduce(ndim);
  }
 
  static constexpr const char* _type_key = "relax.TensorStructInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(TensorStructInfoNode, StructInfoNode);
};
 
class TensorStructInfo : public StructInfo {
 public:
  TVM_DLL TensorStructInfo(Expr shape, DataType dtype, Optional<VDevice> vdevice = NullOpt,
                           Span span = Span());
 
  TVM_DLL TensorStructInfo(DataType dtype, int ndim, Optional<VDevice> vdevice = NullOpt,
                           Span span = Span());
 
  TVM_DEFINE_OBJECT_REF_METHODS(TensorStructInfo, StructInfo, TensorStructInfoNode);
};
 
class TupleStructInfoNode : public StructInfoNode {
 public:
  Array<StructInfo> fields;
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("fields", &fields);
    v->Visit("span", &span);
  }
 
  bool SEqualReduce(const TupleStructInfoNode* other, SEqualReducer equal) const {
    return equal(fields, other->fields);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(fields); }
 
  static constexpr const char* _type_key = "relax.TupleStructInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleStructInfoNode, StructInfoNode);
};
 
class TupleStructInfo : public StructInfo {
 public:
  TVM_DLL TupleStructInfo(Array<StructInfo> fields, Span span = Span());
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TupleStructInfo, StructInfo, TupleStructInfoNode);
};
 
using StructInfoDeriveFunc = TypedEnvFunc<StructInfo(const Call& call, const BlockBuilder& ctx)>;
 
class FuncStructInfoNode : public StructInfoNode {
 public:
  Optional<Array<StructInfo>> params;
  StructInfo ret;
  Optional<StructInfoDeriveFunc> derive_func;
  bool purity;
 
  bool IsOpaque() const { return !params.defined(); }
 
  void VisitAttrs(AttrVisitor* v) {
    v->Visit("params", &params);
    v->Visit("ret", &ret);
    v->Visit("derive_func", &derive_func);
    v->Visit("span", &span);
    v->Visit("purity", &purity);
  }
 
  bool SEqualReduce(const FuncStructInfoNode* other, SEqualReducer equal) const {
    return equal.DefEqual(params, other->params) && equal(ret, other->ret) &&
           equal(purity, other->purity) && equal(derive_func, other->derive_func);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(params);
    hash_reduce(ret);
    hash_reduce(purity);
    hash_reduce(derive_func);
  }
 
  static constexpr const char* _type_key = "relax.FuncStructInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(FuncStructInfoNode, StructInfoNode);
};
 
class FuncStructInfo : public StructInfo {
 public:
  TVM_DLL FuncStructInfo(Array<StructInfo> params, StructInfo ret, bool purity = true,
                         Span span = Span());
 
  TVM_DLL static FuncStructInfo OpaqueFunc(StructInfoDeriveFunc derive_func, bool purity = false,
                                           Span span = Span());
 
  TVM_DLL static FuncStructInfo OpaqueFunc(StructInfo ret = ObjectStructInfo(), bool purity = false,
                                           Span span = Span());
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FuncStructInfo, StructInfo, FuncStructInfoNode);
};
 
template <typename T>
inline Optional<T> MatchStructInfo(const Expr& expr) {
  using TNode = typename T::ContainerType;
  if (const TNode* ptr = expr->struct_info_.as<TNode>()) {
    return GetRef<T>(ptr);
  } else {
    return NullOpt;
  }
}
 
template <typename T>
inline const T* GetStructInfoAs(const Expr& expr) {
  ICHECK(expr->struct_info_.defined())
      << "The struct_info is not populated, check if you have normalized the expr";
  return expr->struct_info_.as<T>();
}
 
inline StructInfo GetStructInfo(const Expr& expr) {
  auto* ptr = expr->struct_info_.as<StructInfoNode>();
  ICHECK(ptr) << "The struct_info is not populated, check if you have normalized the expr";
  return GetRef<StructInfo>(ptr);
}
 
inline bool HasVoidStructInfo(const Expr& expr) {
  auto* ptr = expr->struct_info_.as<TupleStructInfoNode>();
  return ptr != nullptr && ptr->fields.size() == 0;
}
 
TVM_DLL void UpdateStructInfo(Expr expr, StructInfo struct_info);
 
}  // namespace relax
}  // namespace tvm
#endif  // TVM_RELAX_STRUCT_INFO_H_