diff --git a/src/include/kompute/Algorithm.hpp b/src/include/kompute/Algorithm.hpp index 650d8a46..cb416bde 100644 --- a/src/include/kompute/Algorithm.hpp +++ b/src/include/kompute/Algorithm.hpp @@ -226,7 +226,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; diff --git a/test/TestPushConstant.cpp b/test/TestPushConstant.cpp index d575a62f..d0d85ee4 100644 --- a/test/TestPushConstant.cpp +++ b/test/TestPushConstant.cpp @@ -54,6 +54,53 @@ TEST(TestPushConstants, TestConstantsAlgoDispatchOverride) } } +TEST(TestPushConstants, TestConstantsAlgoSetPushConstants) +{ + { + std::string shader(R"( + #version 450 + layout(push_constant) uniform PushConstants { + float x; + float y; + float z; + } pcs; + layout (local_size_x = 1) in; + layout(set = 0, binding = 0) buffer a { float pa[]; }; + void main() { + pa[0] += pcs.x; + pa[1] += pcs.y; + pa[2] += pcs.z; + })"); + + std::vector spirv = compileSource(shader); + + std::shared_ptr sq = nullptr; + + { + kp::Manager mgr; + + std::shared_ptr> tensor = + mgr.tensor({ 0, 0, 0 }); + + std::shared_ptr algo = mgr.algorithm( + { tensor }, spirv, kp::Workgroup({ 1 }), {}, { 0.0, 0.0, 0.0 }); + + sq = mgr.sequence()->eval({ tensor }); + + // We need to run this in sequence to avoid race condition + // We can't use atomicAdd as swiftshader doesn't support it for + // float + algo->setPushConstants(std::vector{ 0.1, 0.2, 0.3 }); + sq->eval(algo); + algo->setPushConstants(std::vector{ 0.3, 0.2, 0.1 }); + sq->eval(algo); + sq->eval({ tensor }); + + EXPECT_EQ(tensor->vector(), std::vector({ 0.4, 0.4, 0.4 })); + } + } +} + TEST(TestPushConstants, TestConstantsAlgoDispatchNoOverride) { {