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/ffi/reflection/registry.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:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ObjectStructInfoNode>();
  }
 
  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;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PrimStructInfoNode>()
        .def_ro("value", &PrimStructInfoNode::value)
        .def_ro("dtype", &PrimStructInfoNode::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; }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ShapeStructInfoNode>()
        .def_ro("values", &ShapeStructInfoNode::values)
        .def_ro("ndim", &ShapeStructInfoNode::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;
  }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TensorStructInfoNode>()
        .def_ro("shape", &TensorStructInfoNode::shape)
        .def_ro("dtype", &TensorStructInfoNode::dtype)
        .def_ro("vdevice", &TensorStructInfoNode::vdevice)
        .def_ro("ndim", &TensorStructInfoNode::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 = std::nullopt,
                           Span span = Span());
 
  TVM_DLL TensorStructInfo(DataType dtype, int ndim, Optional<VDevice> vdevice = std::nullopt,
                           Span span = Span());
 
  TVM_DEFINE_OBJECT_REF_METHODS(TensorStructInfo, StructInfo, TensorStructInfoNode);
};
 
class TupleStructInfoNode : public StructInfoNode {
 public:
  Array<StructInfo> fields;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TupleStructInfoNode>().def_ro("fields", &TupleStructInfoNode::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(); }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FuncStructInfoNode>()
        .def_ro("params", &FuncStructInfoNode::params, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("ret", &FuncStructInfoNode::ret)
        .def_ro("derive_func", &FuncStructInfoNode::derive_func)
        .def_ro("purity", &FuncStructInfoNode::purity);
  }
 
  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 std::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_