Merge pull request #1 from BrianPetkovsek/master

rq

python/src/kp/__init__.py

@@ -0,0 +1 @@
+from .kp import *

setup.py

@@ -5,7 +5,7 @@
 import sysconfig
 import subprocess
 
-from setuptools import setup, Extension
+from setuptools import setup, Extension, find_packages
 from setuptools.command.build_ext import build_ext
 from distutils.version import LooseVersion
 
@@ -83,7 +83,11 @@ def build_extension(self, ext):
     description='Kompute: Blazing fast, mobile-enabled, asynchronous, and optimized for advanced GPU processing usecases.',
     long_description=long_description,
     long_description_content_type='text/markdown',
-    ext_modules=[CMakeExtension('kp')],
+    ext_modules=[CMakeExtension('kp/kp')],
+    packages = find_packages(where="python/src"),
+    package_dir={
+        'kp': 'python/src/kp'
+    },
     install_requires=[
         "numpy<2.0.0"
     ],

src/CMakeLists.txt

@@ -17,7 +17,8 @@ add_library(kompute Algorithm.cpp
     OpTensorSyncLocal.cpp
     Sequence.cpp
     Tensor.cpp
-    Core.cpp)
+    Core.cpp 
+    TypeContainer.cpp)
 
 add_library(kompute::kompute ALIAS kompute)
 

src/OpTensorCopy.cpp

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

src/Tensor.cpp

@@ -6,9 +6,9 @@
 namespace kp {
 
 std::string
-Tensor::toString(const std::type_info& dt)
+Tensor::toString(std::shared_ptr<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,
-               const std::type_info& dataType,
+               std::shared_ptr<ABCTypeContainer> dataType,
                const TensorTypes& tensorType)
   : mDataType(dataType)
 {
@@ -113,7 +113,7 @@ Tensor::memorySize()
     return this->mSize * this->mDataTypeMemorySize;
 }
 
-const std::type_info&
+std::shared_ptr<ABCTypeContainer>
 Tensor::dataType()
 {
     return this->mDataType;

src/TypeContainer.cpp

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include "kompute/TypeContainer.hpp"
+
+size_t IdCounter::counter = 0;

src/include/CMakeLists.txt

@@ -15,6 +15,8 @@ target_sources(kompute PRIVATE
     kompute/Manager.hpp
     kompute/Sequence.hpp
     kompute/Tensor.hpp
+    kompute/TypeContainer.hpp
+    kompute/ABCTypeContainer.hpp
 
     kompute/operations/OpAlgoDispatch.hpp
     kompute/operations/OpBase.hpp

src/include/kompute/ABCTypeContainer.hpp

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <string>
+
+class ABCTypeContainer
+{
+  public:
+    // Pure Virtual Function
+    virtual bool compare(ABCTypeContainer& obj) = 0;
+    virtual std::string name() = 0;
+};

src/include/kompute/Manager.hpp

@@ -9,6 +9,7 @@
 
 #include "kompute/Sequence.hpp"
 #include "logger/Logger.hpp"
+#include <kompute/ABCTypeContainer.hpp>
 
 #define KP_DEFAULT_SESSION "DEFAULT"
 
@@ -107,29 +108,33 @@ class Manager
     std::shared_ptr<Tensor> tensor(
       void* data,
       uint32_t elementTotalCount,
+      uint32_t elementMemorySize = sizeof(T),
+      std::shared_ptr<ABCTypeContainer> dataType = 
+        std::make_shared<TypeContainer<T>>(),
       Tensor::TensorTypes tensorType = Tensor::TensorTypes::eDevice)
     {
-        return tensor(data,
-                      elementTotalCount,
-                      sizeof(T),
-                      typeid(T),
+        return this->tensor(data,
+                      elementTotalCount, 
+                      elementMemorySize,
+                      dataType,
                       tensorType);
     }
 
     std::shared_ptr<Tensor> tensor(
       void* data,
       uint32_t elementTotalCount,
       uint32_t elementMemorySize,
-      const std::type_info& dataType,
+      std::shared_ptr<ABCTypeContainer> dataType,
       Tensor::TensorTypes tensorType = Tensor::TensorTypes::eDevice)
     {
-        std::shared_ptr<Tensor> tensor{ new kp::Tensor(this->mPhysicalDevice,
-                                                       this->mDevice,
-                                                       data,
-                                                       elementTotalCount,
-                                                       elementMemorySize,
-                                                       dataType,
-                                                       tensorType) };
+        std::shared_ptr<Tensor> tensor =
+          std::make_shared<kp::Tensor>(this->mPhysicalDevice,
+                                       this->mDevice,
+                                       data,
+                                       elementTotalCount,
+                                       elementMemorySize,
+                                       dataType,
+                                       tensorType);
 
         if (this->mManageResources) {
             this->mManagedTensors.push_back(tensor);

src/include/kompute/Tensor.hpp

@@ -6,6 +6,7 @@
 #include <memory>
 #include <string>
 #include <typeinfo>
+#include "TypeContainer.hpp"
 
 namespace kp {
 
@@ -33,7 +34,7 @@ class Tensor
         eStorage = 2, ///< Type is Device memory (only)
     };
 
-    static std::string toString(const std::type_info& dt);
+    static std::string toString(std::shared_ptr<ABCTypeContainer> dt);
     static std::string toString(TensorTypes dt);
 
     /**
@@ -51,7 +52,7 @@ class Tensor
            void* data,
            uint32_t elementTotalCount,
            uint32_t elementMemorySize,
-           const std::type_info& dataType,
+           std::shared_ptr<ABCTypeContainer> dataType,
            const TensorTypes& tensorType = TensorTypes::eDevice);
 
     /**
@@ -194,7 +195,7 @@ class Tensor
      *
      * @return Data type of tensor of type kp::Tensor::TensorDataTypes
      */
-    const std::type_info& dataType();
+    std::shared_ptr<ABCTypeContainer> dataType();
 
     /**
      * Retrieve the raw data via the pointer to the memory that contains the raw
@@ -240,7 +241,7 @@ class Tensor
   protected:
     // -------------- ALWAYS OWNED RESOURCES
     TensorTypes mTensorType;
-    const std::type_info& mDataType;
+    std::shared_ptr<ABCTypeContainer> mDataType;
     uint32_t mSize;
     uint32_t mDataTypeMemorySize;
     void* mRawData;
@@ -334,9 +335,9 @@ class TensorT : public Tensor
         Tensor::setRawData(data.data());
     }
 
-    const std::type_info& dataType() 
+    std::shared_ptr<ABCTypeContainer> dataType() 
     { 
-        return typeid(T);
+        return std::make_shared<TypeContainer<T>>();
     }
 };
 

src/include/kompute/TypeContainer.hpp

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include "ABCTypeContainer.hpp"
+#include <typeinfo>
+
+
+struct IdCounter
+{
+    static size_t counter;
+};
+
+template<typename T>
+class TypeContainer : public ABCTypeContainer, IdCounter
+{
+  private:
+    size_t classId()
+    {
+        static size_t id = counter++;
+        return id;
+    }
+
+  public:
+    TypeContainer() : dt(typeid(T)) {}
+
+    bool compare(ABCTypeContainer& other) override
+    {
+        TypeContainer& obj = static_cast<TypeContainer&>(other);
+        return this->classId() == obj.classId();
+    }
+
+    std::string name() override { return this->dt.name(); }
+
+    const std::type_info& dt;
+};