OCCA#

// laplacian.okl

@kernel void laplacian(const int N, const float * in, float * out) {
  for (int bj = 0; bj < N; bj += BLOCK_Y; @outer(1)) {
    for (int bi = 0; bi < N; bi += BLOCK_X; @outer(0)) {
      for (int j = bj; j < bj + BLOCK_Y; ++j; @inner(1)) {
        for (int i = bi; i < bi + BLOCK_X; ++i; @inner(0)) {
          if (i > 0 && i < (N - 1) && j > 0 && j < (N - 1)) {
            const int id = j * N + i;
            out[id] = -4.0f * in[id]
                      + in[id - 1]     // left
                      + in[id + 1]     // right
                      + in[id - N]     // up
                      + in[id + N];    // down
          }
        }
      }
    }
  }
}

// benchmark.cpp
#include <iostream>
#include <iomanip>
#include <cmath>
#include <chrono>
#include <string>

#include <occa.hpp>

struct BenchResult {
  std::string mode;
  double setupTime;
  double kernelTime;
  double bandwidth;
  double gflops;
  double maxError;
};

BenchResult runBenchmark(const std::string &modeJson, int N, int iterations) {
  BenchResult result;
  result.mode = modeJson;

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

  // --- Device & kernel setup ---
  occa::device device;
  device.setup(modeJson);

  occa::json kernelProps;
  kernelProps["defines/BLOCK_X"] = 16;
  kernelProps["defines/BLOCK_Y"] = 16;

  occa::kernel laplacian = device.buildKernel(
    "laplacian.okl", "laplacian", kernelProps
  );

  // --- Allocate and initialize ---
  float *h_in  = new float[N * N];
  float *h_out = new float[N * N];
  float *h_ref = new float[N * N];

  for (int j = 0; j < N; ++j) {
    for (int i = 0; i < N; ++i) {
      float x = (float)i / (N - 1);
      float y = (float)j / (N - 1);
      h_in[j * N + i] = sinf(M_PI * x) * sinf(M_PI * y);
      h_out[j * N + i] = 0.0f;
    }
  }

  occa::memory o_in  = device.malloc<float>(N * N);
  occa::memory o_out = device.malloc<float>(N * N);

  o_in.copyFrom(h_in);
  o_out.copyFrom(h_out);

  auto t1 = std::chrono::high_resolution_clock::now();
  result.setupTime = std::chrono::duration<double, std::milli>(t1 - t0).count();

  // --- Warmup ---
  laplacian(N, o_in, o_out);
  device.finish();

  // --- Timed iterations ---
  auto t2 = std::chrono::high_resolution_clock::now();

  for (int iter = 0; iter < iterations; ++iter) {
    laplacian(N, o_in, o_out);
  }
  device.finish();

  auto t3 = std::chrono::high_resolution_clock::now();
  double elapsed = std::chrono::duration<double, std::milli>(t3 - t2).count();
  result.kernelTime = elapsed / iterations;

  // --- Compute bandwidth and GFLOPS ---
  // Each interior point: 1 read of out, 5 reads of in, 1 write = 7 memory ops
  // (approximate: borders are skipped but negligible for large N)
  long long interiorPoints = (long long)(N - 2) * (N - 2);
  double bytesPerIter = interiorPoints * 7.0 * sizeof(float);
  result.bandwidth = (bytesPerIter / (result.kernelTime * 1e-3)) / 1e9; // GB/s

  // Each interior point: 4 adds + 1 multiply = 5 flops
  double flopsPerIter = interiorPoints * 5.0;
  result.gflops = (flopsPerIter / (result.kernelTime * 1e-3)) / 1e9;

  // --- Verify correctness ---
  o_out.copyTo(h_out);

  // Compute reference on CPU
  for (int j = 1; j < N - 1; ++j) {
    for (int i = 1; i < N - 1; ++i) {
      int id = j * N + i;
      h_ref[id] = -4.0f * h_in[id]
                  + h_in[id - 1] + h_in[id + 1]
                  + h_in[id - N] + h_in[id + N];
    }
  }

  result.maxError = 0.0;
  for (int j = 1; j < N - 1; ++j) {
    for (int i = 1; i < N - 1; ++i) {
      double err = fabs(h_out[j * N + i] - h_ref[j * N + i]);
      if (err > result.maxError) result.maxError = err;
    }
  }

  delete[] h_in;
  delete[] h_out;
  delete[] h_ref;

  return result;
}

void printResult(const BenchResult &r, const std::string &label) {
  std::cout << std::left << std::setw(12) << label
            << " | Setup: " << std::right << std::setw(8) << std::fixed
            << std::setprecision(1) << r.setupTime << " ms"
            << " | Kernel: " << std::setw(8) << std::setprecision(3)
            << r.kernelTime << " ms"
            << " | BW: " << std::setw(6) << std::setprecision(1)
            << r.bandwidth << " GB/s"
            << " | GFLOPS: " << std::setw(6) << std::setprecision(2)
            << r.gflops
            << " | MaxErr: " << std::scientific << std::setprecision(2)
            << r.maxError
            << std::endl;
}

int main(int argc, const char **argv) {
  const int N = 8192;
  const int iterations = 400;

  std::cout << "============================================================\n";
  std::cout << "  OCCA 2D Laplacian Stencil Benchmark\n";
  std::cout << "  Grid: " << N << " x " << N
            << " (" << (N * N * sizeof(float)) / (1024.0 * 1024.0)
            << " MB per field)\n";
  std::cout << "  Iterations: " << iterations << "\n";
  std::cout << "============================================================\n\n";

  // Serial
  auto serial = runBenchmark("{mode: 'Serial'}", N, iterations);
  printResult(serial, "Serial");

  // Metal
  auto metal = runBenchmark("{mode: 'Metal', device_id: 0}", N, iterations);
  printResult(metal, "Metal");

  // OpenCL
  auto opencl = runBenchmark("{mode: 'OpenCL', platform_id: 0, device_id: 0}", N, iterations);
  printResult(opencl, "OpenCL");

  // Summary
  std::cout << "\n--- Speedup over Serial ---\n";
  std::cout << "  Metal  : " << std::fixed << std::setprecision(1)
            << serial.kernelTime / metal.kernelTime << "x\n";
  std::cout << "  OpenCL : " << std::fixed << std::setprecision(1)
            << serial.kernelTime / opencl.kernelTime << "x\n";

  return 0;
}