[REFACTOR][ARITH] Remove legacy compute_expr.h (apache#5738)

Replaces most of the ComptuteReduce using foldl.
trevor-m · Jun 18, 2020 · 3c9d909 · 3c9d909
1 parent fc67c8a
commit 3c9d909
Show file tree

Hide file tree

Showing 23 changed files with 70 additions and 157 deletions.
diff --git a/include/tvm/tir/op.h b/include/tvm/tir/op.h
@@ -463,6 +463,7 @@ TVM_DLL PrimExpr isinf(PrimExpr x);
  * \brief sum of of source expression over axis
  * \param source The source expression.
  * \param axis List of iteration variables that will be used for reduction.
+ * \return The result.
  */
 TVM_DLL PrimExpr sum(PrimExpr source, Array<tir::IterVar> axis);
 
@@ -477,27 +478,31 @@ TVM_DLL PrimExpr all(PrimExpr source, Array<tir::IterVar> axis);
  * \brief logical Or of of source expression over axis
  * \param source The source expression.
  * \param axis List of iteration variables that will be used for reduction.
+ * \return The result.
  */
 TVM_DLL PrimExpr any(PrimExpr source, Array<tir::IterVar> axis);
 
 /*!
  * \brief max of of source expression over axis
  * \param source The source expression.
  * \param axis List of iteration variables that will be used for reduction.
+ * \return The result.
  */
 TVM_DLL PrimExpr max(PrimExpr source, Array<tir::IterVar> axis);
 
 /*!
  * \brief max of of source expression over axis
  * \param source The source expression.
  * \param axis List of iteration variables that will be used for reduction.
+ * \return The result.
  */
 TVM_DLL PrimExpr min(PrimExpr source, Array<tir::IterVar> axis);
 
 /*!
  * \brief product of of source expression over axis
  * \param source The source expression.
  * \param axis List of iteration variables that will be used for reduction.
+ * \return The result.
  */
 TVM_DLL PrimExpr prod(PrimExpr source, Array<tir::IterVar> axis);
 
@@ -658,6 +663,17 @@ inline bool is_zero(const PrimExpr& x) { return is_const_int(x, 0); }
  */
 inline bool is_const(const PrimExpr& x);
 
+/*!
+ * \brief Left fold.
+ * \param freduce The reduction function.
+ * \param init_value The initial value.
+ * \param values The values to be folded.
+ * \return The result.
+ * \tparam FReduce The type of the reduction.
+ */
+template <typename FReduce>
+inline PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array<PrimExpr>& values);
+
 /*!
  * \brief Check whether x is a constant power of two
  * If x is power of two, write the power to the shift.
@@ -762,6 +778,15 @@ inline PrimExpr make_zero(DataType t) {
   }
   return make_const(t, 0);
 }
+
+template <typename FReduce>
+inline PrimExpr foldl(FReduce freduce, PrimExpr init_value, const Array<PrimExpr>& values) {
+  for (PrimExpr val : values) {
+    init_value = freduce(init_value, val);
+  }
+  return init_value;
+}
+
 }  // namespace tir
 
 // additional const expression overloading

diff --git a/src/arith/compute_expr.h b/src/arith/compute_expr.h
diff --git a/src/relay/op/tensor/transform.cc b/src/relay/op/tensor/transform.cc
@@ -38,7 +38,6 @@
 
 #include <vector>
 
-#include "../../../arith/compute_expr.h"
 #include "../../transforms/infer_layout_util.h"
 #include "../../transforms/pattern_util.h"
 #include "../op_common.h"

diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
@@ -42,7 +42,6 @@
 #include <utility>
 #include <vector>
 
-#include "../../arith/compute_expr.h"
 #include "../../runtime/thread_storage_scope.h"
 #include "../../tir/transforms/ir_util.h"
 #include "llvm_common.h"

diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
@@ -25,7 +25,6 @@
 #include <cctype>
 #include <iomanip>
 
-#include "../../arith/compute_expr.h"
 #include "../../arith/pattern_match.h"
 
 namespace tvm {

diff --git a/src/target/spirv/codegen_spirv.cc b/src/target/spirv/codegen_spirv.cc
@@ -25,11 +25,10 @@
 
 #include <tvm/runtime/container.h>
 #include <tvm/tir/expr.h>
+#include <tvm/tir/op.h>
 
 #include <string>
 
-#include "../../arith/compute_expr.h"
-
 namespace tvm {
 namespace codegen {
 

diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc
@@ -34,7 +34,6 @@
 #include <unordered_set>
 #include <utility>
 
-#include "../../arith/compute_expr.h"
 #include "../../arith/interval_set.h"
 #include "../schedule/message_passing.h"
 #include "op_util.h"
@@ -593,8 +592,8 @@ Stmt TransformUpdate(const Stage& stage, const std::unordered_map<IterVar, Range
     }
   }
 
-  return IfThenElseNode::make(arith::ComputeReduce<tir::OrNode>(conds, const_true(1)), update,
-                              body);
+  auto cond = foldl([](PrimExpr a, PrimExpr b) { return a || b; }, const_false(1), conds);
+  return IfThenElseNode::make(cond, update, body);
 }
 
 }  // namespace te

diff --git a/src/te/operation/op_util.cc b/src/te/operation/op_util.cc
@@ -29,7 +29,6 @@
 
 #include <string>
 
-#include "../../arith/compute_expr.h"
 #include "../../runtime/thread_storage_scope.h"
 #include "../schedule/message_passing.h"
 

diff --git a/src/te/operation/tensor_compute_op.cc b/src/te/operation/tensor_compute_op.cc
@@ -29,7 +29,6 @@
 
 #include <unordered_set>
 
-#include "../../arith/compute_expr.h"
 #include "./compute_op.h"
 #include "./op_util.h"
 

diff --git a/src/te/schedule/message_passing.cc b/src/te/schedule/message_passing.cc
@@ -26,8 +26,6 @@
 #include <tvm/arith/analyzer.h>
 #include <tvm/tir/expr.h>
 
-#include "../../arith/compute_expr.h"
-
 namespace tvm {
 namespace te {
 

diff --git a/src/te/schedule/schedule_dataflow_rewrite.cc b/src/te/schedule/schedule_dataflow_rewrite.cc
@@ -22,11 +22,11 @@
  */
 #include <tvm/te/operation.h>
 #include <tvm/te/schedule.h>
