KomputeProject · BrianPetkovsek · Jun 29, 2023 · Jun 29, 2023 · Jun 29, 2023 · Jul 7, 2023
diff --git a/python/src/main.cpp b/python/src/main.cpp
@@ -6,11 +6,11 @@
 #include <string>
 
 #include <kompute/Kompute.hpp>
+#include <kompute/Buffer.hpp>
 
 #include "docstrings.hpp"
 #include "utils.hpp"
 
-typedef unsigned char byte;
 namespace py = pybind11;
 
 // used in Core.hpp
@@ -26,33 +26,33 @@ opAlgoDispatchPyInit(std::shared_ptr<kp::Algorithm>& algorithm,
                  push_consts.size(),
                  std::string(py::str(push_consts.dtype())));
 
-    std::vector<byte> dataVec((byte*)info.ptr,
-                                 ((byte*)info.ptr) + info.size);
+    Buffer dataVec{ info.ptr,
+                    static_cast<size_t>(info.size),
+                    static_cast<size_t>(push_consts.dtype().itemsize()) };
     return std::unique_ptr<kp::OpAlgoDispatch>{ new kp::OpAlgoDispatch(
       algorithm, dataVec) };
 }
 
 class PyTypeContainer : public ABCTypeContainer
 {
   public:
-    PyTypeContainer(py::object clazz)
-      : clazz_(clazz)
-    {}
+    PyTypeContainer(py::object dtype) : dtype_(dtype) {}
 
-    bool compare(ABCTypeContainer& obj) override
+    bool operator==(const ABCTypeContainer& other) const override
     {
         // Implement comparison logic here
-        auto other = dynamic_cast<PyTypeContainer*>(&obj);
-        return other && clazz_.is(other->clazz_);
+        const PyTypeContainer* obj =
+          dynamic_cast<const PyTypeContainer*>(&other);
+        return obj && this->dtype_.is(obj->dtype_);
     }
 
     std::string name() override
     {
         // Return the name here
-        return py::str(clazz_);
+        return py::str(dtype_);
     }
 
-    py::object clazz_;
+    py::object dtype_;
 };
 
 PYBIND11_MODULE(kp, m)
@@ -137,21 +137,23 @@ PYBIND11_MODULE(kp, m)
       m, "Tensor", DOC(kp, Tensor))
       .def(
         "data",
-        [np](kp::Tensor& self) {
+        [](kp::Tensor& self) {
             // Non-owning container exposing the underlying pointer
-            auto frombuffer = np.attr("frombuffer");
-
-            auto dtype =
-              dynamic_cast<PyTypeContainer*>(&(*self.dataType()))->clazz_;
-
-            auto buffer =
-              py::array(self.size() * dtype.attr("itemsize").cast<size_t>(),
-                                    self.data<byte>(),
-                                    py::cast(&self));
+            PyTypeContainer* typeContainer =
+              dynamic_cast<PyTypeContainer*>(self.dataType());
 
+            if (typeContainer == nullptr) {
+                throw std::runtime_error(
+                "Kompute Python data type not supported");
+            }
 
+            py::object dtype = typeContainer->dtype_;
 
-            return frombuffer(buffer, dtype);
+            py::array buffer = py::array(self.size() * dtype.attr("itemsize").cast<size_t>(),
+                                         self.data<unsigned char>(),
+                                         py::cast(&self)); 
+            buffer.attr("dtype") = dtype;
+            return buffer;
         },
         DOC(kp, Tensor, data))
       .def("size", &kp::Tensor::size, DOC(kp, Tensor, size))
@@ -160,7 +162,13 @@ PYBIND11_MODULE(kp, m)
       .def(
         "data_type",
         [](kp::Tensor& self) { 
-              return dynamic_cast<PyTypeContainer*>(&(*self.dataType()))->clazz_;
+            PyTypeContainer* typeContainer =
+              dynamic_cast<PyTypeContainer*>(self.dataType());
+            if (typeContainer == nullptr) {
+                throw std::runtime_error(
+                  "Kompute Python data type not supported");
+            }
+            return typeContainer->dtype_;
         },
         DOC(kp, Tensor, dataType))
       .def("is_init", &kp::Tensor::isInit, DOC(kp, Tensor, isInit))
@@ -241,13 +249,13 @@ PYBIND11_MODULE(kp, m)
                          "float with data size {}",
                          flatdata.size());
 
