implementing distributed linear algebra

lib/orchpy/arrays/dist/__init__.py

@@ -1 +1,2 @@
-from core import DistArray, BLOCK_SIZE
+import random
+from core import *

lib/orchpy/arrays/dist/core.py

@@ -3,12 +3,14 @@ import numpy as np
 import arrays.single as single
 import orchpy as op
 
+__all__ = ["BLOCK_SIZE", "DistArray", "assemble", "zeros", "ones", "copy",
+           "eye", "triu", "tril", "blockwise_dot", "dot", "block_column", "block_row"]
+
 BLOCK_SIZE = 10
 
 class DistArray(object):
-  def construct(self, shape, dtype, objrefs):
+  def construct(self, shape, objrefs):
     self.shape = shape
-    self.dtype = dtype
     self.objrefs = objrefs
     self.ndim = len(shape)
     self.num_blocks = [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in self.shape]
@@ -16,41 +18,49 @@ class DistArray(object):
       raise Exception("The fields `num_blocks` and `objrefs` are inconsistent, `num_blocks` is {} and `objrefs` has shape {}".format(self.num_blocks, list(self.objrefs.shape)))
 
   def deserialize(self, primitives):
-    (shape, dtype_name, objrefs) = primitives
-    self.construct(shape, np.dtype(dtype_name), objrefs)
+    (shape, objrefs) = primitives
+    self.construct(shape, objrefs)
 
   def serialize(self):
-    return (self.shape, self.dtype.__name__, self.objrefs)
+    return (self.shape, self.objrefs)
 
   def __init__(self):
     self.shape = None
-    self.dtype = None
     self.objrefs = None
 
-  def compute_block_lower(self, index):
-    if len(index) != self.ndim:
-      raise Exception("The value `index` equals {}, but `ndim` is {}.".format(index, self.ndim))
+  @staticmethod
+  def compute_block_lower(index, shape):
+    # TODO(rkn): Check that the entries of index are in the correct range.
+    # TODO(rkn): Check that len(index) == len(shape).
     return [elem * BLOCK_SIZE for elem in index]
 
-  def compute_block_upper(self, index):
-    if len(index) != self.ndim:
-      raise Exception("The value `index` equals {}, but `ndim` is {}.".format(index, self.ndim))
+  @staticmethod
+  def compute_block_upper(index, shape):
+    # TODO(rkn): Check that the entries of index are in the correct range.
+    # TODO(rkn): Check that len(index) == len(shape).
     upper = []
-    for i in range(self.ndim):
-      upper.append(min((index[i] + 1) * BLOCK_SIZE, self.shape[i]))
+    for i in range(len(shape)):
+      upper.append(min((index[i] + 1) * BLOCK_SIZE, shape[i]))
     return upper
 
-  def compute_block_shape(self, index):
-    lower = self.compute_block_lower(index)
-    upper = self.compute_block_upper(index)
+  @staticmethod
+  def compute_block_shape(index, shape):
+    lower = DistArray.compute_block_lower(index, shape)
+    upper = DistArray.compute_block_upper(index, shape)
     return [u - l for (l, u) in zip(lower, upper)]
 
+  @staticmethod
+  def compute_num_blocks(shape):
+    return [int(np.ceil(1.0 * a / BLOCK_SIZE)) for a in shape]
+
   def assemble(self):
     """Assemble an array on this node from a distributed array object reference."""
-    result = np.zeros(self.shape)
+    first_block = op.pull(self.objrefs[(0,) * self.ndim])
+    dtype = first_block.dtype
+    result = np.zeros(self.shape, dtype=dtype)
     for index in np.ndindex(*self.num_blocks):
-      lower = self.compute_block_lower(index)
-      upper = self.compute_block_upper(index)
+      lower = DistArray.compute_block_lower(index, self.shape)
+      upper = DistArray.compute_block_upper(index, self.shape)
       result[[slice(l, u) for (l, u) in zip(lower, upper)]] = op.pull(self.objrefs[index])
     return result
 
@@ -58,3 +68,137 @@ class DistArray(object):
     # TODO(rkn): fix this, this is just a placeholder that should work but is inefficient
     a = self.assemble()
     return a[sliced]
