Merge pull request #316 from KomputeProject/adding_storage_tensor

Adding storage tensor functionality with tests

CMakeLists.txt

@@ -97,6 +97,7 @@ kompute_option(KOMPUTE_OPT_BUILD_SHADERS "Rebuilds all compute shaders during co
 
 # External components
 kompute_option(KOMPUTE_OPT_USE_BUILT_IN_SPDLOG "Use the built-in version of Spdlog. Requires 'KOMPUTE_OPT_USE_SPDLOG' to be set to ON in order to have any effect." ON)
+kompute_option(KOMPUTE_OPT_SPDLOG_ASYNC_MODE "If spdlog is enabled this allows for selecting whether the default logger setup creates sync or async logger" OFF)
 kompute_option(KOMPUTE_OPT_USE_BUILT_IN_FMT "Use the built-in version of fmt." ON)
 kompute_option(KOMPUTE_OPT_USE_BUILT_IN_GOOGLE_TEST "Use the built-in version of GoogleTest." ON)
 kompute_option(KOMPUTE_OPT_USE_BUILT_IN_PYBIND11 "Use the built-in version of pybind11." ON)

docs/overview/build-system.rst

@@ -36,10 +36,6 @@ This by default configures without any of the extra build tasks (such as buildin
      - Enable if you want to enable installation.
    * - -DKOMPUTE_OPT_BUILD_PYTHON=ON
      - Enable if you want to build python bindings.
-   * - -DKOMPUTE_OPT_LOG_LEVEL="Default"
-     - Internally we use Spdlog or fmt for logging, depending on the value of 'KOMPUTE_OPT_USE_SPDLOG'. The log level used can be changed here. Possible values: 'Trace', 'Debug', 'Info', 'Warn', 'Error', 'Critical', 'Off', 'Default'. If set to 'Off' logging will be deactivated completely. If set to 'Default', the log level will be set to 'Info' for release builds and 'Debug' else.
-   * - -DKOMPUTE_OPT_USE_SPDLOG=OFF
-     - If enabled, logging via KP_LOG_<DEBUG, INFO, etc...> will happen through Spdlog instead of plan fmt.
    * - -DKOMPUTE_OPT_ANDROID_BUILD=ON
      - Enable android compilation flags required.
    * - -DKOMPUTE_OPT_DISABLE_VK_DEBUG_LAYERS=ON
@@ -48,8 +44,14 @@ This by default configures without any of the extra build tasks (such as buildin
      - Whether to check if your driver supports the Vulkan Header version you are linking against. This might be useful in case you build shared on a different system than you run later.
    * - -DKOMPUTE_OPT_BUILD_SHADERS=OFF
      - Rebuilds all compute shaders during compilation and does not use the already precompiled versions. Requires glslangValidator to be installed on your system.
+   * - -DKOMPUTE_OPT_LOG_LEVEL="Default"
+     - Internally we use Spdlog or fmt for logging, depending on the value of 'KOMPUTE_OPT_USE_SPDLOG'. The log level used can be changed here. Possible values: 'Trace', 'Debug', 'Info', 'Warn', 'Error', 'Critical', 'Off', 'Default'. If set to 'Off' logging will be deactivated completely. If set to 'Default', the log level will be set to 'Info' for release builds and 'Debug' else.
+   * - -DKOMPUTE_OPT_USE_SPDLOG=ON
+     - If enabled, logging via KP_LOG_<DEBUG, INFO, etc...> will happen through Spdlog instead of plan fmt.
    * - -DKOMPUTE_OPT_USE_BUILT_IN_SPDLOG=ON
      - Use the built-in version of Spdlog. Requires 'KOMPUTE_OPT_USE_SPDLOG' to be set to ON in order to have any effect.
+   * - -DKOMPUTE_OPT_SPDLOG_ASYNC_MODE=OFF
+     - If enabled we use the async thread mode of spdlog, but by default using sync for consistent messages.
    * - -DKOMPUTE_OPT_USE_BUILT_IN_FMT=ON
      - Use the built-in version of fmt.
    * - -DKOMPUTE_OPT_USE_BUILT_IN_GOOGLE_TEST=ON

src/OpTensorCopy.cpp

@@ -61,10 +61,20 @@ OpTensorCopy::postEval(const vk::CommandBuffer& /*commandBuffer*/)
 {
     KP_LOG_DEBUG("Kompute OpTensorCopy postEval called");
 
+    // Do not copy on CPU side if source is storage tensor
+    if (this->mTensors[0]->tensorType() == kp::Tensor::TensorTypes::eStorage)
+    {
+        KP_LOG_DEBUG("Kompute OpTensorCopy not copying tensor source given it's of eStorage type");
+        return;
+    }
     void* data = this->mTensors[0]->rawData();
 
     // Copy the data from the first tensor into all the tensors
     for (size_t i = 1; i < this->mTensors.size(); i++) {
+        if (this->mTensors[i]->tensorType() == kp::Tensor::TensorTypes::eStorage) {
+            KP_LOG_DEBUG("Kompute OpTensorCopy not copying to tensor dest given it's of eStorage type");
+            continue;
+        }
         this->mTensors[i]->setRawData(data);
     }
 }

src/Tensor.cpp

@@ -87,9 +87,11 @@ Tensor::rebuild(void* data,
     }
 
     this->allocateMemoryCreateGPUResources();
-    this->mapRawData();
 
-    memcpy(this->mRawData, data, this->memorySize());
+    if (this->tensorType() != Tensor::TensorTypes::eStorage) {
+        this->mapRawData();
+        memcpy(this->mRawData, data, this->memorySize());
+    }
 }
 
 Tensor::TensorTypes
@@ -155,7 +157,7 @@ Tensor::mapRawData()
         hostVisibleMemory = this->mStagingMemory;
     } else {
         KP_LOG_WARN(
-          "Kompute Tensor mapping data not supported on storage tensor");
+          "Kompute Tensor mapping data not supported on {} tensor", toString(this->tensorType()));
         return;
     }
 