+#include <tvm/tir/op.h>
 #include <tvm/tir/stmt_functor.h>
 
 #include <unordered_set>
 
-#include "../../arith/compute_expr.h"
 #include "../../tir/transforms/ir_util.h"
 #include "message_passing.h"
 #include "operation_inline.h"
@@ -89,13 +89,14 @@ PrimExpr InjectPredicate(const Array<PrimExpr>& predicates, PrimExpr body) {
   using tir::SelectNode;
   if (predicates.size() == 0) return body;
   const ReduceNode* reduce = body.as<ReduceNode>();
+  auto fand = [](PrimExpr a, PrimExpr b) { return a && b; };
+
   if (reduce) {
     auto n = make_object<ReduceNode>(*reduce);
-    n->condition = n->condition && arith::ComputeReduce<tir::AndNode>(predicates, PrimExpr());
+    n->condition = foldl(fand, n->condition, predicates);
     return PrimExpr(n);
   }
-  return SelectNode::make(arith::ComputeReduce<tir::AndNode>(predicates, PrimExpr()), body,
-                          make_zero(body.dtype()));
+  return SelectNode::make(foldl(fand, const_true(1), predicates), body, make_zero(body.dtype()));
 }
 
 // Replace data flow appears in all stages given the tensor change.
@@ -707,7 +708,9 @@ Array<Tensor> Schedule::rfactor(const Tensor& tensor, const IterVar& axis, int f
   const ReduceNode* reduce = compute_op->body[idx].as<ReduceNode>();
   CHECK(reduce) << "Can only rfactor non-inline reductions";
   predicates.push_back(reduce->condition);
-  PrimExpr predicate = likely(arith::ComputeReduce<tir::AndNode>(predicates, PrimExpr()));
+  auto fand = [](PrimExpr a, PrimExpr b) { return a && b; };
+
+  PrimExpr predicate = likely(foldl(fand, const_true(1), predicates));
 
   std::unordered_map<const VarNode*, PrimExpr> vsub;
 

diff --git a/src/te/schedule/schedule_postproc_rewrite_for_tensor_core.cc b/src/te/schedule/schedule_postproc_rewrite_for_tensor_core.cc
@@ -36,7 +36,6 @@
 
 #include <unordered_map>
 
-#include "../../arith/compute_expr.h"
 #include "../../runtime/thread_storage_scope.h"
 
 namespace tvm {

diff --git a/src/tir/ir/buffer.cc b/src/tir/ir/buffer.cc
@@ -26,12 +26,11 @@
 #include <tvm/tir/analysis.h>
 #include <tvm/tir/buffer.h>
 #include <tvm/tir/expr.h>
+#include <tvm/tir/op.h>
 
 #include <iterator>
 #include <stack>
 
-#include "../../arith/compute_expr.h"
-
 namespace tvm {
 namespace tir {
 
@@ -367,7 +366,8 @@ PrimExpr Buffer::access_ptr(int access_mask, DataType ptr_type, int content_lane
     int highest_dim = 0;
     extent = self->strides[highest_dim] * self->shape[highest_dim] - offset;
   } else {
-    extent = arith::ComputeReduce<tir::MulNode>(self->shape, PrimExpr()) - offset;
+    auto fmul = [](PrimExpr a, PrimExpr b) { return a * b; };
+    extent = foldl(fmul, make_const(DataType::Int(32), 1), self->shape) - offset;
   }
   PrimExpr elem_offset = self->elem_offset + offset;
   if (content_lanes > 1) {

diff --git a/src/tir/transforms/arg_binder.cc b/src/tir/transforms/arg_binder.cc
@@ -25,8 +25,8 @@
 
 #include <tvm/runtime/device_api.h>
 #include <tvm/tir/expr.h>
+#include <tvm/tir/op.h>
 
-#include "../../arith/compute_expr.h"
 #include "ir_util.h"
 
 namespace tvm {
@@ -225,8 +225,9 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
                    << " expected to be compact array";
     if (conds.size() != 0) {
       auto stride_msg = tvm::tir::StringImmNode::make(stride_err_msg.str());
-      Stmt check = AssertStmtNode::make(arith::ComputeReduce<tir::AndNode>(conds, PrimExpr()),
-                                        stride_msg, EvaluateNode::make(0));
+      auto fand = [](PrimExpr a, PrimExpr b) { return a && b; };
+      Stmt check = AssertStmtNode::make(foldl(fand, const_true(1), conds), stride_msg,
+                                        EvaluateNode::make(0));
       check = IfThenElseNode::make(NotNode::make(is_null), check, Stmt());
       asserts_.emplace_back(SeqStmt({check, EvaluateNode::make(0)}));
     }

diff --git a/src/tir/transforms/inject_double_buffer.cc b/src/tir/transforms/inject_double_buffer.cc
@@ -26,7 +26,6 @@
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
-#include "../../arith/compute_expr.h"
 #include "ir_util.h"
 
 namespace tvm {
@@ -115,8 +114,9 @@ class DoubleBufferInjector : public StmtExprMutator {
   Stmt VisitStmt_(const AllocateNode* op) final {
     auto it = dbuffer_info_.find(op->buffer_var.get());
     if (it != dbuffer_info_.end()) {
+      auto fmul = [](PrimExpr a, PrimExpr b) { return a * b; };
       it->second.stride =
-          arith::ComputeReduce<MulNode>(op->extents, PrimExpr()) * op->dtype.lanes();
+          foldl(fmul, make_const(DataType::Int(32), 1), op->extents) * op->dtype.lanes();
       Stmt stmt = StmtExprMutator::VisitStmt_(op);
       op = stmt.as<AllocateNode>();
       Array<PrimExpr> new_extents{make_const(op->extents[0].dtype(), 2)};

diff --git a/src/tir/transforms/inject_virtual_thread.cc b/src/tir/transforms/inject_virtual_thread.cc
@@ -27,7 +27,6 @@
 
 #include <unordered_set>
 
-#include "../../arith/compute_expr.h"
 #include "ir_util.h"
 
 namespace tvm {
@@ -368,7 +367,9 @@ class VTInjector : public StmtExprMutator {
     // always rewrite if not allow sharing.
     if (touched_var_.count(op->buffer_var.get()) || !allow_share_) {
       // place v on highest dimension.
-      PrimExpr stride = arith::ComputeReduce<MulNode>(op->extents, PrimExpr()) * op->dtype.lanes();
+      auto fmul = [](PrimExpr a, PrimExpr b) { return a * b; };
+      PrimExpr stride =
+          foldl(fmul, make_const(DataType::Int(32), 1), op->extents) * op->dtype.lanes();
       Array<PrimExpr> other;
       other.push_back(make_const(op->extents[0].dtype(), num_threads_));
       for (PrimExpr e : extents) {