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/arith/bound.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>
 
#include <string>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace tvm {
namespace te {
 
using arith::IntSet;
using namespace tvm::tir;
 
// internal node container for Operation
class OperationNode;
class Tensor;
 
class Operation : public ObjectRef {
 public:
  Operation() {}
  explicit Operation(ObjectPtr<Object> n) : ObjectRef(n) {}
  explicit Operation(ffi::UnsafeInit tag) : ObjectRef(tag) {}
  inline const OperationNode* operator->() const;
  TVM_DLL Tensor output(size_t i) const;
  using ContainerType = OperationNode;
};
 
class TensorNode : public DataProducerNode {
 public:
  ffi::Array<PrimExpr> shape;
  DataType dtype;
  Operation op;
  int value_index{0};
 
  static void RegisterReflection();
 
  ffi::Array<PrimExpr> GetShape() const final { return shape; }
 
  DataType GetDataType() const final { return dtype; }
 
  TVM_DLL PrimExpr ToPrimExpr() const final;
 
  TVM_DLL ffi::String GetNameHint() const final;
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Tensor", TensorNode, DataProducerNode);
};
 
class Tensor : public DataProducer {
 private:
  inline PrimExpr IndexTensor(ffi::Array<PrimExpr> indices, bool support_negative_indices) const;
 
 public:
  TVM_DLL Tensor(ffi::Array<PrimExpr> shape, DataType dtype, 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<Var> 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<Var> 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}); }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tensor, DataProducer, TensorNode);
};
 
// 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.defined() || other->op.defined()) {
    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::ObjectPtrHash {};
 
template <>
struct hash<::tvm::te::Tensor> {
  std::size_t operator()(const ::tvm::te::Tensor& k) const {
    ::tvm::ObjectPtrHash hasher;
    if (k.defined() && k->op.defined()) {
      return hasher(k->op);
    } else {
      return hasher(k);
    }
  }
};
}  // namespace std
#endif  // TVM_TE_TENSOR_H_