270 lines
9.8 KiB
C
270 lines
9.8 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file test_metal_gemm_batched.c
|
|
* @brief Tests fm_sgemm_batched against a plain CPU reference GEMM: several
|
|
* batch counts, non-square shapes, padded (non-packed) batch strides,
|
|
* alpha/beta, and rejection of invalid strides.
|
|
*
|
|
* Exits 77 when no Metal device is present so GPU-less CI skips cleanly.
|
|
*
|
|
* @status All tests pass on M4 / macOS 26; each was proven able to fail.
|
|
* @issues None
|
|
* @todo None
|
|
*/
|
|
|
|
#include "lib/test.h"
|
|
#include "metal/fm_metal.h"
|
|
|
|
#include <math.h>
|
|
#include <stdio.h>
|
|
|
|
/* Deterministic, well-conditioned filler — distinct per (matrix, element)
|
|
* so a batch element read at the wrong stride cannot coincidentally match. */
|
|
static float sample_value(int seed, int i)
|
|
{
|
|
return (float)(((seed * 31 + i * 17) % 23) - 11) * 0.25f;
|
|
}
|
|
|
|
/* Plain row-major reference: C[M,N] = alpha * A[M,K] @ B[K,N] + beta * C. */
|
|
static void cpu_gemm(const float *a, const float *b, float *c,
|
|
int M, int N, int K, float alpha, float beta)
|
|
{
|
|
for (int m = 0; m < M; m++) {
|
|
for (int n = 0; n < N; n++) {
|
|
float sum = 0.0f;
|
|
for (int k = 0; k < K; k++)
|
|
sum += a[m * K + k] * b[k * N + n];
|
|
c[m * N + n] = alpha * sum + beta * c[m * N + n];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* One batched-GEMM case: fills A and B, runs the GPU batch and the CPU
|
|
* reference over every batch element, and compares. `pad_matrices`
|
|
* inserts one matrix worth of untouched padding between batch elements,
|
|
* exercising a non-packed stride. */
|
|
static void test_batched_case_against_cpu(int batch, int M, int N, int K,
|
|
float alpha, float beta, bool pad_matrices)
|
|
{
|
|
size_t a_elems = (size_t)M * K, b_elems = (size_t)K * N, c_elems = (size_t)M * N;
|
|
int span = pad_matrices ? 2 : 1;
|
|
size_t a_stride = a_elems * span * sizeof(float);
|
|
size_t b_stride = b_elems * span * sizeof(float);
|
|
size_t c_stride = c_elems * span * sizeof(float);
|
|
|
|
fm_buffer *ba = fm_buffer_alloc(a_stride * (size_t)batch);
|
|
fm_buffer *bb = fm_buffer_alloc(b_stride * (size_t)batch);
|
|
fm_buffer *bc = fm_buffer_alloc(c_stride * (size_t)batch);
|
|
ASSERT_NOT_NULL(ba);
|
|
ASSERT_NOT_NULL(bb);
|
|
ASSERT_NOT_NULL(bc);
|
|
|
|
float *a = (float *)fm_buffer_data(ba);
|
|
float *b = (float *)fm_buffer_data(bb);
|
|
float *c = (float *)fm_buffer_data(bc);
|
|
|
|
for (int e = 0; e < batch; e++) {
|
|
float *ae = a + (size_t)e * a_elems * span;
|
|
float *be = b + (size_t)e * b_elems * span;
|
|
float *ce = c + (size_t)e * c_elems * span;
|
|
for (size_t i = 0; i < a_elems; i++) ae[i] = sample_value(e + 1, (int)i);
|
|
for (size_t i = 0; i < b_elems; i++) be[i] = sample_value(e + 101, (int)i);
|
|
for (size_t i = 0; i < c_elems; i++) ce[i] = sample_value(e + 201, (int)i);
|
|
}
|
|
|
|
/* CPU reference over a private copy of C, before the GPU touches it. */
|
|
float *expected = (float *)malloc(c_elems * (size_t)batch * sizeof(float));
|
|
ASSERT_NOT_NULL(expected);
|
|
for (int e = 0; e < batch; e++) {
|
|
float *ce = expected + (size_t)e * c_elems;
|
|
memcpy(ce, c + (size_t)e * c_elems * span, c_elems * sizeof(float));
|
|
cpu_gemm(a + (size_t)e * a_elems * span,
|
|
b + (size_t)e * b_elems * span,
|
|
ce, M, N, K, alpha, beta);
|
|
}
|
|
|
|
fm_sgemm_batched(ba, bb, bc, batch, M, N, K,
|
|
pad_matrices ? a_stride : 0,
|
|
pad_matrices ? b_stride : 0,
|
|
pad_matrices ? c_stride : 0,
|
|
alpha, beta);
|
|
fm_sync();
|
|
|
|
float *result = (float *)fm_buffer_data(bc);
|
|
for (int e = 0; e < batch; e++) {
|
|
const float *got = result + (size_t)e * c_elems * span;
|
|
const float *want = expected + (size_t)e * c_elems;
|
|
for (size_t i = 0; i < c_elems; i++) {
|
|
float tol = fabsf(want[i]) * 1e-4f + 1e-4f;
|
|
if (fabsf(got[i] - want[i]) > tol)
|
|
printf("\n batch %d elem %zu: got %.6f want %.6f\n",
|
|
e, i, got[i], want[i]);
|
|
ASSERT_TRUE(fabsf(got[i] - want[i]) <= tol);
|
|
}
|
|
}
|
|
|
|
free(expected);
|
|
fm_buffer_free(ba);
|
|
fm_buffer_free(bb);
|
|
fm_buffer_free(bc);
|
|
}
|
|
|
|
/* ================================================================
|
|
* Correctness vs the CPU reference
|
|
* ================================================================ */
|
|
|
|
static void test_batch_of_one_matches_reference(void)
|
|
{
|
|
test_batched_case_against_cpu(1, 3, 5, 4, 1.0f, 0.0f, false);
|
|
}
|
|
|
|
static void test_batch_of_four_non_square(void)
|
|
{
|
|
/* M != N != K, so a transposed or shape-confused encode cannot pass. */
|
|
test_batched_case_against_cpu(4, 3, 5, 4, 1.0f, 0.0f, false);
|
|
}
|
|
|
|
static void test_batch_of_seventeen_non_square(void)
|
|
{
|
|
test_batched_case_against_cpu(17, 6, 2, 9, 1.0f, 0.0f, false);
|
|
}
|
|
|
|
static void test_batch_honours_alpha_and_beta(void)
|
|
{
|
|
test_batched_case_against_cpu(3, 4, 3, 5, 2.0f, -0.5f, false);
|
|
}
|
|
|
|
static void test_batch_with_padded_strides(void)
|
|
{
|
|
/* Twice-packed strides: each batch element is followed by a matrix
|
|
* worth of padding that must be skipped, not consumed. */
|
|
test_batched_case_against_cpu(5, 3, 5, 4, 1.0f, 0.0f, true);
|
|
}
|
|
|
|
static void test_larger_batch_matches_reference(void)
|
|
{
|
|
test_batched_case_against_cpu(8, 32, 24, 40, 1.0f, 0.0f, false);
|
|
}
|
|
|
|
/* ================================================================
|
|
* Agreement with the single-matrix path
|
|
* ================================================================ */
|
|
|
|
static void test_batch_element_matches_fm_sgemm(void)
|
|
{
|
|
int batch = 3, M = 4, N = 6, K = 5;
|
|
size_t a_elems = (size_t)M * K, b_elems = (size_t)K * N, c_elems = (size_t)M * N;
|
|
|
|
fm_buffer *ba = fm_buffer_alloc(a_elems * (size_t)batch * sizeof(float));
|
|
fm_buffer *bb = fm_buffer_alloc(b_elems * (size_t)batch * sizeof(float));
|
|
fm_buffer *bc = fm_buffer_alloc(c_elems * (size_t)batch * sizeof(float));
|
|
ASSERT_NOT_NULL(ba);
|
|
ASSERT_NOT_NULL(bb);
|
|
ASSERT_NOT_NULL(bc);
|
|
|
|
float *a = (float *)fm_buffer_data(ba);
|
|
float *b = (float *)fm_buffer_data(bb);
|
|
for (int e = 0; e < batch; e++) {
|
|
for (size_t i = 0; i < a_elems; i++) a[(size_t)e * a_elems + i] = sample_value(e + 3, (int)i);
|
|
for (size_t i = 0; i < b_elems; i++) b[(size_t)e * b_elems + i] = sample_value(e + 53, (int)i);
|
|
}
|
|
|
|
fm_sgemm_batched(ba, bb, bc, batch, M, N, K, 0, 0, 0, 1.0f, 0.0f);
|
|
fm_sync();
|
|
|
|
/* Re-run the last batch element on its own through the single-matrix
|
|
* path and require the same numbers. */
|
|
fm_buffer *single_a = fm_buffer_alloc(a_elems * sizeof(float));
|
|
fm_buffer *single_b = fm_buffer_alloc(b_elems * sizeof(float));
|
|
fm_buffer *single_c = fm_buffer_alloc(c_elems * sizeof(float));
|
|
ASSERT_NOT_NULL(single_a);
|
|
ASSERT_NOT_NULL(single_b);
|
|
ASSERT_NOT_NULL(single_c);
|
|
|
|
int last = batch - 1;
|
|
memcpy(fm_buffer_data(single_a), a + (size_t)last * a_elems, a_elems * sizeof(float));
|
|
memcpy(fm_buffer_data(single_b), b + (size_t)last * b_elems, b_elems * sizeof(float));
|
|
fm_sgemm(single_a, single_b, single_c, M, N, K, 1.0f, 0.0f);
|
|
fm_sync();
|
|
|
|
const float *batched = (const float *)fm_buffer_data(bc) + (size_t)last * c_elems;
|
|
const float *single = (const float *)fm_buffer_data(single_c);
|
|
for (size_t i = 0; i < c_elems; i++)
|
|
ASSERT_FLT_NEAR(single[i], batched[i], 1e-4f);
|
|
|
|
fm_buffer_free(ba);
|
|
fm_buffer_free(bb);
|
|
fm_buffer_free(bc);
|
|
fm_buffer_free(single_a);
|
|
fm_buffer_free(single_b);
|
|
fm_buffer_free(single_c);
|
|
}
|
|
|
|
/* ================================================================
|
|
* Argument validation
|
|
* ================================================================ */
|
|
|
|
static void test_invalid_arguments_leave_output_untouched(void)
|
|
{
|
|
int batch = 2, M = 2, N = 2, K = 2;
|
|
size_t c_elems = (size_t)M * N;
|
|
|
|
fm_buffer *ba = fm_buffer_alloc((size_t)M * K * batch * sizeof(float));
|
|
fm_buffer *bb = fm_buffer_alloc((size_t)K * N * batch * sizeof(float));
|
|
fm_buffer *bc = fm_buffer_alloc(c_elems * batch * sizeof(float));
|
|
ASSERT_NOT_NULL(ba);
|
|
ASSERT_NOT_NULL(bb);
|
|
ASSERT_NOT_NULL(bc);
|
|
|
|
float *c = (float *)fm_buffer_data(bc);
|
|
for (size_t i = 0; i < c_elems * (size_t)batch; i++) c[i] = -7.0f;
|
|
|
|
/* A stride below the packed size, a stride that is not a multiple of
|
|
* the row bytes, a zero batch, and a NULL operand must all be
|
|
* rejected before any dispatch. */
|
|
fm_sgemm_batched(ba, bb, bc, batch, M, N, K, sizeof(float), 0, 0, 1.0f, 0.0f);
|
|
fm_sgemm_batched(ba, bb, bc, batch, M, N, K, 0, 0,
|
|
c_elems * sizeof(float) + 1, 1.0f, 0.0f);
|
|
fm_sgemm_batched(ba, bb, bc, 0, M, N, K, 0, 0, 0, 1.0f, 0.0f);
|
|
fm_sgemm_batched(NULL, bb, bc, batch, M, N, K, 0, 0, 0, 1.0f, 0.0f);
|
|
fm_sync();
|
|
|
|
const float *result = (const float *)fm_buffer_data(bc);
|
|
for (size_t i = 0; i < c_elems * (size_t)batch; i++)
|
|
ASSERT_FLT_NEAR(-7.0f, result[i], 1e-6f);
|
|
|
|
fm_buffer_free(ba);
|
|
fm_buffer_free(bb);
|
|
fm_buffer_free(bc);
|
|
}
|
|
|
|
/* ================================================================
|
|
* main
|
|
* ================================================================ */
|
|
|
|
int main(void)
|
|
{
|
|
printf("\n Metal Batched GEMM Tests\n");
|
|
printf(" ========================\n\n");
|
|
|
|
if (fm_init() != 0 || !fm_available()) {
|
|
fprintf(stderr, "Metal unavailable\n");
|
|
return 77;
|
|
}
|
|
|
|
RUN_TEST(batch_of_one_matches_reference);
|
|
RUN_TEST(batch_of_four_non_square);
|
|
RUN_TEST(batch_of_seventeen_non_square);
|
|
RUN_TEST(batch_honours_alpha_and_beta);
|
|
RUN_TEST(batch_with_padded_strides);
|
|
RUN_TEST(larger_batch_matches_reference);
|
|
RUN_TEST(batch_element_matches_fm_sgemm);
|
|
RUN_TEST(invalid_arguments_leave_output_untouched);
|
|
|
|
fm_free();
|
|
|
|
TEST_REPORT();
|
|
return __ft_failed > 0 ? 1 : 0;
|
|
}
|