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

Fully functional python

Browse source 
Signed-off-by: Alejandro Saucedo <axsauze@gmail.com>
This commit is contained in:
Alejandro Saucedo 2021-09-12 15:00:07 +01:00
parent b9e40d5028
commit 2d4c2f7333
2 changed files with 69 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

30
python/src/main.cpp
View file

@ -14,6 +14,31 @@ namespace py = pybind11;
//used in Core.hpp
py::object kp_debug, kp_info, kp_warning, kp_error;
std::unique_ptr<kp::OpAlgoDispatch> opAlgoDispatchPyInit(
std::shared_ptr<kp::Algorithm>& algorithm,
const py::array& push_consts) {
const py::buffer_info info = push_consts.request();
KP_LOG_DEBUG("Kompute Python Manager creating tensor_T with push_consts size {} dtype {}",
push_consts.size(), std::string(py::str(push_consts.dtype())));
if (push_consts.dtype() == py::dtype::of<std::float_t>()) {
std::vector<float> dataVec((float*)info.ptr, ((float*)info.ptr) + info.size);
return std::unique_ptr<kp::OpAlgoDispatch>{new kp::OpAlgoDispatch(algorithm, dataVec)};
} else if (push_consts.dtype() == py::dtype::of<std::uint32_t>()) {
std::vector<uint32_t> dataVec((uint32_t*)info.ptr, ((uint32_t*)info.ptr) + info.size);
return std::unique_ptr<kp::OpAlgoDispatch>{new kp::OpAlgoDispatch(algorithm, dataVec)};
} else if (push_consts.dtype() == py::dtype::of<std::int32_t>()) {
std::vector<int32_t> dataVec((int32_t*)info.ptr, ((int32_t*)info.ptr) + info.size);
return std::unique_ptr<kp::OpAlgoDispatch>{new kp::OpAlgoDispatch(algorithm, dataVec)};
} else if (push_consts.dtype() == py::dtype::of<std::double_t>()) {
std::vector<double> dataVec((double*)info.ptr, ((double*)info.ptr) + info.size);
return std::unique_ptr<kp::OpAlgoDispatch>{new kp::OpAlgoDispatch(algorithm, dataVec)};
} else {
throw std::runtime_error("Kompute Python no valid dtype supported");
}
}
PYBIND11_MODULE(kp, m) {
// The logging modules are used in the Kompute.hpp file
@ -51,7 +76,10 @@ PYBIND11_MODULE(kp, m) {
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::arg("algorithm"), py::arg("push_consts") = kp::Constants())
.def(py::init(&opAlgoDispatchPyInit),
DOC(kp, OpAlgoDispatch, OpAlgoDispatch),
py::arg("algorithm"), py::arg("push_consts"));
py::class_<kp::OpMult, std::shared_ptr<kp::OpMult>>(
m, "OpMult", py::base<kp::OpBase>(), DOC(kp, OpMult))

41
python/test/test_kompute.py
View file

@ -197,10 +197,49 @@ def test_pushconsts():
.record(kp.OpTensorSyncDevice([tensor]))
.record(kp.OpAlgoDispatch(algo))
.record(kp.OpAlgoDispatch(algo, [0.3, 0.2, 0.1]))
.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))
assert np.allclose(tensor.data(), np.array([0.7, 0.6, 0.5], dtype=np.float32))
def test_pushconsts_int():
spirv = compile_source("""
#version 450
layout(push_constant) uniform PushConstants {
int x;
int y;
int z;
} pcs;
layout (local_size_x = 1) in;
layout(set = 0, binding = 0) buffer a { int pa[]; };
void main() {
pa[0] += pcs.x;
pa[1] += pcs.y;
pa[2] += pcs.z;
}
""")
mgr = kp.Manager()
tensor = mgr.tensor_t(np.array([0, 0, 0], dtype=np.int32))
spec_consts = np.array([], dtype=np.int32)
push_consts = np.array([-1, -1, -1], dtype=np.int32)
algo = mgr.algorithm_t([tensor], spirv, (1, 1, 1), spec_consts, push_consts)
(mgr.sequence()
.record(kp.OpTensorSyncDevice([tensor]))
.record(kp.OpAlgoDispatch(algo))
.record(kp.OpAlgoDispatch(algo, np.array([-1, -1, -1], dtype=np.int32)))
.record(kp.OpAlgoDispatch(algo, np.array([-1, -1, -1], dtype=np.int32)))
.record(kp.OpTensorSyncLocal([tensor]))
.eval())
assert np.all(tensor.data() == np.array([-3, -3, -3], dtype=np.int32))
def test_workgroup():