Merge remote-tracking branch 'origin' into update_assigned_attributes

.coveragerc

@@ -1,6 +0,0 @@
-[run]
-omit =
-    /private/var/*
-    tmp/*
-    tests/*
-

.github/workflows/ci-test-example.yml

@@ -15,15 +15,24 @@ jobs:
     steps:
       - name: Check out repository
         uses: actions/checkout@v2
+        with:
+          submodules: 'recursive'
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
       - name: Make compile script executable
-        run: chmod +x ./mithril/cores/c/compile.sh
+        run: |
+          chmod +x ./mithril/cores/c/raw_c/compile.sh
+          chmod +x ./mithril/cores/c/ggml/compile.sh
+          chmod +x ./mithril/cores/c/ggml/build_ggml.sh
       - name: Compile C code
         run: |
-          pushd ./mithril/cores/c
+          pushd ./mithril/cores/c/raw_c
+          ./compile.sh
+          popd
+          pushd ./mithril/cores/c/ggml
+          ./build_ggml.sh
           ./compile.sh
           popd
       - name: Install Python dependencies

.github/workflows/ci-test-macos.yaml

@@ -15,15 +15,24 @@ jobs:
     steps:
       - name: Check out repository
         uses: actions/checkout@v2
+        with:
+          submodules: 'recursive'
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Make compile script executable
-        run: chmod +x ./mithril/cores/c/compile.sh
+      - name: Make compile scripts executable
+        run: |
+          chmod +x ./mithril/cores/c/raw_c/compile.sh
+          chmod +x ./mithril/cores/c/ggml/compile.sh
+          chmod +x ./mithril/cores/c/ggml/build_ggml.sh
       - name: Compile C code
         run: |
-          pushd ./mithril/cores/c
+          pushd ./mithril/cores/c/raw_c
+          ./compile.sh
+          popd
+          pushd ./mithril/cores/c/ggml
+          ./build_ggml.sh
           ./compile.sh
           popd
       - name: Install Python dependencies
@@ -38,7 +47,6 @@ jobs:
           python3 -m pip install mypy
           python3 -m pip install pre-commit
           pre-commit run --all-files
-          python3 license_checker.py
       - name: Execute testcase unit tests
         run: pytest --cov --cov-report=xml -s tests/
       - name: Upload results to Codecov

.github/workflows/ci-test-ubuntu.yaml

@@ -13,15 +13,24 @@ jobs:
     steps:
       - name: Check out repository
         uses: actions/checkout@v2
+        with:
+          submodules: 'recursive'
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Make compile script executable
-        run: chmod +x ./mithril/cores/c/compile.sh
+      - name: Make compile scripts executable
+        run: |
+          chmod +x ./mithril/cores/c/raw_c/compile.sh
+          chmod +x ./mithril/cores/c/ggml/compile.sh
+          chmod +x ./mithril/cores/c/ggml/build_ggml.sh
       - name: Compile C code
         run: |
-          pushd ./mithril/cores/c
+          pushd ./mithril/cores/c/raw_c
+          ./compile.sh
+          popd
+          pushd ./mithril/cores/c/ggml
+          ./build_ggml.sh
           ./compile.sh
           popd
       - name: Install Python dependencies
@@ -35,7 +44,6 @@ jobs:
           python3 -m pip install mypy
           python3 -m pip install pre-commit
           pre-commit run --all-files
-          python3 license_checker.py
       - name: Execute testcase unit tests
         run: pytest -s tests/
 

.github/workflows/pypi-release.yml

@@ -15,6 +15,8 @@ jobs:
     steps:
       - name: Check out code
         uses: actions/checkout@v3
+        with:
+          submodules: 'recursive'
 
       - name: Set up Python
         uses: actions/setup-python@v4

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "mithril/cores/c/ggml/ggml"]
+	path = mithril/cores/c/ggml/ggml
+	url = https://github.com/ggml-org/ggml

.pre-commit-config.yaml

@@ -1,3 +1,4 @@
+exclude: '^mithril/cores/c/ggml/ggml/'
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
@@ -15,12 +16,14 @@ repos:
         name: license checker
         entry: python3 license_checker.py
         language: python
+        args: ["--exclude=mithril/cores/c/ggml/ggml"]
       - id: mypy
         name: mypy
         entry: mypy .
         language: system
         always_run: true
         pass_filenames: false
+        args: ["--exclude=mithril/cores/c/ggml/ggml"]
 
   # - repo: https://github.com/pre-commit/mirrors-mypy
   #   rev: v1.12.0

MANIFEST.in

@@ -1,3 +1,3 @@
 include mithril/cores/c/libmithrilc.so
 recursive-include mithril/cores *
-recursive-include mithril/backends/with_manualgrad/c_backend/src *
+recursive-include mithril/backends/with_manualgrad/c_backend/src *

license_checker.py

@@ -14,6 +14,14 @@
 
 
 import os
+import sys
+
+# Take the exclude argument with argparse
+exclude = ""
+for arg in sys.argv:
+    if arg.startswith("--exclude="):
+        exclude = arg.split("=")[1]
+        break
 
 license_py = """# Copyright 2022 Synnada, Inc.
 #
@@ -45,19 +53,27 @@
 // limitations under the License.
 """
 current_directory = os.getcwd()
-# Walk through the directory recursively
 
+
+# Walk through the directory recursively
 for root, _, files in os.walk(current_directory):
     if os.path.basename(root) == "tmp":
         continue
     for filename in files:
         if filename.endswith((".py", ".c", ".h")):  # Check for .py .h and .c files
             file_path = os.path.join(root, filename)
 
+            if exclude in file_path:
+                continue
+
             # Check if it's a file
             if os.path.isfile(file_path):
                 with open(file_path, encoding="utf-8", errors="ignore") as file:
-                    file_license = "".join(next(file) for _ in range(13))
+                    lines = file.readlines()
+                    if len(lines) < 13:
+                        raise Exception(f"No license found in {file_path}")
+
+                    file_license = "".join(lines[:13])
 
                 license = license_py if filename.endswith(".py") else license_ch
 

mithril/__init__.py

@@ -46,6 +46,7 @@
     "MlxBackend",
     "TorchBackend",
     "CBackend",
+    "GGMLBackend",
     "NumpyBackend",
     "compile",
     "DataType",
@@ -92,6 +93,11 @@
 except Exception:
     CBackend = UnavailableBackend  # type: ignore
 
+try:
+    from .backends.with_manualgrad.ggml_backend.backend import GGMLBackend
+except ImportError:
+    GGMLBackend = UnavailableBackend  # type: ignore
+
 try:
     from .backends.with_manualgrad.numpy_backend.backend import NumpyBackend
 except ImportError:

mithril/backends/backend.py

@@ -19,6 +19,7 @@
 from typing import Any, Generic, overload
 
 from .. import types
+from ..common import CGenConfig, PythonGenConfig
 from ..types import DataType
 from .parallel import Parallel
 from .utils import DtypeBits, StaticScalar
@@ -49,6 +50,7 @@ class Backend(ABC, Generic[DataType]):
     registered_primitives: dict[str, Callable[..., DataType]]
     array_creation_funcs: list[str]
     primitive_fn_path: str
+    CODEGEN_CONFIG: PythonGenConfig | CGenConfig
 
     def __init__(self, dtype: types.Dtype = types.float32, device: str = "cpu") -> None:
         # Check if given dtype is a valid one.
@@ -90,10 +92,6 @@ def e(self) -> float:
     def is_manualgrad(self) -> bool:
         raise NotImplementedError("is_manualgrad is not implemented")
 
-    @property
-    def codegen_config(self) -> dict[str, bool]:
-        raise NotImplementedError("codegen_config is not implemented")
-
     def get_backend_array_type(self) -> type[DataType]:
         raise NotImplementedError("get_backend_array_type is not implemented")
 

mithril/backends/with_autograd/jax_backend/backend.py

@@ -27,7 +27,6 @@
 from ...utils import DtypeSubTypes, StaticScalar, process_shape
 from . import utils
 from .parallel import JaxParallel
-from .utils import CODEGEN_CONFIG
 
 __all__ = ["JaxBackend"]
 
@@ -52,6 +51,7 @@ class JaxBackend(ParallelBackend[jax.numpy.ndarray]):
     backend_type = "jax"
     registered_primitives: dict[str, Callable[..., jax.numpy.ndarray]] = {}
     primitive_fn_path = "mithril.cores.python.jax.ops"
+    CODEGEN_CONFIG = utils.CODEGEN_CONFIG
 
     def __init__(
         self,
@@ -107,10 +107,6 @@ def get_device(self) -> Any:
     def DataType(self) -> type[jax.Array]:  # noqa: N802
         return utils.ArrayType
 
-    @property
-    def codegen_config(self) -> dict[str, bool]:
-        return CODEGEN_CONFIG
-
     @staticmethod
     def get_available_devices() -> list[str]:
         """Static method to get a list of available devices.

mithril/backends/with_autograd/jax_backend/utils.py

@@ -19,13 +19,11 @@
 import numpy as np
 
 from .... import types
-from ....common import find_dominant_type
+from ....common import PythonGenConfig, find_dominant_type
 from ....cores.python.jax.utils import dtype_map
 from ...utils import DtypeSubTypes
 
-CODEGEN_CONFIG: dict[str, bool] = {
-    "specify_device": True,
-}
+CODEGEN_CONFIG = PythonGenConfig(SPECIFY_DEVICE=True)
 
 ArrayType = jax.Array
 

mithril/backends/with_autograd/mlx_backend/backend.py

@@ -37,6 +37,7 @@ class MlxBackend(Backend[mx.array]):
     supported_dtypes = [Dtype.float16, Dtype.bfloat16, Dtype.float32]
     registered_primitives: dict[str, Callable[..., mx.array]] = {}
     primitive_fn_path = "mithril.cores.python.mlx.ops"
+    CODEGEN_CONFIG = utils.CODEGEN_CONFIG
 
     def __init__(
         self,
@@ -75,10 +76,6 @@ def nan(self) -> float:
     def device(self) -> Any:
         utils.get_device(self._device)
 
-    @property
-    def codegen_config(self) -> dict[str, bool]:
-        return utils.CODEGEN_CONFIG
-
     def get_device(self) -> Any:
         return self._device
 

mithril/backends/with_autograd/mlx_backend/utils.py

@@ -21,13 +21,12 @@
 import numpy as np
 
 from .... import types
-from ....common import find_dominant_type
+from ....common import PythonGenConfig, find_dominant_type
 from ....cores.python.mlx.utils import dtype_map
 from ...utils import DtypeSubTypes
 
-CODEGEN_CONFIG: dict[str, bool] = {
-    "specify_device": True,
-}
+CODEGEN_CONFIG = PythonGenConfig(SPECIFY_DEVICE=True)
+
 
 ArrayType = mx.array
 

mithril/backends/with_autograd/torch_backend/backend.py

@@ -51,6 +51,7 @@ class TorchBackend(ParallelBackend[torch.Tensor]):
     backend_type = "torch"
     registered_primitives = {}
     primitive_fn_path = "mithril.cores.python.torch.ops"
+    CODEGEN_CONFIG = utils.CODEGEN_CONFIG
 
     def __init__(
         self,
@@ -93,10 +94,6 @@ def nan(self) -> float:
     def DataType(self) -> type[torch.Tensor]:  # noqa: N802
         return utils.ArrayType
 
-    @property
-    def codegen_config(self) -> dict[str, bool]:
-        return utils.CODEGEN_CONFIG
-
     @property
     def device(self) -> torch.device:
         return utils.get_device(self._device)

mithril/backends/with_autograd/torch_backend/utils.py

@@ -30,13 +30,11 @@
 from torch.distributed._tensor import DeviceMesh
 
 from .... import types
-from ....common import find_dominant_type
+from ....common import PythonGenConfig, find_dominant_type
 from ....cores.python.torch.utils import dtype_map
 from ...utils import DtypeSubTypes
 
-CODEGEN_CONFIG: dict[str, bool] = {
-    "specify_device": True,
-}
+CODEGEN_CONFIG = PythonGenConfig(SPECIFY_DEVICE=True)
 
 AVAILABLE_BACKEND_TYPES = ["cpu", "cuda"]
 

mithril/backends/with_manualgrad/c_backend/backend.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import ctypes
 import os
 from typing import Any
 
 import numpy as np
 
 from .... import types
-from ....cores.c import array
 from ....cores.c.array import PyArray
+from ....cores.c.raw_c import array
 from ...backend import Backend
 from ...utils import process_shape
 from . import utils
@@ -29,7 +30,10 @@
 
 class CBackend(Backend[PyArray]):
     backend_type = "c"
-    SRC_PATH = os.path.join(os.path.dirname(__file__), "..", "..", "..", "cores", "c")
+    SRC_PATH = os.path.join(
+        os.path.dirname(__file__), "..", "..", "..", "cores", "c", "raw_c"
+    )
+    CODEGEN_CONFIG = utils.CODEGEN_CONFIG
 
     def __init__(self) -> None:
         self._device = "cpu"
@@ -90,3 +94,6 @@ def array(
         assert dtype is None, "dtype is not supported in CBackend"
         input = input.astype(np.float32)
         return utils.from_numpy(input)
+
+    def get_struct_cls(self) -> type[ctypes.Structure]:
+        return array.Array