Merge pull request #391 from imaginationtech/nocopy_constructor_impro…

…vements Add tests for using uninitialised tensors.
KomputeProject · Aug 21, 2024 · 8f096a7 · 8f096a7
2 parents 1748b82 + 5bd5bef
commit 8f096a7
Show file tree

Hide file tree

Showing 5 changed files with 166 additions and 7 deletions.
diff --git a/src/Tensor.cpp b/src/Tensor.cpp
@@ -139,8 +139,7 @@ Tensor::tensorType()
 bool
 Tensor::isInit()
 {
-    return this->mDevice && this->mPrimaryBuffer && this->mPrimaryMemory &&
-           this->mRawData;
+    return this->mDevice && this->mPrimaryBuffer && this->mPrimaryMemory;
 }
 
 uint32_t
@@ -176,6 +175,9 @@ Tensor::rawData()
 void
 Tensor::setRawData(const void* data)
 {
+    if (!this->mRawData) {
+        this->mapRawData();
+    }
     memcpy(this->mRawData, data, this->memorySize());
 }
 

diff --git a/src/include/kompute/Manager.hpp b/src/include/kompute/Manager.hpp
@@ -95,6 +95,30 @@ class Manager
         return tensor;
     }
 
+    /**
+     * Create a managed tensor that will be destroyed by this manager
+     * if it hasn't been destroyed by its reference count going to zero.
+     *
+     * @param size The number of element in this tensor
+     * @param tensorType The type of tensor to initialize
+     * @returns Shared pointer with initialised tensor
+     */
+    template<typename T>
+    std::shared_ptr<TensorT<T>> tensorT(
+      size_t size,
+      Tensor::TensorTypes tensorType = Tensor::TensorTypes::eDevice)
+    {
+        KP_LOG_DEBUG("Kompute Manager tensor creation triggered");
+
+        std::shared_ptr<TensorT<T>> tensor{ new kp::TensorT<T>(
+          this->mPhysicalDevice, this->mDevice, size, tensorType) };
+
+        if (this->mManageResources) {
+            this->mManagedTensors.push_back(tensor);
+        }
+
+        return tensor;
+    }
     std::shared_ptr<TensorT<float>> tensor(
       const std::vector<float>& data,
       Tensor::TensorTypes tensorType = Tensor::TensorTypes::eDevice)

diff --git a/src/include/kompute/Tensor.hpp b/src/include/kompute/Tensor.hpp
@@ -255,6 +255,10 @@ class Tensor
     template<typename T>
     T* data()
     {
+        if (this->mRawData == nullptr) {
+            this->mapRawData();
+        }
+
         return (T*)this->mRawData;
     }
 
@@ -268,6 +272,10 @@ class Tensor
     template<typename T>
     std::vector<T> vector()
     {
+        if (this->mRawData == nullptr) {
+            this->mapRawData();
+        }
+
         return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
     }
 
@@ -277,7 +285,9 @@ class Tensor
     TensorDataTypes mDataType;
     uint32_t mSize;
     uint32_t mDataTypeMemorySize;
-    void* mRawData;
+    void* mRawData = nullptr;
+
+    void mapRawData();
 
   private:
     // -------------- NEVER OWNED RESOURCES
@@ -318,7 +328,6 @@ class Tensor
     vk::BufferUsageFlags getStagingBufferUsageFlags();
     vk::MemoryPropertyFlags getStagingMemoryPropertyFlags();
 
-    void mapRawData();
     void unmapRawData();
 };
 
@@ -363,6 +372,10 @@ class TensorT : public Tensor
 
     std::vector<T> vector()
     {
+        if (this->mRawData == nullptr) {
+            this->mapRawData();
+        }
+
         return { (T*)this->mRawData, ((T*)this->mRawData) + this->size() };
     }
 

diff --git a/test/TestOpTensorCopy.cpp b/test/TestOpTensorCopy.cpp
@@ -244,3 +244,88 @@ TEST(TestOpTensorCopy, CopyTensorThroughStorageViaAlgorithms)
     // Making sure the GPU holds the same vector
     EXPECT_EQ(tensorIn->vector(), tensorOut->vector());
 }
