bench_restrict.cpp

// bench_restrict.cpp — Effect of __restrict__ on inline GEMM codegen
//
// The __restrict__ qualifier tells the compiler that pointers do not alias,
// enabling it to:
//   - Keep loaded values in registers across loop iterations
//   - Avoid redundant loads after stores
//   - Vectorize more aggressively (no alias checks needed)
//
// We compare the same inline GEMM with and without __restrict__.
// Both versions are marked __attribute__((noinline)) to prevent the
// compiler from seeing the allocation context and deducing non-aliasing.
//
// Compile:
//   g++ -O2 -std=c++17 bench_restrict.cpp -o bench_restrict

#include "common.h"

// ---------------------------------------------------------------------------
// WITHOUT __restrict__: compiler must assume input/output may alias
// ---------------------------------------------------------------------------
__attribute__((noinline)) static void gemm_no_restrict(const float* weight, const float* input, float* output,
                                                       int out_ch, int in_ch, int num_frames)
{
  for (int f = 0; f < num_frames; f++)
  {
    const float* in_col = input + f * in_ch;
    float* out_col = output + f * out_ch;
    for (int o = 0; o < out_ch; o++)
    {
      float acc = 0.0f;
      for (int i = 0; i < in_ch; i++)
        acc += weight[i * out_ch + o] * in_col[i];
      out_col[o] = acc;
    }
  }
}

// ---------------------------------------------------------------------------
// WITH __restrict__: compiler knows pointers don't alias
// ---------------------------------------------------------------------------
__attribute__((noinline)) static void gemm_with_restrict(const float* __restrict__ weight,
                                                         const float* __restrict__ input,
                                                         float* __restrict__ output, int out_ch, int in_ch,
                                                         int num_frames)
{
  for (int f = 0; f < num_frames; f++)
  {
    const float* __restrict__ in_col = input + f * in_ch;
    float* __restrict__ out_col = output + f * out_ch;
    for (int o = 0; o < out_ch; o++)
    {
      float acc = 0.0f;
      for (int i = 0; i < in_ch; i++)
        acc += weight[i * out_ch + o] * in_col[i];
      out_col[o] = acc;
    }
  }
}

// ---------------------------------------------------------------------------
// Fully unrolled 4x8 WITHOUT __restrict__
// ---------------------------------------------------------------------------
__attribute__((noinline)) static void gemm_4x8_no_restrict(const float* weight, const float* input, float* output,
                                                           int num_frames)
{
  const float w00 = weight[0], w10 = weight[1], w20 = weight[2], w30 = weight[3];
  const float w01 = weight[4], w11 = weight[5], w21 = weight[6], w31 = weight[7];
  const float w02 = weight[8], w12 = weight[9], w22 = weight[10], w32 = weight[11];
  const float w03 = weight[12], w13 = weight[13], w23 = weight[14], w33 = weight[15];
  const float w04 = weight[16], w14 = weight[17], w24 = weight[18], w34 = weight[19];
  const float w05 = weight[20], w15 = weight[21], w25 = weight[22], w35 = weight[23];
  const float w06 = weight[24], w16 = weight[25], w26 = weight[26], w36 = weight[27];
  const float w07 = weight[28], w17 = weight[29], w27 = weight[30], w37 = weight[31];

  for (int f = 0; f < num_frames; f++)
  {
    const float i0 = input[f * 8], i1 = input[f * 8 + 1];
    const float i2 = input[f * 8 + 2], i3 = input[f * 8 + 3];
    const float i4 = input[f * 8 + 4], i5 = input[f * 8 + 5];
    const float i6 = input[f * 8 + 6], i7 = input[f * 8 + 7];
    output[f * 4] = w00 * i0 + w01 * i1 + w02 * i2 + w03 * i3 + w04 * i4 + w05 * i5 + w06 * i6 + w07 * i7;
    output[f * 4 + 1] = w10 * i0 + w11 * i1 + w12 * i2 + w13 * i3 + w14 * i4 + w15 * i5 + w16 * i6 + w17 * i7;
    output[f * 4 + 2] = w20 * i0 + w21 * i1 + w22 * i2 + w23 * i3 + w24 * i4 + w25 * i5 + w26 * i6 + w27 * i7;
    output[f * 4 + 3] = w30 * i0 + w31 * i1 + w32 * i2 + w33 * i3 + w34 * i4 + w35 * i5 + w36 * i6 + w37 * i7;
  }
}

