[type] Fix struct-for block dim on bit_structs

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 {

taichi/transforms/flag_access.cpp

@@ -102,7 +102,9 @@ 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->type == SNodeType::bit_array) {
       a = a->parent;
     }
     return a;

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<StructForStmt>(
           snode, std::move(stmt->body), stmt->vectorize, stmt->bit_vectorize,
           stmt->parallelize, stmt->block_dim);