feat: MatMul kernel for Dense <- Dense, Sparse

Author: philipportner

src/runtime/local/kernels/MatMul.cpp

@@ -254,7 +254,8 @@ void MatMul<DenseMatrix<VT>, DenseMatrix<VT>, DenseMatrix<VT>>::apply(DenseMatri
     const auto nr2 = static_cast<int>(transb ? rhs->getNumCols() : rhs->getNumRows());
     const auto nc2 = static_cast<int>(transb ? rhs->getNumRows() : rhs->getNumCols());
     if (nc1 != nr2) {
-        throw std::runtime_error("MatMul - #cols of lhs and #rows of rhs must be the same");
+        throw std::runtime_error("MatMul - #cols of lhs and #rows of rhs must be the same, but got" +
+                                 std::to_string(nc1) + "and" + std::to_string(nr2));
     }
     const VT alpha = 1.0f;
     const VT beta = 0.0f;

src/runtime/local/kernels/MatMul.h

@@ -93,6 +93,56 @@ template <typename VT> struct MatMul<DenseMatrix<VT>, CSRMatrix<VT>, DenseMatrix
     }
 };
 
+// ----------------------------------------------------------------------------
+// DenseMatrix <- DenseMatrix, CSRMatrix
+// ----------------------------------------------------------------------------
+
+template <typename VT> struct MatMul<DenseMatrix<VT>, DenseMatrix<VT>, CSRMatrix<VT>> {
+    static void apply(DenseMatrix<VT> *&res, const DenseMatrix<VT> *lhs, const CSRMatrix<VT> *rhs, bool transa,
+                      bool transb, DCTX(ctx)) {
+        const size_t nr1 = lhs->getNumRows();
+        const size_t nc1 = lhs->getNumCols();
+        const size_t nr2 = rhs->getNumRows();
+        const size_t nc2 = rhs->getNumCols();
+
+        if (nc1 != nr2) {
+            throw std::runtime_error("MatMul - #cols of lhs and #rows of rhs must be the same");
+        }
+        // FIXME: transpose isn't supported atm
+
+        if (res == nullptr)
+            res = DataObjectFactory::create<DenseMatrix<VT>>(nr1, nc2, /*zero=*/true);
+
+        const VT *valuesLhs = lhs->getValues();
+        VT *valuesRes = res->getValues();
+
+        const size_t rowSkipLhs = lhs->getRowSkip();
+        const size_t rowSkipRes = res->getRowSkip();
+
+        // For each row m of lhs
+        for (size_t m = 0; m < nr1; m++) {
+            const size_t rowIdxLhs = m * rowSkipLhs;
+            const size_t rowIdxRes = m * rowSkipRes;
+
+            // For each row n of rhs
+            for (size_t n = 0; n < nr2; n++) {
+                const VT lhsVal = valuesLhs[rowIdxLhs + n];
+
+                // Get non-zeros in row k of rhs
+                const size_t rowNumNonZeros = rhs->getNumNonZeros(n);
+                const size_t *rowColIdxs = rhs->getColIdxs(n);
+                const VT *rowValues = rhs->getValues(n);
+
+                // For each non-zero in row k of rhs
+                for (size_t i = 0; i < rowNumNonZeros; i++) {
+                    const size_t c = rowColIdxs[i];
+                    valuesRes[rowIdxRes + c] += lhsVal * rowValues[i];
+                }
+            }
+        }
+    }
+};
+
 // ----------------------------------------------------------------------------
 // Matrix <- Matrix, Matrix
 // ----------------------------------------------------------------------------

src/runtime/local/kernels/kernels.json

@@ -3695,6 +3695,11 @@
                         ["CSRMatrix", "double"],
                         ["DenseMatrix", "double"]
                     ],
+                    [
+                        ["DenseMatrix", "double"],
+                        ["DenseMatrix", "double"],
+                        ["CSRMatrix", "double"]
+                    ],
                     [
                         ["CSRMatrix", "double"],
                         ["CSRMatrix", "double"],

test/runtime/local/kernels/MatMulTest.cpp

@@ -41,6 +41,15 @@ void checkMatMul(const DT *lhs, const DT *rhs, const DT *exp, DCTX(dctx), bool t
     DataObjectFactory::destroy(res);
 }
 
+template <class DTRes, class DTLhs, class DTRhs>
+void checkMatMulMixed(const DTLhs *lhs, const DTRhs *rhs, const DTRes *exp, DCTX(dctx), bool transa = false,
+                      bool transb = false) {
+    DTRes *res = nullptr;
+    matMul<DTRes, DTLhs, DTRhs>(res, lhs, rhs, transa, transb, dctx);
+    CHECK(*res == *exp);
+    DataObjectFactory::destroy(res);
+}
+
 TEMPLATE_PRODUCT_TEST_CASE("MatMul", TAG_KERNELS, (CSRMatrix, DATA_TYPES), (VALUE_TYPES)) {
     auto dctx = setupContextAndLogger();
 
@@ -143,6 +152,25 @@ TEMPLATE_PRODUCT_TEST_CASE("MatMul", TAG_KERNELS, (CSRMatrix, DATA_TYPES), (VALU
     DataObjectFactory::destroy(m0, m1, m2, m3, m4, m5, m6, v0, v1, v2, v3, v4, v5, v6, v7, v8);
 }
 
+TEMPLATE_TEST_CASE("MatMul Dense x Sparse", TAG_KERNELS, VALUE_TYPES) {
+    using VT = TestType;
+    using LhsDT = DenseMatrix<VT>;
+    using RhsDT = CSRMatrix<VT>;
+    using ResDT = DenseMatrix<VT>;
+
+    auto dctx = setupContextAndLogger();
+
+    // clang-format off
+    auto lhs = genGivenVals<LhsDT>(3, { 1, 0, 2, 3, 4, 0, });
+    auto rhs = genGivenVals<RhsDT>(2, { 0, 5, 0, 0, 6, 0, });
+    auto exp = genGivenVals<ResDT>(3, { 0, 5, 0, 0, 28, 0, 0, 20, 0, });
+    // clang-format on
+
+    checkMatMulMixed<ResDT, LhsDT, RhsDT>(lhs, rhs, exp, dctx.get());
+
+    DataObjectFactory::destroy(lhs, rhs, exp);
+}
+
 TEMPLATE_PRODUCT_TEST_CASE("MatMul Transposed", TAG_KERNELS, (DATA_TYPES), (VALUE_TYPES)) {
     using DT = TestType;
     auto dctx = setupContextAndLogger();