expr.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_EXPR_H_
#define TVM_TIRX_EXPR_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/type.h>
 
#include <string>
 
namespace tvm {
// forward declare shared statement.
class Stmt;
 
namespace tirx {
 
class TensorVar : public Var {
 public:
  TVM_TIRX_DLL explicit TensorVar(ffi::String name, TensorType type, Location loc = UnknownLoc());
 
  explicit TensorVar(Var var) : Var(std::move(var)) {
    TVM_FFI_ICHECK(get() != nullptr && get()->ty.as<TensorTypeNode>())
        << "Expected a non-null Var with TensorType";
  }
 
  Var var() const { return ffi::GetRef<Var>(get()); }
 
  TensorType type() const { return get()->ty.as_or_throw<TensorType>(); }
 
  const ffi::String& name() const { return get()->name; }
 
  const Location& loc() const { return get()->loc; }
 
  Expr data() const { return Call(type()->DataPtrType(), tensor_data_ptr_op(), {var()}); }
 
  ffi::String scope() const { return (*this)->storage_scope; }
 
  TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
 
  PrimType ElementType() const { return (*this)->ElementType(); }
 
  explicit TensorVar(ffi::ObjectPtr<VarNode> node) : Var(std::move(node)) {}
  explicit TensorVar(ffi::UnsafeInit tag) : Var(tag) {}
  TVM_FFI_DEFINE_DEFAULT_COPY_MOVE_AND_ASSIGN(TensorVar);
 
  const TensorTypeNode* operator->() const {
    const auto* var_node = static_cast<const VarNode*>(data_.get());
    TVM_FFI_ICHECK(var_node != nullptr);
    const auto* type_node = var_node->ty.as<TensorTypeNode>();
    TVM_FFI_ICHECK(type_node != nullptr)
        << "Expected a Var with TensorType, but " << var_node->name << " has type " << var_node->ty;
    return type_node;
  }
 
  const VarNode* get() const { return static_cast<const VarNode*>(data_.get()); }
 
  [[maybe_unused]] static constexpr bool _type_is_nullable = false;
  static constexpr bool _type_container_is_exact = false;
  using ContainerType = VarNode;
};
 
// Preserve ObjectRef-style identity comparison for exact TensorVar operands.
// Comparisons widened to Var or Expr continue to build symbolic expressions.
inline bool operator==(const TensorVar& lhs, const TensorVar& rhs) { return lhs.same_as(rhs); }
 
inline bool operator!=(const TensorVar& lhs, const TensorVar& rhs) { return !lhs.same_as(rhs); }
 
inline TensorVar GetTensorVar(const VarNode* var) { return TensorVar(ffi::GetRef<Var>(var)); }
 
inline ffi::ObjectPtr<TensorTypeNode> CopyTensorType(const TensorVar& var) {
  return ffi::make_object<TensorTypeNode>(*var.operator->());
}
 
inline TensorVar RebuildTensorVar(const TensorVar& var, ffi::ObjectPtr<TensorTypeNode> type,
                                  ffi::Optional<ffi::String> name = std::nullopt) {
  return TensorVar(name.value_or(var.name()), TensorType(std::move(type)), var.loc());
}
 
TVM_TIRX_DLL TensorVar decl_tensor(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
                                   ffi::String name = "tensor", ffi::String storage_scope = "",
                                   Location loc = UnknownLoc());
 
TVM_TIRX_DLL tirx::TensorVar TensorWithOffsetAlignment(ffi::Array<PrimExpr> shape, PrimType dtype,
                                                       std::string name, int data_alignment,
                                                       int offset_factor,
                                                       std::string memory_scope = "");
 
TVM_TIRX_DLL TensorLoad MakeTensorLoad(TensorVar tensor, ffi::Array<PrimExpr> indices,
                                       Location loc = UnknownLoc());
 
TVM_TIRX_DLL TensorRegion FullTensorRegion(TensorVar tensor);
TVM_TIRX_DLL TensorRegion TensorRegionFromPoint(TensorVar tensor, ffi::Array<PrimExpr> indices);
 
}  // namespace tirx
}  // namespace tvm
 
namespace tvm::ffi {
 
template <>
inline constexpr bool use_default_type_traits_v<tirx::TensorVar> = false;
 
template <>
struct TypeTraits<tirx::TensorVar> : public ObjectRefTypeTraitsBase<tirx::TensorVar> {
  using Base = ObjectRefTypeTraitsBase<tirx::TensorVar>;
  using Base::CopyFromAnyViewAfterCheck;
  using Base::CopyToAnyView;
  using Base::GetMismatchTypeInfo;
  using Base::MoveFromAnyAfterCheck;
  using Base::MoveToAny;
  using Base::TypeSchema;
  using Base::TypeStr;
 
  TVM_FFI_INLINE static bool CheckAnyStrict(const TVMFFIAny* src) {
    if (src->type_index == TypeIndex::kTVMFFINone) {
      return false;
    }
    if (src->type_index != tvm::VarNode::RuntimeTypeIndex()) {
      return false;
    }
    const auto* var = static_cast<const tvm::VarNode*>(
        details::ObjectUnsafe::ObjectPtrFromUnowned<Object>(src->v_obj).get());
    return details::AnyUnsafe::CheckAnyStrict<tirx::TensorType>(var->ExprNode::ty);
  }
 
  TVM_FFI_INLINE static std::optional<tirx::TensorVar> TryCastFromAnyView(const TVMFFIAny* src) {
    if (CheckAnyStrict(src)) {
      return details::ObjectUnsafe::ObjectRefFromObjectPtr<tirx::TensorVar>(
          details::ObjectUnsafe::ObjectPtrFromUnowned<tvm::VarNode>(src->v_obj));
    }
    return std::nullopt;
  }
};
 
}  // namespace tvm::ffi
 
#endif  // TVM_TIRX_EXPR_H_