Add VTK procedures to runtime

Signed-off-by: Rafael Ravedutti <rafaelravedutti@gmail.com>

code_gen/cgen.py

@@ -11,13 +11,12 @@ from ir.lit import Lit
 from ir.loops import For, Iter, ParticleFor, While
 from ir.math import Ceil, Sqrt
 from ir.memory import Malloc, Realloc
-from ir.properties import Property, PropertyList, RegisterProperty
+from ir.properties import Property, PropertyList, RegisterProperty, UpdateProperty
 from ir.select import Select
 from ir.sizeof import Sizeof
 from ir.utils import Print
 from ir.variables import Var, VarDecl
 from sim.timestep import Timestep
-from sim.vtk import VTKWrite
 from code_gen.printer import Printer
 
 
@@ -47,6 +46,7 @@ class CGen:
         self.print("//---")
         self.print("#include \"runtime/pairs.hpp\"")
         self.print("#include \"runtime/read_from_file.hpp\"")
+        self.print("#include \"runtime/vtk.hpp\"")
         self.print("")
         self.print("using namespace pairs;")
         self.print("")
@@ -113,7 +113,7 @@ class CGen:
 
         if isinstance(ast_node, Call):
             call = self.generate_expression(ast_node)
-            self.print("{call};")
+            self.print(f"{call};")
 
         if isinstance(ast_node, For):
             iterator = self.generate_expression(ast_node.iterator)
@@ -176,18 +176,19 @@ class CGen:
         if isinstance(ast_node, Timestep):
             self.generate_statement(ast_node.block)
 
+        if isinstance(ast_node, UpdateProperty):
+            p = ast_node.property()
+
+            if p.type() != Type_Vector or p.layout() == Layout_Invalid:
+                self.print(f"ps->updateProperty({p.id()}, {p.name()});")
+            else:
+                sizes = ", ".join([str(self.generate_expression(size)) for size in ast_node.sizes()])
+                self.print(f"ps->updateProperty({p.id()}, {p.name()}, {sizes});")
+
         if isinstance(ast_node, VarDecl):
             tkw = CGen.type2keyword(ast_node.var.type())
             self.print(f"{tkw} {ast_node.var.name()} = {ast_node.var.init_value()};")
 
-        if isinstance(ast_node, VTKWrite):
-            nlocal = self.generate_expression(self.sim.nlocal)
-            npbc = self.generate_expression(self.sim.pbc.npbc)
-            nall = self.generate_expression(self.sim.nlocal + self.sim.pbc.npbc)
-            timestep = self.generate_expression(ast_node.timestep)
-            self.generate_vtk_writing(ast_node.vtk_id * 2, f"{ast_node.filename}_local", 0, nlocal, nlocal, timestep)
-            self.generate_vtk_writing(ast_node.vtk_id * 2 + 1, f"{ast_node.filename}_pbc", nlocal, nall, npbc, timestep)
-
         if isinstance(ast_node, While):
             cond = self.generate_expression(ast_node.cond)
             self.print(f"while({cond}) {{")
@@ -259,7 +260,6 @@ class CGen:
 
         if isinstance(ast_node, Lit):
             assert mem is False, "Literal is not lvalue!"
-
             if ast_node.type() == Type_String:
                 return f"\"{ast_node.value}\""
 
@@ -295,59 +295,3 @@ class CGen:
 
         if isinstance(ast_node, Var):
             return ast_node.name()
