From 2d4c2f733371afb4875b712cbd3f8e04ef33da05 Mon Sep 17 00:00:00 2001
From: Alejandro Saucedo <axsauze@gmail.com>
Date: Sun, 12 Sep 2021 15:00:07 +0100
Subject: [PATCH] Fully functional python

Signed-off-by: Alejandro Saucedo <axsauze@gmail.com>
---
 python/src/main.cpp         | 30 ++++++++++++++++++++++++++-
 python/test/test_kompute.py | 41 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/python/src/main.cpp b/python/src/main.cpp
index b0ef31191..82a4bff5e 100644
--- a/python/src/main.cpp
+++ b/python/src/main.cpp
@@ -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))
diff --git a/python/test/test_kompute.py b/python/test/test_kompute.py
index eaf6b28db..385933f26 100644
--- a/python/test/test_kompute.py
+++ b/python/test/test_kompute.py
@@ -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():