-            std::shared_ptr<ABCTypeContainer> typeContainer =
-              std::make_shared<PyTypeContainer>(flatdata.dtype());
-            return self.tensor(info.ptr,
-                               flatdata.size(),
-                               flatdata.dtype().itemsize(),
-                               typeContainer,
-                               tensor_type);
+            ABCTypeContainer* typeContainer =
+              new PyTypeContainer(flatdata.dtype());
+            return self.internal_tensor(info.ptr,
+                                        flatdata.size(),
+                                        flatdata.dtype().itemsize(),
+                                        typeContainer,
+                                        tensor_type);
         },
         DOC(kp, Manager, tensor),
         py::arg("data"),
@@ -265,13 +273,13 @@ PYBIND11_MODULE(kp, m)
                          flatdata.size(),
                          std::string(py::str(flatdata.dtype())));
 
-            std::shared_ptr<ABCTypeContainer> typeContainer =
-              std::make_shared<PyTypeContainer>(flatdata.dtype());
-            return self.tensor(info.ptr,
-                   flatdata.size(),
-                   flatdata.dtype().itemsize(),
-                   typeContainer,
-                   tensor_type);
+            ABCTypeContainer* typeContainer =
+              new PyTypeContainer(flatdata.dtype());
+            return self.internal_tensor(info.ptr,
+                                        flatdata.size(),
+                                        flatdata.dtype().itemsize(),
+                                        typeContainer,
+                                        tensor_type);
         },
         DOC(kp, Manager, tensorT),
         py::arg("data"),
@@ -323,11 +331,14 @@ PYBIND11_MODULE(kp, m)
                          spec_consts.size(),
                          std::string(py::str(spec_consts.dtype())));
 
-            std::vector<byte> specconstsvec(
-              (byte*)specInfo.ptr, ((byte*)specInfo.ptr) + specInfo.size);
+            Buffer pushconstsvec{ pushInfo.ptr,
+                                  static_cast<size_t>(pushInfo.size),
+                           static_cast<size_t>(push_consts.dtype().itemsize()) };
 
-            std::vector<byte> pushconstsvec(
-              (byte*)pushInfo.ptr, ((byte*)pushInfo.ptr) + pushInfo.size);
+            Buffer specconstsvec{ specInfo.ptr,
+                            static_cast<size_t>(specInfo.size),
+                                  static_cast<size_t>(
+                                    spec_consts.dtype().itemsize()) };
 
             return self.algorithm(
               tensors, spirvVec, workgroup, specconstsvec, pushconstsvec);

diff --git a/src/Buffer.cpp b/src/Buffer.cpp
@@ -0,0 +1,3 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kompute/Buffer.hpp"
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -18,7 +18,8 @@ add_library(kompute Algorithm.cpp
     Sequence.cpp
     Tensor.cpp
     Core.cpp 
-    TypeContainer.cpp)
+    TypeContainer.cpp
+    Buffer.cpp)
 
 add_library(kompute::kompute ALIAS kompute)
 

diff --git a/src/OpTensorCopy.cpp b/src/OpTensorCopy.cpp
@@ -17,15 +17,15 @@ OpTensorCopy::OpTensorCopy(const std::vector<std::shared_ptr<Tensor>>& tensors)
           "Kompute OpTensorCopy called with less than 2 tensor");
     }
 