-
-    def generate_vtk_writing(self, id, filename, start, end, n, timestep):
-        # TODO: Do this in a more elegant way, without hard coded stuff
-        header = "# vtk DataFile Version 2.0\n" \
-                 "Particle data\n" \
-                 "ASCII\n" \
-                 "DATASET UNSTRUCTURED_GRID\n"
-
-        filename_var = f"filename{id}"
-        filehandle_var = f"vtk{id}"
-        self.print(f"char {filename_var}[128];")
-        self.print(f"snprintf({filename_var}, sizeof {filename_var}, \"{filename}_%d.vtk\", {timestep});")
-        self.print(f"FILE *{filehandle_var} = fopen({filename_var}, \"w\");")
-        for line in header.split('\n'):
-            if len(line) > 0:
-                self.print(f"fwrite(\"{line}\\n\", 1, {len(line) + 1}, {filehandle_var});")
-
-        # Write positions
-        self.print(f"fprintf({filehandle_var}, \"POINTS %d double\\n\", {n});")
-        self.print(f"for(int i = {start}; i < {end}; i++) {{")
-        self.print.add_ind(4)
-        self.print(f"fprintf({filehandle_var}, \"%.4f %.4f %.4f\\n\", position[i * 3], position[i * 3 + 1], position[i * 3 + 2]);")
-        self.print.add_ind(-4)
-        self.print("}")
-        self.print(f"fwrite(\"\\n\\n\", 1, 2, {filehandle_var});")
-
-        # Write cells
-        self.print(f"fprintf({filehandle_var}, \"CELLS %d %d\\n\", {n}, {n} * 2);")
-        self.print(f"for(int i = {start}; i < {end}; i++) {{")
-        self.print.add_ind(4)
-        self.print(f"fprintf({filehandle_var}, \"1 %d\\n\", i - {start});")
-        self.print.add_ind(-4)
-        self.print("}")
-        self.print(f"fwrite(\"\\n\\n\", 1, 2, {filehandle_var});")
-
-        # Write cell types
-        self.print(f"fprintf({filehandle_var}, \"CELL_TYPES %d\\n\", {n});")
-        self.print(f"for(int i = {start}; i < {end}; i++) {{")
-        self.print.add_ind(4)
-        self.print(f"fwrite(\"1\\n\", 1, 2, {filehandle_var});")
-        self.print.add_ind(-4)
-        self.print("}")
-        self.print(f"fwrite(\"\\n\\n\", 1, 2, {filehandle_var});")
-
-        # Write masses
-        self.print(f"fprintf({filehandle_var}, \"POINT_DATA %d\\n\", {n});")
-        self.print(f"fprintf({filehandle_var}, \"SCALARS mass double\\n\");")
-        self.print(f"fprintf({filehandle_var}, \"LOOKUP_TABLE default\\n\");")
-        self.print(f"for(int i = {start}; i < {end}; i++) {{")
-        self.print.add_ind(4)
-        #self.print(f"fprintf({filehandle_var}, \"%4.f\\n\", mass[i]);")
-        self.print(f"fprintf({filehandle_var}, \"1.0\\n\");")
-        self.print.add_ind(-4)
-        self.print("}")
-        self.print(f"fwrite(\"\\n\\n\", 1, 2, {filehandle_var});")
-        self.print(f"fclose({filehandle_var});")

ir/functions.py

 from ir.bin_op import ASTTerm
-from ir.data_types import Type_Int
+from ir.data_types import Type_Int, Type_Invalid
 from ir.lit import as_lit_ast
 
 
@@ -19,8 +19,15 @@ class Call(ASTTerm):
     def type(self):
         return self.return_type
 
+    def children(self):
+        return self.params
 
 class Call_Int(Call):
     def __init__(self, sim, func_name, parameters):
         super().__init__(sim, func_name, parameters, Type_Int)
-        
+
+
+class Call_Void(Call):
+    def __init__(self, sim, func_name, parameters):
+        super().__init__(sim, func_name, parameters, Type_Invalid)
+        sim.add_statement(self)

ir/properties.py

@@ -114,4 +114,21 @@ class RegisterProperty(ASTNode):
         return self.sizes_list
 
     def __str__(self):
-        return f"Property<{self.prop.name()}>"
+        return f"RegisterProperty<{self.prop.name()}>"
+
+
+class UpdateProperty(ASTNode):
+    def __init__(self, sim, prop, sizes):
+        super().__init__(sim)
+        self.prop = prop
+        self.sizes_list = [as_lit_ast(sim, s) for s in sizes]
+        self.sim.add_statement(self)
+
+    def property(self):
+        return self.prop
+
+    def sizes(self):
+        return self.sizes_list
+
+    def __str__(self):
+        return f"UpdateProperty<{self.prop.name()}>"

runtime/pairs.hpp

@@ -70,6 +70,7 @@ public:
     std::string getName() { return name; }
     void *getPointer() { return ptr; }
     void setPointer(void *ptr_) { ptr = ptr_; }
+    void setSizes(int sx_, int sy_) { sx = sx_, sy = sy_; }
     PropertyType getType() { return type; }
     layout_t getLayout() { return layout; }
 };
@@ -103,6 +104,15 @@ private:
 public:
     PairsSim() {}
     void addProperty(Property prop) { properties.push_back(prop); }
