cuButterfly/examples/ntt_device_api.cpp at main · TruNcat3/cuButterfly · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
#include <cuda_runtime_api.h>

#include <cstdint>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

#include "cuntt/ntt.hpp"

namespace {

void check_cuda(cudaError_t status, const char* operation) {
    if (status != cudaSuccess) {
        throw std::runtime_error(std::string(operation) + ": " + cudaGetErrorString(status));
    }
}

}  // namespace

int main() {
    cudaStream_t stream = nullptr;
    void* input = nullptr;
    void* output = nullptr;
    void* workspace = nullptr;

    try {
        cuntt::PlanConfig config;
        config.log_n                   = 16;
        config.batch                   = 2;
        config.backend                 = cuntt::Backend::Tile256;
        config.word_bits               = 64;
        config.auto_allocate_workspace = false;

        cuntt::Plan plan(config);
        check_cuda(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), "create stream");
        check_cuda(cudaMalloc(&input, plan.data_size()), "allocate input");
        check_cuda(cudaMalloc(&output, plan.data_size()), "allocate output");
        if (plan.workspace_size() != 0) {
            check_cuda(cudaMalloc(&workspace, plan.workspace_size()), "allocate workspace");
            plan.set_workspace(workspace, plan.workspace_size());
        }
        plan.set_stream(stream);

        std::vector<std::uint64_t> host_input(plan.data_size() / sizeof(std::uint64_t), 0);
        std::vector<std::uint64_t> host_output(host_input.size());
        host_input.front() = 1;
        check_cuda(cudaMemcpyAsync(input, host_input.data(), plan.data_size(), cudaMemcpyHostToDevice, stream), "copy input");
        plan.execute_async(static_cast<const std::uint64_t*>(input), static_cast<std::uint64_t*>(output));
        check_cuda(cudaMemcpyAsync(host_output.data(), output, plan.data_size(), cudaMemcpyDeviceToHost, stream), "copy output");
        check_cuda(cudaStreamSynchronize(stream), "wait for transform");

        std::cout << "NTT[0] = " << host_output.front() << '\n';
        cudaFree(workspace);
        cudaFree(output);
        cudaFree(input);
        cudaStreamDestroy(stream);
        return 0;
    } catch (const std::exception& error) {
        cudaFree(workspace);
        cudaFree(output);
        cudaFree(input);
        if (stream != nullptr) {
            cudaStreamDestroy(stream);
        }
        std::cerr << "NTT device API example failed: " << error.what() << '\n';
        return 1;
    }
}