// ---------------------------------------------------------------------------
// Fully unrolled 4x8 WITH __restrict__
// ---------------------------------------------------------------------------
__attribute__((noinline)) static void gemm_4x8_with_restrict(const float* __restrict__ weight,
                                                             const float* __restrict__ input,
                                                             float* __restrict__ output, int num_frames)
{
  const float w00 = weight[0], w10 = weight[1], w20 = weight[2], w30 = weight[3];
  const float w01 = weight[4], w11 = weight[5], w21 = weight[6], w31 = weight[7];
  const float w02 = weight[8], w12 = weight[9], w22 = weight[10], w32 = weight[11];
  const float w03 = weight[12], w13 = weight[13], w23 = weight[14], w33 = weight[15];
  const float w04 = weight[16], w14 = weight[17], w24 = weight[18], w34 = weight[19];
  const float w05 = weight[20], w15 = weight[21], w25 = weight[22], w35 = weight[23];
  const float w06 = weight[24], w16 = weight[25], w26 = weight[26], w36 = weight[27];
  const float w07 = weight[28], w17 = weight[29], w27 = weight[30], w37 = weight[31];

  for (int f = 0; f < num_frames; f++)
  {
    const float i0 = input[f * 8], i1 = input[f * 8 + 1];
    const float i2 = input[f * 8 + 2], i3 = input[f * 8 + 3];
    const float i4 = input[f * 8 + 4], i5 = input[f * 8 + 5];
    const float i6 = input[f * 8 + 6], i7 = input[f * 8 + 7];
    output[f * 4] = w00 * i0 + w01 * i1 + w02 * i2 + w03 * i3 + w04 * i4 + w05 * i5 + w06 * i6 + w07 * i7;
    output[f * 4 + 1] = w10 * i0 + w11 * i1 + w12 * i2 + w13 * i3 + w14 * i4 + w15 * i5 + w16 * i6 + w17 * i7;
    output[f * 4 + 2] = w20 * i0 + w21 * i1 + w22 * i2 + w23 * i3 + w24 * i4 + w25 * i5 + w26 * i6 + w27 * i7;
    output[f * 4 + 3] = w30 * i0 + w31 * i1 + w32 * i2 + w33 * i3 + w34 * i4 + w35 * i5 + w36 * i6 + w37 * i7;
  }
}

static void run_comparison(int out_ch, int in_ch, int frames, int iters)
{
  char title[128];
  snprintf(title, sizeof(title), "__restrict__ on generic GEMM %dx%d, %d frames", out_ch, in_ch, frames);

  std::vector<float> weight(out_ch * in_ch);
  std::vector<float> input(in_ch * frames);
  std::vector<float> output_a(out_ch * frames);
  std::vector<float> output_b(out_ch * frames);

  fill_random(weight.data(), weight.size());
  fill_random(input.data(), input.size());

  auto r_no = bench(
    [&]() {
      gemm_no_restrict(weight.data(), input.data(), output_a.data(), out_ch, in_ch, frames);
      DoNotOptimize(output_a[0]);
    },
    iters);

  auto r_yes = bench(
    [&]() {
      gemm_with_restrict(weight.data(), input.data(), output_b.data(), out_ch, in_ch, frames);
      DoNotOptimize(output_b[0]);
    },
    iters);

  print_comparison(title, "Without __restrict__", r_no, "With __restrict__", r_yes);
}

int main()
{
  printf("=== Benchmark: Effect of __restrict__ on Inline GEMM ===\n");
  printf("Both versions use noinline to prevent the compiler from deducing\n");
  printf("non-aliasing from the calling context.\n\n");

  const int iters = 5000;

  // Generic GEMM with/without restrict
  for (int frames : {48, 2048})
  {
    printf("\n========== %d frames ==========\n", frames);
    run_comparison(4, 4, frames, iters);
    run_comparison(4, 8, frames, iters);
    run_comparison(8, 8, frames, iters);
  }

  // Fully unrolled 4x8 with/without restrict
  for (int frames : {48, 2048})
  {
    char title[128];
    snprintf(title, sizeof(title), "__restrict__ on unrolled 4x8, %d frames", frames);

    std::vector<float> weight(32);
    std::vector<float> input(8 * frames);
    std::vector<float> output_a(4 * frames);
    std::vector<float> output_b(4 * frames);

    fill_random(weight.data(), weight.size());
    fill_random(input.data(), input.size());

    auto r_no = bench(
      [&]() {
        gemm_4x8_no_restrict(weight.data(), input.data(), output_a.data(), frames);
        DoNotOptimize(output_a[0]);
      },
      iters);

    auto r_yes = bench(
      [&]() {
        gemm_4x8_with_restrict(weight.data(), input.data(), output_b.data(), frames);
        DoNotOptimize(output_b[0]);
      },
      iters);

    print_comparison(title, "Without __restrict__", r_no, "With __restrict__", r_yes);
  }

  return 0;
}