From 5db4a585501ff7d9bc3b8d1ba3be37c4d0f5c841 Mon Sep 17 00:00:00 2001
From: Adam Treat <treat.adam@gmail.com>
Date: Fri, 21 Jul 2023 21:54:15 -0400
Subject: [PATCH] Allow to set tensors.

---
 src/Algorithm.cpp                 | 5 +++++
 src/include/kompute/Algorithm.hpp | 9 +++++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/Algorithm.cpp b/src/Algorithm.cpp
index 0e89fdb9b..83df93b91 100644
--- a/src/Algorithm.cpp
+++ b/src/Algorithm.cpp
@@ -403,4 +403,9 @@ Algorithm::getTensors()
     return this->mTensors;
 }
 
+void Algorithm::setTensors(const std::vector<std::shared_ptr<Tensor>>& tensors)
+{
+    this->mTensors = tensors;
+}
+
 }
diff --git a/src/include/kompute/Algorithm.hpp b/src/include/kompute/Algorithm.hpp
index 1917dd37a..45ebfe00c 100644
--- a/src/include/kompute/Algorithm.hpp
+++ b/src/include/kompute/Algorithm.hpp
@@ -197,8 +197,12 @@ class Algorithm
     {
         uint32_t memorySize = sizeof(decltype(pushConstants.back()));
         uint32_t size = pushConstants.size();
-
         this->setPushConstants(pushConstants.data(), size, memorySize);
+
+        this->setWorkgroup(
+          this->mWorkgroup, this->mTensors.size() ? this->mTensors[0]->size() : 1);
+
+        this->createParameters();
     }
 
     /**
@@ -210,7 +214,7 @@ class Algorithm
      * @param size The number of data elements provided in the data
      * @param memorySize The memory size of each of the data elements in bytes.
      */
-    void setPushConstants(void* data, uint32_t size, uint32_t memorySize)
+    void setPushConstants(const void* data, uint32_t size, uint32_t memorySize)
     {
 
         uint32_t totalSize = memorySize * size;
@@ -272,6 +276,7 @@ class Algorithm
      * @returns The list of tensors used in the algorithm.
      */
     const std::vector<std::shared_ptr<Tensor>>& getTensors();
+    void setTensors(const std::vector<std::shared_ptr<Tensor>>& tensors);
 
     void destroy();