gen.scala 2.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  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. var indent = 0
  10. def spacing = " " * indent
  11. def assign(lhs: String, rhs: String) =
  12. spacing + lhs + " = " + rhs + ";\n"
  13. def init(t: String, n: String, v: String) =
  14. assign(t+" "+n, v)
  15. def open_block(s: String = "") = {
  16. val result = (if (s != "") spacing + s else "") + spacing + "{\n"
  17. indent = indent + 1
  18. result
  19. }
  20. def close_block = {
  21. indent = indent - 1
  22. spacing + "}\n"
  23. }
  24. def ar(m: String, i: String) = m+"["+i+"]"
  25. def r(a: String, b: String*) = (a :: b.toList).reduceLeft(_+"_"+_)
  26. def rb(m: Int, n: Int, p: Int) = {
  27. 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")
  28. for (i <- 0 until m)
  29. s += init("t*", r("c", i), "&"+ar("c", "ldc*"+i))
  30. for (i <- 0 until m; j <- 0 until n)
  31. s += init("t", r("c", i, j), ar(r("c", i), j))
  32. def doit(m: Int, n: Int, p: Int) = {
  33. for (i <- 0 until m)
  34. s += init("t*", r("a", i), "&"+ar("a", "lda*"+i))
  35. for (k <- 0 until p)
  36. s += init("t*", r("b", k), "&"+ar("b", "ldb*"+k))
  37. for (k <- 0 until p; i <- 0 until m; j <- 0 until n)
  38. s += assign(r("c", i, j), "fma(" + ar(r("a", i), k) + ", " + ar(r("b", k), j) + ", " + r("c", i, j) + ")")
  39. }
  40. s += open_block("for (t *a = a0, *b = b0; a < a0 + p/RBK*RBK; a += RBK, b += RBK*ldb)\n")
  41. doit(m, n, p)
  42. s += close_block
  43. s += open_block("for (t *a = a0 + p/RBK*RBK, *b = b0 + p/RBK*RBK*ldb; a < a0 + p; a++, b += ldb)\n")
  44. doit(m, n, 1)
  45. s += close_block
  46. for (i <- 0 until m; j <- 0 until n)
  47. s += assign(ar(r("c", i), j), r("c", i, j))
  48. s += close_block
  49. s
  50. }
  51. def gcd(a: Int, b: Int): Int = if (b == 0) a else gcd(b, a%b)
  52. def lcm(a: Int, b: Int): Int = a*b/gcd(a, b)
  53. def lcm(a: Seq[Int]): Int = {
  54. if (a.tail.isEmpty) a.head
  55. else lcm(a.head, lcm(a.tail))
  56. }
  57. def test1(m: Int, n: Int, p: Int, m1: Int, n1: Int, p1: Int) = {
  58. val decl = "static const int RBM = "+m+", RBN = "+n+", RBK = "+p+";\n" +
  59. "static const int CBM = "+m1+", CBN = "+n1+", CBK = "+p1+";\n"
  60. writeFile("rb.h", decl + rb(m, n, p))
  61. //"make"!!
  62. "make run"!
  63. ("cp a.out " + Seq("b", m, n, p, m1, n1, p1, "run").reduce(_+"."+_))!
  64. }
  65. def main(args: Array[String]): Unit = {
  66. test1(4, 5, 6, 24, 25, 24)
  67. //for (i <- 4 to 6; j <- 4 to 6; k <- 4 to 6)
  68. // test1(i, j, k, if (i == 5) 35 else 36, if (j == 5) 35 else 36, if (k == 5) 35 else 36)
  69. }
  70. }