Add variables pruning pass

include/public/marco/Dialect/BaseModelica/Transforms/Passes.h

@@ -42,6 +42,7 @@
 #include "marco/Dialect/BaseModelica/Transforms/Scheduling.h"
 #include "marco/Dialect/BaseModelica/Transforms/SingleValuedInductionElimination.h"
 #include "marco/Dialect/BaseModelica/Transforms/VariablesPromotion.h"
+#include "marco/Dialect/BaseModelica/Transforms/VariablesPruning.h"
 #include "marco/Dialect/BaseModelica/Transforms/ViewAccessFolding.h"
 
 namespace mlir::bmodelica {

include/public/marco/Dialect/BaseModelica/Transforms/Passes.td

@@ -400,6 +400,28 @@ def VariablesPromotionPass
     let constructor = "mlir::bmodelica::createVariablesPromotionPass()";
 }
 
+def VariablesPruningPass : Pass<"variables-pruning", "mlir::ModuleOp">
+{
+    let summary = "Remove the unneeded variables.";
+
+    let description = [{
+        Remove the variables and the associated matched equations that are not
+        needed to compute the output variables.
+    }];
+
+    let dependentDialects = [
+        "mlir::bmodelica::BaseModelicaDialect"
+    ];
+
+    let options = [
+        Option<"outputVariables",
+               "output-variables", "std::string", "",
+               "Comma separated list of output variables">
+    ];
+
+    let constructor = "mlir::bmodelica::createVariablesPruningPass()";
+}
+
 //===---------------------------------------------------------------------===//
 // Model solving
 //===---------------------------------------------------------------------===//

include/public/marco/Dialect/BaseModelica/Transforms/VariablesPruning.h

@@ -0,0 +1,14 @@
+#ifndef MARCO_DIALECT_BASEMODELICA_TRANSFORMS_VARIABLESPRUNING_H
+#define MARCO_DIALECT_BASEMODELICA_TRANSFORMS_VARIABLESPRUNING_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir::bmodelica
+{
+#define GEN_PASS_DECL_VARIABLESPRUNINGPASS
+#include "marco/Dialect/BaseModelica/Transforms/Passes.h.inc"
+
+std::unique_ptr<mlir::Pass> createVariablesPruningPass();
+}
+
+#endif // MARCO_DIALECT_BASEMODELICA_TRANSFORMS_VARIABLESPRUNING_H

include/public/marco/Frontend/CodegenOptions.h

