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

Updated python to align with current configuration

Browse source
This commit is contained in:
Alejandro Saucedo 2021-02-28 14:53:46 +00:00
parent 7dc1f35206
commit 38f356fdae
4 changed files with 50 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

16
python/src/main.cpp
View file

@ -4,6 +4,8 @@
#include <kompute/Kompute.hpp>
#include "fmt/ranges.h"
#include "docstrings.hpp"
namespace py = pybind11;
@ -64,7 +66,8 @@ PYBIND11_MODULE(kp, m) {
.def(py::init<const std::vector<std::shared_ptr<kp::Tensor>>&>());
py::class_<kp::OpAlgoDispatch, std::shared_ptr<kp::OpAlgoDispatch>>(m, "OpAlgoDispatch", py::base<kp::OpBase>())
.def(py::init<const std::shared_ptr<kp::Algorithm>&>());
.def(py::init<const std::shared_ptr<kp::Algorithm>&,const kp::Constants&>(),
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>&>());
@ -73,12 +76,10 @@ PYBIND11_MODULE(kp, m) {
.def("get_tensors", &kp::Algorithm::getTensors)
.def("destroy", &kp::Algorithm::destroy)
.def("get_spec_consts", &kp::Algorithm::getSpecializationConstants)
.def("get_push_consts", &kp::Algorithm::getPushConstants)
.def("is_init", &kp::Algorithm::isInit);
py::class_<kp::Tensor, std::shared_ptr<kp::Tensor>>(m, "Tensor", DOC(kp, Tensor))
.def("data", &kp::Tensor::data, DOC(kp, Tensor, data))
.def("numpy", [](kp::Tensor& self) {
.def("data", [](kp::Tensor& self) {
return py::array(self.data().size(), self.data().data());
}, "Returns stored data as a new numpy array.")
.def("__getitem__", [](kp::Tensor &self, size_t index) -> float { return self.data()[index]; },
@ -150,16 +151,15 @@ 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& push_consts) {
const kp::Constants& spec_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, push_consts);
return self.algorithm(tensors, spirvVec, workgroup, spec_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());
py::arg("tensors"), py::arg("spirv"), py::arg("workgroup") = kp::Workgroup(), py::arg("spec_consts") = kp::Constants());
#ifdef VERSION_INFO
m.attr("__version__") = VERSION_INFO;

4
python/test/test_array_multiplication.py
View file

@ -30,5 +30,5 @@ def test_array_multiplication():
.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])
assert tensor_out.data().tolist() == [2.0, 4.0, 6.0]
assert np.all(tensor_out.data() == [2.0, 4.0, 6.0])

47
python/test/test_kompute.py
View file

@ -69,7 +69,7 @@ void main()
.record(kp.OpTensorSyncLocal(params))
.eval())
assert tensor_out.data() == [2.0, 4.0, 6.0]
assert tensor_out.data().tolist() == [2.0, 4.0, 6.0]
def test_sequence():
"""
@ -116,8 +116,8 @@ def test_sequence():
assert sq.is_init() == False
assert tensor_out.data() == [2.0, 4.0, 6.0]
assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0])
assert tensor_out.data().tolist() == [2.0, 4.0, 6.0]
assert np.all(tensor_out.data() == [2.0, 4.0, 6.0])
tensor_in_a.destroy()
tensor_in_b.destroy()
@ -127,6 +127,39 @@ def test_sequence():
assert tensor_in_b.is_init() == False
assert tensor_out.is_init() == False
def test_pushconsts():
spirv = kp.Shader.compile_source("""
#version 450
layout(push_constant) uniform PushConstants {
float x;
float y;
float z;
} pcs;
layout (local_size_x = 1) in;
layout(set = 0, binding = 0) buffer a { float pa[]; };
void main() {
pa[0] += pcs.x;
pa[1] += pcs.y;
pa[2] += pcs.z;
}
""")
mgr = kp.Manager()
tensor = mgr.tensor([0, 0, 0])
algo = mgr.algorithm([tensor], spirv, (1, 1, 1))
(mgr.sequence()
.record(kp.OpTensorSyncDevice([tensor]))
.record(kp.OpAlgoDispatch(algo, [0.1, 0.2, 0.3]))
.record(kp.OpAlgoDispatch(algo, [0.3, 0.2, 0.1]))
.record(kp.OpTensorSyncLocal([tensor]))
.eval())
assert np.all(tensor.data() == np.array([0.4, 0.4, 0.4], dtype=np.float32))
def test_workgroup():
mgr = kp.Manager(0)
@ -151,9 +184,9 @@ def test_workgroup():
.record(kp.OpTensorSyncLocal([tensor_a, tensor_b]))
.eval())
print(tensor_a.numpy())
print(tensor_b.numpy())
print(tensor_a.data())
print(tensor_b.data())
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())
assert np.all(tensor_a.data() == np.stack([np.arange(16)]*8, axis=1).ravel())
assert np.all(tensor_b.data() == np.stack([np.arange(8)]*16, axis=0).ravel())