Matmul two variables

CMakeLists.txt

@@ -1,6 +1,5 @@
 cmake_minimum_required(VERSION 3.15)
 project(DNLP_Diff_Engine C)
-
 set(CMAKE_C_STANDARD 99)
 
 #Set default build type to Release if not specified
@@ -32,6 +31,7 @@ include_directories(${PROJECT_SOURCE_DIR}/include)
 # Source files - automatically gather all .c files from src/
 file(GLOB_RECURSE SOURCES "src/*.c")
 
+
 # Create core library
 add_library(dnlp_diff ${SOURCES})
 target_link_libraries(dnlp_diff m)

include/bivariate.h

@@ -10,6 +10,9 @@ expr *new_quad_over_lin(expr *left, expr *right);
 expr *new_rel_entr_first_arg_scalar(expr *left, expr *right);
 expr *new_rel_entr_second_arg_scalar(expr *left, expr *right);
 
+/* Matrix multiplication: Z = X @ Y */
+expr *new_matmul(expr *x, expr *y);
+
 /* Left matrix multiplication: A @ f(x) where A is a constant matrix */
 expr *new_left_matmul(expr *u, const CSR_Matrix *A);
 

include/utils/mini_numpy.h

@@ -20,4 +20,7 @@ void tile_int(int *result, const int *a, int len, int tiles);
  */
 void scaled_ones(double *result, int size, double value);
 
+/* Naive implementation of Z = X @ Y, X is m x k, Y is k x n, Z is m x n */
+void mat_mat_mult(const double *X, const double *Y, double *Z, int m, int k, int n);
+
 #endif /* MINI_NUMPY_H */

python/atoms/matmul.h

@@ -0,0 +1,33 @@
+#ifndef ATOM_MATMUL_H
+#define ATOM_MATMUL_H
+
+#include "bivariate.h"
+#include "common.h"
+
+/* Matrix multiplication: Z = X @ Y */
+static PyObject *py_make_matmul(PyObject *self, PyObject *args)
+{
+    PyObject *left_capsule, *right_capsule;
+    if (!PyArg_ParseTuple(args, "OO", &left_capsule, &right_capsule))
+    {
+        return NULL;
+    }
+    expr *left = (expr *) PyCapsule_GetPointer(left_capsule, EXPR_CAPSULE_NAME);
+    expr *right = (expr *) PyCapsule_GetPointer(right_capsule, EXPR_CAPSULE_NAME);
+    if (!left || !right)
+    {
+        PyErr_SetString(PyExc_ValueError, "invalid child capsule");
+        return NULL;
+    }
+
+    expr *node = new_matmul(left, right);
+    if (!node)
+    {
+        PyErr_SetString(PyExc_RuntimeError, "failed to create matmul node");
+        return NULL;
+    }
+    expr_retain(node); /* Capsule owns a reference */
+    return PyCapsule_New(node, EXPR_CAPSULE_NAME, expr_capsule_destructor);
+}
+
+#endif /* ATOM_MATMUL_H */

python/bindings.c

@@ -19,6 +19,7 @@
 #include "atoms/linear.h"
 #include "atoms/log.h"
 #include "atoms/logistic.h"
+#include "atoms/matmul.h"
 #include "atoms/multiply.h"
 #include "atoms/neg.h"
 #include "atoms/power.h"
@@ -73,6 +74,8 @@ static PyMethodDef DNLPMethods[] = {
     {"make_promote", py_make_promote, METH_VARARGS, "Create promote node"},
     {"make_multiply", py_make_multiply, METH_VARARGS,
      "Create elementwise multiply node"},
+    {"make_matmul", py_make_matmul, METH_VARARGS,
+     "Create matrix multiplication node (Z = X @ Y)"},
     {"make_const_scalar_mult", py_make_const_scalar_mult, METH_VARARGS,
      "Create constant scalar multiplication node (a * f(x))"},
     {"make_const_vector_mult", py_make_const_vector_mult, METH_VARARGS,