@@ -182,7 +184,7 @@ Tensor::unmapRawData()
         hostVisibleMemory = this->mStagingMemory;
     } else {
         KP_LOG_WARN(
-          "Kompute Tensor mapping data not supported on storage tensor");
+          "Kompute Tensor mapping data not supported on {} tensor", toString(this->tensorType()));
         return;
     }
 
@@ -520,7 +522,9 @@ Tensor::destroy()
     }
 
     // Unmap the current memory data
-    this->unmapRawData();
+    if (this->tensorType() != Tensor::TensorTypes::eStorage) {
+        this->unmapRawData();
+    }
 
     if (this->mFreePrimaryBuffer) {
         if (!this->mPrimaryBuffer) {

src/logger/CMakeLists.txt

@@ -59,6 +59,10 @@ if(NOT KOMPUTE_OPT_LOG_LEVEL_DISABLED)
         target_compile_definitions(spdlog INTERFACE SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${KOMPUTE_OPT_LOG_LEVEL})
         target_compile_definitions(kp_logger INTERFACE SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${KOMPUTE_OPT_LOG_LEVEL})
         message(STATUS "setting SPDLOG_ACTIVE_LEVEL to SPDLOG_LEVEL_${KOMPUTE_OPT_LOG_LEVEL}")
+
+        if(KOMPUTE_OPT_SPDLOG_ASYNC_MODE)
+            target_compile_definitions(kp_logger INTERFACE KOMPUTE_SPDLOG_ASYNC_LOGGING=1)
+        endif()
     else()
         target_link_libraries(kp_logger PUBLIC fmt::fmt)
     endif()

src/logger/Logger.cpp

@@ -49,13 +49,21 @@ setupLogger()
     console_sink->set_pattern("[%H:%M:%S %z] [%^%=9l%$] [%=15s] %v");
 #endif
     std::vector<spdlog::sink_ptr> sinks{ console_sink };
+    // TODO: Add flag in compile flags
     std::shared_ptr<spdlog::logger> logger =
-      std::make_shared<spdlog::async_logger>(
-        "",
-        sinks.begin(),
-        sinks.end(),
-        spdlog::thread_pool(),
-        spdlog::async_overflow_policy::block);
+#if KOMPUTE_SPDLOG_ASYNC_LOGGING
+          std::make_shared<spdlog::async_logger>(
+            "",
+            sinks.begin(),
+            sinks.end(),
+            spdlog::thread_pool(),
+            spdlog::async_overflow_policy::block);
+#else
+          std::make_shared<spdlog::logger>(
+            "",
+            sinks.begin(),
+            sinks.end());
+#endif
 
     logger->set_level(getLogLevel());
 

test/TestOpTensorCopy.cpp

@@ -157,3 +157,92 @@ TEST(TestOpTensorCopy, SingleTensorShouldFail)
     EXPECT_THROW(mgr.sequence()->eval<kp::OpTensorCopy>({ tensorA }),
                  std::runtime_error);
 }
+
+TEST(TestOpTensorCopy, CopyThroughStorageTensor)
+{
+    kp::Manager mgr;
+
+    std::vector<float> testVecIn{ 9, 1, 3 };
+    std::vector<float> testVecOut{ 0, 0, 0 };
+
+    std::shared_ptr<kp::TensorT<float>> tensorIn = mgr.tensor(testVecIn);
+    std::shared_ptr<kp::TensorT<float>> tensorOut = mgr.tensor(testVecOut);
+    // Tensor storage requires a vector to be passed only to reflect size
+    std::shared_ptr<kp::TensorT<float>> tensorStorage =
+        mgr.tensor({ 0, 0, 0 }, kp::Tensor::TensorTypes::eStorage);
+
+    mgr.sequence()
+        ->eval<kp::OpTensorSyncDevice>({ tensorIn, tensorOut })
+        ->eval<kp::OpTensorCopy>({ tensorIn, tensorStorage })
+        ->eval<kp::OpTensorCopy>({ tensorStorage, tensorOut })
+        ->eval<kp::OpTensorSyncLocal>({ tensorIn, tensorOut });
+
+    // Making sure the GPU holds the same vector
+    EXPECT_EQ(tensorIn->vector(), tensorOut->vector());
+}
+
+TEST(TestOpTensorCopy, CopyTensorThroughStorageViaAlgorithms)
+{
+    kp::Manager mgr;
+
+    std::vector<float> testVecIn{ 9, 1, 3 };
+    std::vector<float> testVecOut{ 0, 0, 0 };
+
+    std::shared_ptr<kp::TensorT<float>> tensorIn = mgr.tensor(testVecIn);
+    std::shared_ptr<kp::TensorT<float>> tensorOut = mgr.tensor(testVecOut);
+    // Tensor storage requires a vector to be passed only to reflect size
+    std::shared_ptr<kp::TensorT<float>> tensorStorage =
+        mgr.tensor({ 0, 0, 0 }, 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(), tensorOut->vector());
+}