Python implementation
This commit is contained in:
Alejandro Saucedo
parent
198fb46eb6
commit
4c4d073b90
18 changed files with 3172 additions and 3349 deletions
|
|
@ -7,25 +7,26 @@ import pyshader as ps
|
|||
|
||||
DIRNAME = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
def test_opalgobase_file():
|
||||
"""
|
||||
Test basic OpMult operation
|
||||
"""
|
||||
|
||||
tensor_in_a = kp.Tensor([2, 2, 2])
|
||||
tensor_in_b = kp.Tensor([1, 2, 3])
|
||||
tensor_out = kp.Tensor([0, 0, 0])
|
||||
|
||||
mgr = kp.Manager()
|
||||
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
|
||||
|
||||
shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv")
|
||||
|
||||
mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)
|
||||
|
||||
mgr.eval_tensor_sync_local_def([tensor_out])
|
||||
|
||||
assert tensor_out.data() == [2.0, 4.0, 6.0]
|
||||
# TODO: Add example with file
|
||||
#def test_opalgobase_file():
|
||||
# """
|
||||
# Test basic OpMult operation
|
||||
# """
|
||||
#
|
||||
# tensor_in_a = kp.Tensor([2, 2, 2])
|
||||
# tensor_in_b = kp.Tensor([1, 2, 3])
|
||||
# tensor_out = kp.Tensor([0, 0, 0])
|
||||
#
|
||||
# mgr = kp.Manager()
|
||||
# mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
|
||||
#
|
||||
# shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv")
|
||||
#
|
||||
# mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)
|
||||
#
|
||||
# mgr.eval_tensor_sync_local_def([tensor_out])
|
||||
#
|
||||
# assert tensor_out.data() == [2.0, 4.0, 6.0]
|
||||
|
||||
|
||||
|
||||
|
|
@ -48,18 +49,23 @@ void main()
|
|||
}
|
||||
"""
|
||||
|
||||
tensor_in_a = kp.Tensor([2, 2, 2])
|
||||
tensor_in_b = kp.Tensor([1, 2, 3])
|
||||
tensor_out = kp.Tensor([0, 0, 0])
|
||||
|
||||
mgr = kp.Manager()
|
||||
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
|
||||
|
||||
spirv = kp.Shader.compile_source(shader)
|
||||
|
||||
mgr.eval_algo_data_def([tensor_in_a, tensor_in_b, tensor_out], spirv)
|
||||
mgr = kp.Manager()
|
||||
|
||||
mgr.eval_tensor_sync_local_def([tensor_out])
|
||||
tensor_in_a = mgr.tensor([2, 2, 2])
|
||||
tensor_in_b = mgr.tensor([1, 2, 3])
|
||||
tensor_out = mgr.tensor([0, 0, 0])
|
||||
|
||||
params = [tensor_in_a, tensor_in_b, tensor_out]
|
||||
|
||||
algo = mgr.algorithm(params, spirv)
|
||||
|
||||
(mgr.sequence()
|
||||
.record(kp.OpTensorSyncLocal(params))
|
||||
.record(kp.OpAlgoDispatch(algo))
|
||||
.record(kp.OpTensorSyncDevice(params))
|
||||
.eval())
|
||||
|
||||
assert tensor_out.data() == [2.0, 4.0, 6.0]
|
||||
|
||||
|
|
@ -67,36 +73,53 @@ def test_sequence():
|
|||
"""
|
||||
Test basic OpAlgoBase operation
|
||||
"""
|
||||
mgr = kp.Manager(0, [2])
|
||||
|
||||
tensor_in_a = kp.Tensor([2, 2, 2])
|
||||
tensor_in_b = kp.Tensor([1, 2, 3])
|
||||
tensor_out = kp.Tensor([0, 0, 0])
|
||||
shader = """
|
||||
#version 450
|
||||
layout(set = 0, binding = 0) buffer tensorLhs {float valuesLhs[];};
|
||||
layout(set = 0, binding = 1) buffer tensorRhs {float valuesRhs[];};
|
||||
layout(set = 0, binding = 2) buffer tensorOutput { float valuesOutput[];};
|
||||
layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
|
||||
void main()
|
||||
{
|
||||
uint index = gl_GlobalInvocationID.x;
|
||||
valuesOutput[index] = valuesLhs[index] * valuesRhs[index];
|
||||
}
|
||||
"""
|
||||
|
||||
shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv"))
|
||||
mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)
|
||||
spirv = kp.Shader.compile_source(shader)
|
||||
|
||||
mgr.eval_await_def()
|
||||
mgr = kp.Manager(0)
|
||||
|
||||
seq = mgr.sequence("op")
|
||||
seq.begin()
|
||||
seq.record_tensor_sync_local([tensor_in_a])
|
||||
seq.record_tensor_sync_local([tensor_in_b])
|
||||
seq.record_tensor_sync_local([tensor_out])
|
||||
seq.end()
|
||||
seq.eval()
|
||||
tensor_in_a = mgr.tensor([2, 2, 2])
|
||||
tensor_in_b = mgr.tensor([1, 2, 3])
|
||||
tensor_out = mgr.tensor([0, 0, 0])
|
||||
|
||||
mgr.destroy("op")
|
||||
params = [tensor_in_a, tensor_in_b, tensor_out]
|
||||
|
||||
assert seq.is_init() == False
|
||||
algo = mgr.algorithm(params, spirv)
|
||||
|
||||
sq = mgr.sequence()
|
||||
|
||||
sq.record(kp.OpTensorSyncLocal(params))
|
||||
sq.record(kp.OpAlgoDispatch(algo))
|
||||
sq.record(kp.OpTensorSyncDevice(params))
|
||||
|
||||
sq.eval()
|
||||
|
||||
assert sq.is_init() == True
|
||||
|
||||
sq.destroy()
|
||||
|
||||
assert sq.is_init() == False
|
||||
|
||||
assert tensor_out.data() == [2.0, 4.0, 6.0]
|
||||
assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0])
|
||||
|
||||
mgr.destroy(tensor_in_a)
|
||||
mgr.destroy([tensor_in_b, tensor_out])
|
||||
tensor_in_a.destroy()
|
||||
tensor_in_b.destroy()
|
||||
tensor_out.destroy()
|
||||
|
||||
assert tensor_in_a.is_init() == False
|
||||
assert tensor_in_b.is_init() == False
|
||||
|
|
@ -105,10 +128,8 @@ def test_sequence():
|
|||
def test_workgroup():
|
||||
mgr = kp.Manager(0)
|
||||
|
||||
tensor_a = kp.Tensor(np.zeros([16,8]))
|
||||
tensor_b = kp.Tensor(np.zeros([16,8]))
|
||||
|
||||
mgr.rebuild([tensor_a, tensor_b])
|
||||
tensor_a = mgr.tensor(np.zeros([16,8]))
|
||||
tensor_b = mgr.tensor(np.zeros([16,8]))
|
||||
|
||||
@ps.python2shader
|
||||
def compute_shader_wg(gl_idx=("input", "GlobalInvocationId", ps.ivec3),
|
||||
|
|
@ -120,17 +141,15 @@ def test_workgroup():
|
|||
data1[i] = f32(gl_idx.x)
|
||||
data2[i] = f32(gl_idx.y)
|
||||
|
||||
seq = mgr.sequence("new")
|
||||
seq.begin()
|
||||
seq.record_algo_data([tensor_a, tensor_b], compute_shader_wg.to_spirv(), workgroup=(16,8,1))
|
||||
seq.end()
|
||||
seq.eval()
|
||||
algo = mgr.algorithm([tensor_a, tensor_b], compute_shader_wg.to_spirv(), (16,8,1), [], [])
|
||||
|
||||
mgr.destroy(seq)
|
||||
(mgr.sequence()
|
||||
.record(kp.OpTensorSyncDevice([tensor_a, tensor_b]))
|
||||
.record(kp.OpAlgoDispatch(algo))
|
||||
.record(kp.OpAlgoTensorSyncLocal([tensor_a, tensor_b]))
|
||||
.eval())
|
||||
|
||||
assert seq.is_init() == False
|
||||
|
||||
mgr.eval_tensor_sync_local_def([tensor_a, tensor_b])
|
||||
assert sq.is_init() == False
|
||||
|
||||
print(tensor_a.numpy())
|
||||
print(tensor_b.numpy())
|
||||
|
|
@ -138,32 +157,3 @@ def test_workgroup():
|
|||
assert np.all(tensor_a.numpy() == np.stack([np.arange(16)]*8, axis=1).ravel())
|
||||
assert np.all(tensor_b.numpy() == np.stack([np.arange(8)]*16, axis=0).ravel())
|
||||
|
||||
mgr.destroy([tensor_a, tensor_b])
|
||||
|
||||
assert tensor_a.is_init() == False
|
||||
assert tensor_b.is_init() == False
|
||||
|
||||
|
||||
def test_tensor_rebuild_backwards_compat():
|
||||
"""
|
||||
Test basic OpMult operation
|
||||
"""
|
||||
|
||||
tensor_in_a = kp.Tensor([2, 2, 2])
|
||||
tensor_in_b = kp.Tensor([1, 2, 3])
|
||||
tensor_out = kp.Tensor([0, 0, 0])
|
||||
|
||||
mgr = kp.Manager()
|
||||
|
||||
mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])
|
||||
|
||||
shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv"))
|
||||
mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)
|
||||
mgr.eval_await_def()
|
||||
|
||||
mgr.eval_tensor_sync_local_def([tensor_out])
|
||||
|
||||
assert tensor_out.data() == [2.0, 4.0, 6.0]
|
||||
assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0])
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue