【CINN】Add IndexExpr (#68617)

liuruyan · web-flow · commit 679a1f0ecc33 · 2024-10-11T10:23:17.000+08:00
* add index_expr

* remove IterExpr.
diff --git a/paddle/cinn/ir/ir.cc b/paddle/cinn/ir/ir.cc
@@ -15,6 +15,7 @@
 #include "paddle/cinn/ir/ir.h"
 
 #include <map>
+#include <optional>
 #include <string>
 #include <vector>
 #include "paddle/cinn/common/cinn_value.h"
@@ -1516,5 +1517,115 @@ void Block::Verify() const {}
 
 void PrimitiveNode::Verify() const {}
 
+IndexExpr &IndexExpr::operator=(const IndexExpr &other) {
+  *static_cast<IrNodeRef *>(this) = *static_cast<const IrNodeRef *>(&other);
+  return *this;
+}
+
+static std::optional<IndexExpr> SimplifyAdd(IndexExpr lhs, IndexExpr rhs) {
+  auto lhsConst = lhs.As<IntImm>();
+  auto rhsConst = rhs.As<IntImm>();
+  if (lhsConst && rhsConst) {
+    return IndexExpr(lhsConst->value + rhsConst->value);
+  }
+
+  if (lhsConst && !rhsConst) {
+    return rhs + lhs;
+  }
+  if (rhsConst && rhsConst->value == 0) {
+    return lhs;
+  }
+
+  return std::nullopt;
+}
+static std::optional<IndexExpr> SimplifySub(IndexExpr lhs, IndexExpr rhs) {
+  auto lhsConst = lhs.As<IntImm>();
+  auto rhsConst = rhs.As<IntImm>();
+
+  if (lhsConst && rhsConst) {
+    return IndexExpr(lhsConst->value - rhsConst->value);
+  }
+
+  if (rhsConst && rhsConst->value == 0) {
+    return lhs;
+  }
+
+  return std::nullopt;
+}
+
+static std::optional<IndexExpr> SimplifyMul(IndexExpr lhs, IndexExpr rhs) {
+  auto lhsConst = lhs.As<IntImm>();
+  auto rhsConst = rhs.As<IntImm>();
+
+  if (lhsConst && rhsConst) {
+    return IndexExpr(lhsConst->value * rhsConst->value);
+  }
+
+  if (lhsConst && !rhsConst) {
+    return rhs * lhs;
+  }
+
+  if (rhsConst) {
+    if (rhsConst->value == 0) {
+      return IndexExpr(0);
+    }
+    if (rhsConst->value == 1) {
+      return lhs;
+    }
+  }
+
+  return std::nullopt;
+}
+
+static std::optional<IndexExpr> SimplifyDiv(IndexExpr lhs, IndexExpr rhs) {
+  auto lhsConst = lhs.As<IntImm>();
+  auto rhsConst = rhs.As<IntImm>();
+
+  if (lhsConst && rhsConst) {
+    return IndexExpr(lhsConst->value / rhsConst->value);
+  }
+
+  if (rhsConst && rhsConst->value == 1) {
+    return lhs;
+  }
+
+  return std::nullopt;
+}
+
+static std::optional<IndexExpr> SimplifyMod(IndexExpr lhs, IndexExpr rhs) {
+  auto lhsConst = lhs.As<IntImm>();
+  auto rhsConst = rhs.As<IntImm>();
+
+  if (lhsConst && rhsConst) {
+    return IndexExpr(lhsConst->value % rhsConst->value);
+  }
+
+  if (rhsConst && rhsConst->value == 1) {
+    return IndexExpr(0);
+  }
+
+  return std::nullopt;
+}
+
+#define DEFINE_BINARY_OPERATOR(op, simplifyFunc, makeFunc)  \
+  IndexExpr IndexExpr::operator op(int64_t v) const {       \
+    return *this op IndexExpr(v);                           \
+  }                                                         \
+  IndexExpr IndexExpr::operator op(int32_t v) const {       \
+    return *this op IndexExpr(v);                           \
+  }                                                         \
+  IndexExpr IndexExpr::operator op(IndexExpr other) const { \
+    if (auto simplified = simplifyFunc(*this, other))       \
+      return simplified.value();                            \
+    return makeFunc(*this, other);                          \
+  }
+
+DEFINE_BINARY_OPERATOR(+, SimplifyAdd, Add::Make)
+DEFINE_BINARY_OPERATOR(-, SimplifySub, Sub::Make)
+DEFINE_BINARY_OPERATOR(*, SimplifyMul, Mul::Make)
+DEFINE_BINARY_OPERATOR(/, SimplifyDiv, Div::Make)
+DEFINE_BINARY_OPERATOR(%, SimplifyMod, Mod::Make)
+
+#undef DEFINE_BINARY_OPERATOR
 }  // namespace ir
 }  // namespace cinn
