From 6dfec4c7b8709ce8aabb95abcda233a06301c9f3 Mon Sep 17 00:00:00 2001
From: Ye Kuang <k-ye@users.noreply.github.com>
Date: Sat, 28 Mar 2020 23:30:53 +0900
Subject: [PATCH] [Lang] Simplify dense.bitmasked to bitmasked (#670)

* [Lang] Simplify dense.bitmasked to bitmasked

* fix tests

* put back

* [skip ci] enforce code format

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
---
 python/taichi/lang/snode.py                   | 11 +++---
 taichi/backends/metal/runtime.cpp             | 11 +++---
 .../metal/shaders/runtime_structs.metal.h     |  2 +-
 taichi/codegen/codegen_llvm.h                 | 10 ++++--
 taichi/codegen/codegen_metal.cpp              |  6 ++--
 taichi/inc/snodes.inc.h                       |  1 +
 taichi/ir/ir.h                                |  2 +-
 taichi/ir/snode.cpp                           |  1 -
 taichi/ir/snode.h                             | 19 +++++++----
 taichi/llvm/llvm_context.cpp                  |  6 ++--
 taichi/python/export_lang.cpp                 |  6 ++--
 taichi/runtime/llvm/node_bitmasked.h          | 34 +++++++++++++++++++
 taichi/runtime/llvm/node_dense.h              | 24 ++-----------
 taichi/runtime/llvm/runtime.cpp               |  1 +
 taichi/struct/struct_llvm.cpp                 |  6 ++--
 taichi/struct/struct_metal.cpp                | 23 +++++++------
 .../transforms/demote_dense_struct_fors.cpp   |  3 --
 taichi/transforms/lower_ast.cpp               |  3 +-
 tests/python/test_bitmasked.py                | 34 +++++++++++++++++--
 tests/python/test_sparse_basics.py            | 26 --------------
 20 files changed, 130 insertions(+), 99 deletions(-)
 create mode 100644 taichi/runtime/llvm/node_bitmasked.h

diff --git a/python/taichi/lang/snode.py b/python/taichi/lang/snode.py
index 088b7ddc97f45d..f6cd83453262fd 100644
--- a/python/taichi/lang/snode.py
+++ b/python/taichi/lang/snode.py
@@ -23,9 +23,10 @@ def dynamic(self, index, dimension, chunk_size=None):
             chunk_size = dimension
         return SNode(self.ptr.dynamic(index[0], dimension, chunk_size))
 
-    def bitmasked(self, val=True):
-        self.ptr.bitmasked(val)
-        return self
+    def bitmasked(self, indices, dimensions):
+        if isinstance(dimensions, int):
+            dimensions = [dimensions] * len(indices)
+        return SNode(self.ptr.bitmasked(indices, dimensions))
 
     def place(self, *args):
         from .expr import Expr
@@ -72,8 +73,6 @@ def deactivate_all(self):
         for c in ch:
             c.deactivate_all()
         import taichi as ti
-        if self.ptr.type == ti.core.SNodeType.pointer or (
-                self.ptr.type == ti.core.SNodeType.dense
-                and self.ptr.is_bitmasked):
+        if self.ptr.type == ti.core.SNodeType.pointer or self.ptr.type == ti.core.SNodeType.bitmasked:
             from .meta import snode_deactivate
             snode_deactivate(self)
diff --git a/taichi/backends/metal/runtime.cpp b/taichi/backends/metal/runtime.cpp
index d74da41fa99f28..51b3026da95311 100644
--- a/taichi/backends/metal/runtime.cpp
+++ b/taichi/backends/metal/runtime.cpp
@@ -509,15 +509,14 @@ class MetalRuntime::Impl {
       rtm_meta->num_slots = sn_meta.num_slots;
       rtm_meta->mem_offset_in_parent = sn_meta.mem_offset_in_parent;
       switch (sn_meta.snode->type) {
+        case SNodeType::dense:
+          rtm_meta->type = SNodeMeta::Dense;
+          break;
         case SNodeType::root:
           rtm_meta->type = SNodeMeta::Root;
           break;
-        case SNodeType::dense:
-          if (sn_meta.snode->_bitmasked) {
-            rtm_meta->type = SNodeMeta::DenseBitmask;
-          } else {
-            rtm_meta->type = SNodeMeta::Dense;
-          }
+        case SNodeType::bitmasked:
+          rtm_meta->type = SNodeMeta::Bitmasked;
           break;
         default:
           TI_ERROR("Unsupported SNode type={}",
diff --git a/taichi/backends/metal/shaders/runtime_structs.metal.h b/taichi/backends/metal/shaders/runtime_structs.metal.h
index 85b757189f385c..e5e57363505d2f 100644
--- a/taichi/backends/metal/shaders/runtime_structs.metal.h
+++ b/taichi/backends/metal/shaders/runtime_structs.metal.h
@@ -48,7 +48,7 @@ STR(
     };
 
     struct SNodeMeta {
-      enum Type { Root = 0, Dense = 1, DenseBitmask = 2 };
+      enum Type { Root = 0, Dense = 1, Bitmasked = 2 };
       int32_t element_stride = 0;
       int32_t num_slots = 0;
       int32_t mem_offset_in_parent = 0;
diff --git a/taichi/codegen/codegen_llvm.h b/taichi/codegen/codegen_llvm.h
index d4a68440407d9a..8f3eb9ca6f5fb4 100644
--- a/taichi/codegen/codegen_llvm.h
+++ b/taichi/codegen/codegen_llvm.h
@@ -186,7 +186,6 @@ class CodeGenLLVM : public IRVisitor, public ModuleBuilder {
     if (snode->type == SNodeType::dense) {
       meta = std::make_unique<RuntimeObject>("DenseMeta", this, builder.get());
       emit_struct_meta_base("Dense", meta->ptr, snode);
-      meta->call("set_bitmasked", tlctx->get_constant(snode->_bitmasked));
       meta->call("set_morton_dim", tlctx->get_constant((int)snode->_morton));
     } else if (snode->type == SNodeType::pointer) {
       meta =
@@ -200,6 +199,10 @@ class CodeGenLLVM : public IRVisitor, public ModuleBuilder {
           std::make_unique<RuntimeObject>("DynamicMeta", this, builder.get());
       emit_struct_meta_base("Dynamic", meta->ptr, snode);
       meta->call("set_chunk_size", tlctx->get_constant(snode->chunk_size));
+    } else if (snode->type == SNodeType::bitmasked) {
+      meta =
+          std::make_unique<RuntimeObject>("BitmaskedMeta", this, builder.get());
+      emit_struct_meta_base("Bitmasked", meta->ptr, snode);
     } else {
       TI_P(snode_type_name(snode->type));
       TI_NOT_IMPLEMENTED;
@@ -1067,6 +1070,8 @@ class CodeGenLLVM : public IRVisitor, public ModuleBuilder {
       return "pointer";
     } else if (snode->type == SNodeType::hash) {
       return "Hash";
+    } else if (snode->type == SNodeType::bitmasked) {
+      return "Bitmasked";
     } else {
       TI_P(snode_type_name(snode->type));
       TI_NOT_IMPLEMENTED
@@ -1141,7 +1146,8 @@ class CodeGenLLVM : public IRVisitor, public ModuleBuilder {
       stmt->value = builder->CreateGEP(parent, stmt->input_index->value);
     } else if (snode->type == SNodeType::dense ||
                snode->type == SNodeType::pointer ||
-               snode->type == SNodeType::dynamic) {
+               snode->type == SNodeType::dynamic ||
+               snode->type == SNodeType::bitmasked) {
       if (stmt->activate) {
         call(snode, stmt->input_snode->value, "activate",
              {stmt->input_index->value});
diff --git a/taichi/codegen/codegen_metal.cpp b/taichi/codegen/codegen_metal.cpp
index 532d349d88851d..04b5c3a92bfa8e 100644
--- a/taichi/codegen/codegen_metal.cpp
+++ b/taichi/codegen/codegen_metal.cpp
@@ -150,7 +150,7 @@ class MetalKernelCodegen : public IRVisitor {
     emit(R"({}_ch {} = {}.children({});)", sn->node_type_name, stmt->raw_name(),
          parent, index_name);
     if (stmt->activate) {
-      TI_ASSERT(sn->type == SNodeType::dense && sn->_bitmasked);
+      TI_ASSERT(sn->type == SNodeType::bitmasked);
       emit("{{");
       {
         ScopedIndent s(line_appender_);
@@ -707,7 +707,7 @@ class MetalKernelCodegen : public IRVisitor {
 
   std::string make_snode_meta_bm(const SNode *sn,
                                  const std::string &var_name) const {
-    TI_ASSERT(sn->type == SNodeType::dense && sn->_bitmasked);
+    TI_ASSERT(sn->type == SNodeType::bitmasked);
     const auto &meta = compiled_snodes_->snode_descriptors.find(sn->id)->second;
     LineAppender la = line_appender_;
     // Keep the indentation settings only
@@ -716,7 +716,7 @@ class MetalKernelCodegen : public IRVisitor {
     la.append("SNodeMeta {};", var_name);
     la.append("{}.element_stride = {};", var_name, meta.element_stride);
     la.append("{}.num_slots = {};", var_name, meta.num_slots);
-    la.append("{}.type = {};", var_name, meta.num_slots);
+    la.append("{}.type = {};", var_name, (int)shaders::SNodeMeta::Bitmasked);
     return la.lines();
   }
 
diff --git a/taichi/inc/snodes.inc.h b/taichi/inc/snodes.inc.h
index 051ae8661b20b4..016483cbdafd05 100644
--- a/taichi/inc/snodes.inc.h
+++ b/taichi/inc/snodes.inc.h
@@ -2,6 +2,7 @@ PER_SNODE(root)
 PER_SNODE(dense)
 PER_SNODE(dynamic)
 PER_SNODE(pointer)
+PER_SNODE(bitmasked)
 PER_SNODE(hash)
 PER_SNODE(place)
 PER_SNODE(undefined)
diff --git a/taichi/ir/ir.h b/taichi/ir/ir.h
index ffffee2338e7a0..181ed7aa836c63 100644
--- a/taichi/ir/ir.h
+++ b/taichi/ir/ir.h
@@ -2073,7 +2073,7 @@ class SNodeOpExpression : public Expression {
       // It should be lowered into a pointer to parent and an index.
       TI_ERROR_IF(
           snode->type != SNodeType::pointer && snode->type != SNodeType::hash &&
-              !(snode->type == SNodeType::dense && snode->_bitmasked),
+              snode->type != SNodeType::bitmasked,
           "ti.is_active only works on pointer, hash or bitmasked nodes.");
       ret.push_back<SNodeOpStmt>(SNodeOpType::is_active, snode, indices_stmt);
     } else {
diff --git a/taichi/ir/snode.cpp b/taichi/ir/snode.cpp
index d7511edda9d19c..7953fa0adb35db 100644
--- a/taichi/ir/snode.cpp
+++ b/taichi/ir/snode.cpp
@@ -222,7 +222,6 @@ SNode::SNode(int depth, SNodeType t) : depth(depth), type(t) {
   has_ambient = false;
   dt = DataType::gen;
   _morton = false;
-  _bitmasked = false;
 
   reader_kernel = nullptr;
   writer_kernel = nullptr;
diff --git a/taichi/ir/snode.h b/taichi/ir/snode.h
index ad5fb0c796322a..f6644b43e79edd 100644
--- a/taichi/ir/snode.h
+++ b/taichi/ir/snode.h
@@ -87,7 +87,6 @@ class SNode {
   std::string node_type_name;
   SNodeType type;
   bool _morton{};
-  bool _bitmasked{};
 
   std::string get_node_type_name() const {
     return fmt::format("S{}", id);
@@ -140,6 +139,19 @@ class SNode {
     return SNode::pointer(std::vector<Index>{index}, size);
   }
 
+  SNode &bitmasked(const std::vector<Index> &indices,
+                   const std::vector<int> &sizes) {
+    return create_node(indices, sizes, SNodeType::bitmasked);
+  }
+
+  SNode &bitmasked(const std::vector<Index> &indices, int sizes) {
+    return create_node(indices, std::vector<int>{sizes}, SNodeType::bitmasked);
+  }
+
+  SNode &bitmasked(const Index &index, int size) {
+    return SNode::bitmasked(std::vector<Index>{index}, size);
+  }
+
   SNode &hash(const std::vector<Index> &indices,
               const std::vector<int> &sizes) {
     return create_node(indices, sizes, SNodeType::hash);
@@ -194,11 +206,6 @@ class SNode {
     return *this;
   }
 
-  SNode &bitmasked(bool val = true) {
-    _bitmasked = val;
-    return *this;
-  }
-
   // for float and double
   void write_float(const std::vector<int> &I, float64);
   float64 read_float(const std::vector<int> &I);
diff --git a/taichi/llvm/llvm_context.cpp b/taichi/llvm/llvm_context.cpp
index d5f446798b54fa..99c04b9524c0fb 100644
--- a/taichi/llvm/llvm_context.cpp
+++ b/taichi/llvm/llvm_context.cpp
@@ -136,9 +136,9 @@ void compile_runtime_bitcode(Arch arch) {
     if (ret) {
       TI_ERROR("LLVMRuntime compilation failed.");
     }
-    std::system(fmt::format("llvm-as {}runtime.ll -o {}{}", runtime_folder,
-                            runtime_folder, get_runtime_fn(arch))
-                    .c_str());
+    cmd = fmt::format("llvm-as {}runtime.ll -o {}{}", runtime_folder,
+                      runtime_folder, get_runtime_fn(arch));
+    std::system(cmd.c_str());
     TI_TRACE("runtime module bitcode compiled.");
     runtime_compiled.insert((int)arch);
   }
diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
index 61a221419f3459..35b40b08172a70 100644
--- a/taichi/python/export_lang.cpp
+++ b/taichi/python/export_lang.cpp
@@ -143,7 +143,6 @@ void export_lang(py::module &m) {
       .def(py::init<>())
       .def_readwrite("parent", &SNode::parent)
       .def_readonly("type", &SNode::type)
-      .def_readonly("is_bitmasked", &SNode::_bitmasked)
       .def("dense",
            (SNode & (SNode::*)(const std::vector<Index> &,
                                const std::vector<int> &))(&SNode::dense),
@@ -158,7 +157,10 @@ void export_lang(py::module &m) {
            py::return_value_policy::reference)
       .def("dynamic", &SNode::dynamic_chunked,
            py::return_value_policy::reference)
-      .def("bitmasked", &SNode::bitmasked)
+      .def("bitmasked",
+           (SNode & (SNode::*)(const std::vector<Index> &,
+                               const std::vector<int> &))(&SNode::bitmasked),
+           py::return_value_policy::reference)
       .def("place", (SNode & (SNode::*)(Expr &))(&SNode::place),
            py::return_value_policy::reference)
       .def("data_type", [](SNode *snode) { return snode->dt; })
diff --git a/taichi/runtime/llvm/node_bitmasked.h b/taichi/runtime/llvm/node_bitmasked.h
new file mode 100644
index 00000000000000..b8358c919f824e
--- /dev/null
+++ b/taichi/runtime/llvm/node_bitmasked.h
@@ -0,0 +1,34 @@
+#pragma once
+
+// Specialized Attributes and functions
+struct BitmaskedMeta : public StructMeta {
+  bool _;
+};
+
+STRUCT_FIELD(BitmaskedMeta, _);
+
+i32 Bitmasked_get_num_elements(Ptr meta, Ptr node) {
+  return ((StructMeta *)meta)->max_num_elements;
+}
+
+void Bitmasked_activate(Ptr meta, Ptr node, int i) {
+  auto smeta = (StructMeta *)meta;
+  auto element_size = StructMeta_get_element_size(smeta);
+  auto num_elements = Bitmasked_get_num_elements(meta, node);
+  auto data_section_size = element_size * num_elements;
+  auto mask_begin = (uint64 *)(node + data_section_size);
+  atomic_or_u64(&mask_begin[i / 64], 1UL << (i % 64));
+}
+
+i32 Bitmasked_is_active(Ptr meta, Ptr node, int i) {
+  auto smeta = (StructMeta *)meta;
+  auto element_size = StructMeta_get_element_size(smeta);
+  auto num_elements = Dense_get_num_elements(meta, node);
+  auto data_section_size = element_size * num_elements;
+  auto mask_begin = node + data_section_size;
+  return i32(bool((mask_begin[i / 8] >> (i % 8)) & 1));
+}
+
+Ptr Bitmasked_lookup_element(Ptr meta, Ptr node, int i) {
+  return node + ((StructMeta *)meta)->element_size * i;
+}
diff --git a/taichi/runtime/llvm/node_dense.h b/taichi/runtime/llvm/node_dense.h
index bf52cf3ebf1e24..4e14bc8ea2b6ee 100644
--- a/taichi/runtime/llvm/node_dense.h
+++ b/taichi/runtime/llvm/node_dense.h
@@ -2,11 +2,9 @@
 
 // Specialized Attributes and functions
 struct DenseMeta : public StructMeta {
-  bool bitmasked;
   int morton_dim;
 };
 
-STRUCT_FIELD(DenseMeta, bitmasked)
 STRUCT_FIELD(DenseMeta, morton_dim)
 
 i32 Dense_get_num_elements(Ptr meta, Ptr node) {
@@ -14,29 +12,11 @@ i32 Dense_get_num_elements(Ptr meta, Ptr node) {
 }
 
 void Dense_activate(Ptr meta, Ptr node, int i) {
-  auto smeta = (StructMeta *)meta;
-  auto dmeta = (DenseMeta *)meta;
-  if (DenseMeta_get_bitmasked(dmeta)) {
-    auto element_size = StructMeta_get_element_size(smeta);
-    auto num_elements = Dense_get_num_elements(meta, node);
-    auto data_section_size = element_size * num_elements;
-    auto mask_begin = (uint64 *)(node + data_section_size);
-    atomic_or_u64(&mask_begin[i / 64], 1UL << (i % 64));
-  }
+  // Dense elements are always active
 }
 
 i32 Dense_is_active(Ptr meta, Ptr node, int i) {
-  auto smeta = (StructMeta *)meta;
-  auto dmeta = (DenseMeta *)meta;
-  if (DenseMeta_get_bitmasked(dmeta)) {
-    auto element_size = StructMeta_get_element_size(smeta);
-    auto num_elements = Dense_get_num_elements(meta, node);
-    auto data_section_size = element_size * num_elements;
-    auto mask_begin = node + data_section_size;
-    return i32(bool((mask_begin[i / 8] >> (i % 8)) & 1));
-  } else {
-    return 1;
-  }
+  return 1;
 }
 
 Ptr Dense_lookup_element(Ptr meta, Ptr node, int i) {
diff --git a/taichi/runtime/llvm/runtime.cpp b/taichi/runtime/llvm/runtime.cpp
index 204bbfca8a6a75..5f9cdc0c7b0786 100644
--- a/taichi/runtime/llvm/runtime.cpp
+++ b/taichi/runtime/llvm/runtime.cpp
@@ -1085,6 +1085,7 @@ i32 linear_thread_idx() {
 #include "node_dynamic.h"
 #include "node_pointer.h"
 #include "node_root.h"
+#include "node_bitmasked.h"
 
 void ListManager::touch_chunk(int chunk_id) {
   if (!chunks[chunk_id]) {
diff --git a/taichi/struct/struct_llvm.cpp b/taichi/struct/struct_llvm.cpp
index 034c22d6164cbd..cc110761f5c0fa 100644
--- a/taichi/struct/struct_llvm.cpp
+++ b/taichi/struct/struct_llvm.cpp
@@ -40,10 +40,10 @@ void StructCompilerLLVM::generate_types(SNode &snode) {
   snode_attr[snode].llvm_element_type = ch_type;
 
   llvm::Type *body_type = nullptr, *aux_type = nullptr;
-  if (type == SNodeType::dense) {
+  if (type == SNodeType::dense || type == SNodeType::bitmasked) {
     TI_ASSERT(snode._morton == false);
     body_type = llvm::ArrayType::get(ch_type, snode.max_num_elements());
-    if (snode._bitmasked) {
+    if (type == SNodeType::bitmasked) {
       aux_type = llvm::ArrayType::get(llvm::Type::getInt32Ty(*llvm_ctx),
                                       (snode.max_num_elements() + 31) / 32);
     }
@@ -206,7 +206,7 @@ std::unique_ptr<StructCompiler> StructCompiler::make(Program *prog, Arch arch) {
 
 bool SNode::need_activation() const {
   return type == SNodeType::pointer || type == SNodeType::hash ||
-         (type == SNodeType::dense && _bitmasked) || type == SNodeType::dynamic;
+         type == SNodeType::bitmasked || type == SNodeType::dynamic;
 }
 
 TLANG_NAMESPACE_END
diff --git a/taichi/struct/struct_metal.cpp b/taichi/struct/struct_metal.cpp
index efcb27d9eabf39..f3e3574bdc63db 100644
--- a/taichi/struct/struct_metal.cpp
+++ b/taichi/struct/struct_metal.cpp
@@ -30,10 +30,6 @@ constexpr size_t kListManagerSize = sizeof(shaders::ListManager);
 constexpr size_t kSNodeMetaSize = sizeof(shaders::SNodeMeta);
 constexpr size_t kSNodeExtractorsSize = sizeof(shaders::SNodeExtractors);
 
-inline bool is_bitmasked(const SNode &sn) {
-  return (sn.type == SNodeType::dense && sn._bitmasked);
-}
-
 inline size_t bitmasks_stride(int n) {
   constexpr int kBitsPerByte = 8;
   const int bytes_needed = iroundup(n, kBitsPerByte) / kBitsPerByte;
@@ -41,6 +37,11 @@ inline size_t bitmasks_stride(int n) {
   return iroundup(bytes_needed, 8);
 }
 
+inline int get_n(const SNode &sn) {
+  // For root, sn.n is 0.
+  return sn.type == SNodeType::root ? 1 : sn.n;
+}
+
 class StructCompiler {
  public:
   StructCompiledResult run(SNode &root) {
@@ -54,7 +55,8 @@ class StructCompiler {
     {
       max_snodes_ = 0;
       for (const auto &sn : snodes_) {
-        if (sn->type == SNodeType::root || sn->type == SNodeType::dense) {
+        if (sn->type == SNodeType::root || sn->type == SNodeType::dense ||
+            sn->type == SNodeType::bitmasked) {
           max_snodes_ = std::max(max_snodes_, sn->id);
         }
       }
@@ -132,12 +134,13 @@ class StructCompiler {
       emit("  device {}* val;", dt_name);
       emit("}};");
     } else if (snode.type == SNodeType::dense ||
-               snode.type == SNodeType::root) {
-      const bool bitmasked = is_bitmasked(snode);
+               snode.type == SNodeType::root ||
+               snode.type == SNodeType::bitmasked) {
+      const bool bitmasked = snode.type == SNodeType::bitmasked;
       const std::string ch_name = fmt::format("{}_ch", node_name);
       emit("struct {} {{", node_name);
       emit("  // {}", snode_type_name(snode.type));
-      const int n = (snode.type == SNodeType::dense) ? snode.n : 1;
+      const int n = get_n(snode);
       emit("  constant static constexpr int n = {};", n);
       if (bitmasked) {
         emit(
@@ -170,7 +173,7 @@ class StructCompiler {
       return metal_data_type_bytes(to_metal_type(sn->dt));
     }
 
-    const int n = (sn->type == SNodeType::dense) ? sn->n : 1;
+    const int n = get_n(*sn);
     size_t ch_size = 0;
     for (const auto &ch : sn->ch) {
       const size_t ch_offset = ch_size;
@@ -186,7 +189,7 @@ class StructCompiler {
     sn_desc.element_stride = ch_size;
     sn_desc.num_slots = n;
     sn_desc.stride = ch_size * n;
-    if (is_bitmasked(*sn)) {
+    if (sn->type == SNodeType::bitmasked) {
       sn_desc.stride += bitmasks_stride(n);
     }
     sn_desc.total_num_elems_from_root = 1;
diff --git a/taichi/transforms/demote_dense_struct_fors.cpp b/taichi/transforms/demote_dense_struct_fors.cpp
index 04a725705aff82..73c27bbcf53b44 100644
--- a/taichi/transforms/demote_dense_struct_fors.cpp
+++ b/taichi/transforms/demote_dense_struct_fors.cpp
@@ -117,9 +117,6 @@ void demote_dense_struct_fors(IRNode *root) {
         if (snode->type != SNodeType::dense) {
           all_dense = false;
         }
-        if (snode->type == SNodeType::dense && snode->_bitmasked) {
-          all_dense = false;
-        }
         snode = snode->parent;
       }
       if (all_dense) {
diff --git a/taichi/transforms/lower_ast.cpp b/taichi/transforms/lower_ast.cpp
index a8185e45842112..4efaed77cd2a70 100644
--- a/taichi/transforms/lower_ast.cpp
+++ b/taichi/transforms/lower_ast.cpp
@@ -349,8 +349,7 @@ class LowerAST : public IRVisitor {
     } else if (stmt->snode->type == SNodeType::pointer ||
                stmt->snode->type == SNodeType::hash ||
                stmt->snode->type == SNodeType::dynamic ||
-               (stmt->snode->type == SNodeType::dense &&
-                stmt->snode->_bitmasked)) {
+               stmt->snode->type == SNodeType::bitmasked) {
       TI_ASSERT(SNodeOpStmt::activation_related(stmt->op_type));
       flattened.push_back<SNodeOpStmt>(stmt->op_type, stmt->snode,
                                        indices_stmt);
diff --git a/tests/python/test_bitmasked.py b/tests/python/test_bitmasked.py
index dc26bc9d9c515c..ae1c149a44fc3c 100644
--- a/tests/python/test_bitmasked.py
+++ b/tests/python/test_bitmasked.py
@@ -1,13 +1,17 @@
 import taichi as ti
 
 
-@ti.archs_excluding(ti.opengl)
+def archs_support_bitmasked(func):
+    return ti.archs_excluding(ti.opengl)(func)
+
+
+@archs_support_bitmasked
 def test_basic():
     x = ti.var(ti.i32)
     c = ti.var(ti.i32)
     s = ti.var(ti.i32)
 
-    bm = ti.root.dense(ti.ij, (3, 6)).bitmasked().dense(ti.i, 5).bitmasked()
+    bm = ti.root.bitmasked(ti.ij, (3, 6)).bitmasked(ti.i, 5)
     bm.place(x)
     ti.root.place(c, s)
 
@@ -28,3 +32,29 @@ def sum():
 
     assert c[None] == 3
     assert s[None] == 42
+
+
+@archs_support_bitmasked
+def test_bitmasked_then_dense():
+    x = ti.var(ti.f32)
+    s = ti.var(ti.i32)
+
+    n = 128
+
+    @ti.layout
+    def place():
+        ti.root.bitmasked(ti.i, n).dense(ti.i, n).place(x)
+        ti.root.place(s)
+
+    @ti.kernel
+    def func():
+        for i in x:
+            s[None] += 1
+
+    x[0] = 1
+    x[127] = 1
+    x[256] = 1
+    x[257] = 1
+
+    func()
+    assert s[None] == 256
diff --git a/tests/python/test_sparse_basics.py b/tests/python/test_sparse_basics.py
index ab89ba8cbfba23..c3800c31e5c56e 100644
--- a/tests/python/test_sparse_basics.py
+++ b/tests/python/test_sparse_basics.py
@@ -1,32 +1,6 @@
 import taichi as ti
 
 
-@ti.archs_support_sparse
-def test_bitmasked():
-    x = ti.var(ti.f32)
-    s = ti.var(ti.i32)
-
-    n = 128
-
-    @ti.layout
-    def place():
-        ti.root.dense(ti.i, n).bitmasked().dense(ti.i, n).place(x)
-        ti.root.place(s)
-
-    @ti.kernel
-    def func():
-        for i in x:
-            s[None] += 1
-
-    x[0] = 1
-    x[127] = 1
-    x[256] = 1
-    x[257] = 1
-
-    func()
-    assert s[None] == 256
-
-
 @ti.archs_support_sparse
 def test_pointer():
     x = ti.var(ti.f32)