From a113199977d295cedbf29909eefa14170b68dc49 Mon Sep 17 00:00:00 2001 From: Yuanming Hu Date: Mon, 11 Jan 2021 20:30:07 -0500 Subject: [PATCH 1/2] [type] Fix struct-for block dim on bit_structs --- taichi/codegen/codegen_llvm.cpp | 8 +++----- taichi/transforms/flag_access.cpp | 3 ++- taichi/transforms/lower_ast.cpp | 8 ++++++++ 3 files changed, 13 insertions(+), 6 deletions(-) diff --git a/taichi/codegen/codegen_llvm.cpp b/taichi/codegen/codegen_llvm.cpp index 693b9ef4c8606..89dc2ed3b84a6 100644 --- a/taichi/codegen/codegen_llvm.cpp +++ b/taichi/codegen/codegen_llvm.cpp @@ -2041,13 +2041,11 @@ void CodeGenLLVM::create_offload_struct_for(OffloadedStmt *stmt, bool spmd) { llvm::Function *body = nullptr; auto leaf_block = stmt->snode; - // When looping over bit_arrays and bit_structs, we generate struct for on - // their parent node (usually "dense") instead of itself for higher + // When looping over bit_arrays, we always vectorize and generate struct for + // on their parent node (usually "dense") instead of itself for higher // performance. Also, note that the loop must be bit_vectorized for // bit_arrays, and their parent must be "dense". - if (leaf_block->type == SNodeType::bit_struct) { - leaf_block = leaf_block->parent; - } else if (leaf_block->type == SNodeType::bit_array) { + if (leaf_block->type == SNodeType::bit_array) { if (leaf_block->parent->type == SNodeType::dense) { leaf_block = leaf_block->parent; } else { diff --git a/taichi/transforms/flag_access.cpp b/taichi/transforms/flag_access.cpp index 377f09695c57a..7c3f37a77a363 100644 --- a/taichi/transforms/flag_access.cpp +++ b/taichi/transforms/flag_access.cpp @@ -102,7 +102,8 @@ class WeakenAccess : public BasicStmtVisitor { } static SNode *least_sparse_ancestor(SNode *a) { - while (a->type == SNodeType::place || a->type == SNodeType::dense) { + while (a->type == SNodeType::place || a->type == SNodeType::dense || + a->type == SNodeType::bit_struct) { a = a->parent; } return a; diff --git a/taichi/transforms/lower_ast.cpp b/taichi/transforms/lower_ast.cpp index 3d7e5b0b3f4cf..82316f9a34afb 100644 --- a/taichi/transforms/lower_ast.cpp +++ b/taichi/transforms/lower_ast.cpp @@ -281,6 +281,14 @@ class LowerAST : public IRVisitor { offsets = snode->index_offsets; snode = snode->parent; } + + // Climb up one more level if inside bit_struct. + // Note that when looping over bit_structs, we generate + // struct for on their parent node instead of itself for + // higher performance. + if (snode->type == SNodeType::bit_struct) + snode = snode->parent; + auto &&new_for = std::make_unique( snode, std::move(stmt->body), stmt->vectorize, stmt->bit_vectorize, stmt->parallelize, stmt->block_dim); From 6cf8948e19e9742f34d91832a08a3345f0128af8 Mon Sep 17 00:00:00 2001 From: Yuanming Hu Date: Mon, 11 Jan 2021 20:36:24 -0500 Subject: [PATCH 2/2] . --- taichi/transforms/flag_access.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/taichi/transforms/flag_access.cpp b/taichi/transforms/flag_access.cpp index 7c3f37a77a363..18f28adf0e6ad 100644 --- a/taichi/transforms/flag_access.cpp +++ b/taichi/transforms/flag_access.cpp @@ -103,7 +103,8 @@ class WeakenAccess : public BasicStmtVisitor { static SNode *least_sparse_ancestor(SNode *a) { while (a->type == SNodeType::place || a->type == SNodeType::dense || - a->type == SNodeType::bit_struct) { + a->type == SNodeType::bit_struct || + a->type == SNodeType::bit_array) { a = a->parent; } return a;