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
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))
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))
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))
