diff --git a/test/TestOpTensorCopy.cpp b/test/TestOpTensorCopy.cpp index 22e8130c..9f8de608 100644 --- a/test/TestOpTensorCopy.cpp +++ b/test/TestOpTensorCopy.cpp @@ -157,3 +157,92 @@ TEST(TestOpTensorCopy, SingleTensorShouldFail) EXPECT_THROW(mgr.sequence()->eval({ tensorA }), std::runtime_error); } + +TEST(TestOpTensorCopy, CopyThroughStorageTensor) +{ + kp::Manager mgr; + + std::vector testVecIn{ 9, 1, 3 }; + std::vector testVecOut{ 0, 0, 0 }; + + std::shared_ptr> tensorIn = mgr.tensor(testVecIn); + std::shared_ptr> tensorOut = mgr.tensor(testVecOut); + // Tensor storage requires a vector to be passed only to reflect size + std::shared_ptr> tensorStorage = + mgr.tensor({ 0, 0, 0 }, kp::Tensor::TensorTypes::eStorage); + + mgr.sequence() + ->eval({ tensorIn, tensorOut }) + ->eval({ tensorIn, tensorStorage }) + ->eval({ tensorStorage, tensorOut }) + ->eval({ tensorIn, tensorOut }); + + // Making sure the GPU holds the same vector + EXPECT_EQ(tensorIn->vector(), tensorOut->vector()); +} + +TEST(TestOpTensorCopy, CopyTensorThroughStorageViaAlgorithms) +{ + kp::Manager mgr; + + std::vector testVecIn{ 9, 1, 3 }; + std::vector testVecOut{ 0, 0, 0 }; + + std::shared_ptr> tensorIn = mgr.tensor(testVecIn); + std::shared_ptr> tensorOut = mgr.tensor(testVecOut); + // Tensor storage requires a vector to be passed only to reflect size + std::shared_ptr> 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({ tensorIn }) + ->eval(algoA) + ->eval(algoB) + ->eval({ tensorOut }); + + // Making sure the GPU holds the same vector + EXPECT_EQ(tensorIn->vector(), tensorOut->vector()); +}