├── tensorflow_matmult.py
└── spark_square_matrix_matmul.scala


/tensorflow_matmult.py:
--------------------------------------------------------------------------------
 1 | import tensorflow as tf
 2 | 
 3 | cluster=tf.train.ClusterSpec({"worker":["10.0.1.8:2222", "10.0.1.9:2222", "10.0.1.10:2222"]})
 4 | 
 5 | # the index should be adjusted for different tasks
 6 | server=tf.train.Server(cluster, job_name="worker", task_index=0)
 7 | 
 8 | matrix = tf.Variable(tf.random_normal([9000, 9000], stddev=0.35, name="matrix"))
 9 | init = tf.initialize_all_variables()
10 | 
11 | with tf.device("/job:worker/task:0"):
12 |     sub_row1=tf.gather(matrix, range(0,3000))
13 |     r1=tf.matmul(sub_row1, matrix)
14 | 
15 | with tf.device("/job:worker/task:1"):
16 |     sub_row2=tf.gather(matrix, range(3000,6000))
17 |     r2=tf.matmul(sub_row2, matrix)
18 | 
19 | with tf.device("/job:worker/task:2"):
20 |     sub_row3=tf.gather(matrix, range(6000,9000))
21 |     r3=tf.matmul(sub_row3, matrix)
22 | 
23 | with tf.device("/job:worker/task:0"):
24 |     all_matrix = tf.concat(0, [r1, r2, r3])
25 | 
26 | with tf.Session("grpc://localhost:2222") as sess:
27 |     sess.run(init)
28 |     output = sess.run(all_matrix)
29 | 
30 | output.shape
31 | 
32 | 


--------------------------------------------------------------------------------
/spark_square_matrix_matmul.scala:
--------------------------------------------------------------------------------
 1 | import org.apache.spark.mllib.linalg.Matrices
 2 | import org.apache.spark.mllib.linalg.distributed.BlockMatrix
 3 | 
 4 | val rowParallelism = 4 
 5 | val colParallelism = 4
 6 | val rowBlockSize = 2048
 7 | val colBlockSize = 2048
 8 | 
 9 | val rows = rowParallelism * rowBlockSize
10 | val cols = colParallelism * colBlockSize
11 | 
12 | val rowBlocks = rows/rowBlockSize
13 | val colBlocks = cols/colBlockSize
14 | 
15 | val left_rdd = sc.parallelize( {for(i <- 0 until rowBlocks; j <- 0 until colBlocks) yield (i, j)}, rowParallelism*colParallelism).map( coord => (coord, Matrices.rand(rowBlockSize, colBlockSize, util.Random.self)))
16 | 
17 | val right_rdd = sc.parallelize( {for(i <- 0 until rowBlocks; j <- 0 until colBlocks) yield (i, j)}, rowParallelism*colParallelism).map( coord => (coord, Matrices.rand(rowBlockSize, colBlockSize, util.Random.self)))
18 | 
19 | val left_bm = new BlockMatrix(left_rdd, rowBlockSize, colBlockSize).cache()
20 | val right_bm = new BlockMatrix(right_rdd, rowBlockSize, colBlockSize).cache()
21 | 
22 | val t = System.nanoTime()
23 | val result = left_bm.multiply(right_bm)
24 | result.validate()
25 | val latency = (System.nanoTime() - t) / 1e9
26 | println(latency)
27 | 
28 | 


--------------------------------------------------------------------------------