Merge pull request #24 from edin-dal/feature/selective-iteration

Selective variable ordering added
edin-dal · Jul 19, 2024 · 33cc294 · 33cc294
2 parents cb78ba1 + 9d891f4
commit 33cc294
Show file tree

Hide file tree

Showing 13 changed files with 309 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -180,6 +180,32 @@ for (int j = 0; j < N; ++j) {
 }
 ```
 
+#### Variable Ordering
+
+You can provide arbitrary variable ordering and include extra non-symbol and unbound variables through a `;` separated mapping named `iters` by providing the mapping from the left-hand side tensor names to a comma separated list of variables. For example:
+
+```
+symbols: N, M, ITERS
+iters: covar -> iter, i, j, k
+covar(j, k) := X(i, j) * X(i, k) * (0 <= iter < ITERS)
+covar:D(i, j) := (0 <= i < N) * (0 <= j < N)
+X:D(i, j) := (0 <= i < M) * (0 <= j < N)
+```
+
+This will generate the following code:
+
+```c++
+for (int iter = 0; iter < ITERS; ++iter) {
+  for (int i = 0; i < M; ++i) {
+    for (int j = 0; j < N; ++j) {
+      for (int k = 0; k < min({(j) + 1, N}); ++k) {
+        covar[j][k] += (X[i][k] * X[i][j]);
+      }
+    }
+  }
+}
+```
+
 ### Generating C++ Code
 
 To generate code for one of the existing `.stur` codes in the `examples/` directory, use the following command:

diff --git a/examples/including-extra-variable.stur b/examples/including-extra-variable.stur
@@ -0,0 +1,5 @@
+symbols: N, M
+iters: covar -> iter, i, j, k
+covar(j, k) := X(i, j) * X(i, k) * (0 <= iter < 1000)
+covar:D(i, j) := (0 <= i < N) * (0 <= j < N)
+X:D(i, j) := (0 <= i < M) * (0 <= j < N)
diff --git a/examples/partial-variable-ordering.stur b/examples/partial-variable-ordering.stur
@@ -0,0 +1,5 @@
+symbols: N, M
+iters: covar -> i
+covar(j, k) := X(i, j) * X(i, k)
+covar:D(i, j) := (0 <= i < N) * (0 <= j < N)
+X:D(i, j) := (0 <= i < M) * (0 <= j < N)
diff --git a/examples/variable-ordering.stur b/examples/variable-ordering.stur
@@ -0,0 +1,5 @@
+symbols: N, M
+iters: covar -> i, j, k
+covar(j, k) := X(i, j) * X(i, k)
+covar:D(i, j) := (0 <= i < N) * (0 <= j < N)
+X:D(i, j) := (0 <= i < M) * (0 <= j < N)
diff --git a/src/main/scala/uk/ac/ed/dal/structtensor/Main.scala b/src/main/scala/uk/ac/ed/dal/structtensor/Main.scala
@@ -115,16 +115,27 @@ object Main extends App {
           lineSeqInit.zipWithIndex.filter(_._1.startsWith("symbols:")).unzip
         val (outputs_lines, outputs_index) =
           lineSeqInit.zipWithIndex.filter(_._1.startsWith("outputs:")).unzip
+        val (iters_lines, iters_index) =
+          lineSeqInit.zipWithIndex.filter(_._1.startsWith("iters:")).unzip
         val symbols = symbols_lines
           .map(e => e.slice(8, e.length))
           .flatMap(_.split(",").map(_.trim).toSeq)
           .map(Variable(_))
         val outputs_names = outputs_lines
           .map(e => e.slice(8, e.length))
           .flatMap(_.split(",").map(_.trim).toSeq)
+        val iters_map = iters_lines
+          .map(e => e.slice(6, e.length))
+          .flatMap(_.split(";").map(_.trim).toSeq)
+          .map(iter_str =>
+            fastparse.parse(iter_str, Parser.iterators(_)).get.value
+          )
+          .toMap
         val lineSeq = lineSeqInit.zipWithIndex
           .filterNot(x =>
-            symbols_index.contains(x._2) || outputs_index.contains(x._2)
+            symbols_index.contains(x._2) ||
+              outputs_index.contains(x._2) ||
+              iters_index.contains(x._2)
           )
           .map(_._1)
         val preprocess_start_index = lineSeq.indexOf("@preprocess_start")
@@ -184,8 +195,9 @@ object Main extends App {
             )
           )((acc, tc) => {
             val inps = getInputs(tc, acc._1, acc._2, acc._3)
+            val iters = iters_map.getOrElse(tc.head.name, Seq())
             val (usRule, rmRule, ccRule) =
-              compile(tc, inps, symbols, outputs_names)
+              compile(tc, inps, symbols, outputs_names, iters)
             val rcRule = Rule(
               ccRule.head,
               SoPTimesSoP(
@@ -226,7 +238,9 @@ object Main extends App {
           )
         )((acc, tc) => {
           val inps = getInputs(tc, acc._1, acc._2, acc._3)
-          val (usRule, rmRule, ccRule) = compile(tc, inps, symbols)
+          val iters = iters_map.getOrElse(tc.head.name, Seq())
+          val (usRule, rmRule, ccRule) =
+            compile(tc, inps, symbols, iters = iters)
           (
             acc._1 + (usRule.head -> usRule),
             acc._2 + (rmRule.head -> rmRule),
@@ -236,17 +250,33 @@ object Main extends App {
         })
 
         val preprocessComputation = ccRuleSeq_preprocess
-          .map(r => Codegen(r, symbols, config.codeLang, Tensor))
+          .map(r => Codegen(r, symbols, config.codeLang, Tensor, iters_map))
           .mkString("\n")
         val ccComputation = outputs_names.isEmpty match {
           case true =>
             ccRuleSeq
-              .map(r => Codegen(r, symbols, config.codeLang, CompressedTensor))
+              .map(r =>
+                Codegen(
+                  r,
+                  symbols,
+                  config.codeLang,
+                  CompressedTensor,
+                  iters_map
+                )
+              )
               .mkString("\n")
           case false =>
             ccRuleSeq
               .filter(r => outputs_names.contains(r.head.name))
-              .map(r => Codegen(r, symbols, config.codeLang, CompressedTensor))
+              .map(r =>
+                Codegen(
+                  r,
+                  symbols,
+                  config.codeLang,
+                  CompressedTensor,
+                  iters_map
+                )
+              )
               .mkString("\n")
         }
 

diff --git a/src/main/scala/uk/ac/ed/dal/structtensor/codegen/Codegen.scala b/src/main/scala/uk/ac/ed/dal/structtensor/codegen/Codegen.scala
@@ -183,10 +183,12 @@ object Codegen {
       accesses: Seq[Access],
       symbols: Seq[Variable],
       codeLang: String,
-      kind: AccessType
+      kind: AccessType,
+      iters: Seq[Variable]
   ): String = {
     val variablesInit =
-      (computationHead.vars ++ accesses.flatMap(_.vars)).distinct
+      (iters ++ computationHead.vars ++ accesses.flatMap(_.vars)).distinct
+
     val variables = reorder(variablesInit, conditions, symbols)
     val (loopNests, restOfConditions, numberOfBrackets1) = variables.reverse.foldLeft(
       Seq[String](),
@@ -583,9 +585,11 @@ object Codegen {
       rule: Rule,
       symbols: Seq[Variable],
       codeLang: String,
-      kind: AccessType
+      kind: AccessType,
+      iters_map: Map[String, Seq[Variable]] = Map()
   ): String = {
     val computationHead = rule.head
+    val iters = iters_map.getOrElse(computationHead.name, Seq())
     rule.body.prods.zipWithIndex
       .map {
         case (prod, ind) => {
@@ -598,7 +602,8 @@ object Codegen {
             accesses,
             symbols,
             codeLang,
-            kind
+            kind,
+            iters
           )
         }
       }

diff --git a/src/main/scala/uk/ac/ed/dal/structtensor/compiler/Compiler.scala b/src/main/scala/uk/ac/ed/dal/structtensor/compiler/Compiler.scala
@@ -1189,7 +1189,8 @@ object Compiler {
       us: Rule,
       rm: Rule,
       cc: Rule,
-      symbols: Seq[Variable]
+      symbols: Seq[Variable],
+      iters: Seq[Variable]
   ): (Rule, Rule, Rule) = {
     val (idempotentOptUS, idempotentOptRM, idempotentOptCC) =
       (setIdempotentOpt(us), setIdempotentOpt(rm), setIdempotentOpt(cc))
@@ -1205,7 +1206,7 @@ object Compiler {
     ) = (
       replaceEqualVariables(removeEmptyProdOptUS, symbols),
       replaceEqualVariables(removeEmptyProdOptRM, symbols),
-      replaceEqualVariables(removeEmptyProdOptCC, symbols)
+      replaceEqualVariables(removeEmptyProdOptCC, symbols, iters)
     )
     if (
       isSoPEquals(replacedEqualVariablesOptUS.body, us.body) && isSoPEquals(
@@ -1223,15 +1224,17 @@ object Compiler {
         replacedEqualVariablesOptUS,
         replacedEqualVariablesOptRM,
         replacedEqualVariablesOptCC,
-        symbols
+        symbols,
+        iters
       )
   }
 
   def compile(
       computation: Rule,
       inputs: Seq[(Rule, Rule, Rule, Rule)],
       symbols: Seq[Variable],
-      outputs_names: Seq[String] = Seq()
+      outputs_names: Seq[String] = Seq(),
+      iters: Seq[Variable] = Seq()
   ): (Rule, Rule, Rule) = {
     val norm = normalize(computation)
     val us_rm_cc_tc_seq = norm.foldLeft(inputs)((ctx, r) => {
@@ -1259,13 +1262,14 @@ object Compiler {
     ) = (
       replaceEqualVariables(removeEmptyProdOptUS, symbols),
       replaceEqualVariables(removeEmptyProdOptRM, symbols),
-      replaceEqualVariables(removeEmptyProdOptCC, symbols)
+      replaceEqualVariables(removeEmptyProdOptCC, symbols, iters)
     )
     val (fixedUS, fixedRM, fixedCC) = fixedPointOpt(
       replacedEqualVariablesOptUS,
       replacedEqualVariablesOptRM,
       replacedEqualVariablesOptCC,
-      symbols
+      symbols,
+      iters
     )
 
     (fixedUS, fixedRM, fixedCC)

diff --git a/src/main/scala/uk/ac/ed/dal/structtensor/compiler/Optimizer.scala b/src/main/scala/uk/ac/ed/dal/structtensor/compiler/Optimizer.scala
@@ -584,7 +584,11 @@ object Optimizer {
     case _ => throw new Exception("Unknown expression")
   }
 
-  def replaceEqualVariables(rule: Rule, symbols: Seq[Variable]): Rule = {
+  def replaceEqualVariables(
+      rule: Rule,
+      symbols: Seq[Variable],
+      iters: Seq[Variable] = Seq()
+  ): Rule = {
     val equalVariablesSet = getEqualVariables(rule)
     val newBody = SoP(rule.body.prods.zip(equalVariablesSet).map {
       case (p, eSet) => {
@@ -593,13 +597,12 @@ object Optimizer {
     })
 
     val base_variables =
-      symbols.map(_.name) ++ rule.head.vars.map(_.name).distinct
+      (iters ++ symbols ++ rule.head.vars).map(_.name).distinct
     val finalBody = SoP(newBody.prods.map(p => {
-      val all_variable_names = base_variables ++ p.exps
+      val all_variable_names = (base_variables ++ p.exps
         .collect { case a: Access => a.vars }
         .flatten
-        .distinct
-        .map(_.name)
+        .map(_.name)).distinct
       Prod(
         p.exps.filter(e =>
           getVariables(e).map(_.name).forall(all_variable_names.contains(_))

diff --git a/src/main/scala/uk/ac/ed/dal/structtensor/parser/Parser.scala b/src/main/scala/uk/ac/ed/dal/structtensor/parser/Parser.scala
@@ -157,4 +157,8 @@ object Parser {
   def program[$: P]: P[Seq[Rule]] = P(rule.rep(sep = "\n"))
 
   def parser[$: P] = P(program ~ End)
+
+  def iterators[$: P]: P[(String, Seq[Variable])] = P(
+    name ~ "->" ~ variableSeq
+  )
 }
diff --git a/src/test/resources/correct_test_outputs/including-extra-variable_wo_body.cpp b/src/test/resources/correct_test_outputs/including-extra-variable_wo_body.cpp
@@ -0,0 +1,47 @@
+
+#include <iostream>
+#include <random>
+#include <algorithm>
+#include <chrono>
+
+using namespace std;
+using namespace std::chrono;
+
+extern "C"
+void fn(double ** covar, double ** X, int N, int M) {
+
+
+long time_computation = 0, start_computation, end_computation;
+start_computation = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+for (int iter = 0; iter < 1000; ++iter) {
+
+for (int i = 0; i < M; ++i) {
+
+for (int j = 0; j < N; ++j) {
+
+for (int k = 0; k < min({(j) + 1, N}); ++k) {
+
+covar[j][k] += (X[i][k] * X[i][j]);
+}
+}
+}
+}
+end_computation = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+time_computation = end_computation - start_computation;
+cout << time_computation << endl;
+long time_reconstruction = 0, start_reconstruction, end_reconstruction;
+start_reconstruction = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+for (int j = 0; j < N; ++j) {
+
+int kp = j;
+for (int k = max({(j) + 1, 0}); k < N; ++k) {
+
+int jp = k;
+covar[j][k] = covar[jp][kp];
+}
+}
+end_reconstruction = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+time_reconstruction = end_reconstruction - start_reconstruction;
+cout << time_reconstruction << endl;
+
+}
diff --git a/src/test/resources/correct_test_outputs/partial-variable-ordering_wo_body.cpp b/src/test/resources/correct_test_outputs/partial-variable-ordering_wo_body.cpp
@@ -0,0 +1,44 @@
+
+#include <iostream>
+#include <random>
+#include <algorithm>
+#include <chrono>
+
+using namespace std;
+using namespace std::chrono;
+
+extern "C"
+void fn(double ** covar, double ** X, int N, int M) {
+
+
+long time_computation = 0, start_computation, end_computation;
+start_computation = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+for (int i = 0; i < M; ++i) {
+
+for (int j = 0; j < N; ++j) {
+
+for (int k = 0; k < min({(j) + 1, N}); ++k) {
+
+covar[j][k] += (X[i][k] * X[i][j]);
+}
+}
+}
+end_computation = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+time_computation = end_computation - start_computation;
+cout << time_computation << endl;
+long time_reconstruction = 0, start_reconstruction, end_reconstruction;
+start_reconstruction = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+for (int j = 0; j < N; ++j) {
+
+int kp = j;
+for (int k = max({(j) + 1, 0}); k < N; ++k) {
+
+int jp = k;
+covar[j][k] = covar[jp][kp];
+}
+}
+end_reconstruction = duration_cast<microseconds>(system_clock::now().time_since_epoch()).count();
+time_reconstruction = end_reconstruction - start_reconstruction;
+cout << time_reconstruction << endl;
+
+}