nomic-ai/kompute
1
0
Fork 0
mirror of https://github.com/nomic-ai/kompute.git synced 2026-05-11 00:49:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added python functionality for shader compile

Browse Source
This commit is contained in:
Alejandro Saucedo
2021-02-20 19:34:16 +00:00
parent 8035bbb527
commit b9ca70b8b1
2 changed files with 37 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
python/src/main.cpp
View File

@@ -31,6 +31,18 @@ PYBIND11_MODULE(kp, m) {
.export_values();
py::class_<kp::Shader>(m, "Shader", "Shader class")
.def_static("compile_source", [](
const std::string& source,
const std::string& entryPoint,
const std::vector<std::pair<std::string,std::string>>& definitions) {
std::vector<uint32_t> 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<std::pair<std::string,std::string>>() )
.def_static("compile_sources", &kp::Shader::compile_sources);
py::class_<kp::Tensor, std::shared_ptr<kp::Tensor>>(m, "Tensor", DOC(kp, Tensor))
.def(py::init(
[np](const py::array_t<float> data, kp::Tensor::TensorTypes tensor_type) {
@@ -127,7 +139,7 @@ PYBIND11_MODULE(kp, m) {
const char *data = reinterpret_cast<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
return self.record<kp::OpAlgoBase>(
tensors, std::vector<char>(data, data + length), workgroup, constants);
tensors, std::vector<uint32_t>((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<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOpDefault<kp::OpAlgoBase>(
tensors, std::vector<char>(data, data + length), workgroup, constants);
tensors, std::vector<uint32_t>((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<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOp<kp::OpAlgoBase>(
tensors, sequenceName, std::vector<char>(data, data + length), workgroup, constants);
tensors, sequenceName, std::vector<uint32_t>((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<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOpAsyncDefault<kp::OpAlgoBase>(
tensors, std::vector<char>(data, data + length), workgroup, constants);
tensors, std::vector<uint32_t>((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<const char *>(info.ptr);
size_t length = static_cast<size_t>(info.size);
self.evalOpAsync<kp::OpAlgoBase>(
tensors, sequenceName, std::vector<char>(data, data + length), workgroup, constants);
tensors, sequenceName, std::vector<uint32_t>((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() );

27
python/test/test_kompute.py
View File

@@ -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])