From c60c70aa708a17687047f1bfa8364439e59f61a7 Mon Sep 17 00:00:00 2001 From: Yi Xu Date: Wed, 26 Oct 2022 18:23:43 +0800 Subject: [PATCH 1/2] [opt] Remove redundant mod for SNode access under packed mode --- taichi/transforms/scalar_pointer_lowerer.cpp | 14 +++++++++++- .../scalar_pointer_lowerer_test.cpp | 22 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/taichi/transforms/scalar_pointer_lowerer.cpp b/taichi/transforms/scalar_pointer_lowerer.cpp index f4e41bd787df6..c4692d9fe9246 100644 --- a/taichi/transforms/scalar_pointer_lowerer.cpp +++ b/taichi/transforms/scalar_pointer_lowerer.cpp @@ -55,6 +55,8 @@ void ScalarPointerLowerer::run() { total_shape[j] *= s->extractors[j].shape; } } + std::array is_first_extraction; + is_first_extraction.fill(true); if (path_length_ == 0) return; @@ -80,7 +82,17 @@ void ScalarPointerLowerer::run() { const int prev = total_shape[k]; total_shape[k] /= snode->extractors[k].shape; const int next = total_shape[k]; - extracted = generate_mod_x_div_y(lowered_, indices_[k_], prev, next); + if (is_first_extraction[k]) { + // Upon first extraction on axis k, "indices_[k_]" is the user + // coordinate on axis k and "prev" is the total shape of axis k. + // Unless it is an invalid out-of-bound access, we can assume + // "indices_[k_] < prev" so we don't need a mod here. + auto const_next = lowered_->push_back(TypedConstant(next)); + extracted = lowered_->push_back(BinaryOpType::div, indices_[k_], const_next); + is_first_extraction[k] = false; + } else { + extracted = generate_mod_x_div_y(lowered_, indices_[k_], prev, next); + } } else { const int end = start_bits[k]; start_bits[k] -= snode->extractors[k].num_bits; diff --git a/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp b/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp index 939c4ffb34b42..f9aef884bf785 100644 --- a/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp +++ b/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp @@ -103,5 +103,27 @@ TEST_F(ScalarPointerLowererTest, Basic) { } } +TEST(ScalarPointerLowerer, EliminateMod) { + const bool kPacked = true; + IRBuilder builder; + VecStatement lowered; + Stmt *index = builder.get_int32(2); + auto root = std::make_unique(/*depth=*/0, SNodeType::root); + SNode *dense_1 = &(root->dense({Axis{2}, Axis{1}}, /*size=*/7, kPacked)); + SNode *dense_2 = &(root->dense({Axis{1}}, /*size=*/3, kPacked)); + SNode *dense_3 = &(dense_2->dense({Axis{0}}, /*size=*/5, kPacked)); + SNode *leaf_1 = &(dense_1->insert_children(SNodeType::place)); + SNode *leaf_2 = &(dense_3->insert_children(SNodeType::place)); + LowererImpl lowerer_1{leaf_1, {index, index}, SNodeOpType::undefined, + /*is_bit_vectorized=*/false, &lowered, kPacked}; + lowerer_1.run(); + LowererImpl lowerer_2{leaf_2, {index}, SNodeOpType::undefined, + /*is_bit_vectorized=*/false, &lowered, kPacked}; + lowerer_2.run(); + for (int i = 0; i < lowered.size(); i++) { + ASSERT_FALSE(lowered[i]->is() && + lowered[i]->as()->op_type == BinaryOpType::mod); + } +} } // namespace } // namespace taichi::lang From decb8c3cb0ce519709a3f948d6ebd5dfe4e03b15 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 26 Oct 2022 10:49:37 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- taichi/transforms/scalar_pointer_lowerer.cpp | 3 ++- .../transforms/scalar_pointer_lowerer_test.cpp | 16 ++++++++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/taichi/transforms/scalar_pointer_lowerer.cpp b/taichi/transforms/scalar_pointer_lowerer.cpp index c4692d9fe9246..e4c27959e9aec 100644 --- a/taichi/transforms/scalar_pointer_lowerer.cpp +++ b/taichi/transforms/scalar_pointer_lowerer.cpp @@ -88,7 +88,8 @@ void ScalarPointerLowerer::run() { // Unless it is an invalid out-of-bound access, we can assume // "indices_[k_] < prev" so we don't need a mod here. auto const_next = lowered_->push_back(TypedConstant(next)); - extracted = lowered_->push_back(BinaryOpType::div, indices_[k_], const_next); + extracted = lowered_->push_back( + BinaryOpType::div, indices_[k_], const_next); is_first_extraction[k] = false; } else { extracted = generate_mod_x_div_y(lowered_, indices_[k_], prev, next); diff --git a/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp b/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp index f9aef884bf785..dbfeabfe0c22c 100644 --- a/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp +++ b/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp @@ -114,11 +114,19 @@ TEST(ScalarPointerLowerer, EliminateMod) { SNode *dense_3 = &(dense_2->dense({Axis{0}}, /*size=*/5, kPacked)); SNode *leaf_1 = &(dense_1->insert_children(SNodeType::place)); SNode *leaf_2 = &(dense_3->insert_children(SNodeType::place)); - LowererImpl lowerer_1{leaf_1, {index, index}, SNodeOpType::undefined, - /*is_bit_vectorized=*/false, &lowered, kPacked}; + LowererImpl lowerer_1{leaf_1, + {index, index}, + SNodeOpType::undefined, + /*is_bit_vectorized=*/false, + &lowered, + kPacked}; lowerer_1.run(); - LowererImpl lowerer_2{leaf_2, {index}, SNodeOpType::undefined, - /*is_bit_vectorized=*/false, &lowered, kPacked}; + LowererImpl lowerer_2{leaf_2, + {index}, + SNodeOpType::undefined, + /*is_bit_vectorized=*/false, + &lowered, + kPacked}; lowerer_2.run(); for (int i = 0; i < lowered.size(); i++) { ASSERT_FALSE(lowered[i]->is() &&