From f74d75d91d6540e64c7e841de0f6b01c7b541809 Mon Sep 17 00:00:00 2001 From: Zhanlue Yang Date: Sun, 24 Dec 2023 16:16:07 +0800 Subject: [PATCH] [Bug] Fix CFG aliasing error with matrix of matrix (#8445) Issue: https://github.com/taichi-dev/taichi/issues/8435 ### Brief Summary copilot:summary ### Walkthrough copilot:walkthrough --- taichi/ir/control_flow_graph.cpp | 67 ++++++++++++++++++++++++++----- taichi/ir/frontend_ir.cpp | 9 ++--- tests/python/test_shared_array.py | 16 ++++++++ 3 files changed, 76 insertions(+), 16 deletions(-) diff --git a/taichi/ir/control_flow_graph.cpp b/taichi/ir/control_flow_graph.cpp index 5f8bd03fa22e9..5cbf8d8852237 100644 --- a/taichi/ir/control_flow_graph.cpp +++ b/taichi/ir/control_flow_graph.cpp @@ -384,6 +384,9 @@ void CFGNode::reaching_definition_analysis(bool after_lower_access) { if (after_lower_access && !((data_source_ptr->is() && data_source_ptr->as()->origin->is()) || + (data_source_ptr->is() && + data_source_ptr->as() + ->origin->is()) || data_source_ptr->is())) { // After lower_access, we only analyze local variables. continue; @@ -555,6 +558,8 @@ void CFGNode::live_variable_analysis(bool after_lower_access) { if (!after_lower_access || (load_ptr->is() && load_ptr->as()->origin->is()) || + (load_ptr->is() && + load_ptr->as()->origin->is()) || (load_ptr->is() || load_ptr->is())) { // After lower_access, we only analyze local variables and stacks. if (!contain_variable(live_kill, load_ptr)) { @@ -581,6 +586,8 @@ void CFGNode::live_variable_analysis(bool after_lower_access) { if (!after_lower_access || (store_ptr->is() && store_ptr->as()->origin->is()) || + (store_ptr->is() && + store_ptr->as()->origin->is()) || (store_ptr->is() || store_ptr->is())) { // After lower_access, we only analyze local variables and stacks. live_kill.insert(store_ptr); @@ -589,32 +596,41 @@ void CFGNode::live_variable_analysis(bool after_lower_access) { } } -static void update_aliased_stmts( +static void recursive_update_aliased_elements( const std::unordered_map> &tensor_to_matrix_ptrs_map, - const std::unordered_map &matrix_ptr_to_tensor_map, std::unordered_map &container, Stmt *key, bool to_erase) { if (tensor_to_matrix_ptrs_map.find(key) != tensor_to_matrix_ptrs_map.end()) { - auto scalars_address = tensor_to_matrix_ptrs_map.at(key); + const auto &elements_address = tensor_to_matrix_ptrs_map.at(key); // Update aliased MatrixPtrStmt for TensorType<>* - for (auto scalar_address : scalars_address) { + for (const auto &element_address : elements_address) { if (to_erase) { - if (container.find(scalar_address) != container.end()) { - TI_ASSERT(container[scalar_address] == - CFGNode::UseDefineStatus::NONE); - container.erase(scalar_address); + if (container.find(element_address) != container.end()) { + container.erase(element_address); } } else { - container[scalar_address] = CFGNode::UseDefineStatus::NONE; + container[element_address] = CFGNode::UseDefineStatus::NONE; + if (element_address->ret_type.ptr_removed()->is()) { + container[element_address] = CFGNode::UseDefineStatus::FULL; + } } + + // Recursively update aliased addresses + recursive_update_aliased_elements(tensor_to_matrix_ptrs_map, container, + element_address, to_erase); } } +} - // Update aliased TensorType<>* for MatrixPtrStmt +static void recursive_update_aliased_parent( + const std::unordered_map &matrix_ptr_to_tensor_map, + std::unordered_map &container, + Stmt *key, + bool to_erase) { if (matrix_ptr_to_tensor_map.find(key) != matrix_ptr_to_tensor_map.end()) { - auto tensor_address = matrix_ptr_to_tensor_map.at(key); + const auto &tensor_address = matrix_ptr_to_tensor_map.at(key); // no matter to_erase or not, the tensor_address is only partially defined // or used if (to_erase) { @@ -624,9 +640,29 @@ static void update_aliased_stmts( } else { container[tensor_address] = CFGNode::UseDefineStatus::PARTIAL; } + + // Recursively update aliased addresses + recursive_update_aliased_parent(matrix_ptr_to_tensor_map, container, + tensor_address, to_erase); } } +static void update_aliased_stmts( + const std::unordered_map> + &tensor_to_matrix_ptrs_map, + const std::unordered_map &matrix_ptr_to_tensor_map, + std::unordered_map &container, + Stmt *key, + bool to_erase) { + // Update aliased MatrixPtrStmt for TensorType<>* + recursive_update_aliased_elements(tensor_to_matrix_ptrs_map, container, key, + to_erase); + + // Update aliased TensorType<>* for MatrixPtrStmt + recursive_update_aliased_parent(matrix_ptr_to_tensor_map, container, key, + to_erase); +} + // Insert or erase "key" to "container". // In case where "key" being MatrixPtrStmt, we also update the aliased original // address. In case where "key" is involved with TensorType, we also update the @@ -648,6 +684,7 @@ static void update_container_with_alias( } else { container[key] = CFGNode::UseDefineStatus::NONE; } + // Recursively update aliased addresses update_aliased_stmts(tensor_to_matrix_ptrs_map, matrix_ptr_to_tensor_map, container, key, to_erase); } @@ -715,6 +752,8 @@ bool CFGNode::dead_store_elimination(bool after_lower_access) { if (!after_lower_access || (store_ptr->is() && store_ptr->as()->origin->is()) || + (store_ptr->is() && + store_ptr->as()->origin->is()) || (store_ptr->is() || store_ptr->is())) { // !may_contain_variable(live_in_this_node, store_ptr): address is not // loaded after this store @@ -816,6 +855,8 @@ bool CFGNode::dead_store_elimination(bool after_lower_access) { if (!after_lower_access || (load_ptr->is() && load_ptr->as()->origin->is()) || + (load_ptr->is() && + load_ptr->as()->origin->is()) || (load_ptr->is() || load_ptr->is())) { // live_load_in_this_node[addr]: tracks the // next load to the same address @@ -844,6 +885,8 @@ bool CFGNode::dead_store_elimination(bool after_lower_access) { if (!after_lower_access || (load_ptr->is() && load_ptr->as()->origin->is()) || + (load_ptr->is() && + load_ptr->as()->origin->is()) || (load_ptr->is() || load_ptr->is())) { // Addr is used in this node, so it's live in this node update_container_with_alias(tensor_to_matrix_ptrs_map, @@ -976,6 +1019,8 @@ void ControlFlowGraph::reaching_definition_analysis(bool after_lower_access) { auto stmt = nodes[i]->block->statements[j].get(); if ((stmt->is() && stmt->as()->origin->is()) || + (stmt->is() && + stmt->as()->origin->is()) || (!after_lower_access && (stmt->is() || stmt->is() || stmt->is() || stmt->is() || diff --git a/taichi/ir/frontend_ir.cpp b/taichi/ir/frontend_ir.cpp index a64a9fc6b2488..5bda730c48096 100644 --- a/taichi/ir/frontend_ir.cpp +++ b/taichi/ir/frontend_ir.cpp @@ -843,12 +843,11 @@ bool IndexExpression::is_local() const { } bool IndexExpression::is_global() const { - // Special case: Indexing into TensorType-element of ExternalPtrStmt - // or GlobalPtrStmt should be treated as global ptrs if (var.is()) { - TI_ASSERT(var.cast()->is_matrix_field() || - var.cast()->is_ndarray()); - return true; + // Special case: Pointer chasing. For example, if we are indexing into + // tensor elements of fields / ndarrays, this index expr should be treated + // as global. + return var.cast()->is_global(); } // Only Ndarray and Field comes outside from a kernel diff --git a/tests/python/test_shared_array.py b/tests/python/test_shared_array.py index 64e499eef2dad..00843b6e6cdfa 100644 --- a/tests/python/test_shared_array.py +++ b/tests/python/test_shared_array.py @@ -166,3 +166,19 @@ def test(): test() assert (y.to_numpy()[0] == [4.0, 8.0, 12.0, 16.0]).all() + + +@test_utils.test(arch=[ti.cuda], debug=True) +def test_shared_array_matrix(): + @ti.kernel + def foo(): + for x in range(10): + shared = ti.simt.block.SharedArray((10,), dtype=ti.math.vec3) + shared[x] = ti.Vector([x + 1, x + 2, x + 3]) + assert shared[x].z == x + 3 + assert (shared[x] == ti.Vector([x + 1, x + 2, x + 3])).all() + + print(shared[x].z) + print(shared[x]) + + foo()