From b9ca70b8b1b1da7437e1f12750f18fb82e86e921 Mon Sep 17 00:00:00 2001 From: Alejandro Saucedo Date: Sat, 20 Feb 2021 19:34:16 +0000 Subject: [PATCH] Added python functionality for shader compile --- python/src/main.cpp | 22 +++++++++++++++++----- python/test/test_kompute.py | 27 ++++++++++++++++++++------- 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/python/src/main.cpp b/python/src/main.cpp index c70863a..7852b38 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -31,6 +31,18 @@ PYBIND11_MODULE(kp, m) { .export_values(); + py::class_(m, "Shader", "Shader class") + .def_static("compile_source", []( + const std::string& source, + const std::string& entryPoint, + const std::vector>& definitions) { + std::vector spirv = kp::Shader::compile_source(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>() ) + .def_static("compile_sources", &kp::Shader::compile_sources); + py::class_>(m, "Tensor", DOC(kp, Tensor)) .def(py::init( [np](const py::array_t data, kp::Tensor::TensorTypes tensor_type) { @@ -127,7 +139,7 @@ PYBIND11_MODULE(kp, m) { const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); return self.record( - tensors, std::vector(data, data + length), workgroup, constants); + tensors, std::vector((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants); }, "Records an operation using a custom shader provided as spirv bytes", py::arg("tensors"), py::arg("bytes"), py::arg("workgroup") = kp::Workgroup(), py::arg("constants") = kp::Constants() ); @@ -205,7 +217,7 @@ PYBIND11_MODULE(kp, m) { const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOpDefault( - tensors, std::vector(data, data + length), workgroup, constants); + tensors, std::vector((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants); }, "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() ) @@ -235,7 +247,7 @@ PYBIND11_MODULE(kp, m) { const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOp( - tensors, sequenceName, std::vector(data, data + length), workgroup, constants); + tensors, sequenceName, std::vector((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants); }, "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() ) @@ -264,7 +276,7 @@ PYBIND11_MODULE(kp, m) { const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOpAsyncDefault( - tensors, std::vector(data, data + length), workgroup, constants); + tensors, std::vector((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants); }, "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() ) @@ -294,7 +306,7 @@ PYBIND11_MODULE(kp, m) { const char *data = reinterpret_cast(info.ptr); size_t length = static_cast(info.size); self.evalOpAsync( - tensors, sequenceName, std::vector(data, data + length), workgroup, constants); + tensors, sequenceName, std::vector((uint32_t*)data, (uint32_t*)(data + length)), workgroup, constants); }, "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() ); diff --git a/python/test/test_kompute.py b/python/test/test_kompute.py index f6e549f..b998532 100644 --- a/python/test/test_kompute.py +++ b/python/test/test_kompute.py @@ -17,23 +17,36 @@ def test_opalgobase_file(): tensor_out = kp.Tensor([0, 0, 0]) mgr = kp.Manager() - mgr.rebuild([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) + shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv") + + mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path) 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_file(): +def test_shader_str(): """ Test basic OpAlgoBase operation """ + shader = """ +#version 450 +layout(set = 0, binding = 0) buffer tensorLhs {float valuesLhs[];}; +layout(set = 0, binding = 1) buffer tensorRhs {float valuesRhs[];}; +layout(set = 0, binding = 2) buffer tensorOutput { float valuesOutput[];}; +layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in; + +void main() +{ + uint index = gl_GlobalInvocationID.x; + valuesOutput[index] = valuesLhs[index] * valuesRhs[index]; +} + """ + tensor_in_a = kp.Tensor([2, 2, 2]) tensor_in_b = kp.Tensor([1, 2, 3]) tensor_out = kp.Tensor([0, 0, 0]) @@ -41,9 +54,9 @@ def test_opalgobase_file(): mgr = kp.Manager() mgr.rebuild([tensor_in_a, tensor_in_b, tensor_out]) - shader_path = os.path.join(DIRNAME, "../../shaders/glsl/opmult.comp.spv") + spirv = kp.Shader.compile_source(shader) - mgr.eval_algo_file_def([tensor_in_a, tensor_in_b, tensor_out], shader_path) + mgr.eval_algo_data_def([tensor_in_a, tensor_in_b, tensor_out], spirv) mgr.eval_tensor_sync_local_def([tensor_out])