diff --git a/paddle/cinn/ir/ir.h b/paddle/cinn/ir/ir.h
@@ -1012,6 +1012,32 @@ struct Block : public ExprNode<Block> {
   static const IrNodeTy _node_type_ = IrNodeTy::Block;
 };
 
+struct IndexExpr : public Expr {
+ public:
+  IndexExpr() = default;
+  IndexExpr(const IndexExpr& other) : Expr(other.ptr()) {}
+  IndexExpr(IrNode* p) : Expr(p) {}            // NOLINT
+  IndexExpr(const Expr& expr) : Expr(expr) {}  // NOLINT
+
+  explicit IndexExpr(int32_t x) : Expr(x) {}
+  explicit IndexExpr(int64_t x) : Expr(x) {}
+
+  IndexExpr& operator=(const IndexExpr& other);
+
+#define DEFINE_OPERATOR(op)               \
+  IndexExpr operator op(int64_t v) const; \
+  IndexExpr operator op(int32_t v) const; \
+  IndexExpr operator op(IndexExpr other) const;
+
+  DEFINE_OPERATOR(+)
+  DEFINE_OPERATOR(-)
+  DEFINE_OPERATOR(*)
+  DEFINE_OPERATOR(/)
+  DEFINE_OPERATOR(%)
+
+#undef DEFINE_OPERATOR
+};
+
 /**
  * \brief IterMark is a special ExprNode, which can be used to mark ther entire
  * ierator. source is a IterSum or iterator. extent is the extent of the
diff --git a/paddle/cinn/ir/ir_base.cc b/paddle/cinn/ir/ir_base.cc
@@ -247,6 +247,41 @@ double Expr::get_constant() const {
 
 bool Expr::is_var() const { return As<_Var_>(); }
 
+bool Expr::is_index() const {
+  switch (node_type()) {
+    case ir::IrNodeTy::_Var_:
+      return true;
+    case ir::IrNodeTy::IntImm: {
+      if (type().is_index_type()) return true;
+    }
+    case ir::IrNodeTy::Add:
+      [[fallthrough]];
+    case ir::IrNodeTy::Sub:
+      [[fallthrough]];
+    case ir::IrNodeTy::Mul:
+      [[fallthrough]];
+    case ir::IrNodeTy::Div:
+      [[fallthrough]];
+    case ir::IrNodeTy::Mod:
+      return p_->operand(0).is_index() && p_->operand(1).is_index();
+  }
+  return false;
+}
+
+const IndexExpr Expr::as_index() const {
+  if (is_index()) {
+    return IndexExpr(*this);
+  }
+  PADDLE_THROW(::common::errors::InvalidType("Expr is not IndexExpr!"));
+}
+
+IndexExpr Expr::as_index() {
+  if (is_index()) {
+    return IndexExpr(*this);
+  }
+  PADDLE_THROW(::common::errors::InvalidType("Expr is not IndexExpr!"));
+}
+
 _Buffer_ *Expr::as_buffer() { return As<_Buffer_>(); }
 const _Buffer_ *Expr::as_buffer() const { return As<_Buffer_>(); }
 Buffer Expr::as_buffer_ref() const { return Buffer(&Reference(as_buffer())); }
diff --git a/paddle/cinn/ir/ir_base.h b/paddle/cinn/ir/ir_base.h
@@ -347,6 +347,7 @@ struct StringImm : public ExprNode<StringImm> {
 };
 
 class Var;
+class IndexExpr;
 /**
  * An expression that represents some value or the result of some operations.
  */
