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

Updated python docstrings

Browse source
This commit is contained in:
Alejandro Saucedo 2021-03-06 12:02:26 +00:00
parent fa5dc43b44
commit f569bae998
2 changed files with 178 additions and 104 deletions
Download patch file Download diff file Expand all files Collapse all files

137
python/src/main.cpp
View file

@ -26,9 +26,9 @@ PYBIND11_MODULE(kp, m) {
py::module_ np = py::module_::import("numpy");
py::enum_<kp::Tensor::TensorTypes>(m, "TensorTypes")
.value("device", kp::Tensor::TensorTypes::eDevice, "Tensor holding data in GPU memory.")
.value("host", kp::Tensor::TensorTypes::eHost, "Tensor used for CPU visible GPU data.")
.value("storage", kp::Tensor::TensorTypes::eStorage, "Tensor with host visible gpu memory.")
.value("device", kp::Tensor::TensorTypes::eDevice, DOC(kp, Tensor, TensorTypes, eDevice))
.value("host", kp::Tensor::TensorTypes::eHost, DOC(kp, Tensor, TensorTypes, eHost))
.value("storage", kp::Tensor::TensorTypes::eStorage, DOC(kp, Tensor, TensorTypes, eStorage))
.export_values();
#if !defined(KOMPUTE_DISABLE_SHADER_UTILS) || !KOMPUTE_DISABLE_SHADER_UTILS
@ -37,51 +37,63 @@ PYBIND11_MODULE(kp, m) {
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);
std::vector<uint32_t> spirv = kp::Shader::compileSource(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>>() )
DOC(kp, Shader, compileSource),
py::arg("source"),
py::arg("entryPoint") = "main",
py::arg("definitions") = std::vector<std::pair<std::string,std::string>>() )
.def_static("compile_sources", [](
const std::vector<std::string>& source,
const std::vector<std::string>& files,
const std::string& entryPoint,
const std::vector<std::pair<std::string,std::string>>& definitions) {
std::vector<uint32_t> spirv = kp::Shader::compile_sources(source, files, entryPoint, definitions);
std::vector<uint32_t> spirv = kp::Shader::compileSources(source, files, entryPoint, definitions);
return py::bytes((const char*)spirv.data(), spirv.size() * sizeof(uint32_t));
},
"Compiles sources provided with file names and returns the value in bytes",
py::arg("sources"), py::arg("files") = std::vector<std::string>(), py::arg("entryPoint") = "main", py::arg("definitions") = std::vector<std::pair<std::string,std::string>>() );
DOC(kp, Shader, compileSources),
py::arg("sources"),
py::arg("files") = std::vector<std::string>(),
py::arg("entryPoint") = "main",
py::arg("definitions") = std::vector<std::pair<std::string,std::string>>() );
#endif // KOMPUTE_DISABLE_SHADER_UTILS
py::class_<kp::OpBase, std::shared_ptr<kp::OpBase>>(m, "OpBase");
py::class_<kp::OpBase, std::shared_ptr<kp::OpBase>>(m, "OpBase", DOC(kp, OpBase));
py::class_<kp::OpTensorSyncDevice, std::shared_ptr<kp::OpTensorSyncDevice>>(m, "OpTensorSyncDevice", py::base<kp::OpBase>())
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>());
py::class_<kp::OpTensorSyncDevice, std::shared_ptr<kp::OpTensorSyncDevice>>(
m, "OpTensorSyncDevice", py::base<kp::OpBase>(), DOC(kp, OpTensorSyncDevice))
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>(), DOC(kp, OpTensorSyncDevice, OpTensorSyncDevice));
py::class_<kp::OpTensorSyncLocal, std::shared_ptr<kp::OpTensorSyncLocal>>(m, "OpTensorSyncLocal", py::base<kp::OpBase>())
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>());
py::class_<kp::OpTensorSyncLocal, std::shared_ptr<kp::OpTensorSyncLocal>>(
m, "OpTensorSyncLocal", py::base<kp::OpBase>(), DOC(kp, OpTensorSyncLocal))
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>(), DOC(kp, OpTensorSyncLocal, OpTensorSyncLocal));
py::class_<kp::OpTensorCopy, std::shared_ptr<kp::OpTensorCopy>>(m, "OpTensorCopy", py::base<kp::OpBase>())
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>());
py::class_<kp::OpTensorCopy, std::shared_ptr<kp::OpTensorCopy>>(
m, "OpTensorCopy", py::base<kp::OpBase>(), DOC(kp, OpTensorCopy))
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>(), DOC(kp, OpTensorCopy, OpTensorCopy));
py::class_<kp::OpAlgoDispatch, std::shared_ptr<kp::OpAlgoDispatch>>(m, "OpAlgoDispatch", py::base<kp::OpBase>())
py::class_<kp::OpAlgoDispatch, std::shared_ptr<kp::OpAlgoDispatch>>(
m, "OpAlgoDispatch", py::base<kp::OpBase>(), DOC(kp, OpAlgoDispatch))
.def(py::init<const std::shared_ptr<kp::Algorithm>&,const kp::Constants&>(),
DOC(kp, OpAlgoDispatch, OpAlgoDispatch),
py::arg("algorithm"), py::arg("push_consts") = kp::Constants());
py::class_<kp::OpMult, std::shared_ptr<kp::OpMult>>(m, "OpMult", py::base<kp::OpBase>())
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&,const std::shared_ptr<kp::Algorithm>&>());
py::class_<kp::OpMult, std::shared_ptr<kp::OpMult>>(
m, "OpMult", py::base<kp::OpBase>(), DOC(kp, OpMult))
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&,const std::shared_ptr<kp::Algorithm>&>(),
DOC(kp, OpMult, OpMult));
py::class_<kp::Algorithm, std::shared_ptr<kp::Algorithm>>(m, "Algorithm")
.def("get_tensors", &kp::Algorithm::getTensors)
.def("destroy", &kp::Algorithm::destroy)
.def("get_spec_consts", &kp::Algorithm::getSpecializationConstants)
.def("is_init", &kp::Algorithm::isInit);
py::class_<kp::Algorithm, std::shared_ptr<kp::Algorithm>>(m, "Algorithm", DOC(kp, Algorithm, Algorithm))
.def("get_tensors", &kp::Algorithm::getTensors, DOC(kp, Algorithm, getTensors))
.def("destroy", &kp::Algorithm::destroy, DOC(kp, Algorithm, destroy))
.def("get_spec_consts", &kp::Algorithm::getSpecializationConstants, DOC(kp, Algorithm, getSpecializationConstants))
.def("is_init", &kp::Algorithm::isInit, DOC(kp, Algorithm, isInit));
py::class_<kp::Tensor, std::shared_ptr<kp::Tensor>>(m, "Tensor", DOC(kp, Tensor))
.def("data", [](kp::Tensor& self) {
return py::array(self.data().size(), self.data().data());
}, "Returns stored data as a new numpy array.")
}, DOC(kp, Tensor, data))
.def("__getitem__", [](kp::Tensor &self, size_t index) -> float { return self.data()[index]; },
"When only an index is necessary")
.def("__setitem__", [](kp::Tensor &self, size_t index, float value) {
@ -91,7 +103,7 @@ PYBIND11_MODULE(kp, m) {
const py::buffer_info info = flatdata.request();
const float* ptr = (float*) info.ptr;
self.setData(std::vector<float>(ptr, ptr+flatdata.size()));
}, "Overrides the data in the local Tensor memory.")
}, DOC(kp, Tensor, setData))
.def("__iter__", [](kp::Tensor &self) {
return py::make_iterator(self.data().begin(), self.data().end());
}, py::keep_alive<0, 1>(), // Required to keep alive iterator while exists
@ -112,35 +124,52 @@ PYBIND11_MODULE(kp, m) {
}
return reversed;
})
.def("size", &kp::Tensor::size, "Retrieves the size of the Tensor data as per the local Tensor memory.")
.def("__len__", &kp::Tensor::size, "Retrieves the size of the Tensor data as per the local Tensor memory.")
.def("tensor_type", &kp::Tensor::tensorType, "Retreves the memory type of the tensor.")
.def("is_init", &kp::Tensor::isInit, "Checks whether the tensor GPU memory has been initialised.")
.def("destroy", &kp::Tensor::destroy, "Destroy tensor GPU resources.");
.def("size", &kp::Tensor::size, DOC(kp, Tensor, size))
.def("__len__", &kp::Tensor::size, DOC(kp, Tensor, size))
.def("tensor_type", &kp::Tensor::tensorType, DOC(kp, Tensor, tensorType))
.def("is_init", &kp::Tensor::isInit, DOC(kp, Tensor, isInit))
.def("destroy", &kp::Tensor::destroy, DOC(kp, Tensor, destroy));
py::class_<kp::Sequence, std::shared_ptr<kp::Sequence>>(m, "Sequence")
.def("record", [](kp::Sequence& self, std::shared_ptr<kp::OpBase> op) { return self.record(op); })
.def("eval", [](kp::Sequence& self) { return self.eval(); })
.def("eval", [](kp::Sequence& self, std::shared_ptr<kp::OpBase> op) { return self.eval(op); })
.def("eval_async", [](kp::Sequence& self) { return self.eval(); })
.def("eval_async", [](kp::Sequence& self, std::shared_ptr<kp::OpBase> op) { return self.evalAsync(op); })
.def("eval_await", [](kp::Sequence& self) { return self.evalAwait(); })
.def("eval_await", [](kp::Sequence& self, uint32_t wait) { return self.evalAwait(wait); })
.def("is_recording", &kp::Sequence::isRecording)
.def("is_running", &kp::Sequence::isRunning)
.def("is_init", &kp::Sequence::isInit)
.def("get_timestamps", &kp::Sequence::getTimestamps)
.def("clear", &kp::Sequence::clear)
.def("destroy", &kp::Sequence::destroy);
py::class_<kp::Sequence, std::shared_ptr<kp::Sequence>>(m, "Sequence", DOC(kp, Sequence))
.def("record", [](kp::Sequence& self, std::shared_ptr<kp::OpBase> op) { return self.record(op); },
DOC(kp, Sequence, record))
.def("eval", [](kp::Sequence& self) { return self.eval(); },
DOC(kp, Sequence, eval))
.def("eval", [](kp::Sequence& self, std::shared_ptr<kp::OpBase> op) { return self.eval(op); },
DOC(kp, Sequence, eval))
.def("eval_async", [](kp::Sequence& self) { return self.eval(); },
DOC(kp, Sequence, evalAsync))
.def("eval_async", [](kp::Sequence& self, std::shared_ptr<kp::OpBase> op) { return self.evalAsync(op); },
DOC(kp, Sequence, evalAsync))
.def("eval_await", [](kp::Sequence& self) { return self.evalAwait(); },
DOC(kp, Sequence, evalAwait))
.def("eval_await", [](kp::Sequence& self, uint32_t wait) { return self.evalAwait(wait); },
DOC(kp, Sequence, evalAwait))
.def("is_recording", &kp::Sequence::isRecording,
DOC(kp, Sequence, isRecording))
.def("is_running", &kp::Sequence::isRunning,
DOC(kp, Sequence, isRunning))
.def("is_init", &kp::Sequence::isInit,
DOC(kp, Sequence, isInit))
.def("clear", &kp::Sequence::clear,
DOC(kp, Sequence, clear))
.def("rerecord", &kp::Sequence::rerecord,
DOC(kp, Sequence, rerecord))
.def("get_timestamps", &kp::Sequence::getTimestamps,
DOC(kp, Sequence, getTimestamps))
.def("destroy", &kp::Sequence::destroy,
DOC(kp, Sequence, destroy));
py::class_<kp::Manager, std::shared_ptr<kp::Manager>>(m, "Manager")
.def(py::init())
.def(py::init<uint32_t>())
py::class_<kp::Manager, std::shared_ptr<kp::Manager>>(m, "Manager", DOC(kp, Manager))
.def(py::init(), DOC(kp, Manager, Manager))
.def(py::init<uint32_t>(), DOC(kp, Manager, Manager_2))
.def(py::init<uint32_t,const std::vector<uint32_t>&,const std::vector<std::string>&>(),
DOC(kp, Manager, Manager_2),
py::arg("device") = 0,
py::arg("family_queue_indices") = std::vector<uint32_t>(),
py::arg("desired_extensions") = std::vector<std::string>())
.def("sequence", &kp::Manager::sequence, py::arg("queue_index") = 0, py::arg("total_timestamps") = 0)
.def("sequence", &kp::Manager::sequence, DOC(kp, Manager, sequence),
py::arg("queue_index") = 0, py::arg("total_timestamps") = 0)
.def("tensor", [np](kp::Manager& self,
const py::array_t<float> data,
kp::Tensor::TensorTypes tensor_type) {
@ -149,7 +178,7 @@ PYBIND11_MODULE(kp, m) {
const float* ptr = (float*) info.ptr;
return self.tensor(std::vector<float>(ptr, ptr+flatdata.size()), tensor_type);
},
"Tensor initialisation function with data and tensor type",
DOC(kp, Manager, tensor),
py::arg("data"), py::arg("tensor_type") = kp::Tensor::TensorTypes::eDevice)
.def("algorithm", [](kp::Manager& self,
const std::vector<std::shared_ptr<kp::Tensor>>& tensors,
@ -163,8 +192,12 @@ PYBIND11_MODULE(kp, m) {
std::vector<uint32_t> spirvVec((uint32_t*)data, (uint32_t*)(data + length));
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("push_consts") = kp::Constants());
DOC(kp, Manager, algorithm),
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;