thek0tyara/llama-cpp-turboquant
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 4

Updated array mult file to align with example

Browse source
This commit is contained in:
Alejandro Saucedo 2020-11-11 07:25:31 +00:00
parent d581d5c36f
commit 5cb2b61426
1 changed files with 17 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

24
python/test/test_array_multiplication.py
View file

@ -4,6 +4,18 @@ import kp
def test_array_multiplication():
# 1. Create Kompute Manager (selects device 0 by default)
mgr = kp.Manager()
# 2. Create Kompute Tensors to hold data
tensor_in_a = kp.Tensor([2, 2, 2])
tensor_in_b = kp.Tensor([1, 2, 3])
tensor_out = kp.Tensor([0, 0, 0])
# 3. Initialise the Kompute Tensors in the GPU
mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])
# 4. Define the multiplication shader code to run on the GPU
@ps.python2shader
def compute_shader_multiply(index=("input", "GlobalInvocationId", ps.ivec3),
data1=("buffer", 0, ps.Array(ps.f32)),
@ -12,14 +24,12 @@ def test_array_multiplication():
i = index.x
data3[i] = data1[i] * data2[i]
tensor_in_a = kp.Tensor([2, 2, 2])
tensor_in_b = kp.Tensor([1, 2, 3])
tensor_out = kp.Tensor([0, 0, 0])
# 5. Run shader code against our previously defined tensors
mgr.eval_algo_data_def(
[tensor_in_a, tensor_in_b, tensor_out],
compute_shader_multiply.to_spirv())
mgr = kp.Manager()
mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])
mgr.eval_algo_data_def([tensor_in_a, tensor_in_b, tensor_out], compute_shader_multiply.to_spirv())
# 6. Sync tensor data from GPU back to local
mgr.eval_tensor_sync_local_def([tensor_out])
assert tensor_out.data() == [2.0, 4.0, 6.0]