Faster matrix multiplication mod p <= 255

doc/source/nmod_mat.rst

@@ -391,6 +391,30 @@ Matrix multiplication
     when FLINT is built with a 32-bit word size.
     Dimensions must be compatible for matrix multiplication.
 
+.. function:: void nmod_mat_mul_u8(nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B)
+
+    Sets `C = AB` using kernels specialized to moduli up to `255`,
+    which the modulus of `C` must not exceed. Moduli `2` and `3` pack
+    the entries into bit representations directly (a
+    method-of-four-Russians leaf under a Strassen recursion,
+    parallelized over a tile grid of shared packed operands when
+    several threads are available); moduli up to `15` use an
+    in-register table-lookup kernel on a byte image of the entries;
+    larger moduli lift to single precision for :func:`flint_sgemm`.
+    Aliasing of the operands is supported. Dimensions must be compatible
+    for matrix multiplication.
+
+.. function:: void _nmod_mat_mul_u8(uint8_t * C, slong Cstride, const uint8_t * A, slong Astride, const uint8_t * B, slong Bstride, slong m, slong k, slong n, nmod_t mod)
+
+    Underlying multiplication on byte matrices with arbitrary row
+    strides: sets the `m \times n` matrix `C` to `AB` where `A` is
+    `m \times k` and `B` is `k \times n`, all entries reduced modulo
+    ``mod.n``, which must not exceed `255`. This entry point skips the
+    ``ulong``-to-byte conversions of :func:`nmod_mat_mul_u8` and is the
+    natural interface for byte-entry matrix types. Exact aliasing of
+    `C` with `A` or `B` is supported (arbitrary partial overlap is
+    not).
+
 .. function:: void nmod_mat_addmul(nmod_mat_t D, const nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B)
 
     Sets `D = C + AB`. `C` and `D` may be aliased with each other but

src/gr/nmod8.c

@@ -808,12 +808,28 @@ _nmod8_poly_mullow(uint8_t * res, const uint8_t * A, slong Alen, const uint8_t *
     return _nmod8_poly_mulmid(res, A, Alen, B, Blen, 0, len, ctx);
 }
 
-/* todo: tuning for rectangular matrices */
+#include "nmod_mat.h"
+
 static int
 _nmod8_mat_mul(gr_mat_t C, const gr_mat_t A, const gr_mat_t B, gr_ctx_t ctx)
 {
-    if (A->r >= 256 && A->c >= 256 && B->c >= 256)
-        return gr_mat_mul_strassen(C, A, B, ctx);
+    slong ar, ac, br, bc;
+
+    ar = gr_mat_nrows(A, ctx);
+    ac = gr_mat_ncols(A, ctx);
+    br = gr_mat_nrows(B, ctx);
+    bc = gr_mat_ncols(B, ctx);
+
+    if (ar >= 8 && ac >= 8 && bc >= 8)
+    {
+        if (ac != br || ar != gr_mat_nrows(C, ctx) || bc != gr_mat_ncols(C, ctx))
+            return GR_DOMAIN;
+
+        _nmod_mat_mul_u8(C->entries, C->stride, A->entries, A->stride, B->entries, B->stride,
+            ar, ac, bc, NMOD8_CTX(ctx));
+
+        return GR_SUCCESS;
+    }
     else
         return gr_mat_mul_classical(C, A, B, ctx);
 }

src/nmod_mat.h

@@ -20,6 +20,7 @@
 #define NMOD_MAT_INLINE static inline
 #endif
 
+#include <stdint.h>
 #include "nmod_types.h"
 
 #ifdef __cplusplus
@@ -179,6 +180,12 @@ void nmod_mat_mul(nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B);
 
 int nmod_mat_mul_blas(nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B);
 
+void _nmod_mat_mul_u8(uint8_t * C, slong Cstride,
+                 const uint8_t * A, slong Astride,
+                 const uint8_t * B, slong Bstride,
+                 slong m, slong k, slong n, nmod_t mod);
+void nmod_mat_mul_u8(nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B);
+
 void nmod_mat_mul_classical(nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B);
 
 void

src/nmod_mat/mul.c

@@ -25,12 +25,23 @@ nmod_mat_mul(nmod_mat_t C, const nmod_mat_t A, const nmod_mat_t B)
     slong n = B->c;
     slong min_dim = FLINT_MIN(FLINT_MIN(m, k), n);
     slong cutoff;
-    slong flint_num_threads = flint_get_num_threads();
 
     FLINT_ASSERT(C->r == A->r);
     FLINT_ASSERT(C->c == B->c);
     FLINT_ASSERT(A->c == B->r);
 
+    /* Todo: optimize nmod_mat_mul_blas. For mod.n >= 17, nmod_mat_mul_blas
+       should be faster up to dim about 1000-2000 as nmod_mat_mul_u8 does
+       the same sgemm but with an extra uint8 roundtrip. Currently
+       mul_blas narrowly loses to mul_u8 due to slow modular reductions. */
+    if (C->mod.n <= 255 && min_dim >= 8)
+    {
+        nmod_mat_mul_u8(C, A, B);
+        return;
+    }
+
+    slong flint_num_threads = flint_get_num_threads();
+
     /*
         tuning is based on several assumptions:
         (1) the gemm used by nmod_mat_mul_blas is at least as parallel