kompute/test/TestOpAlgoLoopsPassingData.cpp

#include "gtest/gtest.h"

#include "kompute/Kompute.hpp"

TEST(TestProcessingIterations, IterateThroughMultipleSumAndCopies) {
    kp::Manager mgr;

    float TOTAL_ITER = 10;

    std::vector<float> testExpectedOutVec = {TOTAL_ITER, TOTAL_ITER, TOTAL_ITER};

    std::shared_ptr<kp::Tensor> tensorA{ new kp::Tensor({ 0, 0, 0 })};
    std::shared_ptr<kp::Tensor> tensorB{ new kp::Tensor({ 0, 0, 0 })};

    std::string shader(R"(
        #version 450

        layout (local_size_x = 1) in;

        layout(set = 0, binding = 0) buffer a { float pa[]; };
        layout(set = 0, binding = 1) buffer b { float pb[]; };

        void main() {
            uint index = gl_GlobalInvocationID.x;
            pb[index] = pa[index] + 1;
        }
    )");

    std::weak_ptr<kp::Sequence> sqWeakPtr = 
        mgr.getOrCreateManagedSequence("default");

    if (std::shared_ptr<kp::Sequence> sq = sqWeakPtr.lock()) 
    {
        sq->begin();

        sq->record<kp::OpTensorCreate>({ tensorA, tensorB });

        sq->end();

        sq->eval();
    }

    std::weak_ptr<kp::Sequence> sqWeakPtr2 = 
        mgr.getOrCreateManagedSequence("run");

    if (std::shared_ptr<kp::Sequence> sq = sqWeakPtr2.lock()) 
    {
        sq->begin();

        sq->record<kp::OpAlgoBase<>>(
                { tensorA, tensorB }, 
                std::vector<char>(shader.begin(), shader.end()));

        sq->record<kp::OpTensorCopy>({tensorB, tensorA});
        sq->end();

        for (size_t i = 0; i < TOTAL_ITER; i++) {
            sq->eval();
        }
    }

    std::weak_ptr<kp::Sequence> sqWeakPtr3 = 
        mgr.getOrCreateManagedSequence("export");

    if (std::shared_ptr<kp::Sequence> sq = sqWeakPtr3.lock()) 
    {
        sq->begin();

        sq->record<kp::OpTensorSyncLocal>({tensorA, tensorB});

        sq->end();

        sq->eval();
    }

    EXPECT_EQ(tensorA->data(), testExpectedOutVec);
}