operation.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_OPERATION_H_
#define TVM_TE_OPERATION_H_
 
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/ir/prim/op.h>
#include <tvm/sym/analyzer.h>
#include <tvm/te/reduction.h>
#include <tvm/te/tensor.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/stmt.h>
 
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
namespace tvm {
namespace te {
 
// Reduce operator
class CommReducerNode : public ffi::Object {
 public:
  ffi::Array<PrimVar> lhs;
  ffi::Array<PrimVar> rhs;
  ffi::Array<PrimExpr> result;
  ffi::Array<PrimExpr> identity_element;
  ffi::Array<PrimExpr> operator()(ffi::Array<PrimExpr> a, ffi::Array<PrimExpr> b) const;
  mutable Location loc = UnknownLoc();
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<CommReducerNode>()
        .def_ro("lhs", &CommReducerNode::lhs, refl::AttachFieldFlag::SEqHashDefPattern())
        .def_ro("rhs", &CommReducerNode::rhs, refl::AttachFieldFlag::SEqHashDefPattern())
        .def_ro("result", &CommReducerNode::result)
        .def_ro("identity_element", &CommReducerNode::identity_element)
        .def_ro("loc", &CommReducerNode::loc, refl::AttachFieldFlag::SEqHashIgnore());
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.CommReducer", CommReducerNode, ffi::Object);
};
 
class CommReducer : public ffi::ObjectRef {
 public:
  TVM_DLL CommReducer(ffi::Array<PrimVar> lhs, ffi::Array<PrimVar> rhs, ffi::Array<PrimExpr> result,
                      ffi::Array<PrimExpr> identity_element, Location loc = UnknownLoc());
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CommReducer, ffi::ObjectRef, CommReducerNode);
};
 
class ReduceNode : public OpaqueExprNode {
 public:
  explicit ReduceNode(CommReducer combiner, PrimExpr condition)
      : combiner(std::move(combiner)), condition(std::move(condition)) {}
  explicit ReduceNode(ffi::UnsafeInit)
      : combiner(ffi::UnsafeInit{}), condition(ffi::UnsafeInit{}) {}
 
  CommReducer combiner;
  ffi::Array<PrimExpr> source;
  ffi::Array<PrimExpr> init;
  ffi::Array<s_tir::IterVar> axis;
  PrimExpr condition;
  int value_index;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ReduceNode>()
        .def_ro("combiner", &ReduceNode::combiner)
        .def_ro("source", &ReduceNode::source)
        .def_ro("init", &ReduceNode::init)
        .def_ro("axis", &ReduceNode::axis)
        .def_ro("condition", &ReduceNode::condition)
        .def_ro("value_index", &ReduceNode::value_index);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.Reduce", ReduceNode, OpaqueExprNode);
};
 
class Reduce : public PrimExpr {
 public:
  TVM_DLL Reduce(CommReducer combiner, ffi::Array<PrimExpr> src, ffi::Array<s_tir::IterVar> rdom,
                 ffi::Optional<PrimExpr> condition, int value_index, ffi::Array<PrimExpr> init,
                 Location loc = UnknownLoc());
  explicit Reduce(ffi::ObjectPtr<ReduceNode> node) : PrimExpr(std::move(node)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Reduce, PrimExpr, ReduceNode);
  static constexpr bool _type_container_is_exact = true;
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ReduceNode);
};
 
struct TensorDom {
  // constructor
  explicit TensorDom(int ndim) : data(ndim) {}
  std::vector<std::vector<IntSet>> data;
};
 
class TVM_DLL OperationNode : public ffi::Object {
 public:
  std::string name;
  std::string tag;
  ffi::Map<ffi::String, ffi::Any> attrs;
  // virtual destructor.
  virtual ~OperationNode() {}
  virtual int num_outputs() const = 0;
  virtual PrimType output_dtype(size_t i) const = 0;
  virtual ffi::Array<PrimExpr> output_shape(size_t i) const = 0;
  virtual ffi::Array<Tensor> InputTensors() const = 0;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OperationNode>()
        .def_ro("name", &OperationNode::name)
        .def_ro("tag", &OperationNode::tag)
        .def_ro("attrs", &OperationNode::attrs);
  }
  TVM_FFI_DECLARE_OBJECT_INFO("te.Operation", OperationNode, ffi::Object);
};
 
