WIP

.github/workflows/tests.yml

@@ -152,14 +152,3 @@ jobs:
 
       - name: Run tests with extras
         run: python -m pytest --pyargs thinc --cov=thinc --cov-report=term -p thinc.tests.enable_tensorflow -p thinc.tests.enable_mxnet
-
-      - name: Run tests for thinc-apple-ops
-        run: |
-          pip uninstall -y tensorflow
-          pip install thinc-apple-ops
-          python -m pytest --pyargs thinc_apple_ops
-        if: matrix.os == 'macos-latest' && matrix.python_version == '3.10'
-
-      - name: Run tests with thinc-apple-ops
-        run: python -m pytest --pyargs thinc
-        if: matrix.os == 'macos-latest' && matrix.python_version == '3.10'

setup.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python
+import platform
 import sys
 from setuptools.command.build_ext import build_ext
 from sysconfig import get_path
@@ -13,14 +14,22 @@
 # http://docs.cython.org/en/latest/src/userguide/source_files_and_compilation.html#compiler-options
 Options.docstrings = True
 
+APPLE_OPS = ["thinc.backends.apple_ops", "thinc.backends.apple_blas"]
 
 PACKAGES = find_packages()
-MOD_NAMES = [
-    "thinc.backends.cblas",
-    "thinc.backends.numpy_ops",
-    "thinc.layers.sparselinear",
-    "thinc.layers.premap_ids",
-]
+MOD_NAMES = (
+    [
+        "thinc.backends.cblas",
+        "thinc.backends.linalg",
+        "thinc.backends.numpy_ops",
+        "thinc.extra.search",
+        "thinc.layers.sparselinear",
+        "thinc.layers.premap_ids",
+    ]
+    + APPLE_OPS
+    if platform.system() == "Darwin"
+    else []
+)
 COMPILE_OPTIONS = {
     "msvc": ["/Ox", "/EHsc"],
     "other": ["-O3", "-Wno-strict-prototypes", "-Wno-unused-function", "-std=c++11"],

thinc/backends/__init__.py

@@ -83,10 +83,6 @@ def use_tensorflow_for_gpu_memory() -> None:  # pragma: no cover
 
 
 def _import_extra_cpu_backends():
-    try:
-        from thinc_apple_ops import AppleOps
-    except ImportError:
-        pass
     try:
         from thinc_bigendian_ops import BigEndianOps
     except ImportError:

thinc/backends/apple_blas.pxd

@@ -0,0 +1,40 @@
+cdef extern from "Accelerate/Accelerate.h":
+    enum CBLAS_ORDER:     CblasRowMajor, CblasColMajor
+    enum CBLAS_TRANSPOSE: CblasNoTrans, CblasTrans, CblasConjTrans
+    enum CBLAS_UPLO:      CblasUpper, CblasLower
+    enum CBLAS_DIAG:      CblasNonUnit, CblasUnit
+    enum CBLAS_SIDE:      CblasLeft, CblasRight
+
+    # BLAS level 1 routines
+
+    void cblas_sswap(int M, float  *x, int incX, float  *y, int incY) nogil
+    void cblas_sscal(int N, float  alpha, float  *x, int incX) nogil
+    void cblas_scopy(int N, float  *x, int incX, float  *y, int incY) nogil
+    void cblas_saxpy(int N, float  alpha, float  *x, int incX, float  *y, int incY ) nogil
+    float cblas_sdot(int N, float  *x, int incX, float  *y, int incY ) nogil
+    float cblas_snrm2(int N, float  *x, int incX) nogil
+    float cblas_sasum(int N, float  *x, int incX) nogil
+    int cblas_isamax(int N, float  *x, int incX) nogil
+
+    # BLAS level 2 routines
+    void cblas_sgemv(CBLAS_ORDER Order, CBLAS_TRANSPOSE TransA, int M, int N,
+                                 float  alpha, float  *A, int lda, float  *x, int incX,
+                                 float  beta, float  *y, int incY) nogil
+
+    void cblas_sger(CBLAS_ORDER Order, int M, int N, float  alpha, float  *x,
+                                int incX, float  *y, int incY, float  *A, int lda) nogil
+
+    # BLAS level 3 routines
+    void cblas_sgemm(CBLAS_ORDER Order, CBLAS_TRANSPOSE TransA,
+                                 CBLAS_TRANSPOSE TransB, int M, int N, int K,
+                                 float  alpha, float  *A, int lda, float  *B, int ldb,
+                                 float  beta, float  *C, int ldc) nogil
+
+
+cdef void sgemm(bint TransA, bint TransB, int M, int N, int K,
+                    float alpha, const float* A, int lda, const float *B,
+                    int ldb, float beta, float* C, int ldc) nogil
+
+
+cdef void saxpy(int N, float alpha, const float* X, int incX,
+                float *Y, int incY) nogil

thinc/backends/apple_blas.pyx

@@ -0,0 +1,74 @@
+from libc.stdint cimport uintptr_t
+cimport numpy as np
+import numpy
+
+
+cpdef np.ndarray gemm(float[:, ::1] A, float[:, ::1] B,
+                      bint trans1=False, bint trans2=False,
+                      np.ndarray out=None):
+    cdef int nM = A.shape[0] if not trans1 else A.shape[1]
+    cdef int nK = A.shape[1] if not trans1 else A.shape[0]
+    cdef int nK_b = B.shape[0] if not trans2 else B.shape[1]
+    cdef int nN = B.shape[1] if not trans2 else B.shape[0]
+
+    cdef float[:, ::1] C = out
+
+    if out is None:
+        out = numpy.empty((nM, nN), dtype="f")
+        C = out
+    else:
+        if C.shape[0] != nM or C.shape[1] != nN:
+            msg = "Shape mismatch for output matrix, was: (%d, %d), expected (%d, %d)"
+            raise ValueError(msg % (C.shape[0], C.shape[1], nM, nN))
+
+
+    if nK != nK_b:
+        msg = "Shape mismatch for gemm: (%d, %d), (%d, %d)"
+        raise ValueError(msg % (nM, nK, nK_b, nN))
+
+    if nM == 0 or nK == 0 or nN == 0:
+        return out
+
+    cblas_sgemm(
+        CblasRowMajor,
+        CblasTrans if trans1 else CblasNoTrans,
+        CblasTrans if trans2 else CblasNoTrans,
+        nM,
+        nN,
+        nK,
+        1.0,
+        &A[0, 0],
+        A.shape[1],
+        &B[0, 0],
+        B.shape[1],
+        0.0,
+        &C[0, 0],
+        C.shape[1]
+    )
+    return out
+
+
+cdef void sgemm(bint TransA, bint TransB, int M, int N, int K,
+                    float alpha, const float* A, int lda, const float *B,
+                    int ldb, float beta, float* C, int ldc) nogil:
+    cblas_sgemm(
+        CblasRowMajor,
+        CblasTrans if TransA else CblasNoTrans,
+        CblasTrans if TransB else CblasNoTrans,
+        M,
+        N,
+        K,
+        alpha,
+        A,
+        lda,
+        B,
+        ldb,
+        beta,
+        C,
+        ldc
+    )
+
+
+cdef void saxpy(int N, float alpha, const float* X, int incX,
+                float *Y, int incY) nogil:
+    cblas_saxpy(N, alpha, X, incX, Y, incY)

thinc/backends/apple_ops.pyx

@@ -0,0 +1,36 @@
+from typing import Optional
+import numpy
+
+from . import apple_blas
+from .apple_blas cimport saxpy, sgemm
+from .cblas cimport CBlas, set_saxpy, set_sgemm
+from .numpy_ops import NumpyOps
+from ..types import Floats2d
+from .. import registry
+
+
+@registry.ops("AppleOps")
+class AppleOps(NumpyOps):
+    """Thinc Ops class that calls into Apple's native libraries for some
+    operations. Other operations fall back to numpy."""
+    name = "apple"
+    xp = numpy
+
+    def cblas(self) -> CBlas:
+        cdef CBlas cblas = CBlas()
+        set_saxpy(cblas, saxpy)
+        set_sgemm(cblas, sgemm)
+        return cblas
+
+    def gemm(
+        self,
+        x: Floats2d,
+        y: Floats2d,
+        out: Optional[Floats2d] = None,
+        trans1: bool = False,
+        trans2: bool = False,
+    ) -> Floats2d:
+        """Perform General Matrix Multiplication (GeMM) and optionally store
+        the result in the specified output variable.
+        """
+        return apple_blas.gemm(x, y, out=out, trans1=trans1, trans2=trans2)

thinc/backends/mps_ops.py

@@ -12,9 +12,9 @@
     # during type checking.
     _Ops = Ops
 else:
-    try:
-        from thinc_apple_ops import AppleOps
+    from .apple_ops import AppleOps
 
+    try:
         _Ops = AppleOps
     except ImportError:
         _Ops = NumpyOps

thinc/tests/backends/test_ops.py

@@ -1403,7 +1403,7 @@ def test_get_ops():
     # If Apple ops are available, "cpu" should return AppleOps or
     # NumpyOps otherwise.
     try:
-        from thinc_apple_ops import AppleOps
+        from thinc.backends.apple_ops import AppleOps
 
         assert isinstance(get_ops("cpu"), AppleOps)
     except ImportError: