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

All python tests pass

Browse source
This commit is contained in:
Alejandro Saucedo 2021-02-28 07:57:36 +00:00
parent 4c4d073b90
commit 91d3b9a223
11 changed files with 158 additions and 169 deletions
Download patch file Download diff file Expand all files Collapse all files

23
python/test/test_array_multiplication.py
View file

@ -9,29 +9,26 @@ def test_array_multiplication():
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])
tensor_in_a = mgr.tensor([2, 2, 2])
tensor_in_b = mgr.tensor([1, 2, 3])
tensor_out = mgr.tensor([0, 0, 0])
# 3. Initialise the Kompute Tensors in the GPU
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
params = [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),
def compute_mult(index=("input", "GlobalInvocationId", ps.ivec3),
data1=("buffer", 0, ps.Array(ps.f32)),
data2=("buffer", 1, ps.Array(ps.f32)),
data3=("buffer", 2, ps.Array(ps.f32))):
i = index.x
data3[i] = data1[i] * data2[i]
# 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())
# 6. Sync tensor data from GPU back to local
mgr.eval_tensor_sync_local_def([tensor_out])
(mgr.sequence()
.record(kp.OpTensorSyncDevice(params))
.record(kp.OpAlgoDispatch(mgr.algorithm(params, compute_mult.to_spirv())))
.record(kp.OpTensorSyncLocal([tensor_out]))
.eval())
assert tensor_out.data() == [2.0, 4.0, 6.0]
assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0])