-    std::shared_ptr<ABCTypeContainer> dataType = this->mTensors[0]->dataType();
+    ABCTypeContainer* dataType = this->mTensors[0]->dataType();
 
     uint32_t size = this->mTensors[0]->size();
     for (const std::shared_ptr<Tensor>& tensor : tensors) {
-        if (!(*dataType).compare(*tensor->dataType())) {
+        if (*dataType != *tensor->dataType()) {
             throw std::runtime_error(fmt::format(
               "Attempting to copy tensors of different types from {} to {}",
-              Tensor::toString(dataType),
-              Tensor::toString(tensor->dataType())));
+              Tensor::toString(*dataType),
+              Tensor::toString(*tensor->dataType())));
         }
         if (tensor->size() != size) {
             throw std::runtime_error(fmt::format(

diff --git a/src/Tensor.cpp b/src/Tensor.cpp
@@ -6,9 +6,9 @@
 namespace kp {
 
 std::string
-Tensor::toString(std::shared_ptr<ABCTypeContainer> dt)
+Tensor::toString(ABCTypeContainer& dt)
 {
-    return (*dt).name();
+    return dt.name();
 }
 
 std::string
@@ -31,7 +31,7 @@ Tensor::Tensor(std::shared_ptr<vk::PhysicalDevice> physicalDevice,
                void* data,
                uint32_t elementTotalCount,
                uint32_t elementMemorySize,
-               std::shared_ptr<ABCTypeContainer> dataType,
+               ABCTypeContainer* dataType,
                const TensorTypes& tensorType)
   : mDataType(dataType)
 {
@@ -113,7 +113,7 @@ Tensor::memorySize()
     return this->mSize * this->mDataTypeMemorySize;
 }
 
-std::shared_ptr<ABCTypeContainer>
+ABCTypeContainer*
 Tensor::dataType()
 {
     return this->mDataType;
@@ -502,6 +502,7 @@ Tensor::destroy()
     this->mRawData = nullptr;
     this->mSize = 0;
     this->mDataTypeMemorySize = 0;
+    free(this->mDataType);
 
     if (!this->mDevice) {
         KP_LOG_WARN(

diff --git a/src/TypeContainer.cpp b/src/TypeContainer.cpp
@@ -1,5 +1,3 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include "kompute/TypeContainer.hpp"
-
-size_t IdCounter::counter = 0;
diff --git a/src/include/CMakeLists.txt b/src/include/CMakeLists.txt
@@ -17,6 +17,7 @@ target_sources(kompute PRIVATE
     kompute/Tensor.hpp
     kompute/TypeContainer.hpp
     kompute/ABCTypeContainer.hpp
+    kompute/Buffer.hpp
 
     kompute/operations/OpAlgoDispatch.hpp
     kompute/operations/OpBase.hpp

diff --git a/src/include/kompute/ABCTypeContainer.hpp b/src/include/kompute/ABCTypeContainer.hpp
@@ -6,6 +6,11 @@ class ABCTypeContainer
 {
   public:
     // Pure Virtual Function
-    virtual bool compare(ABCTypeContainer& obj) = 0;
+    virtual bool operator==(const ABCTypeContainer& other) const = 0;
     virtual std::string name() = 0;
+
+    bool operator!=(const ABCTypeContainer& other) const
+    {
+        return !(*this == other);
+    }
 };
diff --git a/src/include/kompute/Algorithm.hpp b/src/include/kompute/Algorithm.hpp
@@ -6,6 +6,7 @@
 #include "fmt/format.h"
 #include "kompute/Tensor.hpp"
 #include "logger/Logger.hpp"
+#include "kompute/Buffer.hpp"
 
 namespace kp {
 
@@ -41,6 +42,37 @@ class Algorithm
               const Workgroup& workgroup = {},
               const std::vector<S>& specializationConstants = {},
               const std::vector<P>& pushConstants = {})
+    {
+        this->initAlgorithm(device,
+                            tensors,
+                            spirv,
+                            workgroup,
+                            Buffer::from_vector(specializationConstants),
+                            Buffer::from_vector(pushConstants));
+    }
+
+    Algorithm(std::shared_ptr<vk::Device> device,
+              const std::vector<std::shared_ptr<Tensor>>& tensors = {},
+              const std::vector<uint32_t>& spirv = {},
+              const Workgroup& workgroup = {},
+              const Buffer& specializationConstants = { 0, 0, 0 },
+              const Buffer& pushConstants = { 0, 0, 0 })
+    {
+        this->initAlgorithm(device,
+                            tensors,
+                            spirv,
+                            workgroup,
+                            specializationConstants,
+                            pushConstants);
+    }
+
+    private:
+    void initAlgorithm(std::shared_ptr<vk::Device> device,
+              const std::vector<std::shared_ptr<Tensor>>& tensors = {},
+              const std::vector<uint32_t>& spirv = {},
+              const Workgroup& workgroup = {},
+              const Buffer& specializationConstants = {0,0,0},
+              const Buffer& pushConstants = {0,0,0})
     {
         KP_LOG_DEBUG("Kompute Algorithm Constructor with device");
 
@@ -64,6 +96,8 @@ class Algorithm
         }
     }
 
+    public:
+
     /**
      *  Rebuild function to reconstruct algorithm with configuration parameters
      * to create the underlying resources.
@@ -86,6 +120,20 @@ class Algorithm
                  const Workgroup& workgroup = {},
                  const std::vector<S>& specializationConstants = {},
                  const std::vector<P>& pushConstants = {})
+    {
+        this->rebuild(tensors,
+                      spirv,
+                      workgroup,
+                      Buffer::from_vector(specializationConstants), 
+                      Buffer::from_vector(pushConstants));
+    }
+
+    private:
+    void rebuild(const std::vector<std::shared_ptr<Tensor>>& tensors,
+                 const std::vector<uint32_t>& spirv,
+                 const Workgroup& workgroup = {},
+                 const Buffer& specializationConstants = {0,0,0},
+                 const Buffer& pushConstants = {0,0,0})
     {
         KP_LOG_DEBUG("Kompute Algorithm rebuild started");
 
@@ -96,13 +144,12 @@ class Algorithm
             if (this->mSpecializationConstantsData) {
                 free(this->mSpecializationConstantsData);
             }
-            uint32_t memorySize =
-              sizeof(decltype(specializationConstants.back()));
+            uint32_t memorySize = specializationConstants.element_size();
             uint32_t size = specializationConstants.size();
             uint32_t totalSize = size * memorySize;
             this->mSpecializationConstantsData = malloc(totalSize);
             memcpy(this->mSpecializationConstantsData,
-                   specializationConstants.data(),
+                   specializationConstants.begin(),
                    totalSize);
             this->mSpecializationConstantsDataTypeMemorySize = memorySize;
             this->mSpecializationConstantsSize = size;
@@ -112,11 +159,11 @@ class Algorithm
             if (this->mPushConstantsData) {
                 free(this->mPushConstantsData);
             }
-            uint32_t memorySize = sizeof(decltype(pushConstants.back()));
+            uint32_t memorySize = pushConstants.element_size();
             uint32_t size = pushConstants.size();
             uint32_t totalSize = size * memorySize;
             this->mPushConstantsData = malloc(totalSize);
-            memcpy(this->mPushConstantsData, pushConstants.data(), totalSize);
+            memcpy(this->mPushConstantsData, pushConstants.begin(), totalSize);
             this->mPushConstantsDataTypeMemorySize = memorySize;
             this->mPushConstantsSize = size;
         }
@@ -135,6 +182,8 @@ class Algorithm
         this->createPipeline();
     }
 
+    public:
+
     /**
      * Destructor for Algorithm which is responsible for freeing and desroying
      * respective pipelines and owned parameter groups.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		// SPDX-License-Identifier: Apache-2.0

		#include "kompute/Buffer.hpp"