class PlaceholderOpNode : public OperationNode {
 public:
  ffi::Array<PrimExpr> shape;
  PrimType dtype = PrimType::Void();
  // override behavior.
  int num_outputs() const final;
  PrimType output_dtype(size_t i) const final;
  ffi::Array<PrimExpr> output_shape(size_t i) const final;
  ffi::Array<Tensor> InputTensors() const final;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PlaceholderOpNode>()
        .def_ro("shape", &PlaceholderOpNode::shape)
        .def_ro("dtype", &PlaceholderOpNode::dtype);
  }
  TVM_FFI_DECLARE_OBJECT_INFO("te.PlaceholderOp", PlaceholderOpNode, OperationNode);
};
 
class PlaceholderOp : public Operation {
 public:
  TVM_DLL PlaceholderOp(std::string name, ffi::Array<PrimExpr> shape, PrimType dtype);
 
  explicit PlaceholderOp(ffi::UnsafeInit tag, ffi::ObjectPtr<PlaceholderOpNode> node)
      : Operation(tag) {
    data_ = std::move(node);
  }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PlaceholderOp, Operation, PlaceholderOpNode);
};
 
class TVM_DLL BaseComputeOpNode : public OperationNode {
 public:
  ffi::Array<IterVar> axis;
  ffi::Array<IterVar> reduce_axis;
  // override functions
  ffi::Array<PrimExpr> output_shape(size_t idx) const final;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BaseComputeOpNode>()
        .def_ro("axis", &BaseComputeOpNode::axis)
        .def_ro("reduce_axis", &BaseComputeOpNode::reduce_axis);
  }
  TVM_FFI_DECLARE_OBJECT_INFO("te.BaseComputeOp", BaseComputeOpNode, OperationNode);
};
 
class TVM_DLL ComputeOpNode : public BaseComputeOpNode {
 public:
  ffi::Array<PrimExpr> body;
  ComputeOpNode() {}
  // override functions
  int num_outputs() const final;
  PrimType output_dtype(size_t i) const final;
  ffi::Array<Tensor> InputTensors() const final;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ComputeOpNode>().def_ro("body", &ComputeOpNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ComputeOp", ComputeOpNode, BaseComputeOpNode);
};
 
class ComputeOp : public Operation {
 public:
  TVM_DLL ComputeOp(std::string name, std::string tag, ffi::Map<ffi::String, ffi::Any> attrs,
                    ffi::Array<IterVar> axis, ffi::Array<PrimExpr> body);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComputeOp, Operation, ComputeOpNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeOpNode);
};
 
class ScanOpNode : public OperationNode {
 public:
  IterVar scan_axis{ffi::UnsafeInit{}};
  ffi::Array<Tensor> init;
  ffi::Array<Tensor> update;
  ffi::Array<Tensor> state_placeholder;
  ffi::Array<Tensor> inputs;
  ffi::Array<IterVar> spatial_axis_;
  ScanOpNode() {}
  // override behavior.
  int num_outputs() const final;
  PrimType output_dtype(size_t i) const final;
  ffi::Array<PrimExpr> output_shape(size_t i) const final;
  ffi::Array<Tensor> InputTensors() const final;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ScanOpNode>()
        .def_ro("scan_axis", &ScanOpNode::scan_axis)
        .def_ro("init", &ScanOpNode::init)
        .def_ro("update", &ScanOpNode::update)
        .def_ro("state_placeholder", &ScanOpNode::state_placeholder)
        .def_ro("inputs", &ScanOpNode::inputs)
        .def_ro("spatial_axis_", &ScanOpNode::spatial_axis_);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ScanOp", ScanOpNode, OperationNode);
};
 
class ScanOp : public Operation {
 public:
  TVM_DLL ScanOp(std::string name, std::string tag,
                 ffi::Optional<ffi::Map<ffi::String, ffi::Any>> attrs, IterVar axis,
                 ffi::Array<Tensor> init, ffi::Array<Tensor> update,
                 ffi::Array<Tensor> state_placeholder, ffi::Array<Tensor> input);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ScanOp, Operation, ScanOpNode);
};
 
class ExternOpNode : public OperationNode {
 public:
  explicit ExternOpNode(SeqStmt body) : body(std::move(body)) {}
  explicit ExternOpNode(ffi::UnsafeInit) : body(ffi::UnsafeInit{}) {}
 
  ffi::Array<Tensor> inputs;
  ffi::Array<TensorVar> input_placeholders;
  ffi::Array<TensorVar> output_placeholders;
  SeqStmt body;
 