+
+@op.distributed([DistArray], [np.ndarray])
+def assemble(a):
+  return a.assemble()
+
+@op.distributed([List[int], str], [DistArray])
+def zeros(shape, dtype_name):
+  num_blocks = DistArray.compute_num_blocks(shape)
+  objrefs = np.empty(num_blocks, dtype=object)
+  for index in np.ndindex(*num_blocks):
+    objrefs[index] = single.zeros(DistArray.compute_block_shape(index, shape), dtype_name)
+  result = DistArray()
+  result.construct(shape, objrefs)
+  return result
+
+@op.distributed([List[int], str], [DistArray])
+def ones(shape, dtype_name):
+  num_blocks = DistArray.compute_num_blocks(shape)
+  objrefs = np.empty(num_blocks, dtype=object)
+  for index in np.ndindex(*num_blocks):
+    objrefs[index] = single.ones(DistArray.compute_block_shape(index, shape), dtype_name)
+  result = DistArray()
+  result.construct(shape, objrefs)
+  return result
+
+@op.distributed([DistArray], [DistArray])
+def copy(a):
+  num_blocks = DistArray.compute_num_blocks(a.shape)
+  objrefs = np.empty(num_blocks, dtype=object)
+  for index in np.ndindex(*num_blocks):
+    objrefs[index] = single.copy(a.objrefs[index])
+  result = DistArray()
+  result.construct(a.shape, objrefs)
+  return result
+
+@op.distributed([int, str], [DistArray])
+def eye(dim, dtype_name):
+  shape = [dim, dim]
+  num_blocks = DistArray.compute_num_blocks(shape)
+  objrefs = np.empty(num_blocks, dtype=object)
+  for (i, j) in np.ndindex(*num_blocks):
+    if i == j:
+      objrefs[i, j] = single.eye(DistArray.compute_block_shape([i, j], shape)[0], dtype_name)
+    else:
+      objrefs[i, j] = single.zeros(DistArray.compute_block_shape([i, j], shape), dtype_name)
+  result = DistArray()
+  result.construct(shape, objrefs)
+  return result
+
+@op.distributed([DistArray], [DistArray])
+def triu(a):
+  if a.ndim != 2:
+    raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim))
+  objrefs = np.empty(a.num_blocks, dtype=object)
+  for i in range(a.num_blocks[0]):
+    for j in range(a.num_blocks[1]):
+      if i < j:
+        objrefs[i, j] = single.copy(a.objrefs[i, j])
+      elif i == j:
+        objrefs[i, j] = single.triu(a.objrefs[i, j])
+      else:
+        objrefs[i, j] = single.zeros_like(a.objrefs[i, j])
+  result = DistArray()
+  result.construct(a.shape, objrefs)
+  return result
+
+@op.distributed([DistArray], [DistArray])
+def tril(a):
+  if a.ndim != 2:
+    raise Exception("Input must have 2 dimensions, but a.ndim is " + str(a.ndim))
+  objrefs = np.empty(a.num_blocks, dtype=object)
+  for i in range(a.num_blocks[0]):
+    for j in range(a.num_blocks[1]):
+      if i > j:
+        objrefs[i, j] = single.copy(a.objrefs[i, j])
+      elif i == j:
+        objrefs[i, j] = single.tril(a.objrefs[i, j])
+      else:
+        objrefs[i, j] = single.zeros_like(a.objrefs[i, j])
+  result = DistArray()
+  result.construct(a.shape, objrefs)
+  return result
+
+@op.distributed([np.ndarray, None], [np.ndarray])
+def blockwise_dot(*matrices):
+  n = len(matrices)
+  if n % 2 != 0:
+    raise Exception("blockwise_dot expects an even number of arguments, but len(matrices) is {}.".format(n))
+  shape = (matrices[0].shape[0], matrices[n / 2].shape[1])
+  result = np.zeros(shape)
+  for i in range(n / 2):
+    result += np.dot(matrices[i], matrices[n / 2 + i])
+  return result
+
+@op.distributed([DistArray, DistArray], [DistArray])
+def dot(a, b):
+  if a.ndim != 2:
+    raise Exception("dot expects its arguments to be 2-dimensional, but a.ndim = {}.".format(a.ndim))
+  if b.ndim != 2:
+    raise Exception("dot expects its arguments to be 2-dimensional, but b.ndim = {}.".format(b.ndim))
+  if a.shape[1] != b.shape[0]:
+    raise Exception("dot expects a.shape[1] to equal b.shape[0], but a.shape = {} and b.shape = {}.".format(a.shape, b.shape))
+  shape = [a.shape[0], b.shape[1]]
+  num_blocks = DistArray.compute_num_blocks(shape)
+  objrefs = np.empty(num_blocks, dtype=object)
+  for i in range(num_blocks[0]):
+    for j in range(num_blocks[1]):
+      args = list(a.objrefs[i, :]) + list(b.objrefs[:, j])
+      objrefs[i, j] = blockwise_dot(*args)
+  result = DistArray()
+  result.construct(shape, objrefs)
+  return result
+
+# This is not in numpy, should we expose this?
+@op.distributed([DistArray], [DistArray])
+def block_column(a, col):
+  if a.ndim != 2:
+    raise Exception("block_column expects its argument to be 2-dimensional, but a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape))
+  top_block_shape = DistArray.compute_block_shape([0, col])
+  shape = [a.shape[0], top_block_shape[1]]
+  result = DistArray()
+  result.construct(shape, a.objrefs[:, col])
+  return result
+
+# This is not in numpy, should we expose this?
+@op.distributed([DistArray], [DistArray])
+def block_row(a, row):
+  if a.ndim != 2:
+    raise Exception("block_row expects its argument to be 2-dimensional, but a.ndim = {}, a.shape = {}.".format(a.ndim, a.shape))
+  left_block_shape = DistArray.compute_block_shape([row, 0])
+  shape = [left_block_shape[0], a.shape[1]]
+  result = DistArray()
+  result.construct(shape, a.objrefs[row, :])
+  return result

lib/orchpy/arrays/dist/random.py

@@ -0,0 +1,17 @@
+from typing import List
+
+import numpy as np
+import arrays.single as single
+import orchpy as op
+
+from core import *
+
+@op.distributed([List[int]], [DistArray])
+def normal(shape):
+  num_blocks = DistArray.compute_num_blocks(shape)
+  objrefs = np.empty(num_blocks, dtype=object)
+  for index in np.ndindex(*num_blocks):
+    objrefs[index] = single.random.normal(DistArray.compute_block_shape(index, shape))
+  result = DistArray()
+  result.construct(shape, objrefs)
+  return result

lib/orchpy/arrays/single/__init__.py

@@ -1,2 +1,2 @@
 import random, linalg
-from core import zeros, ones, eye, dot, vstack, hstack, subarray, copy, tril, triu
+from core import zeros, zeros_like, ones, eye, dot, vstack, hstack, subarray, copy, tril, triu

lib/orchpy/arrays/single/core.py

@@ -2,17 +2,21 @@ from typing import List
 import numpy as np
 import orchpy as op
 
-@op.distributed([List[int]], [np.ndarray])
-def zeros(shape):
-  return np.zeros(shape)
+@op.distributed([List[int], str], [np.ndarray])
+def zeros(shape, dtype_name):
+  return np.zeros(shape, dtype=np.dtype(dtype_name))
 
-@op.distributed([List[int]], [np.ndarray])
-def ones(shape):
-  return np.ones(shape)
+@op.distributed([np.ndarray], [np.ndarray])
+def zeros_like(x):
+  return np.zeros_like(x)
 
-@op.distributed([int], [np.ndarray])
-def eye(dim):
-  return np.eye(dim)
+@op.distributed([List[int], str], [np.ndarray])
+def ones(shape, dtype_name):
+  return np.ones(shape, dtype=np.dtype(dtype_name))
+
+@op.distributed([int, str], [np.ndarray])
+def eye(dim, dtype_name):
+  return np.eye(dim, dtype=np.dtype(dtype_name))
 
 @op.distributed([np.ndarray, np.ndarray], [np.ndarray])
 def dot(a, b):

lib/orchpy/arrays/single/linalg.py

@@ -2,21 +2,95 @@ from typing import List
 import numpy as np
 import orchpy as op
 
-# TODO(rkn): this should take the same optional "mode" argument as np.linalg.qr, except that the different options sometimes have different numbers of return values, which could be a problem
+__all__ = ["matrix_power", "solve", "tensorsolve", "tensorinv", "inv",
+           "cholesky", "eigvals", "eigvalsh", "pinv", "slogdet", "det",
+           "svd", "eig", "eigh", "lstsq", "norm", "qr", "cond", "matrix_rank",
+           "LinAlgError", "multi_dot"]
+
+@op.distributed([np.ndarray, int], [np.ndarray])
+def matrix_power(M, n):
+  return np.linalg.matrix_power(M, n)
+
+@op.distributed([np.ndarray, np.ndarray], [np.ndarray])
+def solve(a, b):
+  return np.linalg.solve(a, b)
+
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
+def tensorsolve(a):
+  raise NotImplementedError
+
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
+def tensorinv(a):
+  raise NotImplementedError
+
+@op.distributed([np.ndarray], [np.ndarray])
+def inv(a):
+  return np.linalg.inv(a)
+
+@op.distributed([np.ndarray], [np.ndarray])
+def cholesky(a):
+  return np.linalg.cholesky(a)
+
+@op.distributed([np.ndarray], [np.ndarray])
+def eigvals(a):
+  return np.linalg.eigvals(a)
+
+@op.distributed([np.ndarray], [np.ndarray])
+def eigvalsh(a):
+  raise NotImplementedError
+
+@op.distributed([np.ndarray], [np.ndarray])
+def pinv(a):
+  return np.linalg.pinv(a)
+
+@op.distributed([np.ndarray], [int])
+def slogdet(a):
+  raise NotImplementedError
+
+@op.distributed([np.ndarray], [float])
+def det(a):
+  return np.linalg.det(a)
+
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray])
+def svd(a):
+  return np.linalg.svd(a)
+
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
+def eig(a):
+  return np.linalg.eig(a)
+
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray])
+def eigh(a):
+  return np.linalg.eigh(a)
+
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray, int, np.ndarray])
+def lstsq(a, b):
+  return np.linalg.lstsq(a)
+
+@op.distributed([np.ndarray], [float])
+def norm(x):
+  return np.linalg.norm(x)
+
 @op.distributed([np.ndarray], [np.ndarray, np.ndarray])
 def qr(a):
