def plan_Copy(self, ops, legacy=False):
noslice = Ellipsis if legacy else None # LEGACY
copies, ops = split(
ops, lambda op:
(op.src_slice is noslice and op.dst_slice is noslice))
plans = []
if copies:
X = self.all_data[[self.sidx[op.src] for op in copies]]
Y = self.all_data[[self.sidx[op.dst] for op in copies]]
incs = np.array([op.inc for op in copies], dtype=np.int32)
plans.append(plan_copy(self.queue, X, Y, incs))
if ops:
inds = lambda ary, i: np.arange(ary.size, dtype=np.int32)[
Ellipsis if i is None else i]
xinds = [inds(op.src, op.src_slice) for op in ops]
yinds = [inds(op.dst, op.dst_slice) for op in ops]
dupl = lambda s: (s is not None and not (isinstance(
s, np.ndarray) and s.dtype == np.bool) and len(s) != len(set(s)
))
if any(dupl(i) for i in xinds) or any(dupl(i) for i in yinds):
raise NotImplementedError("Duplicates in indices")
X = self.all_data[[self.sidx[op.src] for op in ops]]
Y = self.all_data[[self.sidx[op.dst] for op in ops]]
Xinds = self.RaggedArray(xinds)
Yinds = self.RaggedArray(yinds)
incs = np.array([op.inc for op in ops], dtype=np.int32)
plans.append(plan_slicedcopy(self.queue, X, Y, Xinds, Yinds, incs))
return plans
def plan_Copy(self, ops, legacy=False):
noslice = Ellipsis if legacy else None # LEGACY
copies, ops = split(ops, lambda op: (
op.src_slice is noslice and op.dst_slice is noslice))
plans = []
if copies:
X = self.all_data[[self.sidx[op.src] for op in copies]]
Y = self.all_data[[self.sidx[op.dst] for op in copies]]
incs = np.array([op.inc for op in copies], dtype=np.int32)
plans.append(plan_copy(self.queue, X, Y, incs))
if ops:
inds = lambda ary, i: np.arange(ary.size, dtype=np.int32)[
Ellipsis if i is None else i]
xinds = [inds(op.src, op.src_slice) for op in ops]
yinds = [inds(op.dst, op.dst_slice) for op in ops]
dupl = lambda s: (
s is not None
and not (isinstance(s, np.ndarray) and s.dtype == np.bool)
and len(s) != len(set(s)))
if any(dupl(i) for i in xinds) or any(dupl(i) for i in yinds):
raise NotImplementedError("Duplicates in indices")
X = self.all_data[[self.sidx[op.src] for op in ops]]
Y = self.all_data[[self.sidx[op.dst] for op in ops]]
Xinds = self.RaggedArray(xinds)
Yinds = self.RaggedArray(yinds)
incs = np.array([op.inc for op in ops], dtype=np.int32)
plans.append(plan_slicedcopy(self.queue, X, Y, Xinds, Yinds, incs))
return plans
def plan_Copy(self, ops, legacy=False):
noslice = Ellipsis if legacy else None # LEGACY
copies, ops = split(
ops, lambda op:
(op.src_slice is noslice and op.dst_slice is noslice))
plans = []
if copies:
A = self.all_data[[self.sidx[op.src] for op in copies]]
B = self.all_data[[self.sidx[op.dst] for op in copies]]
incs = np.array([op.inc for op in copies], dtype=np.int32)
plans.append(plan_copy(self.queue, A, B, incs))
if ops:
A = self.all_data[[self.sidx[op.src] for op in ops]]
B = self.all_data[[self.sidx[op.dst] for op in ops]]
inds = lambda ary, i: np.arange(ary.size, dtype=np.int32)[
Ellipsis if i is None else i]
Ainds = self.RaggedArray(
[inds(op.src, op.src_slice) for op in ops])
Binds = self.RaggedArray(
[inds(op.dst, op.dst_slice) for op in ops])
incs = np.array([op.inc for op in ops], dtype=np.int32)
plans.append(plan_slicedcopy(self.queue, A, B, Ainds, Binds, incs))
return plans
def plan_SlicedCopy(self, ops):
A = self.all_data[[self.sidx[op.a] for op in ops]]
B = self.all_data[[self.sidx[op.b] for op in ops]]
Ainds = self.RaggedArray([
np.arange(op.a.size, dtype=np.int32)[op.a_slice] for op in ops])
Binds = self.RaggedArray([
np.arange(op.b.size, dtype=np.int32)[op.b_slice] for op in ops])
incs = self.RaggedArray([
np.array(op.inc, dtype=np.int32) for op in ops])
return [plan_slicedcopy(self.queue, A, B, Ainds, Binds, incs)]
def test_slicedcopy(rng):
sizes = rng.randint(20, 200, size=10)
A = RA([rng.normal(size=size) for size in sizes])
B = RA([rng.normal(size=size) for size in sizes])
incs = RA([rng.randint(0, 2) for _ in sizes])
Ainds = []
Binds = []
for size in sizes:
r = np.arange(size, dtype=np.int32)
u = rng.choice([0, 1, 2])
if u == 0:
Ainds.append(r)
Binds.append(r)
elif u == 1:
Ainds.append(r[:10])
Binds.append(r[-10:])
elif u == 2:
n = rng.randint(2, size - 2)
Ainds.append(rng.permutation(size)[:n])
Binds.append(rng.permutation(size)[:n])
Ainds = RA(Ainds)
Binds = RA(Binds)
queue = cl.CommandQueue(ctx)
clA = CLRA(queue, A)
clB = CLRA(queue, B)
clAinds = CLRA(queue, Ainds)
clBinds = CLRA(queue, Binds)
clincs = CLRA(queue, incs)
# compute on host
for i in range(len(sizes)):
if incs[i]:
B[i][Binds[i]] += A[i][Ainds[i]]
else:
B[i][Binds[i]] = A[i][Ainds[i]]
# compute on device
plan = plan_slicedcopy(queue, clA, clB, clAinds, clBinds, clincs)
plan()
# check result
for y, yy in zip(B, clB.to_host()):
assert np.allclose(y, yy)
def test_slicedcopy(rng):
sizes = rng.randint(20, 200, size=10)
A = RA([rng.normal(size=size) for size in sizes])
B = RA([rng.normal(size=size) for size in sizes])
incs = RA([rng.randint(0, 2) for _ in sizes])
Ainds = []
Binds = []
for size in sizes:
r = np.arange(size, dtype=np.int32)
u = rng.choice([0, 1, 2])
if u == 0:
Ainds.append(r)
Binds.append(r)
elif u == 1:
Ainds.append(r[:10])
Binds.append(r[-10:])
elif u == 2:
n = rng.randint(2, size - 2)
Ainds.append(rng.permutation(size)[:n])
Binds.append(rng.permutation(size)[:n])
Ainds = RA(Ainds)
Binds = RA(Binds)
queue = cl.CommandQueue(ctx)
clA = CLRA(queue, A)
clB = CLRA(queue, B)
clAinds = CLRA(queue, Ainds)
clBinds = CLRA(queue, Binds)
clincs = CLRA(queue, incs)
# compute on host
for i in range(len(sizes)):
if incs[i]:
B[i][Binds[i]] += A[i][Ainds[i]]
else:
B[i][Binds[i]] = A[i][Ainds[i]]
# compute on device
plan = plan_slicedcopy(queue, clA, clB, clAinds, clBinds, clincs)
plan()
# check result
for y, yy in zip(B, clB.to_host()):
assert np.allclose(y, yy)
def plan_SlicedCopy(self, ops):
copies, ops = split(
ops, lambda op: op.a_slice is Ellipsis and op.b_slice is Ellipsis)
plans = []
if copies:
A = self.all_data[[self.sidx[op.a] for op in copies]]
B = self.all_data[[self.sidx[op.b] for op in copies]]
incs = np.array([op.inc for op in copies], dtype=np.int32)
plans.append(plan_copy(self.queue, A, B, incs))
if ops:
A = self.all_data[[self.sidx[op.a] for op in ops]]
B = self.all_data[[self.sidx[op.b] for op in ops]]
inds = lambda ary, i: np.arange(ary.size, dtype=np.int32)[i]
Ainds = self.RaggedArray([inds(op.a, op.a_slice) for op in ops])
Binds = self.RaggedArray([inds(op.b, op.b_slice) for op in ops])
incs = np.array([op.inc for op in ops], dtype=np.int32)
plans.append(plan_slicedcopy(self.queue, A, B, Ainds, Binds, incs))
return plans
def plan_SlicedCopy(self, ops):
copies, ops = split(
ops, lambda op: op.a_slice is Ellipsis and op.b_slice is Ellipsis)
plans = []
if copies:
A = self.all_data[[self.sidx[op.a] for op in copies]]
B = self.all_data[[self.sidx[op.b] for op in copies]]
incs = np.array([op.inc for op in copies], dtype=np.int32)
plans.append(plan_copy(self.queue, A, B, incs))
if ops:
A = self.all_data[[self.sidx[op.a] for op in ops]]
B = self.all_data[[self.sidx[op.b] for op in ops]]
inds = lambda ary, i: np.arange(ary.size, dtype=np.int32)[i]
Ainds = self.RaggedArray([inds(op.a, op.a_slice) for op in ops])
Binds = self.RaggedArray([inds(op.b, op.b_slice) for op in ops])
incs = np.array([op.inc for op in ops], dtype=np.int32)
plans.append(plan_slicedcopy(self.queue, A, B, Ainds, Binds, incs))
return plans