Fix how num_elements is computed in Dense node (#512)

taichi-dev · Feb 22, 2020 · 0cc4f3c · 0cc4f3c
1 parent 73cb534
commit 0cc4f3c
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/taichi/runtime/node_dense.h b/taichi/runtime/node_dense.h
@@ -9,12 +9,16 @@ struct DenseMeta : public StructMeta {
 STRUCT_FIELD(DenseMeta, bitmasked)
 STRUCT_FIELD(DenseMeta, morton_dim)
 
+i32 Dense_get_num_elements(Ptr meta, Ptr node) {
+  return ((StructMeta *)meta)->max_num_elements;
+}
+
 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 = 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));
@@ -26,7 +30,7 @@ i32 Dense_is_active(Ptr meta, Ptr node, int i) {
   auto dmeta = (DenseMeta *)meta;
   if (DenseMeta_get_bitmasked(dmeta)) {
     auto element_size = StructMeta_get_element_size(smeta);
-    auto num_elements = 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));
@@ -38,7 +42,3 @@ i32 Dense_is_active(Ptr meta, Ptr node, int i) {
 Ptr Dense_lookup_element(Ptr meta, Ptr node, int i) {
   return node + ((StructMeta *)meta)->element_size * i;
 }
-
-i32 Dense_get_num_elements(Ptr meta, Ptr node) {
-  return ((StructMeta *)meta)->max_num_elements;
-}