+    void updateProperty(property_t id, void *ptr, int sx = 0, int sy = 0) {
+        auto p = std::find_if(properties.begin(), properties.end(), [id](Property p) {
+            return p.getId() == id;
+        });
+
+        PAIRS_ASSERT(p != std::end(properties));
+        p->setPointer(ptr);
+        p->setSizes(sx, sy);
+    }
 
     Property &getProperty(property_t id) {
         auto p = std::find_if(properties.begin(), properties.end(), [id](Property p) {

runtime/vtk.hpp

+#include <iomanip>
+#include <iostream>
+#include <fstream>
+//---
+#include "pairs.hpp"
+
+#pragma once
+
+namespace pairs {
+
+void vtk_write_data(PairsSim *ps, const char *filename, int start, int end, int timestep) {
+    std::string output_filename(filename);
+    std::ostringstream filename_oss;
+    filename_oss << filename << "_" << timestep << ".vtk";
+    std::ofstream out_file(filename_oss.str());
+    auto masses = ps->getAsFloatProperty(ps->getPropertyByName("mass"));
+    auto positions = ps->getAsVectorProperty(ps->getPropertyByName("position"));
+    const int n = end - start;
+
+    if(out_file.is_open()) {
+        out_file << "# vtk DataFile Version 2.0\n";
+        out_file << "Particle data\n";
+        out_file << "ASCII\n";
+        out_file << "DATASET UNSTRUCTURED_GRID\n";
+        out_file << "POINTS " << n << " double\n";
+
+        for(int i = start; i < end; i++) {
+            out_file << std::fixed << std::setprecision(4) << positions(i, 0) << " ";
+            out_file << std::fixed << std::setprecision(4) << positions(i, 1) << " ";
+            out_file << std::fixed << std::setprecision(4) << positions(i, 2) << "\n";
+        }
+
+        out_file << "\n\n";
+        out_file << "CELLS " << n << " " << (n * 2) << "\n";
+        for(int i = start; i < end; i++) {
+            out_file << "1 " << (i - start) << "\n";
+        }
+
+        out_file << "\n\n";
+        out_file << "CELL_TYPES " << n << "\n";
+        for(int i = start; i < end; i++) {
+            out_file << "1\n";
+        }
+
+        out_file << "\n\n";
+        out_file << "POINT_DATA " << n << "\n";
+        out_file << "SCALARS mass double\n";
+        out_file << "LOOKUP_TABLE default\n";
+        for(int i = start; i < end; i++) {
+            out_file << std::fixed << std::setprecision(4) << masses(i) << "\n";
+        }
+
+        out_file << "\n\n";
+        out_file.close();
+    }
+}
+
+}

sim/particle_simulation.py

@@ -192,12 +192,12 @@ class ParticleSimulation:
             self.kernels.lower()
         ])
 
-        timestep.add(Block(self, VTKWrite(self, self.vtk_file, timestep.timestep() + 1)))
+        timestep.add(VTKWrite(self, self.vtk_file, timestep.timestep() + 1).lower())
 
         body = Block.from_list(self, [
             self.setups.lower(),
             CellListsStencilBuild(self.cell_lists).lower(),
-            Block(self, VTKWrite(self, self.vtk_file, 0)),
+            VTKWrite(self, self.vtk_file, 0).lower(),
             timestep.as_block()
         ])
 

sim/properties.py

@@ -2,7 +2,7 @@ from functools import reduce
 from ir.data_types import Type_Float, Type_Vector
 from ir.loops import ParticleFor
 from ir.memory import Malloc, Realloc
-from ir.properties import RegisterProperty
+from ir.properties import RegisterProperty, UpdateProperty
 from ir.utils import Print
 import operator
 
@@ -27,6 +27,7 @@ class PropertiesAlloc:
 
             if self.realloc:
                 Realloc(self.sim, p, reduce(operator.mul, sizes))
+                UpdateProperty(self.sim, p, sizes)
             else:
                 Malloc(self.sim, p, reduce(operator.mul, sizes), True)
                 RegisterProperty(self.sim, p, sizes)

sim/resize.py

+from functools import reduce
 from ir.branches import Filter
 from ir.data_types import Type_Int, Type_Float, Type_Vector
 from ir.loops import While
 from ir.memory import Realloc
+from ir.properties import UpdateProperty
 from ir.utils import Print
+import operator
+
 
 class Resize:
     def __init__(self, sim, capacity_var, grow_fn=None):
@@ -27,8 +31,9 @@ class Resize:
                     capacity = sum(self.sim.properties.capacities)
                     for p in properties.all():
                         if p.type() == Type_Vector:
-                            sizes = capacity * self.sim.ndims()
+                            sizes = [capacity, self.sim.ndims()]
                         else:
-                            sizes = capacity
+                            sizes = [capacity]
 
-                        Realloc(self.sim, p, sizes)
+                        Realloc(self.sim, p, reduce(operator.mul, sizes))
+                        UpdateProperty(self.sim, p, sizes)

sim/vtk.py

-from ir.lit import as_lit_ast
 from ir.ast_node import ASTNode
+from ir.functions import Call_Void
+from ir.lit import as_lit_ast
 
 
 class VTKWrite(ASTNode):
-    vtk_id = 0
-
     def __init__(self, sim, filename, timestep):
         super().__init__(sim)
-        self.vtk_id = VTKWrite.vtk_id
         self.filename = filename
         self.timestep = as_lit_ast(sim, timestep)
-        VTKWrite.vtk_id += 1
+
+    def lower(self):
+        nlocal = self.sim.nlocal
+        npbc = self.sim.pbc.npbc
+        self.sim.clear_block()
+        nall = nlocal + npbc
+        Call_Void(self.sim, "pairs::vtk_write_data", [self.filename + "_local", 0, nlocal, self.timestep])
+        Call_Void(self.sim, "pairs::vtk_write_data", [self.filename + "_pbc", nlocal, nall, self.timestep])
+        return self.sim.block
 
     def children(self):
         return [self.timestep]