From 77def3f21e7841bc04a28bb8af32ad6512750a3b Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 13:57:31 +0200 Subject: [PATCH 01/49] test: cover packed matmul and Linear cache invalidation --- tests/unit/nn/test_nn_linear.c | 43 ++++++++++++++++++++++++++ tests/unit/tensor/test_tensor_matmul.c | 40 ++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/tests/unit/nn/test_nn_linear.c b/tests/unit/nn/test_nn_linear.c index f9db939..2afcdc7 100644 --- a/tests/unit/nn/test_nn_linear.c +++ b/tests/unit/nn/test_nn_linear.c @@ -156,6 +156,48 @@ static void test_vector_batch_and_no_bias(void) nn_linear_destroy(layer); } +static void test_weight_update_refreshes_packed_state(void) +{ + int input_dims[] = {2, 3}; + float values[] = {1, 2, 3, 4, 5, 6}; + nn_linear* layer = nn_linear_create( + "cached", 3, 2, 0, NN_INIT_ZERO, NN_INIT_ZERO, NULL); + ag_tensor* input = make_ag(2, input_dims, values, 0); + ag_tensor* first; + ag_tensor* second; + + CHECK(layer != NULL && input != NULL); + if (layer == NULL || input == NULL) { + ag_tensor_release(input); + nn_linear_destroy(layer); + return; + } + for (int i = 0; i < 6; ++i) { + layer->weight->value->value->storage->data[i] = 1.0f; + } + tensor_mark_modified(layer->weight->value->value); + first = nn_linear_forward(layer, input); + CHECK(first != NULL); + if (first != NULL) { + CHECK(first->value->storage->data[0] == 6.0f); + CHECK(first->value->storage->data[1] == 6.0f); + } + for (int i = 0; i < 6; ++i) { + layer->weight->value->value->storage->data[i] = 2.0f; + } + tensor_mark_modified(layer->weight->value->value); + second = nn_linear_forward(layer, input); + CHECK(second != NULL); + if (second != NULL) { + CHECK(second->value->storage->data[0] == 12.0f); + CHECK(second->value->storage->data[1] == 12.0f); + } + ag_tensor_release(second); + ag_tensor_release(first); + ag_tensor_release(input); + nn_linear_destroy(layer); +} + static void test_invalid(void) { int bad_dims[] = {2, 4}; @@ -186,6 +228,7 @@ int main(void) test_constructor(); test_forward_and_backward(); test_vector_batch_and_no_bias(); + test_weight_update_refreshes_packed_state(); test_invalid(); if (failures != 0) { fprintf(stderr, "%d Linear contract checks failed\n", failures); diff --git a/tests/unit/tensor/test_tensor_matmul.c b/tests/unit/tensor/test_tensor_matmul.c index e23b081..89c8491 100644 --- a/tests/unit/tensor/test_tensor_matmul.c +++ b/tests/unit/tensor/test_tensor_matmul.c @@ -135,6 +135,45 @@ TEST(test_t_matmul_accepts_transposed_views) { t_free(a_base); } +TEST(test_t_matmul_transposed_2d_uneven_rows_and_columns) { + int a_dims[] = {5, 3}; + int b_dims[] = {4, 3}; + tensor* a = t_alloc(2, a_dims); + tensor* b_base = t_alloc(2, b_dims); + tensor* b = NULL; + tensor* actual = NULL; + + ASSERT_NOT_NULL(a); + ASSERT_NOT_NULL(b_base); + for (int i = 0; i < tensor_numel(a); ++i) { + a->storage->data[i] = (float)(i + 1); + } + for (int i = 0; i < tensor_numel(b_base); ++i) { + b_base->storage->data[i] = (float)(i - 3); + } + b = t_transpose(b_base, 0, 1); + actual = t_matmul(a, b); + ASSERT_NOT_NULL(b); + ASSERT_NOT_NULL(actual); + ASSERT_EQ_INT(actual->ndim, 2); + ASSERT_EQ_INT(actual->dims[0], 5); + ASSERT_EQ_INT(actual->dims[1], 4); + for (int row = 0; row < 5; ++row) { + for (int column = 0; column < 4; ++column) { + float expected = 0.0f; + for (int inner = 0; inner < 3; ++inner) { + expected += a->storage->data[row * 3 + inner] * + b_base->storage->data[column * 3 + inner]; + } + ASSERT_EQ_FLOAT(actual->storage->data[row * 4 + column], expected); + } + } + t_free(actual); + t_free(b); + t_free(b_base); + t_free(a); +} + TEST(test_t_matmul_vector_cases) { int vector_dims[1] = {3}; int matrix_dims[2] = {2, 3}; @@ -479,6 +518,7 @@ int main(void) { RUN_TEST(test_t_matmul_batched_3d); RUN_TEST(test_t_matmul_broadcasts_batch_dimensions); RUN_TEST(test_t_matmul_accepts_transposed_views); + RUN_TEST(test_t_matmul_transposed_2d_uneven_rows_and_columns); RUN_TEST(test_t_matmul_vector_cases); RUN_TEST(test_packed_rhs_snapshots_transposed_view); RUN_TEST(test_packed_rhs_handles_transposed_slices_and_kernel_tails); From a9f0b441c8c91195ac7aea3066c68f37cf0c8761 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:00:08 +0200 Subject: [PATCH 02/49] perf: parallelize packed matmul row blocks --- src/tensor/tensor_matmul.c | 77 ++++++++++++++++++++++++-------------- 1 file changed, 48 insertions(+), 29 deletions(-) diff --git a/src/tensor/tensor_matmul.c b/src/tensor/tensor_matmul.c index 568fdbb..3b6759c 100644 --- a/src/tensor/tensor_matmul.c +++ b/src/tensor/tensor_matmul.c @@ -56,9 +56,10 @@ static int checked_size_multiply(size_t left, size_t right, size_t* result) { #if TENSORLIB_HAS_AVX2_KERNEL static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, - const float* packed_rhs, int rows, - int inner, int columns, tensor* output, - int output_base, int output_batch_ndim); + const float* packed_rhs, int inner, + int columns, tensor* output, + int output_base, int output_batch_ndim, + int row_start, int row_end); #endif static void* matmul_aligned_malloc(size_t bytes) { @@ -234,13 +235,14 @@ static void matmul_2d_packed_rhs_scalar(const tensor* lhs, int columns, tensor* output, int output_base, - int output_batch_ndim) { + int output_batch_ndim, + int row_start, int row_end) { int lhs_row_stride = lhs->strides[lhs->ndim - 2]; int lhs_inner_stride = lhs->strides[lhs->ndim - 1]; int output_row_stride = output->strides[output_batch_ndim]; int output_column_stride = output->strides[output_batch_ndim + 1]; - for (int row = 0; row < lhs_info->rows; ++row) { + for (int row = row_start; row < row_end; ++row) { for (int column = 0; column < columns; ++column) { const float* rhs_panel = packed_rhs + (size_t)(column / TENSORLIB_MATMUL_NR) * @@ -266,7 +268,9 @@ static void matmul_packed_rhs_batch(const tensor* lhs, tensor* output, const int* batch_dims, int batch_ndim, - size_t batch) { + size_t batch, + int row_start, + int row_end) { int coords[TENSORLIB_MATMUL_MAX_NDIM]; batch_index_to_coords(batch, batch_dims, batch_ndim, coords); @@ -281,19 +285,21 @@ static void matmul_packed_rhs_batch(const tensor* lhs, if (lhs_row_stride > 0 && lhs_inner_stride > 0 && matmul_avx2_available() && matmul_2d_packed_rhs_avx2(lhs, lhs_base, packed_rhs, - lhs_info->rows, lhs_info->inner, - rhs->columns, output, output_base, - batch_ndim)) { + lhs_info->inner, rhs->columns, + output, output_base, + batch_ndim, row_start, row_end)) { /* The AVX2 kernel completed this batch. */ } else { matmul_2d_packed_rhs_scalar(lhs, lhs_info, lhs_base, packed_rhs, rhs->columns, - output, output_base, batch_ndim); + output, output_base, batch_ndim, + row_start, row_end); } #else matmul_2d_packed_rhs_scalar(lhs, lhs_info, lhs_base, packed_rhs, rhs->columns, - output, output_base, batch_ndim); + output, output_base, batch_ndim, + row_start, row_end); #endif } @@ -354,25 +360,32 @@ tensor* t_matmul_packed_rhs(const tensor* lhs, return NULL; } - int batch_parallel = (batch_ndim <= TENSORLIB_MATMUL_MAX_NDIM); + int row_block_count = (lhs_info.rows + TENSORLIB_MATMUL_MC - 1) / + TENSORLIB_MATMUL_MC; + int task_count = (int)(batch_count * (size_t)row_block_count); long long batch_flops = 2LL * lhs_info.rows * lhs_info.inner * rhs->columns * (long long)batch_count; int threads = tensorlib_parallel_threads( batch_flops, TENSORLIB_MATMUL_MIN_PARALLEL_FLOPS, - batch_parallel ? (int)batch_count : 1); + batch_ndim <= TENSORLIB_MATMUL_MAX_NDIM ? task_count : 1); if (threads > 1) { #ifdef _OPENMP #pragma omp parallel for schedule(static) num_threads(threads) #endif - for (long long batch = 0; batch < (long long)batch_count; ++batch) { + for (long long task = 0; task < (long long)task_count; ++task) { + size_t batch = (size_t)task / (size_t)row_block_count; + int row_start = (int)((size_t)task % (size_t)row_block_count) * + TENSORLIB_MATMUL_MC; + int row_end = row_start + TENSORLIB_MATMUL_MC; + if (row_end > lhs_info.rows) row_end = lhs_info.rows; matmul_packed_rhs_batch(lhs, &lhs_info, rhs, output, batch_dims, - batch_ndim, (size_t)batch); + batch_ndim, batch, row_start, row_end); } } else { for (size_t batch = 0; batch < batch_count; ++batch) { matmul_packed_rhs_batch(lhs, &lhs_info, rhs, output, batch_dims, - batch_ndim, batch); + batch_ndim, batch, 0, lhs_info.rows); } } @@ -765,9 +778,10 @@ static void matmul_2d_packed_rhs_scalar_tails(const tensor* lhs, int lhs_base, TENSORLIB_AVX2_TARGET static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, - const float* packed_rhs, int rows, - int inner, int columns, tensor* output, - int output_base, int output_batch_ndim) { + const float* packed_rhs, int inner, + int columns, tensor* output, + int output_base, int output_batch_ndim, + int row_start, int row_end) { size_t workspace_values; if (!checked_size_multiply(TENSORLIB_MATMUL_MC, TENSORLIB_MATMUL_KC, &workspace_values)) { @@ -776,15 +790,16 @@ static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, float* packed_lhs = (float*)matmul_aligned_malloc(workspace_values * sizeof(float)); if (packed_lhs == NULL) return 0; - int full_rows = rows - rows % TENSORLIB_MATMUL_MR; + int row_count = row_end - row_start; + int full_rows = row_count - row_count % TENSORLIB_MATMUL_MR; int full_columns = columns - columns % TENSORLIB_MATMUL_NR; int output_stride = output->strides[output_batch_ndim]; for (int row_block = 0; row_block < full_rows; row_block += TENSORLIB_MATMUL_MC) { - int row_end = row_block + TENSORLIB_MATMUL_MC; - if (row_end > full_rows) row_end = full_rows; - int row_count = row_end - row_block; + int block_end = row_block + TENSORLIB_MATMUL_MC; + if (block_end > full_rows) block_end = full_rows; + int block_rows = block_end - row_block; for (int k_start = 0; k_start < inner; k_start += TENSORLIB_MATMUL_KC) { int k_end = k_start + TENSORLIB_MATMUL_KC; @@ -792,7 +807,7 @@ static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, int k_count = k_end - k_start; int accumulate = (k_start != 0); - pack_lhs_block(lhs, lhs_base, row_block, row_count, + pack_lhs_block(lhs, lhs_base, row_start + row_block, block_rows, k_start, k_count, packed_lhs); for (int column = 0; column < full_columns; @@ -802,23 +817,27 @@ static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, TENSORLIB_MATMUL_NR + (size_t)k_start * TENSORLIB_MATMUL_NR; - for (int row = row_block; row < row_end; + for (int row = row_block; row < block_end; row += TENSORLIB_MATMUL_MR) { matmul_4x16_packed_a_kernel( packed_lhs + (size_t)(row - row_block) * (size_t)k_count, rhs_panel, output->storage->data + output_base + - row * output_stride + column, + (row_start + row) * output_stride + column, output_stride, k_count, accumulate); } } } } - matmul_2d_packed_rhs_scalar_tails(lhs, lhs_base, packed_rhs, - rows, inner, columns, + matmul_2d_packed_rhs_scalar_tails( + lhs, + lhs_base + row_start * lhs->strides[lhs->ndim - 2], + packed_rhs, row_count, inner, columns, full_rows, full_columns, - output, output_base, output_batch_ndim); + output, + output_base + row_start * output_stride, + output_batch_ndim); matmul_aligned_free(packed_lhs); return 1; } From 0915da62316e7f3bffa8ba40095ab35f968f2f9b Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:02:18 +0200 Subject: [PATCH 03/49] perf: cache packed Linear weights --- include/tensorlib/autograd_internal.h | 6 ++++++ include/tensorlib/nn.h | 4 ++++ src/autograd/autograd_matmul.c | 18 +++++++++++++++--- src/nn/linear.c | 23 ++++++++++++++++++++++- 4 files changed, 47 insertions(+), 4 deletions(-) diff --git a/include/tensorlib/autograd_internal.h b/include/tensorlib/autograd_internal.h index 8b381ae..b2990fe 100644 --- a/include/tensorlib/autograd_internal.h +++ b/include/tensorlib/autograd_internal.h @@ -13,4 +13,10 @@ ag_tensor* ag_make_result(tensor* output, tensor* ag_full_like(const tensor* reference, float value); +ag_tensor* ag_matmul_packed_rhs( + const ag_tensor* a, + const ag_tensor* b, + const tensor_matmul_packed_rhs* packed_rhs +); + #endif diff --git a/include/tensorlib/nn.h b/include/tensorlib/nn.h index 6a1173e..4f3ac0a 100644 --- a/include/tensorlib/nn.h +++ b/include/tensorlib/nn.h @@ -102,6 +102,10 @@ struct nn_linear { int in_features; int out_features; int use_bias; + + /* Runtime-only packed view of weight^T, refreshed by storage version. */ + tensor_matmul_packed_rhs* packed_weight; + uint64_t packed_weight_version; }; struct nn_embedding { diff --git a/src/autograd/autograd_matmul.c b/src/autograd/autograd_matmul.c index 85a46f5..42b48c2 100644 --- a/src/autograd/autograd_matmul.c +++ b/src/autograd/autograd_matmul.c @@ -75,10 +75,22 @@ static int backward_matmul(const ag_node* node, return 1; } -ag_tensor* ag_matmul(const ag_tensor* a, const ag_tensor* b) { - if (a == NULL || b == NULL) return NULL; - tensor* output = t_matmul(a->value, b->value); +static ag_tensor* make_matmul_result(const ag_tensor* a, + const ag_tensor* b, + tensor* output) { ag_tensor* inputs[2] = {(ag_tensor*)a, (ag_tensor*)b}; return ag_make_result(output, AG_OP_MATMUL, 2, inputs, backward_matmul, NULL, NULL); } + +ag_tensor* ag_matmul(const ag_tensor* a, const ag_tensor* b) { + if (a == NULL || b == NULL) return NULL; + return make_matmul_result(a, b, t_matmul(a->value, b->value)); +} + +ag_tensor* ag_matmul_packed_rhs(const ag_tensor* a, + const ag_tensor* b, + const tensor_matmul_packed_rhs* packed_rhs) { + if (a == NULL || b == NULL || packed_rhs == NULL) return NULL; + return make_matmul_result(a, b, t_matmul_packed_rhs(a->value, packed_rhs)); +} diff --git a/src/nn/linear.c b/src/nn/linear.c index de9f294..57dbd54 100644 --- a/src/nn/linear.c +++ b/src/nn/linear.c @@ -3,6 +3,7 @@ #include #include "nn_internal.h" +#include "../../include/tensorlib/autograd_internal.h" static char* nn_parameter_name(const char* module_name, const char* suffix) { @@ -100,6 +101,7 @@ nn_linear* nn_linear_create(const char* name, void nn_linear_destroy(nn_linear* layer) { if (layer == NULL) return; + t_free_matmul_packed_rhs(layer->packed_weight); nn_module_destroy_base(&layer->base); free(layer); } @@ -110,12 +112,15 @@ ag_tensor* nn_linear_forward(const nn_linear* layer, const ag_tensor* input) ag_tensor* product; ag_tensor* result; tensor* input_value; + uint64_t weight_version; + nn_linear* mutable_layer; if (layer == NULL || input == NULL || layer->weight == NULL || layer->weight->value == NULL) { return NULL; } input_value = input->value; + mutable_layer = (nn_linear*)layer; if (!tensor_has_valid_metadata(input_value) || input_value->ndim < 1 || input_value->dims[input_value->ndim - 1] != layer->in_features) { return NULL; @@ -123,7 +128,23 @@ ag_tensor* nn_linear_forward(const nn_linear* layer, const ag_tensor* input) transposed = ag_transpose(layer->weight->value, 0, 1); if (transposed == NULL) return NULL; - product = ag_matmul(input, transposed); + weight_version = layer->weight->value->value->storage->version; + if (mutable_layer->packed_weight == NULL || + mutable_layer->packed_weight_version != weight_version) { + tensor_matmul_packed_rhs* packed = + t_pack_matmul_rhs(transposed->value); + if (packed != NULL) { + t_free_matmul_packed_rhs(mutable_layer->packed_weight); + mutable_layer->packed_weight = packed; + mutable_layer->packed_weight_version = weight_version; + } + } + product = input_value->ndim >= 2 && + mutable_layer->packed_weight != NULL && + mutable_layer->packed_weight_version == weight_version + ? ag_matmul_packed_rhs(input, transposed, + mutable_layer->packed_weight) + : ag_matmul(input, transposed); ag_tensor_release(transposed); if (product == NULL) return NULL; if (!layer->use_bias) return product; From fed797250d6770fe363f8b31c2de4c3dae5e0e46 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:03:19 +0200 Subject: [PATCH 04/49] test: cover LayerNorm rank and constant-row edges --- tests/unit/nn/test_nn_layer_norm.c | 45 ++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/unit/nn/test_nn_layer_norm.c b/tests/unit/nn/test_nn_layer_norm.c index 0423bdd..2511aea 100644 --- a/tests/unit/nn/test_nn_layer_norm.c +++ b/tests/unit/nn/test_nn_layer_norm.c @@ -97,10 +97,55 @@ static void test_non_affine_and_invalid(void) nn_layer_norm_destroy(NULL); } +static void test_rank_three_constant_rows(void) +{ + int dims[] = {2, 2, 4}; + nn_layer_norm* layer = nn_layer_norm_create("rank3", 4, 1e-5f, 1); + tensor* raw = t_alloc(3, dims); + ag_tensor* input; + ag_tensor* output; + ag_tensor* reduced; + ag_tensor* loss; + + CHECK(layer != NULL && raw != NULL); + if (layer == NULL || raw == NULL) { + t_free(raw); + nn_layer_norm_destroy(layer); + return; + } + for (int i = 0; i < tensor_numel(raw); ++i) { + raw->storage->data[i] = i < 4 ? 3.0f : (float)(i - 3); + } + input = ag_from_owned_tensor(raw, 1); + output = nn_layer_norm_forward(layer, input); + CHECK(output != NULL); + if (output != NULL) { + for (int i = 0; i < 4; ++i) CHECK(fabsf(output->value->storage->data[i]) < 1e-6f); + for (int i = 4; i < tensor_numel(output->value); ++i) { + CHECK(isfinite(output->value->storage->data[i])); + } + reduced = ag_sum(output, 0, 0); + loss = reduced == NULL ? NULL : ag_sum(reduced, 0, 0); + ag_tensor_release(reduced); + reduced = loss; + loss = reduced == NULL ? NULL : ag_sum(reduced, 0, 0); + ag_tensor_release(reduced); + CHECK(loss != NULL && ag_backward(loss) == 0); + if (input->grad != NULL) { + for (int i = 0; i < 4; ++i) CHECK(fabsf(input->grad->storage->data[i]) < 1e-6f); + } + ag_tensor_release(loss); + } + ag_tensor_release(output); + ag_tensor_release(input); + nn_layer_norm_destroy(layer); +} + int main(void) { test_forward_and_backward(); test_non_affine_and_invalid(); + test_rank_three_constant_rows(); if (failures != 0) { fprintf(stderr, "%d LayerNorm checks failed\n", failures); return 1; From 064fb6ee17b91834e0ee31f533c51bbb4705e151 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:06:11 +0200 Subject: [PATCH 05/49] perf: add fused LayerNorm autograd kernel --- src/autograd/autograd_layer_norm.c | 183 +++++++++++++++++++++++++++++ 1 file changed, 183 insertions(+) create mode 100644 src/autograd/autograd_layer_norm.c diff --git a/src/autograd/autograd_layer_norm.c b/src/autograd/autograd_layer_norm.c new file mode 100644 index 0000000..bcc79a1 --- /dev/null +++ b/src/autograd/autograd_layer_norm.c @@ -0,0 +1,183 @@ +#include +#include + +#include "../../include/tensorlib/autograd_internal.h" + +typedef struct { + int rows; + int width; + float* means; + float* inverse_stds; +} layer_norm_context; + +static void free_layer_norm_context(void* opaque) +{ + layer_norm_context* context = (layer_norm_context*)opaque; + if (context == NULL) return; + free(context->means); + free(context->inverse_stds); + free(context); +} + +static int row_base(const tensor* value, int row, int width) +{ + if (is_contiguous((tensor*)value)) return value->offset + row * width; + int base = value->offset; + int remaining = row; + for (int axis = value->ndim - 2; axis >= 0; --axis) { + int coordinate = remaining % value->dims[axis]; + remaining /= value->dims[axis]; + base += coordinate * value->strides[axis]; + } + return base; +} + +static int backward_layer_norm(const ag_node* node, + const tensor* output_gradient, + tensor** input_gradients) +{ + layer_norm_context* context; + tensor* input; + tensor* weight; + int width; + float width_f; + tensor* input_gradient = NULL; + tensor* weight_gradient = NULL; + tensor* bias_gradient = NULL; + + if (node == NULL || input_gradients == NULL || + !tensor_has_valid_metadata(output_gradient)) return 1; + context = (layer_norm_context*)node->context; + if (context == NULL) return 1; + input = node->inputs[0]->value; + weight = node->input_count == 3 ? node->inputs[1]->value : NULL; + width = context->width; + width_f = (float)width; + if (node->inputs[0]->requires_grad) { + input_gradient = t_alloc(input->ndim, input->dims); + if (input_gradient == NULL) goto fail; + } + if (node->input_count == 3 && node->inputs[1]->requires_grad) { + weight_gradient = t_alloc(1, &width); + if (weight_gradient == NULL) goto fail; + } + if (node->input_count == 3 && node->inputs[2]->requires_grad) { + bias_gradient = t_alloc(1, &width); + if (bias_gradient == NULL) goto fail; + } + + for (int row = 0; row < context->rows; ++row) { + int input_base = row_base(input, row, width); + int gradient_base = row_base(output_gradient, row, width); + float sum = 0.0f; + float weighted_sum = 0.0f; + for (int k = 0; k < width; ++k) { + float normalized = (input->storage->data[input_base + k * + input->strides[input->ndim - 1]] - context->means[row]) * + context->inverse_stds[row]; + float upstream = output_gradient->storage->data[gradient_base + k * + output_gradient->strides[output_gradient->ndim - 1]]; + float scale = weight == NULL ? 1.0f : weight->storage->data[k]; + sum += upstream * scale; + weighted_sum += upstream * scale * normalized; + if (weight_gradient != NULL) { + weight_gradient->storage->data[k] += upstream * normalized; + } + if (bias_gradient != NULL) bias_gradient->storage->data[k] += upstream; + } + if (input_gradient != NULL) { + for (int k = 0; k < width; ++k) { + float normalized = (input->storage->data[input_base + k * + input->strides[input->ndim - 1]] - context->means[row]) * + context->inverse_stds[row]; + float upstream = output_gradient->storage->data[gradient_base + k * + output_gradient->strides[output_gradient->ndim - 1]]; + float scale = weight == NULL ? 1.0f : weight->storage->data[k]; + input_gradient->storage->data[row * width + k] = + context->inverse_stds[row] * scale * + (upstream - sum / width_f - normalized * weighted_sum / width_f); + } + } + } + input_gradients[0] = input_gradient; + if (node->input_count == 3) { + input_gradients[1] = weight_gradient; + input_gradients[2] = bias_gradient; + } + return 0; + +fail: + t_free(input_gradient); + t_free(weight_gradient); + t_free(bias_gradient); + return 1; +} + +ag_tensor* ag_layer_norm(const ag_tensor* input, + const ag_tensor* weight, + const ag_tensor* bias, + float epsilon) +{ + tensor* output = NULL; + layer_norm_context* context = NULL; + ag_tensor* inputs[3]; + int width; + float width_f; + + if (input == NULL || !tensor_has_valid_metadata(input->value) || + input->value->ndim < 1 || !isfinite(epsilon) || epsilon <= 0.0f) return NULL; + width = input->value->dims[input->value->ndim - 1]; + width_f = (float)width; + if ((weight == NULL) != (bias == NULL) || + (weight != NULL && (!tensor_has_valid_metadata(weight->value) || + !tensor_has_valid_metadata(bias->value) || + weight->value->ndim != 1 || bias->value->ndim != 1 || + weight->value->dims[0] != width || bias->value->dims[0] != width))) { + return NULL; + } + output = t_alloc(input->value->ndim, input->value->dims); + context = (layer_norm_context*)calloc(1, sizeof(*context)); + if (output == NULL || context == NULL) goto fail; + context->width = width; + context->rows = tensor_numel(input->value) / width; + context->means = (float*)calloc((size_t)context->rows, sizeof(float)); + context->inverse_stds = (float*)calloc((size_t)context->rows, sizeof(float)); + if (context->means == NULL || context->inverse_stds == NULL) goto fail; + + for (int row = 0; row < context->rows; ++row) { + int base = row_base(input->value, row, width); + float mean = 0.0f; + float variance = 0.0f; + for (int k = 0; k < width; ++k) mean += input->value->storage->data[ + base + k * input->value->strides[input->value->ndim - 1]]; + mean /= width_f; + for (int k = 0; k < width; ++k) { + float delta = input->value->storage->data[ + base + k * input->value->strides[input->value->ndim - 1]] - mean; + variance += delta * delta; + } + context->means[row] = mean; + context->inverse_stds[row] = 1.0f / sqrtf(variance / width_f + epsilon); + for (int k = 0; k < width; ++k) { + float normalized = (input->value->storage->data[ + base + k * input->value->strides[input->value->ndim - 1]] - mean) * + context->inverse_stds[row]; + output->storage->data[row * width + k] = + normalized * (weight == NULL ? 1.0f : weight->value->storage->data[k]) + + (bias == NULL ? 0.0f : bias->value->storage->data[k]); + } + } + inputs[0] = (ag_tensor*)input; + if (weight == NULL) return ag_make_result(output, AG_OP_LAYER_NORM, 1, + inputs, backward_layer_norm, + context, free_layer_norm_context); + inputs[1] = (ag_tensor*)weight; + inputs[2] = (ag_tensor*)bias; + return ag_make_result(output, AG_OP_LAYER_NORM, 3, inputs, + backward_layer_norm, context, free_layer_norm_context); + +fail: + t_free(output); + free_layer_norm_context(context); + return NULL; +} From af42553537adc9b93026b4446f80c1c3ba097499 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:06:21 +0200 Subject: [PATCH 06/49] perf: wire LayerNorm through fused autograd --- CMakeLists.txt | 1 + include/tensorlib/autograd.h | 3 +- include/tensorlib/autograd_internal.h | 5 ++ src/nn/layer_norm.c | 66 +++------------------------ 4 files changed, 15 insertions(+), 60 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 99a4a24..8ac5e06 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -49,6 +49,7 @@ set(TENSORLIB_SOURCES src/autograd/autograd_gather.c src/autograd/autograd_reduc.c src/autograd/autograd_matmul.c + src/autograd/autograd_layer_norm.c src/autograd/autograd_backward.c src/init/rng.c src/nn/parameter.c diff --git a/include/tensorlib/autograd.h b/include/tensorlib/autograd.h index e0333fc..b16d3eb 100644 --- a/include/tensorlib/autograd.h +++ b/include/tensorlib/autograd.h @@ -48,7 +48,8 @@ typedef enum { AG_OP_EXPAND, AG_OP_GATHER_ROWS, AG_OP_MUL_SCALAR, - AG_OP_DIV_SCALAR + AG_OP_DIV_SCALAR, + AG_OP_LAYER_NORM } ag_op; diff --git a/include/tensorlib/autograd_internal.h b/include/tensorlib/autograd_internal.h index b2990fe..94d6336 100644 --- a/include/tensorlib/autograd_internal.h +++ b/include/tensorlib/autograd_internal.h @@ -19,4 +19,9 @@ ag_tensor* ag_matmul_packed_rhs( const tensor_matmul_packed_rhs* packed_rhs ); +ag_tensor* ag_layer_norm(const ag_tensor* input, + const ag_tensor* weight, + const ag_tensor* bias, + float epsilon); + #endif diff --git a/src/nn/layer_norm.c b/src/nn/layer_norm.c index 89e26e0..cea593e 100644 --- a/src/nn/layer_norm.c +++ b/src/nn/layer_norm.c @@ -4,6 +4,7 @@ #include #include "nn_internal.h" +#include "../../include/tensorlib/autograd_internal.h" static char* layer_norm_parameter_name(const char* module_name, const char* suffix) @@ -108,30 +109,9 @@ void nn_layer_norm_destroy(nn_layer_norm* layer) free(layer); } -static ag_tensor* scalar(float value) -{ - tensor* raw = t_alloc(0, NULL); - - if (raw == NULL) return NULL; - raw->storage->data[0] = value; - return ag_from_owned_tensor(raw, 0); -} - ag_tensor* nn_layer_norm_forward(const nn_layer_norm* layer, const ag_tensor* input) { - ag_tensor* mean = NULL; - ag_tensor* centered = NULL; - ag_tensor* squared = NULL; - ag_tensor* variance = NULL; - ag_tensor* epsilon = NULL; - ag_tensor* stabilized = NULL; - ag_tensor* deviation = NULL; - ag_tensor* normalized = NULL; - ag_tensor* scaled = NULL; - ag_tensor* result = NULL; - int axis; - if (layer == NULL || input == NULL || !tensor_has_valid_metadata(input->value) || input->value->ndim < 1 || @@ -139,45 +119,13 @@ ag_tensor* nn_layer_norm_forward(const nn_layer_norm* layer, layer->normalized_width) { return NULL; } - axis = input->value->ndim - 1; - mean = ag_mean(input, axis, 1); - if (mean == NULL) goto cleanup; - centered = ag_sub(input, mean); - if (centered == NULL) goto cleanup; - squared = ag_mul(centered, centered); - if (squared == NULL) goto cleanup; - variance = ag_mean(squared, axis, 1); - if (variance == NULL) goto cleanup; - epsilon = scalar(layer->epsilon); - if (epsilon == NULL) goto cleanup; - stabilized = ag_add(variance, epsilon); - if (stabilized == NULL) goto cleanup; - deviation = ag_sqrt(stabilized); - if (deviation == NULL) goto cleanup; - normalized = ag_div(centered, deviation); - if (normalized == NULL) goto cleanup; - if (!layer->affine) { - result = normalized; - normalized = NULL; - goto cleanup; - } if (layer->weight == NULL || layer->bias == NULL || layer->weight->value == NULL || layer->bias->value == NULL) { - goto cleanup; + return layer->affine ? NULL : ag_layer_norm(input, NULL, NULL, + layer->epsilon); } - scaled = ag_mul(normalized, layer->weight->value); - if (scaled == NULL) goto cleanup; - result = ag_add(scaled, layer->bias->value); - -cleanup: - ag_tensor_release(scaled); - ag_tensor_release(normalized); - ag_tensor_release(deviation); - ag_tensor_release(stabilized); - ag_tensor_release(epsilon); - ag_tensor_release(variance); - ag_tensor_release(squared); - ag_tensor_release(centered); - ag_tensor_release(mean); - return result; + return layer->affine + ? ag_layer_norm(input, layer->weight->value, + layer->bias->value, layer->epsilon) + : ag_layer_norm(input, NULL, NULL, layer->epsilon); } From 6c52a4bd07a850de526d4354d4f52395fe3f5fc4 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:06:53 +0200 Subject: [PATCH 07/49] test: cover fused classification edge cases --- tests/unit/nn/test_nn_loss.c | 42 ++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/unit/nn/test_nn_loss.c b/tests/unit/nn/test_nn_loss.c index 8ee0df9..029be65 100644 --- a/tests/unit/nn/test_nn_loss.c +++ b/tests/unit/nn/test_nn_loss.c @@ -152,6 +152,47 @@ static void test_vector_and_rank_three(void) ag_tensor_release(vector_logits); } +static void test_rank_three_softmax_and_repeated_targets(void) +{ + int logit_dims[] = {2, 2, 4}; + int target_dims[] = {2, 2}; + float logits[] = {10000.0f, 9999.0f, -10000.0f, 0.0f, + -3.0f, -2.0f, -1.0f, 0.0f, + 1.0f, 2.0f, 3.0f, 4.0f, + 4.0f, 3.0f, 2.0f, 1.0f}; + float targets_data[] = {1, 1, 1, 1}; + ag_tensor* input = make_ag(3, logit_dims, logits, 1); + tensor* targets = make_tensor(2, target_dims, targets_data); + ag_tensor* probabilities = nn_softmax(input); + ag_tensor* loss = nn_cross_entropy(input, targets); + + CHECK(probabilities != NULL && loss != NULL); + if (probabilities != NULL) { + for (int row = 0; row < 4; ++row) { + float sum = 0.0f; + for (int column = 0; column < 4; ++column) { + float value = probabilities->value->storage->data[row * 4 + column]; + CHECK(isfinite(value)); + sum += value; + } + CHECK_NEAR(sum, 1.0f, 2.0e-6f); + } + } + if (loss != NULL) { + CHECK(ag_backward(loss) == 0); + CHECK(input->grad != NULL); + if (input->grad != NULL) { + for (int row = 0; row < 4; ++row) { + CHECK(input->grad->storage->data[row * 4 + 1] < 0.0f); + } + } + } + ag_tensor_release(loss); + ag_tensor_release(probabilities); + t_free(targets); + ag_tensor_release(input); +} + static void test_invalid(void) { int dims[] = {2, 3}; @@ -184,6 +225,7 @@ int main(void) test_softmax(); test_cross_entropy(); test_vector_and_rank_three(); + test_rank_three_softmax_and_repeated_targets(); test_invalid(); if (failures != 0) { fprintf(stderr, "%d classification checks failed\n", failures); From b81aa44ddbded865580331164c7349efe353bc15 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:09:32 +0200 Subject: [PATCH 08/49] perf: add fused row-wise softmax kernels --- src/autograd/autograd_softmax.c | 141 ++++++++++++++++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 src/autograd/autograd_softmax.c diff --git a/src/autograd/autograd_softmax.c b/src/autograd/autograd_softmax.c new file mode 100644 index 0000000..e979b10 --- /dev/null +++ b/src/autograd/autograd_softmax.c @@ -0,0 +1,141 @@ +#include +#include + +#include "../../include/tensorlib/autograd_internal.h" + +typedef struct { + int rows; + int width; + int log_softmax; + int causal; +} softmax_context; + +static int row_base(const tensor* value, int row, int width) +{ + if (is_contiguous((tensor*)value)) return value->offset + row * width; + int base = value->offset; + int remaining = row; + for (int axis = value->ndim - 2; axis >= 0; --axis) { + int coordinate = remaining % value->dims[axis]; + remaining /= value->dims[axis]; + base += coordinate * value->strides[axis]; + } + return base; +} + +static int allowed_width(const softmax_context* context, int row) +{ + return context->causal ? row % context->width + 1 : context->width; +} + +static int backward_softmax(const ag_node* node, + const tensor* output_gradient, + tensor** input_gradients) +{ + softmax_context* context; + tensor* output; + tensor* gradient; + + if (node == NULL || input_gradients == NULL || + !tensor_has_valid_metadata(output_gradient)) return 1; + context = (softmax_context*)node->context; + if (context == NULL || !node->inputs[0]->requires_grad) return 0; + output = node->output->value; + gradient = t_alloc(output->ndim, output->dims); + if (gradient == NULL) return 1; + + for (int row = 0; row < context->rows; ++row) { + int width = allowed_width(context, row); + int output_base = row_base(output, row, context->width); + int upstream_base = row_base(output_gradient, row, context->width); + float total = 0.0f; + if (context->log_softmax) { + for (int k = 0; k < width; ++k) { + total += output_gradient->storage->data[upstream_base + k * + output_gradient->strides[output_gradient->ndim - 1]]; + } + } else { + for (int k = 0; k < width; ++k) { + total += output_gradient->storage->data[upstream_base + k * + output_gradient->strides[output_gradient->ndim - 1]] * + output->storage->data[output_base + k]; + } + } + for (int k = 0; k < context->width; ++k) { + float upstream = output_gradient->storage->data[upstream_base + k * + output_gradient->strides[output_gradient->ndim - 1]]; + float probability = context->log_softmax + ? expf(output->storage->data[output_base + k]) + : output->storage->data[output_base + k]; + gradient->storage->data[row * context->width + k] = k < width + ? (context->log_softmax + ? upstream - probability * total + : probability * (upstream - total)) + : 0.0f; + } + } + input_gradients[0] = gradient; + return 0; +} + +ag_tensor* ag_softmax_last_dim(const ag_tensor* input, + int log_softmax, + int causal) +{ + tensor* output; + softmax_context* context; + int width; + + if (input == NULL || !tensor_has_valid_metadata(input->value) || + input->value->ndim < 1) return NULL; + width = input->value->dims[input->value->ndim - 1]; + if (causal && (input->value->ndim < 2 || + input->value->dims[input->value->ndim - 2] != width)) return NULL; + output = t_alloc(input->value->ndim, input->value->dims); + context = (softmax_context*)calloc(1, sizeof(*context)); + if (output == NULL || context == NULL) goto fail; + context->width = width; + context->rows = tensor_numel(input->value) / width; + context->log_softmax = log_softmax != 0; + context->causal = causal != 0; + + for (int row = 0; row < context->rows; ++row) { + int input_base = row_base(input->value, row, width); + int output_base = row * width; + int count = allowed_width(context, row); + float maximum = -INFINITY; + float sum = 0.0f; + for (int k = 0; k < count; ++k) { + float value = input->value->storage->data[input_base + k * + input->value->strides[input->value->ndim - 1]]; + if (isnan(value)) maximum = value; + else if (value > maximum) maximum = value; + } + for (int k = 0; k < count; ++k) { + float exponential = expf(input->value->storage->data[input_base + k * + input->value->strides[input->value->ndim - 1]] - maximum); + output->storage->data[output_base + k] = exponential; + sum += exponential; + } + float log_sum = logf(sum); + for (int k = 0; k < width; ++k) { + if (k >= count) output->storage->data[output_base + k] = + context->log_softmax ? -INFINITY : 0.0f; + else if (context->log_softmax) output->storage->data[output_base + k] = + input->value->storage->data[input_base + k * + input->value->strides[input->value->ndim - 1]] - maximum - log_sum; + else output->storage->data[output_base + k] /= sum; + } + } + { + ag_tensor* inputs[1] = {(ag_tensor*)input}; + return ag_make_result(output, + context->log_softmax ? AG_OP_LOG_SOFTMAX : AG_OP_SOFTMAX, + 1, inputs, backward_softmax, context, free); + } + +fail: + t_free(output); + free(context); + return NULL; +} From afba7a41dff4259c13bb65a115083b38aaa3c8d2 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:09:39 +0200 Subject: [PATCH 09/49] perf: add indexed cross-entropy kernel --- src/autograd/autograd_cross_entropy.c | 148 ++++++++++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 src/autograd/autograd_cross_entropy.c diff --git a/src/autograd/autograd_cross_entropy.c b/src/autograd/autograd_cross_entropy.c new file mode 100644 index 0000000..09154ab --- /dev/null +++ b/src/autograd/autograd_cross_entropy.c @@ -0,0 +1,148 @@ +#include +#include + +#include "../../include/tensorlib/autograd_internal.h" + +typedef struct { + int rows; + int classes; + int* targets; +} cross_entropy_context; + +static void free_cross_entropy_context(void* opaque) +{ + cross_entropy_context* context = (cross_entropy_context*)opaque; + if (context == NULL) return; + free(context->targets); + free(context); +} + +static int flat_index(const tensor* value, int flat) +{ + int index = value->offset; + int remaining = flat; + for (int axis = value->ndim - 1; axis >= 0; --axis) { + int coordinate = remaining % value->dims[axis]; + remaining /= value->dims[axis]; + index += coordinate * value->strides[axis]; + } + return index; +} + +static int row_base(const tensor* value, int row, int width) +{ + if (is_contiguous((tensor*)value)) return value->offset + row * width; + int base = value->offset; + int remaining = row; + for (int axis = value->ndim - 2; axis >= 0; --axis) { + int coordinate = remaining % value->dims[axis]; + remaining /= value->dims[axis]; + base += coordinate * value->strides[axis]; + } + return base; +} + +static int backward_cross_entropy(const ag_node* node, + const tensor* output_gradient, + tensor** input_gradients) +{ + cross_entropy_context* context; + tensor* input; + tensor* gradient; + float upstream; + + if (node == NULL || input_gradients == NULL || + !tensor_has_valid_metadata(output_gradient)) return 1; + context = (cross_entropy_context*)node->context; + if (context == NULL || !node->inputs[0]->requires_grad) return 0; + input = node->inputs[0]->value; + gradient = t_alloc(input->ndim, input->dims); + if (gradient == NULL) return 1; + upstream = output_gradient->storage->data[output_gradient->offset] / + (float)context->rows; + for (int row = 0; row < context->rows; ++row) { + int base = row_base(input, row, context->classes); + float maximum = -INFINITY; + float sum = 0.0f; + for (int k = 0; k < context->classes; ++k) { + float value = input->storage->data[base + k * + input->strides[input->ndim - 1]]; + if (isnan(value)) maximum = value; + else if (value > maximum) maximum = value; + } + for (int k = 0; k < context->classes; ++k) { + sum += expf(input->storage->data[base + k * + input->strides[input->ndim - 1]] - maximum); + } + for (int k = 0; k < context->classes; ++k) { + float probability = expf(input->storage->data[base + k * + input->strides[input->ndim - 1]] - maximum) / sum; + gradient->storage->data[row * context->classes + k] = upstream * + (probability - (k == context->targets[row] ? 1.0f : 0.0f)); + } + } + input_gradients[0] = gradient; + return 0; +} + +ag_tensor* ag_cross_entropy(const ag_tensor* logits, const tensor* targets) +{ + tensor* output = NULL; + cross_entropy_context* context = NULL; + int classes; + int rows; + + if (logits == NULL || !tensor_has_valid_metadata(logits->value) || + !tensor_has_valid_metadata(targets) || logits->value->ndim < 1 || + targets->ndim != logits->value->ndim - 1) return NULL; + for (int axis = 0; axis < targets->ndim; ++axis) { + if (targets->dims[axis] != logits->value->dims[axis]) return NULL; + } + classes = logits->value->dims[logits->value->ndim - 1]; + rows = tensor_numel((tensor*)targets); + context = (cross_entropy_context*)calloc(1, sizeof(*context)); + output = t_alloc(0, NULL); + if (context == NULL || output == NULL) goto fail; + context->rows = rows; + context->classes = classes; + context->targets = (int*)malloc((size_t)rows * sizeof(*context->targets)); + if (context->targets == NULL) goto fail; + for (int row = 0; row < rows; ++row) { + float target = targets->storage->data[flat_index(targets, row)]; + if (!isfinite(target) || floorf(target) != target || target < 0.0f || + target >= (float)classes) goto fail; + context->targets[row] = (int)target; + } + + float loss = 0.0f; + for (int row = 0; row < rows; ++row) { + int base = row_base(logits->value, row, classes); + float maximum = -INFINITY; + float sum = 0.0f; + for (int k = 0; k < classes; ++k) { + float value = logits->value->storage->data[base + k * + logits->value->strides[logits->value->ndim - 1]]; + if (isnan(value)) maximum = value; + else if (value > maximum) maximum = value; + } + for (int k = 0; k < classes; ++k) { + sum += expf(logits->value->storage->data[base + k * + logits->value->strides[logits->value->ndim - 1]] - maximum); + } + loss -= logits->value->storage->data[base + + context->targets[row] * logits->value->strides[ + logits->value->ndim - 1]] - maximum - logf(sum); + } + output->storage->data[0] = loss / (float)rows; + { + ag_tensor* inputs[1] = {(ag_tensor*)logits}; + return ag_make_result(output, AG_OP_CROSS_ENTROPY, 1, inputs, + backward_cross_entropy, context, + free_cross_entropy_context); + } + +fail: + t_free(output); + free_cross_entropy_context(context); + return NULL; +} From 589ecdb52737afcfc44d89c9322cf40a87a55233 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:09:48 +0200 Subject: [PATCH 10/49] perf: route classification APIs through fused kernels --- CMakeLists.txt | 2 + include/tensorlib/autograd.h | 5 +- include/tensorlib/autograd_internal.h | 6 ++ src/losses/classification.c | 132 +------------------------- 4 files changed, 16 insertions(+), 129 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ac5e06..0823cb3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -50,6 +50,8 @@ set(TENSORLIB_SOURCES src/autograd/autograd_reduc.c src/autograd/autograd_matmul.c src/autograd/autograd_layer_norm.c + src/autograd/autograd_softmax.c + src/autograd/autograd_cross_entropy.c src/autograd/autograd_backward.c src/init/rng.c src/nn/parameter.c diff --git a/include/tensorlib/autograd.h b/include/tensorlib/autograd.h index b16d3eb..e939df4 100644 --- a/include/tensorlib/autograd.h +++ b/include/tensorlib/autograd.h @@ -49,7 +49,10 @@ typedef enum { AG_OP_GATHER_ROWS, AG_OP_MUL_SCALAR, AG_OP_DIV_SCALAR, - AG_OP_LAYER_NORM + AG_OP_LAYER_NORM, + AG_OP_SOFTMAX, + AG_OP_LOG_SOFTMAX, + AG_OP_CROSS_ENTROPY } ag_op; diff --git a/include/tensorlib/autograd_internal.h b/include/tensorlib/autograd_internal.h index 94d6336..66a7b0b 100644 --- a/include/tensorlib/autograd_internal.h +++ b/include/tensorlib/autograd_internal.h @@ -24,4 +24,10 @@ ag_tensor* ag_layer_norm(const ag_tensor* input, const ag_tensor* bias, float epsilon); +ag_tensor* ag_softmax_last_dim(const ag_tensor* input, + int log_softmax, + int causal); + +ag_tensor* ag_cross_entropy(const ag_tensor* logits, const tensor* targets); + #endif diff --git a/src/losses/classification.c b/src/losses/classification.c index cacd566..4d41846 100644 --- a/src/losses/classification.c +++ b/src/losses/classification.c @@ -1,141 +1,17 @@ -#include - #include "../../include/tensorlib/nn.h" - -static float nn_tensor_at_flat(const tensor* value, int flat) -{ - int index = value->offset; - int remaining = flat; - - for (int dim = value->ndim - 1; dim >= 0; --dim) { - int coordinate = remaining % value->dims[dim]; - remaining /= value->dims[dim]; - index += coordinate * value->strides[dim]; - } - return value->storage->data[index]; -} +#include "../../include/tensorlib/autograd_internal.h" ag_tensor* nn_log_softmax(const ag_tensor* logits) { - ag_tensor* maximum; - ag_tensor* shifted; - ag_tensor* exponentials; - ag_tensor* sum; - ag_tensor* log_sum; - ag_tensor* result; - int axis; - - if (logits == NULL || !tensor_has_valid_metadata(logits->value) || - logits->value->ndim < 1) { - return NULL; - } - axis = logits->value->ndim - 1; - maximum = ag_max(logits, axis, 1); - if (maximum == NULL) return NULL; - shifted = ag_sub(logits, maximum); - ag_tensor_release(maximum); - if (shifted == NULL) return NULL; - exponentials = ag_exp(shifted); - if (exponentials == NULL) { - ag_tensor_release(shifted); - return NULL; - } - sum = ag_sum(exponentials, axis, 1); - ag_tensor_release(exponentials); - if (sum == NULL) { - ag_tensor_release(shifted); - return NULL; - } - log_sum = ag_log(sum); - ag_tensor_release(sum); - if (log_sum == NULL) { - ag_tensor_release(shifted); - return NULL; - } - result = ag_sub(shifted, log_sum); - ag_tensor_release(log_sum); - ag_tensor_release(shifted); - return result; + return ag_softmax_last_dim(logits, 1, 0); } ag_tensor* nn_softmax(const ag_tensor* logits) { - ag_tensor* log_probabilities = nn_log_softmax(logits); - ag_tensor* probabilities; - - if (log_probabilities == NULL) return NULL; - probabilities = ag_exp(log_probabilities); - ag_tensor_release(log_probabilities); - return probabilities; -} - -static int nn_targets_match(const tensor* logits, const tensor* targets) -{ - if (!tensor_has_valid_metadata(targets) || - targets->ndim != logits->ndim - 1) { - return 0; - } - for (int i = 0; i < targets->ndim; ++i) { - if (targets->dims[i] != logits->dims[i]) return 0; - } - return 1; + return ag_softmax_last_dim(logits, 0, 0); } ag_tensor* nn_cross_entropy(const ag_tensor* logits, const tensor* targets) { - ag_tensor* log_probabilities; - ag_tensor* selector; - ag_tensor* selected; - ag_tensor* class_sum; - ag_tensor* loss; - tensor* selector_value; - int classes; - int rows; - - if (logits == NULL || !tensor_has_valid_metadata(logits->value) || - logits->value->ndim < 1 || !nn_targets_match(logits->value, targets)) { - return NULL; - } - classes = logits->value->dims[logits->value->ndim - 1]; - rows = tensor_numel((tensor*)targets); - selector_value = t_alloc(logits->value->ndim, logits->value->dims); - if (selector_value == NULL) return NULL; - for (int i = 0; i < tensor_numel(selector_value); ++i) { - selector_value->storage->data[i] = 0.0f; - } - for (int row = 0; row < rows; ++row) { - float target = nn_tensor_at_flat(targets, row); - int class_index; - if (!isfinite(target) || floorf(target) != target || - target < 0.0f || target >= (float)classes) { - t_free(selector_value); - return NULL; - } - class_index = (int)target; - selector_value->storage->data[row * classes + class_index] = 1.0f; - } - selector = ag_from_owned_tensor(selector_value, 0); - if (selector == NULL) return NULL; - log_probabilities = nn_log_softmax(logits); - if (log_probabilities == NULL) { - ag_tensor_release(selector); - return NULL; - } - selected = ag_mul(log_probabilities, selector); - ag_tensor_release(log_probabilities); - ag_tensor_release(selector); - if (selected == NULL) return NULL; - class_sum = ag_sum(selected, selected->value->ndim - 1, 0); - ag_tensor_release(selected); - if (class_sum == NULL) return NULL; - loss = ag_neg(class_sum); - ag_tensor_release(class_sum); - if (loss == NULL) return NULL; - while (loss->value->ndim > 0) { - ag_tensor* reduced = ag_mean(loss, 0, 0); - ag_tensor_release(loss); - if (reduced == NULL) return NULL; - loss = reduced; - } - return loss; + return ag_cross_entropy(logits, targets); } From a7634c7e379e3859a45dcbfd8a1e0c25525f6a18 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:10:20 +0200 Subject: [PATCH 11/49] test: cover single-token causal attention --- tests/unit/nn/test_nn_multihead_attention.c | 23 +++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/unit/nn/test_nn_multihead_attention.c b/tests/unit/nn/test_nn_multihead_attention.c index 40792ad..6c6613c 100644 --- a/tests/unit/nn/test_nn_multihead_attention.c +++ b/tests/unit/nn/test_nn_multihead_attention.c @@ -218,10 +218,33 @@ static void test_topology_validation_and_eval_rng(void) nn_multihead_attention_destroy(NULL); } +static void test_single_token_causal_attention(void) +{ + nn_rng rng; + float values[2] = {2.0f, -3.0f}; + ag_tensor* input; + ag_tensor* output; + nn_multihead_attention* attention; + + nn_rng_seed(&rng, 44); + attention = identity_attention(&rng); + input = input_from(values, 1, 0); + output = nn_multihead_attention_forward(attention, input); + CHECK(output != NULL); + if (output != NULL) { + CHECK(fabsf(output->value->storage->data[0] - values[0]) < 1e-6f); + CHECK(fabsf(output->value->storage->data[1] - values[1]) < 1e-6f); + } + ag_tensor_release(output); + ag_tensor_release(input); + nn_multihead_attention_destroy(attention); +} + int main(void) { test_exact_forward_causality_and_backward(); test_topology_validation_and_eval_rng(); + test_single_token_causal_attention(); if (failures != 0) { fprintf(stderr, "%d multi-head-attention checks failed\n", failures); return 1; From 0020ee98188c6ed4b014927ab867593e5872cb2b Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:15:55 +0200 Subject: [PATCH 12/49] fix: preserve fused LayerNorm training gradients --- src/autograd/autograd_layer_norm.c | 11 +++++++++-- src/nn/multihead_attention.c | 7 ++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/autograd/autograd_layer_norm.c b/src/autograd/autograd_layer_norm.c index bcc79a1..a0a91ba 100644 --- a/src/autograd/autograd_layer_norm.c +++ b/src/autograd/autograd_layer_norm.c @@ -65,6 +65,12 @@ static int backward_layer_norm(const ag_node* node, bias_gradient = t_alloc(1, &width); if (bias_gradient == NULL) goto fail; } + if (weight_gradient != NULL) { + for (int k = 0; k < width; ++k) weight_gradient->storage->data[k] = 0.0f; + } + if (bias_gradient != NULL) { + for (int k = 0; k < width; ++k) bias_gradient->storage->data[k] = 0.0f; + } for (int row = 0; row < context->rows; ++row) { int input_base = row_base(input, row, width); @@ -94,8 +100,9 @@ static int backward_layer_norm(const ag_node* node, output_gradient->strides[output_gradient->ndim - 1]]; float scale = weight == NULL ? 1.0f : weight->storage->data[k]; input_gradient->storage->data[row * width + k] = - context->inverse_stds[row] * scale * - (upstream - sum / width_f - normalized * weighted_sum / width_f); + context->inverse_stds[row] * + (upstream * scale - sum / width_f - + normalized * weighted_sum / width_f); } } } diff --git a/src/nn/multihead_attention.c b/src/nn/multihead_attention.c index 787f92b..d9f4044 100644 --- a/src/nn/multihead_attention.c +++ b/src/nn/multihead_attention.c @@ -4,6 +4,7 @@ #include #include "nn_internal.h" +#include "../../include/tensorlib/autograd_internal.h" static ag_tensor* attention_module_forward(const nn_module* module, const ag_tensor* input) @@ -235,7 +236,6 @@ ag_tensor* nn_multihead_attention_forward( ag_tensor* transposed_key = NULL; ag_tensor* scores = NULL; ag_tensor* scaled = NULL; - ag_tensor* masked = NULL; ag_tensor* probabilities = NULL; ag_tensor* context = NULL; ag_tensor* transposed_context = NULL; @@ -274,9 +274,7 @@ ag_tensor* nn_multihead_attention_forward( if (scores == NULL) goto cleanup; scaled = ag_div_scalar(scores, sqrtf((float)attention->head_width)); if (scaled == NULL) goto cleanup; - masked = nn_apply_causal_mask(scaled); - if (masked == NULL) goto cleanup; - probabilities = nn_softmax(masked); + probabilities = ag_softmax_last_dim(scaled, 0, 1); if (probabilities == NULL) goto cleanup; context = ag_matmul(probabilities, value); if (context == NULL) goto cleanup; @@ -297,7 +295,6 @@ ag_tensor* nn_multihead_attention_forward( ag_tensor_release(transposed_context); ag_tensor_release(context); ag_tensor_release(probabilities); - ag_tensor_release(masked); ag_tensor_release(scaled); ag_tensor_release(scores); ag_tensor_release(transposed_key); From 147b6deb8544a42ab50cdd54d2f2c6f96aff3cfb Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:16:31 +0200 Subject: [PATCH 13/49] test: cover strided offset and broadcast layouts --- tests/unit/tensor/test_tensor_ops.c | 31 ++++++++++++++++++++++++++++ tests/unit/tensor/test_tensor_view.c | 28 +++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/tests/unit/tensor/test_tensor_ops.c b/tests/unit/tensor/test_tensor_ops.c index 03fb1a7..af4bdc2 100644 --- a/tests/unit/tensor/test_tensor_ops.c +++ b/tests/unit/tensor/test_tensor_ops.c @@ -241,6 +241,36 @@ TEST(test_binary_operations_transposed_inputs) { t_free(div); t_free(mul); t_free(sub); t_free(b); t_free(a); t_free(bb); t_free(ab); } +TEST(test_binary_operations_offset_transpose_and_broadcast) { + int base_dims[] = {3, 4}; + int bias_dims[] = {1, 3}; + tensor* base = t_alloc(2, base_dims); + tensor* transposed; + tensor* bias = t_alloc(2, bias_dims); + tensor* actual; + + ASSERT_NOT_NULL(base); + ASSERT_NOT_NULL(bias); + for (int i = 0; i < 12; ++i) base->storage->data[i] = (float)(i + 1); + for (int i = 0; i < 3; ++i) bias->storage->data[i] = (float)(10 * i); + transposed = t_transpose(base, 0, 1); + actual = t_add(transposed, bias); + ASSERT_NOT_NULL(actual); + ASSERT_EQ_INT(actual->dims[0], 4); + ASSERT_EQ_INT(actual->dims[1], 3); + for (int row = 0; row < 4; ++row) { + for (int column = 0; column < 3; ++column) { + ASSERT_EQ_FLOAT(actual->storage->data[row * 3 + column], + base->storage->data[column * 4 + row] + + bias->storage->data[column]); + } + } + t_free(actual); + t_free(transposed); + t_free(bias); + t_free(base); +} + TEST(test_all_operations_reject_null_inputs) { TEST_LOG("checking every P0 operation returns NULL for NULL tensor input"); unary_op ops[] = {t_neg, t_sqrt, t_exp, t_log, t_relu, t_gelu, t_sigmoid, t_tanh}; @@ -308,6 +338,7 @@ int main(void) { RUN_TEST(test_t_add_null_args_returns_null); RUN_TEST(test_sub_mul_div_contiguous_and_ieee_edges); RUN_TEST(test_binary_operations_transposed_inputs); + RUN_TEST(test_binary_operations_offset_transpose_and_broadcast); RUN_TEST(test_all_operations_reject_null_inputs); RUN_TEST(test_unary_operations_contiguous_values); RUN_TEST(test_unary_operations_transposed_inputs); diff --git a/tests/unit/tensor/test_tensor_view.c b/tests/unit/tensor/test_tensor_view.c index ccaf246..1ca2d9a 100644 --- a/tests/unit/tensor/test_tensor_view.c +++ b/tests/unit/tensor/test_tensor_view.c @@ -98,6 +98,33 @@ TEST(test_t_contiguous_on_transposed_materializes_correct_order) { t_free(a); } +TEST(test_t_contiguous_transposed_offset_non_square) { + int dims[] = {4, 5}; + tensor* base = t_alloc(2, dims); + tensor* slice; + tensor* transposed; + tensor* actual; + + ASSERT_NOT_NULL(base); + for (int i = 0; i < 20; ++i) base->storage->data[i] = (float)i; + slice = t_slice(base, 0, 1, 4); + transposed = t_transpose(slice, 0, 1); + actual = t_contiguous(transposed); + ASSERT_NOT_NULL(actual); + ASSERT_EQ_INT(actual->dims[0], 5); + ASSERT_EQ_INT(actual->dims[1], 3); + for (int row = 0; row < 5; ++row) { + for (int column = 0; column < 3; ++column) { + ASSERT_EQ_FLOAT(actual->storage->data[row * 3 + column], + base->storage->data[(column + 1) * 5 + row]); + } + } + t_free(actual); + t_free(transposed); + t_free(slice); + t_free(base); +} + TEST(test_t_contiguous_null_input) { ASSERT_NULL(t_contiguous(NULL)); } @@ -320,6 +347,7 @@ int main(void) { RUN_TEST(test_t_transpose_null_input); RUN_TEST(test_t_contiguous_on_already_contiguous_returns_equivalent_view); RUN_TEST(test_t_contiguous_on_transposed_materializes_correct_order); + RUN_TEST(test_t_contiguous_transposed_offset_non_square); RUN_TEST(test_t_contiguous_null_input); RUN_TEST(test_t_reshape_contiguous_input_returns_view); RUN_TEST(test_t_reshape_strided_input_materializes_contiguous_tensor); From 79343fe9a4c42bbb04679f0e4d4f8dac976c834d Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:17:36 +0200 Subject: [PATCH 14/49] perf: optimize strided tensor cloning --- src/tensor/tensor_alloc.c | 41 +++++++++++++++++++++++++++++++-------- 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/src/tensor/tensor_alloc.c b/src/tensor/tensor_alloc.c index 9bec3a5..914a96b 100644 --- a/src/tensor/tensor_alloc.c +++ b/src/tensor/tensor_alloc.c @@ -131,19 +131,44 @@ tensor* t_clone(tensor* t) { return a; } - int* coords = NULL; - if (t->ndim > 0) { - coords = (int*)calloc((size_t)t->ndim, sizeof(int)); - if (coords == NULL) { - t_free(a); - return NULL; + if (t->ndim == 2 && t->strides[0] == 1 && + t->strides[1] >= t->dims[0]) { + const int tile = 32; + for (int row_block = 0; row_block < t->dims[0]; row_block += tile) { + int row_end = row_block + tile; + if (row_end > t->dims[0]) row_end = t->dims[0]; + for (int column_block = 0; column_block < t->dims[1]; + column_block += tile) { + int column_end = column_block + tile; + if (column_end > t->dims[1]) column_end = t->dims[1]; + for (int row = row_block; row < row_end; ++row) { + for (int column = column_block; column < column_end; ++column) { + a->storage->data[row * t->dims[1] + column] = + t->storage->data[t->offset + row * t->strides[0] + + column * t->strides[1]]; + } + } + } } + return a; + } + + int* coords = (int*)calloc((size_t)t->ndim, sizeof(int)); + int src_idx = t->offset; + if (t->ndim > 0 && coords == NULL) { + t_free(a); + return NULL; } for (int i = 0; i < total_elements; i++) { - int src_idx = get_flat_index_nd(t, coords); a->storage->data[i] = t->storage->data[src_idx]; - advance_coords(coords, t->dims, t->ndim); + for (int axis = t->ndim - 1; axis >= 0; --axis) { + coords[axis]++; + src_idx += t->strides[axis]; + if (coords[axis] < t->dims[axis]) break; + coords[axis] = 0; + src_idx -= t->dims[axis] * t->strides[axis]; + } } free(coords); return a; From a32d2ef6571027d861f32eaceccbc6c5f6d8224f Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:18:26 +0200 Subject: [PATCH 15/49] perf: use stride cursors for broadcast iteration --- src/tensor/tensor_ops.c | 56 +++++++++++++++++++++++++---------------- 1 file changed, 34 insertions(+), 22 deletions(-) diff --git a/src/tensor/tensor_ops.c b/src/tensor/tensor_ops.c index 220b43b..190ed18 100644 --- a/src/tensor/tensor_ops.c +++ b/src/tensor/tensor_ops.c @@ -107,23 +107,11 @@ static int broadcast_output_shape(const tensor* a, return 1; } -static int input_index_for_broadcast(const tensor* input, - const int* output_coords, - int output_ndim) { - int index = input->offset; - int rank_offset = output_ndim - input->ndim; - - for (int output_axis = 0; output_axis < output_ndim; ++output_axis) { - int input_axis = output_axis - rank_offset; - if (input_axis < 0) continue; - - int coordinate = (input->dims[input_axis] == 1) - ? 0 - : output_coords[output_axis]; - index += coordinate * input->strides[input_axis]; - } - - return index; +static int broadcast_step(const tensor* input, int output_axis, int output_ndim) +{ + int input_axis = output_axis - (output_ndim - input->ndim); + if (input_axis < 0 || input->dims[input_axis] == 1) return 0; + return input->strides[input_axis]; } /* Fast path for element-wise binary ops where the output is contiguous and each @@ -237,21 +225,45 @@ static tensor* apply_binary(tensor* a, /* Handled by the streaming broadcast fast path. */ } else { int* coords = NULL; + int* steps_a = NULL; + int* steps_b = NULL; + int idx_a = a->offset; + int idx_b = b->offset; + int idx_c = c->offset; if (output_ndim > 0) { coords = (int*)calloc((size_t)output_ndim, sizeof(int)); + steps_a = (int*)malloc((size_t)output_ndim * sizeof(int)); + steps_b = (int*)malloc((size_t)output_ndim * sizeof(int)); } - if (output_ndim > 0 && coords == NULL) { + if (output_ndim > 0 && (coords == NULL || steps_a == NULL || + steps_b == NULL)) { + free(steps_b); + free(steps_a); + free(coords); t_free(c); return NULL; } + for (int axis = 0; axis < output_ndim; ++axis) { + steps_a[axis] = broadcast_step(a, axis, output_ndim); + steps_b[axis] = broadcast_step(b, axis, output_ndim); + } for (int i = 0; i < total_elements; i++) { - int idx_a = input_index_for_broadcast(a, coords, output_ndim); - int idx_b = input_index_for_broadcast(b, coords, output_ndim); - int idx_c = get_flat_index_nd(c, coords); c->storage->data[idx_c] = scalar_op(a->storage->data[idx_a], b->storage->data[idx_b]); - advance_coords(coords, c->dims, c->ndim); + for (int axis = output_ndim - 1; axis >= 0; --axis) { + coords[axis]++; + idx_a += steps_a[axis]; + idx_b += steps_b[axis]; + idx_c += c->strides[axis]; + if (coords[axis] < c->dims[axis]) break; + coords[axis] = 0; + idx_a -= c->dims[axis] * steps_a[axis]; + idx_b -= c->dims[axis] * steps_b[axis]; + idx_c -= c->dims[axis] * c->strides[axis]; + } } + free(steps_b); + free(steps_a); free(coords); } From 39dc5a27540832643042e83d15079a10da428d68 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 14:19:01 +0200 Subject: [PATCH 16/49] perf: parallelize large broadcast rows --- src/tensor/tensor_ops.c | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/tensor/tensor_ops.c b/src/tensor/tensor_ops.c index 190ed18..db329f2 100644 --- a/src/tensor/tensor_ops.c +++ b/src/tensor/tensor_ops.c @@ -166,7 +166,12 @@ static int try_contiguous_broadcast(tensor* a, tensor* b, const float* a_data = a->storage->data + a->offset; const float* b_data = b->storage->data + b->offset; float* c_data = c->storage->data + c->offset; + int threads = tensorlib_parallel_threads( + (long long)outer * block, TENSORLIB_OP_MIN_PARALLEL_ELEMENTS, outer); +#ifdef _OPENMP +#pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) +#endif for (int r = 0; r < outer; ++r) { /* Row-major decomposition of r over the outer axes 0..s. */ int remaining = r; From 7e358de65828eaaf0a5350c9a8aabc442cd67d01 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:10:14 +0200 Subject: [PATCH 17/49] Add training phase benchmark breakdown --- benchmarks/bench_nn.c | 82 +++++++++++++++++++++++++++++++++++++-- benchmarks/bench_runner.c | 52 ++++++++++++++++--------- benchmarks/bench_runner.h | 7 ++++ 3 files changed, 118 insertions(+), 23 deletions(-) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index ad70c5e..f7eac32 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -16,8 +16,20 @@ typedef struct { nn_adamw* adamw; nn_rng rng; int train; + double phase_seconds[6]; + long phase_calls; } nn_bench_context; +enum { + PHASE_ZERO_GRAD, + PHASE_FORWARD, + PHASE_LOSS, + PHASE_BACKWARD, + PHASE_ADAMW, + PHASE_GRAPH_RELEASE, + PHASE_COUNT +}; + static ag_tensor* make_input(int ndim, const int* dims, int token_ids) { tensor* raw = t_alloc(ndim, dims); @@ -44,23 +56,49 @@ static tensor* make_targets(int count, int classes) static int nn_operation(void* opaque, double* checksum) { nn_bench_context* context = (nn_bench_context*)opaque; - ag_tensor* output = context->forward(context->model, context->input); + ag_tensor* output; ag_tensor* loss = NULL; int status = 1; + double started; + + if (context->train) { + started = bench_now_seconds(); + if (context->sgd != NULL) { + nn_sgd_zero_grad(context->sgd); + } else if (context->adamw != NULL) { + nn_adamw_zero_grad(context->adamw); + } + context->phase_seconds[PHASE_ZERO_GRAD] += + bench_now_seconds() - started; + } + started = bench_now_seconds(); + output = context->forward(context->model, context->input); + if (context->train) { + context->phase_seconds[PHASE_FORWARD] += + bench_now_seconds() - started; + } if (output == NULL) goto cleanup; if (context->train) { + started = bench_now_seconds(); loss = nn_cross_entropy(output, context->targets); - if (loss == NULL || ag_backward(loss) != 0) goto cleanup; + context->phase_seconds[PHASE_LOSS] += + bench_now_seconds() - started; + if (loss == NULL) goto cleanup; + started = bench_now_seconds(); + if (ag_backward(loss) != 0) goto cleanup; + context->phase_seconds[PHASE_BACKWARD] += + bench_now_seconds() - started; + started = bench_now_seconds(); if (context->sgd != NULL) { if (nn_sgd_step(context->sgd) != 0) goto cleanup; - nn_sgd_zero_grad(context->sgd); } else { if (context->adamw == NULL || nn_adamw_step(context->adamw) != 0) { goto cleanup; } - nn_adamw_zero_grad(context->adamw); } + context->phase_seconds[PHASE_ADAMW] += + bench_now_seconds() - started; *checksum += loss->value->storage->data[loss->value->offset]; } else { *checksum += output->value->storage->data[output->value->offset]; @@ -68,11 +106,44 @@ static int nn_operation(void* opaque, double* checksum) status = 0; cleanup: + if (context->train) started = bench_now_seconds(); ag_tensor_release(loss); ag_tensor_release(output); + if (context->train) { + context->phase_seconds[PHASE_GRAPH_RELEASE] += + bench_now_seconds() - started; + ++context->phase_calls; + } return status; } +static void report_training_phases(const bench_options* options, + FILE* csv, + int requested_threads, + nn_bench_context* context) +{ + static const char* names[PHASE_COUNT] = { + "tiny_lm_zero_grad", "tiny_lm_forward_phase", "tiny_lm_loss", + "tiny_lm_backward", "tiny_lm_adamw", "tiny_lm_graph_release" + }; + if (context->phase_calls <= 0) return; + for (int phase = 0; phase < PHASE_COUNT; ++phase) { + bench_case benchmark = { + "nn_phase", names[phase], "[Bx128]->[Bx128x256]", + "isolated-phase", "ms/call", 0.0, 0, NULL, NULL + }; + bench_measurement result; + result.median_seconds = context->phase_seconds[phase] / + (double)context->phase_calls; + result.p95_seconds = result.median_seconds; + result.checksum = result.median_seconds; + result.iterations_per_sample = 1; + benchmark.context = context; + bench_record_measurement(options, csv, &benchmark, + requested_threads, &result); + } +} + static int run_nn_case(const bench_options* options, FILE* csv, const char* suite, @@ -319,6 +390,9 @@ static int run_decoder_case(const bench_options* options, "[Bx128]->[Bx128x256]", train ? "forward+loss+backward+adamw" : "forward;graph-build", "tokens/s", (double)(batch * time), threads, &context, result); + if (status == 0 && train && strcmp(suite, "nn") == 0) { + report_training_phases(options, csv, threads, &context); + } destroy_context(&context); return status; } diff --git a/benchmarks/bench_runner.c b/benchmarks/bench_runner.c index 5074885..d5bc9e0 100644 --- a/benchmarks/bench_runner.c +++ b/benchmarks/bench_runner.c @@ -45,30 +45,20 @@ static void write_skipped(FILE* csv, } #endif -int bench_execute_case(const bench_options* options, - FILE* csv, - const bench_case* benchmark, - int requested_threads, - bench_measurement* result) +int bench_record_measurement(const bench_options* options, + FILE* csv, + const bench_case* benchmark, + int requested_threads, + const bench_measurement* result) { double value; int threads; - if (options == NULL || benchmark == NULL || result == NULL) return 1; -#ifndef _OPENMP - if (requested_threads > 1) { - write_skipped(csv, options, benchmark, requested_threads); - return 2; - } -#endif + if (options == NULL || benchmark == NULL || result == NULL || + !isfinite(result->median_seconds) || + !isfinite(result->p95_seconds) || + !isfinite(result->checksum)) return 1; threads = actual_threads(requested_threads); - if (bench_measure(benchmark->operation, benchmark->context, - &options->profile, result) != 0 || - !isfinite(result->checksum)) { - fprintf(stderr, "Benchmark failed: %s/%s\n", - benchmark->suite, benchmark->name); - return 1; - } value = benchmark->units_per_call > 0.0 ? benchmark->units_per_call / result->median_seconds : result->median_seconds * 1000.0; @@ -87,6 +77,30 @@ int bench_execute_case(const bench_options* options, return 0; } +int bench_execute_case(const bench_options* options, + FILE* csv, + const bench_case* benchmark, + int requested_threads, + bench_measurement* result) +{ + if (options == NULL || benchmark == NULL || result == NULL) return 1; +#ifndef _OPENMP + if (requested_threads > 1) { + write_skipped(csv, options, benchmark, requested_threads); + return 2; + } +#endif + if (bench_measure(benchmark->operation, benchmark->context, + &options->profile, result) != 0 || + !isfinite(result->checksum)) { + fprintf(stderr, "Benchmark failed: %s/%s\n", + benchmark->suite, benchmark->name); + return 1; + } + return bench_record_measurement(options, csv, benchmark, + requested_threads, result); +} + static int suite_selected(const char* requested, const char* suite) { return strcmp(requested, "all") == 0 || strcmp(requested, suite) == 0; diff --git a/benchmarks/bench_runner.h b/benchmarks/bench_runner.h index 133b265..1010405 100644 --- a/benchmarks/bench_runner.h +++ b/benchmarks/bench_runner.h @@ -33,4 +33,11 @@ int bench_execute_case(const bench_options* options, int requested_threads, bench_measurement* result); +/* Writes a measurement already collected by a custom benchmark. */ +int bench_record_measurement(const bench_options* options, + FILE* csv, + const bench_case* benchmark, + int requested_threads, + const bench_measurement* result); + #endif From 1b09e37f3b08c5deed052d2f105b3d5d912afbb6 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:40:51 +0200 Subject: [PATCH 18/49] fix: size backward contribution workspace by graph arity --- src/autograd/autograd_backward.c | 32 +++++++++++++++++++++++++++++--- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index 2c8a5b5..5b4f6f8 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -144,6 +144,15 @@ static int merge_persistent_gradients(const tensor_list* tensors, tensor** pass_ return 0; } +static void free_contributions(tensor** contributions, int count) +{ + if (contributions == NULL) return; + for (int index = 0; index < count; ++index) { + t_free(contributions[index]); + contributions[index] = NULL; + } +} + int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { if (output == NULL || !output->requires_grad || !tensor_has_valid_metadata(output->value) || @@ -153,11 +162,24 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { tensor_list tensors = {0}; node_list nodes = {0}; tensor** pass_gradients = NULL; + tensor** contributions = NULL; + int contribution_capacity = 0; int status = 1; if (collect_graph(output, &tensors, &nodes) != 0) goto cleanup; if (!graph_versions_match(&nodes)) goto cleanup; + for (int node_index = 0; node_index < nodes.count; ++node_index) { + if (nodes.values[node_index]->input_count > contribution_capacity) { + contribution_capacity = nodes.values[node_index]->input_count; + } + } + if (contribution_capacity > 0) { + contributions = (tensor**)calloc((size_t)contribution_capacity, + sizeof(*contributions)); + if (contributions == NULL) goto cleanup; + } + pass_gradients = (tensor**)calloc((size_t)tensors.count, sizeof(*pass_gradients)); if (pass_gradients == NULL) goto cleanup; @@ -172,10 +194,12 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { int gradient_index = node->output->graph_index; if (gradient_index < 0 || pass_gradients[gradient_index] == NULL) goto cleanup; - tensor* contributions[2] = {NULL, NULL}; + for (int input_index = 0; input_index < node->input_count; ++input_index) { + contributions[input_index] = NULL; + } if (node->backward(node, pass_gradients[gradient_index], contributions) != 0) { - t_free(contributions[0]); t_free(contributions[1]); + free_contributions(contributions, node->input_count); goto cleanup; } @@ -191,7 +215,7 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { accumulate_pass_gradient(&pass_gradients[destination], contributions[input_index], input->value) != 0) { contributions[input_index] = NULL; - t_free(contributions[0]); t_free(contributions[1]); + free_contributions(contributions, node->input_count); goto cleanup; } contributions[input_index] = NULL; @@ -201,6 +225,8 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { status = merge_persistent_gradients(&tensors, pass_gradients); cleanup: + free_contributions(contributions, contribution_capacity); + free(contributions); if (pass_gradients != NULL) { for (int i = 0; i < tensors.count; ++i) t_free(pass_gradients[i]); free(pass_gradients); From 96c29dd3ee7b8fc728df735386add86b23231641 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:42:27 +0200 Subject: [PATCH 19/49] fix: configure benchmark threads before measurement --- benchmarks/bench_nn.c | 2 +- benchmarks/bench_runner.c | 10 +++++----- benchmarks/bench_runner.h | 1 + 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index f7eac32..4127ac3 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -140,7 +140,7 @@ static void report_training_phases(const bench_options* options, result.iterations_per_sample = 1; benchmark.context = context; bench_record_measurement(options, csv, &benchmark, - requested_threads, &result); + requested_threads, requested_threads, &result); } } diff --git a/benchmarks/bench_runner.c b/benchmarks/bench_runner.c index d5bc9e0..4c41439 100644 --- a/benchmarks/bench_runner.c +++ b/benchmarks/bench_runner.c @@ -49,28 +49,27 @@ int bench_record_measurement(const bench_options* options, FILE* csv, const bench_case* benchmark, int requested_threads, + int measured_threads, const bench_measurement* result) { double value; - int threads; if (options == NULL || benchmark == NULL || result == NULL || !isfinite(result->median_seconds) || !isfinite(result->p95_seconds) || !isfinite(result->checksum)) return 1; - threads = actual_threads(requested_threads); value = benchmark->units_per_call > 0.0 ? benchmark->units_per_call / result->median_seconds : result->median_seconds * 1000.0; printf(" %-24s threads=%-3d median=%9.3f ms p95=%9.3f ms " "%s=%9.3f\n", - benchmark->name, threads, result->median_seconds * 1000.0, + benchmark->name, measured_threads, result->median_seconds * 1000.0, result->p95_seconds * 1000.0, benchmark->metric, value); if (csv != NULL) { fprintf(csv, "%s,%s,%s,%s,%s,%d,%d,%.9g,%.9g,%s,%.9g,%d,%.9g,ok\n", benchmark->suite, benchmark->name, benchmark->shape, benchmark->layout, options->profile.profile, - requested_threads, threads, result->median_seconds, + requested_threads, measured_threads, result->median_seconds, result->p95_seconds, benchmark->metric, value, result->iterations_per_sample, result->checksum); } @@ -90,6 +89,7 @@ int bench_execute_case(const bench_options* options, return 2; } #endif + int measured_threads = actual_threads(requested_threads); if (bench_measure(benchmark->operation, benchmark->context, &options->profile, result) != 0 || !isfinite(result->checksum)) { @@ -98,7 +98,7 @@ int bench_execute_case(const bench_options* options, return 1; } return bench_record_measurement(options, csv, benchmark, - requested_threads, result); + requested_threads, measured_threads, result); } static int suite_selected(const char* requested, const char* suite) diff --git a/benchmarks/bench_runner.h b/benchmarks/bench_runner.h index 1010405..0023a86 100644 --- a/benchmarks/bench_runner.h +++ b/benchmarks/bench_runner.h @@ -38,6 +38,7 @@ int bench_record_measurement(const bench_options* options, FILE* csv, const bench_case* benchmark, int requested_threads, + int measured_threads, const bench_measurement* result); #endif From f2f2635ec7a4273665228a0e468586c3bf685507 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:43:35 +0200 Subject: [PATCH 20/49] bench: add measurement lifecycle reset --- benchmarks/bench_harness.c | 18 ++++++++++++++---- benchmarks/bench_harness.h | 6 ++++++ benchmarks/bench_runner.c | 4 ++-- benchmarks/bench_runner.h | 1 + tests/unit/benchmark/test_bench_harness.c | 11 +++++++++-- 5 files changed, 32 insertions(+), 8 deletions(-) diff --git a/benchmarks/bench_harness.c b/benchmarks/bench_harness.c index 9cdc6f0..ed240ac 100644 --- a/benchmarks/bench_harness.c +++ b/benchmarks/bench_harness.c @@ -68,10 +68,11 @@ static int run_iterations(bench_operation operation, return 0; } -int bench_measure(bench_operation operation, - void* context, - const bench_profile* profile, - bench_measurement* result) +int bench_measure_with_reset(bench_operation operation, + void* context, + const bench_profile* profile, + bench_measure_reset reset, + bench_measurement* result) { double* samples; double checksum = 0.0; @@ -96,6 +97,7 @@ int bench_measure(bench_operation operation, if (iterations > 1000000) iterations = 1000000; } + if (reset != NULL) reset(context); samples = (double*)malloc((size_t)profile->sample_count * sizeof(*samples)); if (samples == NULL) return 1; checksum = 0.0; @@ -116,3 +118,11 @@ int bench_measure(bench_operation operation, free(samples); return 0; } + +int bench_measure(bench_operation operation, + void* context, + const bench_profile* profile, + bench_measurement* result) +{ + return bench_measure_with_reset(operation, context, profile, NULL, result); +} diff --git a/benchmarks/bench_harness.h b/benchmarks/bench_harness.h index c80f4e7..61e952d 100644 --- a/benchmarks/bench_harness.h +++ b/benchmarks/bench_harness.h @@ -4,6 +4,7 @@ #include typedef int (*bench_operation)(void* context, double* checksum); +typedef void (*bench_measure_reset)(void* context); typedef struct { const char* profile; @@ -25,5 +26,10 @@ int bench_measure(bench_operation operation, void* context, const bench_profile* profile, bench_measurement* result); +int bench_measure_with_reset(bench_operation operation, + void* context, + const bench_profile* profile, + bench_measure_reset reset, + bench_measurement* result); #endif diff --git a/benchmarks/bench_runner.c b/benchmarks/bench_runner.c index 4c41439..7271f45 100644 --- a/benchmarks/bench_runner.c +++ b/benchmarks/bench_runner.c @@ -90,8 +90,8 @@ int bench_execute_case(const bench_options* options, } #endif int measured_threads = actual_threads(requested_threads); - if (bench_measure(benchmark->operation, benchmark->context, - &options->profile, result) != 0 || + if (bench_measure_with_reset(benchmark->operation, benchmark->context, + &options->profile, benchmark->reset, result) != 0 || !isfinite(result->checksum)) { fprintf(stderr, "Benchmark failed: %s/%s\n", benchmark->suite, benchmark->name); diff --git a/benchmarks/bench_runner.h b/benchmarks/bench_runner.h index 0023a86..10f14ca 100644 --- a/benchmarks/bench_runner.h +++ b/benchmarks/bench_runner.h @@ -24,6 +24,7 @@ typedef struct { int parallel; bench_operation operation; void* context; + bench_measure_reset reset; } bench_case; int bench_run_suites(const bench_options* options, FILE* csv); diff --git a/tests/unit/benchmark/test_bench_harness.c b/tests/unit/benchmark/test_bench_harness.c index 24db24b..12d06f7 100644 --- a/tests/unit/benchmark/test_bench_harness.c +++ b/tests/unit/benchmark/test_bench_harness.c @@ -12,6 +12,12 @@ static int counted_operation(void* opaque, double* checksum) return 0; } +static void reset_counter(void* opaque) +{ + int* calls = (int*)opaque; + *calls = 0; +} + int main(void) { bench_profile quick = bench_profile_named("quick"); @@ -27,11 +33,12 @@ int main(void) return 1; } if (bench_now_seconds() <= 0.0 || - bench_measure(counted_operation, &calls, &smoke, &result) != 0) { + bench_measure_with_reset(counted_operation, &calls, &smoke, + reset_counter, &result) != 0) { fprintf(stderr, "smoke measurement failed\n"); return 1; } - if (calls != 2 || result.iterations_per_sample != 1 || + if (calls != 1 || result.iterations_per_sample != 1 || result.checksum != 1.0 || !isfinite(result.median_seconds) || result.median_seconds < 0.0 || result.p95_seconds < 0.0) { fprintf(stderr, "smoke measurement produced invalid statistics\n"); From e47be90d0d6cec409c8811f3063c7ed78e64ef86 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:44:40 +0200 Subject: [PATCH 21/49] bench: record real training phase distributions --- benchmarks/bench_nn.c | 77 ++++++++++++++++++++++++++++++++----------- 1 file changed, 58 insertions(+), 19 deletions(-) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index 4127ac3..264a6b4 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -1,5 +1,6 @@ #include "bench_runner.h" +#include #include #include @@ -16,8 +17,9 @@ typedef struct { nn_adamw* adamw; nn_rng rng; int train; - double phase_seconds[6]; - long phase_calls; + double* phase_samples[6]; + int phase_sample_count; + int phase_sample_capacity; } nn_bench_context; enum { @@ -60,6 +62,7 @@ static int nn_operation(void* opaque, double* checksum) ag_tensor* loss = NULL; int status = 1; double started; + double phase_times[PHASE_COUNT] = {0}; if (context->train) { started = bench_now_seconds(); @@ -68,27 +71,23 @@ static int nn_operation(void* opaque, double* checksum) } else if (context->adamw != NULL) { nn_adamw_zero_grad(context->adamw); } - context->phase_seconds[PHASE_ZERO_GRAD] += - bench_now_seconds() - started; + phase_times[PHASE_ZERO_GRAD] = bench_now_seconds() - started; } started = bench_now_seconds(); output = context->forward(context->model, context->input); if (context->train) { - context->phase_seconds[PHASE_FORWARD] += - bench_now_seconds() - started; + phase_times[PHASE_FORWARD] = bench_now_seconds() - started; } if (output == NULL) goto cleanup; if (context->train) { started = bench_now_seconds(); loss = nn_cross_entropy(output, context->targets); - context->phase_seconds[PHASE_LOSS] += - bench_now_seconds() - started; + phase_times[PHASE_LOSS] = bench_now_seconds() - started; if (loss == NULL) goto cleanup; started = bench_now_seconds(); if (ag_backward(loss) != 0) goto cleanup; - context->phase_seconds[PHASE_BACKWARD] += - bench_now_seconds() - started; + phase_times[PHASE_BACKWARD] = bench_now_seconds() - started; started = bench_now_seconds(); if (context->sgd != NULL) { if (nn_sgd_step(context->sgd) != 0) goto cleanup; @@ -97,8 +96,7 @@ static int nn_operation(void* opaque, double* checksum) goto cleanup; } } - context->phase_seconds[PHASE_ADAMW] += - bench_now_seconds() - started; + phase_times[PHASE_ADAMW] = bench_now_seconds() - started; *checksum += loss->value->storage->data[loss->value->offset]; } else { *checksum += output->value->storage->data[output->value->offset]; @@ -110,13 +108,30 @@ static int nn_operation(void* opaque, double* checksum) ag_tensor_release(loss); ag_tensor_release(output); if (context->train) { - context->phase_seconds[PHASE_GRAPH_RELEASE] += - bench_now_seconds() - started; - ++context->phase_calls; + phase_times[PHASE_GRAPH_RELEASE] = bench_now_seconds() - started; + if (context->phase_sample_count < context->phase_sample_capacity) { + for (int phase = 0; phase < PHASE_COUNT; ++phase) { + context->phase_samples[phase][context->phase_sample_count] = + phase_times[phase]; + } + ++context->phase_sample_count; + } } return status; } +static void reset_phase_samples(void* opaque) +{ + ((nn_bench_context*)opaque)->phase_sample_count = 0; +} + +static int compare_double(const void* left, const void* right) +{ + double a = *(const double*)left; + double b = *(const double*)right; + return (a > b) - (a < b); +} + static void report_training_phases(const bench_options* options, FILE* csv, int requested_threads, @@ -126,16 +141,19 @@ static void report_training_phases(const bench_options* options, "tiny_lm_zero_grad", "tiny_lm_forward_phase", "tiny_lm_loss", "tiny_lm_backward", "tiny_lm_adamw", "tiny_lm_graph_release" }; - if (context->phase_calls <= 0) return; + if (context->phase_sample_count <= 0) return; for (int phase = 0; phase < PHASE_COUNT; ++phase) { bench_case benchmark = { "nn_phase", names[phase], "[Bx128]->[Bx128x256]", "isolated-phase", "ms/call", 0.0, 0, NULL, NULL }; bench_measurement result; - result.median_seconds = context->phase_seconds[phase] / - (double)context->phase_calls; - result.p95_seconds = result.median_seconds; + qsort(context->phase_samples[phase], + (size_t)context->phase_sample_count, sizeof(double), compare_double); + result.median_seconds = + context->phase_samples[phase][context->phase_sample_count / 2]; + result.p95_seconds = context->phase_samples[phase][ + (95 * context->phase_sample_count - 1) / 100]; result.checksum = result.median_seconds; result.iterations_per_sample = 1; benchmark.context = context; @@ -160,11 +178,15 @@ static int run_nn_case(const bench_options* options, suite, name, shape, layout, metric, items, 1, nn_operation, context }; + benchmark.reset = context->train ? reset_phase_samples : NULL; return bench_execute_case(options, csv, &benchmark, threads, result); } static void destroy_context(nn_bench_context* context) { + for (int phase = 0; phase < PHASE_COUNT; ++phase) { + free(context->phase_samples[phase]); + } nn_sgd_destroy(context->sgd); nn_adamw_destroy(context->adamw); t_free(context->targets); @@ -173,6 +195,19 @@ static void destroy_context(nn_bench_context* context) memset(context, 0, sizeof(*context)); } +static int allocate_phase_samples(nn_bench_context* context, int sample_count) +{ + int capacity = sample_count * 4; + if (capacity < 64) capacity = 64; + context->phase_sample_capacity = capacity; + for (int phase = 0; phase < PHASE_COUNT; ++phase) { + context->phase_samples[phase] = + (double*)calloc((size_t)capacity, sizeof(double)); + if (context->phase_samples[phase] == NULL) return 1; + } + return 0; +} + static ag_tensor* linear_forward(const void* model, const ag_tensor* input) { return nn_linear_forward((const nn_linear*)model, input); @@ -385,6 +420,10 @@ static int run_decoder_case(const bench_options* options, destroy_context(&context); return 1; } + if (train && allocate_phase_samples(&context, options->profile.sample_count) != 0) { + destroy_context(&context); + return 1; + } int status = run_nn_case(options, csv, suite, train ? "tiny_lm_train_step" : "tiny_lm_forward", "[Bx128]->[Bx128x256]", From 7c8c4b5ae76022a4350bb15ebaee4bf5ae2193c0 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:47:33 +0200 Subject: [PATCH 22/49] bench: measure tensor allocation pressure --- benchmarks/bench_nn.c | 47 +++++++++++++++- benchmarks/bench_runner.c | 23 ++++++++ benchmarks/bench_runner.h | 11 ++++ src/tensor/tensor_alloc.c | 79 +++++++++++++++++++++++++++ src/tensor/tensor_alloc_internal.h | 20 +++++++ tests/unit/tensor/test_tensor_alloc.c | 24 ++++++++ 6 files changed, 203 insertions(+), 1 deletion(-) create mode 100644 src/tensor/tensor_alloc_internal.h diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index 264a6b4..b702549 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -4,6 +4,7 @@ #include #include +#include "../src/tensor/tensor_alloc_internal.h" typedef ag_tensor* (*model_forward)(const void* model, const ag_tensor* input); @@ -20,6 +21,8 @@ typedef struct { double* phase_samples[6]; int phase_sample_count; int phase_sample_capacity; + tensor_alloc_stats allocation_stats; + size_t allocation_baseline_bytes; } nn_bench_context; enum { @@ -122,7 +125,35 @@ static int nn_operation(void* opaque, double* checksum) static void reset_phase_samples(void* opaque) { - ((nn_bench_context*)opaque)->phase_sample_count = 0; + nn_bench_context* context = (nn_bench_context*)opaque; + context->phase_sample_count = 0; + tensor_alloc_stats_reset_counters(); +} + +static void report_allocation_stats(const bench_options* options, + FILE* csv, + int requested_threads, + int measured_threads, + const nn_bench_context* context) +{ + double calls = context->phase_sample_count > 0 + ? (double)context->phase_sample_count : 1.0; + const tensor_alloc_stats* stats = &context->allocation_stats; + bench_record_scalar(options, csv, "nn_phase", "tiny_lm_allocations", + "[Bx128]->[Bx128x256]", "isolated-phase", + "alloc/call", requested_threads, measured_threads, + (double)stats->allocations / calls); + bench_record_scalar(options, csv, "nn_phase", "tiny_lm_allocated_bytes", + "[Bx128]->[Bx128x256]", "isolated-phase", + "bytes/call", requested_threads, measured_threads, + (double)stats->allocated_bytes / calls); + bench_record_scalar(options, csv, "nn_phase", "tiny_lm_peak_live_bytes", + "[Bx128]->[Bx128x256]", "isolated-phase", + "bytes", requested_threads, measured_threads, + (double)(stats->peak_live_bytes >= + context->allocation_baseline_bytes + ? stats->peak_live_bytes - + context->allocation_baseline_bytes : 0)); } static int compare_double(const void* left, const void* right) @@ -416,7 +447,12 @@ static int run_decoder_case(const bench_options* options, nn_bench_context context; bench_measurement local_result; bench_measurement* result = measurement == NULL ? &local_result : measurement; + if (train) { + tensor_alloc_stats_enable(1); + tensor_alloc_stats_reset(); + } if (setup_decoder(&context, batch, time, channels, layers, train) != 0) { + if (train) tensor_alloc_stats_enable(0); destroy_context(&context); return 1; } @@ -424,14 +460,23 @@ static int run_decoder_case(const bench_options* options, destroy_context(&context); return 1; } + if (train) { + tensor_alloc_stats_read(&context.allocation_stats); + context.allocation_baseline_bytes = + context.allocation_stats.live_bytes; + tensor_alloc_stats_reset_counters(); + } int status = run_nn_case(options, csv, suite, train ? "tiny_lm_train_step" : "tiny_lm_forward", "[Bx128]->[Bx128x256]", train ? "forward+loss+backward+adamw" : "forward;graph-build", "tokens/s", (double)(batch * time), threads, &context, result); if (status == 0 && train && strcmp(suite, "nn") == 0) { + tensor_alloc_stats_read(&context.allocation_stats); report_training_phases(options, csv, threads, &context); + report_allocation_stats(options, csv, threads, threads, &context); } + if (train) tensor_alloc_stats_enable(0); destroy_context(&context); return status; } diff --git a/benchmarks/bench_runner.c b/benchmarks/bench_runner.c index 7271f45..57d0fe7 100644 --- a/benchmarks/bench_runner.c +++ b/benchmarks/bench_runner.c @@ -76,6 +76,29 @@ int bench_record_measurement(const bench_options* options, return 0; } +int bench_record_scalar(const bench_options* options, + FILE* csv, + const char* suite, + const char* name, + const char* shape, + const char* layout, + const char* metric, + int requested_threads, + int measured_threads, + double value) +{ + if (options == NULL || suite == NULL || name == NULL || + !isfinite(value)) return 1; + printf(" %-24s threads=%-3d %s=%9.3f\n", + name, measured_threads, metric, value); + if (csv != NULL) { + fprintf(csv, "%s,%s,%s,%s,%s,%d,%d,0,0,%s,%.9g,1,%.9g,ok\n", + suite, name, shape, layout, options->profile.profile, + requested_threads, measured_threads, metric, value, value); + } + return 0; +} + int bench_execute_case(const bench_options* options, FILE* csv, const bench_case* benchmark, diff --git a/benchmarks/bench_runner.h b/benchmarks/bench_runner.h index 10f14ca..902b0a2 100644 --- a/benchmarks/bench_runner.h +++ b/benchmarks/bench_runner.h @@ -42,4 +42,15 @@ int bench_record_measurement(const bench_options* options, int measured_threads, const bench_measurement* result); +int bench_record_scalar(const bench_options* options, + FILE* csv, + const char* suite, + const char* name, + const char* shape, + const char* layout, + const char* metric, + int requested_threads, + int measured_threads, + double value); + #endif diff --git a/src/tensor/tensor_alloc.c b/src/tensor/tensor_alloc.c index 914a96b..72bb918 100644 --- a/src/tensor/tensor_alloc.c +++ b/src/tensor/tensor_alloc.c @@ -1,6 +1,82 @@ #include #include #include "../../include/tensorlib/tensor.h" +#include "tensor_alloc_internal.h" + +static int stats_enabled; +static volatile unsigned long long stats_allocations; +static volatile unsigned long long stats_frees; +static volatile size_t stats_allocated_bytes; +static volatile size_t stats_live_bytes; +static volatile size_t stats_peak_live_bytes; + +static void record_storage_alloc(size_t bytes) +{ + if (!stats_enabled) return; +#ifdef _OPENMP +#pragma omp atomic update +#endif + ++stats_allocations; +#ifdef _OPENMP +#pragma omp atomic update +#endif + stats_allocated_bytes += bytes; +#ifdef _OPENMP +#pragma omp atomic update +#endif + stats_live_bytes += bytes; +#ifdef _OPENMP +#pragma omp critical(tensor_alloc_stats_peak) +#endif + if (stats_live_bytes > stats_peak_live_bytes) { + stats_peak_live_bytes = stats_live_bytes; + } +} + +static void record_storage_free(size_t bytes) +{ + if (!stats_enabled) return; +#ifdef _OPENMP +#pragma omp atomic update +#endif + ++stats_frees; +#ifdef _OPENMP +#pragma omp atomic update +#endif + stats_live_bytes -= bytes; +} + +void tensor_alloc_stats_enable(int enabled) +{ + stats_enabled = enabled != 0; +} + +void tensor_alloc_stats_reset(void) +{ + stats_allocations = 0; + stats_frees = 0; + stats_allocated_bytes = 0; + stats_live_bytes = 0; + stats_peak_live_bytes = 0; +} + +void tensor_alloc_stats_reset_counters(void) +{ + stats_allocations = 0; + stats_frees = 0; + stats_allocated_bytes = 0; + stats_peak_live_bytes = stats_live_bytes; +} + +void tensor_alloc_stats_read(tensor_alloc_stats* result) +{ + if (result == NULL) return; + result->allocations = stats_allocations; + result->frees = stats_frees; + result->allocated_bytes = stats_allocated_bytes; + result->live_bytes = stats_live_bytes; + result->peak_live_bytes = stats_peak_live_bytes; +} void add_ref_count(Storage* a, tensor* b) { if (a != NULL && b != NULL) { @@ -23,6 +99,7 @@ Storage* s_alloc(int ndim, const int* dims) { free(s); return NULL; } + record_storage_alloc(count * sizeof(float)); return s; } @@ -64,6 +141,7 @@ void t_free(tensor* t) { if (t->storage->ref_count > 1) { t->storage->ref_count--; } else { + record_storage_free((size_t)t->storage->size * sizeof(float)); free(t->storage->data); free(t->storage); } @@ -95,6 +173,7 @@ int init_t(tensor* c, tensor* ref) { c->storage = NULL; return 1; } + record_storage_alloc((size_t)total_elements * sizeof(float)); if (ref->ndim > 0) { int* strides = (int*)malloc((size_t)ref->ndim * sizeof(int)); diff --git a/src/tensor/tensor_alloc_internal.h b/src/tensor/tensor_alloc_internal.h new file mode 100644 index 0000000..a9cc0bb --- /dev/null +++ b/src/tensor/tensor_alloc_internal.h @@ -0,0 +1,20 @@ +#ifndef TENSORLIB_TENSOR_ALLOC_INTERNAL_H +#define TENSORLIB_TENSOR_ALLOC_INTERNAL_H + +#include +#include + +typedef struct { + uint64_t allocations; + uint64_t frees; + size_t allocated_bytes; + size_t live_bytes; + size_t peak_live_bytes; +} tensor_alloc_stats; + +void tensor_alloc_stats_enable(int enabled); +void tensor_alloc_stats_reset(void); +void tensor_alloc_stats_reset_counters(void); +void tensor_alloc_stats_read(tensor_alloc_stats* result); + +#endif diff --git a/tests/unit/tensor/test_tensor_alloc.c b/tests/unit/tensor/test_tensor_alloc.c index 10be1db..3101658 100644 --- a/tests/unit/tensor/test_tensor_alloc.c +++ b/tests/unit/tensor/test_tensor_alloc.c @@ -1,6 +1,7 @@ #include #include "../../fixtures/test_common.h" #include "../../../include/tensorlib/tensor.h" +#include "../../../src/tensor/tensor_alloc_internal.h" TEST(test_s_alloc_basic) { int dims[2] = {2, 3}; @@ -44,6 +45,28 @@ TEST(test_t_alloc_basic_shape_and_strides) { t_free(t); } +TEST(test_allocation_stats_track_storage_and_views) { + int dims[2] = {2, 3}; + tensor_alloc_stats stats; + tensor* base; + tensor* view; + + tensor_alloc_stats_enable(1); + tensor_alloc_stats_reset(); + base = t_alloc(2, dims); + view = t_transpose(base, 0, 1); + tensor_alloc_stats_read(&stats); + ASSERT_EQ_INT(stats.allocations, 1); + ASSERT_EQ_INT(stats.live_bytes, 6 * sizeof(float)); + ASSERT_EQ_INT(stats.peak_live_bytes, 6 * sizeof(float)); + t_free(view); + t_free(base); + tensor_alloc_stats_read(&stats); + ASSERT_EQ_INT(stats.frees, 1); + ASSERT_EQ_INT(stats.live_bytes, 0); + tensor_alloc_stats_enable(0); +} + TEST(test_t_alloc_scalar_ndim_zero) { tensor* t = t_alloc(0, NULL); ASSERT_NOT_NULL(t); @@ -239,6 +262,7 @@ int main(void) { RUN_TEST(test_s_alloc_rejects_null_dims_with_positive_ndim); RUN_TEST(test_alloc_rejects_nonpositive_and_overflowing_dimensions); RUN_TEST(test_t_alloc_basic_shape_and_strides); + RUN_TEST(test_allocation_stats_track_storage_and_views); RUN_TEST(test_t_alloc_scalar_ndim_zero); RUN_TEST(test_t_alloc_rejects_negative_ndim); RUN_TEST(test_t_alloc_rejects_null_dims_with_positive_ndim); From 7d53d380a014ff04b1f11b7cacdd2b148eeede4a Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:51:13 +0200 Subject: [PATCH 23/49] bench: add backward component baselines --- benchmarks/bench_nn.c | 42 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index b702549..8c0441f 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -18,6 +18,7 @@ typedef struct { nn_adamw* adamw; nn_rng rng; int train; + int component_backward; double* phase_samples[6]; int phase_sample_count; int phase_sample_capacity; @@ -76,6 +77,7 @@ static int nn_operation(void* opaque, double* checksum) } phase_times[PHASE_ZERO_GRAD] = bench_now_seconds() - started; } + if (context->component_backward) nn_module_zero_grad(context->module); started = bench_now_seconds(); output = context->forward(context->model, context->input); @@ -101,6 +103,18 @@ static int nn_operation(void* opaque, double* checksum) } phase_times[PHASE_ADAMW] = bench_now_seconds() - started; *checksum += loss->value->storage->data[loss->value->offset]; + } else if (context->component_backward) { + tensor* seed = t_alloc(output->value->ndim, output->value->dims); + if (seed == NULL) goto cleanup; + for (int index = 0; index < tensor_numel(seed); ++index) { + seed->storage->data[index] = 1.0f; + } + if (ag_backward_with_grad(output, seed) != 0) { + t_free(seed); + goto cleanup; + } + t_free(seed); + *checksum += output->value->storage->data[output->value->offset]; } else { *checksum += output->value->storage->data[output->value->offset]; } @@ -110,6 +124,7 @@ static int nn_operation(void* opaque, double* checksum) if (context->train) started = bench_now_seconds(); ag_tensor_release(loss); ag_tensor_release(output); + if (context->component_backward) nn_module_zero_grad(context->module); if (context->train) { phase_times[PHASE_GRAPH_RELEASE] = bench_now_seconds() - started; if (context->phase_sample_count < context->phase_sample_capacity) { @@ -280,6 +295,7 @@ typedef enum { static int run_component(const bench_options* options, FILE* csv, component_kind kind, + int backward, int batch, int time, int channels) @@ -290,6 +306,7 @@ static int run_component(const bench_options* options, bench_measurement result; const char* name = NULL; const char* shape = "[BxTxC]"; + char backward_name[64]; memset(&context, 0, sizeof(context)); nn_rng_seed(&context.rng, UINT64_C(0xB34C4)); @@ -332,9 +349,16 @@ static int run_component(const bench_options* options, destroy_context(&context); return 1; } + context.component_backward = backward; + if (backward) context.input->requires_grad = 1; nn_module_set_training(context.module, 0); + if (backward) { + snprintf(backward_name, sizeof(backward_name), "%s_backward", name); + name = backward_name; + } int status = run_nn_case(options, csv, "nn", name, shape, - "forward;graph-build", "tokens/s", + backward ? "forward+backward" : "forward;graph-build", + "tokens/s", (double)(batch * time), 1, &context, &result); destroy_context(&context); @@ -538,13 +562,21 @@ int bench_run_nn_suite(const bench_options* options, FILE* csv) int status = 0; printf("Neural-network suite (eager float32)\n"); - status |= run_component(options, csv, COMPONENT_LINEAR, + status |= run_component(options, csv, COMPONENT_LINEAR, 0, + batch, time, channels); + status |= run_component(options, csv, COMPONENT_LINEAR, 1, + batch, time, channels); + status |= run_component(options, csv, COMPONENT_LAYER_NORM, 0, + batch, time, channels); + status |= run_component(options, csv, COMPONENT_LAYER_NORM, 1, + batch, time, channels); + status |= run_component(options, csv, COMPONENT_ATTENTION, 0, batch, time, channels); - status |= run_component(options, csv, COMPONENT_LAYER_NORM, + status |= run_component(options, csv, COMPONENT_ATTENTION, 1, batch, time, channels); - status |= run_component(options, csv, COMPONENT_ATTENTION, + status |= run_component(options, csv, COMPONENT_BLOCK, 0, batch, time, channels); - status |= run_component(options, csv, COMPONENT_BLOCK, + status |= run_component(options, csv, COMPONENT_BLOCK, 1, batch, time, channels); status |= run_mlp(options, csv, smoke ? 2 : 64, 0); status |= run_mlp(options, csv, smoke ? 2 : 64, 1); From 52c9494cef85a93afb5e6452b55bfd8b7d49759d Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:54:18 +0200 Subject: [PATCH 24/49] perf: transfer first persistent gradients without copying --- src/autograd/autograd_backward.c | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index 5b4f6f8..218b565 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -124,9 +124,15 @@ static int merge_persistent_gradients(const tensor_list* tensors, tensor** pass_ for (int i = 0; i < tensors->count; ++i) { ag_tensor* value = tensors->values[i]; if (!value->requires_grad || pass_gradients[i] == NULL) continue; - merged[i] = value->grad == NULL - ? t_clone(pass_gradients[i]) - : t_add(value->grad, pass_gradients[i]); + if (value->grad == NULL && is_contiguous(pass_gradients[i]) && + pass_gradients[i]->offset == 0) { + merged[i] = pass_gradients[i]; + pass_gradients[i] = NULL; + } else { + merged[i] = value->grad == NULL + ? t_clone(pass_gradients[i]) + : t_add(value->grad, pass_gradients[i]); + } if (merged[i] == NULL) { for (int j = 0; j < tensors->count; ++j) t_free(merged[j]); free(merged); From 4679eaf1c87ba930f97fb789a6e1967294dc110d Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:55:05 +0200 Subject: [PATCH 25/49] perf: accumulate compatible pass gradients in place --- src/autograd/autograd_backward.c | 34 ++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index 218b565..cefe19f 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -100,6 +100,36 @@ static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { return current; } +static int add_in_place(tensor* destination, const tensor* source) +{ + if (!tensor_has_valid_metadata(destination) || + !tensor_has_valid_metadata(source) || + !same_shape(destination, (tensor*)source) || + destination->storage == source->storage) return 1; + + int count = tensor_numel(destination); + if (is_contiguous(destination) && is_contiguous((tensor*)source) && + destination->offset == 0 && source->offset == 0) { + for (int index = 0; index < count; ++index) { + destination->storage->data[index] += source->storage->data[index]; + } + return 0; + } + + int* coords = destination->ndim > 0 + ? (int*)calloc((size_t)destination->ndim, sizeof(*coords)) : NULL; + if (destination->ndim > 0 && coords == NULL) return 1; + for (int index = 0; index < count; ++index) { + int destination_index = get_flat_index_nd(destination, coords); + int source_index = get_flat_index_nd((tensor*)source, coords); + destination->storage->data[destination_index] += + source->storage->data[source_index]; + advance_coords(coords, destination->dims, destination->ndim); + } + free(coords); + return 0; +} + static int accumulate_pass_gradient(tensor** destination, tensor* contribution, const tensor* target) { @@ -109,6 +139,10 @@ static int accumulate_pass_gradient(tensor** destination, *destination = reduced; return 0; } + if (add_in_place(*destination, reduced) == 0) { + t_free(reduced); + return 0; + } tensor* sum = t_add(*destination, reduced); t_free(reduced); if (sum == NULL) return 1; From df9a54ae67cf158cc761e1152e946e3955cea46e Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:56:24 +0200 Subject: [PATCH 26/49] perf: fuse broadcast gradient reduction --- src/autograd/autograd_backward.c | 52 ++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 16 deletions(-) diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index cefe19f..5b0a2ce 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -78,26 +78,46 @@ static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { return NULL; } - tensor* current = contribution; - while (current->ndim > target->ndim) { - tensor* reduced = t_sum(current, 0); - t_free(current); - if (reduced == NULL) return NULL; - current = reduced; + for (int target_axis = 0; target_axis < target->ndim; ++target_axis) { + int contribution_axis = contribution->ndim - target->ndim + target_axis; + if (target->dims[target_axis] != 1 && + target->dims[target_axis] != contribution->dims[contribution_axis]) { + t_free(contribution); + return NULL; + } + } + if (contribution->ndim == target->ndim && + same_shape(contribution, (tensor*)target)) return contribution; + + tensor* reduced = t_alloc(target->ndim, target->dims); + int* coords = contribution->ndim > 0 + ? (int*)calloc((size_t)contribution->ndim, sizeof(*coords)) : NULL; + if (reduced == NULL || (contribution->ndim > 0 && coords == NULL)) { + free(coords); + t_free(reduced); + t_free(contribution); + return NULL; + } + for (int index = 0; index < tensor_numel(reduced); ++index) { + reduced->storage->data[reduced->offset + index] = 0.0f; } - for (int axis = 0; axis < target->ndim; ++axis) { - if (current->dims[axis] == target->dims[axis]) continue; - if (target->dims[axis] != 1 || current->dims[axis] == 1) { - t_free(current); - return NULL; + int contribution_count = tensor_numel(contribution); + for (int index = 0; index < contribution_count; ++index) { + int reduced_index = reduced->offset; + for (int target_axis = 0; target_axis < target->ndim; ++target_axis) { + int contribution_axis = contribution->ndim - target->ndim + target_axis; + int coordinate = target->dims[target_axis] == 1 + ? 0 : coords[contribution_axis]; + reduced_index += coordinate * reduced->strides[target_axis]; } - tensor* reduced = t_sum_keepdim(current, axis); - t_free(current); - if (reduced == NULL) return NULL; - current = reduced; + reduced->storage->data[reduced_index] += + contribution->storage->data[get_flat_index_nd(contribution, coords)]; + advance_coords(coords, contribution->dims, contribution->ndim); } - return current; + free(coords); + t_free(contribution); + return reduced; } static int add_in_place(tensor* destination, const tensor* source) From 5364ec93b765684630e0a50c86a9b8ec9e8ab1bf Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:57:26 +0200 Subject: [PATCH 27/49] perf: parallelize softmax and cross-entropy backward --- src/autograd/autograd_cross_entropy.c | 9 +++++++++ src/autograd/autograd_softmax.c | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/src/autograd/autograd_cross_entropy.c b/src/autograd/autograd_cross_entropy.c index 09154ab..6a2f6cd 100644 --- a/src/autograd/autograd_cross_entropy.c +++ b/src/autograd/autograd_cross_entropy.c @@ -2,6 +2,9 @@ #include #include "../../include/tensorlib/autograd_internal.h" +#include "../tensor/parallel.h" + +#define TENSORLIB_CROSS_ENTROPY_MIN_PARALLEL_ELEMENTS (1 << 16) typedef struct { int rows; @@ -60,6 +63,12 @@ static int backward_cross_entropy(const ag_node* node, if (gradient == NULL) return 1; upstream = output_gradient->storage->data[output_gradient->offset] / (float)context->rows; + int threads = tensorlib_parallel_threads( + (long long)context->rows * context->classes, + TENSORLIB_CROSS_ENTROPY_MIN_PARALLEL_ELEMENTS, context->rows); +#ifdef _OPENMP +#pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) +#endif for (int row = 0; row < context->rows; ++row) { int base = row_base(input, row, context->classes); float maximum = -INFINITY; diff --git a/src/autograd/autograd_softmax.c b/src/autograd/autograd_softmax.c index e979b10..5c85d1d 100644 --- a/src/autograd/autograd_softmax.c +++ b/src/autograd/autograd_softmax.c @@ -2,6 +2,9 @@ #include #include "../../include/tensorlib/autograd_internal.h" +#include "../tensor/parallel.h" + +#define TENSORLIB_SOFTMAX_MIN_PARALLEL_ELEMENTS (1 << 16) typedef struct { int rows; @@ -44,6 +47,12 @@ static int backward_softmax(const ag_node* node, gradient = t_alloc(output->ndim, output->dims); if (gradient == NULL) return 1; + int threads = tensorlib_parallel_threads( + (long long)context->rows * context->width, + TENSORLIB_SOFTMAX_MIN_PARALLEL_ELEMENTS, context->rows); +#ifdef _OPENMP +#pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) +#endif for (int row = 0; row < context->rows; ++row) { int width = allowed_width(context, row); int output_base = row_base(output, row, context->width); From d8883992faea1716e541abc50d99240862c24924 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:58:35 +0200 Subject: [PATCH 28/49] perf: parallelize deterministic layer norm backward --- src/autograd/autograd_layer_norm.c | 60 ++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 4 deletions(-) diff --git a/src/autograd/autograd_layer_norm.c b/src/autograd/autograd_layer_norm.c index a0a91ba..e524656 100644 --- a/src/autograd/autograd_layer_norm.c +++ b/src/autograd/autograd_layer_norm.c @@ -2,6 +2,9 @@ #include #include "../../include/tensorlib/autograd_internal.h" +#include "../tensor/parallel.h" + +#define TENSORLIB_LAYER_NORM_MIN_PARALLEL_ELEMENTS (1 << 16) typedef struct { int rows; @@ -44,6 +47,8 @@ static int backward_layer_norm(const ag_node* node, tensor* input_gradient = NULL; tensor* weight_gradient = NULL; tensor* bias_gradient = NULL; + float* weight_partials = NULL; + float* bias_partials = NULL; if (node == NULL || input_gradients == NULL || !tensor_has_valid_metadata(output_gradient)) return 1; @@ -72,7 +77,36 @@ static int backward_layer_norm(const ag_node* node, for (int k = 0; k < width; ++k) bias_gradient->storage->data[k] = 0.0f; } + int threads = tensorlib_parallel_threads( + (long long)context->rows * width, + TENSORLIB_LAYER_NORM_MIN_PARALLEL_ELEMENTS, context->rows); + if (threads > 1 && weight_gradient != NULL) { + weight_partials = (float*)calloc((size_t)threads * (size_t)width, + sizeof(*weight_partials)); + if (weight_partials == NULL) goto fail; + } + if (threads > 1 && bias_gradient != NULL) { + bias_partials = (float*)calloc((size_t)threads * (size_t)width, + sizeof(*bias_partials)); + if (bias_partials == NULL) goto fail; + } + +#ifdef _OPENMP +#pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) +#endif for (int row = 0; row < context->rows; ++row) { + int tid = 0; +#ifdef _OPENMP + tid = omp_get_thread_num(); +#endif + float* local_weight = weight_partials == NULL + ? (weight_gradient == NULL ? NULL : + weight_gradient->storage->data) + : weight_partials + (size_t)tid * (size_t)width; + float* local_bias = bias_partials == NULL + ? (bias_gradient == NULL ? NULL : + bias_gradient->storage->data) + : bias_partials + (size_t)tid * (size_t)width; int input_base = row_base(input, row, width); int gradient_base = row_base(output_gradient, row, width); float sum = 0.0f; @@ -86,10 +120,8 @@ static int backward_layer_norm(const ag_node* node, float scale = weight == NULL ? 1.0f : weight->storage->data[k]; sum += upstream * scale; weighted_sum += upstream * scale * normalized; - if (weight_gradient != NULL) { - weight_gradient->storage->data[k] += upstream * normalized; - } - if (bias_gradient != NULL) bias_gradient->storage->data[k] += upstream; + if (local_weight != NULL) local_weight[k] += upstream * normalized; + if (local_bias != NULL) local_bias[k] += upstream; } if (input_gradient != NULL) { for (int k = 0; k < width; ++k) { @@ -106,6 +138,24 @@ static int backward_layer_norm(const ag_node* node, } } } + if (weight_partials != NULL) { + for (int tid = 0; tid < threads; ++tid) { + for (int k = 0; k < width; ++k) { + weight_gradient->storage->data[k] += + weight_partials[(size_t)tid * (size_t)width + (size_t)k]; + } + } + } + if (bias_partials != NULL) { + for (int tid = 0; tid < threads; ++tid) { + for (int k = 0; k < width; ++k) { + bias_gradient->storage->data[k] += + bias_partials[(size_t)tid * (size_t)width + (size_t)k]; + } + } + } + free(weight_partials); + free(bias_partials); input_gradients[0] = input_gradient; if (node->input_count == 3) { input_gradients[1] = weight_gradient; @@ -114,6 +164,8 @@ static int backward_layer_norm(const ag_node* node, return 0; fail: + free(weight_partials); + free(bias_partials); t_free(input_gradient); t_free(weight_gradient); t_free(bias_gradient); From ab1de0bd3475490d8c7d82741a6dd386571379f0 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 15:59:25 +0200 Subject: [PATCH 29/49] perf: fuse and parallelize gelu backward --- src/autograd/autograd_ops.c | 40 +++++++++++++++++++++++++++++++++++-- 1 file changed, 38 insertions(+), 2 deletions(-) diff --git a/src/autograd/autograd_ops.c b/src/autograd/autograd_ops.c index 63a265d..484d8e3 100644 --- a/src/autograd/autograd_ops.c +++ b/src/autograd/autograd_ops.c @@ -2,6 +2,9 @@ #include #include "./../../include/tensorlib/autograd_internal.h" +#include "../tensor/parallel.h" + +#define TENSORLIB_GELU_MIN_PARALLEL_ELEMENTS (1 << 16) static void free_gradients(tensor** gradients, int count) { for (int i = 0; i < count; ++i) { @@ -406,8 +409,41 @@ static int backward_gelu(const ag_node* node, if (node == NULL || node->input_count != 1 || input_gradients == NULL || !tensor_has_valid_metadata(output_gradient)) return 1; if (!node->inputs[0]->requires_grad) return 0; - input_gradients[0] = apply_unary_derivative(node, output_gradient, - derivative_gelu); + tensor* input = node->inputs[0]->value; + tensor* result = t_alloc(input->ndim, input->dims); + if (result == NULL) return 1; + int count = tensor_numel(input); + int threads = tensorlib_parallel_threads( + count, TENSORLIB_GELU_MIN_PARALLEL_ELEMENTS, 0); + if (is_contiguous(input) && is_contiguous((tensor*)output_gradient)) { + const float* input_data = input->storage->data + input->offset; + const float* upstream_data = output_gradient->storage->data + + output_gradient->offset; +#ifdef _OPENMP +#pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) +#endif + for (int index = 0; index < count; ++index) { + result->storage->data[index] = upstream_data[index] * + derivative_gelu(input_data[index], 0.0f); + } + } else { + int* coords = input->ndim > 0 + ? (int*)calloc((size_t)input->ndim, sizeof(*coords)) : NULL; + if (input->ndim > 0 && coords == NULL) { + t_free(result); + return 1; + } + for (int index = 0; index < count; ++index) { + int input_index = get_flat_index_nd(input, coords); + int upstream_index = get_flat_index_nd((tensor*)output_gradient, coords); + result->storage->data[index] = + output_gradient->storage->data[upstream_index] * + derivative_gelu(input->storage->data[input_index], 0.0f); + advance_coords(coords, input->dims, input->ndim); + } + free(coords); + } + input_gradients[0] = result; return input_gradients[0] == NULL; } From 8a8eb405ee624dd569220743dbb32ed9ec2e0f0b Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:01:00 +0200 Subject: [PATCH 30/49] perf: add retained transpose-aware matmul packs --- include/tensorlib/tensor.h | 2 ++ src/tensor/tensor_matmul.c | 24 ++++++++++++++++++++++++ tests/unit/tensor/test_tensor_matmul.c | 26 ++++++++++++++++++++++++++ 3 files changed, 52 insertions(+) diff --git a/include/tensorlib/tensor.h b/include/tensorlib/tensor.h index 8934f73..13d1b5a 100644 --- a/include/tensorlib/tensor.h +++ b/include/tensorlib/tensor.h @@ -89,7 +89,9 @@ tensor* t_mean_keepdim(tensor* a, int dim); tensor* t_max_keepdim(tensor* a, int dim); tensor* t_matmul(tensor* a, tensor* b); tensor_matmul_packed_rhs* t_pack_matmul_rhs(const tensor* rhs); +tensor_matmul_packed_rhs* t_pack_matmul_rhs_transposed(const tensor* rhs); tensor* t_matmul_packed_rhs(const tensor* lhs, const tensor_matmul_packed_rhs* rhs); +void t_retain_matmul_packed_rhs(tensor_matmul_packed_rhs* rhs); void t_free_matmul_packed_rhs(tensor_matmul_packed_rhs* rhs); #endif //TENSORLIB_TENSOR_H diff --git a/src/tensor/tensor_matmul.c b/src/tensor/tensor_matmul.c index 3b6759c..1c39acd 100644 --- a/src/tensor/tensor_matmul.c +++ b/src/tensor/tensor_matmul.c @@ -38,6 +38,7 @@ enum { #endif struct tensor_matmul_packed_rhs { + int ref_count; int batch_rank; int* batch_dims; int inner; @@ -128,6 +129,7 @@ tensor_matmul_packed_rhs* t_pack_matmul_rhs(const tensor* rhs) { tensor_matmul_packed_rhs* packed = (tensor_matmul_packed_rhs*)calloc(1, sizeof(*packed)); if (packed == NULL) return NULL; + packed->ref_count = 1; packed->batch_rank = batch_rank; packed->inner = rhs->dims[rhs->ndim - 2]; @@ -191,8 +193,30 @@ tensor_matmul_packed_rhs* t_pack_matmul_rhs(const tensor* rhs) { return packed; } +tensor_matmul_packed_rhs* t_pack_matmul_rhs_transposed(const tensor* rhs) +{ + tensor* transposed; + tensor_matmul_packed_rhs* packed; + + if (!tensor_has_valid_metadata(rhs) || rhs->ndim < 2) return NULL; + transposed = t_transpose((tensor*)rhs, rhs->ndim - 2, rhs->ndim - 1); + if (transposed == NULL) return NULL; + packed = t_pack_matmul_rhs(transposed); + t_free(transposed); + return packed; +} + +void t_retain_matmul_packed_rhs(tensor_matmul_packed_rhs* rhs) +{ + if (rhs != NULL && rhs->ref_count > 0) ++rhs->ref_count; +} + void t_free_matmul_packed_rhs(tensor_matmul_packed_rhs* rhs) { if (rhs == NULL) return; + if (rhs->ref_count > 1) { + --rhs->ref_count; + return; + } matmul_aligned_free(rhs->data); free(rhs->batch_dims); free(rhs); diff --git a/tests/unit/tensor/test_tensor_matmul.c b/tests/unit/tensor/test_tensor_matmul.c index 89c8491..0995304 100644 --- a/tests/unit/tensor/test_tensor_matmul.c +++ b/tests/unit/tensor/test_tensor_matmul.c @@ -242,6 +242,31 @@ TEST(test_packed_rhs_snapshots_transposed_view) { t_free(a); } +TEST(test_packed_rhs_transpose_pack_retains_lifetime) { + int left_dims[2] = {1, 3}; + int rhs_dims[2] = {2, 3}; + tensor* left = t_alloc(2, left_dims); + tensor* rhs = t_alloc(2, rhs_dims); + const float left_values[3] = {1, 2, 3}; + const float rhs_values[6] = {1, 2, 3, 4, 5, 6}; + fill_tensor(left, left_values, 3); + fill_tensor(rhs, rhs_values, 6); + + tensor_matmul_packed_rhs* packed = t_pack_matmul_rhs_transposed(rhs); + ASSERT_NOT_NULL(packed); + t_retain_matmul_packed_rhs(packed); + t_free_matmul_packed_rhs(packed); + tensor* result = t_matmul_packed_rhs(left, packed); + ASSERT_NOT_NULL(result); + ASSERT_EQ_FLOAT(result->storage->data[0], 14.0f); + ASSERT_EQ_FLOAT(result->storage->data[1], 32.0f); + + t_free(result); + t_free_matmul_packed_rhs(packed); + t_free(rhs); + t_free(left); +} + TEST(test_packed_rhs_handles_transposed_slices_and_kernel_tails) { int left_dims[2] = {4, 4}; int right_dims[2] = {4, 4}; @@ -521,6 +546,7 @@ int main(void) { RUN_TEST(test_t_matmul_transposed_2d_uneven_rows_and_columns); RUN_TEST(test_t_matmul_vector_cases); RUN_TEST(test_packed_rhs_snapshots_transposed_view); + RUN_TEST(test_packed_rhs_transpose_pack_retains_lifetime); RUN_TEST(test_packed_rhs_handles_transposed_slices_and_kernel_tails); RUN_TEST(test_packed_rhs_handles_non_multiple_kernel_dimensions); RUN_TEST(test_packed_rhs_handles_reshape_contiguous_squeeze_and_unsqueeze); From 3fbdbe7f3039e2cc5d725e4eb98b5459351a2bec Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:02:24 +0200 Subject: [PATCH 31/49] perf: reuse linear weight packs during backward --- include/tensorlib/autograd_internal.h | 6 ++++ include/tensorlib/nn.h | 2 ++ src/autograd/autograd_matmul.c | 48 ++++++++++++++++++++++++--- src/nn/linear.c | 17 ++++++++-- 4 files changed, 66 insertions(+), 7 deletions(-) diff --git a/include/tensorlib/autograd_internal.h b/include/tensorlib/autograd_internal.h index 66a7b0b..a3f92ab 100644 --- a/include/tensorlib/autograd_internal.h +++ b/include/tensorlib/autograd_internal.h @@ -18,6 +18,12 @@ ag_tensor* ag_matmul_packed_rhs( const ag_tensor* b, const tensor_matmul_packed_rhs* packed_rhs ); +ag_tensor* ag_matmul_packed_rhs_with_backward_pack( + const ag_tensor* a, + const ag_tensor* b, + const tensor_matmul_packed_rhs* packed_rhs, + const tensor_matmul_packed_rhs* backward_rhs +); ag_tensor* ag_layer_norm(const ag_tensor* input, const ag_tensor* weight, diff --git a/include/tensorlib/nn.h b/include/tensorlib/nn.h index 4f3ac0a..30f2af8 100644 --- a/include/tensorlib/nn.h +++ b/include/tensorlib/nn.h @@ -106,6 +106,8 @@ struct nn_linear { /* Runtime-only packed view of weight^T, refreshed by storage version. */ tensor_matmul_packed_rhs* packed_weight; uint64_t packed_weight_version; + tensor_matmul_packed_rhs* packed_backward_weight; + uint64_t packed_backward_weight_version; }; struct nn_embedding { diff --git a/src/autograd/autograd_matmul.c b/src/autograd/autograd_matmul.c index 42b48c2..31ab72c 100644 --- a/src/autograd/autograd_matmul.c +++ b/src/autograd/autograd_matmul.c @@ -1,5 +1,19 @@ +#include + #include "./../../include/tensorlib/autograd_internal.h" +typedef struct { + tensor_matmul_packed_rhs* backward_rhs; +} matmul_context; + +static void free_matmul_context(void* opaque) +{ + matmul_context* context = (matmul_context*)opaque; + if (context == NULL) return; + t_free_matmul_packed_rhs(context->backward_rhs); + free(context); +} + static tensor* promote_operand(tensor* value, int left_operand) { if (value->ndim != 1) return t_contiguous(value); return t_unsqueeze(value, left_operand ? 0 : 1); @@ -46,8 +60,12 @@ static int backward_matmul(const ag_node* node, if (promoted_a == NULL || promoted_b == NULL || promoted_gradient == NULL || transposed_a == NULL || transposed_b == NULL) goto fail; + matmul_context* context = (matmul_context*)node->context; if (node->inputs[0]->requires_grad) { - input_gradients[0] = t_matmul(promoted_gradient, transposed_b); + input_gradients[0] = context != NULL && context->backward_rhs != NULL + ? t_matmul_packed_rhs(promoted_gradient, + context->backward_rhs) + : t_matmul(promoted_gradient, transposed_b); if (input_gradients[0] == NULL) goto fail; if (a_vector) { input_gradients[0] = remove_vector_dimension(input_gradients[0], 1); @@ -77,20 +95,40 @@ static int backward_matmul(const ag_node* node, static ag_tensor* make_matmul_result(const ag_tensor* a, const ag_tensor* b, - tensor* output) { + tensor* output, + matmul_context* context) { ag_tensor* inputs[2] = {(ag_tensor*)a, (ag_tensor*)b}; return ag_make_result(output, AG_OP_MATMUL, 2, inputs, - backward_matmul, NULL, NULL); + backward_matmul, context, + context == NULL ? NULL : free_matmul_context); } ag_tensor* ag_matmul(const ag_tensor* a, const ag_tensor* b) { if (a == NULL || b == NULL) return NULL; - return make_matmul_result(a, b, t_matmul(a->value, b->value)); + return make_matmul_result(a, b, t_matmul(a->value, b->value), NULL); } ag_tensor* ag_matmul_packed_rhs(const ag_tensor* a, const ag_tensor* b, const tensor_matmul_packed_rhs* packed_rhs) { + return ag_matmul_packed_rhs_with_backward_pack(a, b, packed_rhs, NULL); +} + +ag_tensor* ag_matmul_packed_rhs_with_backward_pack( + const ag_tensor* a, + const ag_tensor* b, + const tensor_matmul_packed_rhs* packed_rhs, + const tensor_matmul_packed_rhs* backward_rhs) +{ + matmul_context* context = NULL; if (a == NULL || b == NULL || packed_rhs == NULL) return NULL; - return make_matmul_result(a, b, t_matmul_packed_rhs(a->value, packed_rhs)); + if (backward_rhs != NULL) { + context = (matmul_context*)calloc(1, sizeof(*context)); + if (context == NULL) return NULL; + context->backward_rhs = (tensor_matmul_packed_rhs*)backward_rhs; + t_retain_matmul_packed_rhs(context->backward_rhs); + } + return make_matmul_result(a, b, + t_matmul_packed_rhs(a->value, packed_rhs), + context); } diff --git a/src/nn/linear.c b/src/nn/linear.c index 57dbd54..41437f9 100644 --- a/src/nn/linear.c +++ b/src/nn/linear.c @@ -102,6 +102,7 @@ void nn_linear_destroy(nn_linear* layer) { if (layer == NULL) return; t_free_matmul_packed_rhs(layer->packed_weight); + t_free_matmul_packed_rhs(layer->packed_backward_weight); nn_module_destroy_base(&layer->base); free(layer); } @@ -139,11 +140,23 @@ ag_tensor* nn_linear_forward(const nn_linear* layer, const ag_tensor* input) mutable_layer->packed_weight_version = weight_version; } } + if (mutable_layer->packed_backward_weight == NULL || + mutable_layer->packed_backward_weight_version != weight_version) { + tensor_matmul_packed_rhs* packed = + t_pack_matmul_rhs(layer->weight->value->value); + if (packed != NULL) { + t_free_matmul_packed_rhs(mutable_layer->packed_backward_weight); + mutable_layer->packed_backward_weight = packed; + mutable_layer->packed_backward_weight_version = weight_version; + } + } product = input_value->ndim >= 2 && mutable_layer->packed_weight != NULL && mutable_layer->packed_weight_version == weight_version - ? ag_matmul_packed_rhs(input, transposed, - mutable_layer->packed_weight) + ? ag_matmul_packed_rhs_with_backward_pack( + input, transposed, mutable_layer->packed_weight, + mutable_layer->packed_backward_weight_version == weight_version + ? mutable_layer->packed_backward_weight : NULL) : ag_matmul(input, transposed); ag_tensor_release(transposed); if (product == NULL) return NULL; From 702a31f846dd4e371ff4372e6058dbbb816d9a62 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:03:02 +0200 Subject: [PATCH 32/49] perf: specialize AdamW gradient scanning --- src/optim/adamw.c | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/optim/adamw.c b/src/optim/adamw.c index 6f5e659..bc63142 100644 --- a/src/optim/adamw.c +++ b/src/optim/adamw.c @@ -189,11 +189,21 @@ static int gradient_scale(const nn_adamw* optimizer, double* result) optimizer->steps[i] == UINT64_MAX) { return -1; } - for (int element = 0; element < tensor_numel(gradient); ++element) { - double grad = gradient->storage->data[ - tensor_flat_index(gradient, element)]; - if (!isfinite(grad)) return -1; - squared_norm += grad * grad; + int count = tensor_numel(gradient); + if (is_contiguous(gradient) && gradient->offset == 0) { + const float* data = gradient->storage->data; + for (int element = 0; element < count; ++element) { + double grad = data[element]; + if (!isfinite(grad)) return -1; + squared_norm += grad * grad; + } + } else { + for (int element = 0; element < count; ++element) { + double grad = gradient->storage->data[ + tensor_flat_index(gradient, element)]; + if (!isfinite(grad)) return -1; + squared_norm += grad * grad; + } } } if (!isfinite(squared_norm)) return -1; From 8b294982ee097e94e223f0411b026b5a42165c7c Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:06:10 +0200 Subject: [PATCH 33/49] perf: parallelize deterministic AdamW updates --- src/optim/adamw.c | 116 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 107 insertions(+), 9 deletions(-) diff --git a/src/optim/adamw.c b/src/optim/adamw.c index bc63142..9662f1e 100644 --- a/src/optim/adamw.c +++ b/src/optim/adamw.c @@ -1,8 +1,12 @@ #include +#include #include #include #include "../../include/tensorlib/nn.h" +#include "../tensor/parallel.h" + +#define TENSORLIB_ADAMW_MIN_PARALLEL_ELEMENTS (1 << 16) static int tensor_flat_index(const tensor* value, int flat) { @@ -222,27 +226,102 @@ int nn_adamw_step(nn_adamw* optimizer) { const nn_adamw_config* config; double grad_scale; + uint64_t* next_steps = NULL; + double* correction1s = NULL; + double* correction2s = NULL; + long long work = 0; + int parameter_count; + int eligible_count = 0; + int threads; + int update_failed = 0; if (!topology_valid(optimizer) || gradient_scale(optimizer, &grad_scale) != 0) { return -1; } config = &optimizer->config; - for (size_t i = 0; i < optimizer->parameter_count; ++i) { + if (optimizer->parameter_count > (size_t)INT_MAX) return -1; + parameter_count = (int)optimizer->parameter_count; + /* First determine the amount of independent work without allocating + * parallel-only bookkeeping for the normal small serial path. */ + for (int i = 0; i < parameter_count; ++i) { + nn_parameter* parameter = optimizer->parameters[i]; + tensor* value; + + if (!parameter->trainable || parameter->value->grad == NULL) continue; + value = parameter->value->value; + ++eligible_count; + int count = tensor_numel(value); + work = work > LLONG_MAX - count ? LLONG_MAX : work + count; + } + threads = tensorlib_parallel_threads(work, + TENSORLIB_ADAMW_MIN_PARALLEL_ELEMENTS, + eligible_count); + if (threads > 1) { + next_steps = (uint64_t*)calloc((size_t)parameter_count, + sizeof(*next_steps)); + correction1s = (double*)calloc((size_t)parameter_count, + sizeof(*correction1s)); + correction2s = (double*)calloc((size_t)parameter_count, + sizeof(*correction2s)); + if (next_steps == NULL || correction1s == NULL || correction2s == NULL) { + free(correction2s); free(correction1s); free(next_steps); + return -1; + } + /* Validate every value and precompute all step-dependent factors before + * any parameter or optimizer state is modified. */ + for (int i = 0; i < parameter_count; ++i) { + nn_parameter* parameter = optimizer->parameters[i]; + tensor* value; + uint64_t step; + if (!parameter->trainable || parameter->value->grad == NULL) continue; + value = parameter->value->value; + step = optimizer->steps[i] + 1; + correction1s[i] = 1.0 - pow((double)config->beta1, (double)step); + correction2s[i] = 1.0 - pow((double)config->beta2, (double)step); + if (!(correction1s[i] > 0.0) || !(correction2s[i] > 0.0) || + optimizer->steps[i] == UINT64_MAX) { + free(correction2s); free(correction1s); free(next_steps); + return -1; + } + next_steps[i] = step; + int count = tensor_numel(value); + for (int element = 0; element < count; ++element) { + int value_index = tensor_flat_index(value, element); + int first_index = tensor_flat_index(optimizer->first_moments[i], element); + int second_index = tensor_flat_index(optimizer->second_moments[i], element); + if (!isfinite(value->storage->data[value_index]) || + !isfinite(optimizer->first_moments[i]->storage->data[first_index]) || + !isfinite(optimizer->second_moments[i]->storage->data[second_index])) { + free(correction2s); free(correction1s); free(next_steps); + return -1; + } + } + } + } +#ifdef _OPENMP +#pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) +#endif + for (int i = 0; i < parameter_count; ++i) { nn_parameter* parameter = optimizer->parameters[i]; tensor* value; tensor* gradient; - uint64_t step; double correction1; double correction2; if (!parameter->trainable || parameter->value->grad == NULL) continue; value = parameter->value->value; gradient = parameter->value->grad; - step = optimizer->steps[i] + 1; - correction1 = 1.0 - pow((double)config->beta1, (double)step); - correction2 = 1.0 - pow((double)config->beta2, (double)step); - if (!(correction1 > 0.0) || !(correction2 > 0.0)) return -1; + uint64_t step = optimizer->steps[i] + 1; + correction1 = threads > 1 ? correction1s[i] + : 1.0 - pow((double)config->beta1, (double)step); + correction2 = threads > 1 ? correction2s[i] + : 1.0 - pow((double)config->beta2, (double)step); + if (!(correction1 > 0.0) || !(correction2 > 0.0) || + optimizer->steps[i] == UINT64_MAX) { + update_failed = 1; + continue; + } optimizer->steps[i] = step; /* Fast path: when value/gradient/first/second-moment share a * contiguous layout, the flat element index is the storage index, @@ -276,7 +355,11 @@ int nn_adamw_step(nn_adamw* optimizer) (normalized + config->weight_decay * current); if (!isfinite(current) || !isfinite(first) || !isfinite(second) || !isfinite(updated)) { - return -1; +#ifdef _OPENMP +#pragma omp atomic write +#endif + update_failed = 1; + continue; } first_data[element] = (float)first; second_data[element] = (float)second; @@ -308,7 +391,11 @@ int nn_adamw_step(nn_adamw* optimizer) (normalized + config->weight_decay * current); if (!isfinite(current) || !isfinite(first) || !isfinite(second) || !isfinite(updated)) { - return -1; +#ifdef _OPENMP +#pragma omp atomic write +#endif + update_failed = 1; + continue; } optimizer->first_moments[i]->storage->data[first_index] = (float)first; @@ -317,10 +404,21 @@ int nn_adamw_step(nn_adamw* optimizer) value->storage->data[value_index] = (float)updated; } } + } + if (update_failed) { + free(correction2s); free(correction1s); free(next_steps); + return -1; + } + for (int i = 0; i < parameter_count; ++i) { + nn_parameter* parameter = optimizer->parameters[i]; + if (!parameter->trainable || parameter->value->grad == NULL) continue; tensor_mark_modified(optimizer->first_moments[i]); tensor_mark_modified(optimizer->second_moments[i]); - tensor_mark_modified(value); + tensor_mark_modified(parameter->value->value); } + free(correction2s); + free(correction1s); + free(next_steps); return 0; } From e597e77d126f6f1ecae769e28904dc7dd7f2808f Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:13:00 +0200 Subject: [PATCH 34/49] fix: keep serial warning build clean --- src/autograd/autograd_cross_entropy.c | 3 +++ src/autograd/autograd_layer_norm.c | 3 +++ src/autograd/autograd_ops.c | 3 +++ src/autograd/autograd_softmax.c | 3 +++ src/tensor/tensor_ops.c | 3 +++ 5 files changed, 15 insertions(+) diff --git a/src/autograd/autograd_cross_entropy.c b/src/autograd/autograd_cross_entropy.c index 6a2f6cd..524df87 100644 --- a/src/autograd/autograd_cross_entropy.c +++ b/src/autograd/autograd_cross_entropy.c @@ -66,6 +66,9 @@ static int backward_cross_entropy(const ag_node* node, int threads = tensorlib_parallel_threads( (long long)context->rows * context->classes, TENSORLIB_CROSS_ENTROPY_MIN_PARALLEL_ELEMENTS, context->rows); +#ifndef _OPENMP + (void)threads; +#endif #ifdef _OPENMP #pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) #endif diff --git a/src/autograd/autograd_layer_norm.c b/src/autograd/autograd_layer_norm.c index e524656..c0e866d 100644 --- a/src/autograd/autograd_layer_norm.c +++ b/src/autograd/autograd_layer_norm.c @@ -80,6 +80,9 @@ static int backward_layer_norm(const ag_node* node, int threads = tensorlib_parallel_threads( (long long)context->rows * width, TENSORLIB_LAYER_NORM_MIN_PARALLEL_ELEMENTS, context->rows); +#ifndef _OPENMP + (void)threads; +#endif if (threads > 1 && weight_gradient != NULL) { weight_partials = (float*)calloc((size_t)threads * (size_t)width, sizeof(*weight_partials)); diff --git a/src/autograd/autograd_ops.c b/src/autograd/autograd_ops.c index 484d8e3..58094b1 100644 --- a/src/autograd/autograd_ops.c +++ b/src/autograd/autograd_ops.c @@ -415,6 +415,9 @@ static int backward_gelu(const ag_node* node, int count = tensor_numel(input); int threads = tensorlib_parallel_threads( count, TENSORLIB_GELU_MIN_PARALLEL_ELEMENTS, 0); +#ifndef _OPENMP + (void)threads; +#endif if (is_contiguous(input) && is_contiguous((tensor*)output_gradient)) { const float* input_data = input->storage->data + input->offset; const float* upstream_data = output_gradient->storage->data + diff --git a/src/autograd/autograd_softmax.c b/src/autograd/autograd_softmax.c index 5c85d1d..82fabb8 100644 --- a/src/autograd/autograd_softmax.c +++ b/src/autograd/autograd_softmax.c @@ -50,6 +50,9 @@ static int backward_softmax(const ag_node* node, int threads = tensorlib_parallel_threads( (long long)context->rows * context->width, TENSORLIB_SOFTMAX_MIN_PARALLEL_ELEMENTS, context->rows); +#ifndef _OPENMP + (void)threads; +#endif #ifdef _OPENMP #pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) #endif diff --git a/src/tensor/tensor_ops.c b/src/tensor/tensor_ops.c index db329f2..e1e77af 100644 --- a/src/tensor/tensor_ops.c +++ b/src/tensor/tensor_ops.c @@ -168,6 +168,9 @@ static int try_contiguous_broadcast(tensor* a, tensor* b, float* c_data = c->storage->data + c->offset; int threads = tensorlib_parallel_threads( (long long)outer * block, TENSORLIB_OP_MIN_PARALLEL_ELEMENTS, outer); +#ifndef _OPENMP + (void)threads; +#endif #ifdef _OPENMP #pragma omp parallel for if(threads > 1) schedule(static) num_threads(threads) From 456181b0a4a89e42e318eb73a26cd7415bbc4ba4 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:33:31 +0200 Subject: [PATCH 35/49] bench: measure backward phases at requested thread counts --- CMakeLists.txt | 6 ++++++ benchmarks/bench_nn.c | 25 +++++++++++++++---------- benchmarks/bench_runner.c | 4 ++-- benchmarks/bench_runner.h | 1 + 4 files changed, 24 insertions(+), 12 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 0823cb3..e10c972 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -234,6 +234,12 @@ if (TENSORLIB_BUILD_BENCHMARKS) add_test(NAME test_bench_harness COMMAND test_bench_harness) add_test(NAME test_benchmarks_smoke COMMAND bench_tensorlib --smoke --threads 1) + if (TENSORLIB_ENABLE_OPENMP) + add_test(NAME test_benchmark_threads + COMMAND bench_tensorlib --smoke --suite nn --threads 1,2,4) + set_tests_properties(test_benchmark_threads PROPERTIES + PASS_REGULAR_EXPRESSION "tiny_lm_backward +threads=4") + endif() endif() find_path(OPENBLAS_INCLUDE_DIR NAMES cblas.h PATH_SUFFIXES openblas) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index 8c0441f..7f2899e 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -181,6 +181,7 @@ static int compare_double(const void* left, const void* right) static void report_training_phases(const bench_options* options, FILE* csv, int requested_threads, + int measured_threads, nn_bench_context* context) { static const char* names[PHASE_COUNT] = { @@ -204,7 +205,7 @@ static void report_training_phases(const bench_options* options, result.iterations_per_sample = 1; benchmark.context = context; bench_record_measurement(options, csv, &benchmark, - requested_threads, requested_threads, &result); + requested_threads, measured_threads, &result); } } @@ -496,27 +497,31 @@ static int run_decoder_case(const bench_options* options, train ? "forward+loss+backward+adamw" : "forward;graph-build", "tokens/s", (double)(batch * time), threads, &context, result); if (status == 0 && train && strcmp(suite, "nn") == 0) { + int measured_threads = bench_configure_threads(threads); tensor_alloc_stats_read(&context.allocation_stats); - report_training_phases(options, csv, threads, &context); - report_allocation_stats(options, csv, threads, threads, &context); + report_training_phases(options, csv, threads, measured_threads, &context); + report_allocation_stats(options, csv, threads, measured_threads, &context); } if (train) tensor_alloc_stats_enable(0); destroy_context(&context); return status; } -static int run_decoder_cases(const bench_options* options, FILE* csv, - int batch, int threads) +static int run_decoder_cases(const bench_options* options, FILE* csv, int batch) { int smoke = strcmp(options->profile.profile, "smoke") == 0; int time = smoke ? 8 : 128; int channels = smoke ? 24 : 192; int layers = smoke ? 1 : 4; int forward_status = run_decoder_case(options, csv, "nn", batch, time, - channels, layers, 0, threads, NULL); - int train_status = run_decoder_case(options, csv, "nn", batch, time, - channels, layers, 1, threads, NULL); - return forward_status == 1 || train_status == 1; + channels, layers, 0, options->threads[0], NULL); + int status = forward_status == 1; + for (int index = 0; index < options->thread_count; ++index) { + int train_status = run_decoder_case(options, csv, "nn", batch, time, + channels, layers, 1, options->threads[index], NULL); + if (train_status == 1) status = 1; + } + return status; } int bench_run_decoder_scaling(const bench_options* options, FILE* csv) @@ -580,7 +585,7 @@ int bench_run_nn_suite(const bench_options* options, FILE* csv) batch, time, channels); status |= run_mlp(options, csv, smoke ? 2 : 64, 0); status |= run_mlp(options, csv, smoke ? 2 : 64, 1); - status |= run_decoder_cases(options, csv, batch, 1); + status |= run_decoder_cases(options, csv, batch); printf("\n"); return status; } diff --git a/benchmarks/bench_runner.c b/benchmarks/bench_runner.c index 57d0fe7..f02baf2 100644 --- a/benchmarks/bench_runner.c +++ b/benchmarks/bench_runner.c @@ -10,7 +10,7 @@ #include #endif -static int actual_threads(int requested) +int bench_configure_threads(int requested) { #ifdef _OPENMP int actual = 1; @@ -112,7 +112,7 @@ int bench_execute_case(const bench_options* options, return 2; } #endif - int measured_threads = actual_threads(requested_threads); + int measured_threads = bench_configure_threads(requested_threads); if (bench_measure_with_reset(benchmark->operation, benchmark->context, &options->profile, benchmark->reset, result) != 0 || !isfinite(result->checksum)) { diff --git a/benchmarks/bench_runner.h b/benchmarks/bench_runner.h index 902b0a2..6255877 100644 --- a/benchmarks/bench_runner.h +++ b/benchmarks/bench_runner.h @@ -28,6 +28,7 @@ typedef struct { } bench_case; int bench_run_suites(const bench_options* options, FILE* csv); +int bench_configure_threads(int requested_threads); int bench_execute_case(const bench_options* options, FILE* csv, const bench_case* benchmark, From 53488f09710b116d34ef0978cef1bcc74a154ff3 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:35:14 +0200 Subject: [PATCH 36/49] bench: isolate linear backward matrix operations --- CMakeLists.txt | 1 + benchmarks/bench_backward.c | 163 ++++++++++++++++++++++++++++++++++++ benchmarks/bench_runner.c | 2 + 3 files changed, 166 insertions(+) create mode 100644 benchmarks/bench_backward.c diff --git a/CMakeLists.txt b/CMakeLists.txt index e10c972..4e1d70e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -223,6 +223,7 @@ if (TENSORLIB_BUILD_BENCHMARKS) benchmarks/bench_runner.c benchmarks/bench_kernels.c benchmarks/bench_autograd.c + benchmarks/bench_backward.c benchmarks/bench_nn.c benchmarks/bench_scaling.c) target_link_libraries(bench_tensorlib PRIVATE tensorlib ${TENSORLIB_MATH_LIB}) diff --git a/benchmarks/bench_backward.c b/benchmarks/bench_backward.c new file mode 100644 index 0000000..8c3bf48 --- /dev/null +++ b/benchmarks/bench_backward.c @@ -0,0 +1,163 @@ +#include "bench_runner.h" + +#include +#include + +#include + +typedef enum { + LINEAR_DINPUT, + LINEAR_DWEIGHT, + LINEAR_DBIAS, + LINEAR_PACK +} linear_gradient_kind; + +typedef struct { + linear_gradient_kind kind; + tensor* input; + tensor* output_gradient; + tensor* weight; + tensor_matmul_packed_rhs* packed_weight; +} linear_gradient_context; + +typedef struct { + const char* name; + int in_features; + int out_features; +} linear_shape; + +static void fill_tensor(tensor* value, float scale) +{ + int count = tensor_numel(value); + for (int index = 0; index < count; ++index) { + value->storage->data[index] = scale * (float)((index % 31) - 15); + } +} + +static void destroy_linear_context(linear_gradient_context* context) +{ + t_free_matmul_packed_rhs(context->packed_weight); + t_free(context->weight); + t_free(context->output_gradient); + t_free(context->input); + memset(context, 0, sizeof(*context)); +} + +static int init_linear_context(linear_gradient_context* context, + linear_gradient_kind kind, + int in_features, + int out_features) +{ + int input_dims[3] = {4, 128, in_features}; + int output_dims[3] = {4, 128, out_features}; + int weight_dims[2] = {out_features, in_features}; + memset(context, 0, sizeof(*context)); + context->kind = kind; + context->input = t_alloc(3, input_dims); + context->output_gradient = t_alloc(3, output_dims); + context->weight = t_alloc(2, weight_dims); + if (context->input == NULL || context->output_gradient == NULL || + context->weight == NULL) return 1; + fill_tensor(context->input, 0.001f); + fill_tensor(context->output_gradient, 0.002f); + fill_tensor(context->weight, 0.003f); + if (kind == LINEAR_DINPUT) { + context->packed_weight = t_pack_matmul_rhs(context->weight); + if (context->packed_weight == NULL) return 1; + } + return 0; +} + +static int linear_gradient_operation(void* opaque, double* checksum) +{ + linear_gradient_context* context = (linear_gradient_context*)opaque; + tensor* first = NULL; + tensor* second = NULL; + tensor* output = NULL; + if (context->kind == LINEAR_DINPUT) { + output = t_matmul_packed_rhs(context->output_gradient, + context->packed_weight); + } else if (context->kind == LINEAR_DWEIGHT) { + first = t_transpose(context->input, 1, 2); + second = first == NULL ? NULL : + t_matmul(first, context->output_gradient); + output = second == NULL ? NULL : t_sum(second, 0); + } else if (context->kind == LINEAR_DBIAS) { + first = t_sum(context->output_gradient, 0); + output = first == NULL ? NULL : t_sum(first, 0); + } else { + tensor_matmul_packed_rhs* packed = + t_pack_matmul_rhs(context->weight); + if (packed == NULL) return 1; + *checksum += 1.0; + t_free_matmul_packed_rhs(packed); + return 0; + } + if (output == NULL) { + t_free(second); + t_free(first); + return 1; + } + *checksum += output->storage->data[output->offset]; + t_free(output); + t_free(second); + t_free(first); + return 0; +} + +static int run_linear_gradient(const bench_options* options, + FILE* csv, + const linear_shape* shape, + linear_gradient_kind kind, + const char* suffix, + const char* layout) +{ + linear_gradient_context context; + bench_measurement result; + char name[64]; + char dimensions[64]; + snprintf(name, sizeof(name), "%s_%s", shape->name, suffix); + snprintf(dimensions, sizeof(dimensions), "[4x128x%d]->[4x128x%d]", + shape->in_features, shape->out_features); + if (init_linear_context(&context, kind, shape->in_features, + shape->out_features) != 0) { + destroy_linear_context(&context); + return 1; + } + bench_case benchmark = { + "nn_backward", name, dimensions, layout, "ms/call", 0.0, 1, + linear_gradient_operation, &context, NULL + }; + int status = bench_execute_case(options, csv, &benchmark, + options->threads[0], &result); + destroy_linear_context(&context); + return status == 1; +} + +int bench_run_backward_matrix_suite(const bench_options* options, FILE* csv) +{ + static const linear_shape shapes[] = { + {"linear_qkv", 192, 576}, + {"linear_attention_output", 192, 192}, + {"linear_mlp_expand", 192, 768}, + {"linear_mlp_project", 768, 192}, + {"linear_vocabulary", 192, 256} + }; + int status = 0; + printf("Backward matrix components\n"); + for (size_t index = 0; index < sizeof(shapes) / sizeof(shapes[0]); ++index) { + status |= run_linear_gradient(options, csv, &shapes[index], + LINEAR_DINPUT, "dinput", + "packed-weight"); + status |= run_linear_gradient(options, csv, &shapes[index], + LINEAR_DWEIGHT, "dweight", + "batched-matmul+reduce"); + status |= run_linear_gradient(options, csv, &shapes[index], + LINEAR_DBIAS, "dbias", + "two-axis-reduce"); + status |= run_linear_gradient(options, csv, &shapes[index], + LINEAR_PACK, "pack", "packing-only"); + } + printf("\n"); + return status; +} diff --git a/benchmarks/bench_runner.c b/benchmarks/bench_runner.c index f02baf2..212e2d0 100644 --- a/benchmarks/bench_runner.c +++ b/benchmarks/bench_runner.c @@ -132,6 +132,7 @@ static int suite_selected(const char* requested, const char* suite) int bench_run_kernel_suite(const bench_options* options, FILE* csv); int bench_run_autograd_suite(const bench_options* options, FILE* csv); int bench_run_nn_suite(const bench_options* options, FILE* csv); +int bench_run_backward_matrix_suite(const bench_options* options, FILE* csv); int bench_run_scaling_suite(const bench_options* options, FILE* csv); int bench_run_suites(const bench_options* options, FILE* csv) @@ -147,6 +148,7 @@ int bench_run_suites(const bench_options* options, FILE* csv) } if (suite_selected(options->suite, "nn")) { status |= bench_run_nn_suite(options, csv); + status |= bench_run_backward_matrix_suite(options, csv); } if (suite_selected(options->suite, "scaling")) { status |= bench_run_scaling_suite(options, csv); From ef9084c76581c13997dd37d1b4c591da53497a4d Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:36:28 +0200 Subject: [PATCH 37/49] bench: isolate attention backward matrix operations --- benchmarks/bench_backward.c | 94 +++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/benchmarks/bench_backward.c b/benchmarks/bench_backward.c index 8c3bf48..22f5c3a 100644 --- a/benchmarks/bench_backward.c +++ b/benchmarks/bench_backward.c @@ -26,6 +26,13 @@ typedef struct { int out_features; } linear_shape; +typedef struct { + int probability_value; + tensor* left; + tensor* right; + tensor* output_gradient; +} attention_gradient_context; + static void fill_tensor(tensor* value, float scale) { int count = tensor_numel(value); @@ -134,6 +141,91 @@ static int run_linear_gradient(const bench_options* options, return status == 1; } +static void destroy_attention_context(attention_gradient_context* context) +{ + t_free(context->output_gradient); + t_free(context->right); + t_free(context->left); + memset(context, 0, sizeof(*context)); +} + +static int init_attention_context(attention_gradient_context* context, + int probability_value) +{ + int square_dims[3] = {24, 128, 128}; + int head_dims[3] = {24, 128, 32}; + memset(context, 0, sizeof(*context)); + context->probability_value = probability_value; + context->left = t_alloc(3, probability_value ? square_dims : head_dims); + context->right = t_alloc(3, head_dims); + context->output_gradient = + t_alloc(3, probability_value ? head_dims : square_dims); + if (context->left == NULL || context->right == NULL || + context->output_gradient == NULL) return 1; + fill_tensor(context->left, 0.001f); + fill_tensor(context->right, 0.002f); + fill_tensor(context->output_gradient, 0.003f); + return 0; +} + +static int attention_gradient_operation(void* opaque, double* checksum) +{ + attention_gradient_context* context = + (attention_gradient_context*)opaque; + tensor* left_transpose = NULL; + tensor* right_transpose = NULL; + tensor* left_gradient = NULL; + tensor* right_gradient = NULL; + if (context->probability_value) { + right_transpose = t_transpose(context->right, 1, 2); + left_transpose = t_transpose(context->left, 1, 2); + left_gradient = right_transpose == NULL ? NULL : + t_matmul(context->output_gradient, right_transpose); + right_gradient = left_transpose == NULL ? NULL : + t_matmul(left_transpose, context->output_gradient); + } else { + left_transpose = t_transpose(context->output_gradient, 1, 2); + left_gradient = t_matmul(context->output_gradient, context->right); + right_gradient = left_transpose == NULL ? NULL : + t_matmul(left_transpose, context->left); + } + if (left_gradient == NULL || right_gradient == NULL) { + t_free(right_gradient); t_free(left_gradient); + t_free(right_transpose); t_free(left_transpose); + return 1; + } + *checksum += left_gradient->storage->data[left_gradient->offset] + + right_gradient->storage->data[right_gradient->offset]; + t_free(right_gradient); t_free(left_gradient); + t_free(right_transpose); t_free(left_transpose); + return 0; +} + +static int run_attention_gradient(const bench_options* options, + FILE* csv, + int probability_value) +{ + attention_gradient_context context; + bench_measurement result; + if (init_attention_context(&context, probability_value) != 0) { + destroy_attention_context(&context); + return 1; + } + bench_case benchmark = { + "nn_backward", + probability_value ? "attention_probability_value_backward" : + "attention_qk_transpose_backward", + probability_value ? "[24x128x128]@[24x128x32]" : + "[24x128x32]@[24x32x128]", + "two-gradient-matmuls", "ms/call", 0.0, 1, + attention_gradient_operation, &context, NULL + }; + int status = bench_execute_case(options, csv, &benchmark, + options->threads[0], &result); + destroy_attention_context(&context); + return status == 1; +} + int bench_run_backward_matrix_suite(const bench_options* options, FILE* csv) { static const linear_shape shapes[] = { @@ -158,6 +250,8 @@ int bench_run_backward_matrix_suite(const bench_options* options, FILE* csv) status |= run_linear_gradient(options, csv, &shapes[index], LINEAR_PACK, "pack", "packing-only"); } + status |= run_attention_gradient(options, csv, 0); + status |= run_attention_gradient(options, csv, 1); printf("\n"); return status; } From 6714a2c7129b35415194653611b107168e7e6594 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:39:08 +0200 Subject: [PATCH 38/49] bench: attribute autograd backward engine costs --- benchmarks/bench_nn.c | 52 +++++++++++++++ include/tensorlib/autograd_internal.h | 15 +++++ src/autograd/autograd_backward.c | 66 +++++++++++++++++++- tests/unit/autograd/test_autograd_backward.c | 24 +++++++ 4 files changed, 156 insertions(+), 1 deletion(-) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index 7f2899e..4fbf786 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -4,6 +4,7 @@ #include #include +#include #include "../src/tensor/tensor_alloc_internal.h" typedef ag_tensor* (*model_forward)(const void* model, const ag_tensor* input); @@ -23,6 +24,7 @@ typedef struct { int phase_sample_count; int phase_sample_capacity; tensor_alloc_stats allocation_stats; + ag_backward_stats backward_stats; size_t allocation_baseline_bytes; } nn_bench_context; @@ -171,6 +173,45 @@ static void report_allocation_stats(const bench_options* options, context->allocation_baseline_bytes : 0)); } +static void report_backward_stats(const bench_options* options, + FILE* csv, + int requested_threads, + int measured_threads, + const nn_bench_context* context) +{ + static const char* names[AG_BACKWARD_OP_COUNT] = { + "add", "sub", "mul", "div", "neg", "exp", "log", "pow", + "sqrt", "relu", "sigmoid", "tanh", "gelu", "matmul", "sum", + "mean", "max", "reshape", "transpose", "slice", "expand", + "gather_rows", "mul_scalar", "div_scalar", "layer_norm", + "softmax", "log_softmax", "cross_entropy" + }; + const ag_backward_stats* stats = &context->backward_stats; + for (int operation = 0; operation < AG_BACKWARD_OP_COUNT; ++operation) { + if (stats->operation_calls[operation] == 0) continue; + char name[64]; + snprintf(name, sizeof(name), "backward_op_%s", names[operation]); + bench_record_scalar(options, csv, "nn_backward", name, + "[Bx128]->[Bx128x256]", "profiled-single-call", "ms/call", + requested_threads, measured_threads, + stats->operation_seconds[operation] * 1000.0); + } + static const char* engine_names[] = { + "backward_graph_traversal", "backward_shape_reduction", + "backward_gradient_accumulation", "backward_persistent_merge" + }; + const double engine_values[] = { + stats->traversal_seconds, stats->reduction_seconds, + stats->accumulation_seconds, stats->merge_seconds + }; + for (int index = 0; index < 4; ++index) { + bench_record_scalar(options, csv, "nn_backward", engine_names[index], + "[Bx128]->[Bx128x256]", "profiled-single-call", "ms/call", + requested_threads, measured_threads, + engine_values[index] * 1000.0); + } +} + static int compare_double(const void* left, const void* right) { double a = *(const double*)left; @@ -501,6 +542,17 @@ static int run_decoder_case(const bench_options* options, tensor_alloc_stats_read(&context.allocation_stats); report_training_phases(options, csv, threads, measured_threads, &context); report_allocation_stats(options, csv, threads, measured_threads, &context); + ag_backward_stats_reset(); + ag_backward_stats_enable(1); + double checksum = 0.0; + int profile_status = nn_operation(&context, &checksum); + ag_backward_stats_enable(0); + ag_backward_stats_read(&context.backward_stats); + if (profile_status != 0) status = 1; + if (status == 0) { + report_backward_stats(options, csv, threads, measured_threads, + &context); + } } if (train) tensor_alloc_stats_enable(0); destroy_context(&context); diff --git a/include/tensorlib/autograd_internal.h b/include/tensorlib/autograd_internal.h index a3f92ab..a2a3712 100644 --- a/include/tensorlib/autograd_internal.h +++ b/include/tensorlib/autograd_internal.h @@ -3,6 +3,21 @@ #include "../../include/tensorlib/autograd.h" +#define AG_BACKWARD_OP_COUNT ((int)AG_OP_CROSS_ENTROPY + 1) + +typedef struct { + double operation_seconds[AG_BACKWARD_OP_COUNT]; + unsigned long operation_calls[AG_BACKWARD_OP_COUNT]; + double traversal_seconds; + double reduction_seconds; + double accumulation_seconds; + double merge_seconds; +} ag_backward_stats; + +void ag_backward_stats_enable(int enabled); +void ag_backward_stats_reset(void); +void ag_backward_stats_read(ag_backward_stats* output); + ag_tensor* ag_make_result(tensor* output, ag_op operation, int input_count, diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index 5b0a2ce..f968824 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -1,4 +1,6 @@ #include +#include +#include #include "./../../include/tensorlib/autograd_internal.h" @@ -14,6 +16,35 @@ typedef struct { int capacity; } node_list; +static int backward_stats_enabled; +static ag_backward_stats backward_stats; + +static double backward_now(void) +{ + return (double)clock() / (double)CLOCKS_PER_SEC; +} + +static double backward_elapsed(double started) +{ + double elapsed = backward_now() - started; + return elapsed > 0.0 ? elapsed : 0.0; +} + +void ag_backward_stats_enable(int enabled) +{ + backward_stats_enabled = enabled != 0; +} + +void ag_backward_stats_reset(void) +{ + memset(&backward_stats, 0, sizeof(backward_stats)); +} + +void ag_backward_stats_read(ag_backward_stats* output) +{ + if (output != NULL) *output = backward_stats; +} + static int append_tensor(tensor_list* list, ag_tensor* value) { if (list->count == list->capacity) { int capacity = list->capacity == 0 ? 16 : list->capacity * 2; @@ -153,14 +184,25 @@ static int add_in_place(tensor* destination, const tensor* source) static int accumulate_pass_gradient(tensor** destination, tensor* contribution, const tensor* target) { + double started = backward_stats_enabled ? backward_now() : 0.0; tensor* reduced = reduce_to_shape(contribution, target); + if (backward_stats_enabled) { + backward_stats.reduction_seconds += backward_elapsed(started); + started = backward_now(); + } if (reduced == NULL) return 1; if (*destination == NULL) { *destination = reduced; + if (backward_stats_enabled) { + backward_stats.accumulation_seconds += backward_elapsed(started); + } return 0; } if (add_in_place(*destination, reduced) == 0) { t_free(reduced); + if (backward_stats_enabled) { + backward_stats.accumulation_seconds += backward_elapsed(started); + } return 0; } tensor* sum = t_add(*destination, reduced); @@ -168,6 +210,9 @@ static int accumulate_pass_gradient(tensor** destination, if (sum == NULL) return 1; t_free(*destination); *destination = sum; + if (backward_stats_enabled) { + backward_stats.accumulation_seconds += backward_elapsed(started); + } return 0; } @@ -225,9 +270,13 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { tensor** contributions = NULL; int contribution_capacity = 0; int status = 1; + double started = backward_stats_enabled ? backward_now() : 0.0; if (collect_graph(output, &tensors, &nodes) != 0) goto cleanup; if (!graph_versions_match(&nodes)) goto cleanup; + if (backward_stats_enabled) { + backward_stats.traversal_seconds += backward_elapsed(started); + } for (int node_index = 0; node_index < nodes.count; ++node_index) { if (nodes.values[node_index]->input_count > contribution_capacity) { @@ -258,7 +307,18 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { contributions[input_index] = NULL; } - if (node->backward(node, pass_gradients[gradient_index], contributions) != 0) { + started = backward_stats_enabled ? backward_now() : 0.0; + int backward_status = + node->backward(node, pass_gradients[gradient_index], contributions); + if (backward_stats_enabled) { + int operation = (int)node->operation; + if (operation >= 0 && operation < AG_BACKWARD_OP_COUNT) { + backward_stats.operation_seconds[operation] += + backward_elapsed(started); + ++backward_stats.operation_calls[operation]; + } + } + if (backward_status != 0) { free_contributions(contributions, node->input_count); goto cleanup; } @@ -282,7 +342,11 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { } } + started = backward_stats_enabled ? backward_now() : 0.0; status = merge_persistent_gradients(&tensors, pass_gradients); + if (backward_stats_enabled) { + backward_stats.merge_seconds += backward_elapsed(started); + } cleanup: free_contributions(contributions, contribution_capacity); diff --git a/tests/unit/autograd/test_autograd_backward.c b/tests/unit/autograd/test_autograd_backward.c index 954cad9..8e31878 100644 --- a/tests/unit/autograd/test_autograd_backward.c +++ b/tests/unit/autograd/test_autograd_backward.c @@ -1,4 +1,5 @@ #include "./../../../include/tensorlib/autograd.h" +#include "./../../../include/tensorlib/autograd_internal.h" #include "./../../fixtures/test_common.h" static ag_tensor* make_ag(int ndim, const int* dims, const float* values, int requires_grad) { @@ -28,6 +29,28 @@ TEST(test_backward_chain_computes_weight_gradients) { ag_tensor_release(b); ag_tensor_release(a); } +TEST(test_backward_stats_are_disabled_by_default_and_resettable) { + int dims[1] = {2}; + float values[2] = {2.0f, 3.0f}; + ag_tensor* input = make_ag(1, dims, values, 1); + ag_tensor* square = ag_mul(input, input); + ag_tensor* loss = ag_mean(square, 0, 0); + ag_backward_stats stats; + ag_backward_stats_reset(); + ag_backward_stats_enable(1); + ASSERT_EQ_INT(ag_backward(loss), 0); + ag_backward_stats_enable(0); + ag_backward_stats_read(&stats); + ASSERT_EQ_INT((int)stats.operation_calls[AG_OP_MEAN], 1); + ASSERT_EQ_INT((int)stats.operation_calls[AG_OP_MUL], 1); + ASSERT_TRUE(stats.traversal_seconds >= 0.0); + ASSERT_TRUE(stats.merge_seconds >= 0.0); + ag_backward_stats_reset(); + ag_backward_stats_read(&stats); + ASSERT_EQ_INT((int)stats.operation_calls[AG_OP_MUL], 0); + ag_tensor_release(loss); ag_tensor_release(square); ag_tensor_release(input); +} + TEST(test_shared_dag_accumulates_each_branch_once) { int dims[1]={3}; float values[3]={1,2,3}; ag_tensor* x=make_ag(1,dims,values,1), *square=ag_mul(x,x), *joined=ag_add(square,x); @@ -167,6 +190,7 @@ TEST(test_backward_rejects_modified_intermediate_output) { int main(void) { printf("== autograd_backward.c ==\n"); RUN_TEST(test_backward_chain_computes_weight_gradients); + RUN_TEST(test_backward_stats_are_disabled_by_default_and_resettable); RUN_TEST(test_shared_dag_accumulates_each_branch_once); RUN_TEST(test_backward_unbroadcasts_gradient_to_operand_shape); RUN_TEST(test_seeded_backward_supports_non_scalar_output); From 4541da2178798d18fa1f604d324a73541cc13816 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:44:11 +0200 Subject: [PATCH 39/49] perf: restore serial gradient fast paths --- src/autograd/autograd_backward.c | 78 ++++++++++++++++++++++++-------- 1 file changed, 60 insertions(+), 18 deletions(-) diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index f968824..6633504 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -102,13 +102,17 @@ static int graph_versions_match(const node_list* nodes) { return 1; } -static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { +static tensor* reduce_to_shape(tensor* contribution, const tensor* target, + int* coords, int coord_capacity) { if (!tensor_has_valid_metadata(contribution) || !tensor_has_valid_metadata(target) || contribution->ndim < target->ndim) { t_free(contribution); return NULL; } + if (contribution->ndim == target->ndim && + same_shape(contribution, (tensor*)target)) return contribution; + for (int target_axis = 0; target_axis < target->ndim; ++target_axis) { int contribution_axis = contribution->ndim - target->ndim + target_axis; if (target->dims[target_axis] != 1 && @@ -117,14 +121,8 @@ static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { return NULL; } } - if (contribution->ndim == target->ndim && - same_shape(contribution, (tensor*)target)) return contribution; - tensor* reduced = t_alloc(target->ndim, target->dims); - int* coords = contribution->ndim > 0 - ? (int*)calloc((size_t)contribution->ndim, sizeof(*coords)) : NULL; - if (reduced == NULL || (contribution->ndim > 0 && coords == NULL)) { - free(coords); + if (reduced == NULL) { t_free(reduced); t_free(contribution); return NULL; @@ -133,6 +131,33 @@ static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { reduced->storage->data[reduced->offset + index] = 0.0f; } + int suffix_matches = is_contiguous(contribution); + for (int axis = 0; axis < target->ndim && suffix_matches; ++axis) { + int contribution_axis = contribution->ndim - target->ndim + axis; + suffix_matches = target->dims[axis] == contribution->dims[contribution_axis]; + } + int reduced_count = tensor_numel(reduced); + if (suffix_matches && reduced_count > 0) { + int outer_count = tensor_numel(contribution) / reduced_count; + const float* source = contribution->storage->data + contribution->offset; + float* destination = reduced->storage->data + reduced->offset; + for (int outer = 0; outer < outer_count; ++outer) { + for (int index = 0; index < reduced_count; ++index) { + destination[index] += source[outer * reduced_count + index]; + } + } + t_free(contribution); + return reduced; + } + + if (contribution->ndim > coord_capacity || + (contribution->ndim > 0 && coords == NULL)) { + t_free(reduced); + t_free(contribution); + return NULL; + } + memset(coords, 0, (size_t)contribution->ndim * sizeof(*coords)); + int contribution_count = tensor_numel(contribution); for (int index = 0; index < contribution_count; ++index) { int reduced_index = reduced->offset; @@ -146,12 +171,12 @@ static tensor* reduce_to_shape(tensor* contribution, const tensor* target) { contribution->storage->data[get_flat_index_nd(contribution, coords)]; advance_coords(coords, contribution->dims, contribution->ndim); } - free(coords); t_free(contribution); return reduced; } -static int add_in_place(tensor* destination, const tensor* source) +static int add_in_place(tensor* destination, const tensor* source, + int* coords, int coord_capacity) { if (!tensor_has_valid_metadata(destination) || !tensor_has_valid_metadata(source) || @@ -167,9 +192,9 @@ static int add_in_place(tensor* destination, const tensor* source) return 0; } - int* coords = destination->ndim > 0 - ? (int*)calloc((size_t)destination->ndim, sizeof(*coords)) : NULL; - if (destination->ndim > 0 && coords == NULL) return 1; + if (destination->ndim > coord_capacity || + (destination->ndim > 0 && coords == NULL)) return 1; + memset(coords, 0, (size_t)destination->ndim * sizeof(*coords)); for (int index = 0; index < count; ++index) { int destination_index = get_flat_index_nd(destination, coords); int source_index = get_flat_index_nd((tensor*)source, coords); @@ -177,15 +202,17 @@ static int add_in_place(tensor* destination, const tensor* source) source->storage->data[source_index]; advance_coords(coords, destination->dims, destination->ndim); } - free(coords); return 0; } static int accumulate_pass_gradient(tensor** destination, tensor* contribution, - const tensor* target) { + const tensor* target, + int* coords, + int coord_capacity) { double started = backward_stats_enabled ? backward_now() : 0.0; - tensor* reduced = reduce_to_shape(contribution, target); + tensor* reduced = reduce_to_shape(contribution, target, coords, + coord_capacity); if (backward_stats_enabled) { backward_stats.reduction_seconds += backward_elapsed(started); started = backward_now(); @@ -198,7 +225,7 @@ static int accumulate_pass_gradient(tensor** destination, } return 0; } - if (add_in_place(*destination, reduced) == 0) { + if (add_in_place(*destination, reduced, coords, coord_capacity) == 0) { t_free(reduced); if (backward_stats_enabled) { backward_stats.accumulation_seconds += backward_elapsed(started); @@ -268,7 +295,9 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { node_list nodes = {0}; tensor** pass_gradients = NULL; tensor** contributions = NULL; + int* reduction_coords = NULL; int contribution_capacity = 0; + int reduction_coord_capacity = 0; int status = 1; double started = backward_stats_enabled ? backward_now() : 0.0; @@ -278,6 +307,16 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { backward_stats.traversal_seconds += backward_elapsed(started); } + for (int index = 0; index < tensors.count; ++index) { + int ndim = tensors.values[index]->value->ndim; + if (ndim > reduction_coord_capacity) reduction_coord_capacity = ndim; + } + if (reduction_coord_capacity > 0) { + reduction_coords = (int*)calloc((size_t)reduction_coord_capacity, + sizeof(*reduction_coords)); + if (reduction_coords == NULL) goto cleanup; + } + for (int node_index = 0; node_index < nodes.count; ++node_index) { if (nodes.values[node_index]->input_count > contribution_capacity) { contribution_capacity = nodes.values[node_index]->input_count; @@ -333,7 +372,9 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { int destination = input->graph_index; if (destination < 0 || contributions[input_index] == NULL || accumulate_pass_gradient(&pass_gradients[destination], - contributions[input_index], input->value) != 0) { + contributions[input_index], input->value, + reduction_coords, + reduction_coord_capacity) != 0) { contributions[input_index] = NULL; free_contributions(contributions, node->input_count); goto cleanup; @@ -349,6 +390,7 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { } cleanup: + free(reduction_coords); free_contributions(contributions, contribution_capacity); free(contributions); if (pass_gradients != NULL) { From 38f38bf13ceb994d97a9c2c4323674cbce8f0b9b Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:49:48 +0200 Subject: [PATCH 40/49] perf: reduce broadcast gradients at their producers --- include/tensorlib/autograd_internal.h | 1 + src/autograd/autograd_backward.c | 62 +++++++++++++++++++++++++ src/autograd/autograd_ops.c | 9 ++-- tests/unit/autograd/test_autograd_ops.c | 35 ++++++++++++++ 4 files changed, 104 insertions(+), 3 deletions(-) diff --git a/include/tensorlib/autograd_internal.h b/include/tensorlib/autograd_internal.h index a2a3712..77e25a8 100644 --- a/include/tensorlib/autograd_internal.h +++ b/include/tensorlib/autograd_internal.h @@ -27,6 +27,7 @@ ag_tensor* ag_make_result(tensor* output, void (*free_context)(void*)); tensor* ag_full_like(const tensor* reference, float value); +tensor* ag_sum_to_shape(const tensor* source, const tensor* target, float scale); ag_tensor* ag_matmul_packed_rhs( const ag_tensor* a, diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index 6633504..b8b4653 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -102,6 +102,68 @@ static int graph_versions_match(const node_list* nodes) { return 1; } +tensor* ag_sum_to_shape(const tensor* source, const tensor* target, float scale) +{ + if (!tensor_has_valid_metadata(source) || !tensor_has_valid_metadata(target) || + source->ndim < target->ndim) return NULL; + for (int axis = 0; axis < target->ndim; ++axis) { + int source_axis = source->ndim - target->ndim + axis; + if (target->dims[axis] != 1 && + target->dims[axis] != source->dims[source_axis]) return NULL; + } + if (source->ndim == target->ndim && + same_shape((tensor*)source, (tensor*)target)) { + if (scale == 1.0f) return t_clone((tensor*)source); + if (scale == -1.0f) return t_neg((tensor*)source); + return t_mul_scalar((tensor*)source, scale); + } + + tensor* result = t_alloc(target->ndim, target->dims); + if (result == NULL) return NULL; + int result_count = tensor_numel(result); + for (int index = 0; index < result_count; ++index) { + result->storage->data[result->offset + index] = 0.0f; + } + + int suffix_matches = is_contiguous((tensor*)source); + for (int axis = 0; axis < target->ndim && suffix_matches; ++axis) { + int source_axis = source->ndim - target->ndim + axis; + suffix_matches = target->dims[axis] == source->dims[source_axis]; + } + if (suffix_matches && result_count > 0) { + int outer_count = tensor_numel((tensor*)source) / result_count; + const float* values = source->storage->data + source->offset; + float* destination = result->storage->data + result->offset; + for (int outer = 0; outer < outer_count; ++outer) { + for (int index = 0; index < result_count; ++index) { + destination[index] += scale * values[outer * result_count + index]; + } + } + return result; + } + + int* coords = source->ndim > 0 + ? (int*)calloc((size_t)source->ndim, sizeof(*coords)) : NULL; + if (source->ndim > 0 && coords == NULL) { + t_free(result); + return NULL; + } + int source_count = tensor_numel((tensor*)source); + for (int index = 0; index < source_count; ++index) { + int result_index = result->offset; + for (int axis = 0; axis < target->ndim; ++axis) { + int source_axis = source->ndim - target->ndim + axis; + int coordinate = target->dims[axis] == 1 ? 0 : coords[source_axis]; + result_index += coordinate * result->strides[axis]; + } + result->storage->data[result_index] += scale * + source->storage->data[get_flat_index_nd((tensor*)source, coords)]; + advance_coords(coords, source->dims, source->ndim); + } + free(coords); + return result; +} + static tensor* reduce_to_shape(tensor* contribution, const tensor* target, int* coords, int coord_capacity) { if (!tensor_has_valid_metadata(contribution) || !tensor_has_valid_metadata(target) || diff --git a/src/autograd/autograd_ops.c b/src/autograd/autograd_ops.c index 58094b1..5e20218 100644 --- a/src/autograd/autograd_ops.c +++ b/src/autograd/autograd_ops.c @@ -20,7 +20,8 @@ static int backward_add(const ag_node* node, !tensor_has_valid_metadata(output_gradient) || input_gradients == NULL) return 1; for (int i = 0; i < 2; ++i) { if (node->inputs[i]->requires_grad) { - input_gradients[i] = t_clone((tensor*)output_gradient); + input_gradients[i] = ag_sum_to_shape( + output_gradient, node->inputs[i]->value, 1.0f); if (input_gradients[i] == NULL) { free_gradients(input_gradients, 2); return 1; @@ -36,11 +37,13 @@ static int backward_sub(const ag_node* node, if (node == NULL || node->input_count != 2 || !tensor_has_valid_metadata(output_gradient) || input_gradients == NULL) return 1; if (node->inputs[0]->requires_grad) { - input_gradients[0] = t_clone((tensor*)output_gradient); + input_gradients[0] = ag_sum_to_shape( + output_gradient, node->inputs[0]->value, 1.0f); if (input_gradients[0] == NULL) return 1; } if (node->inputs[1]->requires_grad) { - input_gradients[1] = t_neg((tensor*)output_gradient); + input_gradients[1] = ag_sum_to_shape( + output_gradient, node->inputs[1]->value, -1.0f); if (input_gradients[1] == NULL) { free_gradients(input_gradients, 2); return 1; diff --git a/tests/unit/autograd/test_autograd_ops.c b/tests/unit/autograd/test_autograd_ops.c index c95a832..d250322 100644 --- a/tests/unit/autograd/test_autograd_ops.c +++ b/tests/unit/autograd/test_autograd_ops.c @@ -94,6 +94,40 @@ TEST(test_broadcast_forward_and_local_gradient_use_output_shape) { ag_tensor_release(result); ag_tensor_release(b); ag_tensor_release(a); } +TEST(test_add_sub_backward_reduce_broadcast_at_source) { + int matrix_dims[2] = {2, 3}; + int bias_dims[1] = {3}; + int owner_dims[2] = {3, 2}; + float matrix_values[6] = {1, 2, 3, 4, 5, 6}; + float bias_values[3] = {0.5f, 1.0f, 1.5f}; + ag_tensor* matrix = make_ag(2, matrix_dims, matrix_values, 1); + ag_tensor* bias = make_ag(1, bias_dims, bias_values, 1); + tensor* owner = t_alloc(2, owner_dims); + for (int index = 0; index < 6; ++index) { + owner->storage->data[index] = (float)(index + 1); + } + tensor* upstream = t_transpose(owner, 0, 1); + tensor* gradients[2] = {NULL, NULL}; + ag_tensor* added = ag_add(matrix, bias); + ASSERT_EQ_INT(added->creator->backward( + added->creator, upstream, gradients), 0); + ASSERT_EQ_INT(gradients[1]->ndim, 1); + ASSERT_EQ_FLOAT(gradients[1]->storage->data[0], 3.0f); + ASSERT_EQ_FLOAT(gradients[1]->storage->data[1], 7.0f); + ASSERT_EQ_FLOAT(gradients[1]->storage->data[2], 11.0f); + free_local_gradients(gradients, 2); + ag_tensor* subtracted = ag_sub(matrix, bias); + ASSERT_EQ_INT(subtracted->creator->backward( + subtracted->creator, upstream, gradients), 0); + ASSERT_EQ_FLOAT(gradients[1]->storage->data[0], -3.0f); + ASSERT_EQ_FLOAT(gradients[1]->storage->data[1], -7.0f); + ASSERT_EQ_FLOAT(gradients[1]->storage->data[2], -11.0f); + free_local_gradients(gradients, 2); + ag_tensor_release(subtracted); ag_tensor_release(added); + t_free(upstream); t_free(owner); + ag_tensor_release(bias); ag_tensor_release(matrix); +} + TEST(test_binary_rejects_null_and_incompatible_shapes) { int adims[1] = {2}, bdims[1] = {3}; float av[2] = {1,2}, bv[3] = {1,2,3}; @@ -669,6 +703,7 @@ int main(void) { RUN_TEST(test_binary_without_grad_omits_graph); RUN_TEST(test_binary_local_gradients_match_derivatives); RUN_TEST(test_broadcast_forward_and_local_gradient_use_output_shape); + RUN_TEST(test_add_sub_backward_reduce_broadcast_at_source); RUN_TEST(test_binary_rejects_null_and_incompatible_shapes); RUN_TEST(test_scalar_arithmetic_forward_backward_and_lifecycle); RUN_TEST(test_scalar_arithmetic_views_ieee_and_finite_difference); From 92b4c7b63b224ac8c4a53ec07475de005a364139 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:52:56 +0200 Subject: [PATCH 41/49] perf: add direct batched RHS-gradient matmul --- src/tensor/tensor_matmul.c | 87 ++++++++++++++++++++++++++ src/tensor/tensor_matmul_internal.h | 10 +++ tests/unit/tensor/test_tensor_matmul.c | 47 ++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 src/tensor/tensor_matmul_internal.h diff --git a/src/tensor/tensor_matmul.c b/src/tensor/tensor_matmul.c index 1c39acd..7dd035e 100644 --- a/src/tensor/tensor_matmul.c +++ b/src/tensor/tensor_matmul.c @@ -21,6 +21,7 @@ #include "../../include/tensorlib/tensor_matmul.h" #include "parallel.h" +#include "tensor_matmul_internal.h" enum { TENSORLIB_MATMUL_MC = 64, @@ -925,6 +926,92 @@ int broadcast_batch_shape(const tensor* a, return 1; } +tensor* tensor_matmul_backward_rhs(const tensor* lhs, + const tensor* output_gradient, + const tensor* rhs) +{ + if (!tensor_has_valid_metadata(lhs) || + !tensor_has_valid_metadata(output_gradient) || + !tensor_has_valid_metadata(rhs) || lhs->ndim == 0 || rhs->ndim == 0) { + return NULL; + } + matmul_operand_info lhs_info = { + lhs->ndim == 1, lhs->ndim > 1 ? lhs->ndim - 2 : 0, + lhs->ndim == 1 ? 1 : lhs->dims[lhs->ndim - 2], + lhs->dims[lhs->ndim - 1], 0 + }; + matmul_operand_info rhs_info = { + rhs->ndim == 1, rhs->ndim > 1 ? rhs->ndim - 2 : 0, 0, + rhs->ndim == 1 ? rhs->dims[0] : rhs->dims[rhs->ndim - 2], + rhs->ndim == 1 ? 1 : rhs->dims[rhs->ndim - 1] + }; + if (lhs_info.inner != rhs_info.inner) return NULL; + int batch_ndim = lhs_info.batch_rank > rhs_info.batch_rank + ? lhs_info.batch_rank : rhs_info.batch_rank; + if (batch_ndim > TENSORLIB_MATMUL_MAX_NDIM) return NULL; + int batch_dims[TENSORLIB_MATMUL_MAX_NDIM]; + if (!broadcast_batch_shape(lhs, &lhs_info, rhs, &rhs_info, + batch_ndim, batch_dims)) return NULL; + int expected_ndim = batch_ndim + 2 - lhs_info.is_vector - rhs_info.is_vector; + if (output_gradient->ndim != expected_ndim) return NULL; + for (int axis = 0; axis < batch_ndim; ++axis) { + if (output_gradient->dims[axis] != batch_dims[axis]) return NULL; + } + int output_axis = batch_ndim; + if (!lhs_info.is_vector && + output_gradient->dims[output_axis++] != lhs_info.rows) return NULL; + if (!rhs_info.is_vector && + output_gradient->dims[output_axis] != rhs_info.columns) return NULL; + + tensor* result = t_alloc(rhs->ndim, rhs->dims); + if (result == NULL) return NULL; + for (int index = 0; index < tensor_numel(result); ++index) { + result->storage->data[index] = 0.0f; + } + size_t batch_count; + if (!tensor_checked_numel(batch_ndim, batch_dims, &batch_count)) { + t_free(result); + return NULL; + } + int coords[TENSORLIB_MATMUL_MAX_NDIM]; + for (size_t batch = 0; batch < batch_count; ++batch) { + batch_index_to_coords(batch, batch_dims, batch_ndim, coords); + int lhs_base = operand_batch_offset(lhs, &lhs_info, coords, batch_ndim); + int result_base = operand_batch_offset(result, &rhs_info, + coords, batch_ndim); + int gradient_base = output_batch_offset(output_gradient, coords, + batch_ndim); + for (int inner = 0; inner < lhs_info.inner; ++inner) { + for (int column = 0; column < rhs_info.columns; ++column) { + float sum = 0.0f; + for (int row = 0; row < lhs_info.rows; ++row) { + int lhs_index = lhs_info.is_vector + ? lhs->offset + inner * lhs->strides[0] + : lhs_base + row * lhs->strides[lhs->ndim - 2] + + inner * lhs->strides[lhs->ndim - 1]; + int gradient_index = gradient_base; + if (!lhs_info.is_vector) { + gradient_index += row * output_gradient->strides[batch_ndim]; + } + if (!rhs_info.is_vector) { + int column_axis = batch_ndim + !lhs_info.is_vector; + gradient_index += column * + output_gradient->strides[column_axis]; + } + sum += lhs->storage->data[lhs_index] * + output_gradient->storage->data[gradient_index]; + } + int result_index = rhs_info.is_vector + ? result_base + inner * result->strides[0] + : result_base + inner * result->strides[result->ndim - 2] + + column * result->strides[result->ndim - 1]; + result->storage->data[result_index] += sum; + } + } + } + return result; +} + void matmul_2d_strided(const tensor* a, const matmul_operand_info* a_info, int a_base, diff --git a/src/tensor/tensor_matmul_internal.h b/src/tensor/tensor_matmul_internal.h new file mode 100644 index 0000000..5b8cd73 --- /dev/null +++ b/src/tensor/tensor_matmul_internal.h @@ -0,0 +1,10 @@ +#ifndef TENSORLIB_TENSOR_MATMUL_INTERNAL_H +#define TENSORLIB_TENSOR_MATMUL_INTERNAL_H + +#include "../../include/tensorlib/tensor.h" + +tensor* tensor_matmul_backward_rhs(const tensor* lhs, + const tensor* output_gradient, + const tensor* rhs); + +#endif diff --git a/tests/unit/tensor/test_tensor_matmul.c b/tests/unit/tensor/test_tensor_matmul.c index 0995304..a31bbbd 100644 --- a/tests/unit/tensor/test_tensor_matmul.c +++ b/tests/unit/tensor/test_tensor_matmul.c @@ -2,6 +2,7 @@ #include "../../fixtures/test_common.h" #include "../../../include/tensorlib/tensor.h" +#include "../../../src/tensor/tensor_matmul_internal.h" static void fill_tensor(tensor* t, const float* values, int count) { for (int i = 0; i < count; ++i) t->storage->data[i] = values[i]; @@ -537,6 +538,51 @@ TEST(test_t_matmul_allows_aliasing_and_returns_independent_storage) { t_free(a); } +TEST(test_direct_rhs_gradient_handles_broadcast_strides_and_vectors) { + int owner_dims[3] = {2, 3, 2}; + int gradient_dims[3] = {2, 2, 5}; + int rhs_dims[2] = {3, 5}; + tensor* owner = t_alloc(3, owner_dims); + tensor* lhs = t_transpose(owner, 1, 2); + tensor* gradient = t_alloc(3, gradient_dims); + tensor* rhs = t_alloc(2, rhs_dims); + for (int index = 0; index < tensor_numel(owner); ++index) { + owner->storage->data[index] = 0.1f * (float)(index + 1); + } + for (int index = 0; index < tensor_numel(gradient); ++index) { + gradient->storage->data[index] = 0.01f * (float)(index - 7); + } + tensor* lhs_transpose = t_transpose(lhs, 1, 2); + tensor* batched = t_matmul(lhs_transpose, gradient); + tensor* expected = t_sum(batched, 0); + tensor* actual = tensor_matmul_backward_rhs(lhs, gradient, rhs); + assert_same_tensor(actual, expected); + t_free(actual); t_free(expected); t_free(batched); t_free(lhs_transpose); + t_free(rhs); t_free(gradient); t_free(lhs); t_free(owner); + + int vector_dims[1] = {3}; + int matrix_dims[2] = {3, 2}; + int output_dims[1] = {2}; + float lhs_values[3] = {1, 2, 3}; + float gradient_values[2] = {4, 5}; + tensor* lhs_vector = t_alloc(1, vector_dims); + tensor* rhs_matrix = t_alloc(2, matrix_dims); + tensor* vector_gradient = t_alloc(1, output_dims); + fill_tensor(lhs_vector, lhs_values, 3); + fill_tensor(vector_gradient, gradient_values, 2); + actual = tensor_matmul_backward_rhs(lhs_vector, vector_gradient, rhs_matrix); + ASSERT_NOT_NULL(actual); + for (int inner = 0; inner < 3; ++inner) { + for (int column = 0; column < 2; ++column) { + ASSERT_EQ_FLOAT(actual->storage->data[inner * 2 + column], + lhs_values[inner] * gradient_values[column]); + } + } + ASSERT_NULL(tensor_matmul_backward_rhs(lhs_vector, lhs_vector, rhs_matrix)); + t_free(actual); t_free(vector_gradient); t_free(rhs_matrix); + t_free(lhs_vector); +} + int main(void) { printf("== tensor_matmul.c ==\n"); RUN_TEST(test_t_matmul_2d); @@ -556,5 +602,6 @@ int main(void) { RUN_TEST(test_t_matmul_rejects_invalid_shapes); RUN_TEST(test_t_matmul_validates_vector_inner_dimension); RUN_TEST(test_t_matmul_allows_aliasing_and_returns_independent_storage); + RUN_TEST(test_direct_rhs_gradient_handles_broadcast_strides_and_vectors); TEST_SUITE_SUMMARY(); } From 3095f228752b22b5e410d80d51742b595a83f727 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:54:13 +0200 Subject: [PATCH 42/49] perf: vectorize linear weight-gradient GEMM --- benchmarks/bench_backward.c | 9 +++++++++ src/tensor/tensor_matmul.c | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/benchmarks/bench_backward.c b/benchmarks/bench_backward.c index 22f5c3a..0742bd0 100644 --- a/benchmarks/bench_backward.c +++ b/benchmarks/bench_backward.c @@ -4,10 +4,12 @@ #include #include +#include "../src/tensor/tensor_matmul_internal.h" typedef enum { LINEAR_DINPUT, LINEAR_DWEIGHT, + LINEAR_DWEIGHT_DIRECT, LINEAR_DBIAS, LINEAR_PACK } linear_gradient_kind; @@ -89,6 +91,10 @@ static int linear_gradient_operation(void* opaque, double* checksum) second = first == NULL ? NULL : t_matmul(first, context->output_gradient); output = second == NULL ? NULL : t_sum(second, 0); + } else if (context->kind == LINEAR_DWEIGHT_DIRECT) { + first = t_transpose(context->weight, 0, 1); + output = first == NULL ? NULL : tensor_matmul_backward_rhs( + context->input, context->output_gradient, first); } else if (context->kind == LINEAR_DBIAS) { first = t_sum(context->output_gradient, 0); output = first == NULL ? NULL : t_sum(first, 0); @@ -244,6 +250,9 @@ int bench_run_backward_matrix_suite(const bench_options* options, FILE* csv) status |= run_linear_gradient(options, csv, &shapes[index], LINEAR_DWEIGHT, "dweight", "batched-matmul+reduce"); + status |= run_linear_gradient(options, csv, &shapes[index], + LINEAR_DWEIGHT_DIRECT, "dweight_direct", + "flattened-transpose-left"); status |= run_linear_gradient(options, csv, &shapes[index], LINEAR_DBIAS, "dbias", "two-axis-reduce"); diff --git a/src/tensor/tensor_matmul.c b/src/tensor/tensor_matmul.c index 7dd035e..fedac1a 100644 --- a/src/tensor/tensor_matmul.c +++ b/src/tensor/tensor_matmul.c @@ -926,6 +926,35 @@ int broadcast_batch_shape(const tensor* a, return 1; } +static tensor* try_contiguous_linear_rhs_gradient( + const tensor* lhs, const tensor* output_gradient, const tensor* rhs) +{ + if (rhs->ndim != 2 || lhs->ndim < 2 || + lhs->ndim != output_gradient->ndim || + !is_contiguous((tensor*)lhs) || + !is_contiguous((tensor*)output_gradient)) return NULL; + for (int axis = 0; axis < lhs->ndim - 1; ++axis) { + if (lhs->dims[axis] != output_gradient->dims[axis]) return NULL; + } + int inner = lhs->dims[lhs->ndim - 1]; + int columns = output_gradient->dims[output_gradient->ndim - 1]; + if (rhs->dims[0] != inner || rhs->dims[1] != columns) return NULL; + int rows = tensor_numel((tensor*)lhs) / inner; + int lhs_dims[2] = {rows, inner}; + int gradient_dims[2] = {rows, columns}; + tensor* lhs_matrix = t_reshape((tensor*)lhs, 2, lhs_dims); + tensor* gradient_matrix = + t_reshape((tensor*)output_gradient, 2, gradient_dims); + tensor* lhs_transpose = lhs_matrix == NULL ? NULL : + t_transpose(lhs_matrix, 0, 1); + tensor* result = lhs_transpose == NULL || gradient_matrix == NULL ? NULL : + t_matmul(lhs_transpose, gradient_matrix); + t_free(lhs_transpose); + t_free(gradient_matrix); + t_free(lhs_matrix); + return result; +} + tensor* tensor_matmul_backward_rhs(const tensor* lhs, const tensor* output_gradient, const tensor* rhs) @@ -935,6 +964,9 @@ tensor* tensor_matmul_backward_rhs(const tensor* lhs, !tensor_has_valid_metadata(rhs) || lhs->ndim == 0 || rhs->ndim == 0) { return NULL; } + tensor* linear_result = + try_contiguous_linear_rhs_gradient(lhs, output_gradient, rhs); + if (linear_result != NULL) return linear_result; matmul_operand_info lhs_info = { lhs->ndim == 1, lhs->ndim > 1 ? lhs->ndim - 2 : 0, lhs->ndim == 1 ? 1 : lhs->dims[lhs->ndim - 2], From 07b5ec8f20121826b6b71e6473b27a28092effac Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:55:19 +0200 Subject: [PATCH 43/49] perf: verify deterministic weight-gradient parallelism --- benchmarks/bench_backward.c | 23 ++++++++++------ tests/unit/tensor/test_tensor_matmul.c | 38 ++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 9 deletions(-) diff --git a/benchmarks/bench_backward.c b/benchmarks/bench_backward.c index 0742bd0..4b869c3 100644 --- a/benchmarks/bench_backward.c +++ b/benchmarks/bench_backward.c @@ -123,7 +123,8 @@ static int run_linear_gradient(const bench_options* options, const linear_shape* shape, linear_gradient_kind kind, const char* suffix, - const char* layout) + const char* layout, + int threads) { linear_gradient_context context; bench_measurement result; @@ -142,7 +143,7 @@ static int run_linear_gradient(const bench_options* options, linear_gradient_operation, &context, NULL }; int status = bench_execute_case(options, csv, &benchmark, - options->threads[0], &result); + threads, &result); destroy_linear_context(&context); return status == 1; } @@ -246,18 +247,22 @@ int bench_run_backward_matrix_suite(const bench_options* options, FILE* csv) for (size_t index = 0; index < sizeof(shapes) / sizeof(shapes[0]); ++index) { status |= run_linear_gradient(options, csv, &shapes[index], LINEAR_DINPUT, "dinput", - "packed-weight"); + "packed-weight", options->threads[0]); status |= run_linear_gradient(options, csv, &shapes[index], LINEAR_DWEIGHT, "dweight", - "batched-matmul+reduce"); - status |= run_linear_gradient(options, csv, &shapes[index], - LINEAR_DWEIGHT_DIRECT, "dweight_direct", - "flattened-transpose-left"); + "batched-matmul+reduce", options->threads[0]); + for (int thread_index = 0; thread_index < options->thread_count; + ++thread_index) { + status |= run_linear_gradient(options, csv, &shapes[index], + LINEAR_DWEIGHT_DIRECT, "dweight_direct", + "flattened-transpose-left", options->threads[thread_index]); + } status |= run_linear_gradient(options, csv, &shapes[index], LINEAR_DBIAS, "dbias", - "two-axis-reduce"); + "two-axis-reduce", options->threads[0]); status |= run_linear_gradient(options, csv, &shapes[index], - LINEAR_PACK, "pack", "packing-only"); + LINEAR_PACK, "pack", "packing-only", + options->threads[0]); } status |= run_attention_gradient(options, csv, 0); status |= run_attention_gradient(options, csv, 1); diff --git a/tests/unit/tensor/test_tensor_matmul.c b/tests/unit/tensor/test_tensor_matmul.c index a31bbbd..a9dd403 100644 --- a/tests/unit/tensor/test_tensor_matmul.c +++ b/tests/unit/tensor/test_tensor_matmul.c @@ -4,6 +4,10 @@ #include "../../../include/tensorlib/tensor.h" #include "../../../src/tensor/tensor_matmul_internal.h" +#ifdef _OPENMP +#include +#endif + static void fill_tensor(tensor* t, const float* values, int count) { for (int i = 0; i < count; ++i) t->storage->data[i] = values[i]; } @@ -583,6 +587,39 @@ TEST(test_direct_rhs_gradient_handles_broadcast_strides_and_vectors) { t_free(lhs_vector); } +TEST(test_direct_rhs_gradient_is_deterministic_across_threads) { +#ifdef _OPENMP + int lhs_dims[2] = {128, 192}; + int gradient_dims[2] = {128, 192}; + int rhs_dims[2] = {192, 192}; + tensor* lhs = t_alloc(2, lhs_dims); + tensor* gradient = t_alloc(2, gradient_dims); + tensor* rhs = t_alloc(2, rhs_dims); + for (int index = 0; index < tensor_numel(lhs); ++index) { + lhs->storage->data[index] = 0.001f * (float)((index % 17) - 8); + } + for (int index = 0; index < tensor_numel(gradient); ++index) { + gradient->storage->data[index] = 0.002f * (float)((index % 13) - 6); + } + omp_set_dynamic(0); + omp_set_num_threads(1); + tensor* serial = tensor_matmul_backward_rhs(lhs, gradient, rhs); + omp_set_num_threads(4); + tensor* parallel_a = tensor_matmul_backward_rhs(lhs, gradient, rhs); + tensor* parallel_b = tensor_matmul_backward_rhs(lhs, gradient, rhs); + ASSERT_NOT_NULL(serial); ASSERT_NOT_NULL(parallel_a); ASSERT_NOT_NULL(parallel_b); + for (int index = 0; index < tensor_numel(serial); ++index) { + ASSERT_EQ_FLOAT(parallel_a->storage->data[index], + parallel_b->storage->data[index]); + ASSERT_FLOAT_NEAR(serial->storage->data[index], + parallel_a->storage->data[index], 1e-6f); + } + omp_set_num_threads(1); + t_free(parallel_b); t_free(parallel_a); t_free(serial); + t_free(rhs); t_free(gradient); t_free(lhs); +#endif +} + int main(void) { printf("== tensor_matmul.c ==\n"); RUN_TEST(test_t_matmul_2d); @@ -603,5 +640,6 @@ int main(void) { RUN_TEST(test_t_matmul_validates_vector_inner_dimension); RUN_TEST(test_t_matmul_allows_aliasing_and_returns_independent_storage); RUN_TEST(test_direct_rhs_gradient_handles_broadcast_strides_and_vectors); + RUN_TEST(test_direct_rhs_gradient_is_deterministic_across_threads); TEST_SUITE_SUMMARY(); } From 682b614d3c181427eaaeea48d0217353aee173ec Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 16:59:01 +0200 Subject: [PATCH 44/49] perf: route linear backward through direct RHS gradients --- src/autograd/autograd_matmul.c | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/autograd/autograd_matmul.c b/src/autograd/autograd_matmul.c index 31ab72c..9a271d1 100644 --- a/src/autograd/autograd_matmul.c +++ b/src/autograd/autograd_matmul.c @@ -1,6 +1,7 @@ #include #include "./../../include/tensorlib/autograd_internal.h" +#include "../tensor/tensor_matmul_internal.h" typedef struct { tensor_matmul_packed_rhs* backward_rhs; @@ -73,7 +74,9 @@ static int backward_matmul(const ag_node* node, } } if (node->inputs[1]->requires_grad) { - input_gradients[1] = t_matmul(transposed_a, promoted_gradient); + input_gradients[1] = b->ndim == 2 + ? tensor_matmul_backward_rhs(a, output_gradient, b) + : t_matmul(transposed_a, promoted_gradient); if (input_gradients[1] == NULL) goto fail; if (b_vector) { input_gradients[1] = remove_vector_dimension(input_gradients[1], 0); From 9064b4bbcf64d047db45f38c3e3f2be4e9e5c7e0 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 17:01:05 +0200 Subject: [PATCH 45/49] perf: reuse packed matmul workspaces --- src/tensor/tensor_matmul.c | 35 ++++++++++++++++++++++------------- 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/src/tensor/tensor_matmul.c b/src/tensor/tensor_matmul.c index fedac1a..f74dd9c 100644 --- a/src/tensor/tensor_matmul.c +++ b/src/tensor/tensor_matmul.c @@ -61,7 +61,8 @@ static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, const float* packed_rhs, int inner, int columns, tensor* output, int output_base, int output_batch_ndim, - int row_start, int row_end); + int row_start, int row_end, + float* packed_lhs); #endif static void* matmul_aligned_malloc(size_t bytes) { @@ -295,7 +296,8 @@ static void matmul_packed_rhs_batch(const tensor* lhs, int batch_ndim, size_t batch, int row_start, - int row_end) { + int row_end, + float* workspace) { int coords[TENSORLIB_MATMUL_MAX_NDIM]; batch_index_to_coords(batch, batch_dims, batch_ndim, coords); @@ -312,7 +314,7 @@ static void matmul_packed_rhs_batch(const tensor* lhs, matmul_2d_packed_rhs_avx2(lhs, lhs_base, packed_rhs, lhs_info->inner, rhs->columns, output, output_base, - batch_ndim, row_start, row_end)) { + batch_ndim, row_start, row_end, workspace)) { /* The AVX2 kernel completed this batch. */ } else { matmul_2d_packed_rhs_scalar(lhs, lhs_info, lhs_base, @@ -393,6 +395,10 @@ tensor* t_matmul_packed_rhs(const tensor* lhs, int threads = tensorlib_parallel_threads( batch_flops, TENSORLIB_MATMUL_MIN_PARALLEL_FLOPS, batch_ndim <= TENSORLIB_MATMUL_MAX_NDIM ? task_count : 1); + size_t workspace_values = (size_t)TENSORLIB_MATMUL_MC * + (size_t)TENSORLIB_MATMUL_KC; + float* workspaces = (float*)matmul_aligned_malloc( + (size_t)threads * workspace_values * sizeof(float)); if (threads > 1) { #ifdef _OPENMP @@ -404,16 +410,25 @@ tensor* t_matmul_packed_rhs(const tensor* lhs, TENSORLIB_MATMUL_MC; int row_end = row_start + TENSORLIB_MATMUL_MC; if (row_end > lhs_info.rows) row_end = lhs_info.rows; + int thread_index = 0; +#ifdef _OPENMP + thread_index = omp_get_thread_num(); +#endif + float* workspace = workspaces == NULL ? NULL : workspaces + + (size_t)thread_index * workspace_values; matmul_packed_rhs_batch(lhs, &lhs_info, rhs, output, batch_dims, - batch_ndim, batch, row_start, row_end); + batch_ndim, batch, row_start, row_end, + workspace); } } else { for (size_t batch = 0; batch < batch_count; ++batch) { matmul_packed_rhs_batch(lhs, &lhs_info, rhs, output, batch_dims, - batch_ndim, batch, 0, lhs_info.rows); + batch_ndim, batch, 0, lhs_info.rows, + workspaces); } } + matmul_aligned_free(workspaces); free(batch_dims); return output; } @@ -806,13 +821,8 @@ static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, const float* packed_rhs, int inner, int columns, tensor* output, int output_base, int output_batch_ndim, - int row_start, int row_end) { - size_t workspace_values; - if (!checked_size_multiply(TENSORLIB_MATMUL_MC, TENSORLIB_MATMUL_KC, - &workspace_values)) { - return 0; - } - float* packed_lhs = (float*)matmul_aligned_malloc(workspace_values * sizeof(float)); + int row_start, int row_end, + float* packed_lhs) { if (packed_lhs == NULL) return 0; int row_count = row_end - row_start; @@ -863,7 +873,6 @@ static int matmul_2d_packed_rhs_avx2(const tensor* lhs, int lhs_base, output, output_base + row_start * output_stride, output_batch_ndim); - matmul_aligned_free(packed_lhs); return 1; } #else From 7a89fc8977769e2df233cd02d688a347c3f1a9df Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 17:03:28 +0200 Subject: [PATCH 46/49] perf: specialize contiguous slice backward --- src/autograd/autograd_view.c | 26 +++++++++++++++++++++++- tests/unit/autograd/test_autograd_view.c | 25 +++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/src/autograd/autograd_view.c b/src/autograd/autograd_view.c index ec6dfc2..1d9211e 100644 --- a/src/autograd/autograd_view.c +++ b/src/autograd/autograd_view.c @@ -1,4 +1,5 @@ #include +#include #include "./../../include/tensorlib/autograd_internal.h" @@ -44,9 +45,32 @@ static int backward_slice(const ag_node* node, tensor* input = node->inputs[0]->value; slice_context* context = (slice_context*)node->context; tensor* gradient = ag_full_like(input, 0.0f); + if (gradient == NULL) return 1; + if (is_contiguous((tensor*)output_gradient)) { + size_t outer = 1; + size_t inner = 1; + for (int axis = 0; axis < context->dim; ++axis) { + outer *= (size_t)input->dims[axis]; + } + for (int axis = context->dim + 1; axis < input->ndim; ++axis) { + inner *= (size_t)input->dims[axis]; + } + size_t slice = (size_t)output_gradient->dims[context->dim]; + for (size_t block = 0; block < outer; ++block) { + size_t source = (size_t)output_gradient->offset + + block * slice * inner; + size_t destination = block * (size_t)input->dims[context->dim] * + inner + (size_t)context->start * inner; + memcpy(gradient->storage->data + destination, + output_gradient->storage->data + source, + slice * inner * sizeof(float)); + } + input_gradients[0] = gradient; + return 0; + } int* output_coords = (int*)calloc((size_t)output_gradient->ndim, sizeof(int)); int* input_coords = (int*)calloc((size_t)input->ndim, sizeof(int)); - if (gradient == NULL || output_coords == NULL || input_coords == NULL) { + if (output_coords == NULL || input_coords == NULL) { t_free(gradient); free(output_coords); free(input_coords); return 1; } diff --git a/tests/unit/autograd/test_autograd_view.c b/tests/unit/autograd/test_autograd_view.c index 19c24f3..90f73da 100644 --- a/tests/unit/autograd/test_autograd_view.c +++ b/tests/unit/autograd/test_autograd_view.c @@ -57,6 +57,30 @@ TEST(test_slice_backward_scatters_and_zero_fills) { t_free(gradients[0]); t_free(upstream); ag_tensor_release(output); ag_tensor_release(input); } +TEST(test_slice_backward_copies_contiguous_middle_axis_blocks) { + int dims[3] = {2, 4, 3}; + float values[24] = {0}; + ag_tensor* input = make_ag(3, dims, values); + ag_tensor* output = ag_slice(input, 1, 1, 3); + tensor* upstream = sequential_gradient(output->value->ndim, + output->value->dims); + tensor* gradients[1] = {NULL}; + ASSERT_EQ_INT(output->creator->backward(output->creator, upstream, + gradients), 0); + for (int batch = 0; batch < 2; ++batch) { + for (int row = 0; row < 4; ++row) { + for (int column = 0; column < 3; ++column) { + int index = (batch * 4 + row) * 3 + column; + float expected = row < 1 || row >= 3 ? 0.0f : + (float)((batch * 2 + row - 1) * 3 + column + 1); + ASSERT_EQ_FLOAT(gradients[0]->storage->data[index], expected); + } + } + } + t_free(gradients[0]); t_free(upstream); + ag_tensor_release(output); ag_tensor_release(input); +} + TEST(test_expand_backward_emits_output_shaped_contribution) { int input_dims[2] = {1, 3}, output_dims[3] = {2, 4, 3}; float values[3] = {1,2,3}; @@ -86,6 +110,7 @@ int main(void) { RUN_TEST(test_reshape_backward_restores_input_shape); RUN_TEST(test_transpose_backward_applies_inverse_transpose); RUN_TEST(test_slice_backward_scatters_and_zero_fills); + RUN_TEST(test_slice_backward_copies_contiguous_middle_axis_blocks); RUN_TEST(test_expand_backward_emits_output_shaped_contribution); RUN_TEST(test_view_wrappers_reject_invalid_arguments); TEST_SUITE_SUMMARY(); From 43f4358045400ac4d3546e77fe11f807545346d1 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 17:05:11 +0200 Subject: [PATCH 47/49] perf: reuse autograd traversal storage --- src/autograd/autograd_backward.c | 42 ++++++++++++++------------------ 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/src/autograd/autograd_backward.c b/src/autograd/autograd_backward.c index b8b4653..32155ec 100644 --- a/src/autograd/autograd_backward.c +++ b/src/autograd/autograd_backward.c @@ -305,10 +305,9 @@ static int accumulate_pass_gradient(tensor** destination, return 0; } -static int merge_persistent_gradients(const tensor_list* tensors, tensor** pass_gradients) { - tensor** merged = (tensor**)calloc((size_t)tensors->count, sizeof(*merged)); - if (merged == NULL) return 1; - +static int merge_persistent_gradients(const tensor_list* tensors, + tensor** pass_gradients, + tensor** merged) { for (int i = 0; i < tensors->count; ++i) { ag_tensor* value = tensors->values[i]; if (!value->requires_grad || pass_gradients[i] == NULL) continue; @@ -323,7 +322,6 @@ static int merge_persistent_gradients(const tensor_list* tensors, tensor** pass_ } if (merged[i] == NULL) { for (int j = 0; j < tensors->count; ++j) t_free(merged[j]); - free(merged); return 1; } } @@ -334,7 +332,6 @@ static int merge_persistent_gradients(const tensor_list* tensors, tensor** pass_ tensors->values[i]->grad = merged[i]; merged[i] = NULL; } - free(merged); return 0; } @@ -356,8 +353,10 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { tensor_list tensors = {0}; node_list nodes = {0}; tensor** pass_gradients = NULL; + tensor** merged_gradients = NULL; tensor** contributions = NULL; int* reduction_coords = NULL; + void* workspace = NULL; int contribution_capacity = 0; int reduction_coord_capacity = 0; int status = 1; @@ -373,25 +372,21 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { int ndim = tensors.values[index]->value->ndim; if (ndim > reduction_coord_capacity) reduction_coord_capacity = ndim; } - if (reduction_coord_capacity > 0) { - reduction_coords = (int*)calloc((size_t)reduction_coord_capacity, - sizeof(*reduction_coords)); - if (reduction_coords == NULL) goto cleanup; - } - for (int node_index = 0; node_index < nodes.count; ++node_index) { if (nodes.values[node_index]->input_count > contribution_capacity) { contribution_capacity = nodes.values[node_index]->input_count; } } - if (contribution_capacity > 0) { - contributions = (tensor**)calloc((size_t)contribution_capacity, - sizeof(*contributions)); - if (contributions == NULL) goto cleanup; - } - - pass_gradients = (tensor**)calloc((size_t)tensors.count, sizeof(*pass_gradients)); - if (pass_gradients == NULL) goto cleanup; + size_t pointer_count = (size_t)tensors.count * 2u + + (size_t)contribution_capacity; + size_t workspace_bytes = pointer_count * sizeof(tensor*) + + (size_t)reduction_coord_capacity * sizeof(int); + workspace = calloc(1, workspace_bytes); + if (workspace == NULL) goto cleanup; + pass_gradients = (tensor**)workspace; + merged_gradients = pass_gradients + tensors.count; + contributions = merged_gradients + tensors.count; + reduction_coords = (int*)(contributions + contribution_capacity); { int output_index = output->graph_index; @@ -446,19 +441,18 @@ int ag_backward_with_grad(ag_tensor* output, const tensor* output_gradient) { } started = backward_stats_enabled ? backward_now() : 0.0; - status = merge_persistent_gradients(&tensors, pass_gradients); + status = merge_persistent_gradients(&tensors, pass_gradients, + merged_gradients); if (backward_stats_enabled) { backward_stats.merge_seconds += backward_elapsed(started); } cleanup: - free(reduction_coords); free_contributions(contributions, contribution_capacity); - free(contributions); if (pass_gradients != NULL) { for (int i = 0; i < tensors.count; ++i) t_free(pass_gradients[i]); - free(pass_gradients); } + free(workspace); for (int i = 0; i < tensors.count; ++i) tensors.values[i]->graph_index = -1; free(nodes.values); free(tensors.values); From 74770a36e897a6cf48c976b5e5ff24aa475e1074 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 17:19:57 +0200 Subject: [PATCH 48/49] perf: contract projected weight gradients directly --- src/autograd/autograd_matmul.c | 9 ++-- src/tensor/tensor_matmul.c | 67 +++++++++++++++++++++++++- src/tensor/tensor_matmul_internal.h | 3 ++ tests/unit/tensor/test_tensor_matmul.c | 26 ++++++++++ 4 files changed, 100 insertions(+), 5 deletions(-) diff --git a/src/autograd/autograd_matmul.c b/src/autograd/autograd_matmul.c index 9a271d1..41fe5c4 100644 --- a/src/autograd/autograd_matmul.c +++ b/src/autograd/autograd_matmul.c @@ -74,9 +74,12 @@ static int backward_matmul(const ag_node* node, } } if (node->inputs[1]->requires_grad) { - input_gradients[1] = b->ndim == 2 - ? tensor_matmul_backward_rhs(a, output_gradient, b) - : t_matmul(transposed_a, promoted_gradient); + input_gradients[1] = tensor_matmul_backward_rhs_fast(a, + output_gradient, + b); + if (input_gradients[1] == NULL) { + input_gradients[1] = t_matmul(transposed_a, promoted_gradient); + } if (input_gradients[1] == NULL) goto fail; if (b_vector) { input_gradients[1] = remove_vector_dimension(input_gradients[1], 0); diff --git a/src/tensor/tensor_matmul.c b/src/tensor/tensor_matmul.c index f74dd9c..e5c2cbb 100644 --- a/src/tensor/tensor_matmul.c +++ b/src/tensor/tensor_matmul.c @@ -2,6 +2,7 @@ #include #include #include +#include #if defined(_WIN32) #include #endif @@ -964,6 +965,68 @@ static tensor* try_contiguous_linear_rhs_gradient( return result; } +static tensor* try_contiguous_projected_rhs_gradient( + const tensor* lhs, const tensor* output_gradient, const tensor* rhs) +{ + if (lhs->ndim != 4 || rhs->ndim != 4 || output_gradient->ndim != 4 || + lhs->dims[0] != 1 || rhs->dims[1] != 1 || + output_gradient->dims[0] != rhs->dims[0] || + output_gradient->dims[1] != lhs->dims[1] || + output_gradient->dims[2] != lhs->dims[2] || + rhs->dims[2] != lhs->dims[3] || + output_gradient->dims[3] != rhs->dims[3] || + !is_contiguous((tensor*)lhs) || + !is_contiguous((tensor*)output_gradient)) return NULL; + + int inner = lhs->dims[3]; + int columns = rhs->dims[3]; + int rows = tensor_numel((tensor*)lhs) / inner; + int lhs_dims[2] = {rows, inner}; + int gradient_dims[2] = {rows, columns}; + tensor* result = t_alloc(rhs->ndim, rhs->dims); + tensor* lhs_matrix = t_reshape((tensor*)lhs, 2, lhs_dims); + tensor* lhs_transpose = lhs_matrix == NULL ? NULL : + t_transpose(lhs_matrix, 0, 1); + if (result == NULL || lhs_transpose == NULL) goto fail; + + size_t matrix_bytes = (size_t)inner * (size_t)columns * sizeof(float); + for (int projection = 0; projection < rhs->dims[0]; ++projection) { + tensor* gradient_slice = t_slice((tensor*)output_gradient, 0, + projection, projection + 1); + tensor* gradient_matrix = gradient_slice == NULL ? NULL : + t_reshape(gradient_slice, 2, gradient_dims); + tensor* matrix = gradient_matrix == NULL ? NULL : + t_matmul(lhs_transpose, gradient_matrix); + if (matrix == NULL) { + t_free(matrix); t_free(gradient_matrix); t_free(gradient_slice); + goto fail; + } + memcpy(result->storage->data + (size_t)projection * (size_t)inner * + (size_t)columns, + matrix->storage->data + matrix->offset, matrix_bytes); + t_free(matrix); t_free(gradient_matrix); t_free(gradient_slice); + } + t_free(lhs_transpose); t_free(lhs_matrix); + return result; + +fail: + t_free(lhs_transpose); t_free(lhs_matrix); t_free(result); + return NULL; +} + +tensor* tensor_matmul_backward_rhs_fast(const tensor* lhs, + const tensor* output_gradient, + const tensor* rhs) +{ + if (!tensor_has_valid_metadata(lhs) || + !tensor_has_valid_metadata(output_gradient) || + !tensor_has_valid_metadata(rhs)) return NULL; + tensor* result = try_contiguous_linear_rhs_gradient(lhs, output_gradient, + rhs); + return result != NULL ? result : + try_contiguous_projected_rhs_gradient(lhs, output_gradient, rhs); +} + tensor* tensor_matmul_backward_rhs(const tensor* lhs, const tensor* output_gradient, const tensor* rhs) @@ -973,8 +1036,8 @@ tensor* tensor_matmul_backward_rhs(const tensor* lhs, !tensor_has_valid_metadata(rhs) || lhs->ndim == 0 || rhs->ndim == 0) { return NULL; } - tensor* linear_result = - try_contiguous_linear_rhs_gradient(lhs, output_gradient, rhs); + tensor* linear_result = tensor_matmul_backward_rhs_fast( + lhs, output_gradient, rhs); if (linear_result != NULL) return linear_result; matmul_operand_info lhs_info = { lhs->ndim == 1, lhs->ndim > 1 ? lhs->ndim - 2 : 0, diff --git a/src/tensor/tensor_matmul_internal.h b/src/tensor/tensor_matmul_internal.h index 5b8cd73..c81fe6c 100644 --- a/src/tensor/tensor_matmul_internal.h +++ b/src/tensor/tensor_matmul_internal.h @@ -6,5 +6,8 @@ tensor* tensor_matmul_backward_rhs(const tensor* lhs, const tensor* output_gradient, const tensor* rhs); +tensor* tensor_matmul_backward_rhs_fast(const tensor* lhs, + const tensor* output_gradient, + const tensor* rhs); #endif diff --git a/tests/unit/tensor/test_tensor_matmul.c b/tests/unit/tensor/test_tensor_matmul.c index a9dd403..d47f0fe 100644 --- a/tests/unit/tensor/test_tensor_matmul.c +++ b/tests/unit/tensor/test_tensor_matmul.c @@ -620,6 +620,31 @@ TEST(test_direct_rhs_gradient_is_deterministic_across_threads) { #endif } +TEST(test_direct_rhs_gradient_contracts_projected_batch) { + int lhs_dims[4] = {1, 2, 3, 4}; + int rhs_dims[4] = {3, 1, 4, 5}; + int gradient_dims[4] = {3, 2, 3, 5}; + tensor* lhs = t_alloc(4, lhs_dims); + tensor* rhs = t_alloc(4, rhs_dims); + tensor* gradient = t_alloc(4, gradient_dims); + for (int i = 0; i < tensor_numel(lhs); ++i) { + lhs->storage->data[i] = 0.1f * (float)((i % 7) - 3); + } + for (int i = 0; i < tensor_numel(gradient); ++i) { + gradient->storage->data[i] = 0.05f * (float)((i % 11) - 5); + } + tensor* lhs_transpose = t_transpose(lhs, 2, 3); + tensor* batched = t_matmul(lhs_transpose, gradient); + tensor* reduced = t_sum(batched, 1); + tensor* expected = t_unsqueeze(reduced, 1); + tensor* actual = tensor_matmul_backward_rhs_fast(lhs, gradient, rhs); + ASSERT_NOT_NULL(actual); ASSERT_NOT_NULL(expected); + assert_same_tensor(actual, expected); + t_free(actual); t_free(expected); t_free(reduced); t_free(batched); + t_free(lhs_transpose); + t_free(gradient); t_free(rhs); t_free(lhs); +} + int main(void) { printf("== tensor_matmul.c ==\n"); RUN_TEST(test_t_matmul_2d); @@ -641,5 +666,6 @@ int main(void) { RUN_TEST(test_t_matmul_allows_aliasing_and_returns_independent_storage); RUN_TEST(test_direct_rhs_gradient_handles_broadcast_strides_and_vectors); RUN_TEST(test_direct_rhs_gradient_is_deterministic_across_threads); + RUN_TEST(test_direct_rhs_gradient_contracts_projected_batch); TEST_SUITE_SUMMARY(); } From 7c07ccb8d3060d5a4beb0703fe02a0a1cc2f75e1 Mon Sep 17 00:00:00 2001 From: Nisbenz Date: Sun, 6 Sep 2026 17:34:12 +0200 Subject: [PATCH 49/49] test: require zero live benchmark storage --- benchmarks/bench_nn.c | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/benchmarks/bench_nn.c b/benchmarks/bench_nn.c index 4fbf786..79b0a3f 100644 --- a/benchmarks/bench_nn.c +++ b/benchmarks/bench_nn.c @@ -554,8 +554,12 @@ static int run_decoder_case(const bench_options* options, &context); } } - if (train) tensor_alloc_stats_enable(0); destroy_context(&context); + if (train) { + tensor_alloc_stats_read(&context.allocation_stats); + if (context.allocation_stats.live_bytes != 0) status = 1; + tensor_alloc_stats_enable(0); + } return status; }