From 48c127d37750ed7033bf0f74972adad0ce099716 Mon Sep 17 00:00:00 2001 From: Mohamed Koubaa Date: Tue, 10 Dec 2024 00:15:26 -0600 Subject: [PATCH] add datatype enum to python module (#404) * add datatype enum to python module Signed-off-by: koubaa * add datatype enum to python module Signed-off-by: koubaa * add doc Signed-off-by: koubaa --------- Signed-off-by: koubaa Co-authored-by: koubaa --- README.md | 1 + docs/overview/python-examples.rst | 5 ++++- docs/overview/python-reference.rst | 13 ++++++++++++ python/src/main.cpp | 33 ++++++++++++++++++++++++++++++ 4 files changed, 51 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 38f357bf..4ff60ebf 100644 --- a/README.md +++ b/README.md @@ -187,6 +187,7 @@ def kompute(shader): # Explicit type constructor supports uint32, int32, double, float and bool tensor_out_a = mgr.tensor_t(np.array([0, 0, 0], dtype=np.uint32)) tensor_out_b = mgr.tensor_t(np.array([0, 0, 0], dtype=np.uint32)) + assert(t_data.data_type() == kp.DataTypes.uint) params = [tensor_in_a, tensor_in_b, tensor_out_a, tensor_out_b] diff --git a/docs/overview/python-examples.rst b/docs/overview/python-examples.rst index e201fc43..687ca5c2 100644 --- a/docs/overview/python-examples.rst +++ b/docs/overview/python-examples.rst @@ -63,7 +63,10 @@ Similarly you can find the same extended example as above: # Can be initialized with List[] or np.Array tensor_in_a = mgr.tensor([2, 2, 2]) tensor_in_b = mgr.tensor([1, 2, 3]) - tensor_out = mgr.tensor([0, 0, 0]) + + # By default, tensors use a float type, but that can be explicitly specified + tensor_out = mgr.tensor_t([0, 0, 0], dtype=np.float32) + assert(tensor_out.data_type() == kp.DataTypes.float) seq = mgr.sequence() seq.eval(kp.OpTensorSyncDevice([tensor_in_a, tensor_in_b, tensor_out])) diff --git a/docs/overview/python-reference.rst b/docs/overview/python-reference.rst index 097b8fe8..7b4aa6ad 100644 --- a/docs/overview/python-reference.rst +++ b/docs/overview/python-reference.rst @@ -39,3 +39,16 @@ TensorType .. automodule:: kp :members: +MemoryTypes +------- + +.. automodule:: kp + :members: + +DataTypes +------- + +.. automodule:: kp + :members: + + diff --git a/python/src/main.cpp b/python/src/main.cpp index 210fbe04..3bddce88 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -65,6 +65,39 @@ PYBIND11_MODULE(kp, m) py::module_ np = py::module_::import("numpy"); + py::enum_(m, "DataTypes") + .value("bool", + kp::Memory::DataTypes::eBool, + DOC(kp, Memory, DataTypes, eBool)) + .value("int", + kp::Memory::DataTypes::eInt, + DOC(kp, Memory, DataTypes, eInt)) + .value("uint", + kp::Memory::DataTypes::eUnsignedInt, + DOC(kp, Memory, DataTypes, eUnsignedInt)) + .value("float", + kp::Memory::DataTypes::eFloat, + DOC(kp, Memory, DataTypes, eFloat)) + .value("double", + kp::Memory::DataTypes::eDouble, + DOC(kp, Memory, DataTypes, eDouble)) + .value("custom", + kp::Memory::DataTypes::eCustom, + DOC(kp, Memory, DataTypes, eCustom)) + .value("short", + kp::Memory::DataTypes::eShort, + DOC(kp, Memory, DataTypes, eShort)) + .value("ushort", + kp::Memory::DataTypes::eUnsignedShort, + DOC(kp, Memory, DataTypes, eUnsignedShort)) + .value("char", + kp::Memory::DataTypes::eChar, + DOC(kp, Memory, DataTypes, eChar)) + .value("uchar", + kp::Memory::DataTypes::eUnsignedChar, + DOC(kp, Memory, DataTypes, eUnsignedChar)) + .export_values(); + py::enum_(m, "MemoryTypes") .value("device", kp::Memory::MemoryTypes::eDevice,