@@ -20,6 +20,7 @@ struct CodegenOptions : public clang::CodeGenOptions {
   bool outputArraysPromotion = false;
   bool heapToStackPromotion = false;
   bool readOnlyVariablesPropagation = false;
+  bool variablesPruning = false;
   bool variablesToParametersPromotion = false;
   int64_t sccSolvingBySubstitutionMaxIterations = 100;
   int64_t sccSolvingBySubstitutionMaxEquationsInSCC = 5;

include/public/marco/Modeling/ArrayEquationsDependencyGraph.h

@@ -124,6 +124,9 @@ class ArrayEquationsDependencyGraph {
     return result;
   }
 
+  /// Get the writes map of all the equations added to the graph.
+  const WritesMap &getWritesMap() const { return writesMap; }
+
   /// Map each array variable to the equations that write into some of its
   /// scalar positions.
   ///

include/public/marco/Modeling/SingleEntryDigraph.h

@@ -0,0 +1,156 @@
+#ifndef MARCO_MODELING_SINGLEENTRYDIGRAPH_H
+#define MARCO_MODELING_SINGLEENTRYDIGRAPH_H
+
+#include "marco/Modeling/Graph.h"
+#include "llvm/ADT/GraphTraits.h"
+
+namespace marco::modeling::dependency {
+template <typename VP, typename EP = internal::EmptyEdgeProperty>
+class SingleEntryDigraph {
+public:
+  using VertexProperty = VP;
+  using EdgeProperty = EP;
+
+private:
+  using Graph = internal::DirectedGraph<std::unique_ptr<VertexProperty>,
+                                        std::unique_ptr<EdgeProperty>>;
+
+public:
+  using VertexDescriptor = typename Graph::VertexDescriptor;
+  using EdgeDescriptor = typename Graph::EdgeDescriptor;
+
+  using VertexIterator = typename Graph::VertexIterator;
+  using IncidentEdgeIterator = typename Graph::IncidentEdgeIterator;
+  using LinkedVerticesIterator = typename Graph::LinkedVerticesIterator;
+
+  SingleEntryDigraph() : entryNode(graph.addVertex(nullptr)) {}
+
+  VertexProperty &operator[](VertexDescriptor vertex) {
+    assert(vertex != entryNode && "The entry node doesn't have a property");
+    return *graph[vertex];
+  }
+
+  const VertexProperty &operator[](VertexDescriptor vertex) const {
+    assert(vertex != entryNode && "The entry node doesn't have a property");
+    return *graph[vertex];
+  }
+
+  EdgeProperty &operator[](EdgeDescriptor edge) { return *graph[edge]; }
+
+  const EdgeProperty &operator[](EdgeDescriptor edge) const {
+    return *graph[edge];
+  }
+
+  size_t size() const { return graph.verticesCount(); }
+
+  VertexDescriptor getEntryNode() const { return entryNode; }
+
+  VertexDescriptor addVertex(VertexProperty property) {
+    auto descriptor =
+        graph.addVertex(std::make_unique<VertexProperty>(std::move(property)));
+
+    return descriptor;
+  }
+
+  auto verticesBegin() const {
+    return graph.verticesBegin(
+        [](const typename Graph::VertexProperty &vertex) {
+          // Hide the entry point
+          return vertex != nullptr;
+        });
+  }
+
+  auto verticesEnd() const {
+    return graph.verticesEnd([](const typename Graph::VertexProperty &vertex) {
+      // Hide the entry point
+      return vertex != nullptr;
+    });
+  }
+
+  EdgeDescriptor addEdge(VertexDescriptor from, VertexDescriptor to,
+                         EdgeProperty property = EdgeProperty()) {
+    return graph.addEdge(from, to,
+                         std::make_unique<EdgeProperty>(std::move(property)));
+  }
+
+  auto getEdges() const { return graph.getEdges(); }
+
+  auto outgoingEdgesBegin(VertexDescriptor vertex) const {
+    return graph.outgoingEdgesBegin(std::move(vertex));
+  }
+
+  auto outgoingEdgesEnd(VertexDescriptor vertex) const {
+    return graph.outgoingEdgesEnd(std::move(vertex));
+  }
+
+  auto linkedVerticesBegin(VertexDescriptor vertex) const {
+    return graph.linkedVerticesBegin(std::move(vertex));
+  }
+
+  auto linkedVerticesEnd(VertexDescriptor vertex) const {
+    return graph.linkedVerticesEnd(std::move(vertex));
+  }
+
+private:
+  Graph graph;
+  VertexDescriptor entryNode;
+};
+} // namespace marco::modeling::dependency
+
+namespace llvm {
+// We specialize the LLVM's graph traits in order leverage the algorithms
+// that are defined inside LLVM itself. This way we don't have to implement
+// them from scratch.
+template <typename VertexProperty, typename EdgeProperty>
+struct GraphTraits<const marco::modeling::dependency ::SingleEntryDigraph<
+    VertexProperty, EdgeProperty> *> {
+  // The LLVM traits require the class specified as Graph to be copyable.
+  // We use its address to overcome this limitation.
+  using Graph =
+      const marco::modeling::dependency::SingleEntryDigraph<VertexProperty>;
+
+  using GraphPtr = Graph *;
+
+  using NodeRef = typename Graph::VertexDescriptor;
+  using ChildIteratorType = typename Graph::LinkedVerticesIterator;
+
+  static NodeRef getEntryNode(const GraphPtr &graph) {
+    return graph->getEntryNode();
+  }
+
+  static ChildIteratorType child_begin(NodeRef node) {
+    return node.graph->linkedVerticesBegin(node);
+  }
+
+  static ChildIteratorType child_end(NodeRef node) {
+    return node.graph->linkedVerticesEnd(node);
+  }
+
+  using nodes_iterator = typename Graph::VertexIterator;
+
+  static nodes_iterator nodes_begin(GraphPtr *graph) {
+    return (*graph)->verticesBegin();
+  }
+
+  static nodes_iterator nodes_end(GraphPtr *graph) {
+    return (*graph)->verticesEnd();
+  }
+
+  using EdgeRef = typename Graph::EdgeDescriptor;
+  using ChildEdgeIteratorType = typename Graph::IncidentEdgeIterator;
+
+  static ChildEdgeIteratorType child_edge_begin(NodeRef node) {
+    return node.graph->outgoingEdgesBegin(node);
+  }
+
+  static ChildEdgeIteratorType child_edge_end(NodeRef node) {
+    return node.graph->outgoingEdgesEnd(node);
+  }
+
+  static NodeRef edge_dest(EdgeRef edge) { return edge.to; }
+
+  static size_t size(GraphPtr *graph) { return (*graph)->size(); }
+};
+} // namespace llvm
+
+#endif // MARCO_MODELING_SINGLEENTRYDIGRAPH_H

lib/Dialect/BaseModelica/Transforms/CMakeLists.txt

@@ -56,6 +56,7 @@ add_mlir_dialect_library(MLIRBaseModelicaTransforms
     Scheduling.cpp
     SingleValuedInductionElimination.cpp
     VariablesPromotion.cpp
+    VariablesPruning.cpp
     VectorizableOpInterfaceImpl.cpp
     ViewAccessFolding.cpp