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>
 
#include <utility>
 
namespace tvm {
namespace relax {
 
class ObjectStructInfoNode : public StructInfoNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ObjectStructInfoNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ObjectStructInfo", ObjectStructInfoNode, StructInfoNode);
};
 
class ObjectStructInfo : public StructInfo {
 public:
  TVM_DLL ObjectStructInfo(Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ObjectStructInfo, StructInfo, ObjectStructInfoNode);
};
 
class PrimStructInfoNode : public StructInfoNode {
 public:
  ffi::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);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PrimStructInfo", 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_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimStructInfo, StructInfo, PrimStructInfoNode);
};
 
class ShapeStructInfoNode : public StructInfoNode {
 public:
  ffi::Optional<ffi::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);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeStructInfo", ShapeStructInfoNode, StructInfoNode);
};
 
class ShapeStructInfo : public StructInfo {
 public:
  TVM_DLL ShapeStructInfo(ffi::Array<PrimExpr> values, Span span = Span());
  TVM_DLL ShapeStructInfo(int ndim, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeStructInfo, StructInfo, ShapeStructInfoNode);
};
 
class TensorStructInfoNode : public StructInfoNode {
 public:
  ffi::Optional<Expr> shape;
  ffi::Optional<VDevice> vdevice;
  DataType dtype;
  int ndim;
 
  bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
 
  bool IsUnknownDtype() const { return dtype.is_void(); }
 
  ffi::Optional<ffi::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);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorStructInfo", TensorStructInfoNode, StructInfoNode);
};
 
class TensorStructInfo : public StructInfo {
 public:
  TVM_DLL TensorStructInfo(Expr shape, DataType dtype,
                           ffi::Optional<VDevice> vdevice = std::nullopt, Span span = Span());
 
  TVM_DLL TensorStructInfo(DataType dtype, int ndim, ffi::Optional<VDevice> vdevice = std::nullopt,
                           Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TensorStructInfo, StructInfo, TensorStructInfoNode);
};
 
class TupleStructInfoNode : public StructInfoNode {
 public:
  ffi::Array<StructInfo> fields;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TupleStructInfoNode>().def_ro("fields", &TupleStructInfoNode::fields);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TupleStructInfo", TupleStructInfoNode, StructInfoNode);
};
 
class TupleStructInfo : public StructInfo {
 public:
  TVM_DLL TupleStructInfo(ffi::Array<StructInfo> fields, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TupleStructInfo, StructInfo, TupleStructInfoNode);
};
 
using StructInfoDeriveFunc = TypedEnvFunc<StructInfo(const Call& call, const BlockBuilder& ctx)>;
 
class FuncStructInfoNode : public StructInfoNode {
 public:
  ffi::Optional<ffi::Array<StructInfo>> params;
  StructInfo ret;
  ffi::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);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncStructInfo", FuncStructInfoNode, StructInfoNode);
};
 
class FuncStructInfo : public StructInfo {
 public:
  explicit FuncStructInfo(ObjectPtr<FuncStructInfoNode> data) : StructInfo(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_DLL FuncStructInfo(ffi::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_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncStructInfo, StructInfo, FuncStructInfoNode);
};
 
template <typename T>
inline ffi::Optional<T> MatchStructInfo(const Expr& expr) {
  using TNode = typename T::ContainerType;
  if (const TNode* ptr = expr->struct_info_.as<TNode>()) {
    return ffi::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 ffi::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_