matmul_kernels_strassen.cu

matmul_kernels_strassen.cu

#include "matmul_kernels.cuh"

#include <array>

namespace {

inline constexpr int kStrassenProductCount = 7;

struct StrassenTerm {
  int row0_a = 0;
  int col0_a = 0;
  float scale_a = 0.0f;
  int row0_b = 0;
  int col0_b = 0;
  float scale_b = 0.0f;
};

struct StrassenProductSpec {
  StrassenTerm lhs;
  StrassenTerm rhs;
};

struct StrassenOperandView {
  const float *ptr = nullptr;
  int leading_dim = 0;
};

struct StrassenWorkspace {
  float *lhs[kStrassenProductCount] = {};
  float *rhs[kStrassenProductCount] = {};
  float *products = nullptr;
  cudaStream_t streams[kStrassenProductCount] = {};
  cudaEvent_t done_events[kStrassenProductCount] = {};
  size_t lhs_elements = 0;
  size_t rhs_elements = 0;
  size_t product_elements = 0;
  bool streams_initialized = false;
};

__global__ void buildStrassenOperandKernel(const float *__restrict__ base, float *__restrict__ out, int rows, int cols,
                                           int leading_dim, int row0_a, int col0_a, float scale_a, int row0_b,
                                           int col0_b, float scale_b) {
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int elements = rows * cols;
  if (idx >= elements) {
    return;
  }

  const int row = idx / cols;
  const int col = idx - row * cols;
  float value = scale_a * base[(row0_a + row) * leading_dim + col0_a + col];
  if (scale_b != 0.0f) {
    value += scale_b * base[(row0_b + row) * leading_dim + col0_b + col];
  }
  out[idx] = value;
}

__global__ void buildStrassenOperandFloat4Kernel(const float *__restrict__ base, float *__restrict__ out, int rows,
                                                 int cols, int leading_dim, int row0_a, int col0_a, float scale_a,
                                                 int row0_b, int col0_b, float scale_b) {
  const int vec_idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int vec_cols = cols / 4;
  const int elements = rows * vec_cols;
  if (vec_idx >= elements) {
    return;
  }

  const int row = vec_idx / vec_cols;
  const int col = (vec_idx - row * vec_cols) * 4;
  auto value = *reinterpret_cast<const float4 *>(base + (row0_a + row) * leading_dim + col0_a + col);
  value.x *= scale_a;
  value.y *= scale_a;
  value.z *= scale_a;
  value.w *= scale_a;
  if (scale_b != 0.0f) {
    const auto other = *reinterpret_cast<const float4 *>(base + (row0_b + row) * leading_dim + col0_b + col);
    value.x += scale_b * other.x;
    value.y += scale_b * other.y;
    value.z += scale_b * other.z;
    value.w += scale_b * other.w;
  }
  *reinterpret_cast<float4 *>(out + row * cols + col) = value;
}

__global__ void buildStrassenLhsBatchedFloat4Kernel(const float *__restrict__ a, float *__restrict__ lhs1,
                                                    float *__restrict__ lhs2, float *__restrict__ lhs5,
                                                    float *__restrict__ lhs6, float *__restrict__ lhs7, int m2, int k2,
                                                    int leading_dim) {
  const int vec_idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int vec_cols = k2 / 4;
  const int elements = m2 * vec_cols;
  if (vec_idx >= elements) {
    return;
  }

  const int row = vec_idx / vec_cols;
  const int col = (vec_idx - row * vec_cols) * 4;
  const auto a11 = *reinterpret_cast<const float4 *>(a + row * leading_dim + col);
  const auto a12 = *reinterpret_cast<const float4 *>(a + row * leading_dim + k2 + col);
  const auto a21 = *reinterpret_cast<const float4 *>(a + (m2 + row) * leading_dim + col);
  const auto a22 = *reinterpret_cast<const float4 *>(a + (m2 + row) * leading_dim + k2 + col);
  const int out = row * k2 + col;
  *reinterpret_cast<float4 *>(lhs1 + out) = make_float4(a11.x + a22.x, a11.y + a22.y, a11.z + a22.z, a11.w + a22.w);
  *reinterpret_cast<float4 *>(lhs2 + out) = make_float4(a21.x + a22.x, a21.y + a22.y, a21.z + a22.z, a21.w + a22.w);
  *reinterpret_cast<float4 *>(lhs5 + out) = make_float4(a11.x + a12.x, a11.y + a12.y, a11.z + a12.z, a11.w + a12.w);
  *reinterpret_cast<float4 *>(lhs6 + out) = make_float4(a21.x - a11.x, a21.y - a11.y, a21.z - a11.z, a21.w - a11.w);
  *reinterpret_cast<float4 *>(lhs7 + out) = make_float4(a12.x - a22.x, a12.y - a22.y, a12.z - a22.z, a12.w - a22.w);
}

__global__ void buildStrassenRhsBatchedFloat4Kernel(const float *__restrict__ b, float *__restrict__ rhs1,
                                                    float *__restrict__ rhs3, float *__restrict__ rhs4,
                                                    float *__restrict__ rhs6, float *__restrict__ rhs7, int k2, int n2,
                                                    int leading_dim) {
  const int vec_idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int vec_cols = n2 / 4;
  const int elements = k2 * vec_cols;
  if (vec_idx >= elements) {
    return;
  }

  const int row = vec_idx / vec_cols;
  const int col = (vec_idx - row * vec_cols) * 4;
  const auto b11 = *reinterpret_cast<const float4 *>(b + row * leading_dim + col);
  const auto b12 = *reinterpret_cast<const float4 *>(b + row * leading_dim + n2 + col);
  const auto b21 = *reinterpret_cast<const float4 *>(b + (k2 + row) * leading_dim + col);
  const auto b22 = *reinterpret_cast<const float4 *>(b + (k2 + row) * leading_dim + n2 + col);
  const int out = row * n2 + col;
  *reinterpret_cast<float4 *>(rhs1 + out) = make_float4(b11.x + b22.x, b11.y + b22.y, b11.z + b22.z, b11.w + b22.w);
  *reinterpret_cast<float4 *>(rhs3 + out) = make_float4(b12.x - b22.x, b12.y - b22.y, b12.z - b22.z, b12.w - b22.w);
  *reinterpret_cast<float4 *>(rhs4 + out) = make_float4(b21.x - b11.x, b21.y - b11.y, b21.z - b11.z, b21.w - b11.w);
  *reinterpret_cast<float4 *>(rhs6 + out) = make_float4(b11.x + b12.x, b11.y + b12.y, b11.z + b12.z, b11.w + b12.w);
  *reinterpret_cast<float4 *>(rhs7 + out) = make_float4(b21.x + b22.x, b21.y + b22.y, b21.z + b22.z, b21.w + b22.w);
}

__global__ void combineStrassenProductsKernel(const float *__restrict__ p1, const float *__restrict__ p2,
                                              const float *__restrict__ p3, const float *__restrict__ p4,
                                              const float *__restrict__ p5, const float *__restrict__ p6,
                                              const float *__restrict__ p7, float *__restrict__ c, int m2, int n2,
                                              int n) {
  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int elements = m2 * n2;
  if (idx >= elements) {
    return;
  }

  const int row = idx / n2;
  const int col = idx - row * n2;
  const float v1 = p1[idx];
  const float v2 = p2[idx];
  const float v3 = p3[idx];
  const float v4 = p4[idx];
  const float v5 = p5[idx];
  const float v6 = p6[idx];
  const float v7 = p7[idx];

  c[row * n + col] = (v1 + v4) - v5 + v7;
  c[row * n + col + n2] = v3 + v5;
  c[(row + m2) * n + col] = v2 + v4;
  c[(row + m2) * n + col + n2] = (v1 - v2) + v3 + v6;
}

__global__ void combineStrassenProductsFloat4Kernel(const float *__restrict__ p1, const float *__restrict__ p2,
                                                    const float *__restrict__ p3, const float *__restrict__ p4,
                                                    const float *__restrict__ p5, const float *__restrict__ p6,
                                                    const float *__restrict__ p7, float *__restrict__ c, int m2, int n2,
                                                    int n) {
  const int vec_idx = blockIdx.x * blockDim.x + threadIdx.x;
  const int vec_cols = n2 / 4;
  const int elements = m2 * vec_cols;
  if (vec_idx >= elements) {
    return;
  }

  const int row = vec_idx / vec_cols;
  const int col = (vec_idx - row * vec_cols) * 4;
  const int idx = row * n2 + col;
  const auto v1 = *reinterpret_cast<const float4 *>(p1 + idx);
  const auto v2 = *reinterpret_cast<const float4 *>(p2 + idx);
  const auto v3 = *reinterpret_cast<const float4 *>(p3 + idx);
  const auto v4 = *reinterpret_cast<const float4 *>(p4 + idx);
  const auto v5 = *reinterpret_cast<const float4 *>(p5 + idx);
  const auto v6 = *reinterpret_cast<const float4 *>(p6 + idx);
  const auto v7 = *reinterpret_cast<const float4 *>(p7 + idx);

  *reinterpret_cast<float4 *>(c + row * n + col) =
      make_float4((v1.x + v4.x) - v5.x + v7.x, (v1.y + v4.y) - v5.y + v7.y, (v1.z + v4.z) - v5.z + v7.z,
                  (v1.w + v4.w) - v5.w + v7.w);
  *reinterpret_cast<float4 *>(c + row * n + col + n2) = make_float4(v3.x + v5.x, v3.y + v5.y, v3.z + v5.z, v3.w + v5.w);
  *reinterpret_cast<float4 *>(c + (row + m2) * n + col) =
      make_float4(v2.x + v4.x, v2.y + v4.y, v2.z + v4.z, v2.w + v4.w);
  *reinterpret_cast<float4 *>(c + (row + m2) * n + col + n2) =
      make_float4((v1.x - v2.x) + v3.x + v6.x, (v1.y - v2.y) + v3.y + v6.y, (v1.z - v2.z) + v3.z + v6.z,
                  (v1.w - v2.w) + v3.w + v6.w);
}

__device__ __forceinline__ float4 loadStrassenTerm4(const float *__restrict__ base, int leading_dim,
                                                    const StrassenTerm &term, int row, int col) {
  auto value = *reinterpret_cast<const float4 *>(base + (term.row0_a + row) * leading_dim + term.col0_a + col);
  value.x *= term.scale_a;
  value.y *= term.scale_a;
  value.z *= term.scale_a;
  value.w *= term.scale_a;
  if (term.scale_b != 0.0f) {
    const auto other = *reinterpret_cast<const float4 *>(base + (term.row0_b + row) * leading_dim + term.col0_b + col);
    value.x += term.scale_b * other.x;
    value.y += term.scale_b * other.y;
    value.z += term.scale_b * other.z;
    value.w += term.scale_b * other.w;
  }
  return value;
}

__launch_bounds__(kWide256VecRegThreadsPerBlock, 1) __global__
    void matmulStrassenFusedWide256ProductKernel(const float *__restrict__ A, const float *__restrict__ B,
                                                 float *__restrict__ product, int leading_a, int leading_b,
                                                 int product_n, int product_k, StrassenTerm lhs, StrassenTerm rhs) {
  static_assert(kWide256VecRegThreadRows * kWide256VecRegThreadsY == kWide256VecRegTileM,
                "fused Strassen thread mapping must cover product rows");
  static_assert(kWide256VecRegThreadCols * kWide256VecRegThreadsX == kWide256VecRegTileN,
                "fused Strassen thread mapping must cover product columns");

  __shared__ float tile_a[kWide256VecRegTileM][16 + kAsyncVecSharedPad];
  __shared__ float tile_b[16][kWide256VecRegTileN + kAsyncVecSharedPad];
  const int tx = threadIdx.x;
  const int ty = threadIdx.y;
  const int tid = ty * blockDim.x + tx;
  const int block_row = blockIdx.y * kWide256VecRegTileM;
  const int block_col = blockIdx.x * kWide256VecRegTileN;
  const int thread_row = ty * kWide256VecRegThreadRows;
  const int thread_col = tx * kWide256VecRegThreadCols;

  float accum[kWide256VecRegThreadRows][kWide256VecRegThreadCols] = {};

  for (int k0 = 0; k0 < product_k; k0 += 16) {
    for (int vec = tid; vec < (kWide256VecRegTileM * 16) / 4; vec += kWide256VecRegThreadsPerBlock) {
      const int scalar = vec * 4;
      const int row = scalar / 16;
      const int col = scalar - row * 16;
      const auto values = loadStrassenTerm4(A, leading_a, lhs, block_row + row, k0 + col);
      tile_a[row][col + 0] = values.x;
      tile_a[row][col + 1] = values.y;
      tile_a[row][col + 2] = values.z;
      tile_a[row][col + 3] = values.w;
    }
    for (int vec = tid; vec < (16 * kWide256VecRegTileN) / 4; vec += kWide256VecRegThreadsPerBlock) {
      const int scalar = vec * 4;
      const int row = scalar / kWide256VecRegTileN;
      const int col = scalar - row * kWide256VecRegTileN;
      const auto values = loadStrassenTerm4(B, leading_b, rhs, k0 + row, block_col + col);
      tile_b[row][col + 0] = values.x;
      tile_b[row][col + 1] = values.y;
      tile_b[row][col + 2] = values.z;
      tile_b[row][col + 3] = values.w;
    }
    __syncthreads();

    float a_values[kWide256VecRegThreadRows];
#pragma unroll
    for (int i = 0; i < kWide256VecRegThreadRows; ++i) {
      a_values[i] = tile_a[thread_row + i][0];
    }
#pragma unroll
    for (int kk = 0; kk < 16; ++kk) {
      float next_a_values[kWide256VecRegThreadRows];
      if (kk + 1 < 16) {
#pragma unroll
        for (int i = 0; i < kWide256VecRegThreadRows; ++i) {
          next_a_values[i] = tile_a[thread_row + i][kk + 1];
        }
      }
      const auto b0 = *reinterpret_cast<const float4 *>(&tile_b[kk][thread_col]);
      const auto b1 = *reinterpret_cast<const float4 *>(&tile_b[kk][thread_col + 4]);
#pragma unroll
      for (int i = 0; i < kWide256VecRegThreadRows; ++i) {
        auto &acc = accum[i];
        const float a = a_values[i];
        acc[0] = fmaf(a, b0.x, acc[0]);
        acc[1] = fmaf(a, b0.y, acc[1]);
        acc[2] = fmaf(a, b0.z, acc[2]);
        acc[3] = fmaf(a, b0.w, acc[3]);
        acc[4] = fmaf(a, b1.x, acc[4]);
        acc[5] = fmaf(a, b1.y, acc[5]);
        acc[6] = fmaf(a, b1.z, acc[6]);
        acc[7] = fmaf(a, b1.w, acc[7]);
      }
      if (kk + 1 < 16) {
#pragma unroll
        for (int i = 0; i < kWide256VecRegThreadRows; ++i) {
          a_values[i] = next_a_values[i];
        }
      }
    }
    __syncthreads();
  }

#pragma unroll
  for (int i = 0; i < kWide256VecRegThreadRows; ++i) {
    const int row = block_row + thread_row + i;
    auto *out = reinterpret_cast<float4 *>(product + row * product_n + block_col + thread_col);
    out[0] = make_float4(accum[i][0], accum[i][1], accum[i][2], accum[i][3]);
    out[1] = make_float4(accum[i][4], accum[i][5], accum[i][6], accum[i][7]);
  }
}

StrassenWorkspace &workspace() {
  static StrassenWorkspace value;
  return value;
}

std::array<StrassenProductSpec, kStrassenProductCount> strassenProductSpecs(int m2, int n2, int k2) {
  return {{
      {{0, 0, 1.0f, m2, k2, 1.0f}, {0, 0, 1.0f, k2, n2, 1.0f}},
      {{m2, 0, 1.0f, m2, k2, 1.0f}, {0, 0, 1.0f, 0, 0, 0.0f}},
      {{0, 0, 1.0f, 0, 0, 0.0f}, {0, n2, 1.0f, k2, n2, -1.0f}},
      {{m2, k2, 1.0f, 0, 0, 0.0f}, {k2, 0, 1.0f, 0, 0, -1.0f}},
      {{0, 0, 1.0f, 0, k2, 1.0f}, {k2, n2, 1.0f, 0, 0, 0.0f}},
      {{m2, 0, 1.0f, 0, 0, -1.0f}, {0, 0, 1.0f, 0, n2, 1.0f}},
      {{0, k2, 1.0f, m2, k2, -1.0f}, {k2, 0, 1.0f, k2, n2, 1.0f}},
  }};
}

std::string_view strassenInnerVariant(std::string_view variant) {
  if (variant == kStrassenOneLevelSplitK2Variant) {
    return kStrassenOneLevelSplitK2InnerVariant;
  }
  if (variant == kStrassenOneLevelSplitK4Variant) {
    return kStrassenOneLevelSplitK4InnerVariant;
  }
  return kStrassenOneLevelInnerVariant;
}

void ensureWorkspace(int m2, int n2, int k2, bool allocate_operands) {
  auto &ws = workspace();
  const auto lhs_elements = static_cast<size_t>(m2) * static_cast<size_t>(k2);
  const auto rhs_elements = static_cast<size_t>(k2) * static_cast<size_t>(n2);
  const auto product_elements = static_cast<size_t>(m2) * static_cast<size_t>(n2);

  if (allocate_operands && ws.lhs_elements < lhs_elements) {
    for (int product = 0; product < kStrassenProductCount; ++product) {
      cudaFree(ws.lhs[product]);
      cudaMalloc(&ws.lhs[product], lhs_elements * sizeof(float));
    }
    ws.lhs_elements = lhs_elements;
  }
  if (allocate_operands && ws.rhs_elements < rhs_elements) {
    for (int product = 0; product < kStrassenProductCount; ++product) {
      cudaFree(ws.rhs[product]);
      cudaMalloc(&ws.rhs[product], rhs_elements * sizeof(float));
    }
    ws.rhs_elements = rhs_elements;
  }
  if (ws.product_elements < kStrassenProductCount * product_elements) {
    cudaFree(ws.products);
    cudaMalloc(&ws.products, kStrassenProductCount * product_elements * sizeof(float));
    ws.product_elements = kStrassenProductCount * product_elements;
  }
}

void ensureProductStreams() {
  auto &ws = workspace();
  if (ws.streams_initialized) {
    return;
  }
  for (int product = 0; product < kStrassenProductCount; ++product) {
    cudaStreamCreate(&ws.streams[product]);
    cudaEventCreateWithFlags(&ws.done_events[product], cudaEventDisableTiming);
  }
  ws.streams_initialized = true;
}

void buildOperand(const float *base, float *out, int rows, int cols, int leading_dim, int row0_a, int col0_a,
                  float scale_a, int row0_b, int col0_b, float scale_b, cudaStream_t stream) {
  const int threads = 256;
  const int elements = rows * cols;
  const auto *src_a = base + row0_a * leading_dim + col0_a;
  const auto *src_b = base + row0_b * leading_dim + col0_b;
  if (cols % 4 == 0 && leading_dim % 4 == 0 && isAlignedForFloat4(src_a) && isAlignedForFloat4(out) &&
      (scale_b == 0.0f || isAlignedForFloat4(src_b))) {
    const int vec_elements = elements / 4;
    const int blocks = (vec_elements + threads - 1) / threads;
    buildStrassenOperandFloat4Kernel<<<blocks, threads, 0, stream>>>(base, out, rows, cols, leading_dim, row0_a, col0_a,
                                                                     scale_a, row0_b, col0_b, scale_b);
    return;
  }
  const int blocks = (elements + threads - 1) / threads;
  buildStrassenOperandKernel<<<blocks, threads, 0, stream>>>(base, out, rows, cols, leading_dim, row0_a, col0_a,
                                                             scale_a, row0_b, col0_b, scale_b);
}

bool canBuildBatchedStrassenOperands(const MatmulOptions &options, const float *d_a, const float *d_b) {
  const int m2 = options.m / 2;
  const int n2 = options.n / 2;
  const int k2 = options.k / 2;
  return k2 % 4 == 0 && n2 % 4 == 0 && options.k % 4 == 0 && options.n % 4 == 0 && isAlignedForFloat4(d_a) &&
         isAlignedForFloat4(d_a + k2) && isAlignedForFloat4(d_a + m2 * options.k) &&
         isAlignedForFloat4(d_a + m2 * options.k + k2) && isAlignedForFloat4(d_b) && isAlignedForFloat4(d_b + n2) &&
         isAlignedForFloat4(d_b + k2 * options.n) && isAlignedForFloat4(d_b + k2 * options.n + n2);
}

void buildBatchedStrassenOperands(const MatmulOptions &options, const float *d_a, const float *d_b) {
  auto &ws = workspace();
  const int m2 = options.m / 2;
  const int n2 = options.n / 2;
  const int k2 = options.k / 2;
  const int threads = 256;
  const int lhs_vec_elements = (m2 * k2) / 4;
  const int rhs_vec_elements = (k2 * n2) / 4;
  buildStrassenLhsBatchedFloat4Kernel<<<(lhs_vec_elements + threads - 1) / threads, threads>>>(
      d_a, ws.lhs[0], ws.lhs[1], ws.lhs[4], ws.lhs[5], ws.lhs[6], m2, k2, options.k);
  buildStrassenRhsBatchedFloat4Kernel<<<(rhs_vec_elements + threads - 1) / threads, threads>>>(
      d_b, ws.rhs[0], ws.rhs[2], ws.rhs[3], ws.rhs[5], ws.rhs[6], k2, n2, options.n);
}

void launchInnerProduct(int m2, int n2, int k2, const float *lhs, const float *rhs, float *product,
                        std::string_view inner_variant, cudaStream_t stream) {
  MatmulOptions inner;
  inner.m = m2;
  inner.n = n2;
  inner.k = k2;
  inner.variant = std::string(inner_variant);
  inner.tile = kWide256VecRegTileM;
  inner.block_x = kWide256VecRegThreadsX;
  inner.block_y = kWide256VecRegThreadsY;
  if (inner_variant == kStrassenOneLevelInnerVariant) {
    launchWide256APrefetchMatmulOnStream(inner, lhs, rhs, product, stream);
  } else {
    launchVectorizedRegisterMatmul(inner, lhs, rhs, nullptr, product);
  }
}

void launchProduct(const float *a, const float *b, int n, int k, int m2, int n2, int k2, int product_index,
                   const StrassenProductSpec &spec, std::string_view inner_variant, cudaStream_t stream) {
  auto &ws = workspace();
  const auto product_elements = static_cast<size_t>(m2) * static_cast<size_t>(n2);
  buildOperand(a, ws.lhs[product_index], m2, k2, k, spec.lhs.row0_a, spec.lhs.col0_a, spec.lhs.scale_a, spec.lhs.row0_b,
               spec.lhs.col0_b, spec.lhs.scale_b, stream);
  buildOperand(b, ws.rhs[product_index], k2, n2, n, spec.rhs.row0_a, spec.rhs.col0_a, spec.rhs.scale_a, spec.rhs.row0_b,
               spec.rhs.col0_b, spec.rhs.scale_b, stream);
  launchInnerProduct(m2, n2, k2, ws.lhs[product_index], ws.rhs[product_index],
                     ws.products + product_index * product_elements, inner_variant, stream);
}

void launchFallback(const MatmulOptions &options, const float *d_a, const float *d_b, float *d_c) {
  MatmulOptions fallback = options;
  fallback.variant = std::string(kStrassenOneLevelInnerVariant);
  launchVectorizedRegisterMatmul(fallback, d_a, d_b, nullptr, d_c);
}

void launchStrassenProducts(const MatmulOptions &options, const float *d_a, const float *d_b, bool streamed) {
  auto &ws = workspace();
  const int m2 = options.m / 2;
  const int n2 = options.n / 2;
  const int k2 = options.k / 2;
  const auto specs = strassenProductSpecs(m2, n2, k2);
  const std::string_view inner_variant = strassenInnerVariant(options.variant);
  if (streamed) {
    ensureProductStreams();
  }
  for (int product = 0; product < kStrassenProductCount; ++product) {
    const cudaStream_t stream = streamed ? ws.streams[product] : nullptr;
    launchProduct(d_a, d_b, options.n, options.k, m2, n2, k2, product, specs[product], inner_variant, stream);
    if (streamed) {
      cudaEventRecord(ws.done_events[product], stream);
    }
  }
  if (streamed) {
    for (int product = 0; product < kStrassenProductCount; ++product) {
      cudaStreamWaitEvent(nullptr, ws.done_events[product], 0);
    }
  }
}

void launchFusedStrassenProduct(const MatmulOptions &options, const float *d_a, const float *d_b, int m2, int n2,
                                int k2, int product_index, const StrassenProductSpec &spec, cudaStream_t stream) {
  auto &ws = workspace();
  const auto product_elements = static_cast<size_t>(m2) * static_cast<size_t>(n2);
  const dim3 block(kWide256VecRegThreadsX, kWide256VecRegThreadsY);
  const dim3 grid(n2 / kWide256VecRegTileN, m2 / kWide256VecRegTileM);
  matmulStrassenFusedWide256ProductKernel<<<grid, block, 0, stream>>>(
      d_a, d_b, ws.products + product_index * product_elements, options.k, options.n, n2, k2, spec.lhs, spec.rhs);
}

void launchFusedStrassenProducts(const MatmulOptions &options, const float *d_a, const float *d_b, bool streamed) {
  auto &ws = workspace();
  const int m2 = options.m / 2;
  const int n2 = options.n / 2;
  const int k2 = options.k / 2;
  const auto specs = strassenProductSpecs(m2, n2, k2);
  if (streamed) {
    ensureProductStreams();
  }
  for (int product = 0; product < kStrassenProductCount; ++product) {
    const cudaStream_t stream = streamed ? ws.streams[product] : nullptr;
    launchFusedStrassenProduct(options, d_a, d_b, m2, n2, k2, product, specs[product], stream);
    if (streamed) {
      cudaEventRecord(ws.done_events[product], stream);
    }
  }
  if (streamed) {
    for (int product = 0; product < kStrassenProductCount; ++product) {
      cudaStreamWaitEvent(nullptr, ws.done_events[product], 0);
    }
  }
}

bool isDirectOperand(const StrassenTerm &term) { return term.scale_a == 1.0f && term.scale_b == 0.0f; }

StrassenOperandView prepareOperandView(const float *base, float *workspace_ptr, int rows, int cols,
                                       int base_leading_dim, const StrassenTerm &term, cudaStream_t stream) {
  if (isDirectOperand(term)) {
    return {base + term.row0_a * base_leading_dim + term.col0_a, base_leading_dim};
  }
  buildOperand(base, workspace_ptr, rows, cols, base_leading_dim, term.row0_a, term.col0_a, term.scale_a, term.row0_b,
               term.col0_b, term.scale_b, stream);
  return {workspace_ptr, cols};
}

inline StrassenOperandView batchedOperandView(const float *base, float *workspace_ptr, int base_leading_dim,
                                              int workspace_leading_dim, const StrassenTerm &term) {
  if (isDirectOperand(term)) {
    return {base + term.row0_a * base_leading_dim + term.col0_a, base_leading_dim};
  }
  return {workspace_ptr, workspace_leading_dim};
}

void launchStridedStrassenProducts(const MatmulOptions &options, const float *d_a, const float *d_b, bool streamed,
                                   bool batched_operands) {
  auto &ws = workspace();
  const int m2 = options.m / 2;
  const int n2 = options.n / 2;
  const int k2 = options.k / 2;
  const auto specs = strassenProductSpecs(m2, n2, k2);
  const auto product_elements = static_cast<size_t>(m2) * static_cast<size_t>(n2);
  if (streamed) {
    ensureProductStreams();
  }
  const bool use_batched_operands = batched_operands && !streamed && canBuildBatchedStrassenOperands(options, d_a, d_b);
  if (use_batched_operands) {
    buildBatchedStrassenOperands(options, d_a, d_b);
  }
  for (int product = 0; product < kStrassenProductCount; ++product) {
    const cudaStream_t stream = streamed ? ws.streams[product] : nullptr;
    const auto lhs = use_batched_operands
                         ? batchedOperandView(d_a, ws.lhs[product], options.k, k2, specs[product].lhs)
                         : prepareOperandView(d_a, ws.lhs[product], m2, k2, options.k, specs[product].lhs, stream);
    const auto rhs = use_batched_operands
                         ? batchedOperandView(d_b, ws.rhs[product], options.n, n2, specs[product].rhs)
                         : prepareOperandView(d_b, ws.rhs[product], k2, n2, options.n, specs[product].rhs, stream);
    launchWide256APrefetchStridedMatmulOnStream(m2, n2, k2, lhs.ptr, lhs.leading_dim, rhs.ptr, rhs.leading_dim,
                                                ws.products + product * product_elements, n2, stream);
    if (streamed) {
      cudaEventRecord(ws.done_events[product], stream);
    }
  }
  if (streamed) {
    for (int product = 0; product < kStrassenProductCount; ++product) {
      cudaStreamWaitEvent(nullptr, ws.done_events[product], 0);
    }
  }
}

} // namespace

