Skip to content

Commit

Permalink
Update benchmarks
Browse files Browse the repository at this point in the history
  • Loading branch information
@Routhleck
Routhleck committed Dec 13, 2023
1 parent 0ef3782 commit 5e8f2cf
Show file tree
Hide file tree
Showing 2 changed files with 572 additions and 60 deletions.
329 changes: 295 additions & 34 deletions brainpy/_src/math/event/tests/event_csrmv_taichi_VS_event_csrmv.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,250 @@

bm.set_platform('gpu')

s = [1000, 5000, 10000, 15000, 20000, 25000, 30000]

s = [1000, 5000, 10000, 20000, 25000, 30000]
p = [0.1, 0.2, 0.3, 0.4, 0.5]
values_type = ['homo',
'heter']
events_type = ['bool',

shape = [
1000,
2500,
5000,
10000,
25000,
37500,
50000
]



values_type = [
'homo',
'heter'
]
events_type = [
'bool',
'float',
]
transpose = [True,
False]
transpose = [
True,
False
]

print(bm.get_platform())

def test_event_csrmv_cpu(shape, values_type, events_type, transpose):
rng = bm.random.RandomState(seed=1234)
indices, indptr = bp.conn.FixedProb(0.3)(*shape).require('pre2post')
vector = rng.random(shape[0] if transpose else shape[1]) < 0.1
weight = 1.


if events_type == 'float':
vector = vector.astype(bm.float32)
if values_type == 'heter':
heter_data = bm.ones(indices.shape) * weight
weight = heter_data

# groundtruth = bm.as_jax(vector, dtype=float) @ bm.as_jax(dense)

result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
# time.sleep(2)

time0 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time1 = time.time()
# time.sleep(2)

time2 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time3 = time.time()
# time.sleep(2)

time4 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time5 = time.time()
# time.sleep(2)

time6 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time7 = time.time()

time8 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time9 = time.time()

result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
# print(result1[0])
# print(result2)
# print(groundtruth - result1[0])
# print(groundtruth - result2)

# print(result1[0] - result2)
# print(bm.allclose(groundtruth, result1[0]))
# print(bm.allclose(groundtruth, result2))
# assert bm.allclose(result1[0], result2)

time12 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time13 = time.time()
# time.sleep(2)

time14 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time15 = time.time()
# time.sleep(2)

time16 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time17 = time.time()
# time.sleep(2)

time18 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time19 = time.time()

time20 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time21 = time.time()

taichi_aot_time1 = (time1 - time0) * 1000
taichi_aot_time2 = (time3 - time2) * 1000
taichi_aot_time3 = (time5 - time4) * 1000
taichi_aot_time4 = (time7 - time6) * 1000
taichi_aot_time5 = (time9 - time8) * 1000
brainpy_time1 = (time13 - time12) * 1000
brainpy_time2 = (time15 - time14) * 1000
brainpy_time3 = (time17 - time16) * 1000
brainpy_time4 = (time19 - time18) * 1000
brainpy_time5 = (time21 - time20) * 1000

print('shape: ', shape, 'values_type: ', values_type, 'events_type: ', events_type, 'transpose: ', transpose)
print('taichi_aot_1: ', taichi_aot_time1, 'ms')
print('taichi_aot_2: ', taichi_aot_time2, 'ms')
print('taichi_aot_3: ', taichi_aot_time3, 'ms')
print('taichi_aot_4: ', taichi_aot_time4, 'ms')
print('taichi_aot_5: ', taichi_aot_time5, 'ms')
print('brainpylib_cpu_1: ', brainpy_time1, 'ms')
print('brainpylib_cpu_2: ', brainpy_time2, 'ms')
print('brainpylib_cpu_3: ', brainpy_time3, 'ms')
print('brainpylib_cpu_4: ', brainpy_time4, 'ms')
print('brainpylib_cpu_5: ', brainpy_time5, 'ms')
assert(jnp.allclose(result1[0], result2))

speedup = (brainpy_time1 + brainpy_time2 + brainpy_time3 + brainpy_time4 + brainpy_time5) / \
(taichi_aot_time1 + taichi_aot_time2 + taichi_aot_time3 + taichi_aot_time4 + taichi_aot_time5) - 1

return taichi_aot_time1, taichi_aot_time2, taichi_aot_time3, taichi_aot_time4, taichi_aot_time5,\
brainpy_time1, brainpy_time2, brainpy_time3, brainpy_time4, brainpy_time5, speedup

def test_event_csrmv_gpu(shape, values_type, events_type, transpose):
rng = bm.random.RandomState(seed=1234)
indices, indptr = bp.conn.FixedProb(0.3)(*shape).require('pre2post')
vector = rng.random(shape[0] if transpose else shape[1]) < 0.1
weight = 1.


if events_type == 'float':
vector = vector.astype(bm.float32)
if values_type == 'heter':
heter_data = bm.ones(indices.shape) * weight
weight = heter_data

# groundtruth = bm.as_jax(vector, dtype=float) @ bm.as_jax(dense)



result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
# time.sleep(2)

time0 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time1 = time.time()
# time.sleep(2)

time2 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time3 = time.time()
# time.sleep(2)

time4 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time5 = time.time()
# time.sleep(2)

time6 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time7 = time.time()

time8 = time.time()
result1 = jax.block_until_ready(bm.event.csrmv_taichi(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time9 = time.time()

result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
# print(result1[0])
# print(result2)
# print(groundtruth - result1[0])
# print(groundtruth - result2)

print(result1[0] - result2)
# print(bm.allclose(groundtruth, result1[0]))
# print(bm.allclose(groundtruth, result2))
# assert bm.allclose(result1[0], result2)

time12 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time13 = time.time()
# time.sleep(2)

time14 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time15 = time.time()
# time.sleep(2)

def test_event_ell_cpu(s, p, values_type, events_type, transpose):
time16 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time17 = time.time()
# time.sleep(2)

time18 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time19 = time.time()

time20 = time.time()
result2 = jax.block_until_ready(bm.event.csrmv(weight, indices, indptr, vector, shape=shape, transpose=transpose))
time21 = time.time()

taichi_aot_time1 = (time1 - time0) * 1000
taichi_aot_time2 = (time3 - time2) * 1000
taichi_aot_time3 = (time5 - time4) * 1000
taichi_aot_time4 = (time7 - time6) * 1000
taichi_aot_time5 = (time9 - time8) * 1000
brainpy_time1 = (time13 - time12) * 1000
brainpy_time2 = (time15 - time14) * 1000
brainpy_time3 = (time17 - time16) * 1000
brainpy_time4 = (time19 - time18) * 1000
brainpy_time5 = (time21 - time20) * 1000
print('shape: ', shape, 'values_type: ', values_type, 'events_type: ', events_type, 'transpose: ', transpose)
print('taichi_aot_1: ', taichi_aot_time1, 'ms')
print('taichi_aot_2: ', taichi_aot_time2, 'ms')
print('taichi_aot_3: ', taichi_aot_time3, 'ms')
print('taichi_aot_4: ', taichi_aot_time4, 'ms')
print('taichi_aot_5: ', taichi_aot_time5, 'ms')
print('brainpylib_gpu_1: ', brainpy_time1, 'ms')
print('brainpylib_gpu_2: ', brainpy_time2, 'ms')
print('brainpylib_gpu_3: ', brainpy_time3, 'ms')
print('brainpylib_gpu_4: ', brainpy_time4, 'ms')
print('brainpylib_gpu_5: ', brainpy_time5, 'ms')

# assert(jnp.allclose(result1[0], result2))

speedup = (brainpy_time1 + brainpy_time2 + brainpy_time3 + brainpy_time4 + brainpy_time5) / \
(taichi_aot_time1 + taichi_aot_time2 + taichi_aot_time3 + taichi_aot_time4 + taichi_aot_time5) - 1

return taichi_aot_time1, taichi_aot_time2, taichi_aot_time3, taichi_aot_time4, taichi_aot_time5,\
brainpy_time1, brainpy_time2, brainpy_time3, brainpy_time4, brainpy_time5, speedup


def test_event_csrmv_square_cpu(s, p, values_type, events_type, transpose):
print('s: ', s, 'p: ', p)
k = int(s * p)
bm.random.seed(1234)
Expand Down Expand Up @@ -81,8 +310,8 @@ def test_event_ell_cpu(s, p, values_type, events_type, transpose):
result2 = jax.block_until_ready(bm.event.csrmv(weight, csr_indices, csr_indptr, vector, shape=(s, s), transpose=transpose))
# print(result1[0])
# print(result2)
# print(groundtruth - result1[0])
# print(groundtruth - result2)
# print(groundtruth - result1[0])
# print(groundtruth - result2)

# print(result1[0] - result2)
# print(bm.allclose(groundtruth, result1[0]))
Expand Down Expand Up @@ -141,7 +370,7 @@ def test_event_ell_cpu(s, p, values_type, events_type, transpose):
return taichi_aot_time1, taichi_aot_time2, taichi_aot_time3, taichi_aot_time4, taichi_aot_time5,\
brainpy_time1, brainpy_time2, brainpy_time3, brainpy_time4, brainpy_time5, speedup

def test_event_ell_gpu(s, p, values_type, events_type, transpose):
def test_event_csrmv_square_gpu(s, p, values_type, events_type, transpose):
print('s: ', s, 'p: ', p)
k = int(s * p)
bm.random.seed(1234)
Expand Down Expand Up @@ -265,39 +494,71 @@ def test_event_ell_gpu(s, p, values_type, events_type, transpose):
PATH = os.path.dirname(os.path.abspath(__file__))

# init dataframe
df = pd.DataFrame(columns=['s', 'p', 'backend', 'values type', 'events type', 'transpose',
df = pd.DataFrame(columns=['s', 'p', 'shape[0]', 'shape[1]', 'backend', 'values type', 'events type', 'transpose',
'taichi aot time1(ms)', 'taichi aot time2(ms)', 'taichi aot time3(ms)', 'taichi aot time4(ms)', 'taichi aot time5(ms)',
'brainpy time1(ms)', 'brainpy time2(ms)', 'brainpy time3(ms)', 'brainpy time4(ms)', 'brainpy time5(ms)',
'speedup'])

### SQUARE MATRIX

# if (bm.get_platform() == 'cpu'):
# for _s in s:
# for _p in p:
# for _values_type in values_type:
# for _events_type in events_type:
# for _transpose in transpose:
# taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,\
# brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup = test_event_csrmv_square_cpu(_s, _p, _values_type, _events_type, _transpose)
# # append to dataframe
# df.loc[df.shape[0]] = [_s, _p, _s, _s, 'cpu', _values_type, _events_type, _transpose,
# taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,
# brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup]
# df.to_csv(f'{PATH}/event_csrmv_square_cpu.csv', index=False)

# if (bm.get_platform() == 'gpu'):
# for _s in s:
# for _p in p:
# for _values_type in values_type:
# for _events_type in events_type:
# for _transpose in transpose:
# taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,\
# brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup = test_event_csrmv_square_gpu(_s, _p, _values_type, _events_type, _transpose)
# # append to dataframe
# df.loc[df.shape[0]] = [_s, _p, _s, _s, 'gpu', _values_type, _events_type, _transpose,
# taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,
# brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup]
# df.to_csv(f'{PATH}/event_csrmv_square_gpu.csv', index=False)

### RECTANGULAR MATRIX
if (bm.get_platform() == 'cpu'):
for _s in s:
for _p in p:
for _values_type in values_type:
for _events_type in events_type:
for _transpose in transpose:
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,\
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup = test_event_ell_cpu(_s, _p, _values_type, _events_type, _transpose)
# append to dataframe
df.loc[df.shape[0]] = [_s, _p, 'cpu', _values_type, _events_type, _transpose,
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup]
for shape1 in shape:
for shape2 in shape:
for _values_type in values_type:
for _events_type in events_type:
for _transpose in transpose:
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,\
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup = test_event_csrmv_cpu((shape1, shape2), _values_type, _events_type, _transpose)
# append to dataframe
df.loc[df.shape[0]] = [(shape1, shape2), 0.5 , shape1, shape2,'cpu', _values_type, _events_type, _transpose,
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup]
df.to_csv(f'{PATH}/event_csrmv_cpu.csv', index=False)

if (bm.get_platform() == 'gpu'):
for _s in s:
for _p in p:
for _values_type in values_type:
for _events_type in events_type:
for _transpose in transpose:
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,\
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup = test_event_ell_gpu(_s, _p, _values_type, _events_type, _transpose)
# append to dataframe
df.loc[df.shape[0]] = [_s, _p, 'gpu', _values_type, _events_type, _transpose,
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup]
for shape1 in shape:
for shape2 in shape:
for _values_type in values_type:
for _events_type in events_type:
for _transpose in transpose:
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,\
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup = test_event_csrmv_gpu((shape1, shape2), _values_type, _events_type, _transpose)
# append to dataframe
df.loc[df.shape[0]] = [(shape1, shape2), 0.5 , shape1, shape2, 'gpu', _values_type, _events_type, _transpose,
taichi_aot_time_1, taichi_aot_time_2, taichi_aot_time_3, taichi_aot_time_4, taichi_aot_time_5,
brainpy_time_1, brainpy_time_2, brainpy_time_3, brainpy_time_4, brainpy_time_5, speedup]
df.to_csv(f'{PATH}/event_csrmv_gpu.csv', index=False)


# if (bm.get_platform() == 'gpu'):
# for _s in s:
# for _p in p:
Expand Down
Loading

0 comments on commit 5e8f2cf

Please sign in to comment.