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

Updated python to align with new workgrups and constants, as well as removed all string related functions

Browse source
This commit is contained in:
Alejandro Saucedo 2021-02-14 16:47:24 +00:00
parent d9a6abb918
commit 711bd158ba
2 changed files with 90 additions and 165 deletions
Download patch file Download diff file Expand all files Collapse all files

167
python/src/main.cpp
View file

@ -11,6 +11,8 @@ namespace py = pybind11;
py::object kp_debug, kp_info, kp_warning, kp_error;
PYBIND11_MODULE(kp, m) {
spdlog::set_level(
static_cast<spdlog::level::level_enum>(0));
// The logging modules are used in the Kompute.hpp file
py::module_ logging = py::module_::import("logging");
@ -110,38 +112,27 @@ PYBIND11_MODULE(kp, m) {
"Records operation to sync tensor from local memory to GPU memory")
.def("record_tensor_sync_local", &kp::Sequence::record<kp::OpTensorSyncLocal>,
"Records operation to sync tensor(s) from GPU memory to local memory")
.def("record_algo_mult", &kp::Sequence::record<kp::OpMult>,
"Records operation to run multiplication compute shader to two input tensors and an output tensor")
.def("record_algo_file", [](kp::Sequence &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
const std::string& file_path,
std::tuple<uint32_t,uint32_t,uint32_t> work_group) -> bool {
const kp::OpAlgoBase::KomputeWorkgroup wgroup{
std::get<0>(work_group), std::get<1>(work_group), std::get<2>(work_group),
};
return self.record<kp::OpAlgoBase>(tensors, file_path, wgroup);
},
"Records an operation using a custom shader provided from a shader path",
py::arg("tensors"), py::arg("file_path"), py::arg("work_group") = std::make_tuple(0,0,0) )
.def("record_algo_file", &kp::Sequence::record<
kp::OpAlgoBase,
const std::string&,
kp::Workgroup,
kp::Constants>,
"Records an operation using a custom shader provided from a shader path",
py::arg("tensors"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
.def("record_algo_data", [](kp::Sequence &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
py::bytes &bytes,
std::tuple<uint32_t,uint32_t,uint32_t> work_group) -> bool {
kp::Workgroup workgroup,
kp::Constants constants) -> bool {
// Bytes have to be converted into std::vector
py::buffer_info info(py::buffer(bytes).request());
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
const kp::OpAlgoBase::KomputeWorkgroup wgroup{
std::get<0>(work_group), std::get<1>(work_group), std::get<2>(work_group),
};
return self.record<kp::OpAlgoBase>(
tensors, std::vector<char>(data, data + length), wgroup
);
tensors, std::vector<char>(data, data + length), workgroup, constants);
},
"Records an operation using a custom shader provided as spirv bytes",
py::arg("tensors"), py::arg("bytes"), py::arg("work_group") = std::make_tuple(0,0,0) )
.def("record_algo_lro", &kp::Sequence::record<kp::OpAlgoLhsRhsOut>,
"Records operation to run left right out operation with custom shader");
py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() );
py::class_<kp::Manager>(m, "Manager")
@ -199,33 +190,27 @@ PYBIND11_MODULE(kp, m) {
"Evaluates operation to sync tensor from local memory to GPU memory with new anonymous Sequence")
.def("eval_tensor_sync_local_def", &kp::Manager::evalOpDefault<kp::OpTensorSyncLocal>,
"Evaluates operation to sync tensor(s) from GPU memory to local memory with new anonymous Sequence")
.def("eval_algo_mult_def", &kp::Manager::evalOpDefault<kp::OpMult>,
"Evaluates operation to run multiplication compute shader to two input tensors and an output tensor with new anonymous Sequence")
.def("eval_algo_file_def", &kp::Manager::evalOpDefault<kp::OpAlgoBase, std::string>,
"Evaluates an operation using a custom shader provided from a shader path with new anonymous Sequence")
.def("eval_algo_str_def", &kp::Manager::evalOpDefault<kp::OpAlgoBase, std::vector<char>>,
"Evaluates an operation using a custom shader provided as string provided as list of characters with new anonymous Sequence")
.def("eval_algo_str_def", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
const std::string& shader_str){
const std::vector<char> shader_vec(shader_str.begin(), shader_str.end());
self.evalOpDefault<kp::OpAlgoBase>(tensors, shader_vec);
},
"Evaluates an operation using a custom shader provided as string with a new anonymous Sequence")
.def("eval_algo_file_def", &kp::Manager::evalOpDefault<
kp::OpAlgoBase,
const std::string&,
kp::Workgroup,
kp::Constants>,
"Evaluates an operation using a custom shader provided from a shader path with new anonymous Sequence",
py::arg("tensors"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
.def("eval_algo_data_def", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
py::bytes &bytes) {
py::bytes &bytes,
kp::Workgroup workgroup,
kp::Constants constants) {
// Bytes have to be converted into std::vector
py::buffer_info info(py::buffer(bytes).request());
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOpDefault<kp::OpAlgoBase>(
tensors,
std::vector<char>(data, data + length));
tensors, std::vector<char>(data, data + length), workgroup, constants);
},
"Evaluates an operation using a custom shader provided as spirv bytes with new anonymous Sequence")
.def("eval_algo_lro_def", &kp::Manager::evalOpDefault<kp::OpAlgoLhsRhsOut>,
"Evaluates operation to run left right out operation with custom shader with new anonymous Sequence")
"Evaluates an operation using a custom shader provided as spirv bytes with new anonymous Sequence",
py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
// eval
.def("eval_tensor_copy", &kp::Manager::evalOp<kp::OpTensorCopy>,
@ -234,36 +219,28 @@ PYBIND11_MODULE(kp, m) {
"Evaluates operation to sync tensor from local memory to GPU memory with explicitly named Sequence")
.def("eval_tensor_sync_local", &kp::Manager::evalOp<kp::OpTensorSyncLocal>,
"Evaluates operation to sync tensor(s) from GPU memory to local memory with explicitly named Sequence")
.def("eval_algo_mult", &kp::Manager::evalOp<kp::OpMult>,
"Evaluates operation to run multiplication compute shader to two input tensors and an output tensor with explicitly named Sequence")
.def("eval_algo_file", &kp::Manager::evalOp<kp::OpAlgoBase, std::string>,
"Evaluates an operation using a custom shader provided from a shader path with explicitly named Sequence")
.def("eval_algo_str", &kp::Manager::evalOp<kp::OpAlgoBase, std::vector<char>>,
"Evaluates an operation using a custom shader provided as string provided as list of characters with explicitly named Sequence")
.def("eval_algo_str", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
const std::string& sequenceName,
const std::string& shader_str) {
const std::vector<char> shader_vec(shader_str.begin(), shader_str.end());
self.evalOp<kp::OpAlgoBase>(tensors, sequenceName, shader_vec);
},
"Evaluates an operation using a custom shader provided as string with explicitly named Sequence")
.def("eval_algo_file", &kp::Manager::evalOp<
kp::OpAlgoBase,
const std::string&,
kp::Workgroup,
kp::Constants>,
"Evaluates an operation using a custom shader provided from a shader path with explicitly named Sequence",
py::arg("tensors"), py::arg("sequence_name"), py::arg("data"),py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
.def("eval_algo_data", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
std::string sequenceName,
py::bytes &bytes) {
py::bytes &bytes,
kp::Workgroup workgroup,
kp::Constants constants) {
// Bytes have to be converted into std::vector
py::buffer_info info(py::buffer(bytes).request());
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOp<kp::OpAlgoBase>(
tensors,
sequenceName,
std::vector<char>(data, data + length));
tensors, sequenceName, std::vector<char>(data, data + length), workgroup, constants);
},
"Evaluates an operation using a custom shader provided as spirv bytes with explicitly named Sequence")
.def("eval_algo_lro", &kp::Manager::evalOp<kp::OpAlgoLhsRhsOut>,
"Evaluates operation to run left right out operation with custom shader with explicitly named Sequence")
"Evaluates an operation using a custom shader provided as spirv bytes with explicitly named Sequence",
py::arg("tensors"), py::arg("sequence_name"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
// eval async default
.def("eval_async_tensor_copy_def", &kp::Manager::evalOpAsyncDefault<kp::OpTensorCopy>,
@ -272,33 +249,27 @@ PYBIND11_MODULE(kp, m) {
"Evaluates asynchronously operation to sync tensor from local memory to GPU memory with anonymous Sequence")
.def("eval_async_tensor_sync_local_def", &kp::Manager::evalOpAsyncDefault<kp::OpTensorSyncLocal>,
"Evaluates asynchronously operation to sync tensor(s) from GPU memory to local memory with anonymous Sequence")
.def("eval_async_algo_mult_def", &kp::Manager::evalOpAsyncDefault<kp::OpMult>,
"Evaluates asynchronously operation to run multiplication compute shader to two input tensors and an output tensor with anonymous Sequence")
.def("eval_async_algo_file_def", &kp::Manager::evalOpAsyncDefault<kp::OpAlgoBase, std::string>,
"Evaluates asynchronously an operation using a custom shader provided from a shader path with anonymous Sequence")
.def("eval_async_algo_str_def", &kp::Manager::evalOpAsyncDefault<kp::OpAlgoBase, std::vector<char>>,
"Evaluates Asynchronously an operation using a custom shader provided as string provided as list of characters with new anonymous Sequence")
.def("eval_async_algo_str_def", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
const std::string& shader_str) {
const std::vector<char> shader_vec(shader_str.begin(), shader_str.end());
self.evalOpAsyncDefault<kp::OpAlgoBase>(tensors, shader_vec);
},
"Evaluates Asynchronously an operation using a custom shader provided as string with new anonymous Sequence")
.def("eval_async_algo_file_def", &kp::Manager::evalOpAsyncDefault<
kp::OpAlgoBase,
const std::string&,
kp::Workgroup,
kp::Constants>,
"Evaluates asynchronously an operation using a custom shader provided from a shader path with anonymous Sequence",
py::arg("tensors"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
.def("eval_async_algo_data_def", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
py::bytes &bytes) {
py::bytes &bytes,
kp::Workgroup workgroup,
kp::Constants constants) {
// Bytes have to be converted into std::vector
py::buffer_info info(py::buffer(bytes).request());
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOpAsyncDefault<kp::OpAlgoBase>(
tensors,
std::vector<char>(data, data + length));
tensors, std::vector<char>(data, data + length), workgroup, constants);
},
"Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with anonymous Sequence")
.def("eval_async_algo_lro_def", &kp::Manager::evalOpAsyncDefault<kp::OpAlgoLhsRhsOut>,
"Evaluates asynchronously operation to run left right out operation with custom shader with anonymous Sequence")
"Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with anonymous Sequence",
py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
// eval async
.def("eval_async_tensor_copy", &kp::Manager::evalOpAsync<kp::OpTensorCopy>,
@ -307,36 +278,28 @@ PYBIND11_MODULE(kp, m) {
"Evaluates asynchronously operation to sync tensor from local memory to GPU memory with explicitly named Sequence")
.def("eval_async_tensor_sync_local", &kp::Manager::evalOpAsync<kp::OpTensorSyncLocal>,
"Evaluates asynchronously operation to sync tensor(s) from GPU memory to local memory with explicitly named Sequence")
.def("eval_async_algo_mult", &kp::Manager::evalOpAsync<kp::OpMult>,
"Evaluates asynchronously operation to run multiplication compute shader to two input tensors and an output tensor with explicitly named Sequence")
.def("eval_async_algo_file", &kp::Manager::evalOpAsync<kp::OpAlgoBase, std::string>,
"Evaluates asynchronously an operation using a custom shader provided from a shader path with explicitly named Sequence")
.def("eval_async_algo_str", &kp::Manager::evalOpAsync<kp::OpAlgoBase, std::vector<char>>,
"Evaluates Asynchronous an operation using a custom shader provided as string provided as list of characters with explicitly named Sequence")
.def("eval_async_algo_str", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
const std::string& sequenceName,
const std::string& shader_str) {
const std::vector<char> shader_vec(shader_str.begin(), shader_str.end());
self.evalOpAsync<kp::OpAlgoBase>(tensors, sequenceName, shader_vec);
},
"Evaluates Asynchronous an operation using a custom shader provided as string with explicitly named Sequence")
.def("eval_async_algo_file", &kp::Manager::evalOpAsync<
kp::OpAlgoBase,
const std::string&,
kp::Workgroup,
kp::Constants>,
"Evaluates asynchronously an operation using a custom shader provided from a shader path with explicitly named Sequence",
py::arg("tensors"), py::arg("sequence_name"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() )
.def("eval_async_algo_data", [](kp::Manager &self,
std::vector<std::shared_ptr<kp::Tensor>> tensors,
std::string sequenceName,
py::bytes &bytes) {
py::bytes &bytes,
kp::Workgroup workgroup,
kp::Constants constants) {
// Bytes have to be converted into std::vector
py::buffer_info info(py::buffer(bytes).request());
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOpAsync<kp::OpAlgoBase>(
tensors,
sequenceName,
std::vector<char>(data, data + length));
tensors, sequenceName, std::vector<char>(data, data + length), workgroup, constants);
},
"Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with explicitly named Sequence")
.def("eval_async_algo_lro", &kp::Manager::evalOpAsync<kp::OpAlgoLhsRhsOut>,
"Evaluates asynchronously operation to run left right out operation with custom shader with explicitly named Sequence");
"Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with explicitly named Sequence",
py::arg("tensors"), py::arg("sequence_name"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() );
#ifdef VERSION_INFO
m.attr("__version__") = VERSION_INFO;

88
python/test/test_kompute.py
View file

@ -3,10 +3,11 @@ import os
import kp
import numpy as np
import logging
import pyshader as ps
DIRNAME = os.path.dirname(os.path.abspath(__file__))
def test_opmult():
def test_opalgobase_file():
"""
Test basic OpMult operation
"""
@ -19,43 +20,8 @@ def test_opmult():
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
mgr.eval_algo_mult_def([tensor_in_a, tensor_in_b, tensor_out])
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])
def test_opalgobase_data():
"""
Test basic OpAlgoBase 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()
shaderData = """
#version 450
layout (local_size_x = 1) in;
// The input rebuild bind index is relative to index in parameter passed
layout(set = 0, binding = 0) buffer bina { float tina[]; };
layout(set = 0, binding = 1) buffer binb { float tinb[]; };
layout(set = 0, binding = 2) buffer bout { float tout[]; };
void main() {
uint index = gl_GlobalInvocationID.x;
tout[index] = tina[index] * tinb[index];
}
"""
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
mgr.eval_algo_str_def([tensor_in_a, tensor_in_b, tensor_out], shaderData)
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_tensor_sync_local_def([tensor_out])
@ -73,12 +39,11 @@ def test_opalgobase_file():
tensor_out = kp.Tensor([0, 0, 0])
mgr = kp.Manager()
shaderFilePath = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp")
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shaderFilePath)
shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp")
mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path)
mgr.eval_tensor_sync_local_def([tensor_out])
@ -96,8 +61,8 @@ def test_sequence():
mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out])
shaderFilePath = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp")
mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shaderFilePath)
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()
@ -131,27 +96,19 @@ def test_workgroup():
mgr.rebuild([tensor_a, tensor_b])
shader_src = """
#version 450
layout (local_size_x = 1) in;
// The input rebuild bind index is relative to index in parameter passed
layout(set = 0, binding = 0) writeonly buffer bout { float toutx[]; };
layout(set = 0, binding = 1) writeonly buffer bout2 { float touty[]; };
void main() {
uint index = gl_WorkGroupID.x*gl_NumWorkGroups.y + gl_WorkGroupID.y;
toutx[index] = gl_GlobalInvocationID.x;
touty[index] = gl_GlobalInvocationID.y;
}
"""
shader_src = bytes(shader_src, encoding='utf8')
@ps.python2shader
def compute_shader_wg(gl_idx=("input", "GlobalInvocationId", ps.ivec3),
gl_wg_id=("input", "WorkgroupId", ps.ivec3),
gl_wg_num=("input", "NumWorkgroups", ps.ivec3),
data1=("buffer", 0, ps.Array(ps.f32)),
data2=("buffer", 1, ps.Array(ps.f32))):
i = gl_wg_id.x * gl_wg_num.y + gl_wg_id.y
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], shader_src, (16,8,1))
seq.record_algo_data([tensor_a, tensor_b], compute_shader_wg.to_spirv(), workgroup=(16,8,1))
seq.end()
seq.eval()
@ -161,6 +118,9 @@ def test_workgroup():
mgr.eval_tensor_sync_local_def([tensor_a, tensor_b])
print(tensor_a.numpy())
print(tensor_b.numpy())
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())
@ -183,7 +143,9 @@ def test_tensor_rebuild_backwards_compat():
mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out])
mgr.eval_algo_mult_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])