void launchStrassenOneLevelMatmul(const MatmulOptions &options, const float *d_a, const float *d_b, float *d_c) {
  if (!hasStrassenOneLevelInteriorTiles(options) || !isAlignedForFloat4(d_a) || !isAlignedForFloat4(d_b) ||
      !isAlignedForFloat4(d_c)) {
    launchFallback(options, d_a, d_b, d_c);
    return;
  }

  const int m2 = options.m / 2;
  const int n2 = options.n / 2;
  const int k2 = options.k / 2;
  const bool fused = isStrassenOneLevelFusedVariant(options.variant);
  ensureWorkspace(m2, n2, k2, !fused);
  if (isStrassenOneLevelStridedVariant(options.variant)) {
    launchStridedStrassenProducts(options, d_a, d_b, isStrassenOneLevelStridedStreamedVariant(options.variant),
                                  isStrassenOneLevelStridedBatchedVariant(options.variant));
  } else if (fused) {
    launchFusedStrassenProducts(options, d_a, d_b, isStrassenOneLevelFusedStreamedVariant(options.variant));
  } else {
    launchStrassenProducts(options, d_a, d_b, isStrassenOneLevelStreamedVariant(options.variant));
  }

  auto &ws = workspace();
  const auto product_elements = static_cast<size_t>(m2) * static_cast<size_t>(n2);
  const int threads = 256;
  if (n2 % 4 == 0 && options.n % 4 == 0 && isAlignedForFloat4(ws.products) && isAlignedForFloat4(d_c)) {
    const int vec_elements = static_cast<int>(product_elements / 4);
    const int blocks = (vec_elements + threads - 1) / threads;
    combineStrassenProductsFloat4Kernel<<<blocks, threads>>>(
        ws.products, ws.products + product_elements, ws.products + 2 * product_elements,
        ws.products + 3 * product_elements, ws.products + 4 * product_elements, ws.products + 5 * product_elements,
        ws.products + 6 * product_elements, d_c, m2, n2, options.n);
  } else {
    const int blocks = static_cast<int>((product_elements + threads - 1) / threads);
    combineStrassenProductsKernel<<<blocks, threads>>>(
        ws.products, ws.products + product_elements, ws.products + 2 * product_elements,
        ws.products + 3 * product_elements, ws.products + 4 * product_elements, ws.products + 5 * product_elements,
        ws.products + 6 * product_elements, d_c, m2, n2, options.n);
  }
}

KernelResources selectedStrassenOneLevelKernelResources(const MatmulOptions &options) {
  if (isStrassenOneLevelFusedVariant(options.variant) && hasStrassenOneLevelInteriorTiles(options)) {
    return queryKernelResources(matmulStrassenFusedWide256ProductKernel, kWide256VecRegThreadsPerBlock, 0);
  }
  if (isStrassenOneLevelStridedVariant(options.variant)) {
    return selectedWide256APrefetchStridedKernelResources(options);
  }
  MatmulOptions inner = options;
  inner.variant = std::string(strassenInnerVariant(options.variant));
  if (hasStrassenOneLevelInteriorTiles(options)) {
    inner.m = options.m / 2;
    inner.n = options.n / 2;
    inner.k = options.k / 2;
  }
  return selectedVectorizedRegisterKernelResources(inner);
}