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

Updated python

Browse source
This commit is contained in:
Alejandro Saucedo 2021-03-04 08:27:24 +00:00
parent c211c22a78
commit 647f2f1e62
3 changed files with 52 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

7
python/src/main.cpp
View file

@ -151,15 +151,16 @@ PYBIND11_MODULE(kp, m) {
const std::vector<std::shared_ptr<kp::Tensor>>& tensors,
const py::bytes& spirv,
const kp::Workgroup& workgroup,
const kp::Constants& spec_consts) {
const kp::Constants& spec_consts,
const kp::Constants& push_consts) {
py::buffer_info info(py::buffer(spirv).request());
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
std::vector<uint32_t> spirvVec((uint32_t*)data, (uint32_t*)(data + length));
return self.algorithm(tensors, spirvVec, workgroup, spec_consts);
return self.algorithm(tensors, spirvVec, workgroup, spec_consts, push_consts);
},
"Algorithm initialisation function",
py::arg("tensors"), py::arg("spirv"), py::arg("workgroup") = kp::Workgroup(), py::arg("spec_consts") = kp::Constants());
py::arg("tensors"), py::arg("spirv"), py::arg("workgroup") = kp::Workgroup(), py::arg("spec_consts") = kp::Constants(), py::arg("push_consts") = kp::Constants());
#ifdef VERSION_INFO
m.attr("__version__") = VERSION_INFO;