Esempio n. 1
0
def cholesky(A):
  '''
  Cholesky matrix decomposition.
 
  Args:
    A(Expr): matrix to be decomposed
  '''
 
  A = expr.force(A)
  n = int(math.sqrt(len(A.tiles)))
  tile_size = A.shape[0] / n
  for k in range(n):
    # A[k,k] = DPOTRF(A[k,k])
    diag_ex = get_ex(k, k, tile_size, A.shape)
    A = expr.region_map(A, diag_ex, _cholesky_dpotrf_mapper)
    
    if k == n - 1: break
    
    # A[l,k] = DTRSM(A[k,k], A[l,k]) l -> [k+1,n)
    col_ex = extent.create(((k+1)*tile_size, k*tile_size),(n*tile_size, (k+1)*tile_size), A.shape)
    A = expr.region_map(A, col_ex, _cholesky_dtrsm_mapper, fn_kw=dict(diag_ex=diag_ex))
    
    # A[m,m] = DSYRK(A[m,k], A[m,m]) m -> [k+1,n)
    # A[l,m] = DGEMM(A[l,k], A[m,k], A[l,m]) m -> [k+1,n) l -> [m+1,n)
    col_exs = list([extent.create((m*tile_size, m*tile_size), (n*tile_size, (m+1)*tile_size), A.shape) for m in range(k+1,n)])
    A = expr.region_map(A, col_exs, _cholesky_dsyrk_dgemm_mapper, fn_kw=dict(k=k))
  
  
  # update the right corner to 0
  col_exs = list([extent.create((0, m*tile_size),(m*tile_size, (m+1)*tile_size),A.shape) for m in range(1,n)])
  A = expr.region_map(A, col_exs, lambda input, array, ex: np.zeros(input.shape, input.dtype))
  return A
Esempio n. 2
0
  def test_optimization_region_map(self):
    def mapper(tile, ex):
      return tile + 10

    ex = array.extent.create((0, 0), (1, 5), (5, 5))
    a = expr.region_map(expr.ones((5, 5)), ex, mapper) + expr.ones((5, 5))*10

    for child in a.optimized().op.deps:
      Assert.true(not isinstance(child, expr.local.LocalInput))
Esempio n. 3
0
    def test_optimization_region_map(self):
        def mapper(tile, ex):
            return tile + 10

        ex = array.extent.create((0, 0), (1, 5), (5, 5))
        a = expr.region_map(expr.ones((5, 5)), ex, mapper) + expr.ones(
            (5, 5)) * 10

        for child in a.optimized().op.deps:
            Assert.true(not isinstance(child, expr.operator.local.LocalInput))
Esempio n. 4
0
def test_tilesharing(ctx):
  print "#worker:", ctx.num_workers
  N_EXAMPLES = 5 * ctx.num_workers
  x = expr.ones((N_EXAMPLES, 1), tile_hint=(N_EXAMPLES / ctx.num_workers, 1))
  y = expr.region_map(x, extent.create((0, 0), (3, 1), (N_EXAMPLES, 1)), fn=lambda data, ex, a: data+a, fn_kw={'a': 1})

  npx = np.ones((N_EXAMPLES, 1))
  npy = np.ones((N_EXAMPLES, 1))
  npy[0:3, 0] += 1

  assert np.all(np.equal(x.glom(), npx))
  assert np.all(np.equal(y.glom(), npy))