From 711bd158bafb019cd3cc7e44950f7c098418c409 Mon Sep 17 00:00:00 2001 From: Alejandro Saucedo Date: Sun, 14 Feb 2021 16:47:24 +0000 Subject: [PATCH] Updated python to align with new workgrups and constants, as well as removed all string related functions --- python/src/main.cpp | 167 ++++++++++++++---------------------- python/test/test_kompute.py | 88 ++++++------------- 2 files changed, 90 insertions(+), 165 deletions(-) diff --git a/python/src/main.cpp b/python/src/main.cpp index 98609e5cd..c9e5bebf7 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -11,6 +11,8 @@ namespace py = pybind11; py::object kp_debug, kp_info, kp_warning, kp_error; PYBIND11_MODULE(kp, m) { + spdlog::set_level( + static_cast(0)); // The logging modules are used in the Kompute.hpp file py::module_ logging = py::module_::import("logging"); @@ -110,38 +112,27 @@ PYBIND11_MODULE(kp, m) { "Records operation to sync tensor from local memory to GPU memory") .def("record_tensor_sync_local", &kp::Sequence::record, "Records operation to sync tensor(s) from GPU memory to local memory") - .def("record_algo_mult", &kp::Sequence::record, - "Records operation to run multiplication compute shader to two input tensors and an output tensor") - .def("record_algo_file", [](kp::Sequence &self, - std::vector> tensors, - const std::string& file_path, - std::tuple work_group) -> bool { - const kp::OpAlgoBase::KomputeWorkgroup wgroup{ - std::get<0>(work_group), std::get<1>(work_group), std::get<2>(work_group), - }; - return self.record(tensors, file_path, wgroup); - }, - "Records an operation using a custom shader provided from a shader path", - py::arg("tensors"), py::arg("file_path"), py::arg("work_group") = std::make_tuple(0,0,0) ) + .def("record_algo_file", &kp::Sequence::record< + kp::OpAlgoBase, + const std::string&, + kp::Workgroup, + kp::Constants>, + "Records an operation using a custom shader provided from a shader path", + py::arg("tensors"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) .def("record_algo_data", [](kp::Sequence &self, std::vector> tensors, py::bytes &bytes, - std::tuple work_group) -> bool { + kp::Workgroup workgroup, + kp::Constants constants) -> bool { // Bytes have to be converted into std::vector py::buffer_info info(py::buffer(bytes).request()); const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); - const kp::OpAlgoBase::KomputeWorkgroup wgroup{ - std::get<0>(work_group), std::get<1>(work_group), std::get<2>(work_group), - }; return self.record( - tensors, std::vector(data, data + length), wgroup - ); + tensors, std::vector(data, data + length), workgroup, constants); }, "Records an operation using a custom shader provided as spirv bytes", - py::arg("tensors"), py::arg("bytes"), py::arg("work_group") = std::make_tuple(0,0,0) ) - .def("record_algo_lro", &kp::Sequence::record, - "Records operation to run left right out operation with custom shader"); + py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ); py::class_(m, "Manager") @@ -199,33 +190,27 @@ PYBIND11_MODULE(kp, m) { "Evaluates operation to sync tensor from local memory to GPU memory with new anonymous Sequence") .def("eval_tensor_sync_local_def", &kp::Manager::evalOpDefault, "Evaluates operation to sync tensor(s) from GPU memory to local memory with new anonymous Sequence") - .def("eval_algo_mult_def", &kp::Manager::evalOpDefault, - "Evaluates operation to run multiplication compute shader to two input tensors and an output tensor with new anonymous Sequence") - .def("eval_algo_file_def", &kp::Manager::evalOpDefault, - "Evaluates an operation using a custom shader provided from a shader path with new anonymous Sequence") - .def("eval_algo_str_def", &kp::Manager::evalOpDefault>, - "Evaluates an operation using a custom shader provided as string provided as list of characters with new anonymous Sequence") - .def("eval_algo_str_def", [](kp::Manager &self, - std::vector> tensors, - const std::string& shader_str){ - const std::vector shader_vec(shader_str.begin(), shader_str.end()); - self.evalOpDefault(tensors, shader_vec); - }, - "Evaluates an operation using a custom shader provided as string with a new anonymous Sequence") + .def("eval_algo_file_def", &kp::Manager::evalOpDefault< + kp::OpAlgoBase, + const std::string&, + kp::Workgroup, + kp::Constants>, + "Evaluates an operation using a custom shader provided from a shader path with new anonymous Sequence", + py::arg("tensors"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) .def("eval_algo_data_def", [](kp::Manager &self, std::vector> tensors, - py::bytes &bytes) { + py::bytes &bytes, + kp::Workgroup workgroup, + kp::Constants constants) { // Bytes have to be converted into std::vector py::buffer_info info(py::buffer(bytes).request()); const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOpDefault( - tensors, - std::vector(data, data + length)); + tensors, std::vector(data, data + length), workgroup, constants); }, - "Evaluates an operation using a custom shader provided as spirv bytes with new anonymous Sequence") - .def("eval_algo_lro_def", &kp::Manager::evalOpDefault, - "Evaluates operation to run left right out operation with custom shader with new anonymous Sequence") + "Evaluates an operation using a custom shader provided as spirv bytes with new anonymous Sequence", + py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) // eval .def("eval_tensor_copy", &kp::Manager::evalOp, @@ -234,36 +219,28 @@ PYBIND11_MODULE(kp, m) { "Evaluates operation to sync tensor from local memory to GPU memory with explicitly named Sequence") .def("eval_tensor_sync_local", &kp::Manager::evalOp, "Evaluates operation to sync tensor(s) from GPU memory to local memory with explicitly named Sequence") - .def("eval_algo_mult", &kp::Manager::evalOp, - "Evaluates operation to run multiplication compute shader to two input tensors and an output tensor with explicitly named Sequence") - .def("eval_algo_file", &kp::Manager::evalOp, - "Evaluates an operation using a custom shader provided from a shader path with explicitly named Sequence") - .def("eval_algo_str", &kp::Manager::evalOp>, - "Evaluates an operation using a custom shader provided as string provided as list of characters with explicitly named Sequence") - .def("eval_algo_str", [](kp::Manager &self, - std::vector> tensors, - const std::string& sequenceName, - const std::string& shader_str) { - const std::vector shader_vec(shader_str.begin(), shader_str.end()); - self.evalOp(tensors, sequenceName, shader_vec); - }, - "Evaluates an operation using a custom shader provided as string with explicitly named Sequence") + .def("eval_algo_file", &kp::Manager::evalOp< + kp::OpAlgoBase, + const std::string&, + kp::Workgroup, + kp::Constants>, + "Evaluates an operation using a custom shader provided from a shader path with explicitly named Sequence", + py::arg("tensors"), py::arg("sequence_name"), py::arg("data"),py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) .def("eval_algo_data", [](kp::Manager &self, std::vector> tensors, std::string sequenceName, - py::bytes &bytes) { + py::bytes &bytes, + kp::Workgroup workgroup, + kp::Constants constants) { // Bytes have to be converted into std::vector py::buffer_info info(py::buffer(bytes).request()); const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOp( - tensors, - sequenceName, - std::vector(data, data + length)); + tensors, sequenceName, std::vector(data, data + length), workgroup, constants); }, - "Evaluates an operation using a custom shader provided as spirv bytes with explicitly named Sequence") - .def("eval_algo_lro", &kp::Manager::evalOp, - "Evaluates operation to run left right out operation with custom shader with explicitly named Sequence") + "Evaluates an operation using a custom shader provided as spirv bytes with explicitly named Sequence", + py::arg("tensors"), py::arg("sequence_name"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) // eval async default .def("eval_async_tensor_copy_def", &kp::Manager::evalOpAsyncDefault, @@ -272,33 +249,27 @@ PYBIND11_MODULE(kp, m) { "Evaluates asynchronously operation to sync tensor from local memory to GPU memory with anonymous Sequence") .def("eval_async_tensor_sync_local_def", &kp::Manager::evalOpAsyncDefault, "Evaluates asynchronously operation to sync tensor(s) from GPU memory to local memory with anonymous Sequence") - .def("eval_async_algo_mult_def", &kp::Manager::evalOpAsyncDefault, - "Evaluates asynchronously operation to run multiplication compute shader to two input tensors and an output tensor with anonymous Sequence") - .def("eval_async_algo_file_def", &kp::Manager::evalOpAsyncDefault, - "Evaluates asynchronously an operation using a custom shader provided from a shader path with anonymous Sequence") - .def("eval_async_algo_str_def", &kp::Manager::evalOpAsyncDefault>, - "Evaluates Asynchronously an operation using a custom shader provided as string provided as list of characters with new anonymous Sequence") - .def("eval_async_algo_str_def", [](kp::Manager &self, - std::vector> tensors, - const std::string& shader_str) { - const std::vector shader_vec(shader_str.begin(), shader_str.end()); - self.evalOpAsyncDefault(tensors, shader_vec); - }, - "Evaluates Asynchronously an operation using a custom shader provided as string with new anonymous Sequence") + .def("eval_async_algo_file_def", &kp::Manager::evalOpAsyncDefault< + kp::OpAlgoBase, + const std::string&, + kp::Workgroup, + kp::Constants>, + "Evaluates asynchronously an operation using a custom shader provided from a shader path with anonymous Sequence", + py::arg("tensors"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) .def("eval_async_algo_data_def", [](kp::Manager &self, std::vector> tensors, - py::bytes &bytes) { + py::bytes &bytes, + kp::Workgroup workgroup, + kp::Constants constants) { // Bytes have to be converted into std::vector py::buffer_info info(py::buffer(bytes).request()); const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOpAsyncDefault( - tensors, - std::vector(data, data + length)); + tensors, std::vector(data, data + length), workgroup, constants); }, - "Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with anonymous Sequence") - .def("eval_async_algo_lro_def", &kp::Manager::evalOpAsyncDefault, - "Evaluates asynchronously operation to run left right out operation with custom shader with anonymous Sequence") + "Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with anonymous Sequence", + py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) // eval async .def("eval_async_tensor_copy", &kp::Manager::evalOpAsync, @@ -307,36 +278,28 @@ PYBIND11_MODULE(kp, m) { "Evaluates asynchronously operation to sync tensor from local memory to GPU memory with explicitly named Sequence") .def("eval_async_tensor_sync_local", &kp::Manager::evalOpAsync, "Evaluates asynchronously operation to sync tensor(s) from GPU memory to local memory with explicitly named Sequence") - .def("eval_async_algo_mult", &kp::Manager::evalOpAsync, - "Evaluates asynchronously operation to run multiplication compute shader to two input tensors and an output tensor with explicitly named Sequence") - .def("eval_async_algo_file", &kp::Manager::evalOpAsync, - "Evaluates asynchronously an operation using a custom shader provided from a shader path with explicitly named Sequence") - .def("eval_async_algo_str", &kp::Manager::evalOpAsync>, - "Evaluates Asynchronous an operation using a custom shader provided as string provided as list of characters with explicitly named Sequence") - .def("eval_async_algo_str", [](kp::Manager &self, - std::vector> tensors, - const std::string& sequenceName, - const std::string& shader_str) { - const std::vector shader_vec(shader_str.begin(), shader_str.end()); - self.evalOpAsync(tensors, sequenceName, shader_vec); - }, - "Evaluates Asynchronous an operation using a custom shader provided as string with explicitly named Sequence") + .def("eval_async_algo_file", &kp::Manager::evalOpAsync< + kp::OpAlgoBase, + const std::string&, + kp::Workgroup, + kp::Constants>, + "Evaluates asynchronously an operation using a custom shader provided from a shader path with explicitly named Sequence", + py::arg("tensors"), py::arg("sequence_name"), py::arg("data"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ) .def("eval_async_algo_data", [](kp::Manager &self, std::vector> tensors, std::string sequenceName, - py::bytes &bytes) { + py::bytes &bytes, + kp::Workgroup workgroup, + kp::Constants constants) { // Bytes have to be converted into std::vector py::buffer_info info(py::buffer(bytes).request()); const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOpAsync( - tensors, - sequenceName, - std::vector(data, data + length)); + tensors, sequenceName, std::vector(data, data + length), workgroup, constants); }, - "Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with explicitly named Sequence") - .def("eval_async_algo_lro", &kp::Manager::evalOpAsync, - "Evaluates asynchronously operation to run left right out operation with custom shader with explicitly named Sequence"); + "Evaluates asynchronously an operation using a custom shader provided as raw string or spirv bytes with explicitly named Sequence", + py::arg("tensors"), py::arg("sequence_name"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ); #ifdef VERSION_INFO m.attr("__version__") = VERSION_INFO; diff --git a/python/test/test_kompute.py b/python/test/test_kompute.py index 6ce5d88b0..b3f1f2878 100644 --- a/python/test/test_kompute.py +++ b/python/test/test_kompute.py @@ -3,10 +3,11 @@ import os import kp import numpy as np import logging +import pyshader as ps DIRNAME = os.path.dirname(os.path.abspath(__file__)) -def test_opmult(): +def test_opalgobase_file(): """ Test basic OpMult operation """ @@ -19,43 +20,8 @@ def test_opmult(): mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out]) - mgr.eval_algo_mult_def([tensor_in_a, tensor_in_b, tensor_out]) - - mgr.eval_tensor_sync_local_def([tensor_out]) - - assert tensor_out.data() == [2.0, 4.0, 6.0] - assert np.all(tensor_out.numpy() == [2.0, 4.0, 6.0]) - -def test_opalgobase_data(): - """ - Test basic OpAlgoBase operation - """ - - tensor_in_a = kp.Tensor([2, 2, 2]) - tensor_in_b = kp.Tensor([1, 2, 3]) - tensor_out = kp.Tensor([0, 0, 0]) - - mgr = kp.Manager() - - shaderData = """ - #version 450 - - layout (local_size_x = 1) in; - - // The input rebuild bind index is relative to index in parameter passed - layout(set = 0, binding = 0) buffer bina { float tina[]; }; - layout(set = 0, binding = 1) buffer binb { float tinb[]; }; - layout(set = 0, binding = 2) buffer bout { float tout[]; }; - - void main() { - uint index = gl_GlobalInvocationID.x; - tout[index] = tina[index] * tinb[index]; - } - """ - - mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out]) - - mgr.eval_algo_str_def([tensor_in_a, tensor_in_b, tensor_out], shaderData) + shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv")) + mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path) mgr.eval_tensor_sync_local_def([tensor_out]) @@ -73,12 +39,11 @@ def test_opalgobase_file(): tensor_out = kp.Tensor([0, 0, 0]) mgr = kp.Manager() - - shaderFilePath = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp") - mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out]) - mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shaderFilePath) + shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp") + + mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path) mgr.eval_tensor_sync_local_def([tensor_out]) @@ -96,8 +61,8 @@ def test_sequence(): mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out]) - shaderFilePath = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp") - mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shaderFilePath) + shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv")) + mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path) mgr.eval_await_def() @@ -131,27 +96,19 @@ def test_workgroup(): mgr.rebuild([tensor_a, tensor_b]) - shader_src = """ - #version 450 - - layout (local_size_x = 1) in; - - // The input rebuild bind index is relative to index in parameter passed - layout(set = 0, binding = 0) writeonly buffer bout { float toutx[]; }; - layout(set = 0, binding = 1) writeonly buffer bout2 { float touty[]; }; - - void main() { - uint index = gl_WorkGroupID.x*gl_NumWorkGroups.y + gl_WorkGroupID.y; - - toutx[index] = gl_GlobalInvocationID.x; - touty[index] = gl_GlobalInvocationID.y; - } - """ - shader_src = bytes(shader_src, encoding='utf8') + @ps.python2shader + def compute_shader_wg(gl_idx=("input", "GlobalInvocationId", ps.ivec3), + gl_wg_id=("input", "WorkgroupId", ps.ivec3), + gl_wg_num=("input", "NumWorkgroups", ps.ivec3), + data1=("buffer", 0, ps.Array(ps.f32)), + data2=("buffer", 1, ps.Array(ps.f32))): + i = gl_wg_id.x * gl_wg_num.y + gl_wg_id.y + data1[i] = f32(gl_idx.x) + data2[i] = f32(gl_idx.y) seq = mgr.sequence("new") seq.begin() - seq.record_algo_data([tensor_a, tensor_b], shader_src, (16,8,1)) + seq.record_algo_data([tensor_a, tensor_b], compute_shader_wg.to_spirv(), workgroup=(16,8,1)) seq.end() seq.eval() @@ -161,6 +118,9 @@ def test_workgroup(): mgr.eval_tensor_sync_local_def([tensor_a, tensor_b]) + print(tensor_a.numpy()) + print(tensor_b.numpy()) + 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()) @@ -183,7 +143,9 @@ def test_tensor_rebuild_backwards_compat(): mgr.eval_tensor_create_def([tensor_in_a, tensor_in_b, tensor_out]) - mgr.eval_algo_mult_def([tensor_in_a, tensor_in_b, tensor_out]) + shader_path = os.path.abspath(os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv")) + mgr.eval_async_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path) + mgr.eval_await_def() mgr.eval_tensor_sync_local_def([tensor_out])