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

Added python functionality for shader compile

Browse source
This commit is contained in:
Alejandro Saucedo 2021-02-20 19:34:16 +00:00
parent 8035bbb527
commit b9ca70b8b1
2 changed files with 37 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

22
python/src/main.cpp
View file

@ -31,6 +31,18 @@ PYBIND11_MODULE(kp, m) {
.export_values();
py::class_<kp::Shader>(m, "Shader", "Shader class")
.def_static("compile_source", [](
const std::string& source,
const std::string& entryPoint,
const std::vector<std::pair<std::string,std::string>>& definitions) {
std::vector<uint32_t> spirv = kp::Shader::compile_source(source, entryPoint, definitions);
return py::bytes((const char*)spirv.data(), spirv.size() * sizeof(uint32_t));
},
"Compiles string source provided and returns the value in bytes",
py::arg("source"), py::arg("entryPoint") = "main", py::arg("definitions") = std::vector<std::pair<std::string,std::string>>() )
.def_static("compile_sources", &kp::Shader::compile_sources);
py::class_<kp::Tensor, std::shared_ptr<kp::Tensor>>(m, "Tensor", DOC(kp, Tensor))
.def(py::init(
[np](const py::array_t<float> data, kp::Tensor::TensorTypes tensor_type) {
@ -127,7 +139,7 @@ PYBIND11_MODULE(kp, m) {
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
return self.record<kp::OpAlgoBase>(
tensors, std::vector<char>(data, data + length), workgroup, constants);
tensors, std::vector<uint32_t>((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants);
},
"Records an operation using a custom shader provided as spirv bytes",
py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() );
@ -205,7 +217,7 @@ PYBIND11_MODULE(kp, m) {
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), workgroup, constants);
tensors, std::vector<uint32_t>((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants);
},
"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() )
@ -235,7 +247,7 @@ PYBIND11_MODULE(kp, m) {
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), workgroup, constants);
tensors, sequenceName, std::vector<uint32_t>((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants);
},
"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() )
@ -264,7 +276,7 @@ PYBIND11_MODULE(kp, m) {
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), workgroup, constants);
tensors, std::vector<uint32_t>((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants);
},
"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() )
@ -294,7 +306,7 @@ PYBIND11_MODULE(kp, m) {
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), workgroup, constants);
tensors, sequenceName, std::vector<uint32_t>((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants);
},
"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() );