[Metal] Changes to enable bitmasked on Metal!

python/taichi/lang/snode.py

@@ -72,6 +72,8 @@ def deactivate_all(self):
         for c in ch:
             c.deactivate_all()
         import taichi as ti
-        if self.ptr.type == ti.core.SNodeType.pointer:
+        if self.ptr.type == ti.core.SNodeType.pointer or (
+                self.ptr.type == ti.core.SNodeType.dense
+                and self.ptr.is_bitmasked):
             from .meta import snode_deactivate
             snode_deactivate(self)

taichi/ir/ir.h

@@ -2066,8 +2066,9 @@ class SNodeOpExpression : public Expression {
       // is_active cannot be lowered all the way to a global pointer.
       // It should be lowered into a pointer to parent and an index.
       TI_ERROR_IF(
-          snode->type != SNodeType::pointer && snode->type != SNodeType::hash,
-          "ti.is_active only works on hash and pointer nodes.");
+          snode->type != SNodeType::pointer && snode->type != SNodeType::hash &&
+              !(snode->type == SNodeType::dense && snode->_bitmasked),
+          "ti.is_active only works on pointer, hash or bitmasked nodes.");
       ret.push_back<SNodeOpStmt>(SNodeOpType::is_active, snode, indices_stmt);
     } else {
       auto ptr = ret.push_back<GlobalPtrStmt>(snode, indices_stmt);

taichi/platform/metal/metal_runtime.cpp

@@ -434,7 +434,7 @@ class MetalRuntime::Impl {
     params.profiler = profiler_;
     compiled_taichi_kernels_[taichi_kernel_name] =
         std::make_unique<CompiledTaichiKernel>(params);
-    TI_INFO("Registered Taichi kernel <{}>", taichi_kernel_name);
+    TI_DEBUG("Registered Taichi kernel <{}>", taichi_kernel_name);
   }
 
   void launch_taichi_kernel(const std::string &taichi_kernel_name,

taichi/python/export_lang.cpp

@@ -143,6 +143,7 @@ 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),

taichi/struct/struct_metal.cpp

@@ -73,7 +73,7 @@ class StructCompiler {
     result.runtime_size = compute_runtime_size();
     result.max_snodes = max_snodes_;
     result.snode_descriptors = std::move(snode_descriptors_);
-    TI_INFO("Metal: root_size={} runtime_size={}", result.root_size,
+    TI_DEBUG("Metal: root_size={} runtime_size={}", result.root_size,
             result.runtime_size);
     return result;
   }
@@ -220,13 +220,13 @@ class StructCompiler {
   size_t compute_runtime_size() {
     size_t result = (max_snodes_) *
                     (kSNodeMetaSize + kSNodeExtractorsSize + kListManagerSize);
-    TI_INFO("Metal runtime fields size: {} bytes", result);
+    TI_DEBUG("Metal runtime fields size: {} bytes", result);
     int total_items = 0;
     for (const auto &kv : snode_descriptors_) {
       total_items += kv.second.total_num_elems_from_root;
     }
     const size_t list_data_size = total_items * kListgenElementSize;
-    TI_INFO("Metal runtime list data size: {} bytes", list_data_size);
+    TI_DEBUG("Metal runtime list data size: {} bytes", list_data_size);
     result += list_data_size;
     return result;
   }

taichi/transforms/lower_ast.cpp

@@ -348,7 +348,9 @@ class LowerAST : public IRVisitor {
                                        val_stmt);
     } else if (stmt->snode->type == SNodeType::pointer ||
                stmt->snode->type == SNodeType::hash ||
-               stmt->snode->type == SNodeType::dynamic) {
+               stmt->snode->type == SNodeType::dynamic ||
+               (stmt->snode->type == SNodeType::dense &&
+                stmt->snode->_bitmasked)) {
       TI_ASSERT(SNodeOpStmt::activation_related(stmt->op_type));
       flattened.push_back<SNodeOpStmt>(stmt->op_type, stmt->snode,
                                        indices_stmt);

tests/python/test_bitmasked.py

@@ -0,0 +1,30 @@
+import taichi as ti
+
+
+@ti.archs_excluding(ti.opengl)
+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.place(x)
+    ti.root.place(c, s)
+
+    @ti.kernel
+    def run():
+        x[5, 1] = 2
+        x[9, 4] = 20
+        x[0, 3] = 20
+
+    @ti.kernel
+    def sum():
+        for i, j in x:
+            ti.atomic_add(c[None], ti.is_active(bm, [i, j]))
+            ti.atomic_add(s[None], x[i, j])
+
+    run()
+    sum()
+
+    assert c[None] == 3
+    assert s[None] == 42