-  """
-  Suppose (n, m) = a.shape
-  If n >= m:
-    q.shape == (n, m)
-    r.shape == (m, m)
-  If n < m:
-    q.shape == (n, n)
-    r.shape == (n, m)
-  """
   return np.linalg.qr(a)
 
-#@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray])
+@op.distributed([np.ndarray], [float])
+def cond(x):
+  return np.linalg.cond(x)
+
+@op.distributed([np.ndarray], [int])
+def matrix_rank(M):
+  return np.linalg.matrix_rank(M)
+
+@op.distributed([np.ndarray, None], [np.ndarray])
+def multi_dot(a):
+  raise NotImplementedError
+
+
+
+# This isn't in numpy, should we expose it?
+@op.distributed([np.ndarray], [np.ndarray, np.ndarray, np.ndarray])
 def modified_lu(q):
   """
   Algorithm 5 from http://www.eecs.berkeley.edu/Pubs/TechRpts/2013/EECS-2013-175.pdf

lib/orchpy/orchpy/worker.py

@@ -78,8 +78,8 @@ def distributed(arg_types, return_types, worker=global_worker):
       """This is what gets executed remotely on a worker after a distributed function is scheduled by the scheduler."""
       print "Calling function {} with arguments {}".format(func.__name__, arguments)
       result = func(*arguments)
-      if len(return_types) != 1 and len(result) != len(return_types):
-        raise Exception("The @distributed decorator for function {} has {} return values with types {}, but {} returned {} values.".format(func.__name__, len(return_types), return_types, func.__name__, len(result)))
+      check_return_values(func_call, result) # throws an exception if result is invalid
+      print "Finished executing function {} with arguments {}".format(func.__name__, arguments)
       return result
     def func_call(*args):
       """This is what gets run immediately when a worker calls a distributed function."""
@@ -94,6 +94,18 @@ def distributed(arg_types, return_types, worker=global_worker):
     return func_call
   return distributed_decorator
 
+# helper method, this should not be called by the user
+def check_return_values(function, result):
+  if len(function.return_types) == 1:
+    if not isinstance(result, function.return_types[0]):
+      raise Exception("The @distributed decorator for function {} expects one return value with type {}, but {} returned a {}.".format(function.__name__, function.return_types[0], function.__name__, type(result)))
+  else:
+    if len(result) != len(function.return_types):
+      raise Exception("The @distributed decorator for function {} has {} return values with types {}, but {} returned {} values.".format(function.__name__, len(function.return_types), function.return_types, function.__name__, len(result)))
+    for i in range(len(result)):
+      if not isinstance(result[i], function.return_types[i]):
+        raise Exception("The {}th return value for function {} has type {}, but the @distributed decorator expected a return value of type {}.".format(i, function.__name__, type(result[i]), function.return_types[i]))
+
 # helper method, this should not be called by the user
 def check_arguments(function, args):
   # check the number of args
@@ -107,7 +119,7 @@ def check_arguments(function, args):
       expected_type = function.arg_types[i]
     elif i == len(function.arg_types) - 1 and function.arg_types[-1] is not None:
       expected_type = function.arg_types[-1]
-    elif function.arg_types[-1] is None and len(function.arg_types > 1):
+    elif function.arg_types[-1] is None and len(function.arg_types) > 1:
       expected_type = function.arg_types[-2]
     else:
       assert False, "This code should be unreachable."
@@ -136,7 +148,7 @@ def get_arguments_for_execution(function, args, worker=global_worker):
       expected_type = function.arg_types[i]
     elif i == len(function.arg_types) - 1 and function.arg_types[-1] is not None:
       expected_type = function.arg_types[-1]
-    elif function.arg_types[-1] is None and len(function.arg_types > 1):
+    elif function.arg_types[-1] is None and len(function.arg_types) > 1:
       expected_type = function.arg_types[-2]
     else:
       assert False, "This code should be unreachable."
@@ -145,6 +157,7 @@ def get_arguments_for_execution(function, args, worker=global_worker):
       # get the object from the local object store
       print "Getting argument {} for function {}.".format(i, function.__name__)
       argument = worker.get_object(arg)
+      print "Successfully retrieved argument {} for function {}.".format(i, function.__name__)
     else:
       # pass the argument by value
       argument = arg