+
+TEST(TestOpTensorCopy, CopyDeviceToDeviceTensorUninitialised)
+{
+    kp::Manager mgr;
+
+    std::vector<float> testVecA{ 1, 2, 3 };
+
+    std::shared_ptr<kp::Tensor> tensorA = mgr.tensor(testVecA);
+    std::shared_ptr<kp::TensorT<float>> tensorB = mgr.tensorT<float>(testVecA.size());
+
+    EXPECT_TRUE(tensorA->isInit());
+    EXPECT_TRUE(tensorB->isInit());
+
+    mgr.sequence()
+      ->eval<kp::OpTensorSyncDevice>({ tensorA, tensorB })
+      ->eval<kp::OpTensorCopy>({ tensorA, tensorB })
+      ->eval<kp::OpTensorSyncLocal>({ tensorA, tensorB });
+
+    // Making sure the GPU holds the same vector
+    EXPECT_EQ(tensorA->vector<float>(), tensorB->vector());
+}
+
+TEST(TestOpTensorCopy, CopyTensorThroughStorageViaAlgorithmsUninitialisedOutput)
+{
+    kp::Manager mgr;
+
+    std::vector<float> testVecIn{ 9, 1, 3 };
+
+    std::shared_ptr<kp::Tensor> tensorIn = mgr.tensor(testVecIn);
+    std::shared_ptr<kp::TensorT<float>> tensorOut =
+      mgr.tensorT<float>(testVecIn.size());
+    // Tensor storage requires a vector to be passed only to reflect size
+    std::shared_ptr<kp::TensorT<float>> tensorStorage =
+      mgr.tensorT<float>(testVecIn.size(), kp::Tensor::TensorTypes::eStorage);
+
+    EXPECT_TRUE(tensorIn->isInit());
+    EXPECT_TRUE(tensorOut->isInit());
+
+    // Copy to storage tensor through algorithm
+    std::string shaderA = (R"(
+        #version 450
+
+        layout (local_size_x = 1) in;
+
+        // The input tensors bind index is relative to index in parameter passed
+        layout(set = 0, binding = 0) buffer buf_in { float t_in[]; };
+        layout(set = 0, binding = 1) buffer buf_st { float t_st[]; };
+
+        void main() {
+            uint index = gl_GlobalInvocationID.x;
+            t_st[index] = t_in[index];
+        }
+    )");
+
+    auto algoA =
+      mgr.algorithm({ tensorIn, tensorStorage }, compileSource(shaderA));
+
+    // Copy from storage tensor to output tensor
+    std::string shaderB = (R"(
+        #version 450
+
+        layout (local_size_x = 1) in;
+
+        // The input tensors bind index is relative to index in parameter passed
+        layout(set = 0, binding = 0) buffer buf_st { float t_st[]; };
+        layout(set = 0, binding = 1) buffer buf_out { float t_out[]; };
+
+        void main() {
+            uint index = gl_GlobalInvocationID.x;
+            t_out[index] = t_st[index];
+        }
+    )");
+
+    auto algoB =
+      mgr.algorithm({ tensorStorage, tensorOut }, compileSource(shaderB));
+
+    mgr.sequence()
+      ->eval<kp::OpTensorSyncDevice>({ tensorIn })
+      ->eval<kp::OpAlgoDispatch>(algoA)
+      ->eval<kp::OpAlgoDispatch>(algoB)
+      ->eval<kp::OpTensorSyncLocal>({ tensorOut });
+
+    // Making sure the GPU holds the same vector
+    EXPECT_EQ(tensorIn->vector<float>(), tensorOut->vector());
+}
diff --git a/test/TestTensor.cpp b/test/TestTensor.cpp
@@ -5,6 +5,27 @@
 #include "kompute/Kompute.hpp"
 #include "kompute/logger/Logger.hpp"
 
+// Introducing custom struct that can be used for tensors
+struct TestStruct
+{
+    float x;
+    uint32_t y;
+    int32_t z;
+
+    // Creating an == operator overload for the comparison below
+    bool operator==(const TestStruct rhs) const
+    {
+        return this->x == rhs.x && this->y == rhs.y && this->z == rhs.z;
+    }
+};
+// Custom struct needs to be mapped the eCustom datatype
+template<>
+kp::Tensor::TensorDataTypes
+kp::TensorT<TestStruct>::dataType()
+{
+    return Tensor::TensorDataTypes::eCustom;
+}
+
 TEST(TestTensor, ConstructorData)
 {
     kp::Manager mgr;
@@ -18,9 +39,23 @@ TEST(TestTensor, ConstructorData)
 TEST(TestTensor, ReserveData)
 {
     kp::Manager mgr;
-    std::shared_ptr<kp::TensorT<float>> tensor = mgr.tensor(3, sizeof(uint32_t), Tensor::TensorDataType::eUnsignedInt);
-    EXPECT_EQ(tensor->size(), 0);
-    EXPECT_EQ(tensor->dataTypeMemorySize(), sizeof(uint32_t));
+    std::shared_ptr<kp::Tensor> tensor = mgr.tensor(
+      nullptr, 3, sizeof(float), kp::Tensor::TensorDataTypes::eFloat);
+    EXPECT_EQ(tensor->size(), 3);
+    EXPECT_EQ(tensor->dataTypeMemorySize(), sizeof(float));
+
+    std::shared_ptr<kp::Tensor> tensor2 =
+      mgr.tensor(3, sizeof(float), kp::Tensor::TensorDataTypes::eFloat);
+    EXPECT_EQ(tensor2->size(), 3);
+    EXPECT_EQ(tensor2->dataTypeMemorySize(), sizeof(float));
+
+    std::shared_ptr<kp::TensorT<float>> tensor3 = mgr.tensorT<float>(3);
+    EXPECT_EQ(tensor3->size(), 3);
+    EXPECT_EQ(tensor3->dataTypeMemorySize(), sizeof(float));
+
+    std::shared_ptr<kp::TensorT<TestStruct>> tensor4 = mgr.tensorT<TestStruct>(3);
+    EXPECT_EQ(tensor3->size(), 3);
+    EXPECT_EQ(tensor3->dataTypeMemorySize(), sizeof(TestStruct));
 }
 
 TEST(TestTensor, DataTypes)