tensor.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_TE_TENSOR_H_
#define TVM_TE_TENSOR_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/s_tir/iter_var.h>
#include <tvm/sym/bound.h>
#include <tvm/tirx/expr.h>
 
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace tvm {
namespace te {
using s_tir::IterVar;
using s_tir::IterVarNode;
using s_tir::IterVarType;
using s_tir::IterVarType2String;
using s_tir::kCommReduce;
using s_tir::kDataPar;
using s_tir::kOpaque;
using s_tir::kOrdered;
using s_tir::kParallelized;
using s_tir::kTensorized;
using s_tir::kThreadIndex;
using s_tir::kUnrolled;
using s_tir::kVectorized;
using sym::IntSet;
 
using namespace tvm::tirx;
 
// internal node container for Operation
class OperationNode;
class Tensor;
 
class Operation : public ffi::ObjectRef {
 public:
  static constexpr bool _type_is_nullable = false;
  explicit Operation(ffi::UnsafeInit tag) : ffi::ObjectRef(tag) {}
  inline const OperationNode* operator->() const;
  TVM_DLL Tensor output(size_t i) const;
  using ContainerType = OperationNode;
};
 
class TensorNode : public OpaqueExprNode {
 public:
  ffi::Array<PrimExpr> shape;
  PrimType dtype = PrimType::Void();
  ffi::Optional<Operation> op;
  int value_index{0};
 
  static void RegisterReflection();
 
  ffi::Array<PrimExpr> GetShape() const { return shape; }
 
  PrimType GetDataType() const { return dtype; }
 
  TVM_DLL ffi::String GetNameHint() const;
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorNode, OpaqueExprNode);
};
 
class Tensor : public OpaqueExpr {
 private:
  inline PrimExpr IndexTensor(ffi::Array<PrimExpr> indices, bool support_negative_indices) const;
 
 public:
  TVM_DLL Tensor(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::Optional<Operation> op,
                 int value_index);
 
  inline bool operator==(const Tensor& other) const;
  inline bool operator!=(const Tensor& other) const;
  inline size_t ndim() const;
  template <typename... Args>
  inline PrimExpr operator()(Args&&... args) const {
    ffi::Array<PrimExpr> indices{std::forward<Args>(args)...};
    return operator()(indices);
  }
  TVM_DLL PrimExpr operator()(ffi::Array<PrimExpr> indices) const;
  TVM_DLL PrimExpr operator()(ffi::Array<PrimVar> indices) const;
  template <typename... Args>
  TVM_DLL PrimExpr IndexWithNegativeIndices(Args&&... args) const {
    ffi::Array<PrimExpr> indices{std::forward<Args>(args)...};
    return IndexWithNegativeIndices(indices);
  }
  TVM_DLL PrimExpr IndexWithNegativeIndices(ffi::Array<PrimExpr> indices) const;
  TVM_DLL PrimExpr IndexWithNegativeIndices(ffi::Array<PrimVar> indices) const;
 
  class Slice {
   public:
    // construct via tensor and indices
    Slice(const Tensor& tensor, std::vector<PrimExpr> indices)
        : tensor_(tensor), indices_(indices) {}
    inline Slice operator[](PrimExpr i) {
      std::vector<PrimExpr> other = indices_;
      other.emplace_back(i);
      return Slice(tensor_, other);
    }
    inline operator PrimExpr() const { return tensor_(indices_); }
 
   private:
    const Tensor& tensor_;
    std::vector<PrimExpr> indices_;
  };
  inline Slice operator[](PrimExpr i) const { return Slice(*this, {i}); }
 
  explicit Tensor(ffi::ObjectPtr<TensorNode> node) : OpaqueExpr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Tensor, OpaqueExpr, TensorNode);
};
 
TVM_DLL bool IsTensorLoad(const Expr& expr);
 
TVM_DLL Tensor GetTensorFromLoad(const Call& call);
 
TVM_DLL ffi::Array<PrimExpr> GetTensorLoadIndices(const Call& call);
 
// Implementations of inline functions
inline size_t Tensor::ndim() const { return (*this)->shape.size(); }
 
inline bool Tensor::operator==(const Tensor& other) const {
  if (get() == other.get()) return true;
  if (get() == nullptr || other.get() == nullptr) return false;
  if ((*this)->op.has_value() || other->op.has_value()) {
    return (*this)->op == other->op && (*this)->value_index == other->value_index;
  } else {
    return false;
  }
}
 
inline bool Tensor::operator!=(const Tensor& other) const { return !(*this == other); }
 
// macro to turn every operation of slice to expression
#define DEFINE_OVERLOAD_SLICE_UNARY_OP(Op) \
  inline PrimExpr operator Op(const Tensor::Slice& a) { return Op a.operator PrimExpr(); }
 
#define DEFINE_OVERLOAD_SLICE_BINARY_OP(Op)                                     \
  template <typename T>                                                         \
  inline PrimExpr operator Op(const Tensor::Slice& a, const T& b) {             \
    return a.operator PrimExpr() Op b;                                          \
  }                                                                             \
  template <typename T>                                                         \
  inline PrimExpr operator Op(const T& a, const Tensor::Slice& b) {             \
    return a Op b.operator PrimExpr();                                          \
  }                                                                             \
  inline PrimExpr operator Op(const Tensor::Slice& a, const Tensor::Slice& b) { \
    return a.operator PrimExpr() Op b.operator PrimExpr();                      \
  }
 
DEFINE_OVERLOAD_SLICE_UNARY_OP(!);
DEFINE_OVERLOAD_SLICE_UNARY_OP(-);
DEFINE_OVERLOAD_SLICE_BINARY_OP(+);
DEFINE_OVERLOAD_SLICE_BINARY_OP(-);
DEFINE_OVERLOAD_SLICE_BINARY_OP(*);
DEFINE_OVERLOAD_SLICE_BINARY_OP(==);
DEFINE_OVERLOAD_SLICE_BINARY_OP(<=);
DEFINE_OVERLOAD_SLICE_BINARY_OP(>=);
DEFINE_OVERLOAD_SLICE_BINARY_OP(!=);
DEFINE_OVERLOAD_SLICE_BINARY_OP(&&);
DEFINE_OVERLOAD_SLICE_BINARY_OP(||);
DEFINE_OVERLOAD_SLICE_BINARY_OP(>>);
DEFINE_OVERLOAD_SLICE_BINARY_OP(<<);
DEFINE_OVERLOAD_SLICE_BINARY_OP(>);  // NOLINT(*)
DEFINE_OVERLOAD_SLICE_BINARY_OP(<);  // NOLINT(*)
 
}  // namespace te
}  // namespace tvm
 
namespace std {
template <>
struct hash<::tvm::te::Operation> : public ::tvm::ffi::ObjectPtrHash {};
 
template <>
struct hash<::tvm::te::Tensor> {
  std::size_t operator()(const ::tvm::te::Tensor& k) const {
    ::tvm::ffi::ObjectPtrHash hasher;
    if (k.defined() && k->op.has_value()) {
      return hasher(k->op.value());
    } else {
      return hasher(k);
    }
  }
};
}  // namespace std
#endif  // TVM_TE_TENSOR_H_