@@ -433,6 +434,11 @@ struct Expr : public IrNodeRef {
 
   bool is_var() const;
 
+  bool is_index() const;
+
+  IndexExpr as_index();
+  const IndexExpr as_index() const;
+
   operator Var();
 
   Type type() const { return p_->type(); }
diff --git a/paddle/cinn/ir/ir_visitor.cc b/paddle/cinn/ir/ir_visitor.cc
@@ -28,5 +28,12 @@ bool operator==(Expr a, Expr b) {
 
 bool operator!=(Expr a, Expr b) { return !(a == b); }
 
+bool operator==(IndexExpr a, IndexExpr b) {
+  if (a.get() == b.get()) return true;
+  return ir_utils::IRCompare(a, b);
+}
+
+bool operator!=(IndexExpr a, IndexExpr b) { return !(a == b); }
+
 }  // namespace ir
 }  // namespace cinn
diff --git a/paddle/cinn/ir/ir_visitor.h b/paddle/cinn/ir/ir_visitor.h
@@ -102,5 +102,8 @@ struct IRVisitor : public IRVisitorRequireReImpl<void> {
 bool operator==(Expr a, Expr b);
 bool operator!=(Expr a, Expr b);
 
+bool operator==(IndexExpr a, IndexExpr b);
+bool operator!=(IndexExpr a, IndexExpr b);
+
 }  // namespace ir
 }  // namespace cinn
diff --git a/test/cpp/pir/cinn/adt/CMakeLists.txt b/test/cpp/pir/cinn/adt/CMakeLists.txt
@@ -1,6 +1,7 @@
 if(WITH_TESTING AND WITH_CINN)
   paddle_test(map_expr_test SRCS map_expr_test.cc)
   set_tests_properties(map_expr_test PROPERTIES LABELS "RUN_TYPE=CINN")
+  paddle_test(test_index_expr SRCS iter_simplify_test.cc)
   paddle_test(test_iter_simplify SRCS iter_simplify_test.cc)
   paddle_test(merge_block_utils_test SRCS merge_block_utils_test.cc)
 endif()
diff --git a/test/cpp/pir/cinn/adt/Index_expr_test.cc b/test/cpp/pir/cinn/adt/Index_expr_test.cc
@@ -0,0 +1,43 @@
+// Copyright (c) 2024 CINN Authors. All Rights Reserved.
+//
+// Licensed 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
+//
+//     https://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.
+
+#include <glog/logging.h>
+#include <gtest/gtest.h>
+#include "paddle/cinn/ir/op/ir_operators.h"
+
+namespace cinn {
+namespace common {
+TEST(IndexExpr, IndexExpr_0) {
+  ir::IndexExpr a(14);
+  ir::IndexExpr b(7);
+  Expr d(6);
+  ir::Expr c0 = a + b;
+  ir::Expr c1 = a - b;
+  ir::Expr c2 = a * b;
+  ir::Expr c3 = a / b;
+  ir::Expr c4 = a % b;
+
+  ir::Expr c5 = a / d.as_index();
+  ir::Expr c6 = a % d.as_index();
+
+  EXPECT_EQ(c0, Expr(21));
+  EXPECT_EQ(c1, Expr(7));
+  EXPECT_EQ(c2, Expr(98));
+  EXPECT_EQ(c3, Expr(2));
+  EXPECT_EQ(c4, Expr(0));
+  EXPECT_EQ(c5, Expr(2));
+  EXPECT_EQ(c6, Expr(2));
+}
+}  // namespace common
+}  // namespace cinn
