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_TIRX_TYPE_H_
#define TVM_TIRX_TYPE_H_
 
#include <tvm/ir/expr.h>
#include <tvm/ir/type.h>
#include <tvm/tirx/layout.h>
 
#include <optional>
 
namespace tvm::tirx {
 
#ifndef TVM_INDEX_DEFAULT_I64
#define TVM_INDEX_DEFAULT_I64 1
#endif
inline PrimType DefaultIndexPrimType() {
#if TVM_INDEX_DEFAULT_I64
  static const PrimType default_index_ty = PrimType::Int(64);
#else
  static const PrimType default_index_ty = PrimType::Int(32);
#endif
  return default_index_ty;
}
 
inline DLDataType DefaultIndexType() {
#if TVM_INDEX_DEFAULT_I64
  return DLDataType{kDLInt, 64, 1};
#else
  return DLDataType{kDLInt, 32, 1};
#endif
}
 
class TensorTypeNode : public TypeNode {
 public:
  explicit TensorTypeNode(ffi::UnsafeInit tag) : elem_offset(tag) {}
 
  explicit TensorTypeNode(PrimExpr elem_offset) : elem_offset(std::move(elem_offset)) {}
 
  PrimType dtype = PrimType::Void();
  ffi::String storage_scope;
  ffi::Array<PrimExpr> shape;
  ffi::Array<PrimExpr> strides;
  PrimExpr elem_offset;
  int data_alignment;
  int offset_factor;
  ffi::Optional<Layout> layout;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TensorTypeNode>()
        .def_ro("dtype", &TensorTypeNode::dtype)
        .def_ro("storage_scope", &TensorTypeNode::storage_scope)
        .def_ro("shape", &TensorTypeNode::shape, refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("strides", &TensorTypeNode::strides, refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("elem_offset", &TensorTypeNode::elem_offset,
                refl::AttachFieldFlag::SEqHashDefSimple())
        .def_ro("data_alignment", &TensorTypeNode::data_alignment)
        .def_ro("offset_factor", &TensorTypeNode::offset_factor)
        .def_ro("layout", &TensorTypeNode::layout);
  }
 
  DLDataType DefaultIndexType() const {
    return shape.size() != 0 ? shape[0].ty()->dtype : tvm::tirx::DefaultIndexType();
  }
 
  PrimType ElementType() const { return dtype; }
 
  PtrType DataPtrType() const { return PtrType(dtype, storage_scope); }
 
  TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
 
  TVM_TIRX_DLL std::optional<int64_t> ConstantAllocationSize() const;
 
  ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index, bool inner = false) const;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorType", TensorTypeNode, TypeNode);
};
 
class TensorType : public Type {
 public:
  TVM_TIRX_DLL TensorType(ffi::String storage_scope, PrimType dtype, ffi::Array<PrimExpr> shape,
                          ffi::Array<PrimExpr> strides, ffi::Optional<PrimExpr> elem_offset,
                          int data_alignment, int offset_factor,
                          ffi::Optional<Layout> layout = std::nullopt, Location loc = UnknownLoc());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type, TensorTypeNode);
 
  explicit TensorType(ffi::ObjectPtr<TensorTypeNode> n) : Type(ffi::UnsafeInit{}) {
    TVM_FFI_ICHECK(n != nullptr);
    data_ = std::move(n);
  }
};
 
class TensorMapTypeNode : public TypeNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TensorMapTypeNode>();
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TensorMapType", TensorMapTypeNode, TypeNode);
};
 
class TensorMapType : public Type {
 public:
  TVM_TIRX_DLL TensorMapType(Location loc = UnknownLoc());
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorMapType, Type, TensorMapTypeNode);
};
 
}  // namespace tvm::tirx
#endif  // TVM_TIRX_TYPE_H_