[lang] Expose SNode ID to python (#2162)

taichi-dev · Jan 17, 2021 · c728d38 · c728d38
1 parent 4e363f0
commit c728d38
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 0 deletions.
diff --git a/python/taichi/lang/snode.py b/python/taichi/lang/snode.py
@@ -99,6 +99,10 @@ def data_type(self):
     def dim(self):
         return len(self.shape)
 
+    @property
+    def id(self):
+        return self.ptr.id
+
     @property
     def shape(self):
         impl.get_runtime().materialize()

diff --git a/python/taichi/misc/util.py b/python/taichi/misc/util.py
@@ -280,6 +280,9 @@ def print_async_stats(include_kernel_profiler=False):
     print(f'Tasks launched:       {int(counters["launched_tasks"])}')
     print(f'Instructions emitted: {int(counters["codegen_statements"])}')
     print(f'Tasks compiled:       {int(counters["codegen_offloaded_tasks"])}')
+    NUM_FUSED_TASKS_KEY = 'num_fused_tasks'
+    if NUM_FUSED_TASKS_KEY in counters:
+        print(f'Tasks fused:          {int(counters["num_fused_tasks"])}')
     print('=======================')
 
 

diff --git a/taichi/program/state_flow_graph.cpp b/taichi/program/state_flow_graph.cpp
@@ -626,6 +626,7 @@ std::unordered_set<int> StateFlowGraph::fuse_range(int begin, int end) {
     auto *node_b = nodes[b];
     TI_TRACE("Fuse: nodes[{}]({}) <- nodes[{}]({})", a, node_a->string(), b,
              node_b->string());
+    stat.add("num_fused_tasks");
     auto &rec_a = node_a->rec;
     auto &rec_b = node_b->rec;
     rec_a.ir_handle =

diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
@@ -220,6 +220,7 @@ void export_lang(py::module &m) {
       .def(py::init<>())
       .def_readwrite("parent", &SNode::parent)
       .def_readonly("type", &SNode::type)
+      .def_readonly("id", &SNode::id)
       .def("dense",
            (SNode & (SNode::*)(const std::vector<Index> &,
                                const std::vector<int> &))(&SNode::dense),