type.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_TYPE_H_
#define TVM_RELAX_TYPE_H_
 
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
#include <tvm/ir/env_func.h>
#include <tvm/ir/global_info.h>
#include <tvm/ir/type.h>
#include <tvm/tirx/expr.h>
 
#include <string>
#include <utility>
 
namespace tvm {
namespace relax {
 
using Expr = tvm::Expr;
using ExprNode = tvm::ExprNode;
 
class BlockBuilder;
 
static constexpr int kUnknownNDim = -1;
 
using tvm::TupleType;
using tvm::TupleTypeNode;
 
class PackedFuncTypeNode : public TypeNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PackedFuncTypeNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.PackedFuncType", PackedFuncTypeNode, TypeNode);
};
 
class PackedFuncType : public Type {
 public:
  TVM_DLL PackedFuncType(Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PackedFuncType, Type, PackedFuncTypeNode);
};
 
class AnyTypeNode : public TypeNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AnyTypeNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.AnyType", AnyTypeNode, TypeNode);
};
 
class AnyType : public Type {
 public:
  TVM_DLL AnyType(Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AnyType, Type, AnyTypeNode);
};
 
// Compatibility aliases for existing C++ callers.  New code should use AnyType.
using ObjectTypeNode = AnyTypeNode;
using ObjectType = AnyType;
 
class ShapeTypeNode : public TypeNode {
 public:
  ffi::Optional<ffi::Array<PrimExpr>> values;
  int ndim{kUnknownNDim};
 
  bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ShapeTypeNode>()
        .def_ro("values", &ShapeTypeNode::values)
        .def_ro("ndim", &ShapeTypeNode::ndim);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.ShapeType", ShapeTypeNode, TypeNode);
};
 
class ShapeType : public Type {
 public:
  TVM_DLL ShapeType(ffi::Array<PrimExpr> values, Span span = Span());
  TVM_DLL ShapeType(int ndim, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ShapeType, Type, ShapeTypeNode);
};
 
class TensorTypeNode : public TypeNode {
 public:
  ffi::Optional<Expr> shape;
  ffi::Optional<VDevice> vdevice;
  ffi::Optional<tvm::PrimType> dtype{std::nullopt};
  int ndim{kUnknownNDim};
 
  bool IsUnknownNdim() const { return ndim == kUnknownNDim; }
 
  bool IsUnknownDtype() const { return !dtype.has_value(); }
 
  ffi::Optional<ffi::Array<PrimExpr>> GetShape() const {
    if (!shape.has_value()) return {};
    const Expr& shape_expr = this->shape.value();
    if (shape_expr->ty.IsMissing()) return {};
    if (const auto* shape_ty = shape_expr->ty.as<ShapeTypeNode>()) {
      return shape_ty->values;
    }
    return {};
  }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TensorTypeNode>()
        .def_ro("shape", &TensorTypeNode::shape)
        .def_ro("dtype", &TensorTypeNode::dtype)
        .def_ro("vdevice", &TensorTypeNode::vdevice)
        .def_ro("ndim", &TensorTypeNode::ndim);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.TensorType", TensorTypeNode, TypeNode);
};
 
class TensorType : public Type {
 public:
  explicit TensorType(ffi::ObjectPtr<TensorTypeNode> data) : Type(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
 
  TVM_DLL TensorType(Expr shape, ffi::Optional<tvm::PrimType> dtype = std::nullopt,
                     ffi::Optional<VDevice> vdevice = std::nullopt, Span span = Span());
 
  TVM_DLL TensorType(ffi::Optional<tvm::PrimType> dtype, int ndim,
                     ffi::Optional<VDevice> vdevice = std::nullopt, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode);
};
 
using TypeDeriveFunc = TypedEnvFunc<Type(const Call& call, const BlockBuilder& ctx)>;
 
class FuncTypeNode : public TypeNode {
 public:
  ffi::Optional<ffi::Array<Type>> params;
  Type ret = Type::Missing();
  ffi::Optional<TypeDeriveFunc> derive_func;
  bool purity;
 
  bool IsOpaque() const { return !params.has_value(); }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FuncTypeNode>()
        .def_ro("params", &FuncTypeNode::params, refl::AttachFieldFlag::SEqHashDefRecursive())
        .def_ro("ret", &FuncTypeNode::ret)
        .def_ro("derive_func", &FuncTypeNode::derive_func)
        .def_ro("purity", &FuncTypeNode::purity);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FuncType", FuncTypeNode, TypeNode);
};
 
class FuncType : public Type {
 public:
  explicit FuncType(ffi::ObjectPtr<FuncTypeNode> data) : Type(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(data != nullptr);
    data_ = std::move(data);
  }
  TVM_DLL FuncType(ffi::Array<Type> params, Type ret, bool purity = true, Span span = Span());
 
  TVM_DLL static FuncType OpaqueFunc(TypeDeriveFunc derive_func, bool purity = false,
                                     Span span = Span());
 
  TVM_DLL static FuncType OpaqueFunc(Type ret = AnyType(), bool purity = false, Span span = Span());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FuncType, Type, FuncTypeNode);
};
 
template <typename T>
inline ffi::Optional<T> MatchType(const Expr& expr) {
  if (!expr.defined()) {
    return std::nullopt;
  }
  using TNode = typename T::ContainerType;
  if (const TNode* ptr = expr->ty.as<TNode>()) {
    return ffi::GetRef<T>(ptr);
  } else {
    return std::nullopt;
  }
}
 
template <typename T>
inline const T* GetTypeAs(const Expr& expr) {
  TVM_FFI_ICHECK(!expr->ty.IsMissing())
      << "The type is not populated, check if you have normalized the expr";
  return expr->ty.as<T>();
}
 
inline Type GetType(const Expr& expr) {
  TVM_FFI_ICHECK(!expr->ty.IsMissing())
      << "The type is not populated, check if you have normalized the expr";
  return expr->ty;
}
 
inline bool HasVoidType(const Expr& expr) {
  auto* ptr = expr->ty.as<TupleTypeNode>();
  return ptr != nullptr && ptr->fields.size() == 0;
}
 
TVM_DLL void UpdateType(Expr expr, Type ty);
 
}  // namespace relax
}  // namespace tvm
 
#endif  // TVM_RELAX_TYPE_H_