benchmarks/mm/gen.scala

   1 import scala.sys.process._
   2 object MMGen {
   3   implicit def i2s(i: Int) = i.toString
   4   def writeFile(name: String, contents: String) = {
   5     val f = new java.io.FileWriter(name)
   6     f.write(contents)
   7     f.close
   8   }
   9
  10   var indent = 0
  11   def spacing = "  " * indent
  12   def assign(lhs: String, rhs: String) =
  13     spacing + lhs + " = " + rhs + ";\n"
  14   def init(t: String, n: String, v: String) =
  15     assign(t+" "+n, v)
  16   def open_block(s: String = "") = {
  17     val result = (if (s != "") spacing + s else "") + spacing + "{\n"
  18     indent = indent + 1
  19     result
  20   }
  21   def close_block = {
  22     indent = indent - 1
  23     spacing + "}\n"
  24   }
  25
  26   def ar(m: String, i: String) = m+"["+i+"]"
  27   def r(a: String, b: String*) = (a :: b.toList).reduceLeft(_+"_"+_)
  28
  29   def rb(m: Int, n: Int, p: Int) = {
  30     var s = open_block("static inline void kloop(size_t p, t* a0, size_t lda, t* b0, size_t ldb, t* c, size_t ldc)\n")
  31
  32     for (i <- 0 until m)
  33       s += init("t*", r("c", i), "&"+ar("c", "ldc*"+i))
  34     for (i <- 0 until m; j <- 0 until n)
  35       s += init("t", r("c", i, j), ar(r("c", i), j))
  36
  37     def doit(m: Int, n: Int, p: Int) = {
  38       for (i <- 0 until m)
  39         s += init("t*", r("a", i), "&"+ar("a", "lda*"+i))
  40       for (k <- 0 until p)
  41         s += init("t*", r("b", k), "&"+ar("b", "ldb*"+k))
  42       for (k <- 0 until p; i <- 0 until m; j <- 0 until n)
  43         s += assign(r("c", i, j), "fma(" + ar(r("a", i), k) + ", " + ar(r("b", k), j) + ", " + r("c", i, j) + ")")
  44     }
  45
  46     s += open_block("for (t *a = a0, *b = b0; a < a0 + p/RBK*RBK; a += RBK, b += RBK*ldb)\n")
  47     doit(m, n, p)
  48     s += close_block
  49
  50     s += open_block("for (t *a = a0 + p/RBK*RBK, *b = b0 + p/RBK*RBK*ldb; a < a0 + p; a++, b += ldb)\n")
  51     doit(m, n, 1)
  52     s += close_block
  53
  54     for (i <- 0 until m; j <- 0 until n)
  55       s += assign(ar(r("c", i), j), r("c", i, j))
  56     s += close_block
  57
  58     s
  59   }
  60   def gcd(a: Int, b: Int): Int = if (b == 0) a else gcd(b, a%b)
  61   def lcm(a: Int, b: Int): Int = a*b/gcd(a, b)
  62   def lcm(a: Seq[Int]): Int = {
  63     if (a.tail.isEmpty) a.head
  64     else lcm(a.head, lcm(a.tail))
  65   }
  66   def test1(m: Int, n: Int, p: Int, m1: Int, n1: Int, p1: Int) = {
  67     val decl = "static const int RBM = "+m+", RBN = "+n+", RBK = "+p+";\n" +
  68                "static const int CBM = "+m1+", CBN = "+n1+", CBK = "+p1+";\n"
  69     writeFile("rb.h", decl + rb(m, n, p))
  70     //"make"!!
  71
  72     "make run"!
  73
  74     ("cp a.out " + Seq("b", m, n, p, m1, n1, p1, "run").reduce(_+"."+_))!
  75   }
  76   def main(args: Array[String]): Unit = {
  77     test1(4, 5, 6, 24, 25, 24)
  78     //for (i <- 4 to 6; j <- 4 to 6; k <- 4 to 6)
  79     //  test1(i, j, k, if (i == 5) 35 else 36, if (j == 5) 35 else 36, if (k == 5) 35 else 36)
  80   }
  81 }