[test] Add tests for value_diff_ptr_index pass (#2244)

* [test] Add tests for value_diff_ptr_index pass

* comment

taichi/analysis/alias_analysis.cpp

@@ -78,9 +78,9 @@ AliasResult alias_analysis(Stmt *var1, Stmt *var2) {
     bool uncertain = false;
     for (int i = 0; i < (int)ptr1->indices.size(); i++) {
       auto diff = value_diff_ptr_index(ptr1->indices[i], ptr2->indices[i]);
-      if (!diff.first) {
+      if (!diff.is_diff_certain) {
         uncertain = true;
-      } else if (std::abs(diff.second) >=
+      } else if (std::abs(diff.diff_range) >=
                  (1 << snode->extractors[snode->physical_index_position[i]]
                            .trailing_bits)) {
         return AliasResult::different;

taichi/analysis/value_diff.cpp

@@ -5,7 +5,8 @@
 #include "taichi/ir/statements.h"
 #include "taichi/ir/visitors.h"
 
-TLANG_NAMESPACE_BEGIN
+namespace taichi {
+namespace lang {
 
 DiffRange operator+(const DiffRange &a, const DiffRange &b) {
   return DiffRange(a.related_() && b.related_(), a.coeff + b.coeff,
@@ -17,6 +18,8 @@ DiffRange operator-(const DiffRange &a, const DiffRange &b) {
                    a.low - b.high + 1, a.high - b.low);
 }
 
+namespace {
+
 class ValueDiffLoopIndex : public IRVisitor {
  public:
   // first: related, second: offset
@@ -138,7 +141,10 @@ class FindDirectValueBaseAndOffset : public IRVisitor {
   }
 };
 
-namespace irpass::analysis {
+}  // namespace
+
+namespace irpass {
+namespace analysis {
 
 DiffRange value_diff_loop_index(Stmt *stmt, Stmt *loop, int index_id) {
   TI_ASSERT(loop->is<StructForStmt>() || loop->is<OffloadedStmt>());
@@ -156,22 +162,20 @@ DiffRange value_diff_loop_index(Stmt *stmt, Stmt *loop, int index_id) {
   return diff.run();
 }
 
-std::pair<bool, int> value_diff_ptr_index(Stmt *val1, Stmt *val2) {
-  // <first>: whether the difference of the value of two statements is certain.
-  // <second>: the difference of the value of two statements (i.e. val1 - val2)
-  // if <first> is true.
-  if (val1 == val2)
-    return std::make_pair(true, 0);
+DiffPtrResult value_diff_ptr_index(Stmt *val1, Stmt *val2) {
+  if (val1 == val2) {
+    return DiffPtrResult::make_certain(0);
+  }
   auto v1 = FindDirectValueBaseAndOffset::run(val1);
   auto v2 = FindDirectValueBaseAndOffset::run(val2);
   if (!std::get<0>(v1) || !std::get<0>(v2) ||
       std::get<1>(v1) != std::get<1>(v2)) {
-    // uncertain
-    return std::make_pair(false, 0);
+    return DiffPtrResult::make_uncertain();
   }
-  return std::make_pair(true, std::get<2>(v1) - std::get<2>(v2));
+  return DiffPtrResult::make_certain(std::get<2>(v1) - std::get<2>(v2));
 }
 
-}  // namespace irpass::analysis
-
-TLANG_NAMESPACE_END
+}  // namespace analysis
+}  // namespace irpass
+}  // namespace lang
+}  // namespace taichi

taichi/ir/analysis.h

@@ -6,7 +6,8 @@
 #include <unordered_set>
 #include <unordered_map>
 
-TLANG_NAMESPACE_BEGIN
+namespace taichi {
+namespace lang {
 
 class DiffRange {
  private:
@@ -57,7 +58,8 @@ class IRBank;
 class AsyncStateSet;
 
 // IR Analysis
-namespace irpass::analysis {
+namespace irpass {
+namespace analysis {
 
 AliasResult alias_analysis(Stmt *var1, Stmt *var2);
 std::unique_ptr<ControlFlowGraph> build_cfg(IRNode *root);
@@ -128,13 +130,45 @@ bool same_value(
     Stmt *stmt1,
     Stmt *stmt2,
     const std::optional<std::unordered_map<int, int>> &id_map = std::nullopt);
+
 DiffRange value_diff_loop_index(Stmt *stmt, Stmt *loop, int index_id);
-std::pair<bool, int> value_diff_ptr_index(Stmt *val1, Stmt *val2);
+
+/**
+ * Result of the value_diff_ptr_index pass.
+ */
+struct DiffPtrResult {
+  // Whether the difference of the checkd statements is certain.
+  bool is_diff_certain{false};
+  // The difference of the value of two statements (i.e. val1 - val2). This is
+  // meaningful only when |is_diff_certain| is true.
+  int diff_range{0};
+
+  static DiffPtrResult make_certain(int diff) {
+    return DiffPtrResult{/*is_diff_certain=*/true, /*diff_range=*/diff};
+  }
+  static DiffPtrResult make_uncertain() {
+    return DiffPtrResult{/*is_diff_certain=*/false, /*diff_range=*/0};
+  }
+};
+
+/**
+ * Checks if the difference of the value of the two statements is certain.
+ *
+ * @param val1
+ *   The first statement to check
+ *
+ * @param val2
+ *   The second statement to check
+ */
+DiffPtrResult value_diff_ptr_index(Stmt *val1, Stmt *val2);
+
 std::unordered_set<Stmt *> constexpr_prop(
     Block *block,
     std::function<bool(Stmt *)> is_const_seed);
-void verify(IRNode *root);
 
-}  // namespace irpass::analysis
+void verify(IRNode *root);
 
-TLANG_NAMESPACE_END
+}  // namespace analysis
+}  // namespace irpass
+}  // namespace lang
+}  // namespace taichi

taichi/ir/ir_builder.h

@@ -11,11 +11,6 @@ class IRBuilder {
     int position{0};
   };
 
- private:
-  std::unique_ptr<IRNode> root_{nullptr};
-  InsertPoint insert_point_;
-
- public:
   IRBuilder();
 
   // Clear the IR and the insertion point.
@@ -231,6 +226,10 @@ class IRBuilder {
       TI_ERROR("Statement {} is not a global pointer.", ptr->name());
     }
   }
+
+ private:
+  std::unique_ptr<IRNode> root_{nullptr};
+  InsertPoint insert_point_;
 };
 
 TLANG_NAMESPACE_END

tests/cpp_new/analysis/value_diff_test.cpp

@@ -0,0 +1,64 @@
+#include "gtest/gtest.h"
+
+#include "taichi/ir/analysis.h"
+#include "taichi/ir/ir_builder.h"
+#include "taichi/ir/statements.h"
+
+namespace taichi {
+namespace lang {
+namespace irpass {
+namespace analysis {
+
+TEST(ValueDiffPtrIndex, ConstI32) {
+  IRBuilder builder;
+
+  auto *const1 = builder.get_int32(42);
+  auto *const2 = builder.get_int32(2);
+
+  const auto diff = value_diff_ptr_index(const1, const2);
+  EXPECT_TRUE(diff.is_diff_certain);
+  EXPECT_EQ(diff.diff_range, 40);
+}
+
+TEST(ValueDiffPtrIndex, ConstF32) {
+  IRBuilder builder;
+
+  auto *const1 = builder.get_float32(1.0f);
+  auto *const2 = builder.get_float32(1.0f);
+
+  const auto diff = value_diff_ptr_index(const1, const2);
+  // We don't check floating-point numbers since the pass is for pointer indices
+  EXPECT_FALSE(diff.is_diff_certain);
+}
+
+TEST(ValueDiffPtrIndex, BinOp) {
+  IRBuilder builder;
+
+  auto *alloca = builder.create_local_var(PrimitiveType::i32);
+  auto *load = builder.create_local_load(alloca);
+  auto *const1 = builder.get_int32(1);
+  auto *bin1 = builder.create_add(load, const1);
+  auto *bin2 = builder.create_sub(load, const1);
+
+  const auto diff = value_diff_ptr_index(bin1, bin2);
+  EXPECT_TRUE(diff.is_diff_certain);
+  EXPECT_EQ(diff.diff_range, 2);
+}
+
+TEST(ValueDiffPtrIndex, BinOpButDiff) {
+  IRBuilder builder;
+
+  auto *val1_1 = builder.get_int32(1);
+  auto *val1_2 = builder.get_int32(1);
+  auto *val2 = builder.get_int32(3);
+  auto *bin1 = builder.create_add(val1_1, val2);
+  auto *bin2 = builder.create_add(val1_2, val2);
+
+  const auto diff = value_diff_ptr_index(bin1, bin2);
+  EXPECT_FALSE(diff.is_diff_certain);
+}
+
+}  // namespace analysis
+}  // namespace irpass
+}  // namespace lang
+}  // namespace taichi