// N can be quite a bit larger than this on most GPUs (~ 1M on MBP M1)

static constexpr size_t N = 1000;

Context ctx = createContext();

// Host-side data

std::array<float, N> theta1Arr, theta2Arr, v1Arr, v2Arr, lengthArr;

std::fill(v1Arr.begin(), v1Arr.end(), 0.0);

std::fill(v2Arr.begin(), v2Arr.end(), 0.0);

for (size_t i = 0; i < N; ++i) {

theta1Arr[i] = 3.14159 / 2 + i * 3.14159 / 16 / N;

theta2Arr[i] = 3.14159 / 2 + i * 3.14159 / 16 / N - 0.1;

lengthArr[i] = 1.0 - i * 0.5 / N;

}

// GPU buffers

Tensor theta1 = createTensor(ctx, Shape{N}, kf32, theta1Arr.data());

Tensor theta2 = createTensor(ctx, Shape{N}, kf32, theta2Arr.data());

Tensor vel1 = createTensor(ctx, Shape{N}, kf32, v1Arr.data());

Tensor vel2 = createTensor(ctx, Shape{N}, kf32, v2Arr.data());

Tensor length = createTensor(ctx, Shape{N}, kf32, lengthArr.data());

std::array<float, 2 * 2 * N> posArr; // x, y outputs for each pendulum

std::string screen(80 * 40, ' ');

Tensor pos = createTensor(ctx, Shape{N * 4}, kf32);

// Prepare computation

KernelCode kernel{kUpdateSim, 256, kf32};

printf("WGSL code: %s\n", kernel.data.c_str());

Kernel update = createKernel(

ctx, kernel, Bindings{theta1, theta2, vel1, vel2, length, pos},

/* nWorkgroups */ cdiv({N, 1, 1}, kernel.workgroupSize));

// Main simulation update loop

printf("\033[2J\033[H");

while (true) {

auto start = std::chrono::high_resolution_clock::now();

std::promise<void> promise;

std::future<void> future = promise.get_future();

dispatchKernel(ctx, update, promise);

wait(ctx, future);

toCPU(ctx, pos, posArr.data(), sizeof(posArr));

auto end = std::chrono::high_resolution_clock::now();

std::chrono::duration<double> elapsed = end - start;

// N * 2 because there's two objects per pendulum

rasterize(posArr.data(), N * 2, 2.0, 2.0, screen, 80, 40);

printf("\033[1;1H" // reset cursor

"# simulations: %lu\n%s",

N, screen.c_str());

resetCommandBuffer(ctx.device, update); // Prepare kernel command

// buffer for nxt iteration

std::this_thread::sleep_for(std::chrono::milliseconds(8) - elapsed);

}