  // override functions
  int num_outputs() const final;
  PrimType output_dtype(size_t i) const final;
  ffi::Array<PrimExpr> output_shape(size_t i) const final;
  ffi::Array<Tensor> InputTensors() const final;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExternOpNode>()
        .def_ro("inputs", &ExternOpNode::inputs)
        .def_ro("input_placeholders", &ExternOpNode::input_placeholders)
        .def_ro("output_placeholders", &ExternOpNode::output_placeholders)
        .def_ro("body", &ExternOpNode::body);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("te.ExternOp", ExternOpNode, OperationNode);
};
 
class ExternOp : public Operation {
 public:
  TVM_DLL ExternOp(std::string name, std::string tag, ffi::Map<ffi::String, ffi::Any> attrs,
                   ffi::Array<Tensor> inputs, ffi::Array<TensorVar> input_placeholders,
                   ffi::Array<TensorVar> output_placeholders, SeqStmt body);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExternOp, Operation, ExternOpNode);
};
 
TVM_DLL PrimVar var(std::string name_hint, PrimType t = PrimType::Int(32));
 
TVM_DLL IterVar thread_axis(Range dom, std::string tag);
 
TVM_DLL IterVar reduce_axis(Range dom, std::string name = "rv");
 
using FCompute = std::function<PrimExpr(const ffi::Array<PrimVar>& i)>;
 
using FBatchCompute = std::function<ffi::Array<PrimExpr>(const ffi::Array<PrimVar>& i)>;
 
TVM_DLL Tensor placeholder(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
                           std::string name = "placeholder");
 
TVM_DLL Tensor compute(ffi::Array<PrimExpr> shape, FCompute fcompute, std::string name = "tensor",
                       std::string tag = "", ffi::Map<ffi::String, ffi::Any> attrs = {});
 
TVM_DLL ffi::Array<Tensor> compute(ffi::Array<PrimExpr> shape, FBatchCompute fcompute,
                                   std::string name = "tensor", std::string tag = "",
                                   ffi::Map<ffi::String, ffi::Any> attrs = {});
 
TVM_DLL ffi::Array<Tensor> scan(ffi::Array<Tensor> init, ffi::Array<Tensor> update,
                                ffi::Array<Tensor> state_placeholder,
                                ffi::Array<Tensor> inputs = ffi::Array<Tensor>(),
                                std::string name = "scan", std::string tag = "",
                                ffi::Map<ffi::String, ffi::Any> attrs = {});
 
// same as compute, specialized for different fcompute function
inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(PrimVar)> f,
                      std::string name = "tensor", std::string tag = "",
                      ffi::Map<ffi::String, ffi::Any> attrs = {}) {
  FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0]); };
  return compute(shape, fc, name, tag, attrs);
}
inline Tensor compute(ffi::Array<PrimExpr> shape, std::function<PrimExpr(PrimVar, PrimVar)> f,
                      std::string name = "tensor", std::string tag = "",
                      ffi::Map<ffi::String, ffi::Any> attrs = {}) {
  FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0], i[1]); };
  return compute(shape, fc, name, tag, attrs);
}
inline Tensor compute(ffi::Array<PrimExpr> shape,
                      std::function<PrimExpr(PrimVar, PrimVar, PrimVar)> f,
                      std::string name = "tensor", std::string tag = "",
                      ffi::Map<ffi::String, ffi::Any> attrs = {}) {
  FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0], i[1], i[2]); };
  return compute(shape, fc, name, tag, attrs);
}
inline Tensor compute(ffi::Array<PrimExpr> shape,
                      std::function<PrimExpr(PrimVar, PrimVar, PrimVar, PrimVar)> f,
                      std::string name = "tensor", std::string tag = "",
                      ffi::Map<ffi::String, ffi::Any> attrs = {}) {
  FCompute fc = [f](const ffi::Array<PrimVar>& i) { return f(i[0], i[1], i[2], i[3]); };
  return compute(shape, fc, name, tag, attrs);
}
 
// inline function.
inline const OperationNode* Operation::operator->() const {
  return static_cast<const OperationNode*>(get());
}
}  // namespace te
 
namespace ffi {
 
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, te::ReduceNode> = true;
 
}  // namespace ffi
}  // namespace tvm
#endif  // TVM_TE_OPERATION_H_