Optim attempt

CMakeLists.txt

@@ -8,6 +8,8 @@ if (CMAKE_C_COMPILER_ID MATCHES "GNU|Clang")
     set(CMAKE_C_FLAGS_RELEASE "${CMAKE_C_FLAGS_RELEASE} -march=native -mtune=native")
 endif()
 
+find_package(OpenMP REQUIRED)
+
 set(TENSORLIB_SOURCES
     src/tensor/tensor_core.c
     src/tensor/tensor_ops.c
@@ -49,6 +51,7 @@ add_library(tensorlib STATIC ${TENSORLIB_SOURCES}
             include/tensorlib/autograd.h
             include/tensorlib/nn.h)
 target_include_directories(tensorlib PUBLIC include/tensorlib)
+target_link_libraries(tensorlib PUBLIC OpenMP::OpenMP)
 
 if (WIN32)
     set(TENSORLIB_MATH_LIB)

Makefile

@@ -1,5 +1,5 @@
 CC = gcc
-CFLAGS = -O3 -march=native -mtune=native -Wall -Wextra -g -std=c11
+CFLAGS = -O3 -march=native -mtune=native -Wall -Wextra -g -std=c11 -fopenmp
 SRC = src/tensor/tensor_core.c src/tensor/tensor_alloc.c src/tensor/tensor_view.c src/tensor/tensor_ops.c src/tensor/tensor_gather.c src/tensor/tensor_reduc.c src/tensor/tensor_matmul.c src/autograd/autograd_core.c src/autograd/autograd_ops.c src/autograd/autograd_view.c src/autograd/autograd_gather.c src/autograd/autograd_reduc.c src/autograd/autograd_matmul.c src/autograd/autograd_backward.c src/init/rng.c src/nn/parameter.c src/nn/module.c src/nn/linear.c src/nn/embedding.c src/nn/positional_embedding.c src/nn/layer_norm.c src/nn/dropout.c src/nn/multihead_attention.c src/nn/decoder_block.c src/nn/decoder.c src/nn/mlp.c src/losses/classification.c src/nn/causal_mask.c src/optim/sgd.c src/optim/optim_common.c src/optim/adamw.c src/serialization/checkpoint.c
 HEADERS = include/tensorlib/tensor.h include/tensorlib/tensor_matmul.h include/tensorlib/autograd.h include/tensorlib/nn.h tests/fixtures/test_common.h
 INCLUDES = -Iinclude/tensorlib -Itests/fixtures
@@ -138,11 +138,14 @@ $(BIN)/bench_tensor_matmul_reuse: benchmarks/matmul/bench_tensor_matmul_reuse.c
 $(BIN)/bench_tensor_matmul_packed_views: benchmarks/matmul/bench_tensor_matmul_packed_views.c $(SRC) $(HEADERS) | $(BIN)
 	$(CC) $(CFLAGS) $(INCLUDES) -o $@ benchmarks/matmul/bench_tensor_matmul_packed_views.c $(SRC) -lm
 
+OPENBLAS_INCLUDE = -IC:/msys64/ucrt64/include/openblas
+OPENBLAS_LIB = -LC:/msys64/ucrt64/lib -lopenblas
+
 $(BIN)/bench_openblas_matmul: benchmarks/matmul/bench_openblas_matmul.c $(SRC) $(HEADERS) | $(BIN)
-	$(CC) $(CFLAGS) $(INCLUDES) -o $@ benchmarks/matmul/bench_openblas_matmul.c $(SRC) -lopenblas -lm
+	$(CC) $(CFLAGS) $(INCLUDES) $(OPENBLAS_INCLUDE) -o $@ benchmarks/matmul/bench_openblas_matmul.c $(SRC) $(OPENBLAS_LIB) -lm
 
 $(BIN)/bench_openblas_matmul_reuse: benchmarks/matmul/bench_openblas_matmul_reuse.c $(SRC) $(HEADERS) | $(BIN)
-	$(CC) $(CFLAGS) $(INCLUDES) -o $@ benchmarks/matmul/bench_openblas_matmul_reuse.c $(SRC) -lopenblas -lm
+	$(CC) $(CFLAGS) $(INCLUDES) $(OPENBLAS_INCLUDE) -o $@ benchmarks/matmul/bench_openblas_matmul_reuse.c $(SRC) $(OPENBLAS_LIB) -lm
 
 test: $(TESTS)
 	@for t in $(TESTS); do ./$$t || exit 1; echo; done
@@ -156,15 +159,27 @@ benchmark-matmul-reuse: $(BIN)/bench_tensor_matmul_reuse
 benchmark-packed-views: $(BIN)/bench_tensor_matmul_packed_views
 	./$(BIN)/bench_tensor_matmul_packed_views
 
-benchmark-all: benchmark-matmul benchmark-matmul-reuse benchmark-packed-views benchmark-openblas benchmark-openblas-reuse
-
 benchmark-openblas: $(BIN)/bench_openblas_matmul
-	OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 ./$(BIN)/bench_openblas_matmul
+	./$(BIN)/bench_openblas_matmul
 
 benchmark-openblas-reuse: $(BIN)/bench_openblas_matmul_reuse
-	OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 ./$(BIN)/bench_openblas_matmul_reuse
+	./$(BIN)/bench_openblas_matmul_reuse
+
+benchmark-compare:
+	@echo "===== TensorLib (OpenMP) ====="
+	OMP_NUM_THREADS=1 ./$(BIN)/bench_tensor_matmul | head -4
+	@echo ""
+	OMP_NUM_THREADS= ./$(BIN)/bench_tensor_matmul | head -4
+	@echo ""
+	@echo "===== OpenBLAS ====="
+	OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 ./$(BIN)/bench_openblas_matmul | head -4
+	@echo ""
+	OMP_NUM_THREADS= OPENBLAS_NUM_THREADS=0 ./$(BIN)/bench_openblas_matmul | head -4
+	@echo ""
+
+benchmark-all: benchmark-matmul benchmark-matmul-reuse benchmark-packed-views benchmark-openblas benchmark-openblas-reuse
 
 clean:
 	rm -rf $(BIN)
 
-.PHONY: all test example mnist tiny-lm benchmark-all benchmark-matmul benchmark-matmul-reuse benchmark-packed-views benchmark-openblas benchmark-openblas-reuse clean
+.PHONY: all test example mnist tiny-lm benchmark-all benchmark-matmul benchmark-matmul-reuse benchmark-packed-views benchmark-openblas benchmark-openblas-reuse benchmark-compare clean

benchmarks/matmul/bench_openblas_matmul.c

@@ -5,16 +5,18 @@
 
 #include "../../include/tensorlib/tensor.h"
 
-/* Keep this comparison single-threaded, like bench_tensor_matmul. */
 extern void openblas_set_num_threads(int num_threads);
 
 enum {
-    BATCHES = 8,
-    M = 128,
-    K = 128,
-    N = 128,
+    SMALL_BATCHES = 8,
+    SMALL_M = 128,
+    SMALL_K = 128,
+    SMALL_N = 128,
+    LARGE_M = 4096,
+    LARGE_K = 4096,
+    LARGE_N = 4096,
     WARMUP = 3,
-    REPEATS = 50
+    REPEATS = 20
 };
 
 static void fill_input(tensor* t, float scale) {
@@ -24,82 +26,97 @@ static void fill_input(tensor* t, float scale) {
     }
 }
 
-static void openblas_batch_matmul(const tensor* a, const tensor* b, tensor* output) {
-    const int matrix_size_a = M * K;
-    const int matrix_size_b = K * N;
-    const int matrix_size_output = M * N;
-
-    for (int batch = 0; batch < BATCHES; ++batch) {
-        const float* a_data = a->storage->data + batch * matrix_size_a;
-        const float* b_data = b->storage->data + batch * matrix_size_b;
-        float* output_data = output->storage->data + batch * matrix_size_output;
-
-        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
-                    M, N, K,
-                    1.0f, a_data, K,
-                    b_data, N,
-                    0.0f, output_data, N);
-    }
-}
-
-int main(void) {
-    openblas_set_num_threads(1);
-
-    int a_dims[3] = {BATCHES, M, K};
-    int b_dims[3] = {BATCHES, K, N};
-    int output_dims[3] = {BATCHES, M, N};
-    tensor* a = t_alloc(3, a_dims);
-    tensor* b = t_alloc(3, b_dims);
-    if (a == NULL || b == NULL) {
-        t_free(b);
-        t_free(a);
-        fprintf(stderr, "failed to allocate benchmark inputs\n");
-        return 1;
+static int run_benchmark(const char* label, int num_threads,
+                         int batch, int m, int k, int n,
+                         int warmup, int repeats) {
+    openblas_set_num_threads(num_threads);
+
+    int matrix_size_a = m * k;
+    int matrix_size_b = k * n;
+    int matrix_size_c = m * n;
+
+    int a_ndim = (batch > 1) ? 3 : 2;
+    int a_dims[3] = {batch > 1 ? batch : 0, m, k};
+    int b_ndim = (batch > 1) ? 3 : 2;
+    int b_dims[3] = {batch > 1 ? batch : 0, k, n};
+    int c_ndim = (batch > 1) ? 3 : 2;
+    int c_dims[3] = {batch > 1 ? batch : 0, m, n};
+
+    tensor* a = t_alloc(a_ndim, a_ndim == 3 ? a_dims : a_dims + 1);
+    tensor* b = t_alloc(b_ndim, b_ndim == 3 ? b_dims : b_dims + 1);
+    tensor* output = t_alloc(c_ndim, c_ndim == 3 ? c_dims : c_dims + 1);
+    if (a == NULL || b == NULL || output == NULL) {
+        t_free(output); t_free(b); t_free(a);
+        fprintf(stderr, "failed to allocate\n");
+        return 0;
     }
 
     fill_input(a, 0.01f);
     fill_input(b, 0.02f);
 
-    for (int repeat = 0; repeat < WARMUP; ++repeat) {
-        tensor* output = t_alloc(3, output_dims);
-        if (output == NULL) {
-            t_free(b);
-            t_free(a);
-            fprintf(stderr, "failed to allocate benchmark output\n");
-            return 1;
+    for (int r = 0; r < warmup; ++r) {
+        for (int batch_idx = 0; batch_idx < batch; ++batch_idx) {
+            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
+                        m, n, k,
+                        1.0f, a->storage->data + batch_idx * matrix_size_a, k,
+                        b->storage->data + batch_idx * matrix_size_b, n,
+                        0.0f, output->storage->data + batch_idx * matrix_size_c, n);
         }
-        openblas_batch_matmul(a, b, output);
-        t_free(output);
     }
 
     clock_t start = clock();
     volatile float checksum = 0.0f;
-    for (int repeat = 0; repeat < REPEATS; ++repeat) {
-        tensor* output = t_alloc(3, output_dims);
-        if (output == NULL) {
-            t_free(b);
-            t_free(a);
-            fprintf(stderr, "failed to allocate benchmark output\n");
-            return 1;
+    for (int r = 0; r < repeats; ++r) {
+        for (int batch_idx = 0; batch_idx < batch; ++batch_idx) {
+            cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
+                        m, n, k,
+                        1.0f, a->storage->data + batch_idx * matrix_size_a, k,
+                        b->storage->data + batch_idx * matrix_size_b, n,
+                        0.0f, output->storage->data + batch_idx * matrix_size_c, n);
         }
-        openblas_batch_matmul(a, b, output);
         checksum += output->storage->data[0];
         checksum += output->storage->data[tensor_numel(output) - 1];
-        t_free(output);
     }
     clock_t end = clock();
 
     double seconds = (double)(end - start) / (double)CLOCKS_PER_SEC;
-    double operations = 2.0 * (double)BATCHES * (double)M * (double)K *
-                        (double)N * (double)REPEATS;
-    double gflops = (seconds > 0.0) ? operations / seconds / 1.0e9 : 0.0;
+    double ops = 2.0 * (double)batch * (double)m * (double)k * (double)n * (double)repeats;
+    double gflops = (seconds > 0.0) ? ops / seconds / 1.0e9 : 0.0;
 
-    printf("OpenBLAS: batch=%d M=%d K=%d N=%d warmup=%d repeats=%d\n",
-           BATCHES, M, K, N, WARMUP, REPEATS);
-    printf("elapsed_seconds=%.6f gflops=%.6f checksum=%.6f\n",
+    printf("%s: batch=%d M=%d K=%d N=%d threads=%d warmup=%d repeats=%d\n",
+           label, batch, m, k, n, num_threads, warmup, repeats);
+    printf("  elapsed_seconds=%.6f gflops=%.6f checksum=%.6f\n",
            seconds, gflops, (double)checksum);
 
+    t_free(output);
     t_free(b);
     t_free(a);
+    return 1;
+}
+
+int main(void) {
+    printf("=== OpenBLAS Matmul Benchmarks ===\n\n");
+
+    /* Small batched — single-threaded baseline */
+    run_benchmark("small_st", 1,
+                  SMALL_BATCHES, SMALL_M, SMALL_K, SMALL_N,
+                  WARMUP, 50);
+
+    /* Small batched — multi-threaded */
+    run_benchmark("small_mt", 0,
+                  SMALL_BATCHES, SMALL_M, SMALL_K, SMALL_N,
+                  WARMUP, 50);
+
+    /* Large single — single-threaded */
+    run_benchmark("large_st", 1,
+                  1, LARGE_M, LARGE_K, LARGE_N,
+                  WARMUP, REPEATS);
+
+    /* Large single — multi-threaded */
+    run_benchmark("large_mt", 0,
+                  1, LARGE_M, LARGE_K, LARGE_N,
+                  WARMUP, REPEATS);
+
+    printf("\n");
     return 0;
 }

benchmarks/matmul/bench_tensor_matmul.c

@@ -1,15 +1,23 @@
 #include <stdio.h>
 #include <time.h>
+#include <stdint.h>
 
 #include "../../include/tensorlib/tensor.h"
 
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
 enum {
-    BATCHES = 8,
-    M = 128,
-    K = 128,
-    N = 128,
+    SMALL_BATCHES = 8,
+    SMALL_M = 128,
+    SMALL_K = 128,
+    SMALL_N = 128,
+    LARGE_M = 4096,
+    LARGE_K = 4096,
+    LARGE_N = 4096,
     WARMUP = 3,
-    REPEATS = 50
+    REPEATS = 20
 };
 
 static void fill_input(tensor* t, float scale) {
@@ -19,35 +27,37 @@ static void fill_input(tensor* t, float scale) {
     }
 }
 
-int main(void) {
-    int a_dims[3] = {BATCHES, M, K};
-    int b_dims[3] = {BATCHES, K, N};
-    tensor* a = t_alloc(3, a_dims);
-    tensor* b = t_alloc(3, b_dims);
+static int run_benchmark(const char* label, int batch, int m, int k, int n,
+                         int warmup, int repeats) {
+    int a_ndim = (batch > 1) ? 3 : 2;
+    int a_dims[3] = {batch > 1 ? batch : 0, m, k};
+    int b_ndim = (batch > 1) ? 3 : 2;
+    int b_dims[3] = {batch > 1 ? batch : 0, k, n};
+
+    tensor* a = t_alloc(a_ndim, a_ndim == 3 ? a_dims : a_dims + 1);
+    tensor* b = t_alloc(b_ndim, b_ndim == 3 ? b_dims : b_dims + 1);
     if (a == NULL || b == NULL) {
-        t_free(b);
-        t_free(a);
-        fprintf(stderr, "failed to allocate benchmark inputs\n");
-        return 1;
+        t_free(b); t_free(a);
+        fprintf(stderr, "failed to allocate\n");
+        return 0;
     }
 
     fill_input(a, 0.01f);
     fill_input(b, 0.02f);
 
-    for (int repeat = 0; repeat < WARMUP; ++repeat) {
+    for (int r = 0; r < warmup; ++r) {
         tensor* output = t_matmul(a, b);
         t_free(output);
     }
 
     clock_t start = clock();
     volatile float checksum = 0.0f;
-    for (int repeat = 0; repeat < REPEATS; ++repeat) {
+    for (int r = 0; r < repeats; ++r) {
         tensor* output = t_matmul(a, b);
         if (output == NULL) {
-            t_free(b);
-            t_free(a);
-            fprintf(stderr, "matmul failed during benchmark\n");
-            return 1;
+            t_free(b); t_free(a);
+            fprintf(stderr, "matmul failed\n");
+            return 0;
         }
         checksum += output->storage->data[0];
         checksum += output->storage->data[tensor_numel(output) - 1];
@@ -56,15 +66,36 @@ int main(void) {
     clock_t end = clock();
 
     double seconds = (double)(end - start) / (double)CLOCKS_PER_SEC;
-    double operations = 2.0 * (double)BATCHES * (double)M * (double)K * (double)N * (double)REPEATS;
-    double gflops = (seconds > 0.0) ? operations / seconds / 1.0e9 : 0.0;
+    double ops = 2.0 * (double)batch * (double)m * (double)k * (double)n * (double)repeats;
+    double gflops = (seconds > 0.0) ? ops / seconds / 1.0e9 : 0.0;
 
-    printf("t_matmul API: batch=%d M=%d K=%d N=%d warmup=%d repeats=%d\n",
-           BATCHES, M, K, N, WARMUP, REPEATS);
-    printf("elapsed_seconds=%.6f gflops=%.6f checksum=%.6f\n",
+#ifdef _OPENMP
+    int threads = omp_get_max_threads();
+#else
+    int threads = 1;
+#endif
+
+    printf("%s: batch=%d M=%d K=%d N=%d threads=%d warmup=%d repeats=%d\n",
+           label, batch, m, k, n, threads, warmup, repeats);
+    printf("  elapsed_seconds=%.6f gflops=%.6f checksum=%.6f\n",
            seconds, gflops, (double)checksum);
 
     t_free(b);
     t_free(a);
+    return 1;
+}
+
+int main(void) {
+    printf("=== TensorLib Matmul Benchmarks ===\n\n");
+
+    /* Small batched matmul (existing baseline) */
+    run_benchmark("small_batched", SMALL_BATCHES, SMALL_M, SMALL_K, SMALL_N,
+                  WARMUP, 50);
+
+    /* Single large matmul to show OpenMP scaling */
+    run_benchmark("large_single", 1, LARGE_M, LARGE_K, LARGE_N,
+                  WARMUP, REPEATS);
+
+    printf("\n");
     return 0;
 }