You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2025/02/17 16:01:12 UTC
(tvm) 01/09: cleanup schedule c++
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit e443752d7d9e2a2cfac54bb4c809f8aee5a5f76e
Author: Siyuan Feng <hz...@sjtu.edu.cn>
AuthorDate: Mon Feb 17 14:15:16 2025 +0800
cleanup schedule c++
---
include/tvm/driver/driver_api.h | 40 -
include/tvm/relax/analysis.h | 1 +
include/tvm/relax/op_attr_types.h | 1 -
include/tvm/te/autodiff.h | 96 --
include/tvm/te/operation.h | 236 +---
include/tvm/te/schedule.h | 1010 ----------------
include/tvm/te/schedule_pass.h | 112 --
include/tvm/te/tensor_intrin.h | 145 ---
include/tvm/topi/cuda/dense.h | 154 ---
include/tvm/topi/cuda/injective.h | 83 --
include/tvm/topi/cuda/pooling.h | 187 ---
include/tvm/topi/cuda/reduction.h | 199 ----
include/tvm/topi/cuda/softmax.h | 103 --
include/tvm/topi/generic/default.h | 83 --
include/tvm/topi/generic/extern.h | 69 --
include/tvm/topi/generic/injective.h | 77 --
include/tvm/topi/rocm/dense.h | 99 --
include/tvm/topi/rocm/injective.h | 67 --
include/tvm/topi/rocm/pooling.h | 68 --
include/tvm/topi/rocm/softmax.h | 55 -
include/tvm/topi/x86/bnn.h | 131 ---
include/tvm/topi/x86/default.h | 105 --
include/tvm/topi/x86/injective.h | 85 --
python/tvm/te/__init__.py | 2 +-
python/tvm/te/operation.py | 24 +-
python/tvm/te/tensor.py | 10 -
src/driver/driver_api.cc | 80 --
src/relax/analysis/layout_transformation.cc | 1 +
src/target/codegen.cc | 11 -
src/target/llvm/codegen_llvm.cc | 4 +-
src/te/autodiff/ad_simplify.cc | 1239 --------------------
src/te/autodiff/ad_utils.cc | 197 ----
src/te/autodiff/ad_utils.h | 135 ---
src/te/autodiff/adjoint.cc | 157 ---
src/te/autodiff/jacobian.cc | 365 ------
src/te/operation/compute_op.cc | 375 ------
src/te/operation/compute_op.h | 104 --
src/te/operation/create_primfunc.cc | 3 +-
src/te/operation/cross_thread_reduction.cc | 237 ----
src/te/operation/extern_op.cc | 84 --
src/te/operation/graph.cc | 90 ++
src/te/operation/graph.h | 62 +
src/te/operation/hybrid_op.cc | 486 --------
src/te/operation/hybrid_op.h | 94 --
src/te/operation/op_utils.cc | 290 -----
src/te/operation/op_utils.h | 96 --
src/te/operation/placeholder_op.cc | 27 -
src/te/operation/scan_op.cc | 146 ---
src/te/operation/tensor_compute_op.cc | 259 ----
src/te/operation/tensorize.cc | 496 --------
src/te/schedule/auto_inline_elem_wise.cc | 125 --
src/te/schedule/bound.cc | 262 -----
src/te/schedule/graph.cc | 418 -------
src/te/schedule/graph.h | 130 --
src/te/schedule/message_passing.cc | 744 ------------
src/te/schedule/message_passing.h | 121 --
src/te/schedule/operation_inline.cc | 90 --
src/te/schedule/operation_inline.h | 48 -
src/te/schedule/schedule_dataflow_rewrite.cc | 978 ---------------
src/te/schedule/schedule_lang.cc | 1078 -----------------
src/te/schedule/schedule_ops.cc | 437 -------
src/te/schedule/schedule_postproc_to_primfunc.cc | 447 -------
src/te/schedule/verify_compact_buffer.cc | 63 -
src/te/tensor.cc | 61 -
src/tir/transforms/storage_flatten.cc | 1 +
src/topi/schedule.cc | 319 -----
.../tvm/topi/rocm/reduction.h => src/topi/utils.cc | 38 +-
tests/cpp/build_module_test.cc | 67 --
68 files changed, 179 insertions(+), 13528 deletions(-)
diff --git a/include/tvm/driver/driver_api.h b/include/tvm/driver/driver_api.h
index fffcab4966..eaf737088b 100644
--- a/include/tvm/driver/driver_api.h
+++ b/include/tvm/driver/driver_api.h
@@ -35,14 +35,9 @@
#include <tvm/runtime/packed_func.h>
#include <tvm/support/with.h>
#include <tvm/target/target.h>
-#include <tvm/te/schedule_pass.h>
#include <tvm/tir/function.h>
#include <string>
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-#include <vector>
namespace tvm {
using tvm::transform::Pass;
@@ -105,41 +100,6 @@ TVM_DLL IRModule LowerPrimFunc(tvm::tir::PrimFunc func, const std::string& name,
* \return The result module.
*/
-TVM_DLL IRModule LowerSchedule(te::Schedule sch, const Array<te::Tensor>& args,
- const std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds,
- GlobalVarSupply global_var_supply, bool simple_mode = false);
-
-/*!
- * \brief Build an IRModule given a TE schedule, args and binds. This function also applies
- * the lowering passes defined in CreatePassList.
- * \param sch The TE schedule to lower.
- * \param args The arguments to the function (Array of Tensor, Buffer and Vars)
- * \param name The name of the lowered function.
- * \param binds Buffer assignments.
- * \param global_var_supply The GlobalVarSupply to be used in the module.
- * \param simple_mode Disables the loop partition pass. Defaults to false.
- * \return The result module.
- */
-TVM_DLL IRModule LowerSchedule(te::Schedule sch, const Array<ObjectRef>& args,
- const std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds,
- GlobalVarSupply global_var_supply, bool simple_mode = false);
-
-/*!
- * \brief Create an IRModule out of a TE Schedule. It does not apply lowering passes. If you want
- * to apply lowering passes as well, use LowerSchedule.
- * \param sch The schedule
- * \param args The arguments to the function.
- * \param name The name of the lowered function.
- * \param binds Buffer assignments.
- * \param global_var_supply The GlobalVarSupply to be used in the module and when creating
- * GlobalVars.
- * \return The result module.
- */
-IRModule ScheduleToModule(te::Schedule sch, const Array<ObjectRef>& args, const std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds,
- GlobalVarSupply global_var_supply);
/*!
* \brief Build a device and host module for a specific target from an IRModule.
* \param funcs The functions to be built.
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index 2de2f4fd36..b658758e3c 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -31,6 +31,7 @@
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/struct_info.h>
#include <tvm/tir/function.h>
+#include <tvm/tir/index_map.h>
#include <functional>
#include <utility>
diff --git a/include/tvm/relax/op_attr_types.h b/include/tvm/relax/op_attr_types.h
index 434a89a288..e5dd1aaba1 100644
--- a/include/tvm/relax/op_attr_types.h
+++ b/include/tvm/relax/op_attr_types.h
@@ -26,7 +26,6 @@
#include <tvm/relax/expr.h>
#include <tvm/relax/struct_info.h>
-#include <tvm/te/schedule.h>
#include <tvm/te/tensor.h>
namespace tvm {
diff --git a/include/tvm/te/autodiff.h b/include/tvm/te/autodiff.h
deleted file mode 100644
index e2d379969c..0000000000
--- a/include/tvm/te/autodiff.h
+++ /dev/null
@@ -1,96 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/te/autodiff.h
- * \brief Automatic differentiation of tensor expressions.
- */
-
-#ifndef TVM_TE_AUTODIFF_H_
-#define TVM_TE_AUTODIFF_H_
-
-#include <tvm/runtime/object.h>
-#include <tvm/tir/expr.h>
-
-#include "tensor.h"
-
-namespace tvm {
-/*! \brief Tensor expression language DSL. */
-namespace te {
-
-/*!
- * \brief Take the derivative of the expression with respect to the given variable.
- * \param expr The expression to differentiate.
- * \param var The variable to differentiate with respect to.
- * \return The expression for the derivative.
- */
-PrimExpr Derivative(const PrimExpr& expr, const Var& var);
-
-/*!
- * \brief Get the tensor representing the Jacobian of the output with respect to the input.
- *
- * Note that if \p output depends on \p input indirectly (by using some other tensor
- * depending on \p input), this dependency won't contribute to the resulting Jacobian.
- * For such cases use the function ::Gradient.
- *
- * \param output The tensor to differentiate.
- * \param input The input tensor, which \p output should directly use.
- * \return The tensor representing the Jacobian of shape `output.shape + input.shape`.
- */
-Tensor Jacobian(const Tensor& output, const Tensor& input);
-
-/*!
- * \brief The building block for reverse-mode AD.
- *
- * Differentiate \p output wrt \p input and multiply the result by \p head on the left using tensor
- * dot product. \p input must be an immediate dependency of \p output (must be called from within
- * the body of \p output). That is, the function will compute one summand of the adjoint for \p
- * input given the adjoint for \p output (which is called \p head here).
- *
- * \param output The tensor to differentiate.
- * \param input The input tensor, which \p output should directly use.
- * \param head The adjoint of \p output. Must be of shape `prefix + output.shape`
- * \return The tensor of shape `prefix + input.shape`
- * representing the partial adjoint of \p input wrt one of its consumers (output)
- */
-Tensor VectorJacobianProduct(const Tensor& output, const Tensor& input, const Tensor& head);
-
-/*!
- * \brief Perform reverse mode automatic differentiation.
- *
- * Each item of the `result` field of the result is an adjoint for the corresponding item of
- * \p inputs, i.e. \p head multiplied by the Jacobian of \p output with respect to the
- * corresponding item of \p inputs.
- *
- * \param output The tensor to differentiate.
- * \param inputs The array of input tensors. When the array is empty, will perform differentiation
- * wrt all tensors the output depends on.
- * \param head The adjoint of the output, in other words, some tensor, by which the Jacobians
- * will be multiplied (using tensordot axes=`output.shape`).
- * Its shape must be of the form `prefix + output.shape`. If the null pointer is
- * provided, the identity tensor of shape `output.shape + output.shape` will be used. \return An
- * array of adjoints corresponding to \p inputs.
- */
-TVM_DLL Array<Tensor> Gradient(const Tensor& output, const Array<Tensor>& inputs,
- const Tensor& head = Tensor());
-
-} // namespace te
-} // namespace tvm
-
-#endif // TVM_TE_AUTODIFF_H_
diff --git a/include/tvm/te/operation.h b/include/tvm/te/operation.h
index f5753afa56..4b3ce69918 100644
--- a/include/tvm/te/operation.h
+++ b/include/tvm/te/operation.h
@@ -25,7 +25,6 @@
#define TVM_TE_OPERATION_H_
#include <tvm/arith/analyzer.h>
-#include <tvm/te/schedule.h>
#include <tvm/te/tensor.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/expr.h>
@@ -65,11 +64,6 @@ class TVM_DLL OperationNode : public Object {
virtual ~OperationNode() {}
/*! \return number of outputs */
virtual int num_outputs() const = 0;
- /*!
- * \return The list of iteration variable at root
- * \note root_iter_vars decides the shape of the outputs.
- */
- virtual Array<IterVar> root_iter_vars() const = 0;
/*!
* \brief Get data type. i-th output tensor.
* \param i The output index.
@@ -87,59 +81,6 @@ class TVM_DLL OperationNode : public Object {
* \return List of input tensors.
*/
virtual Array<Tensor> InputTensors() const = 0;
- /*!
- * \brief Replace the input of the operation by pattern specified by rmap.
- *
- * \param self The reference to self.
- * \param rmap The replacement map.
- * \return self if nothing is replaced, otherwise return replaced op.
- */
- virtual Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const = 0;
- /*!
- * \brief Propagate the bounds to inputs
- * \param self The reference to self.
- * \param analyzer The analyzer to be used in the function.
- * \param dom_map the domain map of Variables(corresponds to root_iter_vars)
- * \param out_dom_map The output domain.
- * The function is only asked to fill the bounds for Tensors that
- * is already in the out_dom_map
- */
- virtual void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const = 0;
- /*!
- * \brief Gather the bound from output tensor.
- * Set the range of each root_iter_vars in the op to out_dom_map
- *
- * \param self The reference to self.
- * \param tensor_dom Domain map of Tensor->access set of each dimension.
- * \param out_dom_map The output domain map of each IterVar to be setted.
- */
- virtual void GatherBound(const Operation& self,
- const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const = 0;
- /*!
- * \brief Build the Realize statement that realizes
- * the op's output tensors.
- * \param stage the op's stage.
- * \param realize_map The realization domain map of the operators.
- * \param body The body that is going to get
- * \param storage_scope The storage scope associated with this realization
- * \return A realization statement that wraps body.
- */
- virtual Stmt BuildRealize(const Stage& stage,
- const std::unordered_map<IterVar, Range>& realize_map, const Stmt& body,
- String storage_scope = "") const = 0;
- /*!
- * \brief Build the statement that provide the output tensors.
- * \param stage The schedule stage of the op.
- * \param dom_map The domain map of all iteration domains.
- * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1
- * \return A statement that add production and wraps consumer.
- */
- virtual Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const = 0;
static constexpr const char* _type_key = "Operation";
@@ -157,21 +98,9 @@ class PlaceholderOpNode : public OperationNode {
DataType dtype;
// override behavior.
int num_outputs() const final;
- Array<IterVar> root_iter_vars() const final;
DataType output_dtype(size_t i) const final;
Array<PrimExpr> output_shape(size_t i) const final;
Array<Tensor> InputTensors() const final;
- Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const final;
- void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
- void GatherBound(const Operation& self, const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const final;
- Stmt BuildRealize(const Stage& stage, const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope = "") const final;
- Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const final;
void VisitAttrs(AttrVisitor* v) {
v->Visit("name", &name);
@@ -198,8 +127,7 @@ class PlaceholderOp : public Operation {
/*!
* \brief A Compute op that compute a tensor on certain domain.
- * This is the base class for ComputeOp (operating on a scalar at a time) and
- * TensorComputeOp (operating on a TensorSlice at a time)
+ * This is the base class for ComputeOp (operating on a scalar at a time)
*/
class TVM_DLL BaseComputeOpNode : public OperationNode {
public:
@@ -208,13 +136,7 @@ class TVM_DLL BaseComputeOpNode : public OperationNode {
/*! \brief IterVar on each reduction axis, if the body is a Reduce */
Array<IterVar> reduce_axis;
// override functions
- Array<IterVar> root_iter_vars() const final;
Array<PrimExpr> output_shape(size_t idx) const final;
- void GatherBound(const Operation& self, const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const final;
- Stmt BuildRealize(const Stage& stage, const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope = "") const final;
- virtual size_t num_schedulable_dims() const = 0;
static constexpr const char* _type_key = "BaseComputeOp";
TVM_DECLARE_BASE_OBJECT_INFO(BaseComputeOpNode, OperationNode);
@@ -233,14 +155,6 @@ class TVM_DLL ComputeOpNode : public BaseComputeOpNode {
int num_outputs() const final;
DataType output_dtype(size_t i) const final;
Array<Tensor> InputTensors() const final;
- Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const final;
- void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
- Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const final;
- size_t num_schedulable_dims() const final;
void VisitAttrs(AttrVisitor* v) {
v->Visit("name", &name);
@@ -268,66 +182,6 @@ class ComputeOp : public Operation {
TVM_DEFINE_OBJECT_REF_COW_METHOD(ComputeOpNode);
};
-/*!
- * \brief A TenorCompute op that compute a tensor with an tensor intrinsic.
- */
-class TensorComputeOpNode : public BaseComputeOpNode {
- public:
- /*! \brief number of axes that can be scheduled */
- int schedulable_ndim;
- /*! \brief TensorIntrin used to compute */
- TensorIntrin intrin;
- /*! \brief input tensors of intrin */
- Array<Tensor> inputs;
- /*! \brief region of input tensors */
- Array<Region> input_regions;
- /*! \brief scalar expression inputs */
- Array<PrimExpr> scalar_inputs;
- /*! \brief constructor */
- TensorComputeOpNode() {}
- // override functions
- int num_outputs() const final;
- DataType output_dtype(size_t i) const final;
- Array<Tensor> InputTensors() const final;
- Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const final;
- void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
- Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const final;
- size_t num_schedulable_dims() const final;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("name", &name);
- v->Visit("tag", &tag);
- v->Visit("axis", &axis);
- v->Visit("reduce_axis", &reduce_axis);
- v->Visit("schedulable_ndim", &schedulable_ndim);
- v->Visit("intrin", &intrin);
- v->Visit("inputs", &inputs);
- v->Visit("input_regions", &input_regions);
- v->Visit("scalar_inputs", &scalar_inputs);
- }
-
- static constexpr const char* _type_key = "TensorComputeOp";
- TVM_DECLARE_FINAL_OBJECT_INFO(TensorComputeOpNode, BaseComputeOpNode);
-};
-
-/*!
- * \brief Managed reference to TensorComputeOpNode
- * \sa TensorComputeOpNode
- */
-class TensorComputeOp : public Operation {
- public:
- TVM_DLL TensorComputeOp(std::string name, std::string tag, Array<IterVar> axis,
- Array<IterVar> reduce_axis, int schedulable_ndim, TensorIntrin intrin,
- Array<Tensor> tensors, Array<Region> regions,
- Array<PrimExpr> scalar_inputs);
-
- TVM_DEFINE_OBJECT_REF_METHODS(TensorComputeOp, Operation, TensorComputeOpNode);
-};
-
/*!
* \brief Symbolic scan.
*/
@@ -360,21 +214,9 @@ class ScanOpNode : public OperationNode {
ScanOpNode() {}
// override behavior.
int num_outputs() const final;
- Array<IterVar> root_iter_vars() const final;
DataType output_dtype(size_t i) const final;
Array<PrimExpr> output_shape(size_t i) const final;
Array<Tensor> InputTensors() const final;
- Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const final;
- void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
- void GatherBound(const Operation& self, const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const final;
- Stmt BuildRealize(const Stage& stage, const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope = "") const final;
- Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const final;
void VisitAttrs(AttrVisitor* v) {
v->Visit("name", &name);
@@ -423,21 +265,9 @@ class ExternOpNode : public OperationNode {
ExternOpNode() {}
// override functions
int num_outputs() const final;
- Array<IterVar> root_iter_vars() const final;
DataType output_dtype(size_t i) const final;
Array<PrimExpr> output_shape(size_t i) const final;
Array<Tensor> InputTensors() const final;
- Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const final;
- void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
- void GatherBound(const Operation& self, const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const final;
- Stmt BuildRealize(const Stage& stage, const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope = "") const final;
- Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const final;
void VisitAttrs(AttrVisitor* v) {
v->Visit("name", &name);
@@ -466,70 +296,6 @@ class ExternOp : public Operation {
TVM_DEFINE_OBJECT_REF_METHODS(ExternOp, Operation, ExternOpNode);
};
-/*!
- * \brief A computation operator that generated by hybrid script.
- */
-class HybridOpNode : public OperationNode {
- public:
- /*! \brief The input tensors */
- Array<Tensor> inputs;
- /*! \brief Symbolic placeholder representation of outputs */
- Array<Tensor> outputs;
- /*! \brief The axis of iterations */
- Array<IterVar> axis;
- /*! \brief the statement that generates the computation. This is
- * slightly different from the body in ExternOpNode. All the output
- * tensors keep its own name specified by users in the script.
- * However, when compilation, these tensors will be placed by those
- * actual output tensors. */
- Stmt body;
-
- /*! \brief constructor */
- HybridOpNode() {}
- // override functions
- int num_outputs() const final;
- Array<IterVar> root_iter_vars() const final;
- DataType output_dtype(size_t i) const final;
- Array<PrimExpr> output_shape(size_t i) const final;
- Array<Tensor> InputTensors() const final;
- Operation ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const final;
- void PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const final;
- void GatherBound(const Operation& self, const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const final;
- Stmt BuildRealize(const Stage& stage, const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope = "") const final;
- Stmt BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const final;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("name", &name);
- v->Visit("tag", &tag);
- v->Visit("attrs", &attrs);
- v->Visit("inputs", &inputs);
- v->Visit("outputs", &outputs);
- v->Visit("axis", &axis);
- v->Visit("body", &body);
- }
-
- static constexpr const char* _type_key = "HybridOp";
- TVM_DECLARE_FINAL_OBJECT_INFO(HybridOpNode, OperationNode);
-};
-
-/*!
- * \brief Managed reference to HybridOpNode
- * \sa HybridOpNode
- */
-class HybridOp : public Operation {
- public:
- TVM_DLL HybridOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
- Array<Tensor> inputs, Array<Tensor> outputs, Stmt body);
-
- TVM_DEFINE_OBJECT_REF_METHODS(HybridOp, Operation, HybridOpNode);
-};
-
/*!
* \brief Construct a new Var expression
* \param name_hint The name hint for the expression
diff --git a/include/tvm/te/schedule.h b/include/tvm/te/schedule.h
deleted file mode 100644
index 47787b2c99..0000000000
--- a/include/tvm/te/schedule.h
+++ /dev/null
@@ -1,1010 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/te/schedule.h
- * \brief Define a schedule.
- */
-// Acknowledgement: Many schedule primitives originate from Halide and Loopy.
-#ifndef TVM_TE_SCHEDULE_H_
-#define TVM_TE_SCHEDULE_H_
-
-#include <tvm/support/with.h>
-#include <tvm/te/tensor.h>
-#include <tvm/te/tensor_intrin.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/index_map.h>
-
-#include <string>
-#include <unordered_map>
-
-namespace tvm {
-namespace te {
-// Node container for Stage
-class StageNode;
-// Node container for Schedule
-class ScheduleNode;
-// Node container for IterVarRelation
-class IterVarRelationNode;
-// Attribute of itervar.
-class IterVarAttrNode;
-
-/*! \brief the attachment type */
-enum AttachType : int {
- kGroupRoot = 1,
- kInline = 2,
- kInlinedAlready = 3,
- kScope = 4,
- kScanUpdate = 5
-};
-
-/*! \brief Stage, contains scheduling for a stage of computation. */
-class Stage : public ObjectRef {
- public:
- Stage() {}
- explicit Stage(ObjectPtr<Object> n) : ObjectRef(n) {}
- /*!
- * \brief create a new schedule for op.
- * \param op The operator in the schedule
- * \param sch The schedule which current stage belongs to
- */
- explicit Stage(Operation op, const ScheduleNode* sch);
- /*!
- * \brief access the internal node container
- * \return the pointer to the internal node container
- */
- inline const StageNode* operator->() const;
- /*!
- * \brief access the internal node container
- * \return the pointer to the internal node container
- */
- inline StageNode* operator->();
- /*!
- * \brief set the memory scope of the stage
- * \param scope The memory scope.
- */
- TVM_DLL Stage& set_scope(std::string scope); // NOLINT(*)
- /*!
- * \brief specify the schedule to be computed at the parent schedule's scope.
- * \param parent The parent schedule.
- * \param scope The iteration point to carry the schedule.
- * \return reference to self.
- */
- TVM_DLL Stage& compute_at(Stage parent, IterVar scope); // NOLINT(*)
- /*!
- * \brief Compute the function inline.
- * \return reference to self.
- */
- TVM_DLL Stage& compute_inline(); // NOLINT(*)
- /*!
- * \brief Compute the function at group root.
- * \return reference to self.
- */
- TVM_DLL Stage& compute_root(); // NOLINT(*)
- /*!
- * \brief Bind the IterVar to thread index.
- *
- * \param ivar The IterVar to be bound.
- * \param thread_ivar The thread axis to be bound.
- * \return reference to self.
- */
- TVM_DLL Stage& bind(IterVar ivar, IterVar thread_ivar);
- /*!
- * \brief Set the predicate to determine whether a store to the array should be performed.
- * Use this when there are multiple threads performing the same store and we only
- * need one of them to do the store.
- *
- * \note This is a dangerous scheduling primitive that can change behavior of program.
- * Only do when we are certain that thare are duplicated stores.
- * \param predicate The condition to be checked.
- * \return reference to self.
- */
- TVM_DLL Stage& set_store_predicate(PrimExpr predicate);
- /*!
- * \brief Specify environment threads that launched around the group's scope.
- * This can only be used in group stage.
- * \param threads The threads to be launched around the scope.
- * \note Each thread can only appear in one env_threads.
- * This is a beta feature.
- * \return reference to self.
- */
- TVM_DLL Stage& env_threads(Array<IterVar> threads);
- /*!
- * \brief Split the parent by factor, generate
- * \param parent The parent iteration domain.
- * \param factor The split factor of the loop.
- * \param p_outer The result outer domain
- * \param p_inner The result inner domain.
- * \param disable_predication If enabled, don't create a predicate for guarding the
- * loop. This can be useful when splitting with scalable factors that the schedule writer
- * knows are divisible by the loop bound.
- * Warning: enabling this feature may result in incorrect code generation if not used carefully.
- * \return reference to self.
- */
- TVM_DLL Stage& split(IterVar parent, PrimExpr factor, IterVar* p_outer, IterVar* p_inner,
- bool disable_predication = false); // NOLINT(*)
- /*!
- * \brief Split the iteration with given number of parts.
- *
- * \param parent The parent domain.
- * \param nparts The number of parts in the outer domain.
- * \param p_outer The result outer domain.
- * \param p_inner The result inner domain.
- * \param disable_predication If enabled, don't create a predicate for guarding the
- * loop. This can be useful when splitting with scalable factors that the schedule writer
- * knows are divisible by the loop bound.
- * Warning: enabling this feature may result in incorrect code generation if not used carefully.
- * \return reference to self.
- */
- TVM_DLL Stage& split_by_nparts(IterVar parent, PrimExpr nparts, IterVar* p_outer,
- IterVar* p_inner, bool disable_predication = false); // NOLINT(*)
- /*!
- * \brief Fuse the inner outer domain to the target
- * \param outer The outer domain to be fused.
- * \param inner The inner domain to be fused
- * \param p_target The result target domain.
- * \return reference to self.
- */
- TVM_DLL Stage& fuse(IterVar outer, IterVar inner, IterVar* p_target); // NOLINT(*)
- /*!
- * \brief Fuse all the axes together into a single axis.
- *
- * \param axes All the axes to be fused.
- * \param p_target The result target domain.
- *
- * \note axes can be an empty array,
- * in that case, a singleton IterVar is created and
- * inserted to the outermost loop.
- * The fuse of empty array is used to support zero-dimension tensors.
- *
- * \return reference to self.
- */
- TVM_DLL Stage& fuse(const Array<IterVar>& axes, IterVar* p_target); // NOLINT(*)
- /*!
- * \brief Reorder the iteration
- * \param order The order of iteration variable.
- * \return reference to self.
- */
- TVM_DLL Stage& reorder(const Array<IterVar>& order); // NOLINT(*)
- /*!
- * \brief Perform tiling on two dimensions
- * The final loop order from outmost to inner most are
- * [x_outer, y_outer, x_inner, y_inner]
- *
- * \param x_parent The original x dimension
- * \param y_parent The original y dimension
- * \param x_factor The stride factor on x axis
- * \param y_factor The stride factor on y axis
- * \param p_x_outer Outer axis of x dimension
- * \param p_y_outer Outer axis of y dimension
- * \param p_x_inner Inner axis of x dimension
- * \param p_y_inner Inner axis of y dimension
- * \return reference to self.
- */
- TVM_DLL Stage& tile(IterVar x_parent, IterVar y_parent, // NOLINT(*)
- PrimExpr x_factor, PrimExpr y_factor, IterVar* p_x_outer, IterVar* p_y_outer,
- IterVar* p_x_inner, IterVar* p_y_inner);
- /*!
- * \brief Vectorize iteration.
- * \param var The axis to be vectorized.
- * \return reference to self.
- */
- TVM_DLL Stage& vectorize(IterVar var); // NOLINT(*)
- /*!
- * \brief Replace computation of the current stage by tensor intrinsic f.
- * \param var The axis marks beginning of tensorization.
- * Every operations inside the axis(include axis itself is tensorized).
- * \param f The Tensor compute intrinsics.
- * \return reference to self.
- */
- TVM_DLL Stage& tensorize(IterVar var, TensorIntrin f); // NOLINT(*)
- /*!
- * \brief Unroll iteration.
- * \param var The axis to be unrolled.
- * \return reference to self.
- */
- TVM_DLL Stage& unroll(IterVar var); // NOLINT(*)
- /*!
- * \brief Parallelize iteration.
- * \param var The axis to be parallelized.
- * \return reference to self.
- */
- TVM_DLL Stage& parallel(IterVar var); // NOLINT(*)
- /*!
- * \brief Annotate the iteration with pragma
- *
- * \param var The axis to be parallelized.
- * \param pragma_type The pragma type.
- * \param pragma_value The pragma value
- *
- * \return reference to self.
- */
- TVM_DLL Stage& pragma(IterVar var, const std::string& pragma_type,
- const PrimExpr& pragma_value = PrimExpr()); // NOLINT(*)
- /*!
- * \brief Fetch data in advance.
- * \param domain the tensor to be prefetched
- * \param var the iteration point at which to apply prefetching
- * \param offset the number of iterations be to fetched in advance
- * \return reference to self
- */
- TVM_DLL Stage& prefetch(const Tensor& domain, IterVar var, PrimExpr offset); // NOLINT(*)
- /*!
- * \brief Set alignment requirement for specific dimension.
- *
- * Such that stride[axis] == k * factor + offset for some k.
- *
- * \param axis The dimension to be specified for alignment.
- * \param factor The factor multiple of alignment
- * \param offset The required offset factor.
- * \return reference to self
- */
- TVM_DLL Stage& storage_align(IterVar axis, int factor, int offset); // NOLINT(*)
- /*!
- * \brief Compute current stage with double buffering.
- * \return reference to self.
- */
- TVM_DLL Stage& double_buffer(); // NOLINT(*)
- /*!
- * \brief Compute current stage with rolling buffering.
- * \return reference to self.
- */
- TVM_DLL Stage& rolling_buffer(); // NOLINT(*)
- /*!
- * \brief Defines a layout transformation to be applied to the buffer.
- *
- * The map from initial_index to final_index must be an
- * invertible affine transformation.
- *
- * \param initial_indices An array of variables to represent a
- * value's location in the tensor, using the pre-transformation
- * layout. These variables are used as binding occurrences to
- * represent the initial indices when applying the initial->final
- * mapping, and should not occur elsewhere in the
- * Schedule. (i.e. Pass in newly constructed variables, not the
- * initial IterVar::var)
- *
- * \param final_indices An array of expressions, giving the
- * value's location in the tensor, using the post-transformation layout.
- * Expressions should be in terms of the variables given in
- * initial_indices.
- *
- * \param out_iter_vars An optional output location for the updated
- * loop iteration variables.
- *
- * \return reference to self
- */
- TVM_DLL Stage& transform_layout(const Array<Var>& initial_indices,
- const Array<PrimExpr>& final_indices,
- Array<IterVar>* out_iter_vars = nullptr);
- /*! \brief Defines separators between groups of axes.
- *
- * Used to define `BufferNode::axis_separators`, which has
- * additional details.
- *
- * \param axis_separators A list of axis separators.
- */
- TVM_DLL Stage& set_axis_separators(const Array<IntImm>& axis_separators);
- /*!
- * \brief whether the stage has been scheduled.
- * \return whether the stage has been scheduled.
- */
- bool is_scheduled() const;
- /*!
- * \brief Get attachment spec of current stage.
- * If the stage compute at Group root, this function
- * will traverse the group function to get the
- * final spec from the group.
- * \return A stage representing the attach spec of the group.
- */
- Stage GetAttachSpec() const;
- // declare container type
- using ContainerType = StageNode;
-};
-
-/*!
- * \brief Global schedule container
- * For operations and all the operations they depend on.
- * The schedule per Operation is named as stage.
- */
-class Schedule : public ObjectRef {
- public:
- Schedule() {}
- explicit Schedule(ObjectPtr<Object> n) : ObjectRef(n) {}
- /*!
- * \brief Create a schedule for array of ops(and their dependencies).
- * \param ops The ops to be scheduled.
- */
- TVM_DLL explicit Schedule(Array<Operation> ops);
- /*!
- * \brief Get a copy of current schedule.
- * \return The copied schedule.
- */
- Schedule copy() const;
- /*!
- * \brief Get the stage corresponds to the op
- * \param op The operation.
- */
- TVM_DLL Stage operator[](const Operation& op);
- /*!
- * \brief Short hand for getting the stage of tensor's operation.
- * \param tensor The tensor
- * \return The stage corresponding to the tensor's op
- */
- TVM_DLL Stage operator[](const Tensor& tensor) { return this->operator[](tensor->op); }
- /*!
- * \brief Create a new stage group for all intermediate
- * operations between inputs and outputs.
- *
- * \param outputs The output boundary of the group.
- * \param inputs The input boundary of the group.
- * \param include_inputs Whether include inputs if they are reachable from outputs.
- * \return The new grouped stage.
- */
- TVM_DLL Stage create_group(const Array<Tensor>& outputs, const Array<Tensor>& inputs,
- bool include_inputs = false);
- /*!
- * \brief create a cache read of original tensor for readers.
- * This will mutate the body of the readers.
- * A new stage will be created for the tensor.
- * \param tensor The tensor cached.
- * \param scope The scope of the cache.
- * \param readers The readers to redirect to the tensor.
- * \return The created tensor.
- */
- TVM_DLL Tensor cache_read(const Tensor& tensor, const std::string& scope,
- const Array<Operation>& readers);
- /*!
- * \brief Create a cache write tensor for producing tensor.
- * The tensor will take over body of original tensor op.
- *
- * This function can be used to do data layout transformation.
- * If there is a split/fuse/reorder on the data parallel axis of tensor
- * before cache_write is called. The intermediate cache stores
- * the data in the layout as the iteration order of leave axis.
- * The data will be transformed back to the original layout in the original tensor.
- * User can further call compute_inline to inline the original layout and keep
- * the data stored in the transformed layout.
- *
- * \param tensor The tensors to be produced.
- * \param scope The scope of the storage.
- * \return The created tensor.
- */
- TVM_DLL Array<Tensor> cache_write(const Array<Tensor>& tensor, const std::string& scope);
- /*!
- * \brief Create a cache write tensor for producing tensor.
- * The tensor will take over body of original tensor op.
- *
- * This function can be used to do data layout transformation.
- * If there is a split/fuse/reorder on the data parallel axis of tensor
- * before cache_write is called. The intermediate cache stores
- * the data in the layout as the iteration order of leave axis.
- * The data will be transformed back to the original layout in the original tensor.
- * User can further call compute_inline to inline the original layout and keep
- * the data stored in the transformed layout.
- *
- * \param tensor The tensor to be produced.
- * \param scope The scope of the storage.
- * \return The created tensor.
- */
- TVM_DLL Tensor cache_write(const Tensor& tensor, const std::string& scope);
- /*!
- * \brief Factor a reduction axis in tensor's schedule to be an explicit axis.
- * This will create a new stage that generated the new tensor with axis
- * as the first dimension. The tensor's body will be rewritten as a reduction
- * over the factored tensor.
- *
- * P. Suriana, A. Adams and S. Kamil. Parallel associative reductions in halide. CGO'17
- *
- * \param tensor The tensor to be factored.
- * \param axis The reduction axis in tensor's schedule to be factored.
- * \param factor_axis The position where the new axis is placed.
- * \return The created factored tensors.
- */
- TVM_DLL Array<Tensor> rfactor(const Tensor& tensor, const IterVar& axis, int factor_axis = 0);
- /*!
- * \brief Normalize the schedule.
- * This is needed before bound inference.
- * Insert necessary RebaseNode to make sure all leaf_iter_vars
- * are in form [0, extent)
- *
- * \return A normalized schedule, can be same as current one.
- */
- Schedule normalize();
-
- /*!
- * \brief Normalize the schedule for feature extraction in auto-scheduler.
- * This is similar to `Schedule::normalize`, but we do aggressive simplification
- * to the TE compute with const_matrix=True for faster compilation and feature extraction.
- * The resulted schedule may be wrong, but it is good enough for feature extraction
- * purposes.
- *
- * \return A normalized schedule, can be same as current one.
- */
- Schedule normalize_for_feature_extraction();
-
- /*!
- * \brief access the internal node container
- * \return the pointer to the internal node container
- */
- inline const ScheduleNode* operator->() const;
- /*!
- * \brief access the internal node container
- * \return the pointer to the internal node container
- */
- inline ScheduleNode* operator->();
- // declare container type
- using ContainerType = ScheduleNode;
-};
-
-/*!
- * \brief Context helper to collect debug information of Schedule.
- *
- * Attach With<ScheduleContext>(schedule_instance, primitive_name)
- * inside function body of schedule primitives to collect the
- * snapshot of schedule status and corresponding primitive name
- */
-class ScheduleContext {
- private:
- friend class With<ScheduleContext>;
- ScheduleContext(const ScheduleNode* sch_node, String current_primitive_name);
- void EnterWithScope();
- void ExitWithScope();
-
- /*! \brief Schedule instance to store information for debug */
- Schedule sch_;
- /*! \brief String representing which primitive has been applied to sch_ */
- String current_primitive_name_;
-};
-
-/*!
- * \brief The schedule relation between IterVars
- * can be Split, Fuse.
- */
-class IterVarRelation : public ObjectRef {
- public:
- IterVarRelation() {}
- explicit IterVarRelation(ObjectPtr<Object> n) : ObjectRef(n) {}
- /*!
- * \brief access the internal node container
- * \return the pointer to the internal node container
- */
- inline const IterVarRelationNode* operator->() const;
-};
-
-/*!
- * \brief Additional scheduable attributes about IterVar.
- */
-class IterVarAttr : public ObjectRef {
- public:
- IterVarAttr() {}
- explicit IterVarAttr(ObjectPtr<Object> n) : ObjectRef(n) {}
- /*!
- * \brief access the internal node container
- * \return the pointer to the internal node container
- */
- inline const IterVarAttrNode* operator->() const;
-};
-
-/*!
- * \brief represents a stage.
- *
- * relations form a Directed acylic hypergraph in bipartite manner.
- * With each node is represented by a IterVar,
- * and each hyper-edge is represented by a IterVarRelation.
- * The relations connects the IterVars in the graph.
- *
- * Besides typical stage that corresponds to operations.
- * There is also group stage, which groups stages together.
- * Each stage's group(given by group) represent an constraint,
- * the stage can only be attached to stages within the group.
- *
- * The group stage node can be attached to IterVars as in normal stage.
- */
-class StageNode : public Object {
- public:
- /*!
- * \brief The operation of stage, can be different from original op.
- * If it is null, then this stage is a group stage.
- */
- Operation op;
- /*!
- * \brief The original operator.
- * The op field can change during schedule to alternate the dataflow,
- * while origin_op remains fixed.
- */
- Operation origin_op;
- /*! \brief All the nodes in the iter var
- *
- * Each element of all_iter_vars represents an iteration variable
- * that may appear within this stage's computation. Any element
- * of `all_iter_vars` that is in `leaf_iter_vars` represents a
- * variable that is directly defined and usable within the stage's
- * computation. All other elements of `all_iter_vars` represent
- * variables whose value must be computed from the variables in
- * `leaf_iter_vars`. (e.g. Support index k has been split by
- * ``ko, ki = s.split(k, factor=4)``. ko and ki will appear in
- * `leaf_iter_vars`, while k will not, and must be computed as
- * `4*ko + ki`.
- */
- Array<IterVar> all_iter_vars;
- /*! \brief The current active leaf iter vars in the stage.
- *
- * Each element of leaf_iter_vars will either be replaced with the
- * bound index (e.g. threadIdx.x), or will be expanded into a loop
- * over the variable's extent. `leaf_iter_vars` is a subset of
- * `all_iter_vars`.
- */
- Array<IterVar> leaf_iter_vars;
- /*!
- * \brief Specify threads to be launched at the stage.
- * This is only valid for composite ops such as Scan.
- * \note Experimental primitive: used for thread persistence.
- */
- Array<IterVar> env_threads;
- /*!
- * \brief The predicate under which store can happen
- * Use this when there can be duplicated threads doing the same store.
- * \note Experimental primitive: used by cross thread-reduction.
- */
- PrimExpr store_predicate;
- /*! \brief The relation bwteen of IterVars */
- Array<IterVarRelation> relations;
- /*! \brief additional attributes about iter var. */
- Map<IterVar, IterVarAttr> iter_var_attrs;
- /*! \brief The attachment type of the schedule */
- AttachType attach_type{kGroupRoot};
- /*! \brief The attach point of this schedule. */
- IterVar attach_ivar;
- /*! \brief The stage this node attaches to */
- Stage attach_stage;
- /*! \brief The schedule current stage is attached to */
- const ScheduleNode* attach_sch;
- /*! \brief The thread storage scope level of the stage */
- std::string scope;
- /*! \brief Whether this is an output stage */
- bool is_output{false};
- /*! \brief Whether apply double buffer optimization to this stage */
- bool double_buffer{false};
- /*! \brief Whether apply rolling buffer optimization to this stage */
- bool rolling_buffer{false};
- /*! \brief Layout transformations to be applied onto the stage's tensors. */
- Array<IndexMap> layout_transforms;
- /*! \brief List of axes after which to divide physical axes.
- *
- * Used to populate `BufferNode::axis_separators`, which has
- * additional details.
- */
- Array<IntImm> axis_separators;
- /*!
- * \brief The parent group of the current stage.
- * The stage cannot be assigned to stages outside the group.
- */
- Stage group;
- /*! \brief Number of direct child stages, only used for group stage.*/
- int num_child_stages{0};
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("op", &op);
- v->Visit("origin_op", &origin_op);
- v->Visit("all_iter_vars", &all_iter_vars);
- v->Visit("leaf_iter_vars", &leaf_iter_vars);
- v->Visit("env_threads", &env_threads);
- v->Visit("relations", &relations);
- v->Visit("iter_var_attrs", &iter_var_attrs);
- v->Visit("attach_type", &attach_type);
- v->Visit("attach_ivar", &attach_ivar);
- v->Visit("attach_stage", &attach_stage);
- v->Visit("scope", &scope);
- v->Visit("is_output", &is_output);
- v->Visit("double_buffer", &double_buffer);
- v->Visit("layout_transforms", &layout_transforms);
- v->Visit("axis_separators", &axis_separators);
- v->Visit("group", &group);
- v->Visit("num_child_stages", &num_child_stages);
- }
-
- static constexpr const char* _type_key = "Stage";
- TVM_DECLARE_FINAL_OBJECT_INFO(StageNode, Object);
-};
-
-/*! \brief node container for schedule */
-class ScheduleNode : public Object {
- public:
- /*! \brief The output operations in original data flow graph */
- Array<Operation> outputs;
- /*!
- * \brief list of all stages for ops.
- * The stages are sorted in dependency order.
- */
- Array<Stage> stages;
- /*!
- * \brief List of all stage groups.
- */
- Array<Stage> groups;
- /*! \brief map of original operation to the stages */
- Map<Operation, Stage> stage_map;
- /*!
- * \brief Internal stage map to map internal ops to stages.
- * This is created on demand and can be invalidated.
- */
- std::unordered_map<const Object*, Stage> op2stage_cache_;
- /*!
- * \brief list of all transformed schedules
- * User can display the optimization strategy via TEDD step by step to check
- * the order and effect of primitives. Set "te.keep_schedule_record" in
- * PassContext config as true to enable recording.
- */
- Array<Schedule> schedule_record;
- /*!
- * \brief list of all applied primitive names.
- */
- Array<String> primitive_record;
- /*!
- * \brief Flag to keep schedule record or not.
- */
- Optional<Bool> keep_schedule_record;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("outputs", &outputs);
- v->Visit("stages", &stages);
- v->Visit("groups", &groups);
- v->Visit("stage_map", &stage_map);
- v->Visit("schedule_record", &schedule_record);
- v->Visit("primitive_record", &primitive_record);
- v->Visit("keep_schedule_record", &keep_schedule_record);
- }
-
- /*! \brief Initialize temp cache. */
- void InitCache();
- /*! \brief Invalidate temp cache. */
- void InvalidateCache();
-
- /*!
- * \brief Check if the schedule contains an Operation.
- * \param op The candidate Operation.
- * \return true if the schedule has the Operation. Otherwise, false.
- */
- TVM_DLL bool Contain(const Operation& op) const;
-
- /*!
- * \brief Check if the schedule contains a Tensor.
- * \param tensor The candidate tensor.
- * \return true if the schedule has the tensor. Otherwise, false.
- */
- TVM_DLL bool Contain(const Tensor& tensor) const { return Contain(tensor->op); }
-
- static constexpr const char* _type_key = "Schedule";
- TVM_DECLARE_FINAL_OBJECT_INFO(ScheduleNode, Object);
-};
-
-/*!
- * \brief Create a schedule for array of ops(and their dependencies).
- * \param ops The ops to be scheduled.
- * \return sch The created Schedule.
- */
-inline Schedule create_schedule(Array<Operation> ops) { return Schedule(ops); }
-
-/*! \brief node container for IterVar attr */
-class IterVarAttrNode : public Object {
- public:
- /*! \brief The iteration type. */
- IterVarType iter_type{kDataPar};
- /*! \brief The thread this iter Var binds, can be null */
- IterVar bind_thread;
- /*! \brief List of tensor to be prefetched in this loop */
- Array<Tensor> prefetch_data;
- /*! \brief The offset used in each prefetch */
- Array<PrimExpr> prefetch_offset;
- /*!
- * \brief Tensor intrinsic used in tensorization,
- * when the axis is marked as Tensorized
- */
- TensorIntrin tensor_intrin;
- /*! \brief Alignment factor of buffer dimension */
- int dim_align_factor{0};
- /*! \brief Alignment offset of buffer dimension */
- int dim_align_offset{0};
- /*!
- * \brief Additional pragma keys, array of StringImm
- */
- Array<PrimExpr> pragma_keys;
- /*!
- * \brief Additional values of pragma, if any
- */
- Array<PrimExpr> pragma_values;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("iter_type", &iter_type);
- v->Visit("bind_thread", &bind_thread);
- v->Visit("prefetch_data", &prefetch_data);
- v->Visit("prefetch_offset", &prefetch_offset);
- v->Visit("tensor_intrin", &tensor_intrin);
- v->Visit("dim_align_factor", &dim_align_factor);
- v->Visit("dim_align_offset", &dim_align_offset);
- v->Visit("pragma_keys", &pragma_keys);
- v->Visit("pragma_values", &pragma_values);
- }
-
- static constexpr const char* _type_key = "IterVarAttr";
- TVM_DECLARE_FINAL_OBJECT_INFO(IterVarAttrNode, Object);
-};
-
-/*! \brief base node of iteration var */
-class IterVarRelationNode : public Object {
- public:
- static constexpr const char* _type_key = "IterVarRelation";
- TVM_DECLARE_BASE_OBJECT_INFO(IterVarRelationNode, Object);
-};
-
-/*!
- * \brief Split the parent domain into product of
- * outer and iter.
- */
-class SplitNode : public IterVarRelationNode {
- public:
- /*! \brief The parent domain */
- IterVar parent;
- /*! \brief The outer domain */
- IterVar outer;
- /*! \brief The inner domain */
- IterVar inner;
- /*! \brief The split factor */
- PrimExpr factor;
- /*! \brief Number of parts, only factor or nparts can be given */
- PrimExpr nparts;
- /*! \brief Whether to disable generation of predication. */
- bool disable_predication;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("parent", &parent);
- v->Visit("outer", &outer);
- v->Visit("inner", &inner);
- v->Visit("factor", &factor);
- v->Visit("nparts", &nparts);
- v->Visit("disable_predication", &disable_predication);
- }
-
- static constexpr const char* _type_key = "Split";
- TVM_DECLARE_FINAL_OBJECT_INFO(SplitNode, IterVarRelationNode);
-};
-
-/*!
- * \brief Managed reference to SplitNode
- * \sa SplitNode
- */
-class Split : public IterVarRelation {
- public:
- TVM_DLL Split(IterVar parent, IterVar outer, IterVar inner, PrimExpr factor, PrimExpr nparts,
- bool disable_predication);
-
- TVM_DEFINE_OBJECT_REF_METHODS(Split, IterVarRelation, SplitNode);
-};
-
-/*!
- * \brief Fuse two domains into one domain.
- */
-class FuseNode : public IterVarRelationNode {
- public:
- /*! \brief The outer domain */
- IterVar outer;
- /*! \brief The inner domain */
- IterVar inner;
- /*! \brief The target domain */
- IterVar fused;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("outer", &outer);
- v->Visit("inner", &inner);
- v->Visit("fused", &fused);
- }
-
- static constexpr const char* _type_key = "Fuse";
- TVM_DECLARE_FINAL_OBJECT_INFO(FuseNode, IterVarRelationNode);
-};
-
-/*!
- * \brief Managed reference to FuseNode
- * \sa FuseNode
- */
-class Fuse : public IterVarRelation {
- public:
- TVM_DLL Fuse(IterVar outer, IterVar inner, IterVar fused);
-
- TVM_DEFINE_OBJECT_REF_METHODS(Fuse, IterVarRelation, FuseNode);
-};
-
-/*!
- * \brief Rebase the iteration to make min to be 0.
- * This is useful to normalize the Schedule
- * to make every leaf variable's min to be 0.
- */
-class RebaseNode : public IterVarRelationNode {
- public:
- /*! \brief The parent domain */
- IterVar parent;
- /*! \brief The inner domain */
- IterVar rebased;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("parent", &parent);
- v->Visit("rebased", &rebased);
- }
-
- static constexpr const char* _type_key = "Rebase";
- TVM_DECLARE_FINAL_OBJECT_INFO(RebaseNode, IterVarRelationNode);
-};
-
-/*!
- * \brief Managed reference to RebaseNode
- * \sa RebaseNode
- */
-class Rebase : public IterVarRelation {
- public:
- TVM_DLL Rebase(IterVar parent, IterVar rebased);
-
- TVM_DEFINE_OBJECT_REF_METHODS(Rebase, IterVarRelation, RebaseNode);
-};
-
-/*!
- * \brief Singleton iterator [0, 1)
- */
-class SingletonNode : public IterVarRelationNode {
- public:
- /*! \brief The singleton iterator */
- IterVar iter;
-
- void VisitAttrs(AttrVisitor* v) { v->Visit("iter", &iter); }
-
- static constexpr const char* _type_key = "Singleton";
- TVM_DECLARE_FINAL_OBJECT_INFO(SingletonNode, IterVarRelationNode);
-};
-
-/*!
- * \brief Managed reference to SingletonNode
- * \sa SingletonNode
- */
-class Singleton : public IterVarRelation {
- public:
- TVM_DLL explicit Singleton(IterVar iter);
-
- TVM_DEFINE_OBJECT_REF_METHODS(Singleton, IterVarRelation, SingletonNode);
-};
-
-/*!
- * \brief Transform iterator according to some arbitrary expression.
- */
-class TransformNode : public IterVarRelationNode {
- public:
- /*! \brief The loop variables that were replaced by the transformation.
- *
- * Prior to applying a layout transformation, these represent the
- * loops to iterate over a tensor as it is being computed, following
- * a row-major traversal of the tensor's original shape in the
- * compute definition.
- */
- Array<IterVar> original_variables;
-
- /*! \brief The variables generated by the transformation.
- *
- * After to applying a layout transformation, these represent the
- * loops to iterate over a tensor as it is being computed, following
- * a row-major traversal of the transformed shape of the tensor.
- */
- Array<IterVar> transformed_variables;
-
- /*! \brief Map from the original variables to the transformed variables.
- *
- * Used to determine iterator ranges over the transformed variables.
- */
- IndexMap forward_transformation;
-
- /*! \brief Map from transformed variables to the original variables
- *
- * Used to rewrite expressions containing the original loop iterators
- * in terms of the transformed loop iterators.
- */
- IndexMap inverse_transformation;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("original_variables", &original_variables);
- v->Visit("transformed_variables", &transformed_variables);
- v->Visit("forward_transformation", &forward_transformation);
- v->Visit("inverse_transformation", &inverse_transformation);
- }
-
- static constexpr const char* _type_key = "Transform";
- TVM_DECLARE_FINAL_OBJECT_INFO(TransformNode, IterVarRelationNode);
-};
-
-class Transform : public IterVarRelation {
- public:
- TVM_DLL explicit Transform(Array<IterVar> original_variables,
- Array<IterVar> transformed_variables, IndexMap forward_transformation,
- IndexMap inverse_transformation);
-
- TVM_DEFINE_OBJECT_REF_METHODS(Transform, IterVarRelation, TransformNode);
-};
-
-/*! \brief Container for specialization conditions. */
-class SpecializedConditionNode : public Object {
- public:
- /*!
- * \brief List of conditions in conjunctive joint form (CNF).
- * Each condition should be a simple expression, e.g., n > 16, m % 8 == 0, etc.,
- * where n, m are tvm::Var that represents a dimension in the tensor shape.
- */
- Array<PrimExpr> clauses;
-
- void VisitAttrs(AttrVisitor* v) { v->Visit("clauses", &clauses); }
-
- static constexpr const char* _type_key = "SpecializedCondition";
- TVM_DECLARE_FINAL_OBJECT_INFO(SpecializedConditionNode, Object);
-};
-
-/*!
- * \brief Specialized condition to enable op specialization
- */
-class SpecializedCondition : public ObjectRef {
- public:
- /*!
- * \brief construct from conditions
- * \param conditions The clauses in the specialized condition.
- */
- TVM_DLL SpecializedCondition(Array<PrimExpr> conditions); // NOLINT(*)
-
- /*!
- * \brief Get the current specialized condition.
- * \return the current specialized condition.
- */
- TVM_DLL static SpecializedCondition Current();
-
- TVM_DEFINE_OBJECT_REF_METHODS(SpecializedCondition, ObjectRef, SpecializedConditionNode);
- class Internal;
-
- private:
- // enable with syntax.
- friend class Internal;
- friend class With<SpecializedCondition>;
- /*! \brief Push a new specialized condition onto the thread local stack. */
- TVM_DLL void EnterWithScope();
- /*! \brief Pop a specialized condition off the thread local context stack. */
- TVM_DLL void ExitWithScope();
-};
-
-// implementations
-inline const StageNode* Stage::operator->() const { return static_cast<const StageNode*>(get()); }
-inline StageNode* Stage::operator->() { return static_cast<StageNode*>(get_mutable()); }
-
-inline const ScheduleNode* Schedule::operator->() const {
- return static_cast<const ScheduleNode*>(get());
-}
-inline ScheduleNode* Schedule::operator->() { return static_cast<ScheduleNode*>(get_mutable()); }
-
-inline const IterVarRelationNode* IterVarRelation::operator->() const {
- return static_cast<const IterVarRelationNode*>(get());
-}
-
-inline const IterVarAttrNode* IterVarAttr::operator->() const {
- return static_cast<const IterVarAttrNode*>(get());
-}
-
-} // namespace te
-} // namespace tvm
-#endif // TVM_TE_SCHEDULE_H_
diff --git a/include/tvm/te/schedule_pass.h b/include/tvm/te/schedule_pass.h
deleted file mode 100644
index 3f9da5fb5d..0000000000
--- a/include/tvm/te/schedule_pass.h
+++ /dev/null
@@ -1,112 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/te/schedule_pass.h
- * \brief Collection of Schedule pass functions.
- *
- * These passes works on the schedule hyper-graph
- * and infers information such as bounds, check conditions
- * read/write dependencies between the IterVar
- */
-#ifndef TVM_TE_SCHEDULE_PASS_H_
-#define TVM_TE_SCHEDULE_PASS_H_
-
-#include <tvm/te/schedule.h>
-#include <tvm/tir/function.h>
-
-namespace tvm {
-namespace te {
-
-/*!
- * \brief To automatically inline the element-wise operations.
- *
- * \param sch The schedule to be inlined.
- */
-void AutoInlineElemWise(Schedule sch);
-
-/*!
- * \brief To automatically inline the broadcast operations.
- *
- * \param sch The schedule to be inlined.
- */
-void AutoInlineBroarcast(Schedule sch);
-
-/*!
- * \brief To automatically inline operations with injective writes
- * (i.e. writes without reduction or sequential loops). Note
- * that in this case, guarantees about contiguity, transpose, stride,
- * alignemnt and memory footprint in general do not hold.
- *
- * \param sch The schedule to be inlined.
- */
-TVM_DLL void AutoInlineInjective(Schedule sch);
-
-/*!
- * \brief Infer the bound of all iteration variables relates to the schedule.
- *
- * \param sch The root schedule to infer all the bounds.
- * \return the result bound of the iteration Variable
- */
-Map<IterVar, Range> InferBound(const Schedule& sch);
-
-/*!
- * \brief Verify if there is any argument bound to compact buffer.
- *
- * \param stmt The stmt to be verified.
- * \return true if there is any buffer_bind_scope attribute found,
- * otherwise, false.
- */
-bool VerifyCompactBuffer(const Stmt& stmt);
-
-/*!
- * \brief Schedule s' dependent operations.
- *
- * \param s The schedule to be realized
- * \param dom_map The domain of each iter vars.
- * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1 during lowering.
- * This is a debug feature for dataflow/axis analysis.
- * Note: If this is true, The lowered IR may be incorrect,
- * because we will also delete the init part of reduction
- * \return the result Stmt
- */
-Stmt ScheduleOps(Schedule s, Map<IterVar, Range> dom_map, bool debug_keep_trivial_loop);
-
-/*!
- * \brief Postprocessing the Stmt generated by ScheduleOps to create
- * a PrimFunc that can then be used for further TIR optimizations.
- *
- * Perform this translation before running any TIR optimizations.
- *
- * List of actions taken by the function:
- * - Remove occurrences of te::Tensor, te::Operation in the IR
- * and replace them by corresponding IR nodes via tir::Buffer.
- * - Add annotation of extern buffers using the buffer_map field
- * in the PrimFunc type.
- *
- * \param arg_list Array of Tensor/Var/Buffer arguments to the function.
- * \param body The body of the function.
- * \param bindings potential Tensor to Buffer bindings for the Tensors in the body.
- */
-PrimFunc SchedulePostProcToPrimFunc(Array<ObjectRef> arg_list, Stmt body,
- Optional<Map<Tensor, Buffer>> bindings);
-
-} // namespace te
-} // namespace tvm
-#endif // TVM_TE_SCHEDULE_PASS_H_
diff --git a/include/tvm/te/tensor_intrin.h b/include/tvm/te/tensor_intrin.h
deleted file mode 100644
index 22f29defbb..0000000000
--- a/include/tvm/te/tensor_intrin.h
+++ /dev/null
@@ -1,145 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file tvm/te/tensor_intrin.h
- * \brief Tensor intrinsic operations.
- */
-#ifndef TVM_TE_TENSOR_INTRIN_H_
-#define TVM_TE_TENSOR_INTRIN_H_
-
-#include <tvm/te/tensor.h>
-#include <tvm/tir/buffer.h>
-
-#include <string>
-
-namespace tvm {
-namespace te {
-
-/*! \brief Node to represent a Tensor intrinsic operator */
-class TensorIntrinNode : public Object {
- public:
- /*! \brief The name of the intrinsic */
- std::string name;
- /*! \brief The operation this intrinsics is carrying out */
- Operation op;
- /*! \brief List of inputs of operator, placeholder in postdfs order */
- Array<Tensor> inputs;
- /*!
- * \brief Symbolic buffers of each output/input tensor
- * buffers[0:len(inputs)] are buffers of the inputs.
- * buffers[len(inputs):] are buffers of each output.
- *
- * \note When a field in Buffer is Var, it means we can be flexible
- * wrt that field and Var can occur in body.
- * When it is a constant, it means we can only take data in that shape.
- */
- Array<Buffer> buffers;
- /*! \brief List of scalar variables, used in body. These placeholders
- * will be bound to expressions passed in when the TensorIntrin is called
- * from a TensorComputeOp.
- */
- Array<Var> scalar_params;
- /*! \brief The normal statement to execute the intrinsic */
- Stmt body;
- /*!
- * \brief Special statement for reduction op, can be None
- * reset the value of output buffer to identity value.
- */
- Stmt reduce_init;
- /*!
- * \brief Special statement for reduction op, can be None
- * Reduce: do a reduction of current output buffer with the result.
- */
- Stmt reduce_update;
- /*! \brief constructor */
- TensorIntrinNode() {}
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("name", &name);
- v->Visit("op", &op);
- v->Visit("inputs", &inputs);
- v->Visit("buffers", &buffers);
- v->Visit("scalar_params", &scalar_params);
- v->Visit("body", &body);
- v->Visit("reduce_init", &reduce_init);
- v->Visit("reduce_update", &reduce_update);
- }
-
- static constexpr const char* _type_key = "TensorIntrin";
- TVM_DECLARE_FINAL_OBJECT_INFO(TensorIntrinNode, Object);
-};
-
-/*!
- * \brief Managed reference to TensorIntrinNode
- * \sa TensorIntrinNode
- */
-class TensorIntrin : public ObjectRef {
- public:
- TVM_DLL TensorIntrin(std::string name, Operation op, Array<Tensor> inputs, Array<Buffer> buffers,
- Array<Var> scalar_params, Stmt body, Stmt reduce_init, Stmt reduce_update);
-
- TVM_DEFINE_OBJECT_REF_METHODS(TensorIntrin, ObjectRef, TensorIntrinNode);
-};
-
-class TensorIntrinCallNode : public Object {
- public:
- /*! \brief the tensor intrinsic */
- TensorIntrin intrin;
- /*! \brief input tensors of the intrinsic */
- Array<Tensor> tensors;
- /*! \brief regions of input tensors */
- Array<Region> regions;
-
- /*!
- * \brief IterVar on each reduction axis, if the
- * intrin will use the reduce axis
- */
- Array<IterVar> reduce_axis;
-
- /*! \brief scalar expression inputs */
- Array<PrimExpr> scalar_inputs;
-
- void VisitAttrs(AttrVisitor* v) {
- v->Visit("intrin", &intrin);
- v->Visit("tensors", &tensors);
- v->Visit("regions", ®ions);
- v->Visit("reduce_axis", &reduce_axis);
- v->Visit("scalar_inputs", &scalar_inputs);
- }
-
- static constexpr const char* _type_key = "TensorIntrinCall";
- TVM_DECLARE_FINAL_OBJECT_INFO(TensorIntrinCallNode, Object);
-};
-
-/*!
- * \brief Managed reference to TensorIntrinCallNode
- * \sa TensorIntrinCallNode
- */
-class TensorIntrinCall : public ObjectRef {
- public:
- TVM_DLL TensorIntrinCall(TensorIntrin intrin, Array<Tensor> tensors, Array<Region> regions,
- Array<IterVar> reduce_axis, Array<PrimExpr> scalar_inputs);
-
- TVM_DEFINE_OBJECT_REF_METHODS(TensorIntrinCall, ObjectRef, TensorIntrinCallNode);
-};
-
-} // namespace te
-} // namespace tvm
-#endif // TVM_TE_TENSOR_INTRIN_H_
diff --git a/include/tvm/topi/cuda/dense.h b/include/tvm/topi/cuda/dense.h
deleted file mode 100644
index 7fd3107b6c..0000000000
--- a/include/tvm/topi/cuda/dense.h
+++ /dev/null
@@ -1,154 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file cuda/dense.h
- * \brief CUDA schedule for dense operation
- */
-#ifndef TVM_TOPI_CUDA_DENSE_H_
-#define TVM_TOPI_CUDA_DENSE_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/contrib/cublas.h>
-#include <tvm/topi/detail/array_utils.h>
-#include <tvm/topi/generic/extern.h>
-#include <tvm/topi/nn/dense.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace cuda {
-/*!
- * \brief Implementation of dense for CUDA backend
- *
- * \param target The target device
- * \param data Tensor with shape [batch, in_dim]
- * \param weight Tensor with shape [out_dim, in_dim]
- * \param bias Tensor with shape [out_dim]. Optional; to omit bias, pass Tensor()
- * \param out_dtype Output data type. Used for mixed precision.
- *
- * \return Tensor with shape [batch, out_dim]
- */
-inline tvm::te::Tensor dense_cuda(const Target& target, const tvm::te::Tensor& data,
- const tvm::te::Tensor& weight, const tvm::te::Tensor& bias,
- const DataType& out_dtype) {
- ICHECK_EQ(data->shape.size(), 2) << "dense requires 2-D data";
- ICHECK_EQ(weight->shape.size(), 2) << "dense requires 2-D weight";
- if (bias.defined()) {
- ICHECK_EQ(bias->shape.size(), 1) << "dense requires 1-D bias";
- }
-
- auto batch = data->shape[0];
- auto in_dim = data->shape[1];
- auto out_dim = weight->shape[0];
-
- if (target->GetLibs().count("cublas")) {
- ICHECK_EQ(data->dtype, out_dtype) << "Mixed precision not supported.";
- auto mm = topi::contrib::cublas_matmul(data, weight, false, true);
- if (bias.defined()) {
- mm = tvm::te::compute(
- {batch, out_dim}, [&](Var i, Var j) { return mm(i, j) + bias(j); }, "tensor", kBroadcast);
- }
-
- return mm;
- } else {
- return topi::nn::dense(data, weight, bias, out_dtype);
- }
-}
-
-/*!
- * \brief Create a CUDA schedule for dense
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_dense(const Target& target, const Array<Tensor>& outs) {
- if (target->kind->name == "cuda" && target->GetLibs().count("cublas")) {
- return topi::generic::schedule_extern(target, outs);
- }
-
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- auto _schedule = [&](const Tensor& dense) {
- auto num_thread = 64;
- auto k = dense->op.as<ComputeOpNode>()->reduce_axis[0];
- IterVar ko, kf;
- s[dense].split(k, num_thread, &ko, &kf);
- auto dense_f = s.rfactor(dense, kf)[0];
-
- Tensor out;
- if (detail::contains(s->outputs, dense->op)) {
- out = dense;
- } else {
- out = outs[0]->op.output(0);
- s[dense].compute_at(s[out], s[out]->op.as<ComputeOpNode>()->axis[1]);
- }
- s[out].bind(s[out]->op.as<ComputeOpNode>()->axis[0],
- tvm::te::thread_axis(Range(), "blockIdx.y"));
- s[out].bind(s[out]->op.as<ComputeOpNode>()->axis[1],
- tvm::te::thread_axis(Range(), "blockIdx.x"));
-
- auto tx = s[dense]->op.as<ComputeOpNode>()->reduce_axis[0];
- auto thread_x = tvm::te::thread_axis(Range(), "threadIdx.x");
- s[dense].bind(tx, thread_x);
- s[dense_f].compute_at(s[dense], tx);
- s[dense].set_store_predicate(static_cast<PrimExpr>(thread_x) == 0);
- s[out].set_store_predicate(static_cast<PrimExpr>(thread_x) == 0);
- };
-
- std::function<void(Operation)> traverse;
- traverse = [&](const Operation& op) {
- // Inline all one-to-one-mapping operators except the last stage (output)
- if (is_broadcast(op->tag)) {
- if (!detail::contains(s->outputs, op)) {
- s[op].compute_inline();
- }
- for (auto tensor : op->InputTensors()) {
- if (tensor->op->InputTensors().size() > 0) {
- traverse(tensor->op);
- }
- }
- } else if (op->tag == "dense") {
- // If tag starts with global_pool
- auto dense = op.output(0);
- _schedule(dense);
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
- };
-
- traverse(outs[0]->op);
- return s;
-}
-
-} // namespace cuda
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_CUDA_DENSE_H_
diff --git a/include/tvm/topi/cuda/injective.h b/include/tvm/topi/cuda/injective.h
deleted file mode 100644
index 79ec338aae..0000000000
--- a/include/tvm/topi/cuda/injective.h
+++ /dev/null
@@ -1,83 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file cuda/injective.h
- * \brief CUDA schedule for injective operations
- */
-#ifndef TVM_TOPI_CUDA_INJECTIVE_H_
-#define TVM_TOPI_CUDA_INJECTIVE_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace cuda {
-
-/*!
- * \brief Updates an existing schedule for the given injective ops.
- *
- * \param sch The schedule to update.
- * \param out The tensor representing the injective op.
- *
- * \return The updated schedule.
- */
-inline Schedule schedule_injective_from_existing(Schedule sch, const Tensor& out) {
- auto fused = detail::Fuse(sch[out], sch[out]->op.as<ComputeOpNode>()->axis);
- auto target = Target::Current(false);
- int num_thread = target->GetAttr<Integer>("max_num_threads").value().IntValue();
- IterVar bx, tx;
- sch[out].split(fused, num_thread, &bx, &tx);
- sch[out].bind(bx, thread_axis(Range(), "blockIdx.x"));
- sch[out].bind(tx, thread_axis(Range(), "threadIdx.x"));
- return sch;
-}
-
-/*!
- * \brief Create a CUDA schedule for the given output tensors.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_injective(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- tvm::te::AutoInlineInjective(s);
- for (auto out : outs) {
- schedule_injective_from_existing(s, out);
- }
- return s;
-}
-
-} // namespace cuda
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_CUDA_INJECTIVE_H_
diff --git a/include/tvm/topi/cuda/pooling.h b/include/tvm/topi/cuda/pooling.h
deleted file mode 100644
index 92be031236..0000000000
--- a/include/tvm/topi/cuda/pooling.h
+++ /dev/null
@@ -1,187 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file cuda/pooling.h
- * \brief CUDA schedule for pooling operations
- */
-#ifndef TVM_TOPI_CUDA_POOLING_H_
-#define TVM_TOPI_CUDA_POOLING_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/array_utils.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace cuda {
-
-/*!
- * \brief Create a CUDA schedule for pool
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_pool(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- auto _schedule = [&](const Tensor& padded_input, const Tensor& pool) {
- if (padded_input->op->IsInstance<ComputeOpNode>()) {
- s[padded_input].compute_inline();
- }
- int num_thread = target->GetAttr<Integer>("max_num_threads").value().IntValue();
- Tensor out;
- Tensor OL;
- if (detail::contains(s->outputs, pool->op)) {
- out = pool;
- OL = s.cache_write(pool, "local");
- } else {
- out = outs[0]->op.output(0);
- s[pool].set_scope("local");
- }
- auto fused = detail::Fuse(s[out], s[out]->op.as<ComputeOpNode>()->axis);
- IterVar bx, tx;
- s[out].split(fused, num_thread, &bx, &tx);
- s[out].bind(bx, tvm::te::thread_axis(Range(), "blockIdx.x"));
- s[out].bind(tx, tvm::te::thread_axis(Range(), "threadIdx.x"));
- if (detail::contains(s->outputs, pool->op)) {
- s[OL].compute_at(s[out], tx);
- } else {
- s[pool].compute_at(s[out], tx);
- }
- };
-
- std::function<void(Operation)> traverse;
- traverse = [&](const Operation& op) {
- // Inline all one-to-one-mapping operators except the last stage (output)
- if (is_broadcast(op->tag)) {
- if (!detail::contains(s->outputs, op)) {
- s[op].compute_inline();
- }
- for (auto tensor : op->InputTensors()) {
- if (tensor->op->InputTensors().size() > 0) {
- traverse(tensor->op);
- }
- }
- } else if (op->tag.rfind("pool", 0) == 0) {
- // If tag starts with pool
- auto padded_input = op->InputTensors()[0];
- auto pool = op.output(0);
- _schedule(padded_input, pool);
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
- };
-
- traverse(outs[0]->op);
- return s;
-}
-
-/*!
- * \brief Create a CUDA schedule for global_pool
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_global_pool(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- auto _schedule = [&](const Tensor& pool) {
- auto num_thread = 8;
- auto block_x = tvm::te::thread_axis(Range(), "blockIdx.x");
- auto block_y = tvm::te::thread_axis(Range(), "blockIdx.y");
- auto thread_x = tvm::te::thread_axis(Range(0, num_thread), "threadIdx.x");
- auto thread_y = tvm::te::thread_axis(Range(0, num_thread), "threadIdx.y");
- Tensor out;
- Tensor OL;
- if (detail::contains(s->outputs, pool->op)) {
- out = pool;
- OL = s.cache_write(pool, "local");
- } else {
- out = outs[0]->op.output(0);
- s[pool].set_scope("local");
- }
-
- auto i = s[out]->op.as<ComputeOpNode>()->axis[0];
- auto c = s[out]->op.as<ComputeOpNode>()->axis[1];
-
- IterVar by, ty;
- s[out].split(i, num_thread, &by, &ty);
- IterVar bx, tx;
- s[out].split(c, num_thread, &bx, &tx);
- s[out].reorder({by, bx, ty, tx});
- s[out].bind(ty, thread_y);
- s[out].bind(tx, thread_x);
- s[out].bind(by, block_y);
- s[out].bind(bx, block_x);
-
- if (detail::contains(s->outputs, pool->op)) {
- s[OL].compute_at(s[out], tx);
- } else {
- s[pool].compute_at(s[out], tx);
- }
- };
-
- std::function<void(Operation)> traverse;
- traverse = [&](const Operation& op) {
- // Inline all one-to-one-mapping operators except the last stage (output)
- if (is_broadcast(op->tag)) {
- if (!detail::contains(s->outputs, op)) {
- s[op].compute_inline();
- }
- for (auto tensor : op->InputTensors()) {
- if (tensor->op->InputTensors().size() > 0) {
- traverse(tensor->op);
- }
- }
- } else if (op->tag.rfind("global_pool", 0) == 0) {
- // If tag starts with global_pool
- auto pool = op.output(0);
- _schedule(pool);
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
- };
-
- traverse(outs[0]->op);
- return s;
-}
-
-} // namespace cuda
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_CUDA_POOLING_H_
diff --git a/include/tvm/topi/cuda/reduction.h b/include/tvm/topi/cuda/reduction.h
deleted file mode 100644
index b1905d8442..0000000000
--- a/include/tvm/topi/cuda/reduction.h
+++ /dev/null
@@ -1,199 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file cuda/reduction.h
- * \brief CUDA schedule for reduction operations
- */
-#ifndef TVM_TOPI_CUDA_REDUCTION_H_
-#define TVM_TOPI_CUDA_REDUCTION_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace cuda {
-/*!
- * \brief Schedule a given reduce operation.
- *
- * \param target The target to generate a schedule for.
- * \param op The operation representing the injective operation.
- * \param sch The schedule to apply this scheduling to
- * \param is_idx_reduce Pass true to schedule a reduce op that returns
- * an index, such as argmax or argmin.
- *
- * \return The schedule given by sch
- */
-Schedule ScheduleReduce(const Target& target, Operation op, Schedule sch,
- bool is_idx_reduce = false) {
- Tensor data_out;
- Tensor data_in;
-
- if (!is_idx_reduce) {
- data_in = op->InputTensors()[0];
- data_out = op.output(0);
- } else {
- data_out = op->InputTensors()[0];
- }
-
- auto out_stage = sch[data_out];
- ICHECK_GT(out_stage->op.as<ComputeOpNode>()->reduce_axis.size(), 0)
- << "reduce_axis must be greater than zero";
-
- bool all_reduce;
- int num_thread;
- IterVar block_x, thread_x, thread_y;
-
- if (out_stage->op.as<ComputeOpNode>()->axis.size() > 0) {
- all_reduce = false;
- num_thread = 32;
- if (target->kind->name == "opencl" || target->kind->name == "metal") {
- // Without this, CL_INVALID_WORK_GROUP_SIZE occurs with python tests.
- // Don't know why.
- num_thread = 16;
- }
- block_x = tvm::te::thread_axis(Range(), "blockIdx.x");
- thread_x = tvm::te::thread_axis(Range(0, num_thread), "threadIdx.x");
- thread_y = tvm::te::thread_axis(Range(0, num_thread), "threadIdx.y");
- } else {
- all_reduce = true;
- num_thread = target->GetAttr<Integer>("max_num_threads").value().IntValue();
- thread_x = tvm::te::thread_axis(Range(0, num_thread), "threadIdx.x");
- }
-
- auto fused_reduce = detail::Fuse(out_stage, out_stage->op.as<ComputeOpNode>()->reduce_axis);
-
- IterVar ko, ki;
- out_stage.split(fused_reduce, num_thread, &ko, &ki);
- auto data_out_rf = sch.rfactor(data_out, ki)[0];
- auto tx = out_stage->op.as<ComputeOpNode>()->reduce_axis[0];
- out_stage.bind(tx, thread_x);
- sch[data_out_rf].compute_at(out_stage, tx);
-
- Tensor real_output;
- Tensor temp_idx_input, temp_val_input;
- if (is_idx_reduce) {
- real_output = op.output(0);
- temp_idx_input = data_out->op.output(0);
- temp_val_input = data_out->op.output(1);
- } else {
- real_output = data_out;
- }
-
- auto stage_real = sch[real_output];
- if (!all_reduce) {
- // Fuse and split the axis
- auto fused_outer = detail::Fuse(stage_real, stage_real->op.as<ComputeOpNode>()->axis);
- IterVar bx, outer_in;
- stage_real.split(fused_outer, num_thread, &bx, &outer_in);
-
- // Bind the axes to threads and blocks
- stage_real.bind(outer_in, thread_y);
- stage_real.bind(bx, block_x);
- if (is_idx_reduce) {
- sch[temp_idx_input].compute_at(stage_real, outer_in);
- sch[temp_val_input].compute_at(stage_real, outer_in);
- }
- } else {
- if (is_idx_reduce) {
- sch[temp_idx_input].compute_at(stage_real, stage_real->op.as<ComputeOpNode>()->axis[0]);
- sch[temp_val_input].compute_at(stage_real, stage_real->op.as<ComputeOpNode>()->axis[0]);
- }
- }
-
- stage_real.set_store_predicate(static_cast<PrimExpr>(thread_x) == 0);
- return sch;
-}
-
-/*!
- * \brief Recursively traverse operator inputs, setting injective inputs
- * to be computed inline.
- *
- * \param s The schedule we are building
- * \param op The current op in the traversal
- */
-void TraverseBeforeReduce(Schedule s, Operation op) {
- if (op->IsInstance<PlaceholderOpNode>()) {
- return;
- } else if (is_injective(op->tag)) {
- s[op].compute_inline();
- for (auto tensor : op->InputTensors()) {
- TraverseBeforeReduce(s, tensor->op);
- }
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
-}
-
-/*!
- * \brief Schedule a reduce op, then invoke TraverseBeforeReduce on each
- * of the op's inputs.
- *
- * \param target The target to generate a schedule for.
- * \param s The schedule we are building
- * \param op The reduce op
- */
-void TraverseAfterReduce(const Target& target, Schedule s, Operation op) {
- if (is_broadcast(op->tag)) {
- LOG(ERROR) << "Elementwise op after reduce is not yet supported";
- } else if (op->tag == kCommReduce) {
- ScheduleReduce(target, op, s, false);
- for (auto tensor : op->InputTensors()) {
- TraverseBeforeReduce(s, tensor->op);
- }
- } else if (op->tag == kCommReduceIdx) {
- ScheduleReduce(target, op, s, true);
- for (auto tensor : op->InputTensors()[0]->op->InputTensors()) {
- TraverseBeforeReduce(s, tensor->op);
- }
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
-}
-
-/*!
- * \brief Create a CUDA schedule for a reduce operation.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-Schedule schedule_reduce(const Target& target, Array<Tensor> outs) {
- ICHECK_EQ(outs.size(), 1) << "outs must have size 1";
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- TraverseAfterReduce(target, s, outs[0]->op);
- return s;
-}
-
-} // namespace cuda
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_CUDA_REDUCTION_H_
diff --git a/include/tvm/topi/cuda/softmax.h b/include/tvm/topi/cuda/softmax.h
deleted file mode 100644
index 19613cbbdf..0000000000
--- a/include/tvm/topi/cuda/softmax.h
+++ /dev/null
@@ -1,103 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file cuda/injective.h
- * \brief CUDA schedule for injective operations
- */
-#ifndef TVM_TOPI_CUDA_SOFTMAX_H_
-#define TVM_TOPI_CUDA_SOFTMAX_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace cuda {
-
-/*!
- * \brief Create a CUDA schedule for the given softmax output tensors.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_softmax(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- auto softmax = outs[0];
- tvm::te::Tensor max_elem;
- tvm::te::Tensor expsum;
- tvm::te::Tensor exp;
- bool has_exp = false;
-
- auto tag = softmax->op.as<ComputeOpNode>()->tag;
- if (tag == "softmax_output") {
- expsum = softmax->op->InputTensors()[1];
- exp = softmax->op->InputTensors()[0];
- max_elem = s[exp]->op->InputTensors()[1];
- has_exp = true;
- } else if (tag == "log_softmax_output") {
- max_elem = softmax->op->InputTensors()[1];
- expsum = softmax->op->InputTensors()[2];
- } else {
- LOG(ERROR) << "Tag is expected to be softmax_output or log_softmax_output. Got " << tag;
- }
-
- int num_thread = 64;
- auto block_x = tvm::te::thread_axis(Range(), "blockIdx.x");
- auto thread_x = tvm::te::thread_axis(Range(0, num_thread), "threadIdx.x");
-
- if (has_exp) {
- s[exp].bind(exp->op.as<ComputeOpNode>()->axis[0], block_x);
- }
-
- s[max_elem].bind(max_elem->op.as<ComputeOpNode>()->axis[0], block_x);
-
- auto k = expsum->op.as<ComputeOpNode>()->reduce_axis[0];
- IterVar ko, ki;
- s[expsum].split(k, num_thread, &ko, &ki);
- auto EF = s.rfactor(expsum, ki)[0];
- s[expsum].bind(s[expsum]->op.as<ComputeOpNode>()->axis[0], block_x);
- s[expsum].bind(s[expsum]->op.as<ComputeOpNode>()->reduce_axis[0], thread_x);
- s[EF].compute_at(s[expsum], s[expsum]->op.as<ComputeOpNode>()->reduce_axis[0]);
- s[expsum].set_store_predicate(thread_x->var == 0);
-
- IterVar tx, xi;
- s[softmax].split_by_nparts(softmax->op.as<ComputeOpNode>()->axis[1], num_thread, &tx, &xi);
- s[softmax].bind(tx, thread_x);
-
- return s;
-}
-
-} // namespace cuda
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_CUDA_SOFTMAX_H_
diff --git a/include/tvm/topi/generic/default.h b/include/tvm/topi/generic/default.h
deleted file mode 100644
index 752b6ad153..0000000000
--- a/include/tvm/topi/generic/default.h
+++ /dev/null
@@ -1,83 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file generic/default.h
- * \brief Generic default schedule
- */
-#ifndef TVM_TOPI_GENERIC_DEFAULT_H_
-#define TVM_TOPI_GENERIC_DEFAULT_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace generic {
-/*!
- * \brief Create a generic default schedule for the given output tensors.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule default_schedule(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- return s;
-}
-
-/*!
- * \brief Create a generic default schedule for the given output tensors, and apply
- * auto inline
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule default_schedule_auto_inline(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- auto x = outs[0];
- tvm::te::AutoInlineInjective(s);
- auto axis = s[x]->op.as<ComputeOpNode>()->axis;
- if (axis.size() > 0) {
- detail::Fuse(s[x], axis);
- }
- return s;
-}
-
-} // namespace generic
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_GENERIC_DEFAULT_H_
diff --git a/include/tvm/topi/generic/extern.h b/include/tvm/topi/generic/extern.h
deleted file mode 100644
index 0f1f408fdc..0000000000
--- a/include/tvm/topi/generic/extern.h
+++ /dev/null
@@ -1,69 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file generic/extern.h
- * \brief Schedule for extern followed by injective ops
- */
-#ifndef TVM_TOPI_GENERIC_EXTERN_H_
-#define TVM_TOPI_GENERIC_EXTERN_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/generic/injective.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace generic {
-/*!
- * \brief Schedule an extern op followed by injective operations
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the op.
- */
-inline Schedule schedule_extern(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- tvm::te::AutoInlineInjective(s);
- for (auto out : outs) {
- if (out->op->IsInstance<ExternOpNode>()) {
- continue;
- }
- tvm::GenericFunc::Get("schedule_injective_from_existing")(s, out);
- }
-
- return s;
-}
-
-} // namespace generic
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_GENERIC_EXTERN_H_
diff --git a/include/tvm/topi/generic/injective.h b/include/tvm/topi/generic/injective.h
deleted file mode 100644
index c48c03eee0..0000000000
--- a/include/tvm/topi/generic/injective.h
+++ /dev/null
@@ -1,77 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file generic/injective.h
- * \brief Generic schedule for injective operations
- */
-#ifndef TVM_TOPI_GENERIC_INJECTIVE_H_
-#define TVM_TOPI_GENERIC_INJECTIVE_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace generic {
-
-/*!
- * \brief Updates an existing schedule for the given injective ops.
- *
- * \param sch The schedule to update.
- * \param out The tensor representing the injective op.
- *
- * \return The updated schedule.
- */
-inline Schedule schedule_injective_from_existing(Schedule sch, const Tensor& out) {
- detail::Fuse(sch[out], sch[out]->op.as<ComputeOpNode>()->axis);
- return sch;
-}
-
-/*!
- * \brief Create a generic schedule for the given injective ops.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_injective(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- tvm::te::AutoInlineInjective(s);
- auto x = outs[0];
- schedule_injective_from_existing(s, x);
-
- return s;
-}
-
-} // namespace generic
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_GENERIC_INJECTIVE_H_
diff --git a/include/tvm/topi/rocm/dense.h b/include/tvm/topi/rocm/dense.h
deleted file mode 100644
index b861e6c89a..0000000000
--- a/include/tvm/topi/rocm/dense.h
+++ /dev/null
@@ -1,99 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file rocm/dense.h
- * \brief rocm schedule for dense operation
- */
-#ifndef TVM_TOPI_ROCM_DENSE_H_
-#define TVM_TOPI_ROCM_DENSE_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/topi/contrib/rocblas.h>
-#include <tvm/topi/cuda/dense.h>
-#include <tvm/topi/detail/array_utils.h>
-#include <tvm/topi/generic/extern.h>
-#include <tvm/topi/nn/dense.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace rocm {
-/*!
- * \brief Implementation of dense for rocm backend
- *
- * \param target The target device
- * \param data Tensor with shape [batch, in_dim]
- * \param weight Tensor with shape [out_dim, in_dim]
- * \param bias Tensor with shape [out_dim]. Optional; to omit bias, pass Tensor()
- * \param out_dtype Output data type. Used for mixed precision.
- *
- * \return Tensor with shape [batch, out_dim]
- */
-inline tvm::te::Tensor dense_rocm(const Target& target, const tvm::te::Tensor& data,
- const tvm::te::Tensor& weight, const tvm::te::Tensor& bias,
- const DataType& out_dtype) {
- ICHECK_EQ(data->shape.size(), 2) << "dense requires 2-D data";
- ICHECK_EQ(weight->shape.size(), 2) << "dense requires 2-D weight";
- if (bias.defined()) {
- ICHECK_EQ(bias->shape.size(), 1) << "dense requires 1-D bias";
- }
-
- auto batch = data->shape[0];
- auto in_dim = data->shape[1];
- auto out_dim = weight->shape[0];
-
- if (target->GetLibs().count("rocblas")) {
- ICHECK_EQ(data->dtype, out_dtype) << "Mixed precision not supported.";
- auto mm = topi::contrib::rocblas_matmul(data, weight, false, true);
- if (bias.defined()) {
- mm = tvm::te::compute(
- {batch, out_dim}, [&](Var i, Var j) { return mm(i, j) + bias(j); }, "tensor", kBroadcast);
- }
-
- return mm;
- } else {
- return topi::nn::dense(data, weight, bias, out_dtype);
- }
-}
-
-/*!
- * \brief Create a rocm schedule for dense
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_dense(const Target& target, const Array<Tensor>& outs) {
- if (target->kind->name == "rocm" && target->GetLibs().count("rocblas")) {
- return topi::generic::schedule_extern(target, outs);
- }
-
- return topi::cuda::schedule_dense(target, outs);
-}
-
-} // namespace rocm
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_ROCM_DENSE_H_
diff --git a/include/tvm/topi/rocm/injective.h b/include/tvm/topi/rocm/injective.h
deleted file mode 100644
index 295d930e5c..0000000000
--- a/include/tvm/topi/rocm/injective.h
+++ /dev/null
@@ -1,67 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file rocm/injective.h
- * \brief rocm schedule for injective operations
- */
-#ifndef TVM_TOPI_ROCM_INJECTIVE_H_
-#define TVM_TOPI_ROCM_INJECTIVE_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/topi/cuda/injective.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace rocm {
-
-/*!
- * \brief Updates an existing schedule for the given injective ops.
- *
- * \param sch The schedule to update.
- * \param out The tensor representing the injective op.
- *
- * \return The updated schedule.
- */
-inline Schedule schedule_injective_from_existing(Schedule sch, const Tensor& out) {
- return topi::cuda::schedule_injective_from_existing(sch, out);
-}
-
-/*!
- * \brief Create a rocm schedule for the given output tensors.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_injective(const Target& target, const Array<Tensor>& outs) {
- return topi::cuda::schedule_injective(target, outs);
-}
-
-} // namespace rocm
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_ROCM_INJECTIVE_H_
diff --git a/include/tvm/topi/rocm/pooling.h b/include/tvm/topi/rocm/pooling.h
deleted file mode 100644
index 993c32bf36..0000000000
--- a/include/tvm/topi/rocm/pooling.h
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file rocm/pooling.h
- * \brief rocm schedule for pooling operations
- */
-#ifndef TVM_TOPI_ROCM_POOLING_H_
-#define TVM_TOPI_ROCM_POOLING_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/topi/cuda/pooling.h>
-#include <tvm/topi/detail/array_utils.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace rocm {
-
-/*!
- * \brief Create a rocm schedule for pool
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_pool(const Target& target, const Array<Tensor>& outs) {
- return topi::cuda::schedule_pool(target, outs);
-}
-
-/*!
- * \brief Create a rocm schedule for global_pool
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_global_pool(const Target& target, const Array<Tensor>& outs) {
- return topi::cuda::schedule_global_pool(target, outs);
-}
-
-} // namespace rocm
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_ROCM_POOLING_H_
diff --git a/include/tvm/topi/rocm/softmax.h b/include/tvm/topi/rocm/softmax.h
deleted file mode 100644
index a2ffd2c46e..0000000000
--- a/include/tvm/topi/rocm/softmax.h
+++ /dev/null
@@ -1,55 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file rocm/injective.h
- * \brief ROCM schedule for injective operations
- */
-#ifndef TVM_TOPI_ROCM_SOFTMAX_H_
-#define TVM_TOPI_ROCM_SOFTMAX_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/topi/cuda/softmax.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace rocm {
-
-/*!
- * \brief Create a rocm schedule for the given softmax output tensors.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_softmax(const Target& target, const Array<Tensor>& outs) {
- return topi::cuda::schedule_softmax(target, outs);
-}
-
-} // namespace rocm
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_ROCM_SOFTMAX_H_
diff --git a/include/tvm/topi/x86/bnn.h b/include/tvm/topi/x86/bnn.h
deleted file mode 100644
index c8a7235536..0000000000
--- a/include/tvm/topi/x86/bnn.h
+++ /dev/null
@@ -1,131 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file x86/bnn.h
- * \brief x86 schedule for binary operations
- */
-#ifndef TVM_TOPI_X86_BNN_H_
-#define TVM_TOPI_X86_BNN_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace x86 {
-/*!
- * \brief Create a generic schedule for binarize_pack
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_binarize_pack(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- auto _schedule = [&](const Tensor& out) {
- s[out].parallel(out->op.as<ComputeOpNode>()->axis[0]);
- };
-
- std::function<void(Operation)> traverse;
- traverse = [&](const Operation& op) {
- if (op->tag == "binarize_pack") {
- _schedule(op.output(0));
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
- };
-
- traverse(outs[0]->op);
- return s;
-}
-
-/*!
- * \brief Create a generic schedule for binary_dense
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_binary_dense(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
-
- auto _schedule = [&](const Tensor& A, const Tensor& B, const Tensor& C) {
- IterVar co, ci;
- s[C].split(s[C]->op.as<ComputeOpNode>()->reduce_axis[0], 8, &co, &ci);
- s[C].parallel(s[C]->op.as<ComputeOpNode>()->axis[0]);
-
- Tensor out;
- if (detail::contains(s->outputs, C->op)) {
- out = C;
- } else {
- out = outs[0]->op.output(0);
- }
-
- IterVar xo, xi;
- s[out].split(out->op.as<ComputeOpNode>()->axis[1], 8, &xo, &xi);
- s[out].vectorize(xi);
- };
-
- std::function<void(Operation)> traverse;
- traverse = [&](const Operation& op) {
- // Inline all one-to-one-mapping operators except the last stage (output)
- if (is_broadcast(op->tag)) {
- if (!detail::contains(s->outputs, op)) {
- s[op].compute_inline();
- }
- for (auto tensor : op->InputTensors()) {
- if (tensor->op->InputTensors().size() > 0) {
- traverse(tensor->op);
- }
- }
- } else if (op->tag == "binary_dense") {
- auto output = op.output(0);
- auto data = op->InputTensors()[0];
- auto weight = op->InputTensors()[1];
- _schedule(data, weight, output);
- } else {
- LOG(ERROR) << "Unsupported operator " << op->tag;
- }
- };
-
- traverse(outs[0]->op);
- return s;
-}
-
-} // namespace x86
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_X86_BNN_H_
diff --git a/include/tvm/topi/x86/default.h b/include/tvm/topi/x86/default.h
deleted file mode 100644
index 9c98560402..0000000000
--- a/include/tvm/topi/x86/default.h
+++ /dev/null
@@ -1,105 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file x86/default.h
- * \brief default x86 schedule
- */
-#ifndef TVM_TOPI_X86_DEFAULT_H_
-#define TVM_TOPI_X86_DEFAULT_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace x86 {
-/*!
- * \brief Helper to create a default x86 schedule for the given ops.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- * \param auto_inline Whether to apply the auto inline step.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule MakeDefaultSchedule(const Target& target, const Array<Tensor>& outs,
- bool auto_inline) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- auto x = outs[0];
- auto axis = s[x]->op.as<ComputeOpNode>()->axis;
-
- if (auto_inline) {
- tvm::te::AutoInlineInjective(s);
- if (axis.size() > 0) {
- detail::Fuse(s[x], axis);
- }
- return s;
- }
-
- if (axis.size() == 4) {
- auto n = axis[0];
- auto c = axis[1];
- auto fused = detail::Fuse(s[x], {n, c}); // for nhwc layout, fuse n and h
- s[x].parallel(fused);
- } else {
- s[x].parallel(axis[0]);
- }
-
- return s;
-}
-
-/*!
- * \brief Create a default x86 schedule for the given ops.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule default_schedule(const Target& target, const Array<Tensor>& outs) {
- return MakeDefaultSchedule(target, outs, false);
-}
-
-/*!
- * \brief Create a default x86 schedule for the given ops, with auto inline
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule default_schedule_auto_inline(const Target& target, const Array<Tensor>& outs) {
- return MakeDefaultSchedule(target, outs, true);
-}
-
-} // namespace x86
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_X86_DEFAULT_H_
diff --git a/include/tvm/topi/x86/injective.h b/include/tvm/topi/x86/injective.h
deleted file mode 100644
index cc1f48c860..0000000000
--- a/include/tvm/topi/x86/injective.h
+++ /dev/null
@@ -1,85 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file x86/injective.h
- * \brief x86 schedule for injective ops
- */
-#ifndef TVM_TOPI_X86_INJECTIVE_H_
-#define TVM_TOPI_X86_INJECTIVE_H_
-
-#include <tvm/target/generic_func.h>
-#include <tvm/te/operation.h>
-#include <tvm/topi/detail/fuse.h>
-#include <tvm/topi/tags.h>
-
-namespace tvm {
-namespace topi {
-
-using namespace tvm::te;
-
-namespace x86 {
-
-/*!
- * \brief Updates an existing schedule for the given injective ops.
- *
- * \param sch The schedule to update.
- * \param out The tensor representing the injective op.
- *
- * \return The updated schedule.
- */
-inline Schedule schedule_injective_from_existing(Schedule sch, const Tensor& out) {
- auto axis = sch[out]->op.as<ComputeOpNode>()->axis;
- if (axis.size() == 4) {
- auto n = axis[0];
- auto c = axis[1];
- auto fused = detail::Fuse(sch[out], {n, c}); // for nhwc layout, fuse n and h
- sch[out].parallel(fused);
- } else if (!axis.empty()) {
- sch[out].parallel(axis[0]);
- }
- return sch;
-}
-
-/*!
- * \brief Create an x86 schedule for the given injective ops.
- *
- * \param target The target to generate a schedule for.
- * \param outs The output tensors.
- *
- * \return A schedule for the given ops.
- */
-inline Schedule schedule_injective(const Target& target, const Array<Tensor>& outs) {
- Array<Operation> out_ops;
- for (auto t : outs) {
- out_ops.push_back(t->op);
- }
- auto s = create_schedule(out_ops);
- tvm::te::AutoInlineInjective(s);
-
- auto x = outs[0];
- schedule_injective_from_existing(s, x);
-
- return s;
-}
-
-} // namespace x86
-} // namespace topi
-} // namespace tvm
-#endif // TVM_TOPI_X86_INJECTIVE_H_
diff --git a/python/tvm/te/__init__.py b/python/tvm/te/__init__.py
index e7b394ebf7..b31853bea6 100644
--- a/python/tvm/te/__init__.py
+++ b/python/tvm/te/__init__.py
@@ -35,4 +35,4 @@ from .operation import thread_axis, reduce_axis, AXIS_SEPARATOR
from .operation import create_prim_func
from .operation import extern_primfunc
-from .tensor import PlaceholderOp, ComputeOp, TensorComputeOp, ScanOp, ExternOp
+from .tensor import PlaceholderOp, ComputeOp, ScanOp, ExternOp
diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py
index a9681c6df0..8d72fc7940 100644
--- a/python/tvm/te/operation.py
+++ b/python/tvm/te/operation.py
@@ -130,26 +130,10 @@ def compute(shape, fcompute, name="compute", tag="", attrs=None, varargs_names=N
dim_var = [tvm.tir.IterVar((0, s), x, 0) for x, s in zip(arg_names, shape[:out_ndim])]
body = fcompute(*[v.var for v in dim_var])
- if isinstance(body, _tensor.TensorIntrinCall):
- for i, s in enumerate(shape[out_ndim:]):
- var_name = "ax" + str(i)
- dim_var.append(tvm.tir.IterVar((0, s), var_name, 4))
- op_node = _ffi_api.TensorComputeOp(
- name,
- tag,
- dim_var,
- body.reduce_axis,
- out_ndim,
- body.intrin,
- body.tensors,
- body.regions,
- body.scalar_inputs,
- )
- else:
- if not isinstance(body, (list, tuple)):
- body = [body]
- body = convert(body)
- op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body)
+ if not isinstance(body, (list, tuple)):
+ body = [body]
+ body = convert(body)
+ op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body)
num = op_node.num_outputs
outputs = tuple(op_node.output(i) for i in range(num))
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index 53ab9d0b5b..70c1287cc8 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -48,11 +48,6 @@ class TensorSlice(ObjectGeneric, _expr.ExprOp):
return self.tensor.dtype
-@tvm._ffi.register_object
-class TensorIntrinCall(Object):
- """Intermediate structure for calling a tensor intrinsic."""
-
-
@tvm._ffi.register_object
class Tensor(DataProducer, _expr.ExprOp):
"""Tensor object, to construct, see function.Tensor"""
@@ -172,11 +167,6 @@ class ComputeOp(BaseComputeOp):
"""Scalar operation."""
-@tvm._ffi.register_object
-class TensorComputeOp(BaseComputeOp):
- """Tensor operation."""
-
-
@tvm._ffi.register_object
class ScanOp(Operation):
"""Scan operation."""
diff --git a/src/driver/driver_api.cc b/src/driver/driver_api.cc
index 86c1d44f01..83ed8d261d 100644
--- a/src/driver/driver_api.cc
+++ b/src/driver/driver_api.cc
@@ -288,54 +288,6 @@ IRModule ApplyPasses(IRModule mod, transform::Sequential seq) {
return mod;
}
-// Convert te schedule to IRModule
-IRModule ScheduleToModule(te::Schedule sch, const Array<ObjectRef>& args, const std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds,
- GlobalVarSupply global_var_supply) {
- sch = sch.normalize();
-
- transform::PassContext pass_ctx = transform::PassContext::Current();
- bool debug_keep_trivial_loop =
- pass_ctx->GetConfig<Bool>("tir.debug_keep_trivial_loop", Bool(false)).value();
-
- // Before TIR transformation.
- tir::Stmt stmt = te::ScheduleOps(sch, te::InferBound(sch), debug_keep_trivial_loop);
- bool compact = te::VerifyCompactBuffer(stmt);
-
- Map<te::Tensor, tir::Buffer> out_binds;
- Array<ObjectRef> out_arg_list;
- GetBinds(args, compact, binds, &out_binds, &out_arg_list);
-
- // Build the function, converting from te::Tensor to tir::Buffer
- tir::PrimFunc f = te::SchedulePostProcToPrimFunc(out_arg_list, std::move(stmt), out_binds);
- f = WithAttr(std::move(f), "global_symbol", runtime::String(name));
-
- // Mark this schedule as being converted from an TE schedule. Makes sure that
- // the correct TE passes are run.
- f = WithAttr(std::move(f), "from_legacy_te_schedule", Bool(true));
-
- bool noalias = pass_ctx->GetConfig<Bool>("tir.noalias", Bool(true)).value();
-
- if (noalias) {
- f = WithAttr(std::move(f), "tir.noalias", Bool(true));
- }
- GlobalVar global_var = global_var_supply->UniqueGlobalFor(name, false);
- return IRModule(Map<GlobalVar, BaseFunc>({{global_var, f}}));
-}
-
-TVM_REGISTER_GLOBAL("driver.schedule_to_module")
- .set_body_typed([](te::Schedule sch, const Array<ObjectRef>& args, const String& name,
- const Map<te::Tensor, tir::Buffer>& binds) {
- std::unordered_map<te::Tensor, tir::Buffer> c_binds;
- // Check to make sure binds is not null before doing the conversion;
- if (binds.defined()) {
- for (auto kv : binds) {
- c_binds.insert({kv.first, kv.second});
- }
- }
- IRModule mod = ScheduleToModule(std::move(sch), args, name, c_binds, GlobalVarSupply());
- return mod;
- });
IRModule LowerModule(IRModule mod, bool simple_mode) {
Array<transform::Pass> pass_list = CreatePassList(simple_mode);
@@ -367,38 +319,6 @@ TVM_REGISTER_GLOBAL("driver.lower_primfunc")
return LowerPrimFunc(std::move(func), name, simple_mode);
});
-IRModule LowerSchedule(te::Schedule sch, const Array<te::Tensor>& args, const std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds,
- GlobalVarSupply global_var_supply, bool simple_mode) {
- Array<ObjectRef> ref_args;
- for (ObjectRef x : args) {
- ref_args.push_back(x);
- }
- return LowerSchedule(std::move(sch), ref_args, name, binds, global_var_supply, simple_mode);
-}
-
-IRModule LowerSchedule(te::Schedule sch, const Array<ObjectRef>& args, const std::string& name,
- const std::unordered_map<te::Tensor, tir::Buffer>& binds,
- GlobalVarSupply global_var_supply, bool simple_mode) {
- IRModule mod = ScheduleToModule(std::move(sch), args, name, binds, global_var_supply);
- // Get the legacy TE pass list
- Array<transform::Pass> pass_list = CreatePassList(simple_mode);
- return LowerWithPassList(mod, pass_list);
-}
-
-TVM_REGISTER_GLOBAL("driver.lower_schedule")
- .set_body_typed([](te::Schedule sch, const Array<ObjectRef>& args, const String& name,
- const Map<te::Tensor, tir::Buffer>& binds, bool simple_mode) {
- std::unordered_map<te::Tensor, tir::Buffer> c_binds;
- // Check to make sure binds is not null before doing the conversion;
- if (binds.get() != nullptr) {
- for (auto kv : binds) {
- c_binds.insert({kv.first, kv.second});
- }
- }
- return LowerSchedule(std::move(sch), args, name, c_binds, GlobalVarSupply(), simple_mode);
- });
-
/**
* This function takes the input module that contains both the device and host opts.
* Then, it applies transformation on the original module before splitting into separate modules for
diff --git a/src/relax/analysis/layout_transformation.cc b/src/relax/analysis/layout_transformation.cc
index 25ce1a58f2..f0658dabb3 100644
--- a/src/relax/analysis/layout_transformation.cc
+++ b/src/relax/analysis/layout_transformation.cc
@@ -26,6 +26,7 @@
#include <tvm/arith/iter_affine_map.h>
#include <tvm/relax/analysis.h>
#include <tvm/tir/analysis.h>
+#include <tvm/tir/index_map.h>
#include <tvm/tir/stmt_functor.h>
#include "../../support/array.h"
diff --git a/src/target/codegen.cc b/src/target/codegen.cc
index 49c250b8dc..e5e2c4297c 100644
--- a/src/target/codegen.cc
+++ b/src/target/codegen.cc
@@ -43,17 +43,6 @@
namespace tvm {
namespace codegen {
-/*!
- * \brief TIRToRuntime conversion specific to a TargetKind
- *
- * This function is responsible for scanning an IRModule for appropriate Target-specific functions
- and generating a Runtime module representing the compiled output
- *
- * \param ir_module Unified IRModule
- * \param target Target to filter on or retrieve arguments from
- * \return Runtime Module containing compiled functions
- */
-using FTVMTIRToRuntime = tvm::runtime::TypedPackedFunc<runtime::Module(IRModule, Target)>;
runtime::Module Build(IRModule mod, Target target) {
if (transform::PassContext::Current()
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 1f9cc09b1e..043fc26daf 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -1819,7 +1819,7 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const BufferLoadNode* op) {
auto make_load = [this, &loads](TypedPointer buffer_ptr, int /* subelement_i */,
llvm::Value* predicate, int alignment, bool is_volatile) {
llvm::Instruction* load = nullptr;
- if (predicate != NULL) {
+ if (predicate != nullptr) {
ICHECK(!is_volatile)
<< "The masked load intrinsic does not support declaring load as volatile.";
#if TVM_LLVM_VERSION >= 130
@@ -1972,7 +1972,7 @@ void CodeGenLLVM::VisitStmt_(const BufferStoreNode* op) {
to_store = builder_->CreateExtractElement(value, subelement_i);
}
- if (predicate != NULL) {
+ if (predicate != nullptr) {
ICHECK(!is_volatile)
<< "The masked store intrinsic does not support declaring store as volatile.";
#if TVM_LLVM_VERSION >= 110
diff --git a/src/te/autodiff/ad_simplify.cc b/src/te/autodiff/ad_simplify.cc
deleted file mode 100644
index 26047e879e..0000000000
--- a/src/te/autodiff/ad_simplify.cc
+++ /dev/null
@@ -1,1239 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file ad_simplify.cc
- * \brief Simplify tensor compute generated by tensor-level autodiff.
- *
- * The major simplification we do in this file is to eliminate
- * the Jacobian tensor created by autodiff.
- *
- * Jacobian tensor is sparse because one output element usually relates
- * to a small portion of the inputs. For example, element-wise function has a one-to-one mapping
- * between input tensor and output tensor, thus the Jacobian is diagonal.
- *
- * Generally, we have Out_{\beta} = f( In_{A \alpha} ) in which A is a matrix,
- * \alpha and \beta are vectors represent the indices of In and Out respectively.
- * i.e., the non-zero Jacobian indices is a linear combination of the input indices.
- * Thereby we solve linear equations of \beta = A \alpha,
- * as well as linear inequalities of their domain ranges.
- *
- * Refer to Urban S, van der Smagt P. Automatic differentiation for tensor algebras[J].
- * arXiv preprint arXiv:1711.01348, 2017. for more details.
- *
- * Implement-wise, we extract the equations in the compute definition via NonzeronessCondition,
- * replace the compute expression with solved new axes, and create a selection node
- * (non-zero-condition ? new_compute_expression : 0).
- *
- * Due to TVM's restriction, we also lift the reduction to the top of the compute stage.
- *
- */
-#include <tvm/arith/analyzer.h>
-#include <tvm/arith/int_solver.h>
-#include <tvm/runtime/registry.h>
-#include <tvm/te/autodiff.h>
-#include <tvm/tir/analysis.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <iterator>
-#include <memory>
-#include <optional>
-#include <utility>
-
-#include "ad_utils.h"
-
-namespace tvm {
-namespace te {
-
-using arith::DivMode;
-using arith::kFloorDiv;
-using arith::kSimplifyRewriteCanonicalRewrite;
-using arith::kTruncDiv;
-
-// Combine all expressions from the container using &&.
-template <class container>
-PrimExpr All(const container& c) {
- PrimExpr res;
- for (const auto& e : c) {
- if (res.get()) {
- res = res && e;
- } else {
- res = e;
- }
- }
- if (res.get()) {
- return res;
- } else {
- return const_true();
- }
-}
-
-Map<Var, Range> IterVarsToMap(const Array<IterVar>& itervars) {
- Map<Var, Range> res;
- for (const IterVar& v : itervars) {
- res.Set(v->var, v->dom);
- }
- return res;
-}
-
-// Given a map from vars to ranges create an array of itervars
-Array<IterVar> IterVarsFromMap(const Array<Var>& vars, const Map<Var, Range>& vranges,
- IterVarType iter_type = kDataPar, std::string thread_tag = "") {
- Array<IterVar> res;
- for (const Var& v : vars) {
- ICHECK(vranges.count(v)) << "A range for the variable " << v << " was not provided in map "
- << vranges;
- res.push_back(IterVar(vranges[v], v, iter_type, thread_tag));
- }
- return res;
-}
-
-Array<Var> IterVarsToVars(const Array<IterVar>& itervars) {
- Array<Var> res;
- for (const IterVar& v : itervars) {
- res.push_back(v->var);
- }
- return res;
-}
-
-template <typename ValueType>
-bool is_const_value(const PrimExpr& e, ValueType value) {
- static_assert(std::is_integral<ValueType>::value,
- "Comparison to non-integer values is forbidden.");
- if (const tir::IntImmNode* i = e.as<tir::IntImmNode>()) {
- return i->value == value;
- } else if (const tir::FloatImmNode* i = e.as<tir::FloatImmNode>()) {
- return i->value == value;
- } else if (const tir::CastNode* c = e.as<tir::CastNode>()) {
- return is_const_value(c->value, value);
- } else if (const tir::BroadcastNode* b = e.as<tir::BroadcastNode>()) {
- return is_const_value(b->value, value);
- } else {
- return false;
- }
-}
-
-// Return true if this combiner is just a sum.
-bool IsSumCombiner(const CommReducer& combiner, const Map<Var, Range>& vranges) {
- arith::Analyzer analyzer;
- analyzer.Bind(vranges);
- if (combiner->result.size() != 1) {
- return false;
- }
-
- if (!is_const_value(
- analyzer.Simplify(combiner->identity_element[0], kSimplifyRewriteCanonicalRewrite), 0)) {
- return false;
- }
-
- PrimExpr combiner_result =
- analyzer.Simplify(combiner->result[0], kSimplifyRewriteCanonicalRewrite);
-
- return tir::ExprDeepEqual()(combiner_result, combiner->lhs[0] + combiner->rhs[0]) ||
- tir::ExprDeepEqual()(combiner_result, combiner->rhs[0] + combiner->lhs[0]);
-}
-
-bool CanFactorZeroFromCombiner(const CommReducer& combiner, int value_index,
- const Map<Var, Range>& vranges) {
- arith::Analyzer analyzer;
- analyzer.Bind(vranges);
- if (!is_const_value(analyzer.Simplify(combiner->identity_element[value_index],
- kSimplifyRewriteCanonicalRewrite),
- 0)) {
- return false;
- }
-
- PrimExpr zero = make_zero(combiner->result[value_index].dtype());
- PrimExpr in = Substitute(combiner->result[value_index], {{combiner->lhs[value_index], zero},
- {combiner->rhs[value_index], zero}});
- in = analyzer.Simplify(in, kSimplifyRewriteCanonicalRewrite);
-
- return is_const_value(in, 0);
-}
-
-struct NonzeroConditionResult {
- PrimExpr cond;
- PrimExpr value;
-
- PrimExpr to_expr() const { return Select(cond, value, make_zero(value.dtype())); }
-
- friend std::ostream& operator<<(std::ostream& os, const NonzeroConditionResult& r) {
- return os << r.to_expr();
- }
-};
-
-// The implementation of NonzeroCondition
-// transform expression to cond ? value : 0
-class NonzeroConditionFunctor : public ExprFunctor<NonzeroConditionResult(const PrimExpr&)> {
- public:
- NonzeroConditionResult NonzeroCondition(const PrimExpr& e) {
- if (e.dtype().is_bool()) {
- // Boolean expressions are non-zero whenever they are true themselves
- return {e, const_true()};
- } else {
- return VisitExpr(e);
- }
- }
-
- // Most of the cases are implemented using helpers below
- result_type VisitExpr_(const VarNode* op) final { return Default_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const IntImmNode* op) final { return Const_(GetRef<IntImm>(op)); }
- result_type VisitExpr_(const FloatImmNode* op) final { return Const_(GetRef<FloatImm>(op)); }
- result_type VisitExpr_(const StringImmNode* op) final { return Default_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const AddNode* op) final { return BinOpAddLike_(GetRef<Add>(op)); }
- result_type VisitExpr_(const SubNode* op) final { return BinOpAddLike_(GetRef<Sub>(op)); }
- result_type VisitExpr_(const MulNode* op) final { return BinOpMulLike_(GetRef<Mul>(op)); }
- result_type VisitExpr_(const DivNode* op) final { return BinOpDivLike_(GetRef<Div>(op)); }
- result_type VisitExpr_(const ModNode* op) final { return BinOpDivLike_(GetRef<Mod>(op)); }
- result_type VisitExpr_(const FloorDivNode* op) final {
- return BinOpDivLike_(GetRef<FloorDiv>(op));
- }
- result_type VisitExpr_(const FloorModNode* op) final {
- return BinOpDivLike_(GetRef<FloorMod>(op));
- }
- result_type VisitExpr_(const MinNode* op) final { return BinOpAddLike_(GetRef<Min>(op)); }
- result_type VisitExpr_(const MaxNode* op) final { return BinOpAddLike_(GetRef<Max>(op)); }
-
- result_type VisitExpr_(const CastNode* op) final {
- auto nz_a = NonzeroCondition(op->value);
- return {nz_a.cond, Cast(op->dtype, nz_a.value)};
- }
-
- result_type VisitExpr_(const SelectNode* op) final {
- PrimExpr cond = op->condition, true_val = op->true_value, false_val = op->false_value;
- auto nz_a = NonzeroCondition(true_val);
- auto nz_b = NonzeroCondition(false_val);
-
- // If the false part is zero, we can get rid of the select
- if (is_const_value(nz_b.value, 0)) {
- PrimExpr new_cond = analyzer_.Simplify(nz_a.cond && cond, kSimplifyRewriteCanonicalRewrite);
- return {new_cond, nz_a.value};
- }
-
- // If the true part is zero, we can also get rid of the select
- if (is_const_value(nz_a.value, 0)) {
- PrimExpr new_cond = analyzer_.Simplify(nz_b.cond && !cond, kSimplifyRewriteCanonicalRewrite);
- return {new_cond, nz_b.value};
- }
-
- // Otherwise we retain the select and combine the conditions into this
- PrimExpr new_cond = analyzer_.Simplify((cond && nz_a.cond) || (!cond && nz_b.cond),
- kSimplifyRewriteCanonicalRewrite);
- if (nz_a.value.same_as(true_val) && nz_b.value.same_as(false_val)) {
- return {new_cond, GetRef<PrimExpr>(op)};
- } else {
- return {new_cond, Select(cond, nz_a.value, nz_b.value)};
- }
- }
-
- result_type VisitExpr_(const CallNode* op) final {
- if (op->op.same_as(op_if_then_else_)) {
- PrimExpr cond = op->args[0], true_val = op->args[1], false_val = op->args[2];
- auto nz_a = NonzeroCondition(true_val);
- auto nz_b = NonzeroCondition(false_val);
-
- // We don't have as much freedom here as in the select case
- // since the `if` must be preserved in any case
- PrimExpr new_cond = analyzer_.Simplify((cond && nz_a.cond) || (!cond && nz_b.cond),
- kSimplifyRewriteCanonicalRewrite);
- if (nz_a.value.same_as(true_val) && nz_b.value.same_as(false_val)) {
- return {new_cond, GetRef<PrimExpr>(op)};
- } else {
- return {new_cond, if_then_else(cond, nz_a.value, nz_b.value)};
- }
- } else {
- return Default_(GetRef<PrimExpr>(op));
- }
- }
-
- result_type VisitExpr_(const ProducerLoadNode* op) final {
- return Default_(GetRef<PrimExpr>(op));
- }
-
- NonzeroConditionResult Default_(const PrimExpr& e) {
- // This is always correct, so it's the default
- return {const_true(), e};
- }
-
- template <class T>
- NonzeroConditionResult Const_(const T& op) {
- if (op->value == 0) {
- return {const_false(), op};
- } else {
- return {const_true(), op};
- }
- }
-
- template <class T>
- NonzeroConditionResult BinOpAddLike_(const T& op) {
- auto nz_a = NonzeroCondition(op->a);
- auto nz_b = NonzeroCondition(op->b);
-
- // For addition and similar ops the result may be nonzero if either of the arguments is
- // nonzero, so we combine the conditions with Or.
- if (tir::ExprDeepEqual()(nz_a.cond, nz_b.cond)) {
- // If the conditions are the same, we don't need Or
- if (nz_a.value.same_as(op->a) && nz_b.value.same_as(op->b)) {
- return {nz_a.cond, op};
- } else {
- return {nz_a.cond, T(nz_a.value, nz_b.value)};
- }
- } else {
- // Otherwise use Or
- PrimExpr new_cond =
- analyzer_.Simplify(nz_a.cond || nz_b.cond, kSimplifyRewriteCanonicalRewrite);
- // A little optimization: if the combined condition is the same as one of the inner
- // conditions, we don't need to guard the inner value with a select, otherwise
- // we create a select in the `to_expr` call.
- PrimExpr new_a = tir::ExprDeepEqual()(nz_a.cond, new_cond) ? nz_a.value : nz_a.to_expr();
- PrimExpr new_b = tir::ExprDeepEqual()(nz_b.cond, new_cond) ? nz_b.value : nz_b.to_expr();
- PrimExpr new_expr = T(new_a, new_b);
- return {new_cond, new_expr};
- }
- }
-
- template <class T>
- NonzeroConditionResult BinOpMulLike_(const T& op) {
- auto nz_a = NonzeroCondition(op->a);
- auto nz_b = NonzeroCondition(op->b);
-
- // For multiplication and similar ops the result may be nonzero if
- // both the arguments are nonzero, so we combine with And.
- PrimExpr new_cond =
- analyzer_.Simplify(nz_a.cond && nz_b.cond, kSimplifyRewriteCanonicalRewrite);
-
- if (nz_a.value.same_as(op->a) && nz_b.value.same_as(op->b)) {
- return {new_cond, op};
- } else {
- return {new_cond, T(nz_a.value, nz_b.value)};
- }
- }
-
- template <class T>
- NonzeroConditionResult BinOpDivLike_(const T& op) {
- auto nz_a = NonzeroCondition(op->a);
-
- // For Div we simply use the condition of the numerator.
-
- if (nz_a.value.same_as(op->a)) {
- return {nz_a.cond, op};
- } else {
- return {nz_a.cond, T(nz_a.value, op->b)};
- }
- }
-
- private:
- arith::Analyzer analyzer_;
- const Op& op_if_then_else_ = Op::Get("tir.if_then_else");
-};
-
-inline NonzeroConditionResult NonzeronessCondition(const PrimExpr& expr) {
- return NonzeroConditionFunctor().NonzeroCondition(expr);
-}
-
-struct FactorOutAtomicFormulasResult {
- std::vector<PrimExpr> atomic_formulas;
- PrimExpr rest;
-
- PrimExpr to_expr() const {
- PrimExpr res = rest;
- for (const PrimExpr& e : atomic_formulas) {
- res = And(e, res);
- }
- return res;
- }
-
- Array<PrimExpr> to_array() const {
- Array<PrimExpr> res = atomic_formulas;
- res.push_back(rest);
- return res;
- }
-};
-
-// The implementation of FactorOutAtomicFormulas
-class FactorOutAtomicFormulasFunctor
- : public ExprFunctor<FactorOutAtomicFormulasResult(const PrimExpr&)> {
- public:
- result_type Atomic_(const PrimExpr& e) {
- // For atomic expressions the result is the expr itself with True as the residual
- return {{e}, make_const(e.dtype(), 1)};
- }
-
- // This is basically the list of expression kinds that are considered atomic
- result_type VisitExpr_(const VarNode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const CallNode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const IntImmNode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const EQNode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const NENode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const LENode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const LTNode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const GENode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
- result_type VisitExpr_(const GTNode* op) final { return Atomic_(GetRef<PrimExpr>(op)); }
-
- result_type VisitExpr_(const SelectNode* op) final {
- // Select can be rewritten through other logical ops
- PrimExpr expr = (op->condition && op->true_value) || (!op->condition && op->false_value);
- return VisitExpr(expr);
- }
-
- result_type VisitExpr_(const NotNode* op) final {
- // Not should be moved down
- if (const OrNode* or_expr = op->a.as<OrNode>()) {
- PrimExpr expr = !or_expr->a && !or_expr->b;
- return VisitExpr(expr);
- } else if (const AndNode* and_expr = op->a.as<AndNode>()) {
- PrimExpr expr = !and_expr->a || !and_expr->b;
- return VisitExpr(expr);
- } else if (const SelectNode* sel_expr = op->a.as<SelectNode>()) {
- PrimExpr expr = ((!sel_expr->condition || !sel_expr->true_value) &&
- (sel_expr->condition || !sel_expr->false_value));
- return VisitExpr(expr);
- }
- return Atomic_(GetRef<PrimExpr>(op));
- }
-
- result_type VisitExpr_(const AndNode* op) final {
- auto res_a = VisitExpr(op->a);
- auto res_b = VisitExpr(op->b);
-
- // For the And case we return the union of the sets of atomic formulas
- std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_a_set;
- res_a_set.reserve(res_a.atomic_formulas.size());
- std::copy(res_a.atomic_formulas.begin(), res_a.atomic_formulas.end(),
- std::inserter(res_a_set, res_a_set.end()));
-
- std::vector<PrimExpr> res = res_a.atomic_formulas;
- for (const auto& e : res_b.atomic_formulas) {
- if (res_a_set.find(e) == res_a_set.end()) {
- res.emplace_back(e);
- }
- }
- // And the residuals are combined with &&
- return {res, res_a.rest && res_b.rest};
- }
-
- result_type VisitExpr_(const MulNode* op) final {
- // Since we work with bools, for multiplication we do the same thing as for And
- PrimExpr e_and = op->a && op->b;
- return VisitExpr(e_and);
- }
-
- result_type VisitExpr_(const OrNode* op) final {
- auto res_a = VisitExpr(op->a);
- auto res_b = VisitExpr(op->b);
-
- std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_a_set{
- res_a.atomic_formulas.begin(), res_a.atomic_formulas.end()};
- std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_b_set{
- res_b.atomic_formulas.begin(), res_b.atomic_formulas.end()};
-
- // For the Or case we intersect the sets of atomic formulas
- std::unordered_set<PrimExpr, StructuralHash, StructuralEqual> res_set;
- std::vector<PrimExpr> res;
- res_set.reserve(std::min(res_a.atomic_formulas.size(), res_b.atomic_formulas.size()));
- res.reserve(std::min(res_a.atomic_formulas.size(), res_b.atomic_formulas.size()));
- for (const auto& res_b_formula : res_b.atomic_formulas) {
- if (res_a_set.count(res_b_formula)) {
- res_set.insert(res_b_formula);
- res.push_back(res_b_formula);
- }
- }
-
- // Computing the residual is more complex: we have to compute the sets of atomic formulas
- // which are left behind, and then combine them with the residuals into the new residual.
- std::vector<PrimExpr> new_cond_a;
- new_cond_a.reserve(res_a.atomic_formulas.size() - res_set.size());
- for (const auto& formula : res_a.atomic_formulas) {
- if (!res_set.count(formula)) new_cond_a.emplace_back(formula);
- }
-
- std::vector<PrimExpr> new_cond_b;
- new_cond_b.reserve(res_b.atomic_formulas.size() - res_set.size());
- for (const auto& formula : res_b.atomic_formulas) {
- if (!res_set.count(formula)) new_cond_b.emplace_back(formula);
- }
-
- res_a.atomic_formulas = std::move(new_cond_a);
- res_b.atomic_formulas = std::move(new_cond_b);
-
- PrimExpr new_rest = res_a.to_expr() || res_b.to_expr();
-
- return {res, new_rest};
- }
-};
-
-// Transform the given formula into a conjunction of atomic formulas (represented as an array)
-// and a non-atomic residual. Atomic formulas are consts, calls, variables and comparisons (a <= b,
-// etc), i.e. formulas which are not logical operators (||, &&, !) on the top level.
-FactorOutAtomicFormulasResult FactorOutAtomicFormulas(const PrimExpr& e) {
- ICHECK(e.dtype().is_bool());
- return FactorOutAtomicFormulasFunctor().VisitExpr(e);
-}
-
-struct EliminateDivModResult {
- PrimExpr expr;
- Map<Var, PrimExpr> substitution;
- Array<Var> new_variables;
- Array<PrimExpr> conditions;
- Map<Var, Range> ranges;
-};
-
-inline PrimExpr ModImpl(PrimExpr a, PrimExpr b, DivMode mode) {
- if (mode == kTruncDiv) {
- return truncmod(a, b);
- } else {
- ICHECK_EQ(mode, kFloorDiv);
- return floormod(a, b);
- }
-}
-
-inline PrimExpr DivImpl(PrimExpr a, PrimExpr b, DivMode mode) {
- if (mode == kTruncDiv) {
- return truncdiv(a, b);
- } else {
- ICHECK_EQ(mode, kFloorDiv);
- return floordiv(a, b);
- }
-}
-
-class EliminateDivModMutator : public ExprMutator {
- public:
- Map<Var, PrimExpr> substitution;
- Array<Var> new_variables;
- Array<PrimExpr> conditions;
- Map<Var, Range> ranges;
-
- explicit EliminateDivModMutator(Map<Var, Range> ranges) : ranges(std::move(ranges)) {}
-
- virtual PrimExpr VisitExpr_(const DivNode* op) {
- const IntImmNode* imm = op->b.as<IntImmNode>();
- if (imm && imm->value != 0) {
- if (imm->value < 0) {
- // x / -c == -(x/c) for truncated division
- return make_zero(op->dtype) -
- VisitExpr(truncdiv(op->a, make_const(op->dtype, -imm->value)));
- }
-
- // Try to find the already existing variables for this expression
- auto it = expr_to_vars_.find(std::make_tuple(kTruncDiv, op->a, imm->value));
- if (it != expr_to_vars_.end()) {
- return it->second.first;
- }
-
- // Otherwise recursively mutate the left hand side, and create new variables
- PrimExpr mutated_a = VisitExpr(op->a);
- if (auto var_pair_opt = AddNewVarPair(op->a, mutated_a, imm->value, kTruncDiv)) {
- return var_pair_opt.value().first;
- } else {
- return truncdiv(mutated_a, op->b);
- }
- }
-
- return div(VisitExpr(op->a), VisitExpr(op->b));
- }
-
- virtual PrimExpr VisitExpr_(const ModNode* op) {
- const IntImmNode* imm = op->b.as<IntImmNode>();
- if (imm && imm->value != 0) {
- if (imm->value < 0) {
- // x % -c == x % c for truncated division
- return VisitExpr(truncmod(op->a, make_const(op->dtype, -imm->value)));
- }
-
- // Try to find the already existing variables for this expression
- auto it = expr_to_vars_.find(std::make_tuple(kTruncDiv, op->a, imm->value));
- if (it != expr_to_vars_.end()) {
- return it->second.second;
- }
-
- // Otherwise recursively mutate the left hand side, and create new variables
- PrimExpr mutated_a = VisitExpr(op->a);
- if (auto var_pair_opt = AddNewVarPair(op->a, mutated_a, imm->value, kTruncDiv)) {
- return var_pair_opt.value().second;
- } else {
- return truncmod(mutated_a, op->b);
- }
- }
-
- return truncmod(VisitExpr(op->a), VisitExpr(op->b));
- }
-
- virtual PrimExpr VisitExpr_(const FloorDivNode* op) {
- const IntImmNode* imm = op->b.as<IntImmNode>();
- if (imm && imm->value != 0) {
- if (imm->value < 0) {
- // x / -c == (-x) / c for flooring division
- return VisitExpr(
- floordiv(make_zero(op->dtype) - op->a, make_const(op->dtype, -imm->value)));
- }
-
- // Try to find the already existing variables for this expression
- auto it = expr_to_vars_.find(std::make_tuple(kFloorDiv, op->a, imm->value));
- if (it != expr_to_vars_.end()) {
- return it->second.first;
- }
-
- // Otherwise recursively mutate the left hand side, and create new variables
- PrimExpr mutated_a = VisitExpr(op->a);
- if (auto var_pair_opt = AddNewVarPair(op->a, mutated_a, imm->value, kFloorDiv)) {
- return var_pair_opt.value().first;
- } else {
- return floordiv(mutated_a, op->b);
- }
- }
-
- return floordiv(VisitExpr(op->a), VisitExpr(op->b));
- }
-
- virtual PrimExpr VisitExpr_(const FloorModNode* op) {
- const IntImmNode* imm = op->b.as<IntImmNode>();
- if (imm && imm->value != 0) {
- if (imm->value < 0) {
- // x % -c == -(-x % c) for flooring division
- return VisitExpr(make_zero(op->dtype) - floormod(make_zero(op->dtype) - op->a,
- make_const(op->dtype, -imm->value)));
- }
-
- // Try to find the already existing variables for this expression
- auto it = expr_to_vars_.find(std::make_tuple(kFloorDiv, op->a, imm->value));
- if (it != expr_to_vars_.end()) {
- return it->second.second;
- }
-
- // Otherwise recursively mutate the left hand side, and create new variables
- PrimExpr mutated_a = VisitExpr(op->a);
- if (auto var_pair_opt = AddNewVarPair(op->a, mutated_a, imm->value, kFloorDiv)) {
- return var_pair_opt.value().second;
- } else {
- return floormod(mutated_a, op->b);
- }
- }
-
- return floormod(VisitExpr(op->a), VisitExpr(op->b));
- }
-
- private:
- std::optional<std::pair<Var, Var>> AddNewVarPair(const PrimExpr& e, const PrimExpr& mut,
- int64_t val, DivMode mode) {
- using tresult = std::optional<std::pair<Var, Var>>;
-
- // Try to find the variables using the mutated expressions
- if (!e.same_as(mut)) {
- auto it = expr_to_vars_.find(std::make_tuple(mode, mut, val));
- if (it != expr_to_vars_.end()) {
- return tresult(it->second);
- }
- }
-
- PrimExpr val_e = make_const(e.dtype(), val);
- idx_ += 1;
-
- // Convert `ranges` to IntSets
- std::unordered_map<const VarNode*, IntSet> var_intsets;
- for (const auto& p : ranges) {
- var_intsets[p.first.get()] = IntSet::FromRange(p.second);
- }
-
- // Infer ranges for the expressions we want to replace with variables
- Range div_range = EvalSet(DivImpl(mut, val_e, mode), var_intsets).CoverRange(Range());
- Range mod_range = EvalSet(ModImpl(mut, val_e, mode), var_intsets).CoverRange(Range());
-
- // We don't want to add unbounded variables
- if (!div_range.get() || !mod_range.get()) {
- LOG(WARNING) << "EliminateDivMod: won't eliminate " << DivImpl(e, val_e, mode)
- << " because its bounds cannot be inferred";
- return tresult();
- }
- if (!mod_range.get()) {
- LOG(WARNING) << "EliminateDivMod: won't eliminate " << ModImpl(e, val_e, mode)
- << " because its bounds cannot be inferred";
- return tresult();
- }
-
- // Create new variables for the expressions
- auto div = Var((mode == kTruncDiv ? "tdiv" : "fdiv") + std::to_string(idx_), e.dtype());
- auto mod = Var((mode == kTruncDiv ? "tmod" : "fmod") + std::to_string(idx_), e.dtype());
-
- new_variables.push_back(div);
- new_variables.push_back(mod);
-
- // Note that we have to perform substitution to mut because mut may contain new variables
- substitution.Set(div, DivImpl(Substitute(mut, substitution), val_e, mode));
- substitution.Set(mod, ModImpl(Substitute(mut, substitution), val_e, mode));
-
- ranges.Set(div, div_range);
- ranges.Set(mod, mod_range);
-
- // This additional condition works as a definition for the new variables
- conditions.push_back(mut == div * val_e + mod);
-
- if (!analyzer_.CanProve(mod_range->extent <= val_e)) {
- // If we use the C/C++ definition of mod, there may be multiple values of `mod`
- // satisfying the added condition if the expr `e` may change its sign, so we
- // have to add another condition.
- LOG(WARNING) << "EliminateDivMod: cannot fully eliminate div or mod because "
- << ModImpl(e, val_e, mode) << " probably may change its sign";
- conditions.push_back(Select(e >= 0, mod >= 0, mod <= 0));
- }
-
- auto p = std::make_pair(div, mod);
- expr_to_vars_[std::make_tuple(mode, e, val)] = p;
- if (!e.same_as(mut)) {
- expr_to_vars_[std::make_tuple(mode, mut, val)] = p;
- }
- return tresult(p);
- }
-
- class TupleEqual_ {
- public:
- bool operator()(const std::tuple<DivMode, PrimExpr, int64_t>& lhs,
- const std::tuple<DivMode, PrimExpr, int64_t>& rhs) const {
- return std::get<0>(lhs) == std::get<0>(rhs) &&
- tir::ExprDeepEqual()(std::get<1>(lhs), std::get<1>(rhs)) &&
- std::get<2>(lhs) == std::get<2>(rhs);
- }
- };
-
- class TupleHasher_ {
- public:
- size_t operator()(const std::tuple<DivMode, PrimExpr, int64_t>& key) const {
- return ((std::hash<int>()(std::get<0>(key)) ^ (StructuralHash()(std::get<1>(key)) << 1)) >>
- 1) ^
- (std::hash<int64_t>()(std::get<2>(key)) << 1);
- }
- };
-
- // A counter for naming new variables
- int idx_{0};
- // A map from pairs of exprs and numbers (e, n) to pairs of new vars (div, mod)
- // such that `div = e / n` and `mod = e % n`
- std::unordered_map<std::tuple<DivMode, PrimExpr, int64_t>, std::pair<Var, Var>, TupleHasher_,
- TupleEqual_>
- expr_to_vars_;
- arith::Analyzer analyzer_;
-};
-
-// Replace every subexpr of the form e/const and e % const with a new variable.
-// Syntactically equal expressions will be mapped to the same variable.
-EliminateDivModResult EliminateDivMod(const PrimExpr& expr, Map<Var, Range> ranges) {
- EliminateDivModResult res;
- EliminateDivModMutator mutator(ranges);
- res.expr = mutator(expr);
- res.conditions = std::move(mutator.conditions);
- res.new_variables = std::move(mutator.new_variables);
- res.substitution = std::move(mutator.substitution);
- res.ranges = std::move(mutator.ranges);
- return res;
-}
-
-arith::IntConstraintsTransform EliminateDivModFromDomainConditions(
- const arith::IntConstraints& domain) {
- auto elim_res = EliminateDivMod(All(domain->relations), domain->ranges);
-
- Map<Var, Range> new_vranges = elim_res.ranges;
- Array<Var> new_axis = Concat(domain->variables, elim_res.new_variables);
- PrimExpr new_cond = elim_res.expr && All(elim_res.conditions);
-
- arith::IntConstraints new_domain(new_axis, new_vranges,
- FactorOutAtomicFormulas(new_cond).to_array());
-
- Map<Var, PrimExpr> src_to_dst;
- Map<Var, PrimExpr> dst_to_src = elim_res.substitution;
- for (const Var& v : domain->variables) {
- src_to_dst.Set(v, v);
- dst_to_src.Set(v, v);
- }
-
- return arith::IntConstraintsTransform(domain, new_domain, src_to_dst, dst_to_src);
-}
-
-inline arith::IntConstraintsTransform IdentityTransformation(const arith::IntConstraints& domain) {
- Map<Var, PrimExpr> identity_map;
- for (const Var& v : domain->variables) {
- identity_map.Set(v, v);
- }
- return arith::IntConstraintsTransform(domain, domain, identity_map, identity_map);
-}
-
-// Simplify an iteration domain.
-arith::IntConstraintsTransform SimplifyDomain(const arith::IntConstraints& iter_domains,
- bool eliminate_div_mod) {
- arith::IntConstraintsTransform transf = IdentityTransformation(iter_domains);
-
- if (eliminate_div_mod) {
- transf = transf + EliminateDivModFromDomainConditions(transf->dst);
- }
-
- // TODO(sgrechanik-h): Repeating the following steps has a positive effect, however we probably
- // should find a better terminating criterion (like stop when the domain volume stops decreasing)
- // Also 2 steps seems to be slightly better than 3
- for (size_t i = 0; i < 2; ++i) {
- transf = transf + arith::SolveLinearEquations(transf->dst);
- transf = transf + arith::SolveInequalitiesDeskewRange(transf->dst);
- }
-
- return transf;
-}
-
-// Use the condition of a reduction op to simplify its domain (axis)
-PrimExpr SimplifyReductionDomain(const PrimExpr& expr, const Map<Var, Range>& outer_vranges) {
- if (const ReduceNode* red = expr.as<ReduceNode>()) {
- Array<Var> vars = IterVarsToVars(red->axis);
- Map<Var, Range> vranges = Merge(outer_vranges, IterVarsToMap(red->axis));
- Array<PrimExpr> relations = FactorOutAtomicFormulas(red->condition).to_array();
-
- arith::IntConstraints domain(vars, vranges, relations);
- auto res = SimplifyDomain(domain);
-
- Array<PrimExpr> new_source;
- for (const PrimExpr& src : red->source) {
- new_source.push_back(Substitute(src, res->src_to_dst));
- }
-
- Array<IterVar> new_axis = IterVarsFromMap(res->dst->variables, res->dst->ranges, kCommReduce);
-
- // Perform simplification mainly to remove a possibly empty reduction.
- arith::Analyzer analyzer;
- return analyzer.Simplify(Reduce(red->combiner, new_source, new_axis, All(res->dst->relations),
- red->value_index, red->init),
- kSimplifyRewriteCanonicalRewrite);
- } else {
- return expr;
- }
-}
-
-// Extract from cond an implication of cond not containing vars
-std::pair<PrimExpr, PrimExpr> ImplicationNotContainingVars(
- const PrimExpr& cond, const std::unordered_set<const VarNode*>& vars) {
- ICHECK(cond.dtype().is_bool()) << "The type of cond must be bool";
- // TODO(sgrechanik-h): NOTs could be pushed down using De Morgan laws
- // before running this function but this case didn't seem to be important enough.
- if (const AndNode* op = cond.as<AndNode>()) {
- auto pair_a = ImplicationNotContainingVars(op->a, vars);
- auto pair_b = ImplicationNotContainingVars(op->b, vars);
- return {pair_a.first && pair_b.first, pair_a.second && pair_b.second};
- } else if (const OrNode* op = cond.as<OrNode>()) {
- auto pair_a = ImplicationNotContainingVars(op->a, vars);
- auto pair_b = ImplicationNotContainingVars(op->b, vars);
- return {pair_a.first || pair_b.first, (pair_a.first || pair_b.second) &&
- (pair_b.first || pair_a.second) &&
- (pair_a.second || pair_b.second)};
- } else if (!tir::UsesVar(cond, [&vars](const VarNode* var) { return vars.count(var); })) {
- return {cond, const_true()};
- } else {
- return {const_true(), cond};
- }
-}
-
-// Factor conditions out of a reduction by applying Fourier-Motzkin elimination and moving out
-// (in)equalities which do not depend on the reduction variables.
-std::pair<PrimExpr, PrimExpr> LiftConditionsThroughReduction(const PrimExpr& cond,
- const Array<IterVar>& red_axis,
- const Array<IterVar>& outer_axis) {
- // Factor out atomics so that we can consider this as a system of inequalities
- auto factor_atomic_res = FactorOutAtomicFormulas(cond);
- Array<PrimExpr> atomics = factor_atomic_res.atomic_formulas;
- const PrimExpr& rest = factor_atomic_res.rest;
-
- Array<Var> allvars;
- for (const IterVar& v : red_axis) {
- allvars.push_back(v->var);
- }
- for (const IterVar& v : outer_axis) {
- allvars.push_back(v->var);
- }
-
- auto vranges = Merge(IterVarsToMap(red_axis), IterVarsToMap(outer_axis));
- // start from reduction vars, so that input vars don't depend on them
- arith::IntConstraints ineq_to_solve(allvars, vranges, atomics);
- auto res_ineq = arith::SolveLinearInequalities(ineq_to_solve);
- atomics = arith::AsConditions(allvars, res_ineq.first, res_ineq.second);
-
- // Append the rest part
- PrimExpr rewritten_cond = All(atomics) && rest;
-
- std::unordered_set<const VarNode*> vset;
- for (const IterVar& v : red_axis) {
- vset.insert(v->var.get());
- }
-
- // The outer (first) condition does not contain reduction vars,
- // the inner (second) condition is everything else
- auto res = ImplicationNotContainingVars(rewritten_cond, vset);
- return res;
-}
-
-// Convert an array of itervars to an array of inequalities
-Array<PrimExpr> IterVarsToInequalities(const Array<IterVar>& itervars) {
- Array<PrimExpr> res;
- for (const IterVar& v : itervars) {
- res.push_back(GE(v->var, v->dom->min));
- res.push_back(LT(v->var, v->dom->min + v->dom->extent));
- }
- return res;
-}
-
-class RemoveRedundantInequalitiesMutator : public ExprMutator {
- public:
- explicit RemoveRedundantInequalitiesMutator(Array<PrimExpr> known) {
- for (const PrimExpr& cond : known) {
- known_.push_back(analyzer_.Simplify(cond, kSimplifyRewriteCanonicalRewrite));
- }
- }
-
- virtual PrimExpr VisitExpr_(const SelectNode* op) {
- bool has_side_effect = (SideEffect(GetRef<PrimExpr>(op)) > CallEffectKind::kReadState);
- PrimExpr new_cond =
- analyzer_.Simplify(VisitExpr(op->condition), kSimplifyRewriteCanonicalRewrite);
- if (is_one(new_cond) && !has_side_effect) {
- return VisitExpr(op->true_value);
- } else if (is_zero(new_cond) && !has_side_effect) {
- return VisitExpr(op->false_value);
- } else {
- Array<PrimExpr> new_known = known_;
- for (const PrimExpr& atomic : FactorOutAtomicFormulas(new_cond).atomic_formulas) {
- new_known.push_back(atomic);
- }
- RemoveRedundantInequalitiesMutator new_mutator(new_known);
- // Note that we mutate only the true value with the new mutator
- // TODO(sgrechanik-h): Update known conditions for the false value as well
- return Select(new_cond, new_mutator(op->true_value), VisitExpr(op->false_value));
- }
- }
-
- virtual PrimExpr VisitExpr_(const CallNode* op) {
- if (op->op.same_as(op_if_then_else_)) {
- PrimExpr new_cond =
- analyzer_.Simplify(VisitExpr(op->args[0]), kSimplifyRewriteCanonicalRewrite);
- if (is_one(new_cond)) {
- return VisitExpr(op->args[1]);
- } else if (is_zero(new_cond)) {
- return VisitExpr(op->args[2]);
- } else {
- Array<PrimExpr> new_known = known_;
- for (const PrimExpr& atomic : FactorOutAtomicFormulas(new_cond).atomic_formulas) {
- new_known.push_back(atomic);
- }
- RemoveRedundantInequalitiesMutator new_mutator(new_known);
- // Note that we mutate only the true value with the new mutator
- // TODO(sgrechanik-h): Update known conditions for the false value as well
- return if_then_else(new_cond, new_mutator(op->args[1]), VisitExpr(op->args[2]));
- }
- } else {
- return ExprMutator::VisitExpr_(op);
- }
- }
-
- virtual PrimExpr VisitExpr_(const ReduceNode* op) {
- Array<PrimExpr> known_with_axes = known_;
- ICHECK(op->init.empty()) << "Derivative of Reduction with initialization is not implemented";
- for (const PrimExpr& axis_cond : IterVarsToInequalities(op->axis)) {
- known_with_axes.push_back(axis_cond);
- }
- RemoveRedundantInequalitiesMutator mutator_with_axes(known_with_axes);
-
- PrimExpr new_cond = mutator_with_axes(op->condition);
-
- Array<PrimExpr> new_known = known_with_axes;
- for (const PrimExpr& atomic : FactorOutAtomicFormulas(new_cond).atomic_formulas) {
- new_known.push_back(atomic);
- }
- RemoveRedundantInequalitiesMutator new_mutator(new_known);
-
- Array<PrimExpr> new_source;
- for (const PrimExpr& src : op->source) {
- new_source.push_back(new_mutator(src));
- }
-
- return Reduce(op->combiner, new_source, op->axis, new_cond, op->value_index, op->init);
- }
-
- virtual PrimExpr VisitExpr_(const EQNode* op) { return MutateAtomic_(GetRef<PrimExpr>(op)); }
- virtual PrimExpr VisitExpr_(const NENode* op) { return MutateAtomic_(GetRef<PrimExpr>(op)); }
- virtual PrimExpr VisitExpr_(const LTNode* op) { return MutateAtomic_(GetRef<PrimExpr>(op)); }
- virtual PrimExpr VisitExpr_(const LENode* op) { return MutateAtomic_(GetRef<PrimExpr>(op)); }
- virtual PrimExpr VisitExpr_(const GTNode* op) { return MutateAtomic_(GetRef<PrimExpr>(op)); }
- virtual PrimExpr VisitExpr_(const GENode* op) { return MutateAtomic_(GetRef<PrimExpr>(op)); }
-
- virtual PrimExpr VisitExpr_(const AndNode* op) { return VisitExpr(op->a) && VisitExpr(op->b); }
-
- private:
- PrimExpr MutateAtomic_(const PrimExpr& e) {
- PrimExpr simplified = analyzer_.Simplify(e, kSimplifyRewriteCanonicalRewrite);
- for (const PrimExpr& other : known_) {
- if (ExprDeepEqual()(simplified, other)) {
- return const_true();
- }
- }
- return simplified;
- }
-
- Array<PrimExpr> known_;
- arith::Analyzer analyzer_;
- const Op& op_if_then_else_ = Op::Get("tir.if_then_else");
-};
-
-// Propagate information from conditions and remove redundant inequalities
-inline PrimExpr RemoveRedundantInequalities(const PrimExpr& expr, const Array<PrimExpr>& known) {
- return RemoveRedundantInequalitiesMutator(known)(expr);
-}
-
-// Extract the given expr under the given condition as a separate tensor if the volume of the
-// extracted tensor will be less than the volume of the outer_axis
-PrimExpr TrySimplifyCompute(const PrimExpr& expr, const PrimExpr& cond,
- const Array<Var>& outer_axis, const Map<Var, Range>& vranges) {
- // solve cond, e.g., (jac_i0 == i) && (jac_i1 == j)
- arith::IntConstraints domain_to_solve(outer_axis, vranges,
- FactorOutAtomicFormulas(cond).to_array());
- auto res = SimplifyDomain(domain_to_solve);
-
- arith::Analyzer analyzer;
- analyzer.Bind(res->dst->ranges);
- PrimExpr new_expr =
- analyzer.Simplify(Substitute(expr, res->src_to_dst), kSimplifyRewriteCanonicalRewrite);
- // TODO(yzhliu): This is mostly done to simplify if_then_else
- // which is not realized by the canonical simplifier
- new_expr = RemoveRedundantInequalities(new_expr, res->dst->relations);
-
- // Keep only those variables of the new vars which are used in the new_expr
- Array<Var> used_res_variables;
- for (const Var& var : res->dst->variables) {
- if (tir::UsesVar(new_expr, [&var](const VarNode* var_) { return var_ == var.get(); })) {
- ICHECK(res->dst->ranges.count(var)) << "Range of " << var << " cannot be inferred.";
- used_res_variables.push_back(var);
- }
- }
-
- // If the expression does not use vars then it is probably better to keep it inlined
- if (used_res_variables.empty()) {
- // We can return the new_expr here instead of the old expr because it doesn't use variables
- // otherwise we would need to replace the new vars or create a let-expression
- return new_expr;
- }
-
- // If it's already tensor[...] then it will probably be useless to further simplify it.
- if (new_expr.as<ProducerLoadNode>()) {
- return expr;
- }
-
- // Compute volumes before and after
- PrimExpr old_volume = make_const(DataType::Int(64), 1);
- for (const Var& var : outer_axis) {
- ICHECK(vranges.count(var)) << "Range of " << var << " was not provided.";
- old_volume = old_volume * vranges[var]->extent;
- }
-
- PrimExpr new_volume = make_const(DataType::Int(64), 1);
- for (const Var& var : used_res_variables) {
- new_volume = new_volume * res->dst->ranges[var]->extent;
- }
-
- // if we can prove that the old volume is not greater than the new volume then
- // prefer the old expression.
- arith::Analyzer ana_vranges;
- ana_vranges.Bind(vranges);
- if (ana_vranges.CanProve(old_volume <= new_volume)) {
- return expr;
- }
-
- Tensor tensor = TensorFromExpr(new_expr, IterVarsFromMap(used_res_variables, res->dst->ranges),
- "extracted_tensor");
-
- Array<PrimExpr> args;
- for (const Var& var : used_res_variables) {
- args.push_back(res->dst_to_src[var]);
- }
-
- return ProducerLoad(tensor, args);
-}
-
-class ReductionAsTensorAccessMutator : public ExprMutator {
- public:
- explicit ReductionAsTensorAccessMutator(const Array<Var>& outer_axis, Map<Var, Range> vranges,
- std::string name = "extracted_reduction")
- : outer_axis_(outer_axis), vranges_(std::move(vranges)), name_(std::move(name)) {}
-
- PrimExpr VisitExpr_(const ReduceNode* op) final {
- ReductionAsTensorAccessMutator new_mutator(Concat(IterVarsToVars(op->axis), outer_axis_),
- Merge(vranges_, IterVarsToMap(op->axis)), name_);
-
- ICHECK(op->init.empty()) << "Derivative of Reduction with initialization is not implemented";
- Array<PrimExpr> new_source;
- for (const PrimExpr& src : op->source) {
- new_source.push_back(new_mutator(src));
- }
-
- PrimExpr new_reduce =
- Reduce(op->combiner, new_source, op->axis, op->condition, op->value_index, op->init);
-
- Array<Var> undefined_vars = UndefinedVars(new_reduce);
- std::unordered_set<const VarNode*> undefined_var_set;
- for (const Var& var : undefined_vars) {
- undefined_var_set.insert(var.get());
- }
-
- // Vars of the tensor we are going to create for this reduction
- Array<Var> vars;
- for (const Var& v : outer_axis_) {
- // We take variables from the outer_axis_ which are also present in the new reduction
- if (undefined_var_set.count(v.get())) {
- vars.push_back(v);
- }
- }
-
- auto new_axis_vmap_pair = CloneIterVars(IterVarsFromMap(vars, vranges_));
- Array<IterVar> new_axis = new_axis_vmap_pair.first;
- arith::Analyzer analyzer;
- analyzer.Bind(IterVarsToMap(new_axis));
- new_reduce = analyzer.Simplify(Substitute(new_reduce, new_axis_vmap_pair.second),
- kSimplifyRewriteCanonicalRewrite);
-
- Tensor tensor = TensorFromExpr(new_reduce, new_axis, name_, tag_, attrs_);
-
- Array<PrimExpr> args;
- for (const Var& v : vars) {
- args.push_back(v);
- }
-
- return ProducerLoad(tensor, args);
- }
-
- private:
- Array<Var> outer_axis_;
- Map<Var, Range> vranges_;
- std::string name_;
- std::string tag_;
- Map<String, ObjectRef> attrs_;
-};
-
-// Extract reductions as separate tensors.
-inline PrimExpr ReductionAsTensorAccess(const PrimExpr& expr, const Array<Var>& outer_axis,
- const Map<Var, Range>& vranges) {
- return ReductionAsTensorAccessMutator(outer_axis, vranges)(expr);
-}
-
-PrimExpr LiftReductions(const PrimExpr& expr, const Array<Var>& outer_axis,
- const Map<Var, Range>& vranges) {
- if (const ReduceNode* red = expr.as<ReduceNode>()) {
- Array<Var> new_outer_axis = Concat(IterVarsToVars(red->axis), outer_axis);
- Map<Var, Range> new_vranges = Merge(vranges, IterVarsToMap(red->axis));
- Array<PrimExpr> new_source;
- for (const PrimExpr& src : red->source) {
- new_source.push_back(ReductionAsTensorAccess(src, new_outer_axis, new_vranges));
- }
- PrimExpr new_condition = ReductionAsTensorAccess(red->condition, new_outer_axis, new_vranges);
-
- return Reduce(red->combiner, new_source, red->axis, new_condition, red->value_index, red->init);
- } else {
- return ReductionAsTensorAccess(expr, outer_axis, vranges);
- }
-}
-
-PrimExpr RemoveJacobianAndLiftNonzeroCondImpl(const PrimExpr& expr_orig, const Array<IterVar>& axis,
- const Map<Var, Range>& vranges) {
- PrimExpr result;
- Map<Var, Range> combined_vranges = Merge(vranges, IterVarsToMap(axis));
- arith::Analyzer analyzer;
- analyzer.Bind(combined_vranges);
-
- // Simplify the original expression first, mostly to simplify combiners
- PrimExpr expr = analyzer.Simplify(expr_orig, kSimplifyRewriteCanonicalRewrite);
-
- if (const ReduceNode* red = expr.as<ReduceNode>()) {
- ICHECK(red->init.empty()) << "Derivative of Reduction with initialization is not implemented";
- // TODO(sgrechanik-h): There are some other operations which behave like sum
- bool is_sum = IsSumCombiner(red->combiner, vranges);
- if (is_sum || CanFactorZeroFromCombiner(red->combiner, red->value_index, vranges)) {
- PrimExpr new_red = expr;
-
- // Here we simplify the reduction
- PrimExpr cond = red->condition;
- Array<PrimExpr> source = red->source;
-
- // If it is a summation then we can lift nonzeroness conditions from the source
- // and add them to the reduction conditions
- if (is_sum) {
- auto nz = NonzeronessCondition(red->source[red->value_index]);
- cond = nz.cond && cond;
- source.Set(0, nz.value);
- }
-
- new_red = Reduce(red->combiner, source, red->axis, cond, red->value_index, red->init);
- new_red = SimplifyReductionDomain(new_red, combined_vranges);
- // Update original red pointer for later use.
- red = new_red.as<ReduceNode>();
- // If the reduction disappears completely then transform the result as a non-reduction
- if (!red) {
- return RemoveJacobianAndLiftNonzeroCondImpl(new_red, axis, vranges);
- }
-
- Array<PrimExpr> new_source = red->source;
-
- // Partially lift conditions from the reduce condition
- auto [new_outer_cond, new_reduce_cond] =
- LiftConditionsThroughReduction(red->condition, red->axis, axis);
-
- // If it's not sum then we haven't yet lifted nonzeroness cond from the source
- if (!is_sum) {
- auto nz = NonzeronessCondition(red->source[red->value_index]);
- // Append conditions from the reduction
- PrimExpr nz_source = nz.value;
- auto [outer_nz_cond, nz_cond] =
- LiftConditionsThroughReduction(new_reduce_cond && nz.cond, red->axis, axis);
- new_outer_cond = new_outer_cond && outer_nz_cond;
- new_source.Set(red->value_index, Select(nz_cond, nz_source, make_zero(nz_source.dtype())));
- }
-
- PrimExpr new_reduce = Reduce(red->combiner, new_source, red->axis, new_reduce_cond,
- red->value_index, red->init);
- new_reduce =
- TrySimplifyCompute(new_reduce, new_outer_cond, IterVarsToVars(axis), combined_vranges);
- result = Select(new_outer_cond, new_reduce, make_zero(new_reduce.dtype()));
- } else {
- return SimplifyReductionDomain(expr, combined_vranges);
- }
- } else {
- auto nz = NonzeronessCondition(expr);
- PrimExpr new_expr =
- TrySimplifyCompute(nz.value, nz.cond, IterVarsToVars(axis), combined_vranges);
- result = Select(nz.cond, new_expr, make_zero(new_expr.dtype()));
- }
-
- // Note that RemoveRedundantInequalities can sometimes propagate equalities which
- // other simplifiers cannot, like (i % 3) == 0.
- Array<PrimExpr> axis_conds = IterVarsToInequalities(axis);
- result = RemoveRedundantInequalities(result, axis_conds);
-
- // Currently in TVM reductions are only allowed at the top level of compute,
- // we need to extract intermediate inlined reduction as a separate stage (tensor).
- // Sometimes TrySimplifyCompute doesn't perform lift / extraction,
- // so there may be some non-top reductions left, take care of them.
- result = LiftReductions(result, IterVarsToVars(axis), combined_vranges);
- return analyzer.Simplify(result, kSimplifyRewriteCanonicalRewrite);
-}
-
-Tensor RemoveJacobianAndLiftNonzeroCond(const Tensor& tensor, const Map<Var, Range>& vranges) {
- auto transform_func = [&vranges](const PrimExpr& expr, const Array<IterVar>& axis) {
- return RemoveJacobianAndLiftNonzeroCondImpl(expr, axis, vranges);
- };
- return TransformTensorBody(tensor, transform_func);
-}
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/autodiff/ad_utils.cc b/src/te/autodiff/ad_utils.cc
deleted file mode 100644
index 0d1e4927cd..0000000000
--- a/src/te/autodiff/ad_utils.cc
+++ /dev/null
@@ -1,197 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file ad_utils.cc
- * \brief Utility for tensor-level auto-differentiation.
- */
-#include "ad_utils.h"
-
-#include <tvm/tir/expr.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <set>
-#include <string>
-
-#include "../schedule/operation_inline.h"
-
-namespace tvm {
-namespace te {
-
-std::pair<Array<IterVar>, Map<Var, PrimExpr>> CloneIterVars(const Array<IterVar>& vars) {
- Array<IterVar> new_vars;
- Map<Var, PrimExpr> vmap;
- for (const IterVar& iv : vars) {
- IterVar new_v = IterVar(iv->dom, iv->var.copy_with_suffix(""), iv->iter_type, iv->thread_tag);
- new_vars.push_back(new_v);
- vmap.Set(iv->var, new_v->var);
- }
- return std::make_pair(std::move(new_vars), std::move(vmap));
-}
-
-PrimExpr CloneReduction(const PrimExpr& expr) {
- if (const ReduceNode* red = expr.as<ReduceNode>()) {
- auto [new_axis, vmap] = CloneIterVars(red->axis);
-
- Array<PrimExpr> src_with_newaxis;
- for (const auto& src : red->source) {
- src_with_newaxis.push_back(tir::Substitute(src, vmap));
- }
- Array<PrimExpr> init_with_newaxis;
- for (const auto& init : red->init) {
- init_with_newaxis.push_back(tir::Substitute(init, vmap));
- }
-
- return Reduce(red->combiner, src_with_newaxis, new_axis, tir::Substitute(red->condition, vmap),
- red->value_index, init_with_newaxis);
- } else {
- return expr;
- }
-}
-
-Operation ComputeOpFromExprs(const Array<PrimExpr>& exprs, const Array<IterVar>& axis,
- const std::string& name, const std::string& tag,
- const Map<String, ObjectRef>& attrs, bool clone_axis) {
- if (clone_axis) {
- auto [new_axis, vmap] = CloneIterVars(axis);
- Array<PrimExpr> new_exprs;
- for (const PrimExpr& e : exprs) {
- new_exprs.push_back(Substitute(CloneReduction(e), vmap));
- }
- return ComputeOpFromExprs(new_exprs, new_axis, name, tag, attrs, false);
- }
-
- Array<PrimExpr> new_exprs;
-
- // If this is a reduction then we have to replicate it
- if (const ReduceNode* red = exprs[0].as<ReduceNode>()) {
- for (size_t i = 0; i < red->source.size(); ++i) {
- PrimExpr ith_red =
- Reduce(red->combiner, red->source, red->axis, red->condition, i, red->init);
- new_exprs.push_back(ith_red);
- }
- } else {
- new_exprs = exprs;
- }
-
- return ComputeOp(name, tag, attrs, axis, new_exprs);
-}
-
-Tensor TensorFromExpr(const PrimExpr& expr, const Array<IterVar>& axis, const std::string& name,
- const std::string& tag, const Map<String, ObjectRef>& attrs,
- bool clone_axis) {
- int new_value_index = 0;
- if (const ReduceNode* red = expr.as<ReduceNode>()) {
- new_value_index = red->value_index;
- }
- return ComputeOpFromExprs({expr}, axis, name, tag, attrs, clone_axis).output(new_value_index);
-}
-
-Tensor TransformTensorBody(
- const Tensor& tensor,
- const std::function<PrimExpr(const PrimExpr&, const Array<IterVar>&)>& func) {
- if (const ComputeOpNode* op = tensor->op.as<ComputeOpNode>()) {
- // Transform only one body
- PrimExpr new_body = func(op->body[tensor->value_index], op->axis);
-
- // If the body didn't change then we can return the same tensor
- if (new_body.same_as(op->body[tensor->value_index])) {
- return tensor;
- }
-
- return TensorFromExpr(new_body, op->axis, op->name, op->tag, op->attrs);
- } else {
- return tensor;
- }
-}
-
-Tensor TransformTensorBody(const Tensor& tensor,
- const std::function<PrimExpr(const PrimExpr&)>& func) {
- return TransformTensorBody(tensor,
- [func](const PrimExpr& e, const Array<IterVar>&) { return func(e); });
-}
-
-// If expr is a Tensor Access node, perform inlining, otherwise do nothing
-PrimExpr InlineImmediateTensorAccess(const PrimExpr& expr) {
- if (const ProducerLoadNode* op = expr.as<ProducerLoadNode>()) {
- auto tensor = Downcast<te::Tensor>(op->producer);
- if (const ComputeOpNode* op_comp = tensor->op.as<ComputeOpNode>()) {
- Array<Var> tensor_axes;
- for (const auto& var : op_comp->axis) {
- tensor_axes.push_back(var->var);
- }
-
- Stmt inlined =
- Inline(Evaluate(expr), tensor->op, tensor_axes, op_comp->body[tensor->value_index]);
- if (const EvaluateNode* ev = inlined.as<EvaluateNode>()) {
- // If it is a reduction, clone it
- return CloneReduction(ev->value);
- }
- }
- }
- return expr;
-}
-
-// Implements InlineTensors by trying to inline every Call of the given Expr
-class InlineTensorsMutator : public ExprMutator {
- public:
- explicit InlineTensorsMutator(const Array<Tensor>& inlineable, bool inline_reductions = false)
- : inline_reductions_(inline_reductions) {
- for (const Tensor& tensor : inlineable) {
- inlineable_.emplace(tensor->op.operator->(), tensor->value_index);
- }
- }
-
- PrimExpr VisitExpr_(const ProducerLoadNode* op) final {
- auto tensor = Downcast<te::Tensor>(op->producer);
- if (const ComputeOpNode* op_comp = tensor->op.as<ComputeOpNode>()) {
- // Inline only if the array of inlineable tensors is empty or contains this tensor
- if (inlineable_.empty() || inlineable_.count({op_comp, tensor->value_index})) {
- // Inline only compute nodes that are not reductions (unless inline reductions is allowed)
- if (inline_reductions_ || !op_comp->body[0].as<ReduceNode>()) {
- PrimExpr expr = GetRef<PrimExpr>(op);
- // Inline this tensor access and then try to perform further inlining
- return VisitExpr(InlineImmediateTensorAccess(expr));
- }
- }
- }
- // If we cannot inline this call, we should try to do inlining in its arguments
- return ExprMutator::VisitExpr_(op);
- }
-
- private:
- // Tensors which are allowed to be inlined, represented as pairs (op_node, value_index)
- std::set<std::pair<const OperationNode*, int>> inlineable_;
- bool inline_reductions_;
-};
-
-Tensor InlineTensorAccess(const Tensor& tensor, const Array<Tensor>& inlineable,
- bool inline_reductions) {
- auto transformation = [inlineable, inline_reductions](const PrimExpr& e) {
- return InlineTensorsMutator(inlineable, inline_reductions)(e);
- };
- return TransformTensorBody(tensor, transformation);
-}
-
-Tensor InlineTailTensorAccess(const Tensor& tensor) {
- return TransformTensorBody(tensor, InlineImmediateTensorAccess);
-}
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/autodiff/ad_utils.h b/src/te/autodiff/ad_utils.h
deleted file mode 100644
index 56070ef272..0000000000
--- a/src/te/autodiff/ad_utils.h
+++ /dev/null
@@ -1,135 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file ad_utils.h
- * \brief Helper utilities to implement auto-differentiation.
- */
-#ifndef TVM_TE_AUTODIFF_AD_UTILS_H_
-#define TVM_TE_AUTODIFF_AD_UTILS_H_
-
-#include <tvm/arith/int_solver.h>
-#include <tvm/te/operation.h>
-#include <tvm/tir/expr.h>
-
-#include <string>
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-namespace tvm {
-namespace te {
-
-/*!
- * \brief Clone iter vars and return both the new vars and the substitution from old to new.
- *
- * \param vars The original iter vars.
- * \return A pair containing the array of new iter vars and the map from old vars to new ones.
- */
-std::pair<Array<IterVar>, Map<Var, PrimExpr>> CloneIterVars(const Array<IterVar>& vars);
-
-/*!
- * \brief Clone reduction by cloning the axis variables.
- * \param expr A reduction expr to clone. Non-reduction expressions are left intact.
- */
-PrimExpr CloneReduction(const PrimExpr& expr);
-
-/*!
- * \brief Create a tensor from an expression. The expression may be a reduction, in which
- * case its body will be correctly duplicated if it is a multi-valued reduction.
- *
- * \param expr The expr which will be the tensor's body.
- * \param axis The input variables with ranges.
- * \param name The tensor's name.
- * \param tag The tensor's tag.
- * \param attrs The tensor's attrs.
- * \param clone_axis Whether to clone the given axis and perform substitution.
- * \return A tensor.
- */
-Tensor TensorFromExpr(const PrimExpr& expr, const Array<IterVar>& axis,
- const std::string& name = "tensor", const std::string& tag = "",
- const Map<String, ObjectRef>& attrs = {}, bool clone_axis = true);
-
-Tensor TransformTensorBody(
- const Tensor& tensor,
- const std::function<PrimExpr(const PrimExpr&, const Array<IterVar>&)>& func);
-
-Tensor TransformTensorBody(const Tensor& tensor,
- const std::function<PrimExpr(const PrimExpr&)>& func);
-
-/*!
- * \brief Inline tensors access recursively.
- *
- * This function will inline tensors recursively until it reaches a tensor which is impossible to
- * inline (a reduction if \p inline_reductions is false, a non-compute tensor, a tensor which is
- * not from \p inlineable). It won't descend into non-inlinable tensors' bodies.
- *
- * \param tensor The tensor whose body to transform.
- * \param inlineable A list of tensors which are allowed to be inlined. If empty, try
- * to inline all tensors.
- * \param inline_reductions Whether to inline reductions (this may result in top-level reduction
- * nodes).
- *
- * \return An inlined tensor
- */
-TVM_DLL Tensor InlineTensorAccess(const Tensor& tensor,
- const Array<Tensor>& inlineable = Array<Tensor>(),
- bool inline_reductions = false);
-
-/*!
- * \brief Inline tensors access at the tail.
- * \param tensor The tensor whose body to transform.
- * \return An inlined tensor
- */
-TVM_DLL Tensor InlineTailTensorAccess(const Tensor& tensor);
-
-/*!
- * \brief Simplify an iteration domain.
- *
- * An iteration domain is basically an array of variables and a condition. The function will do the
- * following:
- * - Replace div and mod operations with new variables (optional).
- * - Extract (in)equalities from the condition.
- * - Perform Fourier-Motzkin elimination.
- * - Shear the domain of iteration (e.g. if `y <= x <= y + 2` then x will be replaced with `y + d`
- * where `d` is a new variable such that `0 <= d <= 2`).
- * - Remove redundant variables.
- * - Infer new variable ranges (hopefully more precise).
- *
- * \param iter_domains The original domain.
- * \param eliminate_div_mod Whether to eliminate div and mod by introducing new variables.
- */
-TVM_DLL arith::IntConstraintsTransform SimplifyDomain(const arith::IntConstraints& iter_domains,
- bool eliminate_div_mod = true);
-
-/*!
- * \brief Perform lifting of conditions of being possible to be non-zero together with
- * applying some transformations like simplifying the reduction domain. Works only with
- * this particular tensor's body, i.e. doesn't perform inlining.
- *
- * \param tensor The original tensor;
- * \param vranges Optional map from free variables to their value ranges.
- * \return An optimized tensor.
- */
-TVM_DLL Tensor RemoveJacobianAndLiftNonzeroCond(const Tensor& tensor,
- const Map<Var, Range>& vranges = Map<Var, Range>());
-
-} // namespace te
-} // namespace tvm
-#endif // TVM_TE_AUTODIFF_AD_UTILS_H_
diff --git a/src/te/autodiff/adjoint.cc b/src/te/autodiff/adjoint.cc
deleted file mode 100644
index 34d38aa758..0000000000
--- a/src/te/autodiff/adjoint.cc
+++ /dev/null
@@ -1,157 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file adjoint.cc
- * \brief Perform reverse-mode autodiff.
- * Suppose we have f(x) = g(h1(x), h2(x), ..., hn(x)),
- * df/dx = \sum_i df/dhi * dhi/dx
- * We call df/dx as adjoint(x), df/dhi as adjoint(hi), dhi/dx is the Jacobian
- * The idea is to first construct the reverse-dependency {input->outputs} between tensors,
- * start from one input,
- * (1) collect adjoints from all its dependencies (outputs),
- * (2) multiply the Jacobian (PartialAdjoint),
- * (3) and sum them together to get the adjoint of the input itself.
- * The three steps are computed recursively.
- */
-#include <tvm/runtime/registry.h>
-#include <tvm/te/autodiff.h>
-#include <tvm/tir/stmt_functor.h>
-#include <tvm/topi/elemwise.h>
-#include <tvm/topi/transform.h>
-
-#include <memory>
-#include <vector>
-
-#include "ad_utils.h"
-
-namespace tvm {
-namespace te {
-
-Tensor Identity(const Tensor& output) {
- Array<PrimExpr> shape = output->shape;
- for (auto e : output->shape) {
- // add extra dimension for Jacobian
- shape.push_back(e);
- }
- auto func = [&output](const Array<Var>& input_indices) {
- PrimExpr res = const_true();
- for (size_t i = 0; i < output->shape.size(); ++i) {
- res =
- res && (PrimExpr(input_indices[i]) == PrimExpr(input_indices[output->shape.size() + i]));
- }
- return Cast(output->dtype, res);
- };
- return te::compute(shape, func, "identity");
-}
-
-Tensor VectorJacobianProduct(const Tensor& output, const Tensor& input, const Tensor& head) {
- Tensor jac = Jacobian(output, input);
- Tensor result = topi::tensordot(head, jac, /*axes=*/output->shape.size(),
- output->op->name + "." + input->op->name + ".grad");
- result = InlineTensorAccess(result, {jac}, false);
- result = RemoveJacobianAndLiftNonzeroCond(result);
- // inline tail call
- result = InlineTailTensorAccess(result);
- return result;
-}
-
-Array<Tensor> Gradient(const Tensor& output, const Array<Tensor>& inputs,
- const Tensor& head_or_null) {
- // Diagonal identity tensor
- Tensor head = head_or_null.get() ? head_or_null : Identity(output);
-
- // This Map{input -> outputs} maps a tensor to the list of tensors
- // immediately depending on it (using it in their bodies)
- std::unordered_map<Tensor, std::vector<Tensor>> reverse_dependencies;
- std::vector<Tensor> stack({output});
- while (!stack.empty()) {
- Tensor tensor = stack.back();
- stack.pop_back();
- for (const Tensor& input : tensor->op->InputTensors()) {
- if (!reverse_dependencies.count(input)) {
- stack.push_back(input);
- }
- reverse_dependencies[input].push_back(tensor);
- }
- }
-
- // This map maps tensors to the corresponding adjoints (dLoss/dTensor)
- std::unordered_map<Tensor, Tensor> adjoints;
- // head is the adjoint of output by definition
- adjoints[output] = head;
-
- // This is a recursive function that does all the work. It computes the adjoint for a given
- // tensor, adds it to the map, and returns it
- std::function<Tensor(const Tensor&)> compute_adjoint;
- compute_adjoint = [&compute_adjoint, &adjoints, &reverse_dependencies, &head,
- &output](const Tensor& tensor) {
- if (!adjoints.count(tensor)) {
- // Here the adjoint hasn't been computed yet
- Tensor res_adjoint;
- std::vector<Tensor> direct_consumers = reverse_dependencies[tensor];
- if (direct_consumers.empty()) {
- // No reverse dependencies means that the output does not depend on this tensor,
- // return a zero tensor of the appropriate shape
- // (i.e., output shape + tensor shape, aka shape of Jacobian)
- Array<PrimExpr> result_shape(head->shape.begin(), head->shape.end() - output->shape.size());
- for (auto e : tensor->shape) {
- result_shape.push_back(e);
- }
- res_adjoint = topi::full(result_shape, output->dtype, make_zero(output->dtype));
- } else {
- // The new adjoint is computed as a sum of the reverse dependencies' adjoints multiplied
- // by the corresponding "local" jacobians (dDep/dTensor). The computation of the jacobian
- // and the multiplication is done in the function VectorJacobianProduct
- for (const Tensor& direct_consumer : direct_consumers) {
- // part = (adjoint of direct_consumer) * Jacobian(direct_consumer, tensor)
- Tensor part =
- VectorJacobianProduct(direct_consumer, tensor, compute_adjoint(direct_consumer));
- res_adjoint = res_adjoint.get() ? topi::add(res_adjoint, part) : part;
- }
- }
-
- adjoints[tensor] = res_adjoint;
- return res_adjoint;
- } else {
- return adjoints[tensor];
- }
- };
-
- // Adjoints corresponding to inputs
- Array<Tensor> result;
- // Compute an adjoint for each input
- for (const Tensor& input : inputs) {
- result.push_back(compute_adjoint(input));
- }
-
- return result;
-}
-
-TVM_REGISTER_GLOBAL("te.Gradient").set_body([](TVMArgs args, TVMRetValue* ret) {
- LOG(WARNING) << "te.Gradient is an experimental feature.";
- if (args.size() == 2) {
- *ret = Gradient(args[0], args[1]);
- } else if (args.size() == 3) {
- *ret = Gradient(args[0], args[1], args[2]);
- }
-});
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/autodiff/jacobian.cc b/src/te/autodiff/jacobian.cc
deleted file mode 100644
index 78788dbe1a..0000000000
--- a/src/te/autodiff/jacobian.cc
+++ /dev/null
@@ -1,365 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file jacobian.cc
- * \brief Calculate Jacobian of two tensors dY/dX.
- * X must be direct input tensor of Y.
- * The result Jacobian shape will be (Y.shape, X.shape)
- */
-#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
-#include <tvm/te/autodiff.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <memory>
-
-#include "ad_utils.h"
-
-namespace tvm {
-namespace te {
-
-#define NOT_IMPLEMENTED \
- { \
- LOG(FATAL) << "Derivative of this expr is not implemented: " << GetRef<PrimExpr>(op); \
- throw; \
- }
-
-/*! \brief Differentiate an expression wrt a variable or a tensor element */
-class JacobianMutator : public ExprMutator {
- public:
- /*!
- * \brief Differentiate wrt `input(indices)`.
- * \param input The input tensor.
- * \param indices The indices of the element with respect to which to differentiate.
- */
- explicit JacobianMutator(Tensor input, Array<PrimExpr> indices)
- : input_(input), indices_(indices) {}
- /*!
- * \brief Differentiate wrt the input variable.
- * \param input The input variable.
- */
- explicit JacobianMutator(Var input) : input_var_(input) {}
-
- PrimExpr Mutate(PrimExpr e) {
- if (e.dtype().is_int() || e.dtype().is_uint()) {
- LOG(WARNING) << "For now we assume that the derivative of any integer expression is always 0."
- << " e = " << e;
- return make_zero(e.dtype());
- } else {
- return ExprMutator::VisitExpr(e);
- }
- }
-
- PrimExpr VisitExpr_(const VarNode* op) {
- if (input_var_.get() && input_var_.get() == op && op->dtype.is_float()) {
- return FloatImm(op->dtype, 1.0);
- } else {
- return make_zero(op->dtype);
- }
- }
-
- PrimExpr VisitExpr_(const LetNode* op) NOT_IMPLEMENTED;
-
- PrimExpr VisitExpr_(const ProducerLoadNode* op) final {
- auto tensor = Downcast<te::Tensor>(op->producer);
- if (input_.get() && tensor == input_) {
- // Tensor(indices)
- ICHECK_EQ(indices_.size(), op->indices.size());
- PrimExpr condition = const_true();
- for (size_t i = 0; i < input_.ndim(); ++i) {
- condition = And(condition, EQ(indices_[i], op->indices[i]));
- }
- return Cast(op->dtype, condition);
- } else {
- return make_zero(op->dtype);
- }
- }
-
- PrimExpr VisitExpr_(const CallNode* op) {
- PrimExpr expr = GetRef<PrimExpr>(op);
- if (op->op.same_as(op_exp_)) {
- return Mul(Mutate(op->args[0]), expr);
- } else if (op->op.same_as(op_log_)) {
- return Div(Mutate(op->args[0]), op->args[0]);
- } else if (op->op.same_as(op_sigmoid_)) {
- return Mul(Mutate(op->args[0]), Mul(expr, Sub(FloatImm(expr.dtype(), 1.0), expr)));
- } else if (op->op.same_as(op_sqrt_)) {
- return Div(Mutate(op->args[0]), Mul(expr, FloatImm(expr.dtype(), 2.0)));
- } else if (op->op.same_as(op_tanh_)) {
- return Mul(Mutate(op->args[0]), Sub(FloatImm(expr.dtype(), 1.0), Mul(expr, expr)));
- } else if (op->op.same_as(op_pow_)) {
- auto x = op->args[0], y = op->args[1];
- return expr * (Mutate(y) * log(x) + Mutate(x) * y / x);
- } else if (op->op.same_as(op_fabs_)) {
- auto type = op->args[0].dtype();
- return Mul(Mutate(op->args[0]), Select(GE(op->args[0], make_zero(type)), FloatImm(type, 1.0),
- FloatImm(type, -1.0)));
- } else if (op->op.same_as(op_if_then_else_)) {
- Array<PrimExpr> new_args = {op->args[0], Mutate(op->args[1]), Mutate(op->args[2])};
- return Call(op->dtype, op->op, new_args);
- } else if (piecewise_const.count(op->op)) {
- return FloatImm(expr.dtype(), 0.0);
- } else {
- LOG(FATAL) << "Derivative of this intrinsic is not implemented: " << op->op;
- }
- }
-
- PrimExpr VisitExpr_(const AddNode* op) { return Add(Mutate(op->a), Mutate(op->b)); }
-
- PrimExpr VisitExpr_(const SubNode* op) { return Sub(Mutate(op->a), Mutate(op->b)); }
-
- PrimExpr VisitExpr_(const MulNode* op) {
- return Add(Mul(Mutate(op->a), op->b), Mul(op->a, Mutate(op->b)));
- }
-
- PrimExpr VisitExpr_(const DivNode* op) {
- return Div(Sub(Mul(Mutate(op->a), op->b), Mul(op->a, Mutate(op->b))), Mul(op->b, op->b));
- }
-
- PrimExpr VisitExpr_(const ModNode* op) NOT_IMPLEMENTED;
-
- PrimExpr VisitExpr_(const FloorDivNode* op) {
- return FloorDiv(Sub(Mul(Mutate(op->a), op->b), Mul(op->a, Mutate(op->b))), Mul(op->b, op->b));
- }
-
- PrimExpr VisitExpr_(const FloorModNode* op) NOT_IMPLEMENTED;
-
- PrimExpr VisitExpr_(const MinNode* op) {
- return Select(LE(op->a, op->b), Mutate(op->a), Mutate(op->b));
- }
-
- PrimExpr VisitExpr_(const MaxNode* op) {
- return Select(GE(op->a, op->b), Mutate(op->a), Mutate(op->b));
- }
-
- PrimExpr VisitExpr_(const EQNode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const NENode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const LTNode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const LENode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const GTNode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const GENode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const AndNode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const OrNode* op) NOT_IMPLEMENTED;
-
- PrimExpr VisitExpr_(const ReduceNode* op) {
- // This case is relatively difficult because a reduction expression
- // may use an arbitrary combiner.
- // The resulting reduction expression will return a tuple containing
- // both derivatives and the original results (in exactly this order).
- // The order matters when original init value is different from its derivative init value,
- // and they depend on each other during gradient calculation,
- // we must calculate derivatives first (using origin's init value),
- // switching the order (original results first, then derivatives)
- // makes the origin value be replaced before using,
- // produces incorrect results.
-
- // Example of a ReduceNode,
- // reduce(combiner=comm_reducer(result=[(x + y)], lhs=[x], rhs=[y], identity_element=[0f]),
- // source=[A(k)], axis=[iter_var(k, range(min=0, ext=5))], where=(bool)1, value_index=0)
-
- // We have to clone the reduction axes because otherwise the original expression
- // cannot be used together with the derivative (it will lead to errors during lowering)
- PrimExpr expr_with_new_axes = te::CloneReduction(GetRef<PrimExpr>(op));
- const ReduceNode* new_op = expr_with_new_axes.as<ReduceNode>();
-
- ICHECK(new_op->init.empty())
- << "Derivative of Reduction with initialization is not implemented";
-
- // New lhs and rhs variables of the new combiner consist of
- // variables representing derivatives (which are later derived from new_op->source)
- // followed by the original variables.
- Array<Var> new_lhs;
- for (const auto& var : new_op->combiner->lhs) {
- new_lhs.push_back(var.copy_with_suffix(".jac"));
- }
- for (const auto& var : new_op->combiner->lhs) {
- new_lhs.push_back(var);
- }
-
- Array<Var> new_rhs;
- for (const auto& var : new_op->combiner->rhs) {
- new_rhs.push_back(var.copy_with_suffix(".jac"));
- }
- for (const auto& var : new_op->combiner->rhs) {
- new_rhs.push_back(var);
- }
-
- // The new combiner result also consists of the resulting derivatives
- // followed by the original results.
- Array<PrimExpr> new_result;
- for (const auto& res : new_op->combiner->result) {
- // Each resulting derivative is computed as a sum of derivatives
- // wrt lhs and rhs multiplied by the derivatives of lhs and rhs
- PrimExpr new_res = make_zero(res.dtype());
- for (size_t i = 0; i < new_op->combiner->lhs.size(); ++i) {
- PrimExpr res_di = Derivative(res, new_op->combiner->lhs[i]);
- // new_lhs[i] is the derivative of lhs[i] (wrt our input tensor)
- new_res = Add(new_res, Mul(new_lhs[i], res_di));
- }
- for (size_t i = 0; i < new_op->combiner->rhs.size(); ++i) {
- PrimExpr res_di = Derivative(res, new_op->combiner->rhs[i]);
- // new_rhs[i] is the derivative of rhs[i] (wrt our input tensor)
- new_res = Add(new_res, Mul(new_rhs[i], res_di));
- }
- new_result.push_back(new_res);
- }
- // add original results
- for (const auto& res : new_op->combiner->result) {
- new_result.push_back(res);
- }
-
- // The identity is transformed in a similar way
- Array<PrimExpr> new_identity;
- for (const auto& id : new_op->combiner->identity_element) {
- new_identity.push_back(Mutate(id));
- }
- for (const auto& id : new_op->combiner->identity_element) {
- new_identity.push_back(id);
- }
-
- // Same as source
- Array<PrimExpr> new_source;
- for (const auto& src : new_op->source) {
- new_source.push_back(Mutate(src));
- }
- for (const auto& src : new_op->source) {
- new_source.push_back(src);
- }
-
- CommReducer new_combiner = CommReducer(new_lhs, new_rhs, new_result, new_identity);
- // Also simplify the resulting combiner
- // (mostly to get rid of unused components, e.g., the original expressions)
- return analyzer_.Simplify(Reduce(new_combiner, new_source, new_op->axis, new_op->condition,
- new_op->value_index, new_op->init));
- }
-
- PrimExpr VisitExpr_(const CastNode* op) {
- if (op->dtype.is_float()) {
- return Cast(op->dtype, Mutate(op->value));
- } else {
- return make_zero(op->dtype);
- }
- }
-
- PrimExpr VisitExpr_(const NotNode* op) NOT_IMPLEMENTED;
-
- PrimExpr VisitExpr_(const SelectNode* op) {
- return Select(op->condition, Mutate(op->true_value), Mutate(op->false_value));
- }
-
- PrimExpr VisitExpr_(const RampNode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const BroadcastNode* op) NOT_IMPLEMENTED;
- PrimExpr VisitExpr_(const ShuffleNode* op) NOT_IMPLEMENTED;
-
- PrimExpr VisitExpr_(const IntImmNode* op) { return IntImm(op->dtype, 0); }
-
- PrimExpr VisitExpr_(const FloatImmNode* op) { return FloatImm(op->dtype, 0); }
-
- PrimExpr VisitExpr_(const StringImmNode* op) NOT_IMPLEMENTED;
-
- private:
- Tensor input_;
- Array<PrimExpr> indices_;
- Var input_var_;
- arith::Analyzer analyzer_;
-
- const Op& op_exp_ = Op::Get("tir.exp");
- const Op& op_log_ = Op::Get("tir.log");
- const Op& op_sigmoid_ = Op::Get("tir.sigmoid");
- const Op& op_sqrt_ = Op::Get("tir.sqrt");
- const Op& op_tanh_ = Op::Get("tir.tanh");
- const Op& op_pow_ = Op::Get("tir.pow");
- const Op& op_fabs_ = Op::Get("tir.fabs");
- const Op& op_if_then_else_ = Op::Get("tir.if_then_else");
- std::unordered_set<RelayExpr, ObjectPtrHash, ObjectPtrEqual> piecewise_const = {
- Op::Get("tir.floor"), Op::Get("tir.ceil"), Op::Get("tir.trunc"), Op::Get("tir.round")};
-};
-
-PrimExpr Derivative(const PrimExpr& expr, const Var& var) {
- return JacobianMutator(var).Mutate(expr);
-}
-
-PrimExpr Jacobian(const PrimExpr& expr, const Tensor& input, const Array<PrimExpr>& indices) {
- return JacobianMutator(input, indices).Mutate(expr);
-}
-
-Tensor Jacobian(const Tensor& output, const Tensor& input) {
- const ComputeOpNode* op = output->op.as<ComputeOpNode>();
- ICHECK(op) << "Derivative of this operation is not implemented: " << output->op;
- bool is_input_tensor = false;
- for (const Tensor& child : op->InputTensors()) {
- if (input == child) {
- is_input_tensor = true;
- break;
- }
- }
- ICHECK(is_input_tensor) << "Jacobian is called on a pair of tensors such that the output "
- << "does not directly depend on the input.";
-
- // We have to clone the iteration axes because otherwise the original expression
- // cannot be used together with the derivative (it will lead to errors during lowering)
- auto [new_axis, vmap] = te::CloneIterVars(op->axis);
-
- Array<PrimExpr> input_indices;
- size_t i = 0;
- for (PrimExpr ext : input->shape) {
- IterVar new_v =
- IterVar(Range(0, ext), Var("jac_i" + std::to_string(i++)), IterVarType::kDataPar);
- // Append jacobian iter to new_axis
- new_axis.push_back(new_v);
- // Differentiate wrt input[input_indices]
- input_indices.push_back(new_v);
- }
- arith::Analyzer analzyer;
- // Compute Jacobian
- PrimExpr new_body =
- Jacobian(Substitute(op->body[output->value_index], vmap), input, input_indices);
- new_body = analzyer.Simplify(new_body);
-
- int value_index = 0;
- Array<PrimExpr> new_bodies;
-
- // If this is a reduction then it may return a tuple and we have
- // to repeat the body several times
- if (const ReduceNode* red = new_body.as<ReduceNode>()) {
- value_index = red->value_index;
- for (size_t idx = 0; idx < red->source.size(); ++idx) {
- new_bodies.push_back(
- Reduce(red->combiner, red->source, red->axis, red->condition, idx, red->init));
- }
- } else {
- new_bodies.push_back(new_body);
- }
-
- auto new_op = ComputeOp(op->name + ".jacobian", op->tag, op->attrs, new_axis, new_bodies);
-
- // Jacobian shape = output.shape + input.shape
- Array<PrimExpr> new_shape = output->shape;
- for (const auto& e : input->shape) {
- new_shape.push_back(e);
- }
-
- Tensor ret = Tensor(new_shape, output->dtype, new_op, value_index);
- ret = RemoveJacobianAndLiftNonzeroCond(ret);
- return ret;
-}
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc
index fb839c28da..48d10e97a0 100644
--- a/src/te/operation/compute_op.cc
+++ b/src/te/operation/compute_op.cc
@@ -21,7 +21,6 @@
* \brief Compute Op.
* \file compute_op.cc
*/
-#include "compute_op.h"
#include <tvm/arith/analyzer.h>
#include <tvm/runtime/registry.h>
@@ -35,10 +34,6 @@
#include <unordered_set>
#include <utility>
-#include "../../arith/interval_set.h"
-#include "../schedule/message_passing.h"
-#include "op_utils.h"
-
namespace tvm {
namespace te {
using namespace tir;
@@ -79,15 +74,6 @@ static inline void AssertReduceEqual(const tir::ReduceNode* a, const tir::Reduce
int ComputeOpNode::num_outputs() const { return body.size(); }
-Array<IterVar> BaseComputeOpNode::root_iter_vars() const {
- if (reduce_axis.size() == 0) return axis;
- Array<IterVar> ret = axis;
- for (IterVar iv : reduce_axis) {
- ret.push_back(iv);
- }
- return ret;
-}
-
DataType ComputeOpNode::output_dtype(size_t idx) const {
ICHECK_LT(idx, num_outputs());
return body[idx].dtype();
@@ -185,333 +171,8 @@ Array<Tensor> ComputeOpNode::InputTensors() const {
return ret;
}
-Operation ComputeOpNode::ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const {
- ICHECK_EQ(self.operator->(), this);
- VerifyComputeOp(this);
- Array<PrimExpr> arr;
- if (this->body[0]->IsInstance<tir::ReduceNode>()) {
- // Specially handle reduce so the replaced op
- // still share all the components
- PrimExpr new_reduce = te::ReplaceTensor(this->body[0], rmap);
- if (!new_reduce.same_as(this->body[0])) {
- const tir::ReduceNode* r = new_reduce.as<tir::ReduceNode>();
- for (size_t k = 0; k < this->body.size(); ++k) {
- auto n = make_object<tir::ReduceNode>(*r);
- n->value_index = static_cast<int>(k);
- n->dtype = r->source[k].dtype();
- arr.push_back(PrimExpr(n));
- }
- } else {
- arr = this->body;
- }
- } else {
- arr =
- UpdateArray(this->body, [&rmap](const PrimExpr& e) { return te::ReplaceTensor(e, rmap); });
- }
- if (!arr.same_as(this->body)) {
- return ComputeOp(this->name, this->tag, this->attrs, this->axis, arr);
- } else {
- return self;
- }
-}
-
-void ComputeOpNode::PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const {
- ICHECK_EQ(self.operator->(), this);
- auto fvisit = [&dom_map, out_dom_map, analyzer](const ObjectRef& n) {
- if (auto* pload = n.as<tir::ProducerLoadNode>()) {
- Tensor t = Downcast<Tensor>(pload->producer);
- if (t->op.defined() && out_dom_map->count(t)) {
- TensorDom& dom = out_dom_map->at(t);
- for (size_t i = 0; i < t.ndim(); ++i) {
- // We assume that the value of the argument cannot be out of bounds (otherwise it is
- // undefined behaviour), so we can intersect the estimated set of the argument with the
- // range expected by the tensor. However, intersection may result in overly complex
- // expressions, so we perform a more relaxed form of intersection.
- IntSet arg_intset = analyzer->int_set(pload->indices[i], ConvertDomMap(dom_map));
- const arith::IntervalSetNode* arg_interval = arg_intset.as<arith::IntervalSetNode>();
- if (arg_interval) {
- PrimExpr shape_i_min_value = make_zero(t->shape[i].dtype());
- PrimExpr shape_i_max_value = t->shape[i] - 1;
- PrimExpr min_value = arg_interval->min_value;
- PrimExpr max_value = arg_interval->max_value;
- // Prefer the shape bounds only when we can prove they are tighter.
- // We must update bound's ends in pairs. Here is an counter example: shape_i is
- // [0, 0] and arg_interval is [threadIdx.y, threadIdx.y], where threadIdx.y's range is
- // [0, 7]. If we allowed updating one end, the bound would become [threadIdx.y, 0],
- // awkward for further analysis.
- if ((arith::is_pos_inf(max_value) && arith::is_neg_inf(min_value)) ||
- (analyzer->CanProve(shape_i_min_value >= min_value) &&
- analyzer->CanProve(shape_i_max_value <= max_value))) {
- min_value = shape_i_min_value;
- max_value = shape_i_max_value;
- }
- dom.data[i].push_back(IntSet::Interval(min_value, max_value));
- } else {
- dom.data[i].push_back(arg_intset);
- }
- }
- }
- }
- };
- for (auto& e : body) tir::PostOrderVisit(e, fvisit);
-}
-
-void BaseComputeOpNode::GatherBound(const Operation& self,
- const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const {
- ICHECK_EQ(self.operator->(), this);
- const TensorDom& tdom = tensor_dom.at(self.output(0));
- for (size_t i = 0; i < this->axis.size(); ++i) {
- Range r = arith::Union(tdom.data.at(i)).CoverRange(this->axis[i]->dom);
- ICHECK(!out_dom_map->count(this->axis[i]));
- (*out_dom_map)[this->axis[i]] = r;
- }
- for (size_t i = 0; i < this->reduce_axis.size(); ++i) {
- ICHECK(!out_dom_map->count(this->reduce_axis[i]));
- (*out_dom_map)[this->reduce_axis[i]] = this->reduce_axis[i]->dom;
- }
-}
-
-Stmt BaseComputeOpNode::BuildRealize(const Stage& stage,
- const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope) const {
- ICHECK_EQ(stage->op.get(), this);
- Region bounds;
- for (IterVar iv : this->axis) {
- bounds.push_back(realize_map.at(iv));
- }
- Stmt realize = body;
- for (int i = this->num_outputs(); i > 0; --i) {
- Tensor t = stage->op.output(i - 1);
- realize = tir::ProducerRealize(t, bounds, const_true(), realize, storage_scope);
- // alignment requirement, only useful for compute
- for (size_t i = 0; i < num_schedulable_dims(); ++i) {
- auto it = stage->iter_var_attrs.find(this->axis[i]);
- if (it != stage->iter_var_attrs.end()) {
- IterVarAttr attr = (*it).second;
- if (attr->dim_align_factor != 0) {
- Array<PrimExpr> tuple = {static_cast<int>(i), attr->dim_align_factor,
- attr->dim_align_offset};
- realize =
- tir::AttrStmt(t, tir::attr::buffer_dim_align,
- Call(DataType::Handle(), tir::builtin::tvm_tuple(), tuple), realize);
- }
- }
- }
- }
- return realize;
-}
-
-size_t ComputeOpNode::num_schedulable_dims() const { return axis.size(); }
-
-// Build a reduction body.
-void MakeReduction(const ComputeOpNode* op, const Array<Tensor>& tensors, Stmt* init,
- Stmt* provide) {
- Array<PrimExpr> args;
- for (IterVar iv : op->axis) {
- args.push_back(iv->var);
- }
- std::vector<Stmt> inits, provides;
-
- size_t size = op->body.size();
- const ReduceNode* reduce = op->body[0].as<ReduceNode>();
- ICHECK(reduce);
- const CommReducerNode* combiner = reduce->combiner.as<CommReducerNode>();
- ICHECK(combiner);
- Array<PrimExpr> lhs;
- for (size_t i = 0; i < size; ++i) {
- lhs.push_back(tensors[i](args));
- }
- Array<PrimExpr> init_value = combiner->identity_element;
- Array<PrimExpr> update_value = (*combiner)(lhs, reduce->source);
-
- // If an init was passed to ReduceNode, use that for initialization
- // instead of combiner->identity_element
- Array<PrimExpr> reduce_init = reduce->init;
- if (!reduce_init.empty()) {
- init_value = reduce_init;
- }
- for (size_t i = 0; i < size; ++i) {
- Tensor t = tensors[i];
- inits.emplace_back(ProducerStore(t, init_value[i], args));
- provides.emplace_back(ProducerStore(t, update_value[i], args));
- }
- *init = SeqStmt::Flatten(inits);
- *provide = SeqStmt::Flatten(provides);
- if (!is_one(reduce->condition)) {
- *provide = IfThenElse(reduce->condition, *provide);
- }
-}
-
-// Normal computation.
-Stmt MakeProvide(const ComputeOpNode* op, const Tensor& t) {
- Array<PrimExpr> args;
- for (IterVar iv : op->axis) {
- args.push_back(iv->var);
- }
- return ProducerStore(t, op->body[t->value_index], args);
-}
-
-Stmt MakeComputeStmt(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) {
- // grab the nest structure
- ComputeLoopNest n = ComputeLoopNest::Create(self, stage, dom_map, debug_keep_trivial_loop);
- // Normal loop structure
- n.init_nest.emplace_back(MakeIfNest(n.init_predicates));
- n.main_nest.emplace_back(MakeIfNest(n.main_predicates));
- if (self->reduce_axis.size() != 0) {
- // make reduction.
- Stmt init, provide;
- Array<Tensor> source;
- for (size_t i = 0; i < self->body.size(); ++i) {
- source.push_back(stage->op.output(i));
- }
- MakeReduction(self, source, &init, &provide);
- init = MergeNest(n.init_nest, init);
- init = Substitute(init, n.init_vmap);
- // common nest
- std::vector<std::vector<Stmt>> common(n.main_nest.begin(),
- n.main_nest.begin() + n.num_common_loop + 1);
- std::vector<std::vector<Stmt>> reduce(n.main_nest.begin() + n.num_common_loop + 1,
- n.main_nest.end());
- provide = MergeNest(reduce, provide);
- if (debug_keep_trivial_loop) {
- provide = MergeNest(common, provide);
- } else {
- provide = MergeNest(common, SeqStmt::Flatten(init, provide));
- }
- // run substitution in the on the full nest, because loop condition
- // could depend on outer loops.
- return Substitute(provide, n.main_vmap);
- } else {
- std::vector<Stmt> provides;
- for (size_t i = 0; i < self->body.size(); ++i) {
- provides.emplace_back(MakeProvide(self, stage->op.output(i)));
- }
- Stmt provide = SeqStmt::Flatten(provides);
- provide = MergeNest(n.main_nest, provide);
- // run substitution in the on the full nest, because loop condition
- // could depend on outer loops.
- return Substitute(provide, n.main_vmap);
- }
-}
-
enum class ComputeType { kNormal, kCrossThreadReduction, kTensorize };
-ComputeType DetectComputeType(const ComputeOpNode* self, const Stage& stage) {
- // Verify correctness of leaf nest.
- int thread_red = 0, tensorize = 0;
-
- for (IterVar iv : stage->leaf_iter_vars) {
- IterVarAttr attr;
- auto it = stage->iter_var_attrs.find(iv);
- if (it != stage->iter_var_attrs.end()) {
- attr = (*it).second;
- }
- if (attr.defined() && attr->iter_type == kTensorized) {
- ++tensorize;
- }
- if (iv->iter_type == kCommReduce) {
- if (attr.defined() && attr->bind_thread.defined()) {
- ++thread_red;
- }
- } else {
- ICHECK_EQ(thread_red, 0) << "Cross thread reduce cannot swap with normal data axis";
- }
- }
- if (tensorize != 0) {
- ICHECK(thread_red == 0) << "Cannot mix cross thread reduction with Tensorize";
- return ComputeType::kTensorize;
- }
- if (thread_red != 0) {
- return ComputeType::kCrossThreadReduction;
- } else {
- return ComputeType::kNormal;
- }
-}
-
-// implement the provide utility.
-Stmt ComputeOpNode::BuildProvide(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const {
- ICHECK_EQ(stage->op.operator->(), this);
- ComputeType ctype = DetectComputeType(this, stage);
- if (ctype == ComputeType::kCrossThreadReduction) {
- // specially handle cross thread reduction.
- return MakeCrossThreadReduction(this, stage, dom_map, debug_keep_trivial_loop);
- } else if (ctype == ComputeType::kTensorize) {
- return MakeTensorize(this, stage, dom_map, debug_keep_trivial_loop);
- } else {
- return MakeComputeStmt(this, stage, dom_map, debug_keep_trivial_loop);
- }
-}
-
-ComputeLoopNest ComputeLoopNest::Create(const BaseComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) {
- ICHECK_EQ(stage->op.operator->(), self);
- ComputeLoopNest ret;
- // make main loop nest
- ret.main_nest = MakeLoopNest(stage, dom_map, 0, false, std::unordered_set<IterVar>(),
- &ret.main_vmap, debug_keep_trivial_loop);
- ret.main_predicates =
- MakeBoundCheck(stage, dom_map, ret.main_vmap, false, std::unordered_set<IterVar>());
- for (auto& e : ret.main_predicates) {
- e = likely(e);
- }
- if (stage->store_predicate.defined()) {
- ret.main_predicates.push_back(stage->store_predicate);
- }
- if (self->reduce_axis.size() != 0) {
- // try to find the location to insert the initialization.
- // Fuse the initialization and provide loop when possible.
- std::unordered_map<IterVar, int> update_state;
- for (IterVar iv : self->reduce_axis) {
- update_state[iv] = 2;
- }
- for (size_t i = 0; i < self->num_schedulable_dims(); ++i) {
- update_state[self->axis[i]] = 1;
- }
- // find which iter var is related to reduction and which is related to axis.
- te::PassDownBitMaskOr(stage, &update_state);
- auto leaf_iter_vars = stage->leaf_iter_vars;
- // first first loop that is related to reduction.
- size_t begin_loop = leaf_iter_vars.size();
- for (size_t i = 0; i < leaf_iter_vars.size(); ++i) {
- auto iv = leaf_iter_vars[i];
- int flag = update_state.at(iv);
- if ((flag & 2) != 0) {
- begin_loop = i;
- break;
- }
- ret.init_vmap[iv] = ret.main_vmap.at(iv);
- }
- ret.num_common_loop = begin_loop;
- // skip loops that are related to reduction and are unrelated to axis.
- std::unordered_set<IterVar> skip_iter;
- for (auto kv : update_state) {
- int flag = kv.second;
- if (flag == 2) skip_iter.insert(kv.first);
- }
- ret.init_nest = MakeLoopNest(stage, dom_map, begin_loop, true, skip_iter, &(ret.init_vmap),
- debug_keep_trivial_loop);
- ret.init_predicates =
- MakeBoundCheck(stage, dom_map, ret.init_vmap, !stage->rolling_buffer, skip_iter);
- for (auto& e : ret.init_predicates) {
- e = likely(e);
- }
- } else {
- ICHECK_EQ(ret.main_nest.size(), stage->leaf_iter_vars.size() + 1);
- ret.num_common_loop = stage->leaf_iter_vars.size();
- }
- // copy elison here.
- return ret;
-}
-
namespace {
/*!
* \brief Verify if ComputeOp is valid with respect to Reduce operations.
@@ -581,41 +242,5 @@ static void VerifyComputeOp(const ComputeOpNode* op) {
v.Run();
}
-Stmt TransformUpdate(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- const ComputeLoopNest& n, Stmt body, Stmt update) {
- Array<PrimExpr> conds;
- std::unordered_set<const VarNode*> banned;
- for (size_t i = 0; i < stage->leaf_iter_vars.size(); ++i) {
- IterVar iv = stage->leaf_iter_vars[i];
- auto iit = stage->iter_var_attrs.find(iv);
- if (iit != stage->iter_var_attrs.end()) {
- const IterVarAttr& attr = (*iit).second;
- if (attr->iter_type == kTensorized) {
- break;
- }
- }
- if (iv->iter_type == kCommReduce) {
- auto vit = dom_map.find(iv);
- ICHECK(vit != dom_map.end());
- const Range& vrange = vit->second;
- conds.push_back(likely(iv->var > vrange->min));
- banned.insert(iv->var.get());
- }
- }
-
- auto fbanned = [&](const VarNode* node) { return banned.count(node); };
-
- for (const PrimExpr& pred : n.main_predicates) {
- if (tir::UsesVar(pred, fbanned)) {
- LOG(FATAL) << "Tensorize update transform failed, the condition " << pred
- << " has a conflict with the reset condition";
- }
- }
-
- auto cond = foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_or(a, b, span); },
- const_false(1), conds);
- return IfThenElse(cond, update, body);
-}
-
} // namespace te
} // namespace tvm
diff --git a/src/te/operation/compute_op.h b/src/te/operation/compute_op.h
deleted file mode 100644
index 944334a41f..0000000000
--- a/src/te/operation/compute_op.h
+++ /dev/null
@@ -1,104 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Helper utilities to implement compute_op.
- * \file compute_op.h
- */
-#ifndef TVM_TE_OPERATION_COMPUTE_OP_H_
-#define TVM_TE_OPERATION_COMPUTE_OP_H_
-
-#include <tvm/te/operation.h>
-#include <tvm/tir/expr.h>
-
-#include <unordered_map>
-#include <vector>
-
-namespace tvm {
-namespace te {
-// loop nest structure for general compute
-// This the loop nest structured used in compute.
-// Does not include the loop body.
-struct ComputeLoopNest {
- // The common number of loops between init and main
- size_t num_common_loop;
- // predicates for the initialize loop
- std::vector<PrimExpr> init_predicates;
- // Initialization nest involved.
- std::vector<std::vector<Stmt>> init_nest;
- // Value map for the init code
- std::unordered_map<IterVar, PrimExpr> init_vmap;
- // Predicates for the main update loop
- std::vector<PrimExpr> main_predicates;
- // The general loop nest
- std::vector<std::vector<Stmt>> main_nest;
- // Value map for the IterVar.
- std::unordered_map<IterVar, PrimExpr> main_vmap;
-
- /*!
- * \brief constructor to build ComputeOpNest
- * \param self The pointer to compute op.
- * \param stage The scxhedule stage.
- * \param dom_map The domain map.
- * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1
- * \return The constructed loop nest
- */
- static ComputeLoopNest Create(const BaseComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop);
-};
-
-/*!
- * \brief Build body of compute for cross thread reduction pattern.
- * \param self The pointer to ComputeOpNode
- * \param stage The schedule stage.
- * \param dom_map The domain map.
- * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1
- * \return The created statement.
- */
-Stmt MakeCrossThreadReduction(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop);
-
-/*!
- * \brief Build body of compute for tensorization.
- * \param self The pointer to ComputeOpNode
- * \param stage The schedule stage.
- * \param dom_map The domain map.
- * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1
- * \return The created statement.
- */
-Stmt MakeTensorize(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map, bool debug_keep_trivial_loop);
-
-/*!
- * \brief Transform the update part when there is no init func in tensorizing
- * \param stage The stage for tensorizing.
- * \param dom_map The range of each iter var.
- * \param n The loop nest structured used in compute.
- * \param body The body func in tensorize intrin
- * \param update The update func in tensorize intrin
- * \return Transformed result.
- */
-Stmt TransformUpdate(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- const ComputeLoopNest& n, Stmt body, Stmt update);
-} // namespace te
-} // namespace tvm
-
-#endif // TVM_TE_OPERATION_COMPUTE_OP_H_
diff --git a/src/te/operation/create_primfunc.cc b/src/te/operation/create_primfunc.cc
index 2709bd2f94..9a6d30de4a 100644
--- a/src/te/operation/create_primfunc.cc
+++ b/src/te/operation/create_primfunc.cc
@@ -22,6 +22,7 @@
#include <tvm/arith/analyzer.h>
#include <tvm/ir/name_supply.h>
#include <tvm/runtime/registry.h>
+#include <tvm/te/operation.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/data_type_rewriter.h>
#include <tvm/tir/function.h>
@@ -37,7 +38,7 @@
#include "../../support/array.h"
#include "../../tir/ir/functor_common.h"
#include "../../tir/transforms/ir_utils.h"
-#include "../schedule/graph.h"
+#include "graph.h"
namespace tvm {
namespace tir {
diff --git a/src/te/operation/cross_thread_reduction.cc b/src/te/operation/cross_thread_reduction.cc
deleted file mode 100644
index 52e38c7ba2..0000000000
--- a/src/te/operation/cross_thread_reduction.cc
+++ /dev/null
@@ -1,237 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Logics related to cross thread reduction, used by ComputeOpNode.
- * \file cross_thread_reduction.cc
- */
-#include <tvm/tir/builtin.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include "compute_op.h"
-#include "op_utils.h"
-
-namespace tvm {
-namespace te {
-using namespace tir;
-
-//
-// Cross thread reduction transformation.
-//
-// The input loop nest in generic form (single reduction/thread case)
-//
-// let m be the reduction extent
-// let N be the thread extent
-// let input_pred be the predicate on the reduction
-//
-// B[..] = 0
-// for (tid, 0, N)
-// for (i, 0, floordiv(m+N-1, N))
-// if (i + tid * floordiv(m+N-1, N) < m)
-// if (input_pred)
-// B[..] = op(B[..], A[i + tid * floordiv(m+N-1,N)])
-//
-// The threaded reduction looks like
-//
-// (1) normal reductions (leaves)
-// for (i, 0, floordiv(m+N-1, N))
-// if (i + tid * floordiv(m+N-1, N) < m)
-// if (input_pred)
-// B_temp[0] = op(B_temp[0], A[i + tid * floordiv(m+N-1,N)])
-//
-// (2) threaded reduction does not require predicates as an identity
-// element will be filled if out of bounds.
-//
-// tvm_thread_allreduce(size, B_temp, (bool)1, tid)
-//
-// The last step is to write the final reduction variable,
-// which should be predicated by the existing input_pred if any
-// The consequence is that input_pred should be independent of
-// the reduction axis. Otherwise, we need to separate it into
-// dependent part and independent one.
-//
-// (3) write back
-// if (input_pred)
-// B[..] = B_temp[0]
-//
-// In summary, we are going to need two predicates
-//
-// * the original input_pred from reduction itself
-//
-// * the normal reduction axis predicate
-// normal_pred = (i + tid * floordiv(m+N-1,N)) < m
-// this predicate depends on the normal reduction variable.
-//
-// input_pred will be applied to both normal reduction and
-// the writeback step.
-//
-Stmt MakeCrossThreadReduction(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) {
- Array<PrimExpr> args;
- for (IterVar iv : self->axis) {
- args.push_back(iv->var);
- }
- std::unordered_map<IterVar, PrimExpr> value_map;
- auto nest = MakeLoopNest(stage, dom_map, 0, false, std::unordered_set<IterVar>(), &value_map,
- debug_keep_trivial_loop);
-
- size_t size = self->body.size();
- ICHECK_GT(size, 0);
- std::vector<const ReduceNode*> reduces(size);
- for (size_t i = 0; i < size; ++i) {
- const ReduceNode* reduce = self->body[i].as<ReduceNode>();
- ICHECK(reduce);
- ICHECK(reduce->init.empty())
- << "Cannot perform cross_thread_reduction for reductions with init";
- reduces[i] = reduce;
- }
-
- // This computes the bound checking predicates in normal reduction.
- auto normal_preds =
- MakeBoundCheck(stage, dom_map, value_map, false, std::unordered_set<IterVar>());
-
- // normal_pred = input_pred && normal_pred
- PrimExpr input_pred = reduces[0]->condition;
- normal_preds.push_back(input_pred);
- normal_preds.erase(std::remove_if(normal_preds.begin(), normal_preds.end(),
- [](const PrimExpr& e) { return !e.defined(); }),
- normal_preds.end());
-
- std::vector<std::vector<Stmt>> common, normal_red;
- for (size_t i = 0, n = stage->leaf_iter_vars.size(); i < n; ++i) {
- IterVar iv = stage->leaf_iter_vars[i];
- IterVarAttr attr;
- auto it = stage->iter_var_attrs.find(iv);
- if (it != stage->iter_var_attrs.end()) {
- attr = (*it).second;
- }
- if (iv->iter_type == kCommReduce) {
- if (attr.defined() && attr->bind_thread.defined()) {
- common.emplace_back(nest[i + 1]);
- } else {
- normal_red.emplace_back(nest[i + 1]);
- }
- } else {
- common.emplace_back(nest[i + 1]);
- }
- }
-
- // If we load from and then store into the same res_handles in the thread_allreduce intrinsic,
- // something goes wrong, so we use an extra variable here for normal reduction.
- std::vector<Buffer> normal_res_buffers;
- std::vector<Stmt> normal_init, normal_update;
- if (!normal_red.empty()) {
- normal_res_buffers.reserve(size);
- normal_init.reserve(size);
- normal_update.resize(size);
- const CommReducerNode* combiner = reduces[0]->combiner.as<CommReducerNode>();
- ICHECK(combiner);
- Array<PrimExpr> lhs;
- for (size_t i = 0; i < size; ++i) {
- normal_res_buffers.push_back(
- decl_buffer({1}, reduces[i]->dtype, "normal_reduce_temp" + std::to_string(i), "local"));
- lhs.push_back(BufferLoad(normal_res_buffers[i], {0}));
- }
- Array<PrimExpr> init_value = combiner->identity_element;
- Array<PrimExpr> update_value = (*combiner)(lhs, reduces[0]->source);
- for (size_t i = 0; i < size; ++i) {
- normal_init.emplace_back(BufferStore(normal_res_buffers[i], init_value[i], {0}));
- normal_update.emplace_back(BufferStore(normal_res_buffers[i], update_value[i], {0}));
- }
- }
-
- Array<PrimExpr> freduce_args;
- freduce_args.push_back(make_const(DataType::UInt(32), static_cast<uint32_t>(size)));
- for (size_t i = 0; i < size; ++i) {
- if (!normal_red.empty()) {
- freduce_args.push_back(BufferLoad(normal_res_buffers[i], {0}));
- } else {
- freduce_args.push_back(reduces[0]->source[i]);
- }
- }
-
- // No constraints on the thread reduction step. It may have redundent
- // computation for rare cases. TODO(tvm-team): revisit this.
- freduce_args.push_back(const_true(1));
- std::vector<Buffer> res_buffers(size);
- for (size_t idx = 0; idx < size; ++idx) {
- res_buffers[idx] =
- decl_buffer({1}, reduces[idx]->dtype, "reduce_temp" + std::to_string(idx), "local");
- // Make a BufferLoad object so that we can pass the entire Buffer
- // object through to LowerThreadAllreduce. The index here is
- // unused.
- PrimExpr dummy_load = BufferLoad(res_buffers[idx], {0});
- freduce_args.push_back(dummy_load);
- }
-
- // Checks for the thread.
- std::vector<PrimExpr> output_preds;
- if (stage->store_predicate.defined()) {
- output_preds.emplace_back(stage->store_predicate);
- }
-
- for (IterVar iv : stage->leaf_iter_vars) {
- if (iv->iter_type == kCommReduce) {
- auto it = stage->iter_var_attrs.find(iv);
- if (it != stage->iter_var_attrs.end() && (*it).second->bind_thread.defined()) {
- IterVar tv = (*it).second->bind_thread;
- freduce_args.push_back(tv->var);
- output_preds.push_back(tv->var == make_const(tv->var->dtype, 0));
- }
- }
- }
-
- // Apply the existing input predicate if any.
- output_preds.push_back(input_pred);
-
- Stmt reduce_body =
- Evaluate(Call(DataType::Handle(), tir::builtin::tvm_thread_allreduce(), freduce_args));
- reduce_body = AttrStmt(reduces[0]->combiner, tir::attr::reduce_scope,
- make_zero(DataType::Handle()), reduce_body);
-
- if (!normal_red.empty()) {
- Stmt init_body = SeqStmt::Flatten(normal_init);
- Stmt update_body = SeqStmt::Flatten(normal_update);
- update_body = MergeNest(MakeIfNest(normal_preds), update_body);
- update_body = MergeNest(normal_red, update_body);
- reduce_body = SeqStmt::Flatten(init_body, update_body, reduce_body);
- }
-
- std::vector<Stmt> assigns(size);
- for (size_t idx = 0; idx < size; ++idx) {
- assigns[idx] = ProducerStore(stage->op.output(idx), BufferLoad(res_buffers[idx], {0}), args);
- }
- Stmt assign_body = SeqStmt::Flatten(assigns);
- assign_body = MergeNest(MakeIfNest(output_preds), assign_body);
- Stmt body = SeqStmt::Flatten(reduce_body, assign_body);
- for (size_t idx = size; idx != 0; --idx) {
- const auto& res_buffer = res_buffers[idx - 1];
- body = Allocate(res_buffer->data, res_buffer->dtype, res_buffer->shape, const_true(), body);
- if (!normal_red.empty()) {
- const auto& normal_res_buffer = normal_res_buffers[idx - 1];
- body = Allocate(normal_res_buffer->data, normal_res_buffer->dtype, normal_res_buffer->shape,
- const_true(), body);
- }
- }
- body = Substitute(body, value_map);
- return MergeNest(common, body);
-}
-} // namespace te
-} // namespace tvm
diff --git a/src/te/operation/extern_op.cc b/src/te/operation/extern_op.cc
index b602efcfc2..b2d2e91415 100644
--- a/src/te/operation/extern_op.cc
+++ b/src/te/operation/extern_op.cc
@@ -26,10 +26,6 @@
#include <tvm/te/operation.h>
#include <tvm/tir/expr.h>
-#include <unordered_set>
-
-#include "op_utils.h"
-
namespace tvm {
namespace te {
using namespace tir;
@@ -44,8 +40,6 @@ TVM_REGISTER_NODE_TYPE(ExternOpNode);
int ExternOpNode::num_outputs() const { return static_cast<int>(output_placeholders.size()); }
-Array<IterVar> ExternOpNode::root_iter_vars() const { return {}; }
-
DataType ExternOpNode::output_dtype(size_t i) const { return output_placeholders[i]->dtype; }
Array<PrimExpr> ExternOpNode::output_shape(size_t i) const { return output_placeholders[i]->shape; }
@@ -85,83 +79,5 @@ TVM_REGISTER_GLOBAL("te.ExternOp")
Array<Tensor> ExternOpNode::InputTensors() const { return inputs; }
-Operation ExternOpNode::ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const {
- ICHECK_EQ(self.operator->(), this);
- auto n = make_object<ExternOpNode>(*this);
- n->body = ReplaceTensor(this->body, rmap);
- for (size_t i = 0; i < n->inputs.size(); ++i) {
- Tensor t = n->inputs[i];
- if (rmap.count(t)) {
- n->inputs.Set(i, rmap.at(t));
- }
- }
-
- if (body.same_as(n->body) && inputs.same_as(n->inputs)) {
- return self;
- } else {
- return Operation(n);
- }
-}
-
-void ExternOpNode::PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const {
- for (Tensor t : this->inputs) {
- auto it = out_dom_map->find(t);
- if (it == out_dom_map->end()) continue;
- TensorDom& dom = it->second;
- for (size_t i = 0; i < t->shape.size(); ++i) {
- dom.data[i].emplace_back(
- IntSet::FromRange(Range::FromMinExtent(make_const(t->shape[i].dtype(), 0), t->shape[i])));
- }
- }
-}
-
-void ExternOpNode::GatherBound(const Operation& self,
- const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const {}
-
-Stmt ExternOpNode::BuildRealize(const Stage& stage,
- const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope) const {
- ICHECK_EQ(stage->op.get(), this);
- Stmt realize_body = body;
- for (int k = 0; k < num_outputs(); ++k) {
- Tensor t = stage->op.output(k);
- Region bounds;
- for (size_t i = 0; i < t->shape.size(); ++i) {
- bounds.push_back(Range::FromMinExtent(make_const(t->shape[i].dtype(), 0), t->shape[i]));
- }
- realize_body = tir::ProducerRealize(t, bounds, const_true(), realize_body, storage_scope);
- }
- return realize_body;
-}
-
-Stmt ExternOpNode::BuildProvide(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const {
- ICHECK_EQ(stage->op.operator->(), this);
- Stmt ret = AttrStmt(make_zero(DataType::Int(32)), tir::attr::extern_scope, 0, this->body);
- auto f_push_bind = [&ret](Buffer buffer, Tensor tensor) {
- Array<ObjectRef> bind_spec;
- Array<PrimExpr> tuple;
- bind_spec.push_back(buffer);
- bind_spec.push_back(tensor);
- for (size_t k = 0; k < buffer->shape.size(); ++k) {
- tuple.push_back(make_const(buffer->shape[k].dtype(), 0));
- tuple.push_back(buffer->shape[k]);
- }
- ret = AttrStmt(bind_spec, tir::attr::buffer_bind_scope,
- Call(DataType::Handle(), builtin::tvm_tuple(), tuple), ret);
- };
- for (size_t i = output_placeholders.size(); i != 0; --i) {
- f_push_bind(output_placeholders[i - 1], stage->op.output(i - 1));
- }
- for (size_t i = inputs.size(); i != 0; --i) {
- f_push_bind(input_placeholders[i - 1], inputs[i - 1]);
- }
- return ret;
-}
} // namespace te
} // namespace tvm
diff --git a/src/te/operation/graph.cc b/src/te/operation/graph.cc
new file mode 100644
index 0000000000..cddace2a82
--- /dev/null
+++ b/src/te/operation/graph.cc
@@ -0,0 +1,90 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file graph.cc
+ * \brief Utilities to get information about schedule graph.
+ */
+#include "graph.h"
+
+#include <tvm/runtime/registry.h>
+#include <tvm/te/operation.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt_functor.h>
+
+#include <unordered_set>
+
+namespace tvm {
+namespace te {
+
+// construct a read graph that gives readers of each operation
+// that the root depend on
+ReadGraph CreateReadGraph(const Array<Operation>& roots) {
+ ReadGraph rmap;
+ std::vector<Operation> stack;
+ std::unordered_set<const Object*> visited;
+ // initialize the roots
+ for (Operation op : roots) {
+ stack.push_back(op);
+ visited.insert(op.get());
+ }
+
+ while (!stack.empty()) {
+ Operation op = stack.back();
+ stack.pop_back();
+ Array<Tensor> deps = op->InputTensors();
+ rmap.Set(op, deps);
+ for (Tensor t : deps) {
+ if (t->op.defined() && visited.count(t->op.get()) == 0) {
+ visited.insert(t->op.get());
+ stack.push_back(t->op);
+ }
+ }
+ }
+ return rmap;
+}
+
+void PostDFSOrder(const Operation& op, const ReadGraph& g, std::unordered_set<Operation>* visited,
+ Array<Operation>* post_order) {
+ if (visited->count(op)) return;
+ visited->insert(op);
+ for (const auto& t : g.at(op)) {
+ PostDFSOrder(t->op, g, visited, post_order);
+ }
+ post_order->push_back(op);
+}
+
+Array<Operation> PostDFSOrder(const Array<Operation>& roots, const ReadGraph& g) {
+ std::unordered_set<Operation> visited;
+ Array<Operation> post_order;
+ for (Operation op : roots) {
+ PostDFSOrder(op, g, &visited, &post_order);
+ }
+ return post_order;
+}
+
+TVM_REGISTER_GLOBAL("schedule.CreateReadGraph").set_body_typed(CreateReadGraph);
+
+TVM_REGISTER_GLOBAL("schedule.PostDFSOrder")
+ .set_body_typed([](const Array<Operation>& roots, const ReadGraph& g) {
+ return PostDFSOrder(roots, g);
+ });
+
+} // namespace te
+} // namespace tvm
diff --git a/src/te/operation/graph.h b/src/te/operation/graph.h
new file mode 100644
index 0000000000..fbb1241ad5
--- /dev/null
+++ b/src/te/operation/graph.h
@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file graph.h
+ * \brief Utilities to get information about schedule graph.
+ */
+#ifndef TVM_TE_SCHEDULE_GRAPH_H_
+#define TVM_TE_SCHEDULE_GRAPH_H_
+
+#include <tvm/te/operation.h>
+#include <tvm/tir/expr.h>
+
+namespace tvm {
+namespace te {
+
+/*!
+ * \brief data structure of Operation->Tensors it reads
+ */
+using ReadGraph = Map<Operation, Array<Tensor>>;
+
+/*!
+ * \brief Get read graph of each operation to all the
+ * Tensors that it directly depends on.
+ *
+ * The result map contains Operations needed to finish root Operation.
+ * \param roots The root operation.
+ * \return The result map.
+ */
+ReadGraph CreateReadGraph(const Array<Operation>& roots);
+
+/*!
+ * \brief Get a post DFS ordered of operations in the graph.
+ * \param roots The root of the graph.
+ * \param g The read graph.
+ * \return vector order of Operations in PostDFS order.
+ *
+ * \note PostDFSOrder is a special case of Topoligical order,
+ * and can be used when topoligical order is needed.
+ */
+Array<Operation> PostDFSOrder(const Array<Operation>& roots, const ReadGraph& g);
+
+} // namespace te
+} // namespace tvm
+
+#endif // TVM_TE_SCHEDULE_GRAPH_H_
diff --git a/src/te/operation/hybrid_op.cc b/src/te/operation/hybrid_op.cc
deleted file mode 100644
index 85e764fea4..0000000000
--- a/src/te/operation/hybrid_op.cc
+++ /dev/null
@@ -1,486 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Hybrid computation rule.
- * \file hybrid_op.cc
- */
-#include "hybrid_op.h"
-
-#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
-#include <tvm/te/operation.h>
-#include <tvm/tir/analysis.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/op.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <string>
-#include <unordered_set>
-#include <utility>
-
-#include "op_utils.h"
-
-namespace tvm {
-namespace te {
-using namespace tir;
-// HybridOpNode
-TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
- .set_dispatch<HybridOpNode>([](const ObjectRef& node, ReprPrinter* p) {
- auto* op = static_cast<const HybridOpNode*>(node.get());
- p->stream << "hybrid(" << op->name << ", " << op << ")";
- });
-
-TVM_REGISTER_NODE_TYPE(HybridOpNode);
-
-int HybridOpNode::num_outputs() const { return static_cast<int>(outputs.size()); }
-
-Array<IterVar> HybridOpNode::root_iter_vars() const { return this->axis; }
-
-DataType HybridOpNode::output_dtype(size_t i) const { return outputs[i]->dtype; }
-
-Array<PrimExpr> HybridOpNode::output_shape(size_t i) const { return outputs[i]->shape; }
-
-HybridOp::HybridOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
- Array<Tensor> inputs, Array<Tensor> outputs, Stmt body) {
- if (!attrs.defined()) {
- attrs = Map<String, ObjectRef>();
- }
- auto n = make_object<HybridOpNode>();
- n->name = std::move(name);
- n->tag = std::move(tag);
- n->attrs = std::move(attrs);
- n->inputs = std::move(inputs);
- n->outputs = std::move(outputs);
- n->axis = te::GatherLoopVars(body);
- n->body = std::move(body);
- data_ = std::move(n);
-}
-
-TVM_REGISTER_GLOBAL("te.HybridOp")
- .set_body_typed([](std::string name, std::string tag, Map<String, ObjectRef> attrs,
- Array<Tensor> inputs, Array<Tensor> outputs,
- Stmt body) { return HybridOp(name, tag, attrs, inputs, outputs, body); });
-
-Array<Tensor> HybridOpNode::InputTensors() const {
- // Because input tensors could be potentially inlined into hybrid scripts,
- // we need to check if all input tensors are used in the body.
- std::unordered_set<Tensor> orig_inputs;
- for (auto t : inputs) {
- orig_inputs.insert(t);
- }
- std::unordered_set<Tensor> visited;
- Array<Tensor> curr_inputs;
- tir::PostOrderVisit(body, [&curr_inputs, &orig_inputs, &visited](const ObjectRef& n) {
- if (auto* pload = n.as<tir::ProducerLoadNode>()) {
- Tensor t = Downcast<Tensor>(pload->producer);
- if (orig_inputs.count(t) && !visited.count(t)) {
- curr_inputs.push_back(t);
- visited.insert(t);
- }
- }
- });
- return curr_inputs;
-}
-
-Operation HybridOpNode::ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const {
- ICHECK_EQ(self.operator->(), this);
- auto n = make_object<HybridOpNode>(*this);
- n->body = te::ReplaceTensor(this->body, rmap);
- for (size_t i = 0; i < n->inputs.size(); ++i) {
- Tensor t = n->inputs[i];
- if (rmap.count(t)) {
- n->inputs.Set(i, rmap.at(t));
- }
- }
-
- if (body.same_as(n->body) && inputs.same_as(n->inputs)) {
- return self;
- } else {
- return Operation(n);
- }
-}
-
-void HybridOpNode::PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const {
- auto curr_inputs = InputTensors();
- for (Tensor t : curr_inputs) {
- auto it = out_dom_map->find(t);
- if (it == out_dom_map->end()) continue;
- TensorDom& dom = it->second;
- for (size_t i = 0; i < t->shape.size(); ++i) {
- dom.data[i].emplace_back(
- IntSet::FromRange(Range::FromMinExtent(make_const(t->shape[i].dtype(), 0), t->shape[i])));
- }
- }
-}
-
-void HybridOpNode::GatherBound(const Operation& self,
- const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const {
- for (auto iter_var : axis) {
- ICHECK(!out_dom_map->count(iter_var));
- out_dom_map->operator[](iter_var) = iter_var->dom;
- }
-}
-
-Stmt HybridOpNode::BuildRealize(const Stage& stage,
- const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope) const {
- // TODO(@were): Add attribute inject here and remove it from hybrid parser.
- ICHECK_EQ(stage->op.get(), this);
- Stmt realize_body = body;
- for (int k = 0; k < num_outputs(); ++k) {
- Tensor t = stage->op.output(k);
- Region bounds;
- for (size_t i = 0; i < t->shape.size(); ++i) {
- bounds.push_back(Range::FromMinExtent(make_const(t->shape[i].dtype(), 0), t->shape[i]));
- }
- realize_body = tir::ProducerRealize(t, bounds, const_true(), realize_body, storage_scope);
- }
- return realize_body;
-}
-
-Stmt HybridOpNode::BuildProvide(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const {
- ICHECK_EQ(stage->op.operator->(), this);
- Stmt ret = AttrStmt(make_zero(DataType::Int(32)), tir::attr::extern_scope, 0, this->body);
- std::unordered_map<Tensor, Tensor> rmap;
- for (int i = 0; i < this->num_outputs(); ++i) {
- rmap[outputs[i]] = stage->op.output(i);
- }
- auto n = make_object<HybridOpNode>(*this);
- /* This is a story little bit complicated.
- * The following two lines of codes replace output tensors' usage.
- * This is the simplest way I (@were) can come up with to glue
- * hybrid operation node to TVM op system.
- * In hybrid script all the tensors, especially the output tensors,
- * have their own names defined by the users. However, In TVM
- * conventional ops:
- * 1. Output tensors refer the corresponding op node so that the output
- * tensors have the same names as the operation produces them.
- * 2. Once OpNode is wrapped up by an Operation node, it is finalized.
- * Later access will be from a const OpNode*.
- * This is a chicken-egg paradox. It is impossible to put the output
- * tensors into the function body without forming the op node. The
- * function body is immutable after the node is formed.
- *
- * Finally, I decided to resolve this issue "lazily". During the
- * pipeline of compilation, this stage is a very preliminary stage.
- * Technically, it is before Phase 0. The actual tensors will be replaced
- * here.
- * Thus, the operation body is slightly different from the Phase 0 body.
- * This is a major difference that HybridOpNode is NOT the same as
- * ExternOpNode.
- * */
- ret = te::ReplaceTensor(ret, rmap);
- ret = te::ReplaceProvideTensor(ret, rmap);
-
- ret = te::ApplySchedule(stage, dom_map, ret);
- return ret;
-}
-
-Stmt ApplyLoopShapes(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- Stmt stmt) {
- class LoopSpliter : public StmtExprMutator {
- PrimExpr factor;
- const VarNode* parent;
- IterVar inner, outer;
-
- public:
- bool splitted;
- LoopSpliter(const SplitNode* split, const std::unordered_map<IterVar, Range>& dom_map)
- : factor(split->factor), splitted(false) {
- parent = split->parent->var.get();
-
- auto& inner_ = split->inner;
- ICHECK(dom_map.count(inner_));
- auto& inner_dom = dom_map.find(inner_)->second;
- ICHECK(is_const_int(inner_dom->min, 0));
-
- auto& outer_ = split->outer;
- ICHECK(dom_map.count(outer_));
- auto& outer_dom = dom_map.find(outer_)->second;
- ICHECK(is_const_int(outer_dom->min, 0));
-
- inner = IterVar(inner_dom, inner_->var, inner_->iter_type);
- outer = IterVar(outer_dom, outer_->var, outer_->iter_type);
- }
-
- Stmt VisitStmt_(const ForNode* op) final {
- if (op->loop_var.get() == parent) {
- Stmt ret = tir::Substitute(op->body, {{op->loop_var, inner + outer * factor}});
- PrimExpr cond = likely(outer * factor < (op->extent - inner));
- ret = IfThenElse(cond, ret);
- ret = For(inner->var, PrimExpr(0), inner->dom->extent,
- IterVarTypeToForKind(inner->iter_type), ret);
- ret = For(outer->var, PrimExpr(0), outer->dom->extent,
- IterVarTypeToForKind(outer->iter_type), ret);
- splitted = true;
- return ret;
- }
- return StmtExprMutator::VisitStmt_(op);
- }
- };
-
- class LoopFuser : public StmtExprMutator {
- const IterVar& parent;
- const VarNode* inner;
- const VarNode* outer;
- bool under_outer;
- PrimExpr extent;
-
- public:
- bool fused;
- explicit LoopFuser(const FuseNode* fuse_)
- : parent(fuse_->fused),
- inner(fuse_->inner->var.get()),
- outer(fuse_->outer->var.get()),
- under_outer(false),
- extent(0),
- fused(false) {}
-
- // TODO(@were): Handle imperfect loops
- Stmt VisitStmt_(const ForNode* op) final {
- if (op->loop_var.get() == inner) {
- ICHECK(under_outer);
- std::unordered_map<const VarNode*, PrimExpr> rmap;
- rmap[op->loop_var.get()] = indexmod(parent, op->extent);
- extent = op->extent;
- fused = true;
- return tir::Substitute(op->body, rmap);
- } else if (op->loop_var.get() == outer) {
- under_outer = true;
- Stmt body = this->VisitStmt(op->body);
- std::unordered_map<const VarNode*, PrimExpr> rmap;
- rmap[op->loop_var.get()] = indexdiv(parent, extent);
- body = tir::Substitute(body, rmap);
- under_outer = false;
- return For(parent->var, PrimExpr(0), extent * op->extent, op->kind, body,
- op->thread_binding, op->annotations);
- } else if (under_outer) {
- Stmt body = this->VisitStmt(op->body);
- std::unordered_map<const VarNode*, PrimExpr> rmap;
- rmap[op->loop_var.get()] = indexmod(indexdiv(parent, extent), op->extent);
- body = tir::Substitute(body, rmap);
- extent = extent * op->extent;
- return body;
- }
- return StmtExprMutator::VisitStmt_(op);
- }
- };
-
- for (auto& rel : stage->relations) {
- if (const SplitNode* split = rel.as<SplitNode>()) {
- LoopSpliter Spliter(split, dom_map);
- stmt = Spliter(stmt);
- ICHECK(Spliter.splitted);
- } else if (const FuseNode* fuse = rel.as<FuseNode>()) {
- LoopFuser Fuser(fuse);
- stmt = Fuser(stmt);
- ICHECK(Fuser.fused);
- }
- }
-
- return stmt;
-}
-
-Stmt ApplyLoopAnnotations(const Stage& stage, const std::unordered_map<IterVar, IterVar>& rebased,
- Stmt stmt) {
- class LoopAnnotator : public StmtMutator {
- const VarNode* var;
- const IterVarAttr& attr;
-
- public:
- LoopAnnotator(const VarNode* var_, const IterVarAttr& attr_) : var(var_), attr(attr_) {}
-
- Stmt VisitStmt_(const ForNode* op) final {
- tir::ExprDeepEqual expr_equal;
-
- if (op->loop_var.get() == var) {
- if (attr->bind_thread.defined()) {
- const auto& iter_var = attr->bind_thread;
- if (iter_var->dom.defined()) {
- ICHECK(is_const_int(iter_var->dom->min, 0));
- ICHECK(expr_equal(iter_var->dom->extent, op->extent))
- << "Thread extent and loop extent mismatch!\n";
- }
- std::unordered_map<const VarNode*, PrimExpr> rmap;
- rmap[op->loop_var.get()] = iter_var;
- Stmt body = tir::Substitute(op->body, rmap);
- return AttrStmt(iter_var, "thread_extent", op->extent, body);
- } else {
- return For(op->loop_var, op->min, op->extent, IterVarTypeToForKind(attr->iter_type),
- op->body, op->thread_binding, op->annotations);
- }
- }
- return StmtMutator::VisitStmt_(op);
- }
- };
-
- for (auto& iter_var : stage->leaf_iter_vars) {
- bool need_change = false;
- int found = 0;
-
- const IterVar& actual = rebased.count(iter_var) ? rebased.find(iter_var)->second : iter_var;
- const VarNode* var = actual->var.get();
- ForKind expected = IterVarTypeToForKind(iter_var->iter_type);
- IterVarAttr attr;
- if (stage->iter_var_attrs.count(iter_var)) {
- attr = stage->iter_var_attrs[iter_var];
- expected = IterVarTypeToForKind(attr->iter_type);
- }
-
- PostOrderVisit(stmt, [&found, &var, &attr, &expected, &need_change](const ObjectRef& node) {
- if (const ForNode* op = node.as<ForNode>()) {
- if (op->loop_var.get() == var) {
- ++found;
- need_change = expected != op->kind || (attr.defined() && attr->bind_thread.defined());
- }
- }
- });
-
- ICHECK_EQ(found, 1) << " iter var should be found exactly once!";
- if (need_change) {
- stmt = LoopAnnotator(var, attr)(std::move(stmt));
- }
- }
- return stmt;
-}
-
-Stmt ApplyLoopOrder(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- const std::unordered_map<IterVar, IterVar>& rebased, Stmt stmt) {
- std::vector<const VarNode*> current_order;
- PostOrderVisit(stmt, [¤t_order](const ObjectRef& node) {
- if (const ForNode* op = node.as<ForNode>()) current_order.push_back(op->loop_var.get());
- });
- std::reverse(current_order.begin(), current_order.end());
- auto& required_ord = stage->leaf_iter_vars;
- ICHECK_EQ(current_order.size(), required_ord.size()) << "Cannot reorder the loops!";
- std::unordered_map<const VarNode*, IterVar> reorder;
- bool need_reorder = false;
- for (size_t i = 0; i < current_order.size(); ++i) {
- auto& current = current_order[i];
- const IterVar& iter_var = required_ord[i];
- const IterVar& required = rebased.count(iter_var) ? rebased.find(iter_var)->second : iter_var;
- ICHECK(required->dom.defined() || dom_map.count(required)) << required << "\n";
- reorder[current] = required;
- if (current != required->var.get()) {
- need_reorder = true;
- }
- }
-
- class LoopReorder : public StmtMutator {
- const Stage& stage;
- const std::unordered_map<IterVar, Range>& dom_map;
- const std::unordered_map<const VarNode*, IterVar>& reorder;
-
- public:
- LoopReorder(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- const std::unordered_map<const VarNode*, IterVar>& reorder)
- : stage(stage), dom_map(dom_map), reorder(reorder) {}
-
- Stmt VisitStmt_(const ForNode* op) final {
- // Reorder from in to out
- Stmt body_ = this->VisitStmt(op->body);
- ICHECK(reorder.count(op->loop_var.get()));
- auto target = reorder.find(op->loop_var.get())->second;
- if (body_.same_as(op->body) && op->loop_var.get() == target->var.get())
- return GetRef<Stmt>(op);
- const Stmt& body = op->body.same_as(body_) ? op->body : body_;
- ForKind kind = IterVarTypeToForKind(target->iter_type);
- if (stage->iter_var_attrs.count(target)) {
- kind = IterVarTypeToForKind(stage->iter_var_attrs[target]->iter_type);
- }
- const Range& range = target->dom.defined() ? target->dom : dom_map.find(target)->second;
- return For(target->var, range->min, range->extent, kind, body, op->thread_binding,
- op->annotations);
- }
- };
-
- if (need_reorder) return LoopReorder(stage, dom_map, reorder)(stmt);
-
- return stmt;
-}
-
-Stmt ApplySchedule(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- Stmt stmt) {
- // TODO(@were): Eliminate loop rebase in script parser and move the burden here
- // Gather rebased variables
- std::unordered_map<IterVar, IterVar> rebased;
- for (auto rel : stage->relations) {
- if (const auto* rebase = rel.as<RebaseNode>()) {
- rebased[rebase->rebased] = rebase->parent;
- ICHECK(rebase->parent->dom.defined());
- ICHECK(dom_map.count(rebase->rebased));
- }
- }
- stmt = ApplyLoopShapes(stage, dom_map, stmt);
- stmt = ApplyLoopOrder(stage, dom_map, rebased, stmt);
- stmt = ApplyLoopAnnotations(stage, rebased, stmt);
- return stmt;
-}
-
-std::vector<IterVar> GatherLoopVars(Stmt stmt) {
- // TODO(@were): Write a comprehensive pass to analyze iter var types
- std::vector<IterVar> res_;
- PostOrderVisit(stmt, [&res_](const ObjectRef& node) {
- if (const ForNode* op = node.as<ForNode>()) {
- Var loop_var(op->loop_var);
- Range dom = Range::FromMinExtent(op->min, cast(loop_var.dtype(), op->extent));
- res_.push_back(IterVar(dom, loop_var, ForKindToIterVarType(op->kind)));
- }
- });
- std::reverse(res_.begin(), res_.end());
- return res_;
-}
-
-// replacer to replace tensors' usage in Provide
-class ProviderReplacer : public tir::StmtMutator {
- public:
- explicit ProviderReplacer(const std::unordered_map<Tensor, Tensor>& vmap) : vmap_(vmap) {}
-
- Stmt VisitStmt_(const tir::ProducerStoreNode* op) final {
- Tensor t = Downcast<Tensor>(op->producer);
- auto it = vmap_.find(t);
- if (it != vmap_.end()) {
- Stmt ret = tir::ProducerStore(it->second, op->value, op->indices);
- found = true;
- return this->VisitStmt(ret);
- }
- return StmtMutator::VisitStmt_(op);
- }
-
- // whether it is found.
- bool found{false};
-
- private:
- const std::unordered_map<Tensor, Tensor>& vmap_;
-};
-
-Stmt ReplaceProvideTensor(Stmt stmt, const std::unordered_map<Tensor, Tensor>& replace) {
- ProviderReplacer repl(replace);
- Stmt ret = repl(stmt);
- return repl.found ? ret : stmt;
-}
-} // namespace te
-} // namespace tvm
diff --git a/src/te/operation/hybrid_op.h b/src/te/operation/hybrid_op.h
deleted file mode 100644
index 705456850c..0000000000
--- a/src/te/operation/hybrid_op.h
+++ /dev/null
@@ -1,94 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Helper utilities to implement hybrid_op.
- * \file hybrid_op.h
- */
-#ifndef TVM_TE_OPERATION_HYBRID_OP_H_
-#define TVM_TE_OPERATION_HYBRID_OP_H_
-
-#include <tvm/te/schedule.h>
-#include <tvm/tir/expr.h>
-
-#include <unordered_map>
-#include <unordered_set>
-#include <vector>
-
-#include "../../tir/transforms/arg_binder.h"
-#include "../../tir/transforms/ir_utils.h"
-#include "../schedule/message_passing.h"
-
-namespace tvm {
-namespace te {
-
-/*!
- * \brief Find all the iteration variables in the given statement body.
- * \param stmt The body to be inspected.
- */
-std::vector<IterVar> GatherLoopVars(Stmt stmt);
-
-/*!
- * \brief Replace the tensor reference (especially in Provide's) in stmt by the replace map.
- * \param stmt The statement to be processed.
- * \param replace The replacement rule.
- */
-Stmt ReplaceProvideTensor(Stmt stmt, const std::unordered_map<Tensor, Tensor>& replace);
-
-/*!
- * \brief Apply the schedule manipulation on the function body.
- * \param stmt The statement to be processed.
- * \param dom_map The extents of the iterative variables may be used.
- * \param stage The schedule information to be applied.
- */
-Stmt ApplySchedule(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- Stmt stmt);
-
-/*!
- * \brief Apply loop splits and fuses in the schedule on the function body.
- * \param stage The schedule information to be applied.
- * \param dom_map The extents of the iterative variables may be used.
- * \param stmt The statement to be processed.
- */
-Stmt ApplyLoopShapes(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- Stmt stmt);
-
-/*!
- * \brief Apply loop annotation in the schedule on the function body.
- * \param stage The schedule information to be applied.
- * \param rebased The map specifies the rebase, a.k.a rename, relationship of these variables.
- * \param stmt The statement to be processed.
- */
-Stmt ApplyLoopAnnotations(const Stage& stage, const std::unordered_map<IterVar, IterVar>& rebased,
- Stmt stmt);
-
-/*!
- * \brief Apply loop order in the schedule on the function body.
- * \param stage The schedule information to be applied.
- * \param dom_map The extents of the iterative variables may be used.
- * \param rebased The map specifies the rebase, a.k.a rename, relationship of these variables.
- * \param stmt The statement to be processed.
- */
-Stmt ApplyLoopOrder(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- const std::unordered_map<IterVar, IterVar>& rebased, Stmt stmt);
-
-} // namespace te
-} // namespace tvm
-
-#endif // TVM_TE_OPERATION_HYBRID_OP_H_
diff --git a/src/te/operation/op_utils.cc b/src/te/operation/op_utils.cc
deleted file mode 100644
index 7168933a32..0000000000
--- a/src/te/operation/op_utils.cc
+++ /dev/null
@@ -1,290 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Utility to make loop nest.
- * \file op_utils.cc
- */
-#include "op_utils.h"
-
-#include <tvm/te/operation.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <string>
-
-#include "../../runtime/thread_storage_scope.h"
-#include "../schedule/message_passing.h"
-
-namespace tvm {
-namespace te {
-
-using namespace arith;
-using namespace tir;
-
-std::vector<std::vector<Stmt>> MakeLoopNest(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- size_t begin_iter_pos, bool new_loop_var,
- const std::unordered_set<IterVar>& skip_iter,
- std::unordered_map<IterVar, PrimExpr>* p_value_map,
- bool debug_keep_trivial_loop) {
- auto leaf_iter_vars = stage->leaf_iter_vars;
- Stmt no_op = Evaluate(0);
- // create the loop nest
- std::vector<std::vector<Stmt>> nest;
- nest.resize(leaf_iter_vars.size() + 1);
- std::unordered_map<IterVar, PrimExpr>& value_map = *p_value_map;
-
- for (size_t i = begin_iter_pos; i < leaf_iter_vars.size(); ++i) {
- auto iv = leaf_iter_vars[i];
- if (skip_iter.count(iv) || iv->iter_type == kOpaque) {
- // skip this iteration.
- value_map[iv] = iv->var;
- continue;
- }
- // Bind iv could be another thread.
- IterVar bind_iv = iv;
- if (stage->iter_var_attrs.count(iv)) {
- IterVar bind_thread = stage->iter_var_attrs[iv]->bind_thread;
- if (bind_thread.defined()) bind_iv = bind_thread;
- }
-
- Range dom = dom_map.at(iv);
-
- ICHECK(iv->var.dtype() == dom->min.dtype() && iv->var.dtype() == dom->extent.dtype())
- << "iter_var type " << iv->var.dtype() << " and domain types (min:" << dom->min.dtype()
- << ", extent:" << dom->extent.dtype() << ") should all be the same";
-
- // This is a hack to ensure that the replacing expression has the same
- // dtype as the replacing expression. This happens when a thread/block
- // itervar is bound to another itervar. Because the thread/block itervar
- // has no way to know its correct dtype before it is bound, it defaults to
- // int32. Then the itervar it is bound to may have a different dtype. The
- // thread/block dtype really should be promoted to dtype of what it is
- // bound to (in `bind`) but that would require inplace modification of the
- // itervar.
- // XXX: we will get integer overflow if the bound itervar is greater than int32::max.
- auto promote_to_iv_dtype = [type = iv->var.dtype()](PrimExpr e) {
- return type != e.dtype() ? cast(type, e) : e;
- };
-
- // initialize the offset and loop_level
- Var var = bind_iv->var;
-
- // Mark the iter var in the IR, to remember the point
- if (bind_iv->thread_tag.length() == 0) {
- // Only generate new loop if we're not bound to a thread.
- if (new_loop_var) {
- var = Var(iv->var->name_hint + ".init", bind_iv->var.dtype());
- }
-
- ForKind kind = ForKind::kSerial;
- IterVarAttr it_attr;
- if (stage->iter_var_attrs.count(iv)) {
- it_attr = stage->iter_var_attrs[iv];
- }
- if (it_attr.defined()) {
- switch (it_attr->iter_type) {
- case kUnrolled:
- kind = ForKind::kUnrolled;
- break;
- case kVectorized:
- kind = ForKind::kVectorized;
- break;
- case kParallelized:
- kind = ForKind::kParallel;
- break;
- case kDataPar:
- break;
- case kTensorized:
- break;
- default:
- LOG(FATAL) << "Unknown iter type" << it_attr->iter_type << " in the iter_var_attrs";
- }
- ICHECK_EQ(it_attr->pragma_keys.size(), it_attr->pragma_values.size());
- for (size_t k = 0; k < it_attr->pragma_keys.size(); ++k) {
- const std::string& pkey = it_attr->pragma_keys[k].as<StringImmNode>()->value;
- PrimExpr pvalue = it_attr->pragma_values[k];
- if (!pvalue.defined()) {
- pvalue = make_const(DataType::Int(32), 1);
- }
- nest[i + 1].emplace_back(
- AttrStmt(iv, tir::attr::pragma_scope_prefix + pkey, pvalue, no_op));
- }
- }
- if (!debug_keep_trivial_loop && is_one(dom->extent)) {
- nest[i + 1].emplace_back(LetStmt(var, dom->min, no_op));
- value_map[iv] = dom->min;
- } else if (is_zero(dom->min)) {
- nest[i + 1].emplace_back(For(var, 0, dom->extent, kind, no_op));
- value_map[iv] = promote_to_iv_dtype(var);
- } else {
- Var idx(bind_iv->var->name_hint + ".idx", iv->var.dtype());
- nest[i + 1].emplace_back(For(idx, 0, dom->extent, kind, no_op));
- PrimExpr new_value = dom->min + idx;
- value_map[iv] = new_value;
- nest[i + 1].emplace_back(LetStmt(var, new_value, no_op));
- }
- if (it_attr.defined() && it_attr->prefetch_data.size() != 0) {
- ICHECK(!is_one(dom->extent)) << "Cannot prefetch on trivial loop with extent=1";
- ICHECK_EQ(it_attr->prefetch_data.size(), it_attr->prefetch_offset.size());
- for (size_t j = 0; j < it_attr->prefetch_data.size(); ++j) {
- nest[i + 1].emplace_back(AttrStmt(it_attr->prefetch_data[j], tir::attr::prefetch_scope,
- it_attr->prefetch_offset[j], no_op));
- }
- }
- } else if (bind_iv->thread_tag == "vthread" || bind_iv->thread_tag == "cthread") {
- // virtual thread
- // Always restrict threaded IterVar to starts from 0.
- ICHECK(is_zero(dom->min));
- ICHECK(is_positive_const(dom->extent));
- // annotate the extent of the IterVar
- nest[i + 1].emplace_back(AttrStmt(bind_iv, tir::attr::virtual_thread,
- cast(bind_iv->var.dtype(), dom->extent), no_op));
- value_map[iv] = promote_to_iv_dtype(var);
- } else if (bind_iv->thread_tag == "pipeline") {
- // pipeline marker.
- ICHECK(is_zero(dom->min));
- ICHECK(is_one(dom->extent));
- // annotate the extent of the IterVar
- nest[i + 1].emplace_back(AttrStmt(bind_iv, tir::attr::pipeline_exec_scope,
- cast(bind_iv->var.dtype(), dom->extent), no_op));
- value_map[iv] = dom->min;
- } else {
- // Always restrict threaded IterVar to starts from 0.
- ICHECK(is_zero(dom->min)) << "Itervar " << iv << " must start at zero, but it starts at "
- << dom->min;
- // annotate the extent of the IterVar
- nest[i + 1].emplace_back(AttrStmt(bind_iv, tir::attr::thread_extent,
- cast(bind_iv->var.dtype(), dom->extent), no_op));
- if (!debug_keep_trivial_loop && is_one(dom->extent)) {
- value_map[iv] = dom->min;
- } else if (stage->scope == "") {
- value_map[iv] = promote_to_iv_dtype(var);
- } else {
- runtime::ThreadScope ts = runtime::ThreadScope::Create(bind_iv->thread_tag);
- runtime::StorageScope ss = runtime::StorageScope::Create(stage->scope);
- if (static_cast<int>(ss.rank) <= ts.rank) {
- value_map[iv] = promote_to_iv_dtype(var);
- } else if (stage->scope == "warp" && ts.rank == 1) {
- // To determine whether a thread index is inside or outside a warp, we need
- // to know the thread extent. We leave a warning for now.
- if (ts.dim_index == 0) {
- value_map[iv] = promote_to_iv_dtype(var);
- } else {
- LOG(WARNING)
- << "WARNING: threadIdx.y or threadIdx.z accessing warp-scope memory detected. "
- << "TVM assumes only threadIdx.x indicates threads inside a warp, "
- << "while threadIdx.y and threadIdx.z indicates different warps.";
- value_map[iv] = dom->min;
- }
- } else {
- value_map[iv] = dom->min;
- }
- }
- }
- // annotate the extent of the IterVar
- if (!new_loop_var) {
- nest[i + 1].emplace_back(AttrStmt(iv, tir::attr::loop_scope, iv->var, no_op));
- }
- }
- // message passing to get offset of root iter vars.
- te::PassUpIndex(stage, dom_map, &value_map);
- return nest;
-}
-
-std::vector<Stmt> MakeIfNest(const std::vector<PrimExpr>& predicates) {
- Stmt no_op = Evaluate(0);
- std::vector<Stmt> nest;
- for (const PrimExpr& cond : predicates) {
- nest.emplace_back(IfThenElse(cond, no_op));
- }
- return nest;
-}
-
-// replacer to replace tensors
-class TensorReplacer : public tir::StmtExprMutator {
- public:
- explicit TensorReplacer(const std::unordered_map<Tensor, Tensor>& vmap) : vmap_(vmap) {}
-
- PrimExpr VisitExpr_(const tir::ProducerLoadNode* op) final {
- PrimExpr expr = StmtExprMutator::VisitExpr_(op);
- op = expr.as<tir::ProducerLoadNode>();
- ICHECK(op != nullptr);
-
- Tensor t = Downcast<Tensor>(op->producer);
- auto it = vmap_.find(t);
- if (it != vmap_.end()) {
- found = true;
- return tir::ProducerLoad(it->second, op->indices);
- } else {
- return expr;
- }
- }
-
- // whether it is found.
- bool found{false};
-
- private:
- const std::unordered_map<Tensor, Tensor>& vmap_;
-};
-
-Stmt ReplaceTensor(Stmt stmt, const std::unordered_map<Tensor, Tensor>& replace) {
- TensorReplacer repl(replace);
- Stmt ret = repl(stmt);
- return repl.found ? ret : stmt;
-}
-PrimExpr ReplaceTensor(PrimExpr expr, const std::unordered_map<Tensor, Tensor>& replace) {
- TensorReplacer repl(replace);
- PrimExpr ret = repl(expr);
- return repl.found ? ret : expr;
-}
-
-IterVarType ForKindToIterVarType(tir::ForKind kind) {
- switch (kind) {
- case ForKind::kSerial:
- return kDataPar;
- case ForKind::kParallel:
- return kParallelized;
- case ForKind::kVectorized:
- return kVectorized;
- case ForKind::kUnrolled:
- return kUnrolled;
- default:
- return kDataPar;
- }
-}
-
-tir::ForKind IterVarTypeToForKind(IterVarType iter_type) {
- switch (iter_type) {
- case kDataPar:
- return ForKind::kSerial;
- case kParallelized:
- return ForKind::kParallel;
- case kVectorized:
- return ForKind::kVectorized;
- case kUnrolled:
- return ForKind::kUnrolled;
- default:
- return ForKind::kSerial;
- }
-}
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/operation/op_utils.h b/src/te/operation/op_utils.h
deleted file mode 100644
index 72438cb8dd..0000000000
--- a/src/te/operation/op_utils.h
+++ /dev/null
@@ -1,96 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file op_utils.h
- * \brief Common utility used in operator construction.
- */
-#ifndef TVM_TE_OPERATION_OP_UTILS_H_
-#define TVM_TE_OPERATION_OP_UTILS_H_
-
-#include <tvm/te/schedule.h>
-#include <tvm/tir/expr.h>
-
-#include <unordered_map>
-#include <unordered_set>
-#include <vector>
-
-#include "../../tir/transforms/arg_binder.h"
-#include "../../tir/transforms/ir_utils.h"
-#include "../schedule/message_passing.h"
-
-namespace tvm {
-namespace te {
-
-using tir::MergeNest;
-
-/*!
- * \brief Build loop nest for stage.
- *
- * \param stage The stage to create a loop nest.
- * \param dom_map The range of each iter var.
- * \param begin_iter_pos The beginning position of leaf_iter_vars to generate loop.
- * \param new_loop_var Whether create new loop variable.
- * \param skip_iter Whether skip certain iteration.
- * \param p_value_map The result value of each IterVar.
- * \param debug_keep_trivial_loop Whether keep trivial loops with extent of 1
- */
-std::vector<std::vector<Stmt>> MakeLoopNest(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- size_t begin_iter_pos, bool new_loop_var,
- const std::unordered_set<IterVar>& skip_iter,
- std::unordered_map<IterVar, PrimExpr>* p_value_map,
- bool debug_keep_trivial_loop);
-
-/*!
- * \brief Create a nest of if checking the predicates.
- *
- * \param predicates The predicates to be checked.
- * \return List of If nest that checks the predicates.
- */
-std::vector<Stmt> MakeIfNest(const std::vector<PrimExpr>& predicates);
-
-/*!
- * \brief Replace the tensor reference (especially in Call's) in stmt by the replace map.
- * \param stmt The statement to be processed.
- * \param replace The replacement rule.
- */
-Stmt ReplaceTensor(Stmt stmt, const std::unordered_map<Tensor, Tensor>& replace);
-/*!
- * \brief Replace the tensor reference (especially in Call's) in primExpr by the replace map.
- * \param expr The expression to be processed.
- * \param replace The replacement rule.
- */
-PrimExpr ReplaceTensor(PrimExpr expr, const std::unordered_map<Tensor, Tensor>& replace);
-
-/*!
- * \brief Converts Halide ForKind to its corresponding IterVarType
- * \param kind The ForKind to be converted
- */
-IterVarType ForKindToIterVarType(tir::ForKind kind);
-
-/*!
- * \brief Converts IterVarType to its corresponding Halide ForKind
- * \param iter_type The IterVarType to be converted
- */
-tir::ForKind IterVarTypeToForKind(IterVarType iter_type);
-
-} // namespace te
-} // namespace tvm
-#endif // TVM_TE_OPERATION_OP_UTILS_H_
diff --git a/src/te/operation/placeholder_op.cc b/src/te/operation/placeholder_op.cc
index 774a0f8f1f..eaf1a8f854 100644
--- a/src/te/operation/placeholder_op.cc
+++ b/src/te/operation/placeholder_op.cc
@@ -38,8 +38,6 @@ TVM_REGISTER_NODE_TYPE(PlaceholderOpNode);
int PlaceholderOpNode::num_outputs() const { return 1; }
-Array<IterVar> PlaceholderOpNode::root_iter_vars() const { return {}; }
-
DataType PlaceholderOpNode::output_dtype(size_t i) const {
ICHECK_EQ(i, 0U);
return dtype;
@@ -79,30 +77,5 @@ TVM_REGISTER_GLOBAL("te.Placeholder")
Array<Tensor> PlaceholderOpNode::InputTensors() const { return {}; }
-Operation PlaceholderOpNode::ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const {
- return self;
-}
-
-void PlaceholderOpNode::PropBoundToInputs(
- const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const {}
-
-void PlaceholderOpNode::GatherBound(const Operation& self,
- const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const {}
-
-Stmt PlaceholderOpNode::BuildRealize(const Stage& stage,
- const std::unordered_map<IterVar, Range>& realize_map,
- const Stmt& body, String storage_scope) const {
- return body;
-}
-
-Stmt PlaceholderOpNode::BuildProvide(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const {
- return Stmt();
-}
} // namespace te
} // namespace tvm
diff --git a/src/te/operation/scan_op.cc b/src/te/operation/scan_op.cc
index 39689bd965..9f045a0692 100644
--- a/src/te/operation/scan_op.cc
+++ b/src/te/operation/scan_op.cc
@@ -25,9 +25,6 @@
#include <tvm/te/operation.h>
#include <tvm/tir/expr.h>
-#include "../schedule/graph.h"
-#include "op_utils.h"
-
namespace tvm {
namespace te {
using namespace tir;
@@ -40,13 +37,6 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
TVM_REGISTER_NODE_TYPE(ScanOpNode);
int ScanOpNode::num_outputs() const { return static_cast<int>(update.size()); }
-Array<IterVar> ScanOpNode::root_iter_vars() const {
- Array<IterVar> ret{scan_axis};
- for (IterVar iv : spatial_axis_) {
- ret.push_back(iv);
- }
- return ret;
-}
DataType ScanOpNode::output_dtype(size_t i) const { return update[i]->dtype; }
@@ -139,141 +129,5 @@ Array<Tensor> ScanOpNode::InputTensors() const {
return ret;
}
-Operation ScanOpNode::ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const {
- ICHECK_EQ(self.operator->(), this);
- auto n = make_object<ScanOpNode>(*this);
- for (size_t i = 0; i < n->init.size(); ++i) {
- if (rmap.count(n->init[i])) {
- n->init.Set(i, rmap.at(n->init[i]));
- }
- if (rmap.count(n->update[i])) {
- n->update.Set(i, rmap.at(n->update[i]));
- }
- }
- if (!n->init.same_as(init) || !n->update.same_as(update)) {
- return Operation(n);
- } else {
- return self;
- }
-}
-
-void ScanOpNode::PropBoundToInputs(const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const {
- ICHECK_EQ(self.operator->(), this);
- for (size_t i = 0, sp_idx = 0; i < this->init.size(); ++i) {
- TensorDom* init_dom = nullptr;
- TensorDom* update_dom = nullptr;
- if (out_dom_map->count(this->init[i])) {
- init_dom = &out_dom_map->at(this->init[i]);
- }
- if (out_dom_map->count(this->update[i])) {
- update_dom = &out_dom_map->at(this->update[i]);
- }
- // first dimension, always needed.
- if (init_dom) {
- init_dom->data[0].push_back(
- IntSet::FromRange(Range::FromMinExtent(0, this->init[i]->shape[0])));
- }
- if (update_dom) {
- update_dom->data[0].push_back(dom_map.at(this->scan_axis->var.get()));
- }
- // The update dimensions
- for (size_t k = 1; k < this->update[i]->shape.size(); ++k, ++sp_idx) {
- IterVar sp_ax = this->spatial_axis_[sp_idx];
- if (init_dom) {
- init_dom->data[k].push_back(dom_map.at(sp_ax->var.get()));
- }
- if (update_dom) {
- update_dom->data[k].push_back(dom_map.at(sp_ax->var.get()));
- }
- }
- }
-}
-
-void ScanOpNode::GatherBound(const Operation& self,
- const std::unordered_map<Tensor, TensorDom>& tensor_dom,
- std::unordered_map<IterVar, Range>* out_dom_map) const {
- ICHECK_EQ(self.operator->(), this);
- ICHECK(!out_dom_map->count(this->scan_axis));
- std::vector<Tensor> output(this->num_outputs());
- for (size_t i = 0; i < output.size(); ++i) {
- output[i] = self.output(i);
- }
- // Update for time axis.
- std::vector<IntSet> time_dom;
- for (size_t i = 0; i < output.size(); ++i) {
- const TensorDom& d = tensor_dom.at(output[i]);
- time_dom.insert(time_dom.end(), d.data[0].begin(), d.data[0].end());
- }
- ICHECK(!out_dom_map->count(this->scan_axis));
- arith::Analyzer analyzer;
- Range sdom = this->scan_axis->dom;
- Range r = arith::Union(time_dom).CoverRange(sdom);
- (*out_dom_map)[this->scan_axis] =
- Range::FromMinExtent(sdom->min, analyzer.Simplify(r->extent + r->min - sdom->min));
- Map<IterVar, PrimExpr> fix_pt = ScanFixPointAnalysis(self);
- // Update for spatial axis.
- size_t sp_idx = 0;
- for (size_t i = 0; i < output.size(); ++i) {
- const TensorDom& d = tensor_dom.at(output[i]);
- for (size_t k = 1; k < this->update[i]->shape.size(); ++k, ++sp_idx) {
- IterVar sp_ax = this->spatial_axis_[sp_idx];
- ICHECK(!out_dom_map->count(sp_ax));
- ICHECK(fix_pt.count(sp_ax));
- if (fix_pt[sp_ax].as<tir::IntImmNode>()->value) {
- // fix point, we can slice it.
- (*out_dom_map)[sp_ax] = arith::Union(d.data[k]).CoverRange(sp_ax->dom);
- } else {
- // not a fix point, need to include everything.
- (*out_dom_map)[sp_ax] = sp_ax->dom;
- }
- }
- }
-}
-
-Stmt ScanOpNode::BuildRealize(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- const Stmt& body, String storage_scope) const {
- arith::Analyzer analyzer;
- ICHECK_EQ(stage->op.get(), this);
- Range sdom = dom_map.at(this->scan_axis);
- Range tdom = Range::FromMinExtent(0, analyzer.Simplify(sdom->extent + sdom->min));
- Stmt ret = body;
- size_t sp_idx = 0;
- for (size_t i = 0; i < update.size(); ++i) {
- Tensor t = stage->op.output(i);
- ICHECK_EQ(static_cast<size_t>(t->value_index), i);
- Region bounds;
- bounds.push_back(tdom);
- for (size_t k = 1; k < this->update[i]->shape.size(); ++k, ++sp_idx) {
- IterVar sp_ax = this->spatial_axis_[sp_idx];
- bounds.push_back(dom_map.at(sp_ax));
- }
- ret = tir::ProducerRealize(t, bounds, const_true(), ret, storage_scope);
- }
- return ret;
-}
-
-Stmt ScanOpNode::BuildProvide(const Stage& stage, const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const {
- ICHECK_EQ(stage->op.operator->(), this);
- Stmt provide =
- AttrStmt(stage->op, tir::attr::scan_update_scope, this->scan_axis->var, Evaluate(0));
- Stmt init = AttrStmt(stage->op, tir::attr::scan_init_scope, 0, Evaluate(0));
- size_t begin_scan = 0;
- for (size_t i = 0; i < stage->leaf_iter_vars.size(); ++i) {
- if (stage->leaf_iter_vars[i]->iter_type == kThreadIndex) {
- ICHECK_EQ(begin_scan, i);
- begin_scan = i + 1;
- }
- }
- std::unordered_map<IterVar, PrimExpr> vmap;
- std::unordered_set<IterVar> empty;
- auto nest = MakeLoopNest(stage, dom_map, 0, false, empty, &vmap, debug_keep_trivial_loop);
- nest[begin_scan].push_back(init);
- nest.push_back(MakeIfNest(MakeBoundCheck(stage, dom_map, vmap, false, empty)));
- return MergeNest(nest, provide);
-}
} // namespace te
} // namespace tvm
diff --git a/src/te/operation/tensor_compute_op.cc b/src/te/operation/tensor_compute_op.cc
deleted file mode 100644
index 00f751c58a..0000000000
--- a/src/te/operation/tensor_compute_op.cc
+++ /dev/null
@@ -1,259 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Tensor Compute Op.
- * \file tensor_compute_op.cc
- */
-#include <tvm/arith/analyzer.h>
-#include <tvm/runtime/registry.h>
-#include <tvm/te/operation.h>
-#include <tvm/tir/builtin.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <unordered_set>
-
-#include "./compute_op.h"
-#include "./op_utils.h"
-
-namespace tvm {
-namespace te {
-using namespace tir;
-// TensorComputeOpNode
-TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
- .set_dispatch<TensorComputeOpNode>([](const ObjectRef& node, ReprPrinter* p) {
- auto* op = static_cast<const TensorComputeOpNode*>(node.get());
- p->stream << "tensor_compute_op(" << op->name << ", " << op << ")";
- });
-
-TVM_REGISTER_NODE_TYPE(TensorComputeOpNode);
-
-int TensorComputeOpNode::num_outputs() const {
- return static_cast<int>(this->intrin->buffers.size() - this->inputs.size());
-}
-
-DataType TensorComputeOpNode::output_dtype(size_t i) const {
- return this->intrin->buffers[this->inputs.size() + i]->dtype;
-}
-
-TensorComputeOp::TensorComputeOp(std::string name, std::string tag, Array<IterVar> axis,
- Array<IterVar> reduce_axis, int schedulable_ndim,
- TensorIntrin intrin, Array<Tensor> tensors, Array<Region> regions,
- Array<PrimExpr> scalar_inputs) {
- auto n = make_object<TensorComputeOpNode>();
- n->name = std::move(name);
- n->tag = std::move(tag);
- n->axis = std::move(axis);
- n->reduce_axis = std::move(reduce_axis);
- n->schedulable_ndim = std::move(schedulable_ndim);
- n->intrin = std::move(intrin);
- n->inputs = std::move(tensors);
- n->input_regions = std::move(regions);
- n->scalar_inputs = std::move(scalar_inputs);
- data_ = std::move(n);
-}
-
-TVM_REGISTER_GLOBAL("te.TensorComputeOp")
- .set_body_typed([](std::string name, std::string tag, Array<IterVar> axis,
- Array<IterVar> reduce_axis, int schedulable_ndim, TensorIntrin intrin,
- Array<Tensor> tensors, Array<Region> regions,
- Array<PrimExpr> scalar_inputs) {
- return TensorComputeOp(name, tag, axis, reduce_axis, schedulable_ndim, intrin, tensors,
- regions, scalar_inputs);
- });
-
-Array<Tensor> TensorComputeOpNode::InputTensors() const { return inputs; }
-
-Operation TensorComputeOpNode::ReplaceInputs(const Operation& self,
- const std::unordered_map<Tensor, Tensor>& rmap) const {
- ICHECK_EQ(self.operator->(), this);
- auto n = make_object<TensorComputeOpNode>(*this);
- auto intrin = make_object<TensorIntrinNode>(*(this->intrin.operator->()));
- intrin->body = ReplaceTensor(this->intrin->body, rmap);
- if (intrin->reduce_init.defined()) {
- intrin->reduce_init = ReplaceTensor(this->intrin->reduce_init, rmap);
- }
- if (intrin->reduce_update.defined()) {
- intrin->reduce_update = ReplaceTensor(this->intrin->reduce_update, rmap);
- }
- for (size_t i = 0; i < n->inputs.size(); ++i) {
- Tensor t = n->inputs[i];
- if (rmap.count(t)) {
- n->inputs.Set(i, rmap.at(t));
- }
- }
-
- if (intrin->body.same_as(n->intrin->body) &&
- intrin->reduce_init.same_as(n->intrin->reduce_init) &&
- intrin->reduce_update.same_as(n->intrin->reduce_update) && inputs.same_as(n->inputs)) {
- return self;
- } else {
- n->intrin = TensorIntrin(intrin);
- return Operation(n);
- }
-}
-
-void TensorComputeOpNode::PropBoundToInputs(
- const Operation& self, arith::Analyzer* analyzer,
- const std::unordered_map<const VarNode*, IntSet>& dom_map,
- std::unordered_map<Tensor, TensorDom>* out_dom_map) const {
- for (size_t i = 0; i < this->inputs.size(); ++i) {
- Tensor t = this->inputs[i];
- Region region = input_regions[i];
-
- auto it = out_dom_map->find(t);
- if (it == out_dom_map->end()) continue;
- TensorDom& dom = it->second;
- for (size_t j = 0; j < t.ndim(); ++j) {
- dom.data[j].emplace_back(EvalSet(region[j], dom_map));
- }
- }
-}
-
-size_t TensorComputeOpNode::num_schedulable_dims() const { return schedulable_ndim; }
-
-Stmt TensorComputeOpNode::BuildProvide(const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) const {
- ICHECK_EQ(stage->op.operator->(), this);
-
- // Start bind data.
- Stmt nop = Evaluate(0);
- std::vector<Stmt> input_bind_nest, output_bind_nest;
- Array<Tensor> inputs = this->InputTensors();
-
- // input binding
- size_t num_inputs = inputs.size();
- for (size_t i = 0; i < num_inputs; ++i) {
- Tensor tensor = inputs[i];
- Region region = this->input_regions[i];
- Buffer buffer = this->intrin->buffers[i];
- Array<ObjectRef> bind_spec{buffer, tensor};
-
- Array<PrimExpr> tuple;
- for (size_t i = 0; i < region.size(); ++i) {
- tuple.push_back(region[i]->min);
- tuple.push_back(region[i]->extent);
- }
- input_bind_nest.emplace_back(
- AttrStmt(bind_spec, tir::attr::buffer_bind_scope,
- Call(DataType::Handle(), tir::builtin::tvm_tuple(), tuple), nop));
- }
-
- // output binding
- for (int i = 0; i < this->num_outputs(); ++i) {
- Tensor tensor = stage->op.output(i);
- Buffer buffer = this->intrin->buffers[num_inputs + i];
- Array<ObjectRef> bind_spec{buffer, tensor};
-
- Array<PrimExpr> tuple;
- for (size_t i = 0; i < this->axis.size(); ++i) {
- auto ivar = this->axis[i];
- if (i < static_cast<size_t>(this->schedulable_ndim)) {
- tuple.push_back(ivar->var);
- tuple.push_back(1);
- } else {
- Range dom = ivar->dom;
- tuple.push_back(dom->min);
- tuple.push_back(dom->extent);
- }
- }
-
- output_bind_nest.emplace_back(
- AttrStmt(bind_spec, tir::attr::buffer_bind_scope,
- Call(DataType::Handle(), tir::builtin::tvm_tuple(), tuple), nop));
- }
-
- // Check variable remap
- std::unordered_map<const VarNode*, PrimExpr> vmap;
- tir::ArgBinder binder(&vmap);
-
- // Map the expressions passed in the call to the TensorIntrin, to the placeholder
- // variables
- Array<PrimExpr> user_expr = this->scalar_inputs;
- Array<Var> scalar_params = this->intrin->scalar_params;
- Array<PrimExpr> sp_expr;
- for (auto sp : scalar_params) {
- PrimExpr esp = sp;
- sp_expr.push_back(esp);
- }
- ICHECK_EQ(sp_expr.size(), user_expr.size());
- // TODO(jdavies-huawei): what name should be used here?
- binder.BindArray(sp_expr, user_expr, this->name);
-
- size_t tloc = stage->leaf_iter_vars.size();
- ComputeLoopNest n = ComputeLoopNest::Create(this, stage, dom_map, debug_keep_trivial_loop);
-
- if (this->reduce_axis.size() == 0) {
- std::vector<std::vector<Stmt>> nest(n.main_nest.begin(), n.main_nest.begin() + tloc + 1);
- nest.emplace_back(MakeIfNest(n.main_predicates));
- ICHECK_EQ(n.init_predicates.size(), 0U);
- ICHECK(this->intrin->body.defined())
- << "Normal store op for intrin " << this << " is not defined";
- Stmt body = MergeNest(output_bind_nest, this->intrin->body);
- body = MergeNest(input_bind_nest, body);
- body = tir::Substitute(body, vmap);
- body = MergeNest(binder.asserts(), body);
- body = te::Substitute(body, n.main_vmap);
- Stmt ret = MergeNest(nest, body);
- return ret;
- } else {
- // Need to split reduction
- ICHECK(this->intrin->reduce_update.defined()) << "Reduction update op is not defined";
- // Need init and update steps
- ICHECK_NE(this->reduce_axis.size(), 0U);
- std::vector<std::vector<Stmt>> common(n.main_nest.begin(),
- n.main_nest.begin() + n.num_common_loop + 1);
- std::vector<std::vector<Stmt>> update_nest(n.main_nest.begin() + n.num_common_loop + 1,
- n.main_nest.begin() + tloc + 1);
- update_nest.emplace_back(MakeIfNest(n.main_predicates));
-
- if (this->intrin->reduce_init.defined()) {
- // init nest
- std::vector<std::vector<Stmt>> init_nest(n.init_nest.begin(), n.init_nest.begin() + tloc + 1);
- init_nest.emplace_back(MakeIfNest(n.init_predicates));
- Stmt init = MergeNest(output_bind_nest, this->intrin->reduce_init);
- init = te::Substitute(init, n.init_vmap);
- init = MergeNest(init_nest, init);
- // The update
- Stmt update = MergeNest(output_bind_nest, this->intrin->reduce_update);
- update = MergeNest(input_bind_nest, update);
- update = tir::Substitute(update, vmap);
- update = MergeNest(binder.asserts(), update);
- update = te::Substitute(update, n.main_vmap);
- update = MergeNest(update_nest, update);
- return MergeNest(common, SeqStmt::Flatten(init, update));
- } else {
- // When init op is not available, use body op for reset in the first iter.
- ICHECK(this->intrin->body.defined()) << "Normal body op is not defined";
- Stmt update =
- TransformUpdate(stage, dom_map, n, this->intrin->body, this->intrin->reduce_update);
- update = MergeNest(output_bind_nest, update);
- update = MergeNest(input_bind_nest, update);
- update = tir::Substitute(update, vmap);
- update = MergeNest(binder.asserts(), update);
- update = te::Substitute(update, n.main_vmap);
- update = MergeNest(update_nest, update);
- return MergeNest(common, update);
- }
- }
-}
-} // namespace te
-} // namespace tvm
diff --git a/src/te/operation/tensorize.cc b/src/te/operation/tensorize.cc
deleted file mode 100644
index 138aeeb37f..0000000000
--- a/src/te/operation/tensorize.cc
+++ /dev/null
@@ -1,496 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \brief Logics related to tensorize, used by ComputeOpNode.
- * \file tensorize.cc
- */
-#include <tvm/runtime/registry.h>
-#include <tvm/tir/analysis.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include "../schedule/message_passing.h"
-#include "compute_op.h"
-#include "op_utils.h"
-
-namespace tvm {
-namespace te {
-
-using namespace tir;
-
-// Detect the region of input and output to be tensorized.
-// out_dom: the domain of root iter vars in output op
-// in_region: region of each input tensor.
-// return The location of the tensorized scope start.
-size_t InferTensorizeRegion(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- std::unordered_map<IterVar, Range>* out_dom,
- std::unordered_map<Tensor, Array<Range>>* in_region) {
- // Get the bound of the tensorized scope.
- bool found_point = false;
- size_t loc_scope = 0;
- std::unordered_map<IterVar, IntSet> up_state;
- // Loop over the leafs
- for (size_t i = stage->leaf_iter_vars.size(); i != 0; --i) {
- IterVar iv = stage->leaf_iter_vars[i - 1];
- ICHECK(iv->iter_type == kDataPar || iv->iter_type == kCommReduce);
- auto vit = dom_map.find(iv);
- ICHECK(vit != dom_map.end());
- const Range& vrange = vit->second;
- if (is_one(vrange->extent)) {
- up_state[iv] = IntSet::SinglePoint(vrange->min);
- } else if (found_point) {
- ICHECK(is_zero(vrange->min));
- up_state[iv] = IntSet::SinglePoint(iv->var);
- } else {
- up_state[iv] = IntSet::FromRange(vrange);
- }
- auto iit = stage->iter_var_attrs.find(iv);
- if (iit != stage->iter_var_attrs.end()) {
- const IterVarAttr& attr = (*iit).second;
- if (!found_point) {
- ICHECK(!attr->bind_thread.defined()) << "Do not allow thread in tensorize scope";
- }
- if (attr->iter_type == kTensorized) {
- ICHECK(!found_point) << "Do not allow two tensorized point";
- found_point = true;
- loc_scope = i - 1;
- }
- }
- }
- ICHECK(found_point);
- // Get domain of the tensorized scope.
- te::PassUpDomain(stage, dom_map, &up_state);
- // Get domains if inputs
- std::unordered_map<Tensor, TensorDom> in_dom;
- std::unordered_map<const VarNode*, IntSet> temp_dmap;
- arith::Analyzer analyzer;
- Array<Tensor> inputs = self->InputTensors();
- for (Tensor t : inputs) {
- in_dom.emplace(t, TensorDom(t.ndim()));
- }
- for (IterVar iv : self->root_iter_vars()) {
- IntSet iset = up_state.at(iv);
- Range iv_range = iset.CoverRange(dom_map.at(iv));
- (*out_dom)[iv] = iv_range;
- analyzer.Bind(iv->var, iv_range);
- temp_dmap[iv->var.get()] = iset;
- }
- // Input domains
- self->PropBoundToInputs(stage->op, &analyzer, temp_dmap, &in_dom);
- Range none;
- for (const auto& kv : in_dom) {
- Array<Range> vec;
- const Tensor& t = kv.first;
- for (size_t i = 0; i < t.ndim(); ++i) {
- Range r = arith::Union(kv.second.data.at(i)).CoverRange(none);
- ICHECK(r.defined()) << "cannot deduce region of tensorized scope for input " << t;
- vec.push_back(std::move(r));
- }
- (*in_region)[t] = std::move(vec);
- }
- return loc_scope;
-}
-
-void VerifyTensorizeLoopNest(const ComputeOpNode* self, const Stage& stage,
- const ComputeLoopNest& n, size_t tloc) {
- // Veirfication step.
- std::unordered_set<const VarNode*> banned;
- ICHECK_EQ(n.main_nest.size(), stage->leaf_iter_vars.size() + 1);
- ICHECK(n.init_nest.size() == stage->leaf_iter_vars.size() + 1 || n.init_nest.size() == 0);
- auto f_push_banned = [&banned](const Stmt& s) {
- if (const ForNode* op = s.as<ForNode>()) {
- banned.insert(op->loop_var.get());
- } else if (const AttrStmtNode* op = s.as<AttrStmtNode>()) {
- if (const IterVarNode* iv = op->node.as<IterVarNode>()) {
- banned.insert(iv->var.get());
- }
- } else if (const LetStmtNode* op = s.as<LetStmtNode>()) {
- banned.insert(op->var.get());
- }
- };
- for (size_t i = tloc; i < stage->leaf_iter_vars.size(); ++i) {
- for (const Stmt& s : n.main_nest[i + 1]) {
- f_push_banned(s);
- }
- if (n.init_nest.size() != 0) {
- for (const Stmt& s : n.init_nest[i + 1]) {
- f_push_banned(s);
- }
- }
- }
-
- auto fbanned = [&](const VarNode* node) { return banned.count(node); };
-
- for (const PrimExpr& pred : n.main_predicates) {
- if (tir::UsesVar(pred, fbanned)) {
- LOG(FATAL) << "Tensorize failed, split condition " << pred
- << " relies on var defined inside tensorize scope";
- }
- }
- for (const PrimExpr& pred : n.init_predicates) {
- if (tir::UsesVar(pred, fbanned)) {
- LOG(FATAL) << "Tensorize failed, split condition " << pred
- << " relies on var defined inside tensorize scope";
- }
- }
-}
-
-// Remap the tensor placeholder, index and inline things.
-class TensorIntrinMatcher final : public StmtExprMutator {
- public:
- PrimExpr VisitExpr_(const ProducerLoadNode* op) final {
- PrimExpr expr = StmtExprMutator::VisitExpr_(op);
- op = expr.as<ProducerLoadNode>();
- auto t = Downcast<Tensor>(op->producer);
- auto it = in_remap_.find(t);
- if (it != in_remap_.end()) {
- const InputEntry& e = it->second;
- ICHECK_EQ(op->indices.size(), e.region.size());
- Array<PrimExpr> indices;
- for (size_t i = e.start; i < e.region.size(); ++i) {
- indices.push_back(op->indices[i] - e.region[i]->min);
- }
- return ProducerLoad(e.tensor, indices);
- }
- return expr;
- }
-
- PrimExpr VisitExpr_(const VarNode* op) final {
- auto it = var_remap_.find(op);
- if (it != var_remap_.end()) {
- return it->second;
- } else {
- return GetRef<PrimExpr>(op);
- }
- }
-
- PrimExpr VisitExpr_(const ReduceNode* op) final {
- PrimExpr expr = StmtExprMutator::VisitExpr_(op);
- op = expr.as<ReduceNode>();
- Array<IterVar> axis;
- for (size_t i = 0; i < op->axis.size(); ++i) {
- auto it = axis_remap_.find(op->axis[i]);
- if (it != axis_remap_.end()) {
- axis.push_back(it->second);
- }
- }
- return Reduce(op->combiner, op->source, axis, op->condition, op->value_index, op->init);
- }
-
- void Init(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- const std::unordered_map<IterVar, Range>& out_dom,
- const std::unordered_map<Tensor, Array<Range>>& in_region, const TensorIntrin& intrin,
- Map<Var, Range>* compute_intrin_iter_space) {
- ICHECK(self == stage->op.get());
-
- for (size_t i = 0; i < stage->leaf_iter_vars.size(); ++i) {
- IterVar iv = stage->leaf_iter_vars[i];
- auto vit = dom_map.find(iv);
- if (vit != dom_map.end()) {
- const Range vrange = vit->second;
- compute_intrin_iter_space->Set(iv->var, vrange);
- }
- }
- analyzer_.Bind(*compute_intrin_iter_space);
-
- // input remap.
- Array<Tensor> inputs = self->InputTensors();
- ICHECK_EQ(inputs.size(), intrin->inputs.size());
- for (size_t i = 0; i < inputs.size(); ++i) {
- InputEntry e;
- e.tensor = intrin->inputs[i];
- e.region = Array<Range>(in_region.at(inputs[i]));
- ICHECK_GE(e.region.size(), e.tensor.ndim());
- // Enable fuzzy matching, to match [1, n, m] to [n, m]
- e.start = e.region.size() - e.tensor.ndim();
- for (size_t j = 0; j < e.start; ++j) {
- auto canonical_extent = analyzer_.Simplify(e.region[j]->extent);
- ICHECK(is_one(canonical_extent))
- << "Tensorize " << intrin->name << ":"
- << " Input dimension mismatch with tensor intrin "
- << " expected shape=" << e.tensor->shape << ", given region=" << e.region;
- }
- in_remap_[inputs[i]] = e;
- }
- // output remap
- const ComputeOpNode* intrin_compute = intrin->op.as<ComputeOpNode>();
- ICHECK(intrin_compute) << "Only support compute intrinsic for now";
- ICHECK_GE(self->axis.size(), intrin_compute->axis.size())
- << "Tensorize: Output mismatch with tensor intrin ";
- // Enable fuzzy matching, to match [1, n, m] to [n, m]
- size_t axis_start = self->axis.size() - intrin_compute->axis.size();
- for (size_t i = 0; i < axis_start; ++i) {
- Range r = out_dom.at(self->axis[i]);
- ICHECK(is_one(r->extent)) << "Tensorize: Output mismatch with tensor intrin "
- << " intrin-dim=" << intrin_compute->axis.size()
- << ", tensorize-dim=" << self->axis.size();
- var_remap_[self->axis[i]->var.get()] = r->min;
- }
- // Assume we tensorize at regin axis i [min, min + extent)
- // The corresponding intrinsic axis is j [0, extent)
- // Remap index i to j + min
- for (size_t i = axis_start; i < self->axis.size(); ++i) {
- IterVar iv = self->axis[i];
- IterVar target_iv = intrin_compute->axis[i - axis_start];
- Range r = out_dom.at(iv);
- var_remap_[iv->var.get()] = target_iv->var + r->min;
- axis_remap_[iv] = target_iv;
- compute_intrin_iter_space->Set(target_iv->var, target_iv->dom);
- }
- // Remap reduction axis
- ICHECK_GE(self->reduce_axis.size(), intrin_compute->reduce_axis.size())
- << "Tensorize: Reduction dimension mismatch with tensor intrin";
- axis_start = self->reduce_axis.size() - intrin_compute->reduce_axis.size();
- for (size_t i = 0; i < axis_start; ++i) {
- Range r = out_dom.at(self->reduce_axis[i]);
- ICHECK(is_one(r->extent)) << "Tensorize: Reduction mismatch with tensor intrin "
- << " intrin-dim=" << intrin_compute->reduce_axis.size()
- << ", tensorize-dim=" << self->reduce_axis.size();
- var_remap_[self->reduce_axis[i]->var.get()] = r->min;
- }
- for (size_t i = axis_start; i < self->reduce_axis.size(); ++i) {
- IterVar iv = self->reduce_axis[i];
- IterVar target_iv = intrin_compute->reduce_axis[i - axis_start];
- Range r = out_dom.at(iv);
- var_remap_[iv->var.get()] = target_iv->var + r->min;
- axis_remap_[iv] = target_iv;
- compute_intrin_iter_space->Set(target_iv->var, target_iv->dom);
- }
- }
-
- private:
- // Input entry
- struct InputEntry {
- Tensor tensor;
- size_t start;
- Array<Range> region;
- };
- // input data remap
- std::unordered_map<Tensor, InputEntry> in_remap_;
- // variable remap.
- std::unordered_map<const VarNode*, PrimExpr> var_remap_;
- // IterVar remap.
- std::unordered_map<IterVar, IterVar> axis_remap_;
- // arith analyzer
- arith::Analyzer analyzer_;
-};
-
-// Try to match tensor dataflow of the stage with the intrinsic
-Array<PrimExpr> MatchTensorizeBody(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- const std::unordered_map<IterVar, Range>& out_dom,
- const std::unordered_map<Tensor, Array<Range>>& in_region,
- const TensorIntrin& intrin,
- Map<Var, Range>* compute_intrin_iter_space) {
- TensorIntrinMatcher matcher;
- matcher.Init(self, stage, dom_map, out_dom, in_region, intrin, compute_intrin_iter_space);
- Array<PrimExpr> ret;
- for (PrimExpr expr : self->body) {
- ret.push_back(matcher(expr));
- }
- return ret;
-}
-
-void VerifyTensorizeBody(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, PrimExpr>& value_map,
- const std::unordered_map<IterVar, Range>& dom_map,
- const std::unordered_map<IterVar, Range>& out_dom,
- const std::unordered_map<Tensor, Array<Range>>& in_region,
- const TensorIntrin& intrin) {
- StructuralEqual expr_equal;
- Map<Var, Range> compute_intrin_iter_space;
- Array<PrimExpr> body = MatchTensorizeBody(self, stage, dom_map, out_dom, in_region, intrin,
- &compute_intrin_iter_space);
- const ComputeOpNode* intrin_compute = intrin->op.as<ComputeOpNode>();
- ICHECK(intrin_compute) << "Only support compute intrinsic for now";
- ICHECK_EQ(body.size(), intrin_compute->body.size()) << "Tensorize failed: body size mismatch";
- arith::Analyzer ana;
- ana.Bind(compute_intrin_iter_space);
-
- for (size_t i = 0; i < body.size(); ++i) {
- PrimExpr lhs = ana.Simplify(Substitute(body[i], value_map));
- // run substitution because the intrin body could depend on outer loop vars.
- PrimExpr rhs = ana.Simplify(Substitute(intrin_compute->body[i], value_map));
- if (lhs.dtype() != rhs.dtype()) {
- LOG(FATAL) << "Failed to match the data type with TensorIntrin " << intrin->name
- << "'s declaration "
- << " provided=" << lhs.dtype() << ", intrin=" << rhs.dtype();
- }
- ICHECK(expr_equal(lhs, rhs)) << "Failed to match the compute with TensorIntrin " << intrin->name
- << "'s declaration "
- << " provided= " << lhs << ", intrin= " << rhs
- << ", running this stage: " << stage;
- }
-}
-
-Stmt MakeTensorize(const ComputeOpNode* self, const Stage& stage,
- const std::unordered_map<IterVar, Range>& dom_map,
- bool debug_keep_trivial_loop) {
- std::unordered_map<IterVar, Range> out_dom;
- std::unordered_map<Tensor, Array<Range>> in_region;
- size_t tloc = InferTensorizeRegion(self, stage, dom_map, &out_dom, &in_region);
- TensorIntrin intrin = stage->iter_var_attrs.at(stage->leaf_iter_vars[tloc])->tensor_intrin;
- ICHECK(intrin.defined());
- ComputeLoopNest n = ComputeLoopNest::Create(self, stage, dom_map, debug_keep_trivial_loop);
- VerifyTensorizeLoopNest(self, stage, n, tloc);
- VerifyTensorizeBody(self, stage, n.main_vmap, dom_map, out_dom, in_region, intrin);
- // Start bind data.
- Stmt nop = Evaluate(0);
- std::vector<Stmt> input_bind_nest, output_bind_nest;
- Array<Tensor> inputs = self->InputTensors();
- ICHECK_EQ(inputs.size(), intrin->inputs.size()) << "Tensorize failed: input size mismatch ";
- // input binding
- for (size_t i = 0; i < intrin->inputs.size(); ++i) {
- Tensor tensor = inputs[i];
- Buffer buffer = intrin->buffers[i];
- Array<ObjectRef> bind_spec{buffer, tensor};
- auto it = in_region.find(tensor);
- ICHECK(it != in_region.end());
- const Array<Range>& region = it->second;
- Array<PrimExpr> tuple;
- for (const Range r : region) {
- tuple.push_back(r->min);
- tuple.push_back(r->extent);
- }
- input_bind_nest.emplace_back(
- AttrStmt(bind_spec, tir::attr::buffer_bind_scope,
- Call(DataType::Handle(), tir::builtin::tvm_tuple(), tuple), nop));
- }
- // output binding
- const ComputeOpNode* intrin_compute = intrin->op.as<ComputeOpNode>();
- ICHECK(intrin_compute) << "Only support compute intrinsic for now";
- ICHECK_EQ(intrin->inputs.size() + intrin_compute->body.size(), intrin->buffers.size());
- ICHECK_EQ(intrin_compute->body.size(), self->body.size());
- Array<PrimExpr> tuple;
- for (IterVar iv : self->axis) {
- auto it = out_dom.find(iv);
- ICHECK(it != out_dom.end());
- tuple.push_back(it->second->min);
- tuple.push_back(it->second->extent);
- }
- for (size_t i = intrin->inputs.size(); i < intrin->buffers.size(); ++i) {
- Tensor tensor = stage->op.output(i - intrin->inputs.size());
- Buffer buffer = intrin->buffers[i];
- Array<ObjectRef> bind_spec{buffer, tensor};
- output_bind_nest.emplace_back(
- AttrStmt(bind_spec, tir::attr::buffer_bind_scope,
- Call(DataType::Handle(), tir::builtin::tvm_tuple(), tuple), nop));
- }
- // Check variable remap
- std::unordered_map<const VarNode*, PrimExpr> vmap;
- tir::ArgBinder binder(&vmap);
- ICHECK_GE(self->reduce_axis.size(), intrin_compute->reduce_axis.size())
- << "Tensorization fail: reduction axis size do not match";
- size_t start = self->reduce_axis.size() - intrin_compute->reduce_axis.size();
- for (size_t i = 0; i < start; ++i) {
- IterVar iv = self->reduce_axis[i];
- auto it = out_dom.find(iv);
- ICHECK(it != out_dom.end());
- ICHECK(is_one(it->second->extent)) << "Tensorization fail: reduction axis size do not match";
- }
- for (size_t i = start; i < self->reduce_axis.size(); ++i) {
- IterVar iv = self->reduce_axis[i];
- IterVar target = intrin_compute->reduce_axis[i - start];
- auto it = out_dom.find(iv);
- ICHECK(it != out_dom.end());
- binder.Bind(target->dom->min, make_const(iv->dom->min.dtype(), 0),
- "tensir_intrin.reduction.min");
- binder.Bind(target->dom->extent, it->second->extent, "tensir_intrin.reduction.extent");
- }
- if (tloc <= n.num_common_loop) {
- // Do no need to split reduction
- std::vector<std::vector<Stmt>> nest(n.main_nest.begin(), n.main_nest.begin() + tloc + 1);
- nest.emplace_back(MakeIfNest(n.main_predicates));
- ICHECK_EQ(n.init_predicates.size(), 0U);
- ICHECK(intrin->body.defined()) << "Normal store op for intrin " << intrin << " is not defined";
- Stmt body = MergeNest(output_bind_nest, intrin->body);
- body = MergeNest(input_bind_nest, body);
- body = tir::Substitute(body, vmap);
- body = MergeNest(binder.asserts(), body);
- body = te::Substitute(body, n.main_vmap);
- return MergeNest(nest, body);
- } else {
- // Need to split reduction
- ICHECK(intrin->reduce_update.defined())
- << "Reduction update op for intrin " << intrin << " is not defined";
- // Need init and update steps
- ICHECK_NE(self->reduce_axis.size(), 0U);
- std::vector<std::vector<Stmt>> common(n.main_nest.begin(),
- n.main_nest.begin() + n.num_common_loop + 1);
- std::vector<std::vector<Stmt>> update_nest(n.main_nest.begin() + n.num_common_loop + 1,
- n.main_nest.begin() + tloc + 1);
- update_nest.emplace_back(MakeIfNest(n.main_predicates));
-
- if (intrin->reduce_init.defined()) {
- // init nest
- std::vector<std::vector<Stmt>> init_nest(n.init_nest.begin(), n.init_nest.begin() + tloc + 1);
- init_nest.emplace_back(MakeIfNest(n.init_predicates));
- Stmt init = MergeNest(output_bind_nest, intrin->reduce_init);
- init = te::Substitute(init, n.init_vmap);
- init = MergeNest(init_nest, init);
- // The update
- Stmt update = MergeNest(output_bind_nest, intrin->reduce_update);
- update = MergeNest(input_bind_nest, update);
- update = tir::Substitute(update, vmap);
- update = MergeNest(binder.asserts(), update);
- update = te::Substitute(update, n.main_vmap);
- update = MergeNest(update_nest, update);
- return MergeNest(common, SeqStmt::Flatten(init, update));
- } else {
- // When init op is not available, use body op for reset in the first iter.
- ICHECK(intrin->body.defined()) << "Normal body op for intrin " << intrin << " is not defined";
- Stmt update = TransformUpdate(stage, dom_map, n, intrin->body, intrin->reduce_update);
- update = MergeNest(output_bind_nest, update);
- update = MergeNest(input_bind_nest, update);
- update = tir::Substitute(update, vmap);
- update = MergeNest(binder.asserts(), update);
- update = te::Substitute(update, n.main_vmap);
- update = MergeNest(update_nest, update);
- return MergeNest(common, update);
- }
- }
-}
-
-// Register functions for unittests
-TVM_REGISTER_GLOBAL("test.op.InferTensorizeRegion").set_body([](TVMArgs args, TVMRetValue* ret) {
- Stage stage = args[0];
- Map<IterVar, Range> dmap = args[1];
- std::unordered_map<IterVar, Range> out_dom;
- std::unordered_map<Tensor, Array<Range>> in_region;
- ICHECK(stage->op.as<ComputeOpNode>());
- InferTensorizeRegion(stage->op.as<ComputeOpNode>(), stage, as_unordered_map(dmap), &out_dom,
- &in_region);
- *ret = Array<ObjectRef>{Map<IterVar, Range>(out_dom), Map<Tensor, Array<Range>>(in_region)};
-});
-
-TVM_REGISTER_GLOBAL("test.op.MatchTensorizeBody").set_body([](TVMArgs args, TVMRetValue* ret) {
- Stage stage = args[0];
- Map<IterVar, Range> out_dom = args[1];
- Map<Tensor, Array<Range>> in_region = args[2];
- TensorIntrin intrin = args[3];
- Map<Var, Range> vrange;
- ICHECK(stage->op.as<ComputeOpNode>());
- *ret = MatchTensorizeBody(stage->op.as<ComputeOpNode>(), stage, {{}}, as_unordered_map(out_dom),
- as_unordered_map(in_region), intrin, &vrange);
-});
-} // namespace te
-} // namespace tvm
diff --git a/src/te/schedule/auto_inline_elem_wise.cc b/src/te/schedule/auto_inline_elem_wise.cc
deleted file mode 100644
index bf584df258..0000000000
--- a/src/te/schedule/auto_inline_elem_wise.cc
+++ /dev/null
@@ -1,125 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file auto_inline_elem_wise.cc
- */
-#include <tvm/runtime/registry.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-#include <tvm/tir/expr_functor.h>
-
-namespace tvm {
-namespace te {
-
-using namespace tir;
-
-class ElemWiseDetector : public tir::ExprVisitor {
- public:
- explicit ElemWiseDetector(Array<IterVar> axis) : axis_(axis) {}
-
- void VisitExpr(const PrimExpr& e) final {
- if (!is_elem_wise_) return;
- ExprVisitor::VisitExpr(e);
- }
-
- void VisitExpr_(const ProducerLoadNode* op) final {
- Array<PrimExpr> indices = op->indices;
- if (axis_.size() != indices.size()) {
- is_elem_wise_ = false;
- return;
- }
-
- for (size_t i = 0; i < axis_.size(); ++i) {
- if (!indices[i].same_as(axis_[i]->var)) {
- is_elem_wise_ = false;
- return;
- }
- }
- ExprVisitor::VisitExpr_(op);
- }
-
- bool is_elem_wise_{true};
-
- private:
- Array<IterVar> axis_;
-};
-
-bool IsElemWise(const Operation& op) {
- if (const ComputeOpNode* compute = op.as<ComputeOpNode>()) {
- ElemWiseDetector v = ElemWiseDetector(compute->axis);
- for (auto& e : compute->body) v(e);
- return v.is_elem_wise_;
- }
- return false;
-}
-
-void AutoInlineElemWise(Schedule sch) {
- for (Stage s : sch->stages) {
- if (!s.is_scheduled() && IsElemWise(s->op) && !s->is_output) {
- s.compute_inline();
- }
- }
-}
-
-bool IsBroadcast(const Operation& op) {
- if (const ComputeOpNode* compute = op.as<ComputeOpNode>()) {
- if (compute->reduce_axis.size()) {
- return false;
- }
- constexpr auto kBroadcast = "broadcast";
- // broadcast op in topi has tag `broadcast`
- if (op->tag == kBroadcast) {
- return true;
- }
- }
- return false;
-}
-
-void AutoInlineBroadcast(Schedule sch) {
- for (Stage s : sch->stages) {
- if (!s.is_scheduled() && IsBroadcast(s->op) && !s->is_output) {
- s.compute_inline();
- }
- }
-}
-
-bool IsInjective(const Operation& op) {
- if (const ComputeOpNode* compute = op.as<ComputeOpNode>()) {
- return compute->reduce_axis.size() == 0;
- }
- return false;
-}
-
-void AutoInlineInjective(Schedule sch) {
- for (Stage s : sch->stages) {
- if (!s.is_scheduled() && IsInjective(s->op) && !s->is_output) {
- s.compute_inline();
- }
- }
-}
-
-TVM_REGISTER_GLOBAL("schedule.AutoInlineElemWise").set_body_typed(AutoInlineElemWise);
-
-TVM_REGISTER_GLOBAL("schedule.AutoInlineBroadcast").set_body_typed(AutoInlineBroadcast);
-
-TVM_REGISTER_GLOBAL("schedule.AutoInlineInjective").set_body_typed(AutoInlineInjective);
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/schedule/bound.cc b/src/te/schedule/bound.cc
deleted file mode 100644
index d8abffd6aa..0000000000
--- a/src/te/schedule/bound.cc
+++ /dev/null
@@ -1,262 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file bound.cc
- * \brief The bound inference logic.
- */
-#include <tvm/runtime/registry.h>
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule_pass.h>
-
-#include <unordered_map>
-#include <unordered_set>
-
-#include "../../runtime/thread_storage_scope.h"
-#include "graph.h"
-#include "message_passing.h"
-
-namespace tvm {
-namespace te {
-
-using runtime::StorageRank;
-using runtime::StorageScope;
-using runtime::ThreadScope;
-
-/*! \brief The graph context used during bound inference. */
-struct GraphContext {
- /*! \brief The feed graph */
- FeedGraph feed_graph;
- /*! \brief Attachment path */
- AttachPath attach_path;
- /*! \brief The bind map */
- std::unordered_map<IterVar, IterVar> bind_map;
- /*! \brief map from op to stage */
- std::unordered_map<const Object*, Stage> op2stage_;
-};
-
-bool NeedRelax(const IterVar& iv, bool found_attach,
- const std::unordered_map<IterVar, IterVar>& bind_map,
- const runtime::StorageScope& scope) {
- auto it = bind_map.find(iv);
- const std::string& tag = (it != bind_map.end() ? it->second->thread_tag : iv->thread_tag);
- if (tag.length() == 0 || tag == "pipeline") {
- return !found_attach;
- }
- ThreadScope ts = ThreadScope::Create(tag);
-
- // When there is warp memory
- // threadIdx.x must be set to be warp index.
- if (scope.rank == StorageRank::kWarp && ts.rank == 1 && ts.dim_index == 0) {
- return true;
- }
- return static_cast<int>(scope.rank) <= ts.rank;
-}
-
-// infer storage scope, if not given
-StorageScope InferStorageScope(const Stage& stage, const GraphContext& ctx) {
- if (stage->scope.length() != 0) {
- return StorageScope::Create(stage->scope);
- }
- int max_rank = -1;
- for (IterVar iv : ctx.attach_path.at(stage->op)) {
- auto it = ctx.bind_map.find(iv);
- const std::string& tag = (it != ctx.bind_map.end() ? it->second->thread_tag : iv->thread_tag);
- if (tag != "pipeline" && tag.length() != 0) {
- max_rank = std::max(max_rank, ThreadScope::Create(tag).rank);
- }
- }
- StorageScope s;
- s.rank = runtime::DefaultStorageRank(max_rank);
- return s;
-}
-
-void InferRootBound(const Stage& stage, const GraphContext& ctx,
- std::unordered_map<IterVar, Range>* rmap) {
- ICHECK_NE(stage->attach_type, kInline) << "call schedule.normalize before scheduleops";
- if (stage->attach_type == kInlinedAlready) return;
- if (stage->is_output) {
- // verify correctness.
- ICHECK_EQ(stage.GetAttachSpec()->attach_type, kGroupRoot) << "Output must be attached at root";
- }
- if (stage->is_output || stage->op.as<PlaceholderOpNode>()) {
- for (auto iv : stage->op->root_iter_vars()) {
- ICHECK(iv->dom.defined());
- ICHECK(!rmap->count(iv));
- (*rmap)[iv] = iv->dom;
- }
- return;
- }
- // The tensor domain.
- std::unordered_map<Tensor, TensorDom> tmap;
- // The consumers of the op.
- std::unordered_set<Operation> consumers;
- for (int i = 0; i < stage->op->num_outputs(); ++i) {
- Tensor t = stage->op.output(i);
- tmap.emplace(t, TensorDom(static_cast<int>(t.ndim())));
- auto it = ctx.feed_graph.find(t);
- if (it != ctx.feed_graph.end()) {
- for (const Operation& op : it->second) {
- consumers.insert(op);
- }
- } else {
- LOG(INFO) << "not in feed graph consumer = " << stage->op;
- }
- }
- // storage scope.
- runtime::StorageScope scope = InferStorageScope(stage, ctx);
- // Bound prop by other consumers.
- // - Compute bound by relaxation rules: NeedRelax
- // - For normal index, use relative location of loop nest./
- // - For thread index, use the thread scope.
- //
- Array<IterVar> stage_attach = ctx.attach_path.at(stage->op);
- // The parent set.
- for (const Operation& op : consumers) {
- Map<Var, IntSet> relax_set;
- std::unordered_map<IterVar, IntSet> up_state;
- bool found_attach = false;
- ICHECK(ctx.op2stage_.count(op.get()));
- const Stage& op_stage = ctx.op2stage_.at(op.get());
- // Consumer nest
- for (size_t i = op_stage->leaf_iter_vars.size(); i != 0; --i) {
- IterVar iv = op_stage->leaf_iter_vars[i - 1];
- if (stage_attach.size() != 0 && iv == stage_attach[0]) {
- found_attach = true;
- }
- auto it = rmap->find(iv);
- ICHECK(it != rmap->end());
- const Range& vrange = it->second;
- if (is_one(vrange->extent)) {
- up_state[iv] = IntSet::SinglePoint(vrange->min);
- } else if (!NeedRelax(iv, found_attach, ctx.bind_map, scope)) {
- ICHECK(is_zero(vrange->min)) << "InferBound requires every leaf iter var's min equals 0, "
- << " call schedule.normalize to achieve this. ";
- if (ctx.bind_map.count(iv)) {
- up_state[iv] = IntSet::SinglePoint(ctx.bind_map.at(iv)->var);
- } else {
- up_state[iv] = IntSet::SinglePoint(iv->var);
- }
- } else {
- up_state[iv] = IntSet::FromRange(vrange);
- }
- }
- // Consumer's attach nest
- for (IterVar iv : ctx.attach_path.at(op)) {
- if (stage_attach.size() != 0 && iv == stage_attach[0]) {
- found_attach = true;
- }
- Range vrange = rmap->at(iv);
- ICHECK(is_zero(vrange->min)) << "InferBound requires every leaf iter var's min equals 0, "
- << "call schedule.normalize to achieve this.";
- if (NeedRelax(iv, found_attach, ctx.bind_map, scope)) {
- relax_set.Set(iv->var, IntSet::FromRange(vrange));
- if (ctx.bind_map.count(iv)) {
- relax_set.Set(ctx.bind_map.at(iv)->var, IntSet::FromRange(vrange));
- }
- }
- }
- ICHECK(found_attach || stage_attach.size() == 0)
- << "Invalid Schedule, cannot find the producer " << stage->op
- << " along the loop nest specified by compute_at of consumer " << op;
- // Get the domain of the consumer
- PassUpDomain(op_stage, *rmap, &up_state);
- // Relax if needed.
- std::unordered_map<const VarNode*, IntSet> dom_map;
- arith::Analyzer analyzer;
- for (auto entry : *rmap) {
- analyzer.Bind(entry.first->var, entry.second);
- }
- for (auto iv : op->root_iter_vars()) {
- Range r;
- if (up_state.count(iv)) {
- r = up_state.at(iv).CoverRange(iv->dom);
- } else {
- r = iv->dom;
- }
- if (relax_set.size() != 0) {
- dom_map[iv->var.get()] =
- IntSet::Interval(analyzer.int_set(r->min, relax_set).min(),
- analyzer.int_set(r->min + r->extent - 1, relax_set).max());
- } else {
- dom_map[iv->var.get()] = IntSet::FromRange(r);
- }
- analyzer.Bind(iv->var, r, true);
- }
- op->PropBoundToInputs(op, &analyzer, dom_map, &tmap);
- }
- stage->op->GatherBound(stage->op, tmap, rmap);
-}
-
-Map<IterVar, Range> InferBound(const Schedule& sch) {
- // Prepare context
- GraphContext ctx;
- Array<Operation> roots;
- arith::Analyzer analyzer;
-
- for (Operation op : sch->outputs) {
- roots.push_back(sch->stage_map[op]->op);
- }
- ctx.feed_graph = CreateFeedGraph(CreateReadGraph(roots));
-
- for (Stage stage : sch->stages) {
- for (auto kv : stage->iter_var_attrs) {
- if (kv.second->bind_thread.defined()) {
- ICHECK(!ctx.bind_map.count(kv.first));
- ctx.bind_map[kv.first] = kv.second->bind_thread;
- }
- }
- ctx.op2stage_[stage->op.get()] = stage;
- }
- ctx.attach_path = CreateAttachPath(sch);
- // Run inference.
- std::unordered_map<IterVar, Range> ret;
- for (size_t i = sch->stages.size(); i != 0; --i) {
- const Stage& stage = sch->stages[i - 1];
- InferRootBound(stage, ctx, &ret);
-
- // bind bound of root iter vars.
- for (auto iv : stage->op->root_iter_vars()) {
- auto it = ret.find(iv);
- if (it != ret.end()) {
- analyzer.Bind(iv->var, it->second);
- }
- }
-
- // pass down to get bound of all iter vars.
- PassDownDomain(stage, &ret, &analyzer);
- for (IterVar iv : stage->env_threads) {
- ICHECK(iv->dom.defined());
- ret[iv] = iv->dom;
- }
- }
- for (auto it = ret.begin(); it != ret.end(); it++) {
- DataType var_type = it->first->var.dtype();
- it->second = Range::FromMinExtent(
- // The range associated with each itervar must have the same dtype as the var
- analyzer.Simplify(cast(var_type, it->second->min)),
- analyzer.Simplify(cast(var_type, it->second->extent)));
- }
- return Map<IterVar, Range>(ret.begin(), ret.end());
-}
-
-TVM_REGISTER_GLOBAL("schedule.InferBound").set_body_typed(InferBound);
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/schedule/graph.cc b/src/te/schedule/graph.cc
deleted file mode 100644
index 502753284d..0000000000
--- a/src/te/schedule/graph.cc
+++ /dev/null
@@ -1,418 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file graph.cc
- * \brief Utilities to get information about schedule graph.
- */
-#include "graph.h"
-
-#include <tvm/runtime/registry.h>
-#include <tvm/te/operation.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/stmt_functor.h>
-
-#include <unordered_map>
-#include <unordered_set>
-#include <utility>
-
-namespace tvm {
-namespace te {
-// key to specific tensor dimension.
-struct TensorDimKey {
- Operation op;
- int value_index;
- int dim;
- TensorDimKey() {}
- TensorDimKey(const Tensor& t, int dim) : op(t->op), value_index(t->value_index), dim(dim) {}
- TensorDimKey(const Tensor& t, size_t dim)
- : op(t->op), value_index(t->value_index), dim(static_cast<int>(dim)) {}
- inline bool operator==(const TensorDimKey& other) const {
- return op == other.op && value_index == other.value_index && dim == other.dim;
- }
- inline bool operator!=(const TensorDimKey& other) const { return !operator==(other); }
-};
-} // namespace te
-} // namespace tvm
-
-namespace std {
-template <>
-struct hash<::tvm::te::TensorDimKey> {
- std::size_t operator()(const ::tvm::te::TensorDimKey& k) const {
- size_t lhs = ::tvm::ObjectPtrHash()(k.op);
- size_t rhs = static_cast<size_t>(k.value_index) << 16UL | static_cast<size_t>(k.dim);
- lhs ^= rhs + 0x9e3779b9 + (lhs << 6) + (lhs >> 2);
- return lhs;
- }
-};
-} // namespace std
-
-namespace tvm {
-namespace te {
-
-// construct a read graph that gives readers of each operation
-// that the root depend on
-ReadGraph CreateReadGraph(const Array<Operation>& roots) {
- ReadGraph rmap;
- std::vector<Operation> stack;
- std::unordered_set<const Object*> visited;
- // initialize the roots
- for (Operation op : roots) {
- stack.push_back(op);
- visited.insert(op.get());
- }
-
- while (!stack.empty()) {
- Operation op = stack.back();
- stack.pop_back();
- Array<Tensor> deps = op->InputTensors();
- rmap.Set(op, deps);
- for (Tensor t : deps) {
- if (t->op.defined() && visited.count(t->op.get()) == 0) {
- visited.insert(t->op.get());
- stack.push_back(t->op);
- }
- }
- }
- return rmap;
-}
-
-// Do DFS visit to get the subgraph.
-// Return if op is inside the subgraph.
-bool GetSubGraphByPostDFS_(const Operation& op, const std::unordered_set<const Object*>& boundary,
- bool include_bounary, std::unordered_map<const Object*, bool>* visited,
- Array<Operation>* result) {
- if (visited->count(op.get())) {
- return visited->at(op.get());
- }
- if (boundary.count(op.get())) {
- (*visited)[op.get()] = true;
- if (include_bounary) {
- result->push_back(op);
- }
- return true;
- }
- // mark to avoid loop
- // Not necessary for DAG.
- (*visited)[op.get()] = false;
- // check if we can reach boundary.
- bool reach_boundary = false;
- for (Tensor t : op->InputTensors()) {
- if (GetSubGraphByPostDFS_(t->op, boundary, include_bounary, visited, result)) {
- reach_boundary = true;
- }
- }
- (*visited)[op.get()] = reach_boundary;
- if (reach_boundary) {
- result->push_back(op);
- }
- return reach_boundary;
-}
-
-Array<Operation> GetSubGraph(const Array<Tensor>& outputs, const Array<Tensor>& inputs,
- bool include_inputs) {
- Array<Operation> result;
- std::unordered_set<const Object*> boundary;
- for (Tensor t : inputs) {
- boundary.insert(t->op.get());
- }
- std::unordered_map<const Object*, bool> visited;
- for (Tensor t : outputs) {
- GetSubGraphByPostDFS_(t->op, boundary, include_inputs, &visited, &result);
- }
- return result;
-}
-
-void PostDFSOrder(const Operation& op, const ReadGraph& g, std::unordered_set<Operation>* visited,
- Array<Operation>* post_order) {
- if (visited->count(op)) return;
- visited->insert(op);
- for (const auto& t : g.at(op)) {
- PostDFSOrder(t->op, g, visited, post_order);
- }
- post_order->push_back(op);
-}
-
-Array<Operation> PostDFSOrder(const Array<Operation>& roots, const ReadGraph& g) {
- std::unordered_set<Operation> visited;
- Array<Operation> post_order;
- for (Operation op : roots) {
- PostDFSOrder(op, g, &visited, &post_order);
- }
- return post_order;
-}
-
-FeedGraph CreateFeedGraph(const ReadGraph& g) {
- FeedGraph fg;
- for (auto kv : g) {
- for (Tensor t : kv.second) {
- fg[t].push_back(kv.first);
- }
- }
- return fg;
-}
-
-AttachPath CreateAttachPath(Schedule sch) {
- AttachPath ret;
- for (Stage stage : sch->stages) {
- std::unordered_set<const Object*> visited;
- Array<IterVar> path;
- for (Stage s = stage; s.defined();) {
- ICHECK(!visited.count(s.get())) << "Find loop in compute_at attach group";
- visited.insert(s.get());
- Stage spec = s.GetAttachSpec();
- bool start_attach;
- IterVar attach_ivar;
- if (spec->attach_type == kScope) {
- attach_ivar = spec->attach_ivar;
- s = spec->attach_stage;
- start_attach = false;
- ICHECK(attach_ivar.defined());
- } else if (spec->attach_type == kScanUpdate) {
- s = spec->attach_stage;
- start_attach = true;
- } else {
- break;
- }
- ICHECK(s.defined());
- for (size_t i = s->leaf_iter_vars.size(); i != 0; --i) {
- IterVar iv = s->leaf_iter_vars[i - 1];
- if (!start_attach && iv.same_as(attach_ivar)) {
- start_attach = true;
- }
- if (start_attach) path.push_back(iv);
- }
- ICHECK(start_attach) << "Invalid Schedule: cannot find attach point " << attach_ivar
- << " in the schedule of " << s->op;
- }
- if (!ret.count(stage->op)) {
- ret.Set(stage->op, path);
- }
- }
- return ret;
-}
-
-// graph of push reach relation of tensor dimensions
-using ReachGraph = std::unordered_map<TensorDimKey, std::vector<TensorDimKey>>;
-
-ReachGraph GetReachGraph(const Array<Operation>& ops) {
- ReachGraph reach;
- std::unordered_set<const Object*> bset;
- for (size_t i = 0; i < ops.size(); ++i) {
- bset.insert(ops[i].get());
- }
-
- for (Operation op : ops) {
- if (const auto* scan_op = op.as<ScanOpNode>()) {
- const auto& update = scan_op->update;
- const auto& init = scan_op->init;
- for (size_t i = 0; i < update.size(); ++i) {
- Tensor t = op.output(i);
- for (int k = 1; k < static_cast<int>(update[i]->shape.size()); ++k) {
- reach[TensorDimKey(t, k)].emplace_back(TensorDimKey(update[i], k));
- reach[TensorDimKey(t, k)].emplace_back(TensorDimKey(init[i], k));
- }
- }
- } else if (const auto* compute_op = op.as<ComputeOpNode>()) {
- std::unordered_map<const Object*, TensorDimKey> vmap;
- const auto& axis = compute_op->axis;
- Tensor t = op.output(0);
- for (size_t i = 0; i < axis.size(); ++i) {
- vmap[axis[i]->var.get()] = TensorDimKey(t, i);
- reach[TensorDimKey(t, i)] = {};
- }
- auto fvisit = [&vmap, &reach, &bset](const ObjectRef& n) {
- if (auto* pload = n.as<tir::ProducerLoadNode>()) {
- Tensor t = Downcast<Tensor>(pload->producer);
- if (!bset.count(t->op.get())) return;
- for (size_t i = 0; i < pload->indices.size(); ++i) {
- TensorDimKey dkey(t, static_cast<int>(i));
- auto fpush = [&dkey, &vmap, &reach](const ObjectRef& node) {
- const VarNode* v = node.as<VarNode>();
- auto it = vmap.find(v);
- if (it != vmap.end()) {
- reach[it->second].push_back(dkey);
- }
- };
- tir::PostOrderVisit(pload->indices[i], fpush);
- }
- }
- };
- for (auto& e : compute_op->body) {
- tir::PostOrderVisit(e, fvisit);
- }
- }
- }
- return reach;
-}
-
-Array<Operation> ScanGetBody(const Operation& scan_op) {
- const ScanOpNode* scan = scan_op.as<ScanOpNode>();
- // Get the body.
- Array<Tensor> inputs;
- for (Tensor t : scan->state_placeholder) {
- inputs.push_back(t);
- }
- for (Tensor t : scan->inputs) {
- inputs.push_back(t);
- }
- return GetSubGraph(scan->update, inputs, false);
-}
-
-Map<IterVar, PrimExpr> ScanFixPointAnalysis(const Operation& scan_op) {
- const ScanOpNode* scan = scan_op.as<ScanOpNode>();
- Array<Operation> body = ScanGetBody(scan_op);
-
- std::unordered_map<TensorDimKey, const Object*> exact_reach;
- std::unordered_set<const Object*> fail_set;
-
- for (size_t i = 0, sp_idx = 0; i < scan->update.size(); ++i) {
- for (size_t k = 1; k < scan->update[i]->shape.size(); ++k, ++sp_idx) {
- TensorDimKey key(scan->state_placeholder[i], k);
- exact_reach[key] = scan->spatial_axis_[sp_idx].get();
- }
- }
- // merge exact reach
- auto f_merge_key = [&exact_reach, &fail_set](const TensorDimKey& dst, const TensorDimKey& src) {
- auto sit = exact_reach.find(src);
- if (sit == exact_reach.end()) return;
- auto dit = exact_reach.find(dst);
- if (dit == exact_reach.end()) {
- exact_reach[dst] = sit->second;
- } else {
- if (dit->second != sit->second) {
- fail_set.insert(dit->second);
- fail_set.insert(sit->second);
- }
- }
- };
- // prop exact reach back.
- for (size_t i = 0; i < body.size(); ++i) {
- const Operation& op = body[i];
- if (const auto* scan_op = op.as<ScanOpNode>()) {
- const auto& update = scan_op->update;
- const auto& init = scan_op->init;
- for (size_t i = 0; i < update.size(); ++i) {
- Tensor t = op.output(i);
- for (size_t k = 1; k < update[i]->shape.size(); ++k) {
- f_merge_key(TensorDimKey(t, k), TensorDimKey(update[i], k));
- f_merge_key(TensorDimKey(t, k), TensorDimKey(init[i], k));
- }
- }
- } else if (const auto* compute_op = op.as<ComputeOpNode>()) {
- std::unordered_map<const Object*, std::vector<TensorDimKey>> vmap;
- const auto& axis = compute_op->axis;
- for (size_t i = 0; i < axis.size(); ++i) {
- std::vector<TensorDimKey> keys;
- for (int j = 0; j < op->num_outputs(); ++j) {
- keys.emplace_back(op.output(j), i);
- }
- vmap[axis[i]->var.get()] = std::move(keys);
- }
- auto fvisit = [&vmap, &f_merge_key, &exact_reach, &fail_set](const ObjectRef& n) {
- if (auto* pload = n.as<tir::ProducerLoadNode>()) {
- Tensor t = Downcast<Tensor>(pload->producer);
- for (size_t i = 0; i < pload->indices.size(); ++i) {
- auto it = vmap.find(pload->indices[i].get());
- TensorDimKey src(t, static_cast<int>(i));
- if (it != vmap.end()) {
- const std::vector<TensorDimKey>& keys = it->second;
- for (const auto& key : keys) {
- f_merge_key(key, src);
- }
- } else {
- if (exact_reach.count(src)) {
- fail_set.insert(exact_reach.at(src));
- }
- }
- }
- }
- };
- for (auto& e : compute_op->body) {
- tir::PostOrderVisit(e, fvisit);
- }
- }
- }
- ReachGraph reach;
- Map<IterVar, PrimExpr> ret;
- std::unordered_set<TensorDimKey> place_holder_ref;
- for (size_t i = 0; i < scan->state_placeholder.size(); ++i) {
- for (size_t k = 0; k < scan->state_placeholder[i]->shape.size(); ++k) {
- place_holder_ref.insert(TensorDimKey(scan->state_placeholder[i], k));
- }
- }
-
- for (size_t i = 0, sp_idx = 0; i < scan->update.size(); ++i) {
- for (size_t k = 1; k < scan->update[i]->shape.size(); ++k, ++sp_idx) {
- TensorDimKey key(scan->update[i], k);
- TensorDimKey target(scan->state_placeholder[i], k);
- IterVar sp_iv = scan->spatial_axis_[sp_idx];
- if (fail_set.count(sp_iv.get()) || !exact_reach.count(key) ||
- exact_reach.at(key) != sp_iv.get()) {
- ret.Set(sp_iv, make_const(DataType::Int(32), 0));
- } else {
- // now we proved exact match, need to prove no interference with other graph.
- if (reach.size() == 0) reach = GetReachGraph(body);
- // do a DFS
- std::unordered_set<TensorDimKey> visited;
- std::vector<TensorDimKey> stack{key};
- visited.insert(key);
- while (!stack.empty()) {
- TensorDimKey k = stack.back();
- if (k != target && place_holder_ref.count(k)) break;
- stack.pop_back();
- if (!reach.count(k)) {
- LOG(FATAL) << "cannot find reach of " << k.op << "-" << k.dim;
- }
-
- for (TensorDimKey kk : reach.at(k)) {
- if (visited.count(kk)) {
- continue;
- }
- visited.insert(kk);
- stack.push_back(kk);
- }
- }
- if (!stack.empty()) {
- // failed the prove.
- ret.Set(sp_iv, make_const(DataType::Int(32), 0));
- } else {
- ret.Set(sp_iv, make_const(DataType::Int(32), 1));
- }
- }
- }
- }
- return ret;
-}
-
-TVM_REGISTER_GLOBAL("schedule.CreateReadGraph").set_body_typed(CreateReadGraph);
-
-TVM_REGISTER_GLOBAL("schedule.PostDFSOrder")
- .set_body_typed([](const Array<Operation>& roots, const ReadGraph& g) {
- return PostDFSOrder(roots, g);
- });
-
-TVM_REGISTER_GLOBAL("schedule.CreateAttachPath").set_body_typed(CreateAttachPath);
-
-TVM_REGISTER_GLOBAL("schedule.ScanGetBody").set_body_typed(ScanGetBody);
-
-TVM_REGISTER_GLOBAL("schedule.ScanFixPointAnalysis").set_body_typed(ScanFixPointAnalysis);
-
-} // namespace te
-} // namespace tvm
diff --git a/src/te/schedule/graph.h b/src/te/schedule/graph.h
deleted file mode 100644
index d31473d1b5..0000000000
--- a/src/te/schedule/graph.h
+++ /dev/null
@@ -1,130 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file graph.h
- * \brief Utilities to get information about schedule graph.
- */
-#ifndef TVM_TE_SCHEDULE_GRAPH_H_
-#define TVM_TE_SCHEDULE_GRAPH_H_
-
-#include <tvm/te/operation.h>
-#include <tvm/te/schedule.h>
-#include <tvm/tir/expr.h>
-
-#include <unordered_map>
-#include <unordered_set>
-#include <vector>
-
-namespace tvm {
-namespace te {
-
-/*!
- * \brief data structure of Operation->Tensors it reads
- */
-using ReadGraph = Map<Operation, Array<Tensor>>;
-
-/*!
- * \brief AttachPath maps op-> a list of IterVar
- */
-using AttachPath = Map<Operation, Array<IterVar>>;
-
-/*!
- * \brief The map between tensor and operation it feeds to.
- */
-using FeedGraph = std::unordered_map<Tensor, std::vector<Operation>>;
-
-/*!
- * \brief Get read graph of each operation to all the
- * Tensors that it directly depends on.
- *
- * The result map contains Operations needed to finish root Operation.
- * \param roots The root operation.
- * \return The result map.
- */
-ReadGraph CreateReadGraph(const Array<Operation>& roots);
-
-/*!
- * \brief Get minimum subgraph between outputs and inputs.
- * The operations contains node which input-reachable from any inputs
- * output reachable to any outputs.
- *
- * The inputs won't be included in the subgraph, the outputs will be included.
- *
- * \param outputs The outputs of the subgraph
- * \param inputs The inputs to the subgraph.
- * \param include_inputs Whether to include inputs
- *
- * \return The subgraph.
- */
-Array<Operation> GetSubGraph(const Array<Tensor>& outputs, const Array<Tensor>& inputs,
- bool include_inputs);
-
-/*!
- * \brief Get a post DFS ordered of operations in the graph.
- * \param roots The root of the graph.
- * \param g The read graph.
- * \return vector order of Operations in PostDFS order.
- *
- * \note PostDFSOrder is a special case of Topoligical order,
- * and can be used when topoligical order is needed.
- */
-Array<Operation> PostDFSOrder(const Array<Operation>& roots, const ReadGraph& g);
-
-/*!
- * \brief Create feedgraph for given Schedule
- * \param g The read graph.
- * \return The created feedgraph.
- */
-FeedGraph CreateFeedGraph(const ReadGraph& g);
-
-/*!
- * \brief Create AttachPath that maps op-> a list of IterVar
- * That represents the loop nest op sits in from inner most to outermost
- * Also inserts attach_stage for scan updates when needed.
- *
- * \param sch The schedule.
- * \return The attach path.
- */
-AttachPath CreateAttachPath(Schedule sch);
-
-/*!
- * \brief Get all operations inside the recursion of scan.
- * \param scan_op The scan node ops.
- * \return The body operations, in read dependency order.
- */
-Array<Operation> ScanGetBody(const Operation& scan_op);
-
-/*!
- * \brief Analyze each spatial dimension of scan's result.
- * Give check on whether each dimension is fix point,
- * An axis is a fixed point if it only refers back to itself in recursion
- * and it is not used in axis of other recursion field.
- *
- * next_state[t, ..., axis, ...] = f(prev_state[t-1, ...,axis,...]
- *
- * \param scan The scan node.
- * \return Map of spatial_axis -> IntImm
- */
-Map<IterVar, PrimExpr> ScanFixPointAnalysis(const Operation& scan);
-
-} // namespace te
-} // namespace tvm
-
-#endif // TVM_TE_SCHEDULE_GRAPH_H_
diff --git a/src/te/schedule/message_passing.cc b/src/te/schedule/message_passing.cc
deleted file mode 100644
index e8f0d9332a..0000000000
--- a/src/te/schedule/message_passing.cc
+++ /dev/null
@@ -1,744 +0,0 @@
-/*
- * 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.
- */
-
-/*!
- * \file message_passing.cc
- * \brief The message passing domain.
- */
-#include "message_passing.h"
-
-#include <tvm/arith/analyzer.h>
-#include <tvm/tir/expr.h>
-
-namespace tvm {
-namespace te {
-
-using namespace tir;
-
-void Update(std::unordered_map<IterVar, Range>* p_state, const IterVar& iv, Range r,
- arith::Analyzer* analyzer) {
- auto it = p_state->find(iv);
- if (it == p_state->end()) {
- (*p_state)[iv] = r;
- analyzer->Bind(iv->var, r);
- } else {
- bool match =
- is_zero(it->second->min) && analyzer->CanProve(r->extent - it->second->extent == 0);
- ICHECK(match) << iv << " domain already inferred,"
- << " cannot prove their extents are the same " << it->second->extent << " vs "
- << r->extent;
- }
-}
-
-/*!
- * \param Upward propagating whether an IterVar derives at least one leaf IterVar that binds to
- * a thread.
- *
- * \param stage The stage to operate on.
- * \param p_state The propagation result of each IterVar.
- */
-void PassUpThreadBinding(const Stage& stage, std::unordered_map<IterVar, bool>* p_state) {
- auto bound_to_thread = [&stage](const IterVar& iv) {
- bool bound = false;
- auto it = stage->iter_var_attrs.find(iv);
- if (it != stage->iter_var_attrs.end()) {
- bound = (*it).second->bind_thread.defined();
- }
- return bound;
- };
-
- auto& state = *p_state;
- // Fill p_state with leaf itervars
- for (const IterVar& iv : stage->leaf_iter_vars) {
- state[iv] = bound_to_thread(iv);
- }
- // Traverse the graph bottom-up to propagate thread binding information
- for (size_t i = stage->relations.size(); i != 0; --i) {
- IterVarRelation rel = stage->relations[i - 1];
- if (const SplitNode* s = rel.as<SplitNode>()) {
- state[s->parent] = state[s->inner] || state[s->outer];
- } else if (const FuseNode* s = rel.as<FuseNode>()) {
- state[s->inner] = state[s->fused];
- state[s->outer] = state[s->fused];
- } else if (const RebaseNode* s = rel.as<RebaseNode>()) {
- state[s->parent] = state[s->rebased];
- } else if (rel.as<SingletonNode>()) {
- } else if (const TransformNode* s = rel.as<TransformNode>()) {
- // Currently, this marks all original iter vars as deriving from
- // a thread bind if any of the transformed variables are bound,
- // even if the inverse expression for that iter var doesn't
- // depend on the bound variable.
-
- // TODO(Lunderberg): For each of original variable, check
- // whether any variable in the inverse expression for it has a
- // thread binding.
- bool is_thread_binding = false;
- for (const auto& iter_var : s->transformed_variables) {
- is_thread_binding = is_thread_binding || state[iter_var];
- }
- for (const auto& iter_var : s->original_variables) {
- state[iter_var] = is_thread_binding;
- }
- } else {
- LOG(FATAL) << "unknown relation type";
- }
- }
-}
-
-void PassDownDomain(const Stage& stage, std::unordered_map<IterVar, Range>* p_state,
- arith::Analyzer* actx, bool allow_missing) {
- auto ceil_div = [actx](const PrimExpr& a, const PrimExpr& b) {
- if (actx->CanProve(indexmod(a, b) == 0)) {
- return actx->Simplify(indexdiv(a, b));
- }
- return actx->Simplify(indexdiv(a + (b - 1), b));
- };
-
- auto minimum_or_later = [actx](const PrimExpr& a, const PrimExpr& b) {
- if (actx->CanProve(a < b)) {
- return actx->Simplify(a);
- }
- return actx->Simplify(b);
- };
-
- std::unordered_map<IterVar, bool> dominating_thread;
- PassUpThreadBinding(stage, &dominating_thread);
-
- auto& state = *p_state;
... 4479 lines suppressed ...