Implementation of Common Subexpression Elimination for TIR

include/tvm/tir/transform.h

@@ -458,6 +458,14 @@ TVM_DLL Pass FlattenBuffer();
  */
 TVM_DLL Pass TextureFlatten();
 
+/*!
+ * \brief Implements a Common Subexpression Elimination (CSE) for TIR
+ *        which introduces let-in bindings for duplicated sub-expressions.
+ * \param enable_cse_tir Whether common subexpression elimination is enabled.
+ * \return The pass.
+ */
+TVM_DLL Pass CommonSubexprElimTIR(bool enable_cse_tir = true);
+
 /*!
  * \brief Unify all the thread bindings for "blockIdx.x/y/z", "threadIdx.x/y/z", and
  *        "vthread.x/y/z". Before the unification, two vars that are bound to a thread axis (e.g.,

python/tvm/tir/transform/transform.py

@@ -311,6 +311,17 @@ def BF16TypeLowering():
     return _ffi_api.BF16TypeLowering()  # type: ignore
 
 
+def CommonSubexprElimTIR(enable_cse_tir: bool = True):
+    """Replace redundant computations by new variables.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.CommonSubexprElimTIR(enable_cse_tir)  # type: ignore
+
+
 def RewriteUnsafeSelect():
     """Detect and rewrite unsafe select that contains memory access.
 

src/driver/driver_api.cc

@@ -44,6 +44,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.detect_global_barrier", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_bound_checkers", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_assert", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_vectorize", Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_cse_tir", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_storage_rewrite", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.is_entry_func", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.add_lower_pass", Array<Array<ObjectRef>>);
@@ -196,6 +197,7 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
       pass_ctx->GetConfig<Bool>("tir.disable_storage_rewrite", Bool(false)).value();
   bool instrument_bound_checkers =
       pass_ctx->GetConfig<Bool>("tir.instrument_bound_checkers", Bool(false)).value();
+  bool disable_cse_tir = pass_ctx->GetConfig<Bool>("tir.disable_cse_tir", Bool(false)).value();
 
   // Get any user-added passes
   Array<Array<ObjectRef>> add_lower_pass =
@@ -283,6 +285,9 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   if (instrument_bound_checkers) {
     pass_list.push_back(tir::transform::InstrumentBoundCheckers());
   }
+
+  pass_list.push_back(tir::transform::CommonSubexprElimTIR(!disable_cse_tir));
+
   return pass_list;
 }
 

src/tir/analysis/check_contains.cc

@@ -0,0 +1,98 @@
+/*
+ * 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 check_contains.cc
+  * \brief Implementation of the analysis that tells if an expression contains
+            a node that satisfies a given predicate.
+  */
+
+#include "check_contains.h"
+
+#include <tvm/tir/expr.h>
+
+#include <vector>
+
+namespace tvm {
+namespace tir {
+
+/*!
+ * \brief Toplevel (static) function that tells if an expression contains a subexpression that
+          satisfies a given predicate.
+ * \param expr The expression to check
+ * \param predicate The predicate that must be satisfied
+ * \return Whether `expr` contains a subexpression that satisfies `predicate`
+ */
+bool CheckContains::ExprContains(const PrimExpr& expr,
+                                 std::function<bool(const PrimExpr&)> predicate) {
+  CheckContains check_contains(predicate);
+  check_contains.VisitExpr(expr);
+  return check_contains.contains_it_;
+}
+
+/*!
+ * \brief Toplevel (static) function that tells if a statement contains a subexpression that
+          satisfies a given predicate.
+ * \param stmt The statement to check
+ * \param predicate The predicate that must be satisfied
+ * \return Whether `stmt` contains a subexpression that satisfies `predicate`
+ */
+bool CheckContains::StmtContains(const Stmt& stmt, std::function<bool(const PrimExpr&)> predicate) {
+  CheckContains check_contains(predicate);
+  check_contains.VisitStmt(stmt);
+  return check_contains.contains_it_;
+}
+
+/*!
+ * \brief Protected constructor of CheckContains.
+ * \param predicate The predicate that must be satisfied
+ */
+CheckContains::CheckContains(std::function<bool(const PrimExpr&)> predicate)
+    : predicate_(predicate) {}
+
+/*!
+ * \brief The method which overrides the generic dispatcher of StmtExprVisitor for expressions.
+ * \param expr The expression to visit
+ */
+void CheckContains::VisitExpr(const PrimExpr& expr) {
+  // If the predicate holds on `expr`, we know `expr` contains something which makes
+  // the predicate hold
+  if (predicate_(expr)) {
+    contains_it_ = true;
+  } else {
+    // Otherwise we continue to look for it recursively by calling the dispatcher
+    StmtExprVisitor::VisitExpr(expr);
+  }
+}
+
+/*!
+ * \brief The method which overrides the generic dispatcher of StmtExprVisitor for statements.
+ * \param stmt The statement to visit
+ */
+void CheckContains::VisitStmt(const Stmt& stmt) {
+  // We keep exploring only if `contains_it_` is false
+  if (!contains_it_) {
+    // and in order to do that we call the general dispatcher
+    StmtExprVisitor::VisitStmt(stmt);
+  }
+  // As otherwise we already have our answer
+}
+
+}  // namespace tir
+}  // namespace tvm

src/tir/analysis/check_contains.h

@@ -0,0 +1,60 @@
+/*
+ * 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 check_contains.h
+ * \brief Interface of the analysis that tells if an expression contains
+           a node that satisfies a given predicate.
+ */
+
+#ifndef TVM_TIR_ANALYSIS_CHECK_CONTAINS_H_
+#define TVM_TIR_ANALYSIS_CHECK_CONTAINS_H_
+
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt_functor.h>  // For the class StmtExprVisitor
+
+namespace tvm {
+namespace tir {
+
+/*!
+ * \brief Visitor which tells if a given expression or statement contains a subexpression
+          that satisfies a given predicate
+ */
+class CheckContains : public StmtExprVisitor {
+ public:
+  // Toplevel (static) functions
+  static bool ExprContains(const PrimExpr& expr, std::function<bool(const PrimExpr&)> predicate);
+  static bool StmtContains(const Stmt& stmt, std::function<bool(const PrimExpr&)> predicate);
+
+ protected:
+  // Constructor
+  explicit CheckContains(std::function<bool(const PrimExpr&)> predicate);
+
+  void VisitExpr(const PrimExpr& expr) override;
+  void VisitStmt(const Stmt& stmt) override;
+
+ private:
+  std::function<bool(const PrimExpr&)> predicate_;
+  bool contains_it_ = false;
+};
+
+}  // namespace tir
+}  // namespace tvm
+
+#endif  // TVM_TIR_ANALYSIS_CHECK_CONTAINS_H_