Function main

example/rocprim/device/example_device_search.cpp:23–63

#include "../../example_utils.hpp"

int main()
{
    // Prepare input and output
    common::device_ptr<int> input(std::vector{6, 3, 5, 4, 1, 8, 2, 5, 4, 1});
    common::device_ptr<int> key(std::vector{5, 4, 1});
    common::device_ptr<int> output(1);

    // Get required size of the temporary storage
    size_t temp_storage_size;
    HIP_CHECK(rocprim::search(nullptr,
                              temp_storage_size,
                              input.get(),
                              key.get(),
                              output.get(),
                              input.size(),
                              key.size()));

    // Allocate temporary storage
    common::device_ptr<void> temp_storage(temp_storage_size);

    // Perform search
    HIP_CHECK(rocprim::search(temp_storage.get(),
                              temp_storage_size,
                              input.get(),
                              key.get(),
                              output.get(),
                              input.size(),
                              key.size()));

    // Check for any errors
    HIP_CHECK(hipGetLastError());

    // Wait for the algorithm to finish
    HIP_CHECK(hipDeviceSynchronize());

    // Copy output to host
    auto result = output.load();

    // Check that the result is correct
    ASSERT_TRUE(result[0] == 2);
}