[lang] Add SNode.snode_tree_id

python/taichi/lang/snode.py

@@ -215,6 +215,10 @@ def _id(self):
         """
         return self.ptr.id
 
+    @property
+    def _snode_tree_id(self):
+        return self.ptr.get_snode_tree_id()
+
     @property
     def shape(self):
         """Gets the number of elements from root in each axis of `self`.

taichi/python/export_lang.cpp

@@ -500,6 +500,7 @@ void export_lang(py::module &m) {
       .def_readwrite("parent", &SNode::parent)
       .def_readonly("type", &SNode::type)
       .def_readonly("id", &SNode::id)
+      .def("get_snode_tree_id", &SNode::get_snode_tree_id)
       .def_readonly("offset", &SNode::index_offsets)
       .def("dense",
            (SNode & (SNode::*)(const std::vector<Axis> &,

tests/python/test_snode.py

@@ -0,0 +1,21 @@
+import taichi as ti
+from tests import test_utils
+
+
+@ti.kernel
+def some_kernel(_: ti.template()): ...
+
+
+@test_utils.test(cpu_max_num_threads=1)
+def test_get_snode_tree_id():
+    s = ti.field(int, shape=())
+    some_kernel(s)
+    assert s.snode._snode_tree_id == 0
+
+    s = ti.field(int, shape=())
+    some_kernel(s)
+    assert s.snode._snode_tree_id == 1
+
+    s = ti.field(int, shape=())
+    some_kernel(s)
+    assert s.